From 830afe4e583db79e310091b27f9f92abe33b735f Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Thu, 30 Jul 2026 00:23:58 -0700 Subject: [PATCH 001/177] [None][feat] Add DSpark confidence-head verification scheduling DSpark drafts a whole block every step and verifies all of it. The draft model already carries a confidence head that scores each drafted position's chance of surviving verification, but the head was dead code: the only consumer computed a prefix length with the wrong algorithm, synced to the host, and assumed batch size 1. This wires the head into a real verification scheduler. The block is still drafted in full; only the number of positions sent to the target changes, so acceptance -- and therefore the output distribution -- is untouched. What is here ------------ * Capture-safe confidence forward plus sequential temperature scaling. The STS buffer is written in place; rebinding it would be invisible to an already-captured graph. * Scheduler: prefix survival is a cumulative product of per-position probabilities, and the batch's budget is handed out by global survival rank with a value-independent tie-break, so the whole path is sync-free. * Planner: the verify budget comes from an argmax of tau/cost over a profiled step-cost table. The search is a global argmax rather than a greedy first descent, because the cost curve is a staircase and Theta is not unimodal on it. Without a profiled table the planner refuses to trim rather than optimize against a model that says every extra token is free. * Ragged (per-request) verification behind enable_ragged_verify: each request gets its own window, packed into a flat token axis with a qo_indptr, padded up to a captured token bucket. * Four config knobs on DSparkDecodingConfig, with the golden manifest regenerated. Three pre-existing bugs fixed along the way ------------------------------------------- * dspark.py baked in a stride of max_draft_len + 1 when reading captured target hidden states. Under any dynamic draft length that reads the wrong rows. This one is independent of this feature and would be worth taking on its own. * The overlap scheduler's KV-length correction subtracted a batch-wide token count from a per-request quantity, and the overlap branch of _prepare_tp_inputs appended a fixed number of position_ids per request. Both are silent wrong-answer paths once request lengths differ. * The incremental-update fast path reuses the previous step's attention and spec metadata, which is only valid while sequence lengths are unchanged; it now falls back to a full prepare when the verify windows move. Every change is gated so a batch without per-request windows takes the original code path unchanged: the derived totals are algebraically identical to the previous multiplicative formulas, and py_verify_len is unset on every non-DSpark path. Testing ------- 208 unit tests pass, plus the existing non-DSpark speculative tests as a regression check. Not yet verified ---------------- Accuracy, CUDA-graph and ADP acceptance all need a checkpoint with trained confidence-head weights, which is not available yet -- the model directory is still an empty placeholder. The GSM8K test added here asserts the same reference accuracy as the baseline and is the intended gate. Two specific gaps to be aware of when reviewing: * The ragged CUDA-graph path captures len(tiers) times more graphs than the uniform ladder. SGLang avoids this by keying the graph on the token count alone and deriving the row count from it, so their graph count is unchanged. Matching that means changing what the batch-size axis means in the padding machinery, which is not something to do without being able to run it. * The per-token sampling-parameter fix is only exercised when a batch is not all-greedy, so a temperature=0 accuracy run cannot cover it. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa.py | 106 ++- tensorrt_llm/_torch/models/dspark/draft.py | 46 +- tensorrt_llm/_torch/models/dspark/heads.py | 90 ++- tensorrt_llm/_torch/models/modeling_dspark.py | 28 +- .../_torch/pyexecutor/cuda_graph_runner.py | 46 ++ tensorrt_llm/_torch/pyexecutor/llm_request.py | 19 + .../_torch/pyexecutor/model_engine.py | 603 +++++++++++++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 121 ++- tensorrt_llm/_torch/speculative/dspark.py | 156 +++- .../_torch/speculative/dspark_planner.py | 288 ++++++++ .../_torch/speculative/dspark_ragged.py | 498 +++++++++++++ .../_torch/speculative/dspark_schedule.py | 189 +++++ .../_torch/speculative/dspark_verify.py | 232 ++++++ tensorrt_llm/_torch/speculative/interface.py | 152 +++- .../_torch/speculative/spec_sampler_base.py | 24 +- tensorrt_llm/llmapi/llm_args.py | 112 ++- .../usage/llm_args_golden_manifest.json | 21 + .../defs/accuracy/test_llm_api_pytorch.py | 51 ++ .../test_dspark_confidence_schedule.py | 696 ++++++++++++++++++ .../hw_agnostic/test_dspark_cuda_graph.py | 102 +++ .../hw_agnostic/test_dspark_draft.py | 84 ++- .../hw_agnostic/test_dspark_heads.py | 73 +- .../hw_agnostic/test_dspark_overlap_layout.py | 291 ++++++++ .../hw_agnostic/test_dspark_ragged.py | 676 +++++++++++++++++ .../hw_agnostic/test_dspark_worker.py | 164 ++++- tests/unittest/llmapi/test_llm_args.py | 48 ++ 26 files changed, 4683 insertions(+), 233 deletions(-) create mode 100644 tensorrt_llm/_torch/speculative/dspark_planner.py create mode 100644 tensorrt_llm/_torch/speculative/dspark_ragged.py create mode 100644 tensorrt_llm/_torch/speculative/dspark_schedule.py create mode 100644 tensorrt_llm/_torch/speculative/dspark_verify.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index a57fdadf0776..a790c6c07699 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -634,6 +634,31 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): # Number of compressed KV tokens for context requests num_ctx_kv_tokens: int = 0 gen_indexer_kv_lens_cuda_runtime: Optional[torch.Tensor] = None + # Ragged verification: how many query tokens each generation request + # contributes this step, host-side and in batch order. None (the default, + # and every non-DSpark path) means the uniform `1 + max_draft_tokens` per + # request, which is what the strided expansions below assume. Kept on the + # host because the expansions are built host-side and copied in; reading + # the repeats off a device tensor would sync. + ragged_verify_lens: Optional[List[int]] = None + + @property + def is_ragged_verify(self) -> bool: + """Whether generation requests have per-request query lengths.""" + return self.ragged_verify_lens is not None + + def gen_token_repeats(self, device=None) -> torch.Tensor: + """Per-generation-request query-token counts as a repeat vector. + + Uniform batches get `1 + max_draft_tokens` for every request, so the + callers below stay bit-identical; ragged batches get the scheduler's + per-request windows. + """ + if self.ragged_verify_lens is None: + repeats = [1 + self.max_draft_tokens] * self.num_generations + else: + repeats = self.ragged_verify_lens + return torch.tensor(repeats, dtype=torch.long, device=device) def __init__(self, *args, **kwargs): """Initialize DSA metadata with SM count and indexer chunk size.""" @@ -894,13 +919,13 @@ def on_update_kv_lens(self): self.scheduler_metadata_buffer_full_next_n.copy_( scheduler_metadata_buffer_full_next_n, non_blocking=True) if self.use_expanded_buffers_for_mtp: - num_draft_tokens = 1 + self.max_draft_tokens - num_tokens = self.num_generations * num_draft_tokens - kv_lens_expanded = torch.stack([gen_indexer_kv_lens] * - num_draft_tokens, - dim=0) + # One row per query token; the repeat count is per-request + # under ragged verification and a constant otherwise. + token_repeats = self.gen_token_repeats( + device=gen_indexer_kv_lens.device) + num_tokens = int(token_repeats.sum()) self.kv_lens_expanded_cuda[:num_tokens] = \ - kv_lens_expanded.transpose(0, 1).contiguous().flatten() + gen_indexer_kv_lens.repeat_interleave(token_repeats) scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), _DG_SCHEDULE_BLOCK_KV, self.num_sms) @@ -1475,20 +1500,28 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): # - No distinction between sm90 and sm100 is needed once MTP3 is supported on sm90. # - Remove this once fp8_paged_mqa_logits supports an arbitrary number of MTP draft tokens. use_dsl = self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits - self.use_expanded_buffers_for_mtp = (not use_dsl and ( - (self.max_draft_tokens > 1 and get_sm_version() == 90) or - ((self.max_draft_tokens == 2 or self.max_draft_tokens > 3) - and get_sm_version() >= 100))) + if self.is_ragged_verify: + # The DSL paged-MQA path reshapes (num_gen, next_n) -> (num_gen * + # factor, atom) and broadcasts a single 1-D kv_lens per request, so + # it is structurally uniform: there is no next_n that describes a + # ragged batch. The expanded path already materializes one row per + # token, which is exactly the ragged layout, so route there. + use_dsl = False + self.use_expanded_buffers_for_mtp = ( + not use_dsl + and (self.is_ragged_verify or + (self.max_draft_tokens > 1 and get_sm_version() == 90) or + ((self.max_draft_tokens == 2 or self.max_draft_tokens > 3) + and get_sm_version() >= 100))) if self.use_expanded_buffers_for_mtp: - # Expand kv_lens_cuda (only generation) - num_tokens = self.num_generations * (1 + self.max_draft_tokens) + # Expand kv_lens_cuda (only generation): one row per query token, + # repeating each request's kv_len as many times as it has tokens. + token_repeats = self.gen_token_repeats() + num_tokens = int(token_repeats.sum()) gen_kv_lens = self.get_indexer_kv_lens( kv_lens[self.num_contexts:self.num_seqs]) - gen_kv_lens_expanded = torch.stack([gen_kv_lens] * - (1 + self.max_draft_tokens), - dim=0) - gen_kv_lens_expanded = gen_kv_lens_expanded.transpose( - 0, 1).contiguous().flatten() + gen_kv_lens_expanded = gen_kv_lens.repeat_interleave( + token_repeats.to(gen_kv_lens.device)) self.kv_lens_expanded_host[:num_tokens].copy_(gen_kv_lens_expanded) self.kv_lens_expanded_cuda[:num_tokens].copy_( self.kv_lens_expanded_host[:num_tokens], non_blocking=True) @@ -1501,7 +1534,7 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): gen_block_tensor = self.host_indexer_k_cache_block_offsets[ self.num_contexts:self.num_seqs, :max_len] expanded_blocks = gen_block_tensor.repeat_interleave( - 1 + self.max_draft_tokens, dim=0) + token_repeats.to(gen_block_tensor.device), dim=0) self.host_block_table_expanded[:num_tokens, :max_len].copy_( expanded_blocks, non_blocking=True) self.block_table_expanded[:num_tokens].copy_( @@ -2709,9 +2742,14 @@ def sparse_attn_indexer( # Get decode lengths per request (from seq_lens) for validation gen_seq_lens = metadata.seq_lens[num_contexts:num_contexts + num_generations] - max_decode_len = gen_seq_lens.max().item() - min_decode_len = gen_seq_lens.min().item() - assert max_decode_len == min_decode_len, "max_decode_len != min_decode_len, we need padding" + if not metadata.is_ragged_verify: + # Ragged verification makes unequal decode lengths the whole + # point; the expanded path below handles them natively (one + # kernel row per query token), so this uniformity check only + # applies to the strided paths. + max_decode_len = gen_seq_lens.max().item() + min_decode_len = gen_seq_lens.min().item() + assert max_decode_len == min_decode_len, "max_decode_len != min_decode_len, we need padding" # Reshape q for decode phase: [num_gen_tokens, ...] -> [batch_size, next_n, ...] q_decode = q_fp8[num_ctx_tokens:num_ctx_tokens + num_gen_tokens, @@ -2721,7 +2759,14 @@ def sparse_attn_indexer( # Because fp8_paged_mqa_logits can only support next_n == 1/2/4 on sm100, and # next_n == 1/2 on sm90, for other next_n, we need to flatten the q_decode tensor # and expand the corresponding metadata. - if not metadata.use_expanded_buffers_for_mtp or next_n == 1: + # Ragged batches have no single next_n, so they always take the + # expanded (one row per token) branch. + if metadata.is_ragged_verify: + take_strided_path = False + else: + take_strided_path = (not metadata.use_expanded_buffers_for_mtp + or next_n == 1) + if take_strided_path: q_decode = q_decode.view(num_generations, -1, *q_fp8.shape[1:]) # 2D context_lens slice from the pre-allocated buffer; matches # q_decode's (batch, next_n) layout required by the new @@ -2750,7 +2795,15 @@ def sparse_attn_indexer( block_table = metadata.block_table_expanded[:num_tokens] scheduler_metadata_buffer = metadata.scheduler_metadata_buffer_expanded - assert num_gen_tokens == batch_size * next_n + # Ragged verification packs a different number of query tokens per + # request, so the only invariant left is that the expanded buffers + # cover exactly the tokens handed to the kernel. + if metadata.is_ragged_verify: + assert num_gen_tokens == sum(metadata.ragged_verify_lens), ( + f"ragged gen tokens {num_gen_tokens} != " + f"sum(verify_lens) {sum(metadata.ragged_verify_lens)}") + else: + assert num_gen_tokens == batch_size * next_n weights_decode = weights[num_ctx_tokens:num_ctx_tokens + num_gen_tokens, ...] @@ -2761,7 +2814,12 @@ def sparse_attn_indexer( self.layer_idx) indexer_max_seq_len = metadata.get_indexer_max_seq_len() - if self.use_cute_dsl_paged_mqa_logits: + # The DSL path broadcasts one kv_len per request across all next_n + # positions and requires factor * atom == next_n, neither of which + # a ragged batch can satisfy. prepare_for_spec_decode already + # routed such batches to the expanded buffers; keep the forward + # consistent with that decision. + if self.use_cute_dsl_paged_mqa_logits and not metadata.is_ragged_verify: # DSL kernel design: 1 atom per q (atom = real next_n positions), # kNumNextNAtoms = 1 for any real next_n. The matching schedule # is `scheduler_metadata_buffer` — built in `Indexer.prepare()` diff --git a/tensorrt_llm/_torch/models/dspark/draft.py b/tensorrt_llm/_torch/models/dspark/draft.py index 1b47f4922965..8aa3df8016b5 100644 --- a/tensorrt_llm/_torch/models/dspark/draft.py +++ b/tensorrt_llm/_torch/models/dspark/draft.py @@ -24,7 +24,7 @@ - ``build_draft_input_ids``: ``[bonus_token, noise, noise, ...]`` block input. - ``dspark_propose``: given the per-position backbone ``base_logits`` and the Markov / confidence heads, run the autoregressive Markov refinement to sample - the block tokens and apply the static confidence-threshold truncation. + the block tokens and score them with the confidence head. The backbone (3 V4 blocks producing ``block_hidden``) lives in the model module; this file is the part fully specified by the reference and validated against it. @@ -35,8 +35,6 @@ import torch from torch import nn -from .heads import confident_prefix_length - def build_draft_input_ids( bonus_token_ids: torch.Tensor, *, block_size: int, noise_token_id: int @@ -61,25 +59,32 @@ def dspark_propose( confidence_head: Optional[nn.Module], block_size: int, temperature: float = 0.0, - confidence_threshold: float = 0.0, + return_confidence: bool = False, return_logits: bool = False, ) -> tuple: """Produce DSpark draft tokens for one block (functional-first, static length). + Always proposes the *full* block: the confidence head does not shorten the + draft (the block is produced by a single parallel backbone pass, so there is + nothing to save there). It scores the block so the verification scheduler can + decide how many of those tokens are worth sending to the target. + Args: base_logits: ``[batch, block_size, vocab]`` from the backbone + lm_head. bonus_token_ids: ``[batch]`` the token preceding the first draft position. block_hidden: ``[batch, block_size, hidden]`` backbone hidden (feeds the confidence head, and the RNN-head variant). markov_head / confidence_head: the validated DSpark heads (may be None). + return_confidence: compute the per-position confidence logits. This is a + run-constant flag (read once from the decoding config), never a + per-step decision, so it cannot make the captured graph diverge. Returns: draft_tokens: ``[batch, block_size]`` sampled tokens (full block; callers keep the tensor fixed-width for CUDA-graph safety). - num_proposed: ``[batch]`` int32 — how many leading tokens survive the - static confidence-threshold truncation (== block_size when no head / - threshold<=0). + confidence: ``[batch, block_size]`` fp32 *raw* confidence logits, or None + when disabled / no head. Calibrate with ``confidence_head.apply_sts`` + before taking the cumulative product. """ - batch = base_logits.shape[0] # ``draft_logits`` are the per-position distributions the draft token is drawn # from (markov-corrected when a head is present, else the raw base logits). # Surfaced under ``return_logits`` for the §7.9 probabilistic-acceptance @@ -98,15 +103,10 @@ def dspark_propose( draft_tokens = greedy_or_sample(base_logits, temperature) - # Scaffolding: confidence-based dynamic drafting is NOT enabled in this PR. - # The worker always calls with confidence_threshold=0.0, so the block below is - # inert and num_proposed stays == block_size (the full block is proposed). The - # returned num_proposed is intentionally not yet consumed by the speculative - # scheduler/verifier; wiring it through is a follow-up (see PR description). - num_proposed = torch.full( - (batch,), int(block_size), dtype=torch.int32, device=base_logits.device - ) - if confidence_head is not None and confidence_threshold > 0.0: + # Confidence scoring: branch-free and fixed-shape (no ``.item()``, no + # data-dependent shapes), so the whole block stays CUDA-graph capturable. + confidence = None + if return_confidence and confidence_head is not None: # prev token at position k is [bonus, draft_0, ..., draft_{k-1}] prev_ids = torch.cat([bonus_token_ids.unsqueeze(1), draft_tokens[:, :-1]], dim=1) prev_emb = ( @@ -114,17 +114,11 @@ def dspark_propose( if (markov_head is not None and getattr(confidence_head, "with_markov", False)) else None ) - conf_logits = ( + confidence = ( confidence_head(block_hidden, prev_embeddings=prev_emb) if prev_emb is not None else confidence_head(block_hidden) ) - # Per-request prefix truncation (batch handled row-wise to stay simple; - # functional-first scope typically runs batch=1 for the draft). - for b in range(batch): - num_proposed[b] = confident_prefix_length( - conf_logits[b : b + 1], block_size=block_size, threshold=confidence_threshold - ) if return_logits: - return draft_tokens, num_proposed, draft_logits - return draft_tokens, num_proposed + return draft_tokens, confidence, draft_logits + return draft_tokens, confidence diff --git a/tensorrt_llm/_torch/models/dspark/heads.py b/tensorrt_llm/_torch/models/dspark/heads.py index c49e35fbafa0..94529e4288f7 100644 --- a/tensorrt_llm/_torch/models/dspark/heads.py +++ b/tensorrt_llm/_torch/models/dspark/heads.py @@ -24,9 +24,11 @@ applied autoregressively across the ``block_size`` draft positions (the cheap "sequential" half of DSpark's "semi-parallel" drafting). RNN variant carries a GRU-style recurrent state across positions. - - Confidence head: predicts a per-position acceptance probability; the cumulative - product over positions estimates prefix-acceptance and is used only to - *truncate* the proposed draft length (NOT to decide acceptance). + - Confidence head: predicts a per-position *conditional* acceptance probability + (``P(accept_k | accept_1..k-1)``); the cumulative product over positions is the + prefix-survival probability the verification scheduler budgets against. It only + decides *how many* tokens are verified, never *whether* one is accepted, so + scheduling stays lossless. See ``_torch/speculative/dspark_schedule.py``. This file deliberately depends on ``torch`` only so it can be unit-tested in isolation (token-for-token) against the DeepSpec reference. @@ -214,21 +216,49 @@ class DSparkConfidenceHead(nn.Module): """Per-position acceptance-confidence predictor (DeepSpec AcceptRatePredictor). Input features are the backbone hidden state, optionally concatenated with the - Markov head's previous-token embedding. Output is a single logit per position. + Markov head's previous-token embedding. Output is a single *raw logit* per + position; callers turn it into a probability with :meth:`apply_sts`, which + folds in the sequential-temperature-scaling (STS) calibration. + + Calibration matters here because the scheduler consumes the *cumulative + product* of per-position probabilities: a per-position bias is amplified + geometrically along the block, so an uncalibrated head systematically + over-estimates prefix survival. ``sts_temperatures`` defaults to all-ones, + which makes :meth:`apply_sts` a plain sigmoid (identity calibration). """ - def __init__(self, *, hidden_size: int, markov_rank: int = 0, with_markov: bool = False): + def __init__( + self, + *, + hidden_size: int, + markov_rank: int = 0, + with_markov: bool = False, + bias: bool = False, + block_size: int = 0, + ): super().__init__() self.with_markov = bool(with_markov) input_dim = int(hidden_size) + (int(markov_rank) if with_markov else 0) # The checkpoint stores ``proj`` as a bias-free bf16 weight, but the # confidence score is computed in fp32 (mirrors the DeepSpec reference # ``Linear(input_dim, 1, dtype=torch.float32)`` with the fp32 matmul). - self.proj = nn.Linear(input_dim, 1, bias=False, dtype=torch.float32) + self.proj = nn.Linear(input_dim, 1, bias=bool(bias), dtype=torch.float32) + # Per-position temperatures, broadcast against ``[batch, block]``. Shape + # ``(block_size,)`` when known, else ``(1,)``; both broadcast correctly. + # ``persistent=False`` keeps it out of ``state_dict`` so it never + # participates in checkpoint loading. It is sized at construction and + # only ever updated in place, so a captured CUDA graph keeps seeing the + # same storage (rebinding the attribute would be invisible to the graph). + self.register_buffer( + "sts_temperatures", + torch.ones(max(int(block_size), 1), dtype=torch.float32), + persistent=False, + ) def forward( self, hidden_states: torch.Tensor, prev_embeddings: Optional[torch.Tensor] = None ) -> torch.Tensor: + """``[*, block, hidden] -> [*, block]`` raw (uncalibrated) logits.""" if self.with_markov: assert prev_embeddings is not None features = torch.cat([hidden_states, prev_embeddings.to(hidden_states.dtype)], dim=-1) @@ -237,18 +267,40 @@ def forward( # fp32 matmul for a stable confidence score (mirrors the reference). return self.proj(features.float()).squeeze(-1) + def apply_sts(self, confidence_logits: torch.Tensor) -> torch.Tensor: + """Raw logits -> calibrated per-position acceptance probabilities in [0, 1].""" + return torch.sigmoid(confidence_logits.float() / self.sts_temperatures) + + @torch.no_grad() + def load_sts_temperatures(self, temperatures: torch.Tensor) -> None: + """In-place update of the STS table (never rebind: see ``__init__``).""" + flat = temperatures.reshape(-1).to(device=self.sts_temperatures.device, dtype=torch.float32) + if flat.numel() != self.sts_temperatures.numel(): + raise ValueError( + f"STS temperature table has {flat.numel()} entries but the confidence " + f"head expects {self.sts_temperatures.numel()} (one per block position)" + ) + if not bool(torch.all(flat > 0.0)): + raise ValueError("STS temperatures must be strictly positive") + self.sts_temperatures.copy_(flat) -def confident_prefix_length( - confidence_logits: torch.Tensor, *, block_size: int, threshold: float -) -> int: - """First position k where ``sigmoid(confidence_k) < threshold``. + def load_weights(self, weights: list) -> None: + """Strict loader that refuses to silently drop a checkpoint bias. - Returns ``block_size`` when threshold<=0 (no truncation) or all positions - are confident. Assumes batch size 1 (functional-first scope). - """ - if threshold <= 0.0: - return int(block_size) - below = confidence_logits.sigmoid() < threshold - if not bool(below[0].any().item()): - return int(block_size) - return int(torch.nonzero(below[0], as_tuple=False)[0].item()) + The generic DeepseekV4 loader copies by iterating ``named_parameters()``, + so a ``proj.bias`` present in the checkpoint but absent from the module + would be dropped without any error -- shifting every confidence score. + Fail loudly instead. + """ + (module_weights,) = weights + expected = dict(self.named_parameters()) + unexpected = [k for k in module_weights if k not in expected] + if unexpected: + raise ValueError( + f"DSpark confidence head checkpoint has unsupported key(s) {sorted(unexpected)}; " + f"the module exposes {sorted(expected)}. A checkpoint bias requires constructing " + f"DSparkConfidenceHead(bias=True) -- silently dropping it would bias every " + f"per-position confidence score." + ) + for name, param in expected.items(): + param.data.copy_(module_weights[name][:]) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 4efadd01784d..60761ee6003a 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -171,6 +171,7 @@ def __init__( stage_id: int, num_stages: int, num_capture_layers: int, + block_size: int = 0, ): # The inherited attention uses a draft-local layer index, while the # decoder layer keeps its model-level index for weights and captures. @@ -236,6 +237,8 @@ def __init__( # markov_rank <= 0); otherwise dspark_propose passes no # prev_embeddings and DSparkConfidenceHead.forward would assert. with_markov=self.markov_rank > 0, + # Sizes the per-position STS calibration table once, up front. + block_size=int(block_size), ) @property @@ -326,6 +329,7 @@ def __init__( stage_id=s, num_stages=self.num_stages, num_capture_layers=self.num_capture_layers, + block_size=self.block_size, ) for s in range(self.num_stages) ] @@ -857,7 +861,7 @@ def forward( *, kv_windows: Optional[torch.Tensor] = None, temperature: float = 0.0, - confidence_threshold: float = 0.0, + return_confidence: bool = False, return_logits: bool = False, all_rank_num_tokens: Optional[List[int]] = None, ) -> tuple: @@ -877,7 +881,7 @@ def forward( worker; updated in place each call. ``None`` allocates fresh zero windows (single-shot golden / test path). Returns: - ``(draft_tokens [T, block], num_proposed [T])`` from ``forward_head``. + ``(draft_tokens [T, block], confidence [T, block] or None)``. """ assert start_pos > 0, "DSpark draft runs at generation (start_pos > 0)" if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: @@ -908,7 +912,7 @@ def forward( h, bonus_token_ids, temperature=temperature, - confidence_threshold=confidence_threshold, + return_confidence=return_confidence, return_logits=return_logits, ) @@ -921,7 +925,7 @@ def forward_batched( kv_windows: torch.Tensor, slots: torch.Tensor, temperature: float = 0.0, - confidence_threshold: float = 0.0, + return_confidence: bool = False, return_logits: bool = False, all_rank_num_tokens: Optional[List[int]] = None, ) -> tuple: @@ -933,9 +937,9 @@ def forward_batched( graph). ``start_pos`` is a ``[G]`` tensor (one absolute decode position per gen request); the rolling captured-context windows are written/read through ``slots`` into the worker-owned ``kv_windows`` buffer; RoPE phases are - gathered from a fixed table. ``forward_head`` is run with - ``confidence_threshold == 0`` (the worker proposes the full block), which is - the graph-safe branch of :func:`dspark_propose`. + gathered from a fixed table. ``forward_head`` always proposes the full + block; ``return_confidence`` only adds a fixed-shape scoring matmul, so + both settings stay capturable. Args: main_hidden: ``[G, num_capture * hidden]`` captured target context. @@ -945,7 +949,7 @@ def forward_batched( windows; written in place through ``slots``. slots: ``[G]`` int tensor mapping each request to its ``kv_windows`` row. Returns: - ``(draft_tokens [G, block], num_proposed [G])`` from ``forward_head``. + ``(draft_tokens [G, block], confidence [G, block] or None)``. """ if getattr(self.mtp_layers[0], "_dspark_attn", None) is None: raise RuntimeError( @@ -976,7 +980,7 @@ def forward_batched( h, bonus_token_ids, temperature=temperature, - confidence_threshold=confidence_threshold, + return_confidence=return_confidence, return_logits=return_logits, ) @@ -1025,13 +1029,13 @@ def forward_head( bonus_token_ids: torch.Tensor, *, temperature: float = 0.0, - confidence_threshold: float = 0.0, + return_confidence: bool = False, return_logits: bool = False, ) -> tuple: """Block-draft head: hc_head + norm + lm_head -> markov refine + confidence. ``block_hidden`` is the last stage's mHC residual ``[*, block, hc_mult, hidden]``. - Returns (draft_tokens [*, block], num_proposed [*]); with ``return_logits`` + Returns (draft_tokens [*, block], confidence [*, block] or None); with ``return_logits`` also returns the per-position draft logits [*, block, vocab] (§7.9 1-TV). """ last = self.mtp_layers[-1] @@ -1046,7 +1050,7 @@ def forward_head( confidence_head=last.confidence_head, block_size=self.block_size, temperature=temperature, - confidence_threshold=confidence_threshold, + return_confidence=return_confidence, return_logits=return_logits, ) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index c2a6a5bb4de1..5a9a9c6e5131 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -135,6 +135,10 @@ def __init__(self, config: CUDAGraphRunnerConfig): self.spec_config = config.spec_config self.sparse_config = config.sparse_attention_config self.is_encoder_decoder = config.is_encoder_decoder + # Window each ragged padding row carries, in py_verify_len units. Set + # by ModelEngine.fit_ragged_verify_lens when it picks the token bucket; + # the padding rows are one shared dummy object so they all share it. + self.ragged_pad_verify_len = 0 self.graphs: Dict[KeyType, torch.cuda.CUDAGraph] = {} self.graph_outputs: Dict[KeyType, @@ -265,8 +269,38 @@ def get_graph_key( set(draft_len_list)) == 1, "All draft lengths must be the same" key = (batch_size, draft_len, False, short_seq_len_mode, is_all_greedy_sample) + # Ragged verification breaks the link between the row count and the + # token count: the block is still drafted in full (so draft_len is + # unchanged) but the number of positions actually verified varies, + # and that flat token count is baked into the captured graph. Add it + # as a second axis. The element is absent for every uniform batch, + # so those keys stay exactly what they were. + # + # Derived from the batch rather than plumbed in, because this runs + # inside pad_batch: the padding rows are already present and must be + # counted, and the scheduler has already fitted the total onto a + # captured bucket (see ModelEngine.fit_ragged_verify_lens). + verify_bucket = self._ragged_verify_bucket(batch) + if verify_bucket is not None: + key = key + (verify_bucket, ) return key + @staticmethod + def _ragged_verify_bucket(batch: ScheduledRequests) -> Optional[int]: + """Total verified token count when the batch uses ragged windows. + + None unless *every* generation request carries a window, so a batch the + scheduler declined to make ragged -- or any non-DSpark batch -- keeps + the original key. + """ + total = 0 + for request in batch.generation_requests: + verify_len = getattr(request, "py_verify_len", None) + if verify_len is None: + return None + total += 1 + int(verify_len) + return total if total else None + @staticmethod def _get_mrope_position_delta(request: Any) -> Optional[Any]: mrope_position_delta = getattr(request, "py_mrope_position_delta", None) @@ -615,6 +649,18 @@ def _get_padded_batch(self, batch: ScheduledRequests, self.padding_dummy_requests[runtime_draft_len] = dummy_request padding_dummy_request = self.padding_dummy_requests[runtime_draft_len] + # Under ragged verification the padding rows are part of the token + # budget the scheduler fitted onto a captured bucket. The scheduler + # decided their (uniform) window when it chose the bucket -- it has to, + # because these rows are one shared object and cannot carry different + # values -- so read it back rather than assume. Cleared otherwise: the + # dummy is cached across steps and would carry a stale window into a + # uniform batch. + ragged = any( + getattr(request, "py_verify_len", None) is not None + for request in batch.generation_requests) + padding_dummy_request.py_verify_len = (self.ragged_pad_verify_len + if ragged else None) batch.generation_requests.extend([padding_dummy_request] * padding_size) return padding_size diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 2b1b63725fb8..469384fa3713 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1265,3 +1265,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/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 694be1a14a74..193199c1526a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -67,6 +67,7 @@ restore_attn_metadata_after_draft_replay, update_spec_config_from_loaded_model) from ..speculative.drafting_loops import BaseDraftingLoopWrapper +from ..speculative.dspark_ragged import ragged_gather_index_lists from ..speculative.eagle3 import Eagle3ResourceManager, Eagle3SpecMetadata from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..utils import (get_model_extra_attrs, @@ -81,7 +82,8 @@ from .kv_cache_manager_v2 import KVCacheManagerV2 from .layerwise_nvtx_marker import LayerwiseNvtxMarker from .llm_request import (LlmRequest, LlmRequestState, get_draft_token_length, - get_multimodal_embedding_lengths) + get_multimodal_embedding_lengths, + get_request_tokens_per_gen_step) from .mamba_cache_manager import MambaHybridCacheManager from .model_loader import ModelLoader, _construct_checkpoint_loader from .resource_manager import (BaseResourceManager, KVCacheManager, @@ -599,6 +601,16 @@ def __init__( self.is_warmup = False self.previous_request_ids = [] + # Per-request verify windows seen by the last full _prepare_tp_inputs + # pass. All None without ragged verification, so the incremental-update + # gate below never trips for uniform speculation. + self.previous_verify_lens = [] + # Captured token bucket this step's ragged batch was fitted onto, or + # None when the batch is uniform. Set once per step by the scheduler. + self.ragged_verify_bucket: Optional[int] = None + # Window every CUDA-graph padding row carries this step (they share one + # dummy object, so they all get the same one). In py_verify_len units. + self.ragged_pad_verify_len: int = 0 self.has_previous_device_draft = False self.previous_accepted_tokens_cuda = torch.empty((self.batch_size, ), dtype=torch.int, @@ -697,6 +709,17 @@ def __init__( self.previous_batch_indices_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') + # Ragged verification windows. These are read *inside* the captured + # graph (one-engine acceptance runs in the target's graph), so they must + # live at a stable address and be written in place -- rebinding + # spec_metadata.verify_lens to a freshly allocated tensor each step + # would leave the graph replaying against the first step's buffer. + self.ragged_verify_lens_cuda = torch.empty((self.batch_size + 1, ), + dtype=torch.int, + device='cuda') + self.ragged_qo_indptr_cuda = torch.empty((self.batch_size + 2, ), + dtype=torch.int, + device='cuda') self.input_ids_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') @@ -1794,6 +1817,46 @@ def _get_graphs_to_capture( f"Capturing {len(graphs)} graphs: {graphs}") return graphs + # Case 2b: DSpark confidence-scheduled verification. + # Unlike Case 2, the draft length here is NOT a function of the batch + # size -- at a fixed batch size the scheduler may still pick a different + # length as the batch's confidence changes. So the ladder is captured for + # every batch size, as a cross product. + # + # This is the expensive part of the feature and the reason the ladder is + # kept short: each captured graph costs ~10-23 MB of metadata outside the + # shared pool, and that memory is taken directly out of KV cache (the KV + # pool is sized from what is left after capture). + if (self.spec_config is not None and getattr( + self.spec_config, "enable_confidence_scheduling", False)): + tiers = sorted({ + int(t) + for t in self.spec_config.verify_len_tiers if int(t) >= 0 + }) + if getattr(self.spec_config, "enable_ragged_verify", False): + # Ragged always drafts the full block -- only verification is + # trimmed -- so draft_len is pinned to the top tier and the + # ladder moves to the *token* axis: one bucket per tier, i.e. + # exactly the token totals the uniform ladder would produce. + # Same graph count, different second axis. + top_tier = tiers[-1] + graphs = [(bs, top_tier, bs * (t + 1)) + for bs in cuda_graph_batch_sizes for t in tiers] + logger.info( + f"DSpark ragged verification: capturing {len(graphs)} graphs " + f"({len(cuda_graph_batch_sizes)} batch sizes x {len(tiers)} " + f"token buckets from tiers {tiers}; draft_len pinned to " + f"{top_tier}).") + return graphs + graphs = [(bs, draft_len) for bs in cuda_graph_batch_sizes + for draft_len in tiers] + logger.info( + f"DSpark confidence scheduling: capturing {len(graphs)} graphs " + f"({len(cuda_graph_batch_sizes)} batch sizes x {len(tiers)} " + f"draft-length tiers {tiers}). Every extra tier costs captured " + f"graph memory that would otherwise be KV cache.") + return graphs + # Case 3: Target model (two-model) or one-model without dynamic draft # Match the runtime_draft_len semantics enforced in _prepare_tp_inputs: # logical K for linear-tree modes, total tree tokens for tree decoding. @@ -1949,7 +2012,11 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: # below will keep it False on every subsequent iteration. spec_metadata.is_all_greedy_sample = False try: - for bs, draft_len in graphs_to_capture: + for entry in graphs_to_capture: + # Ragged verification adds a token-count axis, so entries + # may be (bs, draft_len) or (bs, draft_len, token_bucket). + bs, draft_len = entry[0], entry[1] + verify_bucket = entry[2] if len(entry) > 2 else None if bs > self.batch_size: continue @@ -1958,6 +2025,9 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: resource_manager, bs, draft_len, max_seq_len) with self._release_batch_context( warmup_request, resource_manager) as batch: + if batch is not None and verify_bucket is not None: + self._set_warmup_ragged_windows( + batch, verify_bucket, draft_len) if batch is None: # No KV cache space for this batch size. During KV # cache estimation this makes the profiling peak @@ -3265,6 +3335,18 @@ def _can_use_incremental_update( if self.previous_request_ids != request_ids: return False + # This path reuses the attention/spec metadata built by the last full + # prepare, which is only valid while every request's sequence length is + # unchanged. Uniform speculation guarantees that; ragged verification + # re-picks each window every step, so fall back to a full prepare + # whenever the windows moved. + verify_lens = [ + getattr(request, "py_verify_len", None) + for request in scheduled_requests.generation_requests + ] + if self.previous_verify_lens != verify_lens: + return False + has_current_device_draft = next_draft_tokens_device is not None return (self.is_draft_model and self.model_is_wrapped) or ( has_current_device_draft and self.has_previous_device_draft) @@ -3503,18 +3585,269 @@ def _apply_incremental_update_draft( return inputs, self.gather_ids_cuda[:num_generation_tokens] + def _set_warmup_ragged_windows(self, batch, verify_bucket: int, + draft_len: int) -> None: + """Give a capture batch per-request windows summing to ``verify_bucket``. + + A captured ragged graph is keyed on its token total, so the warmup batch + has to hit that total exactly -- otherwise the graph is captured under a + key no runtime batch will ever produce, and every ragged step silently + falls back to eager. Any split with the right sum captures the right + shape (the raggedness lives in the *contents* of qo_indptr, not the + shape), so reuse the same spreading rule the runtime fit uses. + """ + requests = batch.generation_requests + if not requests: + return + from ..speculative.dspark_ragged import RaggedVerifyLayout + + lens = torch.ones(len(requests), dtype=torch.int32) + layout = RaggedVerifyLayout.from_verify_lens( + lens, + graph_num_tokens=int(verify_bucket), + total_verify_tokens=len(requests)) + filled = layout.fill_bucket(max_verify_len=1 + int(draft_len)) + for request, tokens in zip(requests, filled.verify_lens.tolist()): + request.py_verify_len = int(tokens) - 1 + + def ragged_verify_token_buckets(self, padded_bs: int) -> List[int]: + """Captured token totals for a ragged batch of ``padded_bs`` rows. + + One bucket per verify-length tier, i.e. exactly the token counts the + uniform ladder would have produced at this batch size. A ragged batch + rounds its own total up to one of these, so the captured graph set does + not grow beyond the uniform one on the token axis. + """ + if self.spec_config is None: + return [] + tiers = sorted({ + int(t) + for t in getattr(self.spec_config, "verify_len_tiers", None) or [] + }) + return [int(padded_bs) * (t + 1) for t in tiers] + + def fit_ragged_verify_lens( + self, + generation_requests, + verify_lens: List[int], + peer_stats: Optional[List[List[int]]] = None) -> Optional[int]: + """Round this batch's verify total onto a captured bucket, in place. + + The graph key gains a token axis under ragged verification: two batches + with the same row count and the same maximum draft length can still + carry very different token totals, and the flat token count is baked + into the captured graph. So the total has to land on a captured value + *before* the key is built, and the slack is spent on real verification + rather than thrown away (see ``RaggedVerifyLayout.fill_bucket``). + + Writes the fitted window back onto each request's ``py_verify_len`` and + returns the bucket, or None when no captured shape fits -- in which + case the caller should fall back to uniform scheduling rather than drop + out of graph replay. + + ``peer_stats`` is every attention-DP rank's ``[num_requests, + total_tokens]``, including this rank's. It is required under ADP + because ``_get_padded_batch`` pads to the *maximum* batch size across + ranks: sizing the bucket from the local count alone makes each rank + pick a different shape, which is the graph-key divergence this whole + path exists to avoid. + """ + runner = self.cuda_graph_runner + if not runner.enabled or not verify_lens: + return None + tiers = sorted({ + int(t) + for t in getattr(self.spec_config, "verify_len_tiers", None) or [] + }) + if not tiers: + return None + max_verify_len = 1 + tiers[-1] + token_lens = [1 + int(v) for v in verify_lens] + + from ..speculative.dspark_ragged import (RaggedVerifyLayout, + choose_ragged_capture_shape) + + bs_buckets = sorted({int(b) for b in runner.supported_batch_sizes}) + if not bs_buckets: + return None + # The token bucket grid depends on the row count, which itself depends + # on the widest rank, so resolve the rows first and then the tokens. + widest_rows = max( + [int(s[0]) for s in peer_stats], + default=len(token_lens)) if peer_stats else len(token_lens) + padded_bs = runner._round_up_batch_size(widest_rows) + if padded_bs == 0: + return None + buckets = self.ragged_verify_token_buckets(padded_bs) + if not buckets: + return None + + try: + shape = choose_ragged_capture_shape( + num_real_requests=len(token_lens), + total_verify_tokens=sum(token_lens), + bs_buckets=bs_buckets, + token_buckets=buckets, + peer_stats=peer_stats) + except ValueError as exc: + logger.debug(f"DSpark ragged shape selection failed, falling back " + f"to uniform scheduling: {exc}") + return None + padded_bs, bucket = shape.padded_bs, shape.bucket + + # The CUDA-graph padding rows are appended later by _get_padded_batch, + # which appends a *single shared dummy object*, so every pad row + # necessarily carries the same window. Their contribution therefore has + # to be decided here, not left to the fill: a total the padding step + # then contradicts would key into a graph that was never captured. + # + # Pad rows take the minimum (one token) so real requests get the slack, + # and only grow when the real rows cannot absorb the bucket on their + # own -- SGLang spreads the leftover across pad rows for the same + # reason, just at per-row granularity it can afford because it builds + # the pad block on device. + n_real = len(token_lens) + n_pad = padded_bs - n_real + real_capacity = n_real * max_verify_len + floor_tokens = sum(token_lens) + if n_pad == 0: + pad_len = 1 + else: + # pad_len has to leave the real rows a target they can actually + # hit: at least what they already need, at most their capacity. + # bucket - n_pad*pad_len >= floor_tokens -> upper bound + # bucket - n_pad*pad_len <= real_capacity -> lower bound + lo = max(1, -(-(bucket - real_capacity) // n_pad)) # ceil division + hi = min(max_verify_len, (bucket - floor_tokens) // n_pad) + if lo > hi: + logger.debug( + f"DSpark ragged: bucket {bucket} admits no pad-row window " + f"for {n_real} real ({floor_tokens}..{real_capacity} " + f"tokens) + {n_pad} pad rows; falling back to uniform " + f"scheduling") + return None + pad_len = lo + real_target = bucket - n_pad * pad_len + if real_target < floor_tokens or real_target > real_capacity: + logger.debug( + f"DSpark ragged: bucket {bucket} leaves {real_target} tokens " + f"for {n_real} real requests ({n_pad} pad rows at {pad_len}), " + f"outside [{floor_tokens}, {real_capacity}]; falling back " + f"to uniform scheduling") + return None + # Communicated to _get_padded_batch, which stamps it on the dummy. + runner.ragged_pad_verify_len = pad_len - 1 + + lens_tensor = torch.tensor(token_lens, dtype=torch.int32) + layout = RaggedVerifyLayout.from_verify_lens( + lens_tensor, + graph_num_tokens=real_target, + total_verify_tokens=sum(token_lens)) + try: + filled = layout.fill_bucket(max_verify_len=max_verify_len) + except ValueError as exc: + # The guards are about caller consistency, not impossibility; a + # mismatch means this batch has no captured shape, so degrade to + # uniform rather than run eager. + logger.debug(f"DSpark ragged bucket fit failed, falling back to " + f"uniform scheduling: {exc}") + return None + + for request, tokens in zip(generation_requests, + filled.verify_lens.tolist()): + request.py_verify_len = int(tokens) - 1 + return int(bucket) + + def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, + generation_requests) -> None: + """Publish this step's per-request verify windows to both metadatas. + + Only the DSpark confidence scheduler sets ``py_verify_len``; every other + path leaves the fields None and the uniform strides downstream stay + live. Both metadatas need it and for different reasons: acceptance has + to know each request's window to stop its cumulative product at the + right place, and the DSA indexer has to expand kv_lens/block tables one + row per query token instead of a fixed count per request. + """ + verify_lens = [ + getattr(request, "py_verify_len", None) + for request in generation_requests + ] + if not verify_lens or any(v is None for v in verify_lens): + spec_metadata.verify_lens = None + spec_metadata.qo_indptr = None + spec_metadata.total_verify_tokens = None + if attn_metadata is not None and hasattr(attn_metadata, + "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = None + return + + from ..speculative.dspark_ragged import build_qo_indptr + + # py_verify_len counts draft positions; a request also verifies its + # bonus position, so its token window is one wider. + token_lens = [1 + int(v) for v in verify_lens] + n = len(token_lens) + # Write through the persistent buffers rather than binding a fresh + # tensor: a captured graph baked in the address it saw at capture time, + # so a new allocation here would be invisible to every replay. + lens_view = self.ragged_verify_lens_cuda[:n] + lens_view.copy_(torch.tensor(token_lens, + dtype=torch.int32, + pin_memory=prefer_pinned()), + non_blocking=True) + indptr_view = self.ragged_qo_indptr_cuda[:n + 1] + indptr_view.copy_(build_qo_indptr(lens_view), non_blocking=True) + spec_metadata.verify_lens = lens_view + spec_metadata.qo_indptr = indptr_view + spec_metadata.total_verify_tokens = sum(token_lens) + if attn_metadata is not None and hasattr(attn_metadata, + "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = token_lens + + def _ragged_gather_indices( + self, slots: List[int], + counts: List[int]) -> Tuple[torch.Tensor, torch.Tensor]: + """Row/column index pair for gathering a ragged block out of a + [num_slots, max_width] device tensor. + + ``slots[i]`` contributes ``counts[i]`` entries taken from columns + ``0..counts[i]-1``, concatenated in batch order. This replaces the + ``tensor[slots, :width]`` strided gather used when every request + verifies the same number of positions. + """ + rows, cols = ragged_gather_index_lists(slots, counts) + rows_host = torch.tensor(rows, + dtype=torch.long, + pin_memory=prefer_pinned()) + cols_host = torch.tensor(cols, + dtype=torch.long, + pin_memory=prefer_pinned()) + return (rows_host.to('cuda', non_blocking=True), + cols_host.to('cuda', non_blocking=True)) + def _update_target_input_tensors( - self, num_accepted_tokens_device: torch.Tensor, + self, + num_accepted_tokens_device: torch.Tensor, new_tokens_device: torch.Tensor, next_draft_tokens_device: torch.Tensor, - new_tokens_lens_device: torch.Tensor, previous_slots: torch.Tensor, - total_num_tokens: int, num_extend_reqeust_wo_dummy: int, + new_tokens_lens_device: torch.Tensor, + previous_slots: torch.Tensor, + total_num_tokens: int, + num_extend_reqeust_wo_dummy: int, num_tokens_per_extend_request: int, - previous_batch_draft_tokens: int): + previous_batch_draft_tokens: int, + tokens_per_extend_request: Optional[List[int]] = None, + previous_batch_slots: Optional[List[int]] = None): """ This function performs in-place updates on position_ids, num_accepted_draft_tokens, input_ids, draft_tokens, and offset tensors for speculative decoding extend context operations. + + ``tokens_per_extend_request`` carries the per-request token counts under + ragged verification; it is None (and every gather below stays strided + on ``num_tokens_per_extend_request``) whenever the batch is uniform. """ + is_ragged = tokens_per_extend_request is not None # Prepare position_ids idx_accepted_tokens = self.idx_accepted_tokens_cache[:total_num_tokens] @@ -3533,33 +3866,62 @@ def _update_target_input_tensors( # CRITICAL: Only extract the needed tokens based on num_tokens_per_extend_request # new_tokens_device shape: [batch, 1 + max_draft_len] # We need: [previous_batch, num_tokens_per_extend_request] - new_tokens = new_tokens_device.transpose( - 0, 1)[previous_slots, :num_tokens_per_extend_request].flatten() + if is_ragged: + token_rows, token_cols = self._ragged_gather_indices( + previous_batch_slots, tokens_per_extend_request) + new_tokens = new_tokens_device.transpose(0, + 1)[token_rows, + token_cols].flatten() + else: + new_tokens = new_tokens_device.transpose( + 0, 1)[previous_slots, :num_tokens_per_extend_request].flatten() self.input_ids_cuda[:total_num_tokens].copy_(new_tokens, non_blocking=True) # Prepare draft tokens - num_draft_tokens_per_extend_request = num_tokens_per_extend_request - 1 - self.draft_tokens_cuda[:previous_batch_draft_tokens].copy_( - next_draft_tokens_device[ - previous_slots, :num_draft_tokens_per_extend_request].flatten(), - non_blocking=True) + if is_ragged: + if previous_batch_draft_tokens > 0: + draft_rows, draft_cols = self._ragged_gather_indices( + previous_batch_slots, + [t - 1 for t in tokens_per_extend_request]) + self.draft_tokens_cuda[:previous_batch_draft_tokens].copy_( + next_draft_tokens_device[draft_rows, draft_cols].flatten(), + non_blocking=True) + else: + num_draft_tokens_per_extend_request = num_tokens_per_extend_request - 1 + self.draft_tokens_cuda[:previous_batch_draft_tokens].copy_( + next_draft_tokens_device[ + previous_slots, :num_draft_tokens_per_extend_request]. + flatten(), + non_blocking=True) - # Compute kv_len_offsets and update offset tensors - previous_pos_indices = previous_slots.repeat_interleave( - num_tokens_per_extend_request) + # Compute kv_len_offsets and update offset tensors. + # kv_len_offsets pairs with num_cached_tokens_per_seq (past_seen + + # tokens_this_step) computed by the caller: the two cancel to + # past_seen + accepted, so both must use the same per-request count. + if is_ragged: + tokens_per_request_device = torch.tensor( + tokens_per_extend_request, + dtype=torch.long, + pin_memory=prefer_pinned()).to(new_tokens_lens_device.device, + non_blocking=True) + previous_pos_indices = torch.repeat_interleave( + previous_slots, tokens_per_request_device) + previous_kv_len_offsets = (new_tokens_lens_device[previous_slots] - + tokens_per_request_device) + else: + previous_pos_indices = previous_slots.repeat_interleave( + num_tokens_per_extend_request) + kv_len_offsets_device = new_tokens_lens_device - num_tokens_per_extend_request + previous_kv_len_offsets = kv_len_offsets_device[previous_slots] self.previous_pos_indices_cuda[:total_num_tokens].copy_( previous_pos_indices, non_blocking=True) - kv_len_offsets_device = new_tokens_lens_device - num_tokens_per_extend_request - self.previous_pos_id_offsets_cuda[:num_extend_reqeust_wo_dummy * - num_tokens_per_extend_request].copy_( - new_tokens_lens_device[ - self. - previous_pos_indices_cuda[: - total_num_tokens]], - non_blocking=True) + self.previous_pos_id_offsets_cuda[:total_num_tokens].copy_( + new_tokens_lens_device[ + self.previous_pos_indices_cuda[:total_num_tokens]], + non_blocking=True) self.previous_kv_lens_offsets_cuda[:num_extend_reqeust_wo_dummy].copy_( - kv_len_offsets_device[previous_slots], non_blocking=True) + previous_kv_len_offsets, non_blocking=True) def _apply_incremental_update_target( self, @@ -3581,7 +3943,6 @@ def _apply_incremental_update_target( spec_config = self.spec_config num_tokens_per_extend_request = self.get_runtime_tokens_per_gen_step( self.runtime_draft_len) - runtime_draft_token_buffer_width = num_tokens_per_extend_request - 1 prompt_lengths = torch.empty(num_extend_requests, dtype=torch.int, @@ -3599,6 +3960,11 @@ def _apply_incremental_update_target( request_accepted_path = {} num_extend_dummy_requests = 0 num_previous_batch = 0 + # Per-request token counts, in previous-batch order. Stays all-equal to + # num_tokens_per_extend_request unless a ragged scheduler assigned each + # request its own verify window. + tokens_per_extend_request = [] + previous_batch_slots = [] use_extend_ctx = (self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( @@ -3609,11 +3975,17 @@ def _apply_incremental_update_target( request.py_num_accepted_draft_tokens_indices base_past_seen = request.max_beam_num_tokens - 1 + # Under ragged verification every length derived below has to come + # from this one per-request value; mixing it with the batch-wide + # count desynchronizes the flat token layout from the KV-length + # correction applied in _preprocess_inputs. + req_tokens_per_gen_step = get_request_tokens_per_gen_step( + request, num_tokens_per_extend_request) if use_extend_ctx: # We're treating the prompt lengths as context requests here, so # the prompt lens should not include the cached tokens. - prompt_lengths[idx] = num_tokens_per_extend_request + prompt_lengths[idx] = req_tokens_per_gen_step else: prompt_lengths[idx] = request.py_prompt_len @@ -3630,17 +4002,21 @@ def _apply_incremental_update_target( previous_batch_indices[ num_previous_batch] = request.py_batch_idx num_previous_batch += 1 + previous_batch_slots.append(request.py_batch_idx) + tokens_per_extend_request.append(req_tokens_per_gen_step) request.cached_tokens = (base_past_seen + - num_tokens_per_extend_request) + req_tokens_per_gen_step) num_cached_tokens_per_seq[idx] = ( - base_past_seen + num_tokens_per_extend_request - + base_past_seen + req_tokens_per_gen_step - request.py_num_compressed_tokens) request.py_batch_idx = request.py_seq_slot num_extend_reqeust_wo_dummy = num_extend_requests - num_extend_dummy_requests - total_num_tokens = num_extend_reqeust_wo_dummy * num_tokens_per_extend_request + is_ragged_gen = any(tokens != num_tokens_per_extend_request + for tokens in tokens_per_extend_request) + total_num_tokens = sum(tokens_per_extend_request) previous_slots = self.previous_batch_indices_cuda[:num_previous_batch] previous_slots.copy_(previous_batch_indices[:num_previous_batch], @@ -3649,8 +4025,8 @@ def _apply_incremental_update_target( prompt_lengths = prompt_lengths.tolist() num_cached_tokens_per_seq = num_cached_tokens_per_seq.tolist() - previous_batch_draft_tokens = (num_extend_reqeust_wo_dummy * - runtime_draft_token_buffer_width) + previous_batch_draft_tokens = (total_num_tokens - + num_extend_reqeust_wo_dummy) self._update_target_input_tensors( num_accepted_tokens_device=num_accepted_tokens_device, @@ -3661,10 +4037,16 @@ def _apply_incremental_update_target( total_num_tokens=total_num_tokens, num_extend_reqeust_wo_dummy=num_extend_reqeust_wo_dummy, num_tokens_per_extend_request=num_tokens_per_extend_request, - previous_batch_draft_tokens=previous_batch_draft_tokens) - - # Prepare spec_metadata - num_generation_tokens = num_extend_requests * num_tokens_per_extend_request + previous_batch_draft_tokens=previous_batch_draft_tokens, + tokens_per_extend_request=(tokens_per_extend_request + if is_ragged_gen else None), + previous_batch_slots=previous_batch_slots) + + # Prepare spec_metadata. Dummy requests are padded to the batch-wide + # window, so only the real requests contribute ragged token counts. + num_generation_tokens = ( + total_num_tokens + + num_extend_dummy_requests * num_tokens_per_extend_request) if spec_metadata is not None: total_draft_lens = self.max_total_draft_tokens * num_extend_requests spec_metadata.draft_tokens = self.draft_tokens_cuda[: @@ -4187,6 +4569,16 @@ def append_cross_attention_state(request: LlmRequest, # will contain previous batch indices of generation requests previous_batch_indices = [] previous_pos_indices = [] + # Token count each previous-batch request contributes, in batch order. + # Uniform speculation makes every entry runtime_tokens_per_gen_step; + # ragged verification does not, and the device-side gathers below + # switch to an index-list layout when they disagree. + previous_batch_tokens_per_request = [] + # Flat generation-token offset of the first previous-batch request, + # captured while walking the batch so the ragged layout does not have + # to assume a fixed stride. + previous_batch_token_start = None + extend_tokens_emitted = 0 runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step( self.runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 @@ -4236,6 +4628,7 @@ def append_cross_attention_state(request: LlmRequest, num_cached_tokens_per_seq.append( past_seen_token_num - request.py_num_compressed_tokens) request.cached_tokens = past_seen_token_num + extend_tokens_emitted += 1 + num_draft_tokens # update batch index request.py_batch_idx = request.py_seq_slot else: @@ -4243,33 +4636,46 @@ def append_cross_attention_state(request: LlmRequest, previous_batch_idx = request.py_batch_idx request.py_batch_idx = request.py_seq_slot - sequence_lengths.append(runtime_tokens_per_gen_step) + if previous_batch_token_start is None: + previous_batch_token_start = extend_tokens_emitted + + # Under ragged verification each request gets its own window; + # every length below has to come from the same per-request + # value, otherwise the flat token layout and the KV-length + # correction in _preprocess_inputs disagree. + req_tokens_per_gen_step = get_request_tokens_per_gen_step( + request, runtime_tokens_per_gen_step) + + sequence_lengths.append(req_tokens_per_gen_step) num_accepted_draft_tokens.append( request.py_num_accepted_draft_tokens) past_seen_token_num = request.max_beam_num_tokens - 1 - draft_lens.append(runtime_draft_token_buffer_width) + draft_lens.append(req_tokens_per_gen_step - 1) gather_ids.extend( list( range(len(position_ids), - len(position_ids) + runtime_tokens_per_gen_step))) + len(position_ids) + req_tokens_per_gen_step))) position_ids.extend( list( - range(past_seen_token_num, past_seen_token_num + - runtime_tokens_per_gen_step))) + range(past_seen_token_num, + past_seen_token_num + req_tokens_per_gen_step))) # previous tensor previous_batch_indices.append(previous_batch_idx) previous_pos_indices.extend([previous_batch_idx] * - runtime_tokens_per_gen_step) + req_tokens_per_gen_step) + previous_batch_tokens_per_request.append( + req_tokens_per_gen_step) + extend_tokens_emitted += req_tokens_per_gen_step num_cached_tokens_per_seq.append( - past_seen_token_num + runtime_tokens_per_gen_step - + past_seen_token_num + req_tokens_per_gen_step - request.py_num_compressed_tokens) request.cached_tokens = (past_seen_token_num + - runtime_tokens_per_gen_step) + req_tokens_per_gen_step) if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( self.attn_backend) and spec_config.is_linear_tree: - prompt_lengths.append(runtime_tokens_per_gen_step) + prompt_lengths.append(req_tokens_per_gen_step) else: prompt_lengths.append(request.py_prompt_len) @@ -4628,30 +5034,76 @@ def previous_seq_slots_device(): if previous_batch_len > 0: previous_slots = previous_seq_slots_device() + # Ragged verification gives each request its own window, so the + # fixed [previous_batch_len, runtime_tokens_per_gen_step] block + # gathers below no longer describe the layout. Detect that here + # and fall back to explicit (row, col) index lists; a uniform + # batch -- every non-DSpark model, and DSpark whenever the + # scheduler picks one window for the whole batch -- keeps the + # cheaper strided path. + is_ragged_gen = any( + tokens != runtime_tokens_per_gen_step + for tokens in previous_batch_tokens_per_request) # previous input ids - previous_batch_tokens = (previous_batch_len * - runtime_tokens_per_gen_step) - new_tokens = new_tokens_device.transpose( - 0, - 1)[previous_slots, :runtime_tokens_per_gen_step].flatten() + previous_batch_tokens = len(previous_pos_indices) + if is_ragged_gen: + token_rows, token_cols = self._ragged_gather_indices( + previous_batch_indices, + previous_batch_tokens_per_request) + new_tokens = new_tokens_device.transpose( + 0, 1)[token_rows, token_cols].flatten() + else: + new_tokens = new_tokens_device.transpose(0, 1)[ + previous_slots, :runtime_tokens_per_gen_step].flatten() self.input_ids_cuda[num_tokens:num_tokens + previous_batch_tokens].copy_( new_tokens, non_blocking=True) # previous draft tokens - previous_batch_draft_tokens = (previous_batch_len * - runtime_draft_token_buffer_width) - if runtime_draft_token_buffer_width > 0: - self.draft_tokens_cuda[ - num_draft_tokens:num_draft_tokens + - previous_batch_draft_tokens].copy_( - next_draft_tokens_device[ - previous_slots, : - runtime_draft_token_buffer_width].flatten(), - non_blocking=True) - # prepare data for the preprocess inputs - kv_len_offsets_device = (new_tokens_lens_device - - runtime_tokens_per_gen_step) + if is_ragged_gen: + previous_batch_draft_tokens = (previous_batch_tokens - + previous_batch_len) + if previous_batch_draft_tokens > 0: + draft_rows, draft_cols = self._ragged_gather_indices( + previous_batch_indices, + [t - 1 for t in previous_batch_tokens_per_request]) + self.draft_tokens_cuda[ + num_draft_tokens:num_draft_tokens + + previous_batch_draft_tokens].copy_( + next_draft_tokens_device[draft_rows, + draft_cols].flatten(), + non_blocking=True) + else: + previous_batch_draft_tokens = ( + previous_batch_len * runtime_draft_token_buffer_width) + if runtime_draft_token_buffer_width > 0: + self.draft_tokens_cuda[ + num_draft_tokens:num_draft_tokens + + previous_batch_draft_tokens].copy_( + next_draft_tokens_device[ + previous_slots, : + runtime_draft_token_buffer_width].flatten(), + non_blocking=True) + # prepare data for the preprocess inputs. + # kv_len_offsets pairs with the num_cached_tokens_per_seq the + # host wrote above (past_seen + tokens_this_step): the two + # cancel to past_seen + accepted, so both sides must use the + # SAME per-request token count or the KV length is off by + # exactly their difference. + if is_ragged_gen: + tokens_per_previous_request = torch.tensor( + previous_batch_tokens_per_request, + dtype=torch.long, + pin_memory=prefer_pinned()).to( + new_tokens_lens_device.device, non_blocking=True) + previous_kv_len_offsets = ( + new_tokens_lens_device[previous_slots] - + tokens_per_previous_request) + else: + kv_len_offsets_device = (new_tokens_lens_device - + runtime_tokens_per_gen_step) + previous_kv_len_offsets = kv_len_offsets_device[ + previous_slots] previous_pos_indices_host = torch.tensor( previous_pos_indices, dtype=torch.int, @@ -4675,10 +5127,19 @@ def previous_seq_slots_device(): num_extend_reqeust_wo_dummy = len(extend_requests) - len( extend_dummy_requests) + if is_ragged_gen: + # previous_pos_id_offsets_cuda is indexed by flat + # generation-token position, which is no longer + # request_index * runtime_tokens_per_gen_step. Use the + # offset accumulated while walking the batch. + pos_offsets_start = previous_batch_token_start + else: + pos_offsets_start = ( + num_extend_reqeust_wo_dummy - + previous_batch_len) * runtime_tokens_per_gen_step self.previous_pos_id_offsets_cuda[ - (num_extend_reqeust_wo_dummy - previous_batch_len) * - runtime_tokens_per_gen_step:num_extend_reqeust_wo_dummy * - runtime_tokens_per_gen_step].copy_( + pos_offsets_start:pos_offsets_start + + previous_batch_tokens].copy_( new_tokens_lens_device[self.previous_pos_indices_cuda[ 0:previous_batch_tokens]], non_blocking=True) @@ -4686,8 +5147,7 @@ def previous_seq_slots_device(): self.previous_kv_lens_offsets_cuda[ num_extend_reqeust_wo_dummy - previous_batch_len:num_extend_reqeust_wo_dummy].copy_( - kv_len_offsets_device[previous_slots], - non_blocking=True) + previous_kv_len_offsets, non_blocking=True) elif new_tokens_device is not None: seq_slots_device = previous_seq_slots_device() @@ -4941,6 +5401,9 @@ def previous_seq_slots_device(): spec_metadata.seq_lens = sequence_lengths spec_metadata.num_accepted_draft_tokens = self.num_accepted_draft_tokens_cuda[:len( num_accepted_draft_tokens)] + self._attach_ragged_verify_layout( + spec_metadata, attn_metadata, + scheduled_requests.generation_requests) if isinstance(spec_metadata, Eagle3SpecMetadata): spec_metadata.request_accepted_path = request_accepted_path # No-op for non 1-model @@ -4983,6 +5446,10 @@ def previous_seq_slots_device(): if not self.is_warmup: self.previous_request_ids = all_gen_request_ids + self.previous_verify_lens = [ + getattr(request, "py_verify_len", None) + for request in scheduled_requests.generation_requests + ] self.has_previous_device_draft = next_draft_tokens_device is not None # Record the steady-state generation cache when this pass handled diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 34ebd340a5ee..267aef9dd966 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3186,6 +3186,97 @@ def wait_on_pp_send_handles(self, send_handles, microbatch_id): send_handles[microbatch_id].wait() send_handles[microbatch_id] = None + def _dspark_confidence_draft_len(self, + scheduled_batch: ScheduledRequests) -> int: + """Draft length for this iteration, from the DSpark confidence planner. + + Runs entirely outside CUDA-graph capture, which is what makes the + device->host relay legal: the worker writes confidence device-side inside + the captured graph, and this hook -- called before the next forward -- + first *reads* the previously staged snapshot, then stages a fresh one. + That gives a one-iteration lag for free, with no synchronization. + + Falls back to the full block on any missing piece, so an engine without a + confidence-capable checkpoint behaves exactly as before. + """ + max_draft_len = self.model_engine.max_draft_len + worker = self.model_engine._get_spec_worker() + planner = getattr(worker, "verify_planner", None) + if planner is None: + return max_draft_len + + num_gen_requests = len(scheduled_batch.generation_requests) + + # Cross-rank agreement: draft_len is part of the CUDA-graph key but is + # NOT covered by the attention-DP consistency allgather, so ranks that + # pick different lengths select different graphs -- one replays, one + # falls back to eager -- and their collectives diverge. Reduce here. + # + # This costs one extra blocking host collective per decode step under + # ADP. ``_can_queue`` already runs ``tp_allgather(batch_size)`` just + # before this hook, so the two could be fused into a single + # ``[batch_size, local_draft_len]`` gather -- but that means computing + # the local draft length inside ``_can_queue``, which is a hot generic + # path shared by every speculation mode. Left separate deliberately; + # fuse only if the extra collective shows up in a profile. + all_rank_max = None + if self.dist is not None and getattr(self.dist, "tp_size", 1) > 1: + + def all_rank_max(value: int) -> int: + return max(int(v) for v in self.dist.tp_allgather(int(value))) + + ragged_lens = None + if getattr(self.model_engine.spec_config, "enable_ragged_verify", + False): + ragged_lens = planner.decide_verify_lens( + num_gen_requests=num_gen_requests, all_rank_max=all_rank_max) + + # Land the batch's token total on a captured bucket before the graph key + # is built, spending the rounding slack on real verification. A batch + # with no captured shape gets None back and falls through to uniform, + # which always has a graph -- running ragged without one would cost far + # more than the trimmed tokens save. + bucket = None + if ragged_lens is not None: + # Under ADP the batch is padded to the widest rank's request count, + # so the shape has to be chosen from every rank's (rows, tokens) -- + # sizing it locally makes each rank pick a different bucket, i.e. + # a different graph. + peer_stats = None + if self.dist is not None and getattr(self.dist, "tp_size", 1) > 1: + peer_stats = self.dist.tp_allgather( + [len(ragged_lens), + sum(1 + int(v) for v in ragged_lens)]) + bucket = self.model_engine.fit_ragged_verify_lens( + scheduled_batch.generation_requests, + ragged_lens, + peer_stats=peer_stats) + self.model_engine.ragged_verify_bucket = bucket + + if bucket is not None: + # The block is always drafted in full -- only verification is + # trimmed -- so the batch-wide draft length stays at the top tier. + # That keeps the drafted-token buffer and the padded acceptance + # rectangle a fixed width; each request's own (smaller) verify + # window rides along on py_verify_len. + runtime_draft_len = max( + int(t) for t in self.model_engine.spec_config.verify_len_tiers) + else: + # Clear any window left over from a previous ragged step, otherwise + # a fallback iteration would keep trimming on a stale split. + for request in scheduled_batch.generation_requests: + request.py_verify_len = None + runtime_draft_len = planner.decide_draft_len( + num_gen_requests=num_gen_requests, all_rank_max=all_rank_max) + + # Stage this step's confidence for the *next* decision. Non-blocking; if + # it has not landed by then the planner just verifies the full block. + rows = worker.staged_confidence_rows() + if rows is not None and rows.shape[0] > 0: + planner.stage_confidence(rows, rows.shape[0]) + + return int(runtime_draft_len) + def _handle_dynamic_draft_len(self, scheduled_batch: ScheduledRequests) -> None: """Handle dynamic draft length for the current batch. @@ -3213,19 +3304,31 @@ def _handle_dynamic_draft_len(self, self.model_engine.runtime_draft_len = 0 return - if (self.model_engine.spec_config is not None - and self.model_engine.spec_config.draft_len_schedule is not None - and self.model_engine.spec_config.spec_dec_mode. - support_dynamic_draft_len()): + spec_config = self.model_engine.spec_config + schedule_driven = ( + spec_config is not None + and spec_config.draft_len_schedule is not None + and spec_config.spec_dec_mode.support_dynamic_draft_len()) + # DSpark drives the draft length from the batch's confidence rather than + # from a static batch-size schedule, but consumes the exact same + # pad/truncate machinery below. + confidence_driven = (spec_config is not None and getattr( + spec_config, "enable_confidence_scheduling", False)) + + if schedule_driven or confidence_driven: from tensorrt_llm._torch.speculative.utils import \ get_draft_len_for_batch_size - spec_dec_mode = self.model_engine.spec_config.spec_dec_mode + spec_dec_mode = spec_config.spec_dec_mode - # 1. Resolve runtime draft length from schedule - runtime_draft_len = get_draft_len_for_batch_size( - self.model_engine.spec_config.draft_len_schedule, - scheduled_batch.batch_size, self.model_engine.max_draft_len) + # 1. Resolve runtime draft length + if confidence_driven: + runtime_draft_len = self._dspark_confidence_draft_len( + scheduled_batch) + else: + runtime_draft_len = get_draft_len_for_batch_size( + spec_config.draft_len_schedule, scheduled_batch.batch_size, + self.model_engine.max_draft_len) # 2. Pad or truncate draft tokens to the resolved length DRAFT_BUFFER_PAD = 0 # Buffer sentinel, not PARD mask_token_id. rejection_on = getattr(self.model_engine.spec_config, diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 7f43f103ec4e..a54a89c6544a 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -35,6 +35,10 @@ if TYPE_CHECKING: from ...llmapi.llm_args import DSparkDecodingConfig +# Raw-logit fill for an unscored confidence slot: sigmoid(30) rounds to 1.0 in +# fp32, so an unwritten row schedules as "verify the whole block". +_NEUTRAL_CONFIDENCE_LOGIT = 30.0 + @dataclass class DSparkSpecMetadata(SpecMetadata): @@ -189,8 +193,10 @@ class DSparkWorker(SpecWorkerBase): (``DSparkDraftModel.forward``): it projects the captured target-layer hidden states (``main_proj`` + ``main_norm``) into the draft's captured-context attention, runs the ``num_stages`` DSpark blocks over a rolling captured - window, refines the per-position logits with the Markov head, and predicts a - per-position acceptance confidence used to truncate the proposed prefix. + window, refines the per-position logits with the Markov head, and (when + ``enable_confidence_scheduling`` is set) scores each position with the + confidence head. The block is always proposed in full; the confidence only + feeds the verification scheduler's budget, never the acceptance decision. Unlike DFlash, the draft does NOT use the paged KV cache or mask-token cross-attention: its attention K/V come from the worker-owned rolling window @@ -240,6 +246,18 @@ def __init__( # never handed out through ``_free_slots``. self._scratch_slot = 0 + # Confidence-scheduled verification. ``return_confidence`` is read once + # here (never per step) so the captured draft graph cannot diverge. + # ``_confidence_logits`` is the slot-indexed handoff buffer to the + # verification scheduler; see ``_lazy_init`` for its neutral value. + self.return_confidence = bool(getattr(spec_config, "enable_confidence_scheduling", False)) + self._confidence_logits: Optional[torch.Tensor] = None # [max_batch+1, block] + self._last_gen_slots: Optional[torch.Tensor] = None + # Host-side coordinator that turns the confidence snapshot into this + # iteration's draft length. Built in ``_lazy_init`` (it needs the draft + # model's block size and calibration head). + self.verify_planner = None + # The generation draft path is the batched, host-sync-free # ``_draft_gen_block_batched`` + ``DSparkDraftModel.forward_batched`` + # ``dspark_attention_forward_batched``: it is correct in eager mode AND safe @@ -252,6 +270,18 @@ def __init__( f"use_separate_draft_kv_cache={use_separate_draft_kv_cache}" ) + def staged_confidence_rows(self) -> Optional[torch.Tensor]: + """Confidence rows for the requests scored by the last draft, or None. + + Indexes through ``_last_gen_slots`` rather than taking the buffer's + leading rows: the buffer is slot-indexed, and a batch rarely occupies + slots ``0..G-1``. Reading the leading rows would score a stale, unrelated + set of requests. + """ + if self._confidence_logits is None or self._last_gen_slots is None: + return None + return self._confidence_logits[self._last_gen_slots] + @property def max_draft_len(self) -> int: return self.spec_config.max_draft_len @@ -299,6 +329,19 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: ) self._ctx_len = torch.zeros(num_rows, dtype=torch.long, device="cuda") self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") + if self.return_confidence: + # Neutral fill = a large positive logit, i.e. sigmoid ~ 1.0. A slot + # that has not been scored yet therefore reads as "certain to be + # accepted", which makes the scheduler fall back to verifying the + # full block -- the safe direction (never verifies fewer tokens than + # today's static behavior on stale/unwritten rows). + self._confidence_logits = torch.full( + (num_rows, block_size), + _NEUTRAL_CONFIDENCE_LOGIT, + dtype=torch.float32, + device="cuda", + ) + self._build_verify_planner(draft_model, block_size, max_batch) self._free_slots = deque(range(max_batch)) self._req_to_slot = {} self._win_inited = True @@ -308,6 +351,82 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: f"({max_batch} request slots + 1 scratch row)" ) + def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> None: + """Construct the host-side verify planner (once, at lazy init). + + ``max_batch`` is only used to derive the verify-length tier ladder when + the user did not configure one; see :func:`derive_verify_len_tiers` for + why the ladder is a function of batch size. + """ + import json + + from .dspark_planner import SpsCostTable + from .dspark_schedule import DSparkScheduleConfig + from .dspark_verify import DSparkVerifyPlanner + + cfg = DSparkScheduleConfig(block_size=block_size, min_verify_len=1) + + cost_table = None + table_path = getattr(self.spec_config, "confidence_sps_table_path", None) + if table_path: + with open(table_path, encoding="utf-8") as f: + raw = json.load(f) + cost_table = SpsCostTable( + token_counts=[int(v) for v in raw["token_counts"]], + step_time_ms=[float(v) for v in raw["step_time_ms"]], + ) + else: + # Explicitly flat: the planner treats this as "unprofiled" and keeps + # verifying the full block rather than trimming on a cost model that + # says every extra token is free. + cost_table = SpsCostTable.flat() + logger.warning( + "DSpark confidence scheduling is enabled but no " + "confidence_sps_table_path was provided. Without a profiled step-cost " + "curve the planner cannot tell a cheap verify token from an expensive " + "one, so it will keep verifying the full block (no scheduling gain)." + ) + + # Calibration lives on the draft's confidence head; fall back to a plain + # sigmoid when the head is absent (no confidence to calibrate anyway). + head = getattr(getattr(draft_model, "mtp_layers", [None])[-1], "confidence_head", None) + apply_calibration = head.apply_sts if head is not None else None + sts_path = getattr(self.spec_config, "confidence_sts_path", None) + if sts_path and head is not None: + with open(sts_path, encoding="utf-8") as f: + temps = json.load(f)["sts_temperatures"] + head.load_sts_temperatures(torch.tensor(temps, dtype=torch.float32)) + logger.info(f"DSpark: loaded STS calibration from {sts_path}") + + tiers = getattr(self.spec_config, "verify_len_tiers", None) + if not tiers: + # Not configured: derive the ladder from the measured cost curve + # instead of hardcoding one. Each tier is a captured graph whose + # memory comes out of KV cache, so the derivation is capped -- and + # it is only zero-loss at the batch size it was derived for (the + # shelf right edges move with batch size), hence the steady-state + # modal batch size is the right input. A flat table yields just + # [min, block_size], which is the no-scheduling ladder. + from .dspark_planner import derive_verify_len_tiers + + tiers = derive_verify_len_tiers( + cost_table=cost_table, + num_requests=max(int(max_batch), 1), + block_size=block_size, + min_verify_len=cfg.min_verify_len, + ) + logger.info(f"DSpark: derived verify-length tiers {tiers}") + self.verify_planner = DSparkVerifyPlanner( + cfg=cfg, + cost_table=cost_table, + tiers=tiers, + apply_calibration=apply_calibration, + ) + logger.info( + f"DSpark verify planner: tiers={self.verify_planner.tiers}, " + f"profiled_cost_table={not cost_table.is_flat}" + ) + def _assign_slot(self, req_id: int, reset: bool) -> int: """Get (or refresh) the slot for a request; reset clears its window.""" if reset and req_id in self._req_to_slot: @@ -389,8 +508,21 @@ def _draft_gen_block_batched( Confidence truncation stays disabled — the full block is proposed. """ num_gens = batch_size - num_contexts + # Two different widths meet here, and conflating them corrupts the gather: + # K - the DRAFT block width. Fixed by the checkpoint; the draft + # always proposes a full block regardless of how much of it + # the target was asked to verify. + # Kp1 - the TARGET's per-gen-request token stride in ``captured``. + # That is ``runtime_draft_len + 1``, which shrinks when the + # verification scheduler trims the batch's draft length. + # Using max_draft_len for the stride reads into the *next* request's + # hidden states whenever runtime_draft_len < max_draft_len, silently + # drafting from the wrong context. K = self.max_draft_len - Kp1 = K + 1 + runtime_draft_len = getattr(spec_metadata, "runtime_draft_len", None) + if runtime_draft_len is None: + runtime_draft_len = K + Kp1 = int(runtime_draft_len) + 1 device = accepted_tokens.device if num_gens == 0: @@ -420,6 +552,10 @@ def _draft_gen_block_batched( # Fixed-size ([G, K]) masked back-fill of the intermediate accepted tokens # (everything but the bonus) into the rolling window — same frames as the # eager path (old+1 .. old+nacc-1), with j >= nacc-1 masked out. + # The width stays at the static K even when the runtime stride is smaller: + # ``nacc <= runtime_draft_len + 1`` bounds the *valid* j to within this + # request's own block, so the extra columns are always masked out. Keeping + # the width static also keeps this tensor's shape independent of K. old = self._ctx_len[slots] # [G] pre-increment decode position j = torch.arange(K, device=device) # [K] interim_valid = j.unsqueeze(0) < (nacc.unsqueeze(1) - 1) # [G, K] @@ -440,17 +576,27 @@ def _draft_gen_block_batched( # Surface the per-position corrected block logits ([num_gens, K, vocab]) # and let SpecWorkerBase.sample_draft_tokens do the (greedy or rejection) # sampling + TP gather + draft_probs scatter, rather than argmaxing here. - _toks, _num_proposed, block_logits = draft_model.forward_batched( + _toks, confidence, block_logits = draft_model.forward_batched( main_hidden, bonus, start_pos, kv_windows=self._kv_windows, slots=slots, temperature=0.0, - confidence_threshold=0.0, + return_confidence=self.return_confidence, return_logits=True, all_rank_num_tokens=all_rank_num_tokens, ) + # Stash the raw [G, K] confidence logits for the verification scheduler. + # Slot-indexed (not batch-position-indexed) so it survives batch + # reshuffles; written in place so a captured graph keeps the same storage. + # ``_last_gen_slots`` is a *view* into ``_batch_to_slot``, so it keeps + # pointing at whichever slots the current batch occupies on replay -- the + # scheduler must read the rows for the requests that were actually + # scored, not slots 0..G-1. + if confidence is not None: + self._confidence_logits[slots] = confidence.detach() + self._last_gen_slots = slots return block_logits def _forward_impl( diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py new file mode 100644 index 000000000000..52ba92824bb2 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -0,0 +1,288 @@ +# 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 verification-budget planner: how many tokens are worth verifying. + +The scheduler maximizes *system* goodput, not per-request speed:: + + tau(n) = num_gen_requests + sum of the n largest prefix-survival values + Theta(n) = tau(n) / step_cost_ms(num_gen_requests + n) + budget = argmax_n Theta(n) + +``tau(n)`` is the expected number of tokens the step emits if it verifies ``n`` +tokens beyond the per-request floor: every generating request contributes one +bonus token for free, and each admitted draft position contributes its survival +probability. Dividing by the step cost turns that into expected tokens per +millisecond. + +Two things are easy to get wrong here: + +**Take the global argmax, not the first descent.** The marginal cost of a verify +token is a staircase -- flat while a token slots into the current kernel waves, +then a jump when it starts a new one -- so ``Theta`` is not unimodal. A greedy +loop that stops at the first non-improvement parks at the end of the first +shelf and leaves most of the win behind. Both SGLang and vLLM PR #47808 scan the +whole curve; so do we, and :mod:`tests` pins it against a brute-force scan. + +**Cost is per step, not per token.** ``SpsCostTable`` therefore stores measured +step time against the *total* verified token count, and interpolates nothing +across a riser. + +The planner runs on the host: its output selects a CUDA graph, and graph +selection is a host decision in TensorRT-LLM. Keeping it here (rather than on +device) is deliberate -- a device-side budget would have to be copied back +before the batch could be shaped, which is exactly the sync we are avoiding. +""" + +import bisect +from dataclasses import dataclass +from typing import List, Optional, Sequence + +import numpy as np + +__all__ = [ + "SpsCostTable", + "compute_verify_token_budget", + "budget_argmax_over_uniform_lens", + "derive_verify_len_tiers", +] + +# Fixed per-step overhead folded into every cost lookup (sampling, bookkeeping, +# launch latency). Matches vLLM PR #47808's ``_FIXED_OVERHEAD_MS``. +_FIXED_OVERHEAD_MS = 1.0 + + +@dataclass(frozen=True) +class SpsCostTable: + """Measured decode step time as a function of total verified tokens. + + Attributes: + token_counts: strictly increasing total-token breakpoints. + step_time_ms: measured step time at each breakpoint. + + Lookup is a floor (staircase) lookup, never an interpolation: between two + measured points the cost is genuinely flat until a new kernel wave starts, + and smoothing that away is what makes a planner over-spend right before a + riser. Queries above the largest breakpoint clamp to the last entry, and the + caller is expected to bound the budget so this does not happen silently. + """ + + token_counts: Sequence[int] + step_time_ms: Sequence[float] + + def __post_init__(self) -> None: + if len(self.token_counts) != len(self.step_time_ms): + raise ValueError( + f"token_counts ({len(self.token_counts)}) and step_time_ms " + f"({len(self.step_time_ms)}) must have the same length" + ) + if not self.token_counts: + raise ValueError("SpsCostTable requires at least one measured point") + if any(b <= a for a, b in zip(self.token_counts, self.token_counts[1:])): + raise ValueError("token_counts must be strictly increasing") + if any(t <= 0.0 for t in self.step_time_ms): + raise ValueError("step_time_ms entries must be positive") + + def step_time(self, num_tokens: int) -> float: + """Step time (ms) for a step that verifies ``num_tokens`` tokens in total.""" + idx = bisect.bisect_right(self.token_counts, int(num_tokens)) - 1 + idx = min(max(idx, 0), len(self.step_time_ms) - 1) + return float(self.step_time_ms[idx]) + _FIXED_OVERHEAD_MS + + def step_times(self, num_tokens: np.ndarray) -> np.ndarray: + """Vectorized :meth:`step_time`.""" + counts = np.asarray(self.token_counts) + idx = np.clip(np.searchsorted(counts, num_tokens, side="right") - 1, 0, len(counts) - 1) + return np.asarray(self.step_time_ms, dtype=np.float64)[idx] + _FIXED_OVERHEAD_MS + + @classmethod + def flat(cls, step_time_ms: float = 1.0) -> "SpsCostTable": + """A cost model with no staircase at all. + + Useful as an explicit "I have not profiled this machine" marker. With a + flat cost every extra token is free, so the planner will always spend the + entire budget -- i.e. it degenerates to verify-all. Callers should warn + rather than silently ship this. + """ + return cls(token_counts=(0,), step_time_ms=(float(step_time_ms),)) + + @property + def is_flat(self) -> bool: + return len(set(self.step_time_ms)) <= 1 + + +def compute_verify_token_budget( + *, + survival: np.ndarray, + num_gen_requests: int, + cost_table: SpsCostTable, + min_verify_len: int = 1, + max_verify_len: Optional[int] = None, +) -> int: + """Return the batch-wide verify-token budget (tokens *above* the floor). + + Args: + survival: ``[bs, K]`` prefix-survival probabilities. Only the columns at + or beyond ``min_verify_len`` are schedulable; earlier positions are + already granted to every request. + num_gen_requests: generating requests this step. Each contributes one + guaranteed output token, which is the constant term of ``tau``. + cost_table: measured step cost; see :class:`SpsCostTable`. + min_verify_len / max_verify_len: per-request bounds, mirroring + :class:`~.dspark_schedule.DSparkScheduleConfig`. + + Returns: + ``n`` maximizing ``tau(n) / step_cost(floor_tokens + n)``, in ``[0, N]`` + where ``N`` is the number of schedulable candidates. + """ + if survival.ndim != 2: + raise ValueError(f"survival must be [bs, K], got shape {survival.shape}") + bs, block_size = survival.shape + cap = min(int(max_verify_len or block_size), block_size) + schedulable = max(cap - int(min_verify_len), 0) + if bs == 0 or num_gen_requests <= 0 or schedulable <= 0: + return 0 + + surv = np.asarray(survival, dtype=np.float64) + candidates = np.sort(surv[:, min_verify_len : min_verify_len + schedulable].reshape(-1))[::-1] + if candidates.size == 0: + return 0 + + # tau(n) = expected tokens emitted when n candidates are admitted: + # * one bonus token per generating request, emitted regardless; + # * the floor positions, which every request verifies for free -- their + # survival is part of the yield even though no budget buys them. Leaving + # this out shifts the whole curve and can move the argmax, since + # argmax (C + f(n)) / g(n) depends on C. + # * the n admitted candidates, best-first. + base = float(num_gen_requests) + float(surv[:, :min_verify_len].sum()) + tau = base + np.concatenate(([0.0], np.cumsum(candidates))) + floor_tokens = bs * int(min_verify_len) + total_tokens = floor_tokens + np.arange(tau.size) + theta = tau / cost_table.step_times(total_tokens) + return int(np.argmax(theta)) + + +def derive_verify_len_tiers( + *, + cost_table: SpsCostTable, + num_requests: int, + block_size: int, + min_verify_len: int = 1, + max_tiers: int = 3, +) -> List[int]: + """Derive the verify lengths worth capturing a CUDA graph for. + + Restricting K to a handful of captured values sounds like it must cost + something versus a freely-chosen K. **At a fixed batch size it does not**: + + Within one cost shelf the step time is constant while ``tau`` strictly + increases with each admitted token, so ``Theta = tau / cost`` is strictly + increasing across the shelf. The optimum therefore always sits at a + shelf's *right edge* -- never in its interior. For that batch size, a + tier set of the shelf right edges contains the continuous optimum + exactly. + + **The zero-loss property does not survive across batch sizes.** The cost + shelves live in token space, so a shelf's right edge in *length* space is + ``(breakpoint - 1) // batch_size`` -- a function of the batch size. A + deployment captures graphs for many batch sizes but can only capture one + ladder, so a single tier set cannot sit on every batch size's right edges. + Deriving tiers is more robust than hardcoding a pair like ``{1, 5}``, but it + is still a lossy discretization once the batch size moves; the residual loss + has to be measured, not assumed away. Callers that capture one ladder should + derive it at the steady-state modal batch size and check the worst-case loss + across the others. + + This returns the right-edge length for each shelf the given batch size can + reach, always including ``min_verify_len`` and the full block. + + Args: + num_requests: batch size the tiers are derived for. Tiers are a function + of it (total tokens = ``num_requests * length``), so a deployment + that captures several batch sizes derives a set per batch size. + max_tiers: cap on how many lengths to return. Each extra tier is another + captured graph, and captured graphs consume memory that would + otherwise be KV cache -- so this is a real budget, not a formality. + The full block and ``min_verify_len`` are kept first. + """ + lo, hi = int(min_verify_len), int(block_size) + if num_requests <= 0 or hi < lo: + return [lo] + + edges = set() + for breakpoint_tokens in cost_table.token_counts[1:]: + # Largest length whose total token count still sits on the shelf below + # this breakpoint. + length = (int(breakpoint_tokens) - 1) // int(num_requests) + if lo <= length <= hi: + edges.add(length) + # The last shelf extends past the final breakpoint, so the whole block is + # always a right edge; the floor is always runnable. + edges.update({lo, hi}) + + tiers = sorted(edges) + if len(tiers) <= max_tiers: + return tiers + # Keep the endpoints, then spread the remaining slots over the interior + # edges rather than clustering them at one end of the curve. + interior = tiers[1:-1] + step = len(interior) / float(max_tiers - 2) if max_tiers > 2 else 0 + keep = [interior[int(i * step)] for i in range(max_tiers - 2)] if max_tiers > 2 else [] + return sorted({tiers[0], *keep, tiers[-1]}) + + +def budget_argmax_over_uniform_lens( + *, + survival: np.ndarray, + num_gen_requests: int, + cost_table: SpsCostTable, + allowed_lens: Sequence[int], + min_verify_len: int = 1, +) -> int: + """Pick the best verify length from a discrete set (e.g. captured graph tiers). + + TensorRT-LLM can only run one draft length per step, and only lengths that + have a captured CUDA graph. Optimizing over the continuous budget and then + rounding is not the same as optimizing over what is actually runnable: the + rounded-down point can sit on the far side of a cost riser. This evaluates + ``Theta`` directly at each achievable length and returns the best one. + + Args: + allowed_lens: verify lengths with a captured graph. Values below + ``min_verify_len`` are ignored. + Returns: + The chosen per-request verify length. Falls back to ``min_verify_len`` + when no candidate is admissible. + """ + if survival.ndim != 2: + raise ValueError(f"survival must be [bs, K], got shape {survival.shape}") + bs, block_size = survival.shape + lens: List[int] = sorted( + {int(v) for v in allowed_lens if min_verify_len <= int(v) <= block_size} + ) + if bs == 0 or num_gen_requests <= 0 or not lens: + return int(min_verify_len) + + surv = np.asarray(survival, dtype=np.float64) + best_len, best_theta = lens[0], -np.inf + for length in lens: + # Every request verifies ``length`` positions: the expected yield is the + # bonus token plus the survival of *every* verified position (including + # the floor ones), summed over the batch. + tau = float(num_gen_requests) + float(surv[:, :length].sum()) + theta = tau / cost_table.step_time(bs * length) + if theta > best_theta: + best_len, best_theta = length, theta + return int(best_len) diff --git a/tensorrt_llm/_torch/speculative/dspark_ragged.py b/tensorrt_llm/_torch/speculative/dspark_ragged.py new file mode 100644 index 000000000000..69c233e1f127 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_ragged.py @@ -0,0 +1,498 @@ +# 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. + +Uniform-K scheduling picks one verify length for the whole batch. That leaves +real budget on the table, because requests differ a lot in how deep their draft +survives: giving a request with survival ``0.55, 0.22, 0.07`` the same window as +one sitting at ``0.98, 0.95, 0.91`` wastes verification on the first and starves +the second. Ragged verification gives each request its own length. + +The packing follows the shape every variable-length attention path already uses: +a flat token axis plus an index pointer. + + verify_lens [bs] per-request number of verified positions + qo_indptr [bs + 1] exclusive prefix sum; request r owns + [qo_indptr[r], qo_indptr[r + 1]) + extend_start_loc [bs] == qo_indptr[:-1], the offset form some + attention backends want directly + +Two invariants make the rest of the system safe: + +* ``verify_lens >= 1`` for every request. Position 0 carries the bonus/anchor + token; a request that verifies nothing makes no progress at all and would + stall forever. +* ``sum(verify_lens) == graph_num_tokens`` after :meth:`fill_bucket`. The batch + is padded up to a captured CUDA-graph token bucket, so the packed shape is one + of a small fixed set even though the per-request split is arbitrary. That is + what lets a graph capture a ragged batch: the *shape* is constant, the + raggedness lives in the *contents* of ``verify_lens`` / ``qo_indptr``. + +Why the token count has to land exactly on the bucket +----------------------------------------------------- +``seq_lens.sum()`` becomes ``attn_metadata.num_tokens``, which is what gets +all-gathered across attention-DP ranks and drives the MoE's chunk count. If the +packed token count under-reports the rows actually in the batch, attention and +the MoE disagree about how many tokens are in flight -- and nothing raises. + +Reclaiming the bucket padding +----------------------------- +Rounding up to a bucket would normally waste the difference. :meth:`fill_bucket` +spends it in two stages: real requests first (up to ``max_verify_len``, so those +tokens verify real draft positions the step is already paying for), then pad rows +added to reach the captured batch size. Same idea as SGLang's +``align_verify_tokens_to_graph_tier``. + +The real requests cannot always absorb everything -- a request verifies at most +``max_verify_len`` positions, so the batch's capacity is +``padded_bs * max_verify_len``. When the bucket exceeds what the real requests +can take, the remainder must go to pad rows; when it exceeds even the padded +capacity, :meth:`fill_bucket` raises rather than returning a short batch. +""" + +import bisect +from dataclasses import dataclass +from typing import List, Optional, Sequence, Tuple + +import torch + +__all__ = [ + "RaggedVerifyLayout", + "build_qo_indptr", + "round_up_to_bucket", + "row_ids_from_lens", + "scatter_ragged_to_padded", + "count_accepted_ragged", + "fill_padded_rows_onehot", + "ragged_gather_index_lists", + "RaggedCaptureShape", + "choose_ragged_capture_shape", +] + + +def round_up_to_bucket(total: int, buckets: Sequence[int]) -> int: + """Smallest captured bucket that fits ``total``. + + Raises rather than clamping: a batch that exceeds the largest captured + bucket has no graph, and silently running it eager costs far more than the + trimming saves. Callers must reject such a batch *before* selecting a + layout (see :func:`exceeds_captured_buckets`). + """ + if not buckets: + raise ValueError("round_up_to_bucket requires a non-empty bucket list") + if total > buckets[-1]: + raise ValueError( + f"total {total} exceeds the largest captured bucket {buckets[-1]}; " + f"the caller must reject this batch before building a layout" + ) + return buckets[bisect.bisect_left(buckets, total)] + + +def exceeds_captured_buckets(total: int, buckets: Sequence[int]) -> bool: + """Whether ``total`` has no captured bucket (check before building a layout).""" + return bool(buckets) and total > buckets[-1] + + +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: Optional[int] = None + num_real_requests: Optional[int] = 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: Optional[int] = None, + buckets: Optional[Sequence[int]] = None, + total_verify_tokens: Optional[int] = None, + ) -> "RaggedVerifyLayout": + """Pack ``[bs]`` per-request lengths, rounding up to a captured bucket. + + ``graph_num_tokens`` may be given directly (the usual path -- the caller + already knows the bucket because it also had to pick the graph), or + derived from ``buckets`` plus a host-side ``total_verify_tokens``. + Deriving it from the device tensor would need a sync. + """ + lens = verify_lens.to(torch.int32) + if graph_num_tokens is None: + if buckets is None or total_verify_tokens is None: + raise ValueError( + "give graph_num_tokens, or buckets + total_verify_tokens; " + "deriving the total from the device tensor would sync" + ) + graph_num_tokens = round_up_to_bucket(int(total_verify_tokens), buckets) + 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: Optional[int] = None + ) -> "RaggedVerifyLayout": + """Pad the batch to the captured shape, spending the slack on real work. + + Two things have to line up before the batch can run: + + * the row count must equal the captured batch size (``padded_bs``), and + * the token count must equal the captured bucket + (``graph_num_tokens``) -- because ``seq_lens.sum()`` becomes + ``attn_metadata.num_tokens``, which is what gets all-gathered and + drives the MoE's chunk count. If it under-reports, attention and MoE + disagree about how many tokens are in flight. + + So the slack is distributed in two stages: + + 1. **Real requests first**, up to ``max_verify_len``. The step is + already paying for the bucket, so these tokens verify real draft + positions instead of nothing. + 2. **Pad rows** take whatever is left. They exist only to reach + ``padded_bs``; every pad row needs at least one token because an + empty range breaks ``qo_indptr``'s per-row slicing. + + Raises when the bucket cannot be hit exactly -- silently returning a + short batch is what would desynchronize attention from the MoE. + """ + if self.total_verify_tokens is None: + raise ValueError("fill_bucket needs the host-side total_verify_tokens") + real: List[int] = [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 over real rows first -- spreading keeps the extra + # verification where survival is still plausible for several requests + # instead of pushing one request deep into its low-survival tail. Only + # once every real row is saturated does the remainder go to pad rows, + # where it is genuinely wasted. + 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: Optional[int] = None, + exact_fill: bool = False, + ) -> None: + """Assert the invariants the rest of the pipeline relies on. + + ``exact_fill`` additionally requires the token count to hit the bucket + exactly -- the post-:meth:`fill_bucket` contract. Check it on anything + heading for a captured graph: a short batch does not raise anywhere + downstream, it just makes ``seq_lens.sum()`` disagree with the row count + the graph was captured for. + + Host-syncing, so this is for tests and debug builds -- not the hot path. + """ + 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 RaggedCaptureShape: + """The ``(rows, tokens)`` pair a ragged batch is padded up to. + + Attributes: + padded_bs: captured batch size (number of rows, real + pad requests). + bucket: captured token count, i.e. ``sum(verify_lens)`` after padding. + """ + + padded_bs: int + bucket: int + + +def choose_ragged_capture_shape( + *, + num_real_requests: int, + total_verify_tokens: int, + bs_buckets: Sequence[int], + token_buckets: Sequence[int], + peer_stats: Optional[Sequence[Sequence[int]]] = None, +) -> RaggedCaptureShape: + """Pick the captured ``(padded_bs, bucket)`` for a ragged batch. + + Uniform speculation needs one number -- the batch size -- because the token + count follows from it (``bs * (K + 1)``). Ragged verification breaks that + link: two batches with the same row count can hold very different token + totals, so the captured shape is genuinely two-dimensional and both axes + have to be agreed on before the graph is looked up. + + Under attention DP every rank must land on the *same* shape or the ranks + replay different graphs and their collectives diverge. Ranks do not have + equal batches, so agreement cannot come from each rank rounding its own + numbers; it comes from rounding a reduction. Pass every rank's + ``(num_real_requests, total_verify_tokens)`` as ``peer_stats`` (including + this rank's own) and each rank computes an identical answer: + + padded_bs = round_up_bs(max_r num_real_r) + bucket = round_up_bucket(padded_bs + max_r slack_r) + + where ``slack_r = total_r - num_real_r`` is the drafted-position count on + rank ``r``. The ``padded_bs +`` term is the floor every row contributes + (each request verifies at least its bonus position), so the bucket is large + enough for the widest rank's drafts on top of the widest rank's rows. + + Raises when no captured bucket fits, rather than clamping: a batch with no + graph runs eager, which costs far more than the trimming saves, so the + caller has to see it and shrink the batch instead. + """ + if not bs_buckets: + raise ValueError("choose_ragged_capture_shape requires bs_buckets") + if num_real_requests < 0 or total_verify_tokens < num_real_requests: + raise ValueError( + f"total_verify_tokens {total_verify_tokens} cannot be below " + f"num_real_requests {num_real_requests} (every request verifies " + f"at least one position)" + ) + + stats = list(peer_stats) if peer_stats else [(num_real_requests, total_verify_tokens)] + max_real = max(int(real) for real, _ in stats) + max_slack = max(int(total) - int(real) for real, total in stats) + + sorted_bs = sorted({int(b) for b in bs_buckets}) + if max_real > sorted_bs[-1]: + raise ValueError( + f"{max_real} requests exceed the largest captured batch size " + f"{sorted_bs[-1]}; the caller must shrink the batch" + ) + padded_bs = sorted_bs[bisect.bisect_left(sorted_bs, max_real)] + + needed = padded_bs + max_slack + bucket = round_up_to_bucket(needed, sorted({int(t) for t in token_buckets})) + return RaggedCaptureShape(padded_bs=padded_bs, bucket=bucket) + + +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 taken from columns + ``0..counts[i]-1``, concatenated in batch order. This is the ragged + replacement for the ``tensor[slots, :width]`` strided gather the overlap + scheduler uses when every request verifies the same number of positions; + with a constant ``counts`` the two produce the identical flat sequence. + """ + 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) -> torch.Tensor: + """``[bs]`` lengths -> ``[total]`` owning-request id for each packed token. + + ``repeat_interleave`` with a *tensor* of repeats is what makes the packing + ragged; the uniform paths use a scalar repeat and get a fixed stride. + """ + return torch.repeat_interleave( + torch.arange(verify_lens.numel(), device=verify_lens.device), verify_lens.to(torch.long) + ) + + +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. + + The verify path is packed (one flat token axis) because that is what the + attention kernels want, but the acceptance logic is naturally rectangular + (compare draft vs target position by position). This bridges the two + without a host sync: token ``t`` belongs to request ``row_ids[t]`` at column + ``t - qo_indptr[row_ids[t]]``. + """ + bs = int(verify_lens.numel()) + rows = row_ids_from_lens(verify_lens) + 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: + """Make ``[bs, max_len, vocab]`` padding rows a valid one-hot distribution. + + Scattering a ragged batch into a rectangle leaves all-zero probability rows + past each request's window. Those rows are never read -- the accepted count + is clamped to the window -- but a rejection-sampling kernel still walks them + when every real position was accepted, and sampling from an all-zero + distribution can index out of bounds. Putting all the mass on token 0 keeps + the kernel well-defined without affecting any position that is read. + + Modifies ``probs`` in place and returns it. + """ + positions = torch.arange(probs.shape[1], device=probs.device) + invalid = positions.unsqueeze(0) >= verify_lens.to(positions.dtype).unsqueeze(1) + probs[..., 0] = torch.where(invalid, torch.ones_like(probs[..., 0]), probs[..., 0]) + 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. + + Args: + draft_tokens: ``[bs, max_len]`` padded drafted tokens. + target_tokens: ``[bs, max_len]`` padded target tokens at the same + positions. + verify_lens: ``[bs]`` how many positions each request actually verified. + + Positions at or beyond a request's ``verify_len`` were never sent to the + target, so whatever sits there is stale padding. Masking them to "no match" + before the cumulative product makes the run of matches stop at the window + edge -- without the mask a stale slot could compare equal and credit an + acceptance the target never made, which is a silent correctness bug rather + than a throughput one. + """ + 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/dspark_schedule.py b/tensorrt_llm/_torch/speculative/dspark_schedule.py new file mode 100644 index 000000000000..b130b94c5275 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_schedule.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""DSpark confidence-scheduled verification: survival + per-request allocation. + +Given the draft's calibrated per-position acceptance probabilities, decide how +many of the proposed tokens are worth sending to the target for verification. + +The algorithm (DSpark, arXiv:2607.05147 §5; cross-checked against SGLang's +``dspark_schedule.py`` and vLLM PR #47808's ``adaptive_verification.py``): + + 1. ``survival[r][j] = prod(conf[r][:j+1])`` -- the probability that draft + position ``j`` of request ``r`` is reached *and* accepted. The confidence + head predicts a *conditional* acceptance probability, which is exactly what + makes this cumulative product the right prefix statistic. + 2. Rank every ``(request, position)`` candidate by survival, globally across + the batch, and admit the best ``budget`` of them. + +Step 2 needs no explicit prefix constraint: ``survival`` is non-increasing along +each request's positions, so any global top-k is automatically a per-request +*prefix*. That is why the allocation only has to *count* admitted candidates per +row rather than track which ones they were. + +Nothing here decides acceptance -- only how many tokens get verified -- so +scheduling is lossless with respect to the target distribution. + +Determinism note: ties are broken by ``(position, request)``, never by value. +Two TP ranks that compute bitwise-different confidences must still choose the +same verify lengths, or their batch shapes (and therefore their collectives) +diverge. See :func:`schedule_verify_lens_topk`. +""" + +from dataclasses import dataclass + +import torch + +__all__ = [ + "DSparkScheduleConfig", + "compute_survival", + "schedule_verify_lens_topk", +] + + +@dataclass(frozen=True) +class DSparkScheduleConfig: + """Bounds for the per-request verify length. + + Attributes: + block_size: draft block length ``K`` (the number of speculative tokens + the draft proposes per step). + min_verify_len: floor on every request's verify length. Must stay >= 1: + position 0 carries the bonus/anchor token, and a request that + verifies nothing makes no progress at all. Budget is allocated + *above* this floor. + max_verify_len: cap on every request's verify length; 0 means + ``block_size``. Never exceeds ``block_size`` -- a request cannot + verify more tokens than the draft proposed. + survival_eps: candidates whose survival falls below this are dropped + before ranking. Purely a numerical floor, not a tuning knob: it + keeps hopeless positions out of the ordering. + """ + + block_size: int + min_verify_len: int = 1 + max_verify_len: int = 0 + survival_eps: float = 1e-6 + + def __post_init__(self) -> None: + if self.block_size < 1: + raise ValueError(f"block_size must be >= 1, got {self.block_size}") + if self.min_verify_len < 1: + raise ValueError( + f"min_verify_len must be >= 1 (every request verifies its anchor " + f"token), got {self.min_verify_len}" + ) + if self.resolved_max_verify_len < self.min_verify_len: + raise ValueError( + f"max_verify_len ({self.resolved_max_verify_len}) < min_verify_len " + f"({self.min_verify_len})" + ) + if not 0.0 <= self.survival_eps < 1.0: + raise ValueError(f"survival_eps must be in [0, 1), got {self.survival_eps}") + + @property + def resolved_max_verify_len(self) -> int: + """``max_verify_len`` with 0 meaning "the whole block", clamped to ``block_size``.""" + cap = self.max_verify_len or self.block_size + return min(int(cap), int(self.block_size)) + + @property + def schedulable_per_request(self) -> int: + """Positions a single request can win from the budget, above the floor.""" + return self.resolved_max_verify_len - self.min_verify_len + + +def compute_survival(confidence: torch.Tensor) -> torch.Tensor: + """``[bs, K]`` conditional acceptance probabilities -> prefix survival. + + ``out[r][j] = prod_{i <= j} confidence[r][i]``, computed in fp32. The result + is non-increasing along ``j``, which :func:`schedule_verify_lens_topk` relies + on to keep its allocation prefix-shaped. + """ + if confidence.dim() != 2: + raise ValueError(f"confidence must be [bs, K], got shape {tuple(confidence.shape)}") + return torch.cumprod(confidence.to(torch.float32), dim=1) + + +def schedule_verify_lens_topk( + *, + survival: torch.Tensor, + budget: int, + cfg: DSparkScheduleConfig, +) -> torch.Tensor: + """Allocate ``budget`` verify tokens across the batch by global survival rank. + + Args: + survival: ``[bs, K]`` prefix-survival probabilities from + :func:`compute_survival`. + budget: total verify tokens to hand out *above* the per-request floor. + cfg: bounds; see :class:`DSparkScheduleConfig`. + Returns: + ``[bs]`` int32 verify lengths, each in + ``[min_verify_len, resolved_max_verify_len]``, summing to at most + ``bs * min_verify_len + budget``. + + Determinism: ranking uses a value-independent tie-break, so equal survivals + always resolve the same way regardless of float noise or backend. Positions + are ordered before requests, which biases ties toward *earlier* positions -- + the ones every downstream prefix depends on. + """ + if survival.dim() != 2: + raise ValueError(f"survival must be [bs, K], got shape {tuple(survival.shape)}") + bs, num_positions = survival.shape + if num_positions != cfg.block_size: + raise ValueError( + f"survival has {num_positions} positions but block_size is {cfg.block_size}" + ) + + device = survival.device + floor = int(cfg.min_verify_len) + schedulable = int(cfg.schedulable_per_request) + verify_lens = torch.full((bs,), floor, dtype=torch.int32, device=device) + budget = int(budget) + if bs == 0 or schedulable <= 0 or budget <= 0: + return verify_lens + + # Only positions above the floor are up for grabs: positions [0, floor) are + # already granted to every request, so scoring them would let a request win + # budget for something it is getting for free. + candidates = survival[:, floor : floor + schedulable].to(torch.float32) + flat = candidates.reshape(-1) + eligible = flat >= cfg.survival_eps + + # Rank by survival, then break ties deterministically. Sorting the tie-break + # key first and using a *stable* sort on the values makes the final order a + # pure function of (survival values, position, request) -- no dependence on + # the sort implementation's handling of equal keys. + # + # Ineligible candidates are scored -1 so they sort below every eligible one; + # the ``& eligible`` below then drops them even when ``budget`` is larger + # than the eligible count. That keeps the whole path free of any device->host + # sync (no ``.item()``), so it stays usable inside a captured graph. + scores = torch.where(eligible, flat, torch.full_like(flat, -1.0)) + tie_break = torch.arange(schedulable, device=device).repeat(bs) * bs + torch.arange( + bs, device=device + ).repeat_interleave(schedulable) + order = torch.argsort(tie_break, stable=True) + order = order[torch.argsort(scores[order], descending=True, stable=True)] + + rank = torch.empty_like(order) + rank[order] = torch.arange(order.numel(), device=device) + selected = (rank < budget) & eligible + + # survival is non-increasing along positions, so the chosen set for a request + # is necessarily a prefix -- counting per row recovers the full allocation. + granted = selected.view(bs, schedulable).sum(dim=1).to(torch.int32) + verify_lens += granted + return torch.clamp(verify_lens, min=floor, max=cfg.resolved_max_verify_len) diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py new file mode 100644 index 000000000000..e3d522464869 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -0,0 +1,232 @@ +# 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. +"""Host-side coordinator turning DSpark confidence into a per-iteration draft length. + +Three constraints shape this component, and all three point the same way: + +1. **The decision must be on the host.** TensorRT-LLM picks a decode CUDA graph + from ``(batch_size, draft_len, ...)`` -- every input to that choice is a host + value. A device-resident budget would have to be copied back before the batch + could be shaped, which is the synchronization we are trying to avoid. +2. **The decision must land on a captured draft length.** A length with no + captured graph does not raise; ``maybe_get_cuda_graph`` returns ``None`` and + the step silently runs eager, which costs far more than the tokens saved. + Hence :meth:`decide_draft_len` only ever returns a value from ``tiers``. +3. **Every rank must decide identically.** ``draft_len`` is part of the graph key + but is *not* part of the attention-DP consistency allgather, so two ranks that + pick different lengths select different graphs -- one replays, one falls back + to eager -- and their collectives diverge. :meth:`decide_draft_len` therefore + reduces across ranks before returning. + +Confidence for step ``t`` only becomes readable on the host at ``t+1`` (or ``t+2`` +with the overlap scheduler), so the decision always runs on a lagged snapshot. +That is fine -- it is a throughput heuristic, not a correctness input -- but the +snapshot has to be *taken* without blocking, and a snapshot that is not ready yet +must degrade to verifying everything rather than to a stale guess. +""" + +from typing import Callable, List, Optional, Sequence + +import numpy as np +import torch + +from ..._utils import prefer_pinned +from .dspark_planner import SpsCostTable, budget_argmax_over_uniform_lens +from .dspark_schedule import DSparkScheduleConfig, compute_survival, schedule_verify_lens_topk + +__all__ = ["DSparkVerifyPlanner"] + + +class DSparkVerifyPlanner: + """Stages confidence off the device and turns it into a draft length. + + Args: + cfg: verify-length bounds. + cost_table: measured step cost. A flat table means "not profiled"; the + planner then refuses to trim (see :meth:`decide_draft_len`). + tiers: draft lengths that have a captured CUDA graph. The returned length + is always one of these. + apply_calibration: maps raw confidence logits to probabilities. Normally + ``confidence_head.apply_sts``; defaults to a plain sigmoid. + all_rank_max: cross-rank reduction, returning the max of ``value`` over + ranks. ``None`` for single-rank runs. + """ + + def __init__( + self, + *, + cfg: DSparkScheduleConfig, + cost_table: Optional[SpsCostTable] = None, + tiers: Optional[Sequence[int]] = None, + apply_calibration: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + all_rank_max: Optional[Callable[[int], int]] = None, + ): + self.cfg = cfg + self.cost_table = cost_table + self.tiers: List[int] = sorted({int(t) for t in (tiers or [cfg.resolved_max_verify_len])}) + self.apply_calibration = apply_calibration or torch.sigmoid + self.all_rank_max = all_rank_max + + self._host_buffer: Optional[torch.Tensor] = None + self._copy_event: Optional[torch.cuda.Event] = None + self._snapshot_valid = False + # Observability: silent degradation is the main failure mode here, so + # count every path that gives up on trimming. + self.stats = { + "decisions": 0, + "fallback_no_snapshot": 0, + "fallback_flat_cost": 0, + "fallback_no_confidence": 0, + } + + @property + def max_tier(self) -> int: + return self.tiers[-1] if self.tiers else self.cfg.resolved_max_verify_len + + def stage_confidence(self, confidence_logits: torch.Tensor, num_rows: int) -> None: + """Start a non-blocking device->host copy of this step's confidence. + + Never synchronizes: the copy is issued on the current stream and an event + is recorded. :meth:`decide_draft_len` polls that event and simply does not + use the snapshot if it has not landed. + """ + if confidence_logits is None or num_rows <= 0: + return + rows = confidence_logits[:num_rows] + if self._host_buffer is None or self._host_buffer.shape != rows.shape: + self._host_buffer = torch.empty( + rows.shape, dtype=torch.float32, device="cpu", pin_memory=prefer_pinned() + ) + self._copy_event = torch.cuda.Event() + self._host_buffer.copy_(rows, non_blocking=True) + self._copy_event.record() + self._snapshot_valid = True + + def _ready_snapshot(self) -> Optional[torch.Tensor]: + if not self._snapshot_valid or self._host_buffer is None: + return None + # query() is non-blocking; a not-yet-landed copy just means this step + # decides without it. + if self._copy_event is not None and not self._copy_event.query(): + return None + return self._host_buffer + + def decide_draft_len( + self, + *, + num_gen_requests: int, + all_rank_max: Optional[Callable[[int], int]] = None, + ) -> int: + """Choose this iteration's draft length, agreed across all ranks. + + Degrades to the largest captured tier (i.e. verify everything, today's + behavior) whenever the inputs are not trustworthy: no confidence + snapshot yet, or an unprofiled (flat) cost model under which every extra + verified token looks free. + + ``all_rank_max`` overrides the constructor's reduction; the caller that + owns the distributed handle usually supplies it here. + """ + self.stats["decisions"] += 1 + chosen = self._decide_local(num_gen_requests=num_gen_requests) + all_rank_max = all_rank_max or self.all_rank_max + if all_rank_max is not None: + # Max, not min: a rank that wanted to trim more simply verifies a few + # extra tokens. Disagreeing on the graph key is the unrecoverable + # outcome, so agreement matters more than the saving. + chosen = int(all_rank_max(int(chosen))) + if chosen not in self.tiers: + chosen = min((t for t in self.tiers if t >= chosen), default=self.max_tier) + return int(chosen) + + def decide_verify_lens( + self, + *, + num_gen_requests: int, + all_rank_max: Optional[Callable[[int], int]] = None, + ) -> Optional[List[int]]: + """Per-request verify lengths for a ragged step, or None to stay uniform. + + The budget still comes from the cost-model argmax -- that is what knows + how many verified tokens the step can afford -- but instead of splitting + it evenly it is handed to the global survival top-k, so a request whose + draft is still alive at depth 5 can take positions from one whose draft + died at depth 1. + + Returns None whenever the uniform path would have fallen back (no + snapshot, flat cost table, empty batch): raggedness is a throughput + heuristic and an untrustworthy input must degrade to verifying + everything, not to a stale split. + + The batch-wide *maximum* is still reduced across ranks, because the + drafted-token buffer width and the per-request padding are sized from + it; the individual lengths stay local. + """ + if num_gen_requests <= 0: + return None + if self.cost_table is None or self.cost_table.is_flat: + self.stats["fallback_flat_cost"] += 1 + return None + snapshot = self._ready_snapshot() + if snapshot is None: + self.stats["fallback_no_snapshot"] += 1 + return None + rows = snapshot[:num_gen_requests] + if rows.numel() == 0: + self.stats["fallback_no_confidence"] += 1 + return None + + probs = self.apply_calibration(rows) + survival = compute_survival(probs) + uniform_len = budget_argmax_over_uniform_lens( + survival=survival.numpy().astype(np.float64), + num_gen_requests=int(num_gen_requests), + cost_table=self.cost_table, + allowed_lens=self.tiers, + min_verify_len=self.cfg.min_verify_len, + ) + # Same total the uniform decision would have spent, redistributed. + budget = int(num_gen_requests) * (int(uniform_len) - self.cfg.min_verify_len) + lens = schedule_verify_lens_topk(survival=survival, budget=budget, cfg=self.cfg).tolist() + + all_rank_max = all_rank_max or self.all_rank_max + if all_rank_max is not None: + agreed_max = int(all_rank_max(int(max(lens)))) + lens = [min(int(v), agreed_max) for v in lens] + return [int(v) for v in lens] + + def _decide_local(self, *, num_gen_requests: int) -> int: + if num_gen_requests <= 0: + return self.max_tier + if self.cost_table is None or self.cost_table.is_flat: + self.stats["fallback_flat_cost"] += 1 + return self.max_tier + snapshot = self._ready_snapshot() + if snapshot is None: + self.stats["fallback_no_snapshot"] += 1 + return self.max_tier + rows = snapshot[:num_gen_requests] + if rows.numel() == 0: + self.stats["fallback_no_confidence"] += 1 + return self.max_tier + probs = self.apply_calibration(rows) + survival = compute_survival(probs).numpy().astype(np.float64) + return budget_argmax_over_uniform_lens( + survival=survival, + num_gen_requests=int(num_gen_requests), + cost_table=self.cost_table, + allowed_lens=self.tiers, + min_verify_len=self.cfg.min_verify_len, + ) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 50e23d087ca7..691b2decd236 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -518,6 +518,31 @@ 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. ``runtime_draft_len`` above is one + # scalar for the whole iteration, which cannot express "this request + # verifies 5 positions and that one verifies 1". When confidence-scheduled + # verification runs in ragged mode these carry the per-request split: + # + # 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, so downstream code that only knows + # about ``runtime_draft_len`` keeps working unchanged. Built by + # ``_torch/speculative/dspark_ragged.RaggedVerifyLayout``. + verify_lens: Optional[torch.Tensor] = None + qo_indptr: Optional[torch.Tensor] = None + # Host-side ``sum(verify_lens)``. Kept as a Python int because the packed + # buffers have to be sliced before any kernel runs, and reading the total + # off ``verify_lens`` 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.""" + 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) and we can take # the argmax fast-path. False if any request needs sampling. @@ -716,7 +741,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.sampling_utils import \ GREEDY_TEMPERATURE_THRESHOLD from tensorrt_llm.sampling_params import SamplingParams @@ -784,8 +810,16 @@ def _normalize_request_sampling_params( tk_val = _first_or_none(sampling_config.top_k) tp_val = _first_or_none(sampling_config.top_p) - # Context requests have no draft tokens yet. - num_tokens = 1 + self.runtime_draft_len if request.state == LlmRequestState.GENERATION_IN_PROGRESS else 1 + # Context requests have no draft tokens yet. Generation requests + # contribute one row per verified position, which under ragged + # verification is per-request: these buffers are indexed by flat + # token position, so a fixed stride here would apply request r's + # temperature/top_k/top_p to request r+1's tokens. Note this whole + # block is skipped for all-greedy batches, so a greedy-only + # accuracy run (GSM8K) cannot exercise it. + num_tokens = get_request_tokens_per_gen_step( + request, 1 + self.runtime_draft_len + ) if request.state == LlmRequestState.GENERATION_IN_PROGRESS else 1 ( temp_val, @@ -997,6 +1031,30 @@ def populate_sampling_params_for_one_model( (tk for tk in request_top_ks if 0 < tk < _disable_topk), default=0) +def _padded_gen_draft_tokens(spec_metadata, num_gens: int, + runtime_draft_len: int) -> torch.Tensor: + """Unpack ragged-packed draft tokens into ``[num_gens, runtime_draft_len]``. + + The overlap path packs draft tokens by each request's own window, so the + buffer holds ``sum(verify_lens) - num_gens`` entries rather than a + ``num_gens x runtime_draft_len`` rectangle. Acceptance is naturally + rectangular, so scatter back out; positions past a request's window are + padding and get masked by ``count_accepted_ragged``. + """ + from .dspark_ragged import build_qo_indptr, scatter_ragged_to_padded + + # verify_lens counts tokens (bonus position + drafts); the draft buffer + # only holds the drafts. + draft_lens = spec_metadata.verify_lens - 1 + packed_len = spec_metadata.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. @@ -1373,16 +1431,37 @@ 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) - accepted_tokens[num_contexts:, :runtime_draft_len + - 1] = gen_target_tokens + if spec_metadata is not None and spec_metadata.is_ragged_verify: + from .dspark_ragged import (count_accepted_ragged, + scatter_ragged_to_padded) + verify_lens = spec_metadata.verify_lens + total = spec_metadata.total_verify_tokens + gen_target_tokens = scatter_ragged_to_padded( + target_tokens[num_contexts:num_contexts + total], + verify_lens=verify_lens, + qo_indptr=spec_metadata.qo_indptr, + max_len=runtime_draft_len + 1) + accepted_tokens[num_contexts:, :runtime_draft_len + + 1] = gen_target_tokens + # 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) + 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) + # 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( @@ -1391,6 +1470,15 @@ 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 _accept_draft_tokens(self, logits, draft_tokens, num_contexts, @@ -1552,6 +1640,14 @@ 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: + # The buffer is packed by each request's own draft window, so it is + # not a [num_gens, runtime_draft_len] rectangle. Reshaping it would + # not even fail loudly -- fill_bucket pads the total to a captured + # bucket, which can happen to divide -- it would just 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, @@ -1674,7 +1770,12 @@ 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 @@ -1690,9 +1791,19 @@ def _sample_and_accept_draft_tokens_rejection( target_probs_flat = compute_probs_from_logits( gen_logits, temperatures, top_ks, top_ps) - target_probs = target_probs_flat.reshape(num_gens, - runtime_draft_len + 1, - vocab_size) + if is_ragged: + from .dspark_ragged import (fill_padded_rows_onehot, + scatter_ragged_to_padded) + target_probs = scatter_ragged_to_padded( + target_probs_flat, + verify_lens=spec_metadata.verify_lens, + qo_indptr=spec_metadata.qo_indptr, + max_len=runtime_draft_len + 1) + fill_padded_rows_onehot(target_probs, + verify_lens=spec_metadata.verify_lens) + else: + target_probs = target_probs_flat.reshape( + num_gens, runtime_draft_len + 1, vocab_size) draft_vocab_size = draft_probs.shape[-1] assert draft_probs.shape[0] == num_gens, ( @@ -1776,6 +1887,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 _update_advance_draft_sampling_seed(self, device): diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 3d841aaa4a38..bdc598873c15 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -175,6 +175,26 @@ def _request_common_handling( request.py_draft_tokens = next_draft_tokens[request.py_seq_slot][:runtime_draft_len] request.py_decoding_iter += 1 + @staticmethod + def _verified_len(request: LlmRequest, runtime_draft_len: int) -> int: + """How many draft positions THIS request was given to verify. + + Uniform scheduling gives every request the batch-wide + ``runtime_draft_len``; ragged (per-request) verification does not, and + the KV rewind below is computed per request: + + rewind = verified_positions - accepted_positions + + Using the batch-wide value under ragged verification rewinds the wrong + amount for every request that got a shorter window -- silently dropping + or keeping KV entries, i.e. wrong output with no error anywhere. + + ``py_verify_len`` is set only by the ragged path, so this is exactly + today's behavior for every other speculation mode. + """ + per_request = getattr(request, "py_verify_len", None) + return runtime_draft_len if per_request is None else int(per_request) + def update_requests( self, state: SampleStateSpec, @@ -208,7 +228,9 @@ def update_requests( ): break req.py_num_accepted_draft_tokens = num_new_tokens - 1 - req.py_rewind_len = runtime_draft_len - req.py_num_accepted_draft_tokens + req.py_rewind_len = ( + self._verified_len(req, runtime_draft_len) - req.py_num_accepted_draft_tokens + ) self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) def sample_async( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 68d9fdbeb207..05b74ebbb779 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2759,13 +2759,13 @@ class DSparkDecodingConfig(DecodingBaseConfig): layers as cross-attention context and drafts a whole block in a single backbone forward, but it additionally refines the per-position draft logits with a lightweight sequential head (a low-rank Markov head, optionally an RNN - head) and predicts an acceptance-confidence per position to truncate the - proposed prefix. + head) and predicts a per-position acceptance confidence used to budget how + much verification the batch is worth. Key features: - Target-dependent: captures hidden states from ``target_layer_ids``. - Semi-parallel: one block backbone forward + cheap sequential head refine. - - Confidence head: truncates the proposed draft length (NOT the accept rule; + - Confidence head: sizes the *verification* budget (NOT the accept rule; acceptance stays standard target verification, preserving greedy parity). Reference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec). @@ -2802,16 +2802,108 @@ class DSparkDecodingConfig(DecodingBaseConfig): "read from the draft model config (dspark_markov_head_type), " "defaulting to \"vanilla\".") - # NOTE: confidence-based dynamic drafting (the draft model's confidence head - # that truncates the proposed block) is NOT enabled in this PR. The user-facing - # ``enable_confidence_head`` / ``confidence_threshold`` knobs are intentionally - # omitted and will be added when the feature is actually wired into the - # speculative scheduling/verification path. The confidence head module and its - # internal plumbing remain as scaffolding (see DSparkConfidenceHead / - # dspark_propose). + enable_confidence_scheduling: bool = Field( + default=False, + description= + "Score each drafted position with the draft's confidence head and let " + "the verification scheduler size the batch's verify budget from those " + "scores. The block is always drafted in full; only the number of tokens " + "sent to the target for verification changes, so acceptance (and hence " + "output distribution) is unaffected. Requires a checkpoint with trained " + "confidence-head weights. Gains only appear at high concurrency -- at " + "small batch sizes the step cost is nearly flat in the verified token " + "count, so trimming saves nothing and costs acceptance length.", + status="prototype") + + confidence_sts_path: Optional[str] = Field( + default=None, + description= + "Path to a JSON file with per-position sequential-temperature-scaling " + "(STS) values used to calibrate the confidence head. The scheduler " + "consumes a cumulative product of per-position probabilities, so a " + "per-position calibration error compounds geometrically along the " + "block. If None, the raw sigmoid is used (identity calibration).", + status="prototype") + + confidence_sps_table_path: Optional[str] = Field( + default=None, + description= + "Path to a JSON file with the profiled decode step cost as a function of " + "the total verified token count. Without it the planner has a flat cost " + "model, under which every extra verified token looks free and the budget " + "degenerates to verify-all (no scheduling gain).", + status="prototype") + + confidence_verify_len_tiers: Optional[List[PositiveInt]] = Field( + default=None, + description= + "Draft lengths to capture CUDA graphs for. The scheduler may only pick " + "from this ladder: a length with no captured graph does not raise, it " + "silently drops the step out of graph replay into eager execution, which " + "costs far more than the trimmed tokens save. Each extra tier is another " + "captured graph, and captured graphs consume memory that would otherwise " + "be KV cache, so keep the ladder short. If None, defaults to " + "[1, ceil(max_draft_len/2), max_draft_len].", + status="prototype") + + enable_ragged_verify: bool = Field( + default=False, + description= + "Give each request its own verify length instead of one length for the " + "whole batch. The batch's verify budget is then allocated by global " + "survival rank, so a request whose draft is still likely alive at depth " + "5 can take positions from one whose draft already died at depth 1. " + "Requires enable_confidence_scheduling. The packed batch is padded to a " + "captured (batch size, token count) pair, so this captures a second " + "graph axis and costs more graph memory than uniform scheduling.", + status="prototype") decoding_type: Literal["DSpark"] = Field(default="DSpark") + @model_validator(mode="after") + def validate_ragged_verify(self): + if self.enable_ragged_verify and not self.enable_confidence_scheduling: + raise ValueError( + "enable_ragged_verify requires enable_confidence_scheduling=True: " + "per-request verify lengths are chosen from confidence-head " + "survival, which is only computed when scheduling is on") + return self + + @model_validator(mode="after") + def validate_confidence_scheduling(self): + if not self.enable_confidence_scheduling: + if (self.confidence_sts_path or self.confidence_sps_table_path + or self.confidence_verify_len_tiers): + raise ValueError( + "confidence_sts_path / confidence_sps_table_path / " + "confidence_verify_len_tiers require " + "enable_confidence_scheduling=True") + return self + + if self.max_draft_len is not None: + tiers = self.confidence_verify_len_tiers or [ + 1, max(1, (self.max_draft_len + 1) // 2), self.max_draft_len + ] + tiers = sorted({int(t) for t in tiers}) + if tiers[-1] > self.max_draft_len: + raise ValueError( + f"confidence_verify_len_tiers {tiers} exceeds max_draft_len " + f"({self.max_draft_len}); a request cannot verify more tokens " + f"than the draft proposes") + # The full block must stay reachable: it is the fallback used whenever + # the confidence snapshot is stale or the cost model is unprofiled. + if tiers[-1] != self.max_draft_len: + tiers.append(self.max_draft_len) + self.confidence_verify_len_tiers = tiers + return self + + @property + def verify_len_tiers(self) -> List[int]: + """Captured draft-length ladder; ``[max_draft_len]`` when scheduling is off.""" + if not self.enable_confidence_scheduling: + return [self.max_draft_len] + return list(self.confidence_verify_len_tiers or [self.max_draft_len]) + @model_validator(mode="after") def set_max_total_draft_tokens(self): self.max_total_draft_tokens = self.max_draft_len diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index f20e02169d62..73ad68b249ea 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1657,6 +1657,13 @@ "kind": "value", "path": "speculative_config.block_size" }, + { + "allowed_values": [], + "annotation": "Optional[List[Annotated[int, Gt(gt=0)]]]", + "converter": "", + "kind": "value", + "path": "speculative_config.confidence_verify_len_tiers" + }, { "allowed_values": [ "AUTO", @@ -1717,6 +1724,13 @@ "kind": "value", "path": "speculative_config.eagle_choices" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "speculative_config.enable_confidence_scheduling" + }, { "allowed_values": [], "annotation": "", @@ -1724,6 +1738,13 @@ "kind": "value", "path": "speculative_config.enable_global_pool" }, + { + "allowed_values": [], + "annotation": "", + "converter": "", + "kind": "value", + "path": "speculative_config.enable_ragged_verify" + }, { "allowed_values": [], "annotation": "", diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index ad192d34d9a5..33e14f8b4b8b 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4001,6 +4001,57 @@ def test_gsm8k_dep8_megamoe_deepgemm(self): f"GSM8K accuracy {score:.3f} is below recorded reference " f"{acc_params.ref_accuracy:.3f}") + @pytest.mark.skip_less_mpi_world_size(8) + def test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling(self): + """Same accuracy bar with confidence-scheduled verification enabled. + + The scheduler only decides how MANY drafted tokens are sent to the + target; acceptance still goes through the unchanged target verify, so + the output distribution must be untouched. Any drop here means the + dynamic draft length broke an invariant (KV rewind amount, the draft + hidden-state gather stride, or the per-token sampling-parameter stride) + rather than that the scheduler made a bad throughput call -- those are + silent-wrong-answer bugs, which is why this asserts the SAME reference + accuracy as the baseline above rather than a relaxed one. + + Runs under attention-DP + TP8, so it also covers the cross-rank + agreement on the draft length: if ranks disagreed they would select + different CUDA graphs and the collectives would diverge. + """ + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.5) + spec_config = DSparkDecodingConfig(max_draft_len=5, + speculative_model=self.MODEL_PATH, + enable_confidence_scheduling=True) + with LLM(self.MODEL_PATH, + attn_backend="TRTLLM", + tensor_parallel_size=8, + moe_expert_parallel_size=8, + enable_attention_dp=True, + moe_config=MoeConfig(backend="MEGAMOE_DEEPGEMM"), + max_batch_size=DEEPSEEKV4_TEST_MAX_BATCH_SIZE, + max_seq_len=4096, + max_num_tokens=4096, + kv_cache_config=kv_cache_config, + enable_chunked_prefill=False, + disable_overlap_scheduler=True, + custom_tokenizer="deepseek_v4", + speculative_config=spec_config) as llm: + task = GSM8K(self.MODEL_NAME) + acc_params = task.get_hypothesis_testing_params( + dtype=llm.args.dtype, + quant_algo=llm.args.quant_config.quant_algo, + kv_cache_quant_algo=llm.args.quant_config.kv_cache_quant_algo, + spec_dec_algo=llm.args.speculative_config.decoding_type) + assert acc_params.num_samples == GSM8K.NUM_SAMPLES + with mock.patch.dict(os.environ, {"INTEGRATION_TEST": "0"}): + score = task.evaluate( + llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + assert score >= acc_params.ref_accuracy, ( + f"GSM8K accuracy {score:.3f} with confidence scheduling is below " + f"the recorded reference {acc_params.ref_accuracy:.3f}; the " + f"scheduler must not change the output distribution") + @pytest.mark.timeout(14400) @pytest.mark.skip_less_device_memory(140000) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py new file mode 100644 index 000000000000..33658caf198e --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py @@ -0,0 +1,696 @@ +# 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. +"""Unit tests for DSpark confidence-scheduled verification (hardware-agnostic, CPU). + +These cover the parts where a silent error would be expensive and invisible at +runtime: the survival semantics, the budget argmax over a staircase cost curve, +the prefix property of the allocation, and cross-rank determinism. +""" + +import itertools + +import numpy as np +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, + budget_argmax_over_uniform_lens, + compute_verify_token_budget, + derive_verify_len_tiers, +) +from tensorrt_llm._torch.speculative.dspark_schedule import ( + DSparkScheduleConfig, + compute_survival, + schedule_verify_lens_topk, +) + +BLOCK = 7 + + +def _cfg(**kw) -> DSparkScheduleConfig: + return DSparkScheduleConfig(block_size=BLOCK, **kw) + + +# -------------------------------------------------------------------------- +# survival +# -------------------------------------------------------------------------- + + +def test_survival_is_cumulative_product(): + conf = torch.tensor([[0.9, 0.8, 0.5, 1.0, 1.0, 1.0, 1.0]]) + surv = compute_survival(conf) + assert torch.allclose(surv[0, :3], torch.tensor([0.9, 0.72, 0.36]), atol=1e-6) + + +def test_survival_is_non_increasing(): + """The allocator's prefix property depends on this holding for any input.""" + torch.manual_seed(0) + conf = torch.rand(16, BLOCK) + surv = compute_survival(conf) + assert torch.all(surv[:, 1:] <= surv[:, :-1] + 1e-7) + + +def test_survival_rejects_wrong_rank(): + with pytest.raises(ValueError, match=r"\[bs, K\]"): + compute_survival(torch.rand(BLOCK)) + + +# -------------------------------------------------------------------------- +# budget argmax +# -------------------------------------------------------------------------- + + +def _brute_force_budget(survival, num_gen, table, min_len=1, max_len=None): + """Independent O(N) reimplementation used as the oracle for the argmax. + + Spells out the objective: expected emitted tokens per millisecond, where the + expectation counts one bonus token per request, the always-verified floor + positions, and the admitted candidates. + """ + bs, blk = survival.shape + cap = min(max_len or blk, blk) + cand = np.sort(survival[:, min_len:cap].reshape(-1))[::-1] + base = num_gen + float(survival[:, :min_len].sum()) + best_n, best = 0, -np.inf + for n in range(cand.size + 1): + theta = (base + float(cand[:n].sum())) / table.step_time(bs * min_len + n) + if theta > best: + best, best_n = theta, n + return best_n + + +def test_tau_counts_the_always_verified_floor_positions(): + """The floor is verified for free, but its yield still belongs in tau. + + argmax((C + f(n)) / g(n)) depends on C, so dropping the floor's contribution + silently moves the optimum. Two batches that differ only in how good their + floor positions are must be able to choose different budgets. + """ + table = SpsCostTable(token_counts=(0, 6), step_time_ms=(1.0, 3.0)) + bs = 4 + good_floor = np.cumprod(np.full((bs, BLOCK), 0.99), axis=1) + bad_floor = good_floor.copy() + bad_floor[:, 0] = 0.01 # a hopeless first position drags the whole row down + bad_floor = np.cumprod( + np.concatenate([bad_floor[:, :1], np.full((bs, BLOCK - 1), 0.99)], axis=1), axis=1 + ) + for surv in (good_floor, bad_floor): + assert compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=table + ) == _brute_force_budget(surv, bs, table) + + +def test_budget_matches_brute_force_on_staircase(): + """Theta is not unimodal on a staircase; a greedy first-descent scan is wrong.""" + rng = np.random.default_rng(7) + # A cost curve with genuine shelves and risers. + table = SpsCostTable( + token_counts=(0, 64, 128, 192, 256, 320), + step_time_ms=(4.0, 4.05, 6.5, 6.55, 9.0, 9.05), + ) + for trial in range(25): + bs = int(rng.integers(1, 40)) + conf = rng.uniform(0.55, 0.99, size=(bs, BLOCK)) + surv = np.cumprod(conf, axis=1) + got = compute_verify_token_budget(survival=surv, num_gen_requests=bs, cost_table=table) + want = _brute_force_budget(surv, bs, table) + assert got == want, f"trial {trial}: bs={bs} got={got} want={want}" + + +def test_greedy_first_descent_would_be_wrong(): + """Pin the specific failure mode the global argmax exists to avoid. + + Theta rises along a shelf, drops at a riser, then rises again along the next + shelf. A greedy loop that stops at the first non-improvement (as the paper's + pseudocode does) parks at the end of the *first* shelf; here the second + shelf's optimum is almost 3x better. + """ + bs = 4 + # One riser at 6 total tokens, then a long shelf out to the end of the range. + table = SpsCostTable(token_counts=(0, 6), step_time_ms=(1.0, 3.0)) + # Uniformly high confidence: every extra candidate adds nearly a full token, + # so tau grows ~7x across the range while cost only doubles. + surv = np.cumprod(np.full((bs, BLOCK), 0.995), axis=1) + n_star = compute_verify_token_budget(survival=surv, num_gen_requests=bs, cost_table=table) + + cand = np.sort(surv[:, 1:].reshape(-1))[::-1] + + def theta(n): + return (bs + cand[:n].sum()) / table.step_time(bs + n) + + greedy = 0 + while greedy + 1 <= cand.size and theta(greedy + 1) > theta(greedy): + greedy += 1 + + assert theta(n_star) >= theta(greedy) + assert n_star > greedy, "expected the staircase to defeat a first-descent scan" + assert theta(n_star) > 2 * theta(greedy), "the gap should be large, not marginal" + + +def test_flat_cost_table_degenerates_to_verify_all(): + """A flat cost model makes every token free -> spend the whole budget.""" + table = SpsCostTable.flat(2.0) + assert table.is_flat + surv = np.cumprod(np.full((6, BLOCK), 0.9), axis=1) + budget = compute_verify_token_budget(survival=surv, num_gen_requests=6, cost_table=table) + assert budget == 6 * (BLOCK - 1) + + +def test_low_confidence_yields_small_budget(): + table = SpsCostTable(token_counts=(0, 8, 16, 24, 32), step_time_ms=(2.0, 3.0, 4.0, 5.0, 6.0)) + hi = np.cumprod(np.full((8, BLOCK), 0.98), axis=1) + lo = np.cumprod(np.full((8, BLOCK), 0.30), axis=1) + b_hi = compute_verify_token_budget(survival=hi, num_gen_requests=8, cost_table=table) + b_lo = compute_verify_token_budget(survival=lo, num_gen_requests=8, cost_table=table) + assert b_lo < b_hi + + +def test_budget_is_zero_for_empty_batch(): + table = SpsCostTable.flat() + assert ( + compute_verify_token_budget( + survival=np.zeros((0, BLOCK)), num_gen_requests=0, cost_table=table + ) + == 0 + ) + + +def test_cost_table_is_a_staircase_not_an_interpolation(): + table = SpsCostTable(token_counts=(0, 100), step_time_ms=(1.0, 5.0)) + # Everything below the next breakpoint costs the same as the shelf it is on. + assert table.step_time(1) == table.step_time(99) + assert table.step_time(100) > table.step_time(99) + # Above the last breakpoint we clamp rather than extrapolate. + assert table.step_time(10_000) == table.step_time(100) + + +def test_cost_table_rejects_malformed_input(): + with pytest.raises(ValueError, match="strictly increasing"): + SpsCostTable(token_counts=(0, 0), step_time_ms=(1.0, 2.0)) + with pytest.raises(ValueError, match="same length"): + SpsCostTable(token_counts=(0, 1), step_time_ms=(1.0,)) + with pytest.raises(ValueError, match="positive"): + SpsCostTable(token_counts=(0,), step_time_ms=(0.0,)) + + +# -------------------------------------------------------------------------- +# per-request allocation +# -------------------------------------------------------------------------- + + +def test_allocation_is_a_prefix_and_respects_bounds(): + torch.manual_seed(3) + cfg = _cfg(min_verify_len=1) + surv = compute_survival(torch.rand(12, BLOCK)) + lens = schedule_verify_lens_topk(survival=surv, budget=20, cfg=cfg) + assert lens.dtype == torch.int32 + assert torch.all(lens >= cfg.min_verify_len) + assert torch.all(lens <= cfg.resolved_max_verify_len) + # Budget is spent above the floor and never exceeded. + assert int((lens - cfg.min_verify_len).sum()) <= 20 + + +def test_allocation_prefers_the_most_confident_requests(): + cfg = _cfg(min_verify_len=1) + conf = torch.tensor( + [ + [1.0, 0.99, 0.99, 0.99, 0.99, 0.99, 0.99], # confident + [1.0, 0.10, 0.10, 0.10, 0.10, 0.10, 0.10], # not + ] + ) + lens = schedule_verify_lens_topk(survival=compute_survival(conf), budget=6, cfg=cfg) + assert lens[0] > lens[1] + + +def test_zero_budget_gives_everyone_the_floor(): + cfg = _cfg(min_verify_len=2) + surv = compute_survival(torch.rand(5, BLOCK)) + lens = schedule_verify_lens_topk(survival=surv, budget=0, cfg=cfg) + assert torch.equal(lens, torch.full((5,), 2, dtype=torch.int32)) + + +def test_huge_budget_saturates_at_the_cap(): + cfg = _cfg(min_verify_len=1, max_verify_len=4) + surv = compute_survival(torch.full((5, BLOCK), 0.99)) + lens = schedule_verify_lens_topk(survival=surv, budget=10**6, cfg=cfg) + assert torch.equal(lens, torch.full((5,), 4, dtype=torch.int32)) + + +def test_survival_eps_excludes_hopeless_positions(): + """Even with unlimited budget, sub-eps candidates are never admitted.""" + cfg = _cfg(min_verify_len=1, survival_eps=1e-3) + conf = torch.tensor([[1.0, 1.0, 1e-9, 1.0, 1.0, 1.0, 1.0]]) + lens = schedule_verify_lens_topk(survival=compute_survival(conf), budget=10**6, cfg=cfg) + # positions 0,1 survive; from position 2 on, survival collapses below eps. + assert int(lens[0]) == 2 + + +def test_allocation_is_deterministic_under_tied_confidences(): + """Two ranks with identical input must produce identical verify lengths.""" + cfg = _cfg(min_verify_len=1) + surv = compute_survival(torch.full((9, BLOCK), 0.9)) + first = schedule_verify_lens_topk(survival=surv, budget=13, cfg=cfg) + for _ in range(8): + assert torch.equal(schedule_verify_lens_topk(survival=surv, budget=13, cfg=cfg), first) + + +def test_ties_break_toward_earlier_positions(): + """A tie must not hand a later position to one request over an earlier one. + + Every request has identical confidence, so all candidates at a given + position tie. The allocation must stay balanced (a prefix front), not give + one request a deep suffix while another gets nothing. + """ + cfg = _cfg(min_verify_len=1) + surv = compute_survival(torch.full((4, BLOCK), 1.0)) + lens = schedule_verify_lens_topk(survival=surv, budget=4, cfg=cfg) + assert int(lens.max()) - int(lens.min()) <= 1 + + +def test_empty_batch_returns_empty(): + cfg = _cfg() + lens = schedule_verify_lens_topk(survival=torch.zeros(0, BLOCK), budget=5, cfg=cfg) + assert lens.numel() == 0 + + +def test_schedule_config_rejects_bad_bounds(): + with pytest.raises(ValueError, match="min_verify_len must be >= 1"): + DSparkScheduleConfig(block_size=BLOCK, min_verify_len=0) + with pytest.raises(ValueError, match="block_size must be >= 1"): + DSparkScheduleConfig(block_size=0) + with pytest.raises(ValueError, match="< min_verify_len"): + DSparkScheduleConfig(block_size=BLOCK, min_verify_len=5, max_verify_len=3) + + +# -------------------------------------------------------------------------- +# budget conservation: allocation must never exceed what the planner sized +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bs,min_len", list(itertools.product([1, 3, 17], [1, 2]))) +def test_planner_and_allocator_agree_on_the_budget(bs, min_len): + rng = np.random.default_rng(bs * 100 + min_len) + table = SpsCostTable(token_counts=(0, 16, 32, 64, 128), step_time_ms=(3.0, 3.1, 5.0, 5.1, 9.0)) + conf = rng.uniform(0.4, 0.99, size=(bs, BLOCK)) + surv_np = np.cumprod(conf, axis=1) + budget = compute_verify_token_budget( + survival=surv_np, num_gen_requests=bs, cost_table=table, min_verify_len=min_len + ) + lens = schedule_verify_lens_topk( + survival=torch.from_numpy(surv_np).float(), + budget=budget, + cfg=_cfg(min_verify_len=min_len), + ) + assert int((lens - min_len).sum()) <= budget + assert torch.all(lens >= min_len) + + +# -------------------------------------------------------------------------- +# uniform-K projection (what TRT-LLM's graph key can actually express today) +# -------------------------------------------------------------------------- + + +def test_discrete_argmax_beats_round_then_snap_across_a_riser(): + """Optimizing over runnable lengths is not the same as rounding the budget. + + The continuous optimum sits just past a cost riser; snapping down lands on a + length whose Theta is worse than a length the discrete search finds. + """ + table = SpsCostTable(token_counts=(0, 8, 24, 40), step_time_ms=(2.0, 2.05, 2.1, 20.0)) + surv = np.cumprod(np.full((8, BLOCK), 0.97), axis=1) + chosen = budget_argmax_over_uniform_lens( + survival=surv, + num_gen_requests=8, + cost_table=table, + allowed_lens=[1, 3, 5, 7], + ) + # 8*7=56 tokens lands on the 20ms riser; the search must avoid it. + assert chosen < 7 + + def theta(length): + tau = 8 + float(surv[:, 1:length].sum()) + return tau / table.step_time(8 * length) + + assert theta(chosen) == max(theta(v) for v in (1, 3, 5, 7)) + + +def test_derived_tiers_lose_nothing_versus_continuous_k(): + """The shelf-right-edge property: derived tiers contain the true optimum. + + Within a cost shelf the step time is constant while tau strictly increases, + so Theta rises monotonically across the shelf and the optimum can only sit + at a right edge. A tier set built from the right edges must therefore match + an exhaustive search over *every* length. + """ + rng = np.random.default_rng(11) + table = SpsCostTable( + token_counts=(0, 12, 30, 48, 90, 160), + step_time_ms=(2.0, 2.4, 3.9, 4.0, 7.5, 12.0), + ) + for bs in (3, 8, 16, 31): + surv = np.cumprod(rng.uniform(0.5, 0.995, size=(bs, BLOCK)), axis=1) + tiers = derive_verify_len_tiers( + cost_table=table, num_requests=bs, block_size=BLOCK, max_tiers=99 + ) + + def theta(length): + return (bs + float(surv[:, :length].sum())) / table.step_time(bs * length) + + best_any = max(theta(v) for v in range(1, BLOCK + 1)) + best_tier = max(theta(v) for v in tiers) + assert best_tier == pytest.approx(best_any), f"bs={bs} tiers={tiers}" + + +def test_derived_tiers_respect_the_capture_budget(): + table = SpsCostTable( + token_counts=(0, 5, 9, 14, 20, 27, 35), step_time_ms=(1.0, 2, 3, 4, 5, 6, 7) + ) + tiers = derive_verify_len_tiers(cost_table=table, num_requests=2, block_size=BLOCK, max_tiers=3) + assert len(tiers) <= 3 + # Endpoints are always kept: the floor is always runnable and the full block + # is always a right edge (the last shelf is unbounded). + assert tiers[0] == 1 and tiers[-1] == BLOCK + + +def test_derived_tiers_are_a_function_of_batch_size(): + """Tiers move with bs because total tokens = bs * length. + + With a riser at 40 tokens: at bs=8 the shelf ends at length 4, so 4 is a + tier; at bs=32 even length 2 is already past the riser, so the shelf yields + nothing and only the endpoints remain. A tier set derived once and reused + across batch sizes would be wrong for one of them. + """ + table = SpsCostTable(token_counts=(0, 40), step_time_ms=(1.0, 5.0)) + assert derive_verify_len_tiers( + cost_table=table, num_requests=8, block_size=BLOCK, max_tiers=99 + ) == [1, 4, BLOCK] + assert derive_verify_len_tiers( + cost_table=table, num_requests=32, block_size=BLOCK, max_tiers=99 + ) == [1, BLOCK] + + +def test_flat_cost_table_derives_only_the_endpoints(): + """No shelves to find -> nothing to derive beyond floor and full block.""" + tiers = derive_verify_len_tiers( + cost_table=SpsCostTable.flat(), num_requests=8, block_size=BLOCK, max_tiers=99 + ) + assert tiers == [1, BLOCK] + + +def test_discrete_argmax_falls_back_when_nothing_is_allowed(): + table = SpsCostTable.flat() + surv = np.cumprod(np.full((4, BLOCK), 0.9), axis=1) + assert ( + budget_argmax_over_uniform_lens( + survival=surv, num_gen_requests=4, cost_table=table, allowed_lens=[] + ) + == 1 + ) + + +# -------------------------------------------------------------------------- +# host-side planner: relay, fallbacks, cross-rank agreement +# -------------------------------------------------------------------------- + + +def _planner(**kw): + from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + + kw.setdefault("cfg", _cfg(min_verify_len=1)) + kw.setdefault( + "cost_table", + SpsCostTable(token_counts=(0, 8, 24, 48), step_time_ms=(2.0, 2.1, 4.0, 9.0)), + ) + kw.setdefault("tiers", [1, 3, BLOCK]) + return DSparkVerifyPlanner(**kw) + + +def test_planner_returns_max_tier_without_a_snapshot(): + """No confidence yet must mean verify-all, never a stale guess.""" + p = _planner() + assert p.decide_draft_len(num_gen_requests=8) == BLOCK + assert p.stats["fallback_no_snapshot"] == 1 + + +def test_planner_refuses_to_trim_on_an_unprofiled_cost_model(): + """A flat cost table makes every token look free -- trimming would be blind.""" + p = _planner(cost_table=SpsCostTable.flat()) + assert p.decide_draft_len(num_gen_requests=8) == BLOCK + assert p.stats["fallback_flat_cost"] == 1 + p2 = _planner(cost_table=None) + assert p2.decide_draft_len(num_gen_requests=8) == BLOCK + + +def test_planner_only_ever_returns_a_captured_tier(): + """Returning an uncaptured length silently drops the step out of the graph.""" + torch.manual_seed(5) + tiers = [2, 4] + for trial in range(12): + p = _planner(tiers=tiers) + conf = torch.rand(6, BLOCK) + p._host_buffer = conf.float() + p._copy_event = None + p._snapshot_valid = True + assert p.decide_draft_len(num_gen_requests=6) in tiers, f"trial {trial}" + + +def test_planner_trims_when_confidence_is_low(): + hi, lo = torch.full((8, BLOCK), 6.0), torch.full((8, BLOCK), -6.0) + out = [] + for logits in (hi, lo): + p = _planner() + p._host_buffer, p._copy_event, p._snapshot_valid = logits, None, True + out.append(p.decide_draft_len(num_gen_requests=8)) + assert out[1] <= out[0], f"low confidence should not verify more: {out}" + + +def test_planner_cross_rank_reduction_forces_agreement(): + """Ranks must converge on one length even with different local confidence.""" + peer_choices = {} + + def make(rank, logits): + p = _planner(all_rank_max=lambda v, r=rank: max(v, peer_choices.get(1 - r, v))) + p._host_buffer, p._copy_event, p._snapshot_valid = logits, None, True + return p + + r0 = _planner() + r0._host_buffer, r0._copy_event, r0._snapshot_valid = torch.full((4, BLOCK), -6.0), None, True + peer_choices[0] = r0._decide_local(num_gen_requests=4) + r1 = _planner() + r1._host_buffer, r1._copy_event, r1._snapshot_valid = torch.full((4, BLOCK), 6.0), None, True + peer_choices[1] = r1._decide_local(num_gen_requests=4) + + a = make(0, torch.full((4, BLOCK), -6.0)).decide_draft_len(num_gen_requests=4) + b = make(1, torch.full((4, BLOCK), 6.0)).decide_draft_len(num_gen_requests=4) + assert a == b, f"ranks disagreed on the graph key: {a} vs {b}" + assert a == max(peer_choices.values()) + + +def test_planner_cross_rank_result_is_snapped_back_onto_a_tier(): + """A reduction can land off-ladder; the result must still be captured.""" + p = _planner(tiers=[1, 3, BLOCK], all_rank_max=lambda v: 5) # 5 is not a tier + p._host_buffer, p._copy_event, p._snapshot_valid = torch.full((4, BLOCK), 0.0), None, True + assert p.decide_draft_len(num_gen_requests=4) == BLOCK + + +def test_planner_empty_batch_is_max_tier(): + assert _planner().decide_draft_len(num_gen_requests=0) == BLOCK + + +def test_planner_uses_the_supplied_calibration(): + """STS must be applied before the cumprod, not skipped.""" + seen = {} + + def calib(x): + seen["called"] = True + return torch.sigmoid(x / 4.0) + + p = _planner(apply_calibration=calib) + p._host_buffer, p._copy_event, p._snapshot_valid = torch.full((4, BLOCK), 2.0), None, True + p.decide_draft_len(num_gen_requests=4) + assert seen.get("called") + + +# -------------------------------------------------------------------------- +# Acceptance criterion (2): CUDA-graph compatibility +# +# The failure mode this guards is silent, not loud: a runtime draft length with +# no captured graph does not raise -- maybe_get_cuda_graph returns None and the +# step runs eager, costing far more than the trimmed tokens save. So the capture +# set and the planner's output set must be provably identical. +# -------------------------------------------------------------------------- + + +def _capture_set(batch_sizes, tiers, max_draft_len=BLOCK): + """Invoke PyTorchModelEngine._get_graphs_to_capture with a stub engine.""" + import types + + from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + + spec_config = types.SimpleNamespace( + enable_confidence_scheduling=True, + verify_len_tiers=list(tiers), + draft_len_schedule=None, + max_concurrency=None, + is_linear_tree=True, + spec_dec_mode=types.SimpleNamespace( + support_dynamic_draft_len=lambda: False, use_one_engine=lambda: False + ), + ) + engine = types.SimpleNamespace( + is_draft_model=False, + spec_config=spec_config, + max_draft_len=max_draft_len, + max_total_draft_tokens=max_draft_len, + original_max_draft_len=max_draft_len, + _dynamic_draft_len_mapping=None, + _cuda_graph_batch_sizes=list(batch_sizes), + ) + return PyTorchModelEngine._get_graphs_to_capture(engine, list(batch_sizes), None) + + +def test_capture_is_the_batch_size_x_tier_cross_product(): + """K is NOT a function of batch size here -- every bs needs every tier.""" + bss, tiers = [1, 8, 64], [1, 3, BLOCK] + graphs = _capture_set(bss, tiers) + assert set(graphs) == {(bs, k) for bs in bss for k in tiers} + assert len(graphs) == len(bss) * len(tiers) + + +def test_every_length_the_planner_can_pick_has_a_captured_graph(): + """The anti-silent-eager invariant, checked against the planner itself.""" + from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + + bss, tiers = [1, 16, 64], [1, 3, BLOCK] + captured = {k for _bs, k in _capture_set(bss, tiers)} + + torch.manual_seed(11) + table = SpsCostTable(token_counts=(0, 8, 24, 64), step_time_ms=(2.0, 2.2, 4.0, 9.0)) + for trial in range(40): + p = DSparkVerifyPlanner(cfg=_cfg(min_verify_len=1), cost_table=table, tiers=tiers) + p._host_buffer = torch.rand(8, BLOCK) * 12 - 6 + p._copy_event, p._snapshot_valid = None, True + k = p.decide_draft_len(num_gen_requests=8) + assert k in captured, f"trial {trial}: planner picked K={k}, captured={captured}" + + +def test_planner_fallbacks_also_land_on_a_captured_graph(): + """Every degraded path (no snapshot, flat cost, empty batch) must stay captured.""" + from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + + tiers = [1, 3, BLOCK] + captured = {k for _bs, k in _capture_set([1, 64], tiers)} + cases = [ + ( + "no snapshot", + DSparkVerifyPlanner( + cfg=_cfg(), + cost_table=SpsCostTable(token_counts=(0, 8), step_time_ms=(2.0, 5.0)), + tiers=tiers, + ), + 8, + ), + ( + "flat cost", + DSparkVerifyPlanner(cfg=_cfg(), cost_table=SpsCostTable.flat(), tiers=tiers), + 8, + ), + ( + "empty batch", + DSparkVerifyPlanner( + cfg=_cfg(), + cost_table=SpsCostTable(token_counts=(0, 8), step_time_ms=(2.0, 5.0)), + tiers=tiers, + ), + 0, + ), + ] + for name, p, n in cases: + assert p.decide_draft_len(num_gen_requests=n) in captured, name + + +# -------------------------------------------------------------------------- +# Acceptance criterion (3): attention-DP + TP +# +# draft_len is part of the CUDA-graph key but is NOT covered by the ADP +# consistency allgather, so ranks that choose differently select different +# graphs -- one replays, one falls back to eager -- and their collectives +# diverge. Agreement has to be forced before the forward. +# -------------------------------------------------------------------------- + + +def _simulated_ranks(local_confidences, tiers, table, batch_sizes=None): + """Run the planner on each rank with a real max-reduction across them.""" + from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + + batch_sizes = batch_sizes or [c.shape[0] for c in local_confidences] + planners = [] + for conf in local_confidences: + p = DSparkVerifyPlanner(cfg=_cfg(min_verify_len=1), cost_table=table, tiers=tiers) + p._host_buffer, p._copy_event, p._snapshot_valid = conf, None, True + planners.append(p) + locals_ = [p._decide_local(num_gen_requests=bs) for p, bs in zip(planners, batch_sizes)] + peak = max(locals_) + return [ + p.decide_draft_len(num_gen_requests=bs, all_rank_max=lambda _v, m=peak: m) + for p, bs in zip(planners, batch_sizes) + ], locals_ + + +def test_adp_ranks_agree_despite_different_confidence(): + tiers = [1, 3, BLOCK] + table = SpsCostTable(token_counts=(0, 16, 48, 128), step_time_ms=(2.0, 2.2, 4.0, 9.0)) + confs = [ + torch.full((8, BLOCK), -6.0), + torch.full((8, BLOCK), 6.0), + torch.full((8, BLOCK), 0.0), + torch.full((8, BLOCK), 3.0), + ] + chosen, locals_ = _simulated_ranks(confs, tiers, table) + assert len(set(chosen)) == 1, f"ranks disagreed: {chosen} (locals {locals_})" + assert chosen[0] in tiers + + +def test_adp_ranks_agree_despite_different_batch_sizes(): + """Under ADP each rank has its own batch; agreement must survive that too.""" + tiers = [1, 3, BLOCK] + table = SpsCostTable(token_counts=(0, 16, 48, 128), step_time_ms=(2.0, 2.2, 4.0, 9.0)) + torch.manual_seed(3) + confs = [torch.rand(b, BLOCK) * 8 - 4 for b in (2, 9, 17, 33)] + chosen, locals_ = _simulated_ranks(confs, tiers, table, batch_sizes=[2, 9, 17, 33]) + assert len(set(chosen)) == 1, f"ranks disagreed: {chosen} (locals {locals_})" + + +def test_adp_rank_with_empty_batch_does_not_break_agreement(): + """An idle DP rank still has to enter the same graph as everyone else.""" + tiers = [1, 3, BLOCK] + table = SpsCostTable(token_counts=(0, 16, 48, 128), step_time_ms=(2.0, 2.2, 4.0, 9.0)) + confs = [torch.zeros(0, BLOCK), torch.full((8, BLOCK), 6.0), torch.full((8, BLOCK), -6.0)] + chosen, _ = _simulated_ranks(confs, tiers, table, batch_sizes=[0, 8, 8]) + assert len(set(chosen)) == 1, f"idle rank diverged: {chosen}" + + +def test_cross_rank_reduction_is_max_not_min(): + """Max is the safe direction: a rank that wanted less simply verifies more. + + Taking the min would let one pessimistic rank starve the whole batch, and -- + worse -- could drop below another rank's min_verify_len floor. + """ + tiers = [1, 3, BLOCK] + table = SpsCostTable(token_counts=(0, 16, 48, 128), step_time_ms=(2.0, 2.2, 4.0, 9.0)) + confs = [torch.full((8, BLOCK), -6.0), torch.full((8, BLOCK), 6.0)] + chosen, locals_ = _simulated_ranks(confs, tiers, table) + assert chosen[0] == max(locals_) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py index f66a37834fab..11969a8ca0f9 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py @@ -250,3 +250,105 @@ def run(persist): graph.replay() torch.cuda.synchronize() torch.testing.assert_close(out, eager, rtol=2e-2, atol=2e-2) + + +# -------------------------------------------------------------------------- +# Confidence-scheduled verification: the scoring path must be graph-capturable. +# +# DSpark is a one-engine drafter, so its worker forward runs INSIDE the target's +# CUDA graph. Anything the confidence head adds to that forward has to capture. +# The head itself is a fixed-shape Linear, but the surrounding code is where +# capture breaks: the previous implementation followed it with a `.item()` and a +# `torch.nonzero`, both of which are illegal under capture. +# -------------------------------------------------------------------------- + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +@pytest.mark.parametrize("with_markov", [False, True]) +def test_confidence_head_is_cuda_graph_capturable(with_markov): + from tensorrt_llm._torch.models.dspark.heads import DSparkConfidenceHead + + B, BLK, HID, RANK = 4, 5, 32, 16 + head = ( + DSparkConfidenceHead( + hidden_size=HID, markov_rank=RANK, with_markov=with_markov, block_size=BLK + ) + .cuda() + .eval() + ) + hid = torch.randn(B, BLK, HID, device="cuda") + prev = torch.randn(B, BLK, RANK, device="cuda") if with_markov else None + out = torch.empty(B, BLK, device="cuda", dtype=torch.float32) + + def run(): + return head.apply_sts(head(hid, prev_embeddings=prev)) + + # Warm up on a side stream, as CUDA graph capture requires. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out.copy_(run()) + + eager = run() + g.replay() + torch.cuda.synchronize() + assert torch.allclose(out, eager, atol=1e-5), "replay diverged from eager" + assert torch.all((out >= 0.0) & (out <= 1.0)), "calibrated confidence must be a probability" + + # New inputs must flow through on replay -- a graph that captured constants + # would keep returning the old scores and silently freeze the scheduler. + hid.copy_(torch.randn_like(hid)) + fresh = run() + g.replay() + torch.cuda.synchronize() + assert torch.allclose(out, fresh, atol=1e-5), "replay ignored updated inputs" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +def test_sts_table_update_is_visible_to_a_captured_graph(): + """STS must be updated in place; rebinding the buffer would be invisible. + + ``nn.Module.__setattr__`` on a registered buffer replaces the tensor object, + changing its data_ptr. A graph captured before that keeps reading the OLD + storage, so calibration would silently not apply. + """ + from tensorrt_llm._torch.models.dspark.heads import DSparkConfidenceHead + + B, BLK, HID = 2, 5, 32 + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK).cuda().eval() + hid = torch.randn(B, BLK, HID, device="cuda") + out = torch.empty(B, BLK, device="cuda", dtype=torch.float32) + + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + head.apply_sts(head(hid)) + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out.copy_(head.apply_sts(head(hid))) + + g.replay() + torch.cuda.synchronize() + before = out.clone() + + storage_before = head.sts_temperatures.data_ptr() + head.load_sts_temperatures(torch.full((BLK,), 4.0)) + assert head.sts_temperatures.data_ptr() == storage_before, ( + "load_sts_temperatures must write in place; rebinding hides it from the graph" + ) + + g.replay() + torch.cuda.synchronize() + assert not torch.allclose(out, before), ( + "the captured graph did not see the new STS temperatures" + ) + assert torch.allclose(out, torch.sigmoid(head(hid) / 4.0), atol=1e-5) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py index fed0da3a2803..80c8c3bb7884 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py @@ -40,7 +40,7 @@ def test_dspark_propose_full_block_no_confidence(): bonus = torch.randint(0, VOCAB, (B,)) hid = torch.randn(B, BLK, HID) with torch.no_grad(): - tokens, num = dspark_propose( + tokens, confidence = dspark_propose( base, bonus_token_ids=bonus, block_hidden=hid, @@ -49,8 +49,7 @@ def test_dspark_propose_full_block_no_confidence(): block_size=BLK, ) assert tokens.shape == (B, BLK) - # No confidence head -> propose the full block. - assert torch.all(num == BLK) + assert confidence is None # Tokens match the markov head's own greedy block sampling. ref_tokens, _ = markov.sample_block_tokens( base, first_prev_token_ids=bonus, hidden_states=hid, temperature=0.0 @@ -58,42 +57,67 @@ def test_dspark_propose_full_block_no_confidence(): assert torch.equal(tokens, ref_tokens) -def test_dspark_propose_confidence_truncates(): - torch.manual_seed(1) - markov = build_markov_head( - markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID - ).eval() - conf = DSparkConfidenceHead(hidden_size=HID).eval() - # The confidence proj is bias-free, so drive the logit via a constant weight - # against a constant hidden: logit = weight_val * HID per position. +def _propose_with_confidence(conf, markov, *, return_confidence): base = torch.randn(1, BLK, VOCAB) bonus = torch.randint(0, VOCAB, (1,)) hid = torch.ones(1, BLK, HID) with torch.no_grad(): - conf.proj.weight.fill_(5.0 / HID) # logit ~ 5 -> sigmoid ~ 0.993, all confident - with torch.no_grad(): - _, num = dspark_propose( + return dspark_propose( base, bonus_token_ids=bonus, block_hidden=hid, markov_head=markov, confidence_head=conf, block_size=BLK, - confidence_threshold=0.5, + return_confidence=return_confidence, ) - # All-confident -> full block proposed. - assert int(num[0]) == BLK - # Now make the head output low confidence everywhere -> truncate to 0... but - # confident_prefix_length returns first sub-threshold index (0 here). + + +def test_dspark_propose_scores_without_shortening_the_block(): + """The block is always proposed in full; confidence only scores it.""" + torch.manual_seed(1) + markov = build_markov_head( + markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID + ).eval() + conf = DSparkConfidenceHead(hidden_size=HID, block_size=BLK).eval() + # The confidence proj is bias-free, so drive the logit via a constant weight + # against a constant hidden: logit = weight_val * HID per position. with torch.no_grad(): - conf.proj.weight.fill_(-5.0 / HID) # logit ~ -5 -> sigmoid ~ 0.0067 < 0.5 - _, num2 = dspark_propose( - base, - bonus_token_ids=bonus, - block_hidden=hid, - markov_head=markov, - confidence_head=conf, - block_size=BLK, - confidence_threshold=0.5, - ) - assert int(num2[0]) == 0 + conf.proj.weight.fill_(-5.0 / HID) # sigmoid ~ 0.0067, i.e. hopeless + + tokens, confidence = _propose_with_confidence(conf, markov, return_confidence=True) + assert tokens.shape == (1, BLK) + assert confidence.shape == (1, BLK) + # Low confidence must NOT shorten the proposal -- that decision belongs to + # the verification scheduler, not the drafter. + assert torch.all(confidence < 0.0) + assert torch.all(conf.apply_sts(confidence) < 0.5) + + +def test_dspark_propose_confidence_is_opt_in(): + torch.manual_seed(2) + markov = build_markov_head( + markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID + ).eval() + conf = DSparkConfidenceHead(hidden_size=HID, block_size=BLK).eval() + _, confidence = _propose_with_confidence(conf, markov, return_confidence=False) + assert confidence is None + + +def test_dspark_propose_is_free_of_host_syncs(): + """No ``.item()``/``nonzero`` on this path: it runs inside the target's graph.""" + import inspect + import io + import tokenize + + from tensorrt_llm._torch.models.dspark import draft as draft_mod + + src = inspect.getsource(draft_mod.dspark_propose) + # Strip comments and string literals so the check reads the code, not the + # prose describing it (the implementation comments name these very calls). + code = "".join( + tok.string if tok.type not in (tokenize.COMMENT, tokenize.STRING) else " " + for tok in tokenize.generate_tokens(io.StringIO(src).readline) + ) + for banned in (".item(", "nonzero", "range("): + assert banned not in code, f"dspark_propose must stay capture-safe: found {banned!r}" diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py index 16b6a01acd17..e6c121d0f857 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py @@ -22,7 +22,6 @@ RNNHead, VanillaMarkov, build_markov_head, - confident_prefix_length, ) VOCAB, RANK, HID, B, BLK = 257, 16, 32, 3, 5 @@ -93,18 +92,70 @@ def test_build_markov_head_rank_zero_returns_none(): ) -def test_confidence_head_and_prefix_truncation(): - head = DSparkConfidenceHead(hidden_size=HID) +def test_confidence_head_emits_raw_logits(): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) conf = head(torch.randn(B, BLK, HID)) assert conf.shape == (B, BLK) - # threshold 0 disables truncation. - assert confident_prefix_length(conf, block_size=BLK, threshold=0.0) == BLK - # First sub-threshold position truncates the prefix. - logits = torch.tensor([[10.0, 10.0, -10.0, 10.0, 10.0]]) - assert confident_prefix_length(logits, block_size=BLK, threshold=0.5) == 2 - # All-confident -> full block. - logits_hi = torch.full((1, BLK), 10.0) - assert confident_prefix_length(logits_hi, block_size=BLK, threshold=0.5) == BLK + assert conf.dtype == torch.float32 + + +def test_apply_sts_is_identity_by_default(): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + raw = torch.tensor([[2.0, 0.0, -2.0, 1.0, -1.0]]) + assert torch.allclose(head.apply_sts(raw), torch.sigmoid(raw), atol=1e-6) + assert torch.all((head.apply_sts(raw) >= 0.0) & (head.apply_sts(raw) <= 1.0)) + + +def test_apply_sts_uses_per_position_temperatures(): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + temps = torch.tensor([1.0, 2.0, 4.0, 1.0, 1.0]) + head.load_sts_temperatures(temps) + raw = torch.full((1, BLK), 4.0) + got = head.apply_sts(raw) + # A larger temperature pulls the probability toward 0.5. + assert got[0, 0] > got[0, 1] > got[0, 2] + assert torch.allclose(got, torch.sigmoid(raw / temps), atol=1e-6) + + +def test_load_sts_temperatures_updates_in_place(): + """Rebinding the buffer would be invisible to an already-captured graph.""" + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + before = head.sts_temperatures + head.load_sts_temperatures(torch.full((BLK,), 2.0)) + assert head.sts_temperatures is before + assert torch.allclose(before, torch.full((BLK,), 2.0)) + + +def test_load_sts_temperatures_validates(): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + with pytest.raises(ValueError, match="one per block position"): + head.load_sts_temperatures(torch.ones(BLK + 1)) + with pytest.raises(ValueError, match="strictly positive"): + head.load_sts_temperatures(torch.zeros(BLK)) + + +def test_sts_buffer_is_not_persistent(): + """It is calibration, not a checkpoint weight -- it must not affect loading.""" + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + assert "sts_temperatures" not in head.state_dict() + + +def test_confidence_head_load_weights_rejects_a_dropped_bias(): + """A checkpoint bias silently ignored would shift every confidence score.""" + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + good = {"proj.weight": torch.randn(1, HID)} + head.load_weights([good]) + assert torch.allclose(head.proj.weight, good["proj.weight"]) + + with pytest.raises(ValueError, match="bias=True"): + head.load_weights([{**good, "proj.bias": torch.zeros(1)}]) + + +def test_confidence_head_with_bias_accepts_bias_weights(): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK, bias=True) + w = {"proj.weight": torch.randn(1, HID), "proj.bias": torch.randn(1)} + head.load_weights([w]) + assert torch.allclose(head.proj.bias, w["proj.bias"]) def test_confidence_head_with_markov_concat_dim(): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py new file mode 100644 index 000000000000..a586c44a8c9b --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py @@ -0,0 +1,291 @@ +# 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. +"""Overlap-scheduler input layout under ragged (per-request) verification. + +The overlap path builds this iteration's inputs from the *previous* iteration's +device tensors, and does it with a fixed stride: every request is assumed to +contribute ``runtime_tokens_per_gen_step`` tokens. Two places break silently +when that assumption is dropped: + + * the flat gathers out of ``new_tokens`` / ``next_draft_tokens``, and the + ``position_ids`` / ``gather_ids`` / ``seq_lens`` the host appends -- a wrong + count shifts every following request, so RoPE phases and attention windows + go wrong without any shape error; and + * ``kv_len_offsets``, which is subtracted from a host-side estimate that was + built with the *same* count. If only one of the two is made per-request the + KV length is off by exactly their difference. + +These tests pin the arithmetic of both, and -- most importantly -- that a +uniform batch reproduces the strided layout exactly, since that is the path +every non-DSpark model keeps taking. +""" + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.llm_request import get_request_tokens_per_gen_step +from tensorrt_llm._torch.speculative.dspark_ragged import ragged_gather_index_lists + + +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 + + +# -------------------------------------------------------------------------- +# get_request_tokens_per_gen_step +# -------------------------------------------------------------------------- + + +def test_missing_attribute_yields_the_batch_wide_count(): + # Every non-DSpark request lands here: the attribute does not exist at all. + assert get_request_tokens_per_gen_step(_Req(), 4) == 4 + + +def test_none_verify_len_yields_the_batch_wide_count(): + req = _Req() + req.py_verify_len = None + assert get_request_tokens_per_gen_step(req, 4) == 4 + + +def test_verify_len_is_a_draft_length_not_a_token_count(): + # py_verify_len counts draft positions, matching runtime_draft_len; the + # accepted token makes it one more. + assert get_request_tokens_per_gen_step(_Req(2), 6) == 3 + + +def test_zero_verify_len_still_carries_the_accepted_token(): + assert get_request_tokens_per_gen_step(_Req(0), 6) == 1 + + +# -------------------------------------------------------------------------- +# ragged_gather_index_lists +# -------------------------------------------------------------------------- + + +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] + + +def test_zero_count_contributes_nothing(): + rows, cols = ragged_gather_index_lists([1, 2, 3], [2, 0, 1]) + assert rows == [1, 1, 3] + assert cols == [0, 1, 0] + + +def test_empty_batch(): + assert ragged_gather_index_lists([], []) == ([], []) + + +def test_mismatched_lengths_raise(): + with pytest.raises(ValueError, match="same length"): + ragged_gather_index_lists([0, 1], [3]) + + +def test_negative_count_raises(): + with pytest.raises(ValueError, match="negative gather count"): + ragged_gather_index_lists([0], [-1]) + + +def test_total_gathered_equals_the_sum_of_counts(): + counts = [4, 1, 3, 2] + rows, cols = ragged_gather_index_lists([0, 1, 2, 3], counts) + assert len(rows) == len(cols) == sum(counts) + + +# -------------------------------------------------------------------------- +# Host-side layout arithmetic (mirrors the extend-request loop) +# -------------------------------------------------------------------------- + + +def _build_overlap_layout(requests, past_seens, runtime_tokens_per_gen_step): + """Reproduce the per-request bookkeeping the overlap branch emits. + + Mirrors _prepare_tp_inputs' previous-batch branch: one entry per request in + seq_lens/draft_lens/cached, and ``tokens`` flat entries per request in + position_ids/gather_ids/previous_pos_indices. + """ + seq_lens, draft_lens, cached, tokens_each = [], [], [], [] + position_ids, gather_ids, pos_indices = [], [], [] + for slot, (req, past_seen) in enumerate(zip(requests, past_seens)): + tokens = get_request_tokens_per_gen_step(req, runtime_tokens_per_gen_step) + seq_lens.append(tokens) + draft_lens.append(tokens - 1) + tokens_each.append(tokens) + gather_ids.extend(range(len(position_ids), len(position_ids) + tokens)) + position_ids.extend(range(past_seen, past_seen + tokens)) + pos_indices.extend([slot] * tokens) + cached.append(past_seen + tokens) + return { + "seq_lens": seq_lens, + "draft_lens": draft_lens, + "cached": cached, + "tokens_each": tokens_each, + "position_ids": position_ids, + "gather_ids": gather_ids, + "pos_indices": pos_indices, + } + + +def test_uniform_batch_reproduces_the_fixed_stride_layout(): + n = 4 + reqs = [_Req(), _Req(), _Req()] + past = [100, 250, 7] + got = _build_overlap_layout(reqs, past, n) + + assert got["seq_lens"] == [n] * 3 + assert got["draft_lens"] == [n - 1] * 3 + assert len(got["position_ids"]) == 3 * n + assert got["pos_indices"] == [0] * n + [1] * n + [2] * n + assert got["gather_ids"] == list(range(3 * n)) + for slot, base in enumerate(past): + window = got["position_ids"][slot * n : (slot + 1) * n] + assert window == list(range(base, base + n)) + + +def test_ragged_batch_keeps_each_request_positions_contiguous(): + # The bug this pins: emitting a fixed count per request shifts every later + # request's position_ids, which is a silent RoPE-phase error. + reqs = [_Req(3), _Req(0), _Req(1)] + past = [100, 250, 7] + got = _build_overlap_layout(reqs, past, 6) + + assert got["tokens_each"] == [4, 1, 2] + assert got["seq_lens"] == [4, 1, 2] + assert got["draft_lens"] == [3, 0, 1] + assert got["position_ids"] == [100, 101, 102, 103, 250, 7, 8] + assert got["gather_ids"] == list(range(7)) + assert got["pos_indices"] == [0, 0, 0, 0, 1, 2, 2] + + +def test_flat_token_total_matches_the_sum_of_windows(): + reqs = [_Req(3), _Req(0), _Req(1)] + got = _build_overlap_layout(reqs, [0, 0, 0], 6) + assert len(got["position_ids"]) == sum(got["tokens_each"]) + assert len(got["pos_indices"]) == sum(got["tokens_each"]) + + +# -------------------------------------------------------------------------- +# kv_len_offsets: the host estimate and the correction must use one count +# -------------------------------------------------------------------------- + + +def _corrected_kv_lens( + requests, past_seens, accepted, runtime_tokens_per_gen_step, correction_tokens +): + """Host estimate (past_seen + tokens) plus the device correction. + + ``correction_tokens`` is what kv_len_offsets subtracts; passing the + batch-wide value while the estimate used a per-request one is exactly the + A2 bug. + """ + out = [] + for req, past_seen, acc, corr in zip(requests, past_seens, accepted, correction_tokens): + tokens = get_request_tokens_per_gen_step(req, runtime_tokens_per_gen_step) + estimate = past_seen + tokens + out.append(estimate + (acc - corr)) + return out + + +def test_paired_counts_cancel_to_past_seen_plus_accepted(): + reqs = [_Req(3), _Req(0), _Req(1)] + past = [100, 250, 7] + accepted = [2, 1, 2] + tokens_each = [get_request_tokens_per_gen_step(r, 6) for r in reqs] + + got = _corrected_kv_lens(reqs, past, accepted, 6, tokens_each) + + # The token count drops out entirely -- that is what makes the correction + # correct regardless of which window the scheduler picked. + assert got == [p + a for p, a in zip(past, accepted)] + + +def test_batch_wide_correction_under_ragged_is_off_by_the_window_difference(): + reqs = [_Req(3), _Req(0), _Req(1)] + past = [100, 250, 7] + accepted = [2, 1, 2] + n = 6 + + wrong = _corrected_kv_lens(reqs, past, accepted, n, [n] * 3) + right = [p + a for p, a in zip(past, accepted)] + + # Off by (per-request tokens - batch-wide tokens) for every request. + deltas = [w - r for w, r in zip(wrong, right)] + assert deltas == [4 - n, 1 - n, 2 - n] + assert all(d != 0 for d in deltas) + + +def test_uniform_batch_is_unaffected_by_the_pairing_rule(): + reqs = [_Req(), _Req(), _Req()] + past = [100, 250, 7] + accepted = [2, 1, 2] + n = 4 + + per_request = _corrected_kv_lens(reqs, past, accepted, n, [n] * 3) + assert per_request == [p + a for p, a in zip(past, accepted)] + + +# -------------------------------------------------------------------------- +# Derived totals the callers compute from the per-request counts +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "verify_lens", + [ + [None, None, None], + [3, 3, 3], + [3, 0, 1], + [0, 0, 0], + ], +) +def test_draft_token_total_is_tokens_minus_one_per_request(verify_lens): + # previous_batch_draft_tokens = total_num_tokens - num_requests, which has + # to agree with summing (tokens - 1) request by request. + reqs = [_Req(v) for v in verify_lens] + tokens = [get_request_tokens_per_gen_step(r, 4) for r in reqs] + total_num_tokens = sum(tokens) + assert total_num_tokens - len(reqs) == sum(t - 1 for t in tokens) + + +def test_uniform_totals_match_the_multiplicative_formulas(): + # The formulas the uniform path used before this change. + n, num_requests = 4, 5 + reqs = [_Req() for _ in range(num_requests)] + tokens = [get_request_tokens_per_gen_step(r, n) for r in reqs] + assert sum(tokens) == num_requests * n + assert sum(tokens) - num_requests == num_requests * (n - 1) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py new file mode 100644 index 000000000000..132d01620c10 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py @@ -0,0 +1,676 @@ +# 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. +"""Unit tests for the ragged (per-request) verification layout.""" + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_ragged import ( + RaggedVerifyLayout, + build_qo_indptr, + choose_ragged_capture_shape, + count_accepted_ragged, + exceeds_captured_buckets, + fill_padded_rows_onehot, + round_up_to_bucket, + row_ids_from_lens, + scatter_ragged_to_padded, +) +from tensorrt_llm._torch.speculative.dspark_schedule import ( + DSparkScheduleConfig, + compute_survival, + schedule_verify_lens_topk, +) + +BLOCK = 7 + + +def _layout(lens, bucket=None, buckets=None): + t = torch.tensor(lens, dtype=torch.int32) + return RaggedVerifyLayout.from_verify_lens( + t, graph_num_tokens=bucket, buckets=buckets, total_verify_tokens=sum(lens) + ) + + +# -------------------------------------------------------------------------- +# packing +# -------------------------------------------------------------------------- + + +def test_qo_indptr_is_the_exclusive_prefix_sum(): + lens = torch.tensor([3, 1, 4, 2], dtype=torch.int32) + assert build_qo_indptr(lens).tolist() == [0, 3, 4, 8, 10] + + +def test_qo_indptr_slices_recover_each_request(): + """Request r must own exactly [qo_indptr[r], qo_indptr[r+1]).""" + lens = [3, 1, 4, 2] + lay = _layout(lens, bucket=10) + flat = torch.arange(sum(lens)) + for r, n in enumerate(lens): + lo, hi = int(lay.qo_indptr[r]), int(lay.qo_indptr[r + 1]) + assert hi - lo == n + assert flat[lo:hi].numel() == n + assert int(lay.qo_indptr[-1]) == sum(lens) + + +def test_extend_start_loc_matches_indptr_head(): + lay = _layout([2, 5, 1], bucket=8) + assert torch.equal(lay.extend_start_loc, lay.qo_indptr[:-1]) + + +def test_empty_batch_indptr(): + assert build_qo_indptr(torch.zeros(0, dtype=torch.int32)).tolist() == [0] + + +def test_build_qo_indptr_rejects_2d(): + with pytest.raises(ValueError, match="1-D"): + build_qo_indptr(torch.ones(2, 3, dtype=torch.int32)) + + +# -------------------------------------------------------------------------- +# bucket selection +# -------------------------------------------------------------------------- + + +def test_round_up_picks_the_smallest_fitting_bucket(): + b = [8, 16, 32, 64] + assert round_up_to_bucket(1, b) == 8 + assert round_up_to_bucket(8, b) == 8 + assert round_up_to_bucket(9, b) == 16 + assert round_up_to_bucket(64, b) == 64 + + +def test_round_up_raises_past_the_largest_bucket(): + """Clamping would silently drop the step out of graph replay into eager.""" + with pytest.raises(ValueError, match="exceeds the largest captured bucket"): + round_up_to_bucket(65, [8, 16, 32, 64]) + assert exceeds_captured_buckets(65, [8, 16, 32, 64]) + assert not exceeds_captured_buckets(64, [8, 16, 32, 64]) + + +def test_layout_derives_the_bucket_from_the_host_total(): + lay = RaggedVerifyLayout.from_verify_lens( + torch.tensor([3, 1, 4], dtype=torch.int32), buckets=[8, 16, 32], total_verify_tokens=8 + ) + assert lay.graph_num_tokens == 8 + + +def test_layout_refuses_to_sync_for_the_total(): + with pytest.raises(ValueError, match="would sync"): + RaggedVerifyLayout.from_verify_lens(torch.tensor([3, 1], dtype=torch.int32)) + + +# -------------------------------------------------------------------------- +# padding reclamation +# -------------------------------------------------------------------------- + + +def test_fill_bucket_converts_padding_into_real_verification(): + """The step already pays for the bucket -- spare slots must not be wasted.""" + lay = _layout([2, 1, 1], bucket=12) # 4 used, 8 spare + filled = lay.fill_bucket(max_verify_len=BLOCK) + assert filled.total_verify_tokens == 12 + assert int(filled.verify_lens.sum()) == 12 + 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_is_a_noop_when_exact(): + lay = _layout([3, 2, 3], bucket=8) + assert lay.fill_bucket(max_verify_len=BLOCK) is not None + assert int(lay.fill_bucket(max_verify_len=BLOCK).verify_lens.sum()) == 8 + + +def test_fill_bucket_adds_pad_rows_to_reach_the_captured_batch_size(): + """Row count must hit padded_bs; the leftovers land on the 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 are saturated before any slack is wasted on pad rows. + lens = filled.verify_lens.tolist() + assert lens[0] >= 3 and lens[1] >= 2 + assert min(lens[2:]) >= 1, "a pad row with 0 tokens breaks qo_indptr slicing" + + +def test_fill_bucket_prefers_real_rows_over_pad_rows(): + lay = _layout([1, 1], bucket=10) + filled = lay.fill_bucket(max_verify_len=BLOCK, padded_bs=4) + lens = filled.verify_lens.tolist() + assert sum(lens) == 10 + # 8 spare beyond the baseline (1+1 real + 1+1 pad): real rows take theirs + # first, so both real rows outrank both pad rows. + assert min(lens[:2]) >= max(lens[2:]) + + +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) + + +def test_fill_bucket_raises_when_the_bucket_is_too_small(): + lay = _layout([5, 5], bucket=8) # 10 real tokens already > 8 + with pytest.raises(ValueError, match="too small"): + lay.fill_bucket(max_verify_len=BLOCK) + + +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() + + +# -------------------------------------------------------------------------- +# end-to-end: confidence -> per-request lengths -> packed layout +# -------------------------------------------------------------------------- + + +def test_scheduler_output_packs_into_a_valid_layout(): + """The whole point: heterogeneous confidence must yield heterogeneous lengths.""" + conf = torch.tensor( + [ + [0.98, 0.97, 0.96, 0.95, 0.94, 0.93, 0.92], # deep, worth verifying + [0.55, 0.40, 0.30, 0.25, 0.20, 0.15, 0.10], # shallow + [0.90, 0.85, 0.70, 0.55, 0.45, 0.35, 0.25], + [0.60, 0.45, 0.35, 0.28, 0.22, 0.18, 0.12], + ] + ) + cfg = DSparkScheduleConfig(block_size=BLOCK, min_verify_len=1) + lens = schedule_verify_lens_topk(survival=compute_survival(conf), budget=8, cfg=cfg) + + assert lens[0] > lens[1], "the confident request should verify deeper" + total = int(lens.sum()) + lay = RaggedVerifyLayout.from_verify_lens( + lens, buckets=[8, 16, 32, 64], total_verify_tokens=total + ) + lay.validate(max_verify_len=BLOCK) + assert lay.graph_num_tokens >= total + + +def test_ragged_beats_uniform_at_equal_budget(): + """Pins the reason ragged exists: same tokens verified, more expected yield.""" + conf = torch.tensor( + [ + [0.98, 0.97, 0.96, 0.95, 0.94, 0.93, 0.92], + [0.55, 0.40, 0.30, 0.25, 0.20, 0.15, 0.10], + [0.90, 0.85, 0.70, 0.55, 0.45, 0.35, 0.25], + [0.60, 0.45, 0.35, 0.28, 0.22, 0.18, 0.12], + ] + ) + surv = compute_survival(conf) + cfg = DSparkScheduleConfig(block_size=BLOCK, min_verify_len=1) + budget, bs = 8, conf.shape[0] + + ragged_lens = schedule_verify_lens_topk(survival=surv, budget=budget, cfg=cfg) + tau_ragged = sum(float(surv[r, : int(ragged_lens[r])].sum()) for r in range(bs)) + + # Best uniform length that fits the same total token spend. + budget_tokens = bs * cfg.min_verify_len + budget + tau_uniform = max( + float(surv[:, :L].sum()) for L in range(1, BLOCK + 1) if bs * L <= budget_tokens + ) + + assert int(ragged_lens.sum()) <= budget_tokens + assert tau_ragged > tau_uniform, ( + f"ragged {tau_ragged:.2f} should beat uniform {tau_uniform:.2f} at the same token spend" + ) + + +# -------------------------------------------------------------------------- +# 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).tolist() == [0, 0, 0, 1, 2, 2] + + +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]] + + +def test_scatter_round_trips_through_qo_indptr(): + torch.manual_seed(4) + lens_list = [4, 1, 7, 2, 5] + lens = torch.tensor(lens_list, dtype=torch.int32) + lay = _layout(lens_list, bucket=32) + flat = torch.randint(0, 1000, (sum(lens_list),)) + out = scatter_ragged_to_padded(flat, verify_lens=lens, qo_indptr=lay.qo_indptr, max_len=BLOCK) + for r, n in enumerate(lens_list): + lo = int(lay.qo_indptr[r]) + assert out[r, :n].tolist() == flat[lo : lo + n].tolist() + + +def test_accept_count_stops_at_the_first_mismatch(): + draft = torch.tensor([[1, 2, 3, 4]]) + target = torch.tensor([[1, 2, 9, 4]]) + lens = torch.tensor([4]) + assert count_accepted_ragged( + draft_tokens=draft, target_tokens=target, verify_lens=lens + ).tolist() == [2] + + +def test_accept_count_respects_each_request_window(): + """The core ragged invariant: a request only gets credit inside its window.""" + draft = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]]) + target = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]]) + # Row 0 verified all 4; row 1 verified only 2 even though everything matches. + lens = torch.tensor([4, 2]) + got = count_accepted_ragged(draft_tokens=draft, target_tokens=target, verify_lens=lens) + assert got.tolist() == [4, 2] + + +def test_accept_count_ignores_stale_padding_beyond_the_window(): + """Stale slots must not be able to credit an acceptance the target never made. + + Positions past verify_len were never sent to the target. If they happen to + hold matching values from a previous iteration and are not masked, the + cumulative product runs straight through them and reports more accepted + tokens than were actually verified -- wrong output, no error. + """ + draft = torch.tensor([[1, 2, 3, 4]]) + target = torch.tensor([[1, 2, 3, 4]]) # everything "matches", incl. stale tail + lens = torch.tensor([2]) # but only 2 were really verified + got = count_accepted_ragged(draft_tokens=draft, target_tokens=target, verify_lens=lens) + assert got.tolist() == [2], "unmasked padding leaked into the accept count" + + +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]), + ) + + +# --------------------------------------------------------------------------- +# choose_ragged_capture_shape +# --------------------------------------------------------------------------- + +BS_BUCKETS = [1, 2, 4, 8, 16] +TOK_BUCKETS = [8, 16, 32, 64, 128] + + +def _shape(num_real, total, peers=None): + return choose_ragged_capture_shape( + num_real_requests=num_real, + total_verify_tokens=total, + bs_buckets=BS_BUCKETS, + token_buckets=TOK_BUCKETS, + peer_stats=peers, + ) + + +def test_shape_rounds_both_axes_up_to_captured_values(): + got = _shape(num_real=5, total=20) + assert got.padded_bs == 8 # 5 -> 8 + # slack = 20 - 5 = 15; needed = 8 + 15 = 23 -> 32 + assert got.bucket == 32 + + +def test_bucket_covers_the_padded_rows_floor(): + # Every row verifies at least its bonus position, so the bucket can never + # be smaller than the row count. + got = _shape(num_real=3, total=3) + assert got.padded_bs == 4 + assert got.bucket >= got.padded_bs + + +def test_exact_captured_values_are_not_rounded_further(): + # 8 requests, slack 8 -> needed 16, both already captured. + got = _shape(num_real=8, total=16) + assert got.padded_bs == 8 + assert got.bucket == 16 + + +def test_every_rank_computes_the_same_shape(): + # The ADP invariant: ranks with different batches must agree. Each rank + # passes the same peer list and must get an identical answer. + peers = [(3, 9), (6, 30), (5, 11)] + shapes = {_shape(num_real=real, total=total, peers=peers) for real, total in peers} + assert len(shapes) == 1 + + +def test_shape_is_driven_by_the_widest_rank_not_the_local_one(): + # This rank is small (3 requests, 9 tokens) but must still size for the + # widest peer: rows from max real (6 -> 8), slack from max (30 - 6 = 24). + peers = [(3, 9), (6, 30), (5, 11)] + got = _shape(num_real=3, total=9, peers=peers) + assert got.padded_bs == 8 + assert got.bucket == 32 # 8 + 24 = 32, exactly a captured bucket + + +def test_widest_axes_can_come_from_different_ranks(): + # rank A has the most rows, rank B the most slack. + peers = [(7, 8), (2, 20)] + got = _shape(num_real=7, total=8, peers=peers) + assert got.padded_bs == 8 # from rank A + # slack: A = 1, B = 18 -> needed = 8 + 18 = 26 -> 32 + assert got.bucket == 32 + + +def test_shape_raises_when_the_batch_exceeds_captured_rows(): + with pytest.raises(ValueError, match="exceed the largest captured batch"): + _shape(num_real=17, total=17) + + +def test_shape_raises_when_no_token_bucket_fits(): + with pytest.raises(ValueError, match="exceeds the largest captured bucket"): + _shape(num_real=16, total=1000) + + +def test_shape_rejects_fewer_tokens_than_requests(): + with pytest.raises(ValueError, match="cannot be below"): + _shape(num_real=5, total=4) + + +def test_shape_requires_bs_buckets(): + with pytest.raises(ValueError, match="requires bs_buckets"): + choose_ragged_capture_shape( + num_real_requests=1, total_verify_tokens=1, bs_buckets=[], token_buckets=TOK_BUCKETS + ) + + +def test_uniform_batch_shape_matches_the_one_dimensional_rule(): + # With a uniform K the token total is bs * (K + 1), so the joint rule must + # not pick anything larger than the uniform path would have needed. + k_plus_1 = 4 + for num_real in (1, 3, 8): + got = _shape(num_real=num_real, total=num_real * k_plus_1) + assert got.padded_bs >= num_real + assert got.bucket >= num_real * k_plus_1 + + +# --------------------------------------------------------------------------- +# fill_padded_rows_onehot +# --------------------------------------------------------------------------- + + +def test_padding_rows_become_a_valid_distribution(): + probs = torch.zeros(2, 3, 5) + 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])) + # Padding positions sum to 1 instead of 0. + 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)) + + +def test_full_length_rows_are_left_alone(): + probs = torch.full((2, 3, 4), 0.25) + before = probs.clone() + fill_padded_rows_onehot(probs, verify_lens=torch.tensor([3, 3])) + assert torch.equal(probs, 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_hits_the_bucket_exactly(): + # 3 real requests wanting 2+3+1=6 tokens, padded to 4 rows, bucket 12. + got = _fit([2, 3, 1], bucket=12, padded_bs=4, max_verify_len=6) + assert int(got.verify_lens.sum()) == 12 + assert got.bs == 4 + + +def test_fit_spends_slack_on_real_requests_before_pad_rows(): + got = _fit([2, 2], bucket=8, padded_bs=4, max_verify_len=6) + lens = got.verify_lens.tolist() + # Two real rows grew; the two pad rows stay at the 1-token minimum. + assert sum(lens) == 8 + assert lens[2] == 1 and lens[3] == 1 + assert lens[0] + lens[1] == 6 + + +def test_fit_pad_rows_are_never_empty(): + # An empty range breaks qo_indptr's per-row slicing. Baseline here is + # 2*4 real + 6 pad rows = 14, so the bucket has to be at least that. + got = _fit([4, 4], bucket=16, padded_bs=8, max_verify_len=4) + assert all(v >= 1 for v in got.verify_lens.tolist()) + + +def test_fit_rejects_a_bucket_that_cannot_hold_the_pad_rows(): + # The same batch against a bucket that leaves no room for the pad rows: + # 8 real tokens + 6 pad rows needs 14, not 10. Raising here is the point -- + # a short batch would desync attention from the MoE. + with pytest.raises(ValueError, match="too small"): + _fit([4, 4], bucket=10, padded_bs=8, max_verify_len=4) + + +def test_fit_marks_which_rows_are_real(): + got = _fit([2, 3], bucket=10, padded_bs=4, max_verify_len=6) + assert got.num_real_requests == 2 + assert got.num_pad_requests == 2 + + +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 + + +def test_capture_buckets_match_the_uniform_token_totals(): + # The ragged ladder is one bucket per tier at each batch size, which is + # exactly what the uniform ladder would have produced -- so enabling ragged + # does not grow the captured graph count. + tiers = [1, 3, 5] + for bs in (1, 4, 16): + buckets = [bs * (t + 1) for t in tiers] + assert buckets == sorted(buckets) + assert len(buckets) == len(tiers) + # Every bucket is reachable: bs rows at >= 1 token, <= max tier + 1. + for b in buckets: + assert bs <= b <= bs * (tiers[-1] + 1) + + +def test_warmup_split_reaches_any_captured_bucket(): + # Capture has to hit the bucket exactly or the graph is keyed on a total no + # runtime batch produces, and every ragged step silently runs eager. + tiers = [1, 3, 5] + for bs in (1, 2, 8): + for t in tiers: + bucket = bs * (t + 1) + got = _fit([1] * bs, bucket=bucket, padded_bs=bs, max_verify_len=1 + tiers[-1]) + assert int(got.verify_lens.sum()) == bucket + + +def _pad_len(*, bucket, n_real, n_pad, max_verify_len, floor=None): + """Mirror of the pad-row window ModelEngine.fit_ragged_verify_lens derives. + + Pad rows take the one-token minimum so real requests keep the slack, and + only grow when the real rows physically cannot absorb the bucket. Returns + None when no integer window leaves the real rows a reachable target -- the + scheduler then declines and the step stays uniform. + """ + real_capacity = n_real * max_verify_len + floor = n_real if floor is None else floor + if n_pad == 0: + pad_len = 1 + else: + lo = max(1, -(-(bucket - real_capacity) // n_pad)) # ceil division + hi = min(max_verify_len, (bucket - floor) // n_pad) + if lo > hi: + return None + pad_len = lo + real_target = bucket - n_pad * pad_len + if real_target < floor or real_target > real_capacity: + return None + return pad_len + + +def test_pad_rows_take_the_minimum_when_real_rows_can_absorb_the_bucket(): + # 3 real rows capped at 6 hold up to 18; bucket 24 minus 5 pad rows leaves + # 19 -- one too many, so the pad rows have to grow. Contrast with below. + padded_bs, bucket, max_verify_len = 8, 18, 6 + real = [2, 3, 2] + n_pad = padded_bs - len(real) + pad_len = _pad_len(bucket=bucket, n_real=len(real), n_pad=n_pad, max_verify_len=max_verify_len) + assert pad_len == 1 + + real_target = bucket - n_pad * pad_len + got = _fit(real, bucket=real_target, padded_bs=len(real), max_verify_len=max_verify_len) + assert int(got.verify_lens.sum()) + n_pad * pad_len == bucket + + +def test_pad_rows_grow_when_the_real_rows_cannot_absorb_the_bucket(): + # This is the case that used to have no solution: 3 real rows cap at 18, + # but bucket 24 with 5 one-token pad rows would demand 19 from them. + padded_bs, bucket, max_verify_len = 8, 24, 6 + real = [2, 3, 2] + n_pad = padded_bs - len(real) + pad_len = _pad_len(bucket=bucket, n_real=len(real), n_pad=n_pad, max_verify_len=max_verify_len) + assert pad_len == 2, "pad rows must absorb what the real rows cannot" + + real_target = bucket - n_pad * pad_len + 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 * pad_len == bucket + + +def test_pad_length_never_yields_an_infeasible_real_target(): + """Exhaustive sweep of every shape the scheduler can actually produce. + + A wrong pad window is the worst failure mode here: the batch's token total + silently misses every captured bucket and the step drops out of graph + replay into eager. So sweep the real bucket grid (one bucket per tier at + each captured batch size), pick the bucket the way the scheduler does, and + require that the derived split either reconstructs the bucket exactly or + declines outright. + """ + 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 + pad_len = _pad_len( + bucket=bucket, + n_real=n_real, + n_pad=n_pad, + max_verify_len=max_verify_len, + floor=floor, + ) + if pad_len is None: + declined += 1 + continue + real_target = bucket - n_pad * pad_len + 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={pad_len} -> {real_target}" + ) + assert real_target + n_pad * 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" diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 4625d1b0115a..1ecf85b8f332 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -93,17 +93,26 @@ def test_metadata_prepare_batch_indices(): assert meta.batch_indices_cuda[:3].tolist() == [0, 1, 2] -def _make_worker(): +def _make_worker(enable_confidence_scheduling=False, verify_len_tiers=None): cfg = types.SimpleNamespace( max_draft_len=5, spec_dec_mode=SpeculativeDecodingMode.DSPARK, - confidence_threshold=0.5, + enable_confidence_scheduling=enable_confidence_scheduling, + verify_len_tiers=verify_len_tiers or [1, 3, 5], + confidence_sps_table_path=None, + confidence_sts_path=None, ) from tensorrt_llm.mapping import Mapping return DSparkWorker(cfg, Mapping()) +def test_worker_reads_confidence_flag_once(): + """The flag must be run-constant: a per-step read could diverge the graph.""" + assert _make_worker().return_confidence is False + assert _make_worker(enable_confidence_scheduling=True).return_confidence is True + + def _fake_draft_model(num_stages=3, window_size=128, head_dim=64): return types.SimpleNamespace( num_stages=num_stages, @@ -454,3 +463,154 @@ def fake_sample_draft_tokens(gl, sm, bs, *, num_contexts): # Verified tokens are surfaced unchanged. assert torch.equal(out["new_tokens"], accepted) assert torch.equal(out["new_tokens_lens"], num_accepted) + + +def _stride_probe_draft_model(block_size=5, num_stages=3, window_size=128, head_dim=64): + """Draft model stub that records the ``main_hidden`` it was handed.""" + seen = {} + + def forward_batched(main_hidden, bonus, start_pos, **kwargs): + seen["main_hidden"] = main_hidden.clone() + g = main_hidden.shape[0] + return ( + torch.zeros(g, block_size, dtype=torch.int32, device="cuda"), + None, + torch.zeros(g, block_size, 8, device="cuda"), + ) + + dm = types.SimpleNamespace( + num_stages=num_stages, + block_size=block_size, + _attn_params={"window_size": window_size, "head_dim": head_dim}, + write_context_windows_batched=lambda *a, **k: None, + forward_batched=forward_batched, + ) + return dm, seen + + +@pytest.mark.parametrize("runtime_draft_len", [5, 3, 1]) +def test_gen_draft_gathers_hidden_with_the_runtime_stride(runtime_draft_len): + """The gather stride must follow runtime_draft_len, not max_draft_len. + + The target lays out ``runtime_draft_len + 1`` tokens per generation request. + Striding by ``max_draft_len + 1`` instead walks into the *next* request's + hidden states as soon as the scheduler trims the draft length -- the draft + then conditions on another request's context, with no error anywhere. + """ + num_gens, hidden = 4, HIDDEN * NCAP + worker = _make_worker() + dm, seen = _stride_probe_draft_model() + meta = _make_metadata(max_num_requests=8, max_num_tokens=256) + worker._lazy_init(dm, meta) + worker._batch_to_slot[:num_gens] = torch.arange(num_gens, device="cuda") + + # Row r of request g is tagged with g so a cross-request read is visible. + stride = runtime_draft_len + 1 + captured = torch.zeros(num_gens * stride, hidden, device="cuda") + for g in range(num_gens): + captured[g * stride : (g + 1) * stride] = float(g) + + meta.runtime_draft_len = runtime_draft_len + meta.get_hidden_states = lambda _n: captured + attn = types.SimpleNamespace(num_ctx_tokens=0, num_seqs=num_gens, num_contexts=0) + + # Everyone accepts exactly one token, so the bonus sits at offset 0 of each + # request's block: main_hidden[g] must be exactly g. + nacc = torch.ones(num_gens, dtype=torch.int32, device="cuda") + accepted = torch.zeros(num_gens, worker.max_draft_len + 1, dtype=torch.int32, device="cuda") + + worker._draft_gen_block_batched( + dm, + meta, + attn, + accepted, + nacc, + num_contexts=0, + batch_size=num_gens, + total_target_tokens=captured.shape[0], + ) + + got = seen["main_hidden"][:, 0].tolist() + assert got == [float(g) for g in range(num_gens)], ( + f"runtime_draft_len={runtime_draft_len}: drafted from the wrong request's " + f"hidden states, got {got}" + ) + + +def test_gen_draft_stride_falls_back_to_max_draft_len(): + """Metadata without runtime_draft_len keeps the pre-existing behavior.""" + num_gens, hidden = 3, HIDDEN * NCAP + worker = _make_worker() + dm, seen = _stride_probe_draft_model() + meta = _make_metadata(max_num_requests=8, max_num_tokens=256) + worker._lazy_init(dm, meta) + worker._batch_to_slot[:num_gens] = torch.arange(num_gens, device="cuda") + + stride = worker.max_draft_len + 1 + captured = torch.zeros(num_gens * stride, hidden, device="cuda") + for g in range(num_gens): + captured[g * stride : (g + 1) * stride] = float(g) + + meta.runtime_draft_len = None + meta.get_hidden_states = lambda _n: captured + attn = types.SimpleNamespace(num_ctx_tokens=0, num_seqs=num_gens, num_contexts=0) + nacc = torch.ones(num_gens, dtype=torch.int32, device="cuda") + accepted = torch.zeros(num_gens, worker.max_draft_len + 1, dtype=torch.int32, device="cuda") + + worker._draft_gen_block_batched( + dm, + meta, + attn, + accepted, + nacc, + num_contexts=0, + batch_size=num_gens, + total_target_tokens=captured.shape[0], + ) + assert seen["main_hidden"][:, 0].tolist() == [float(g) for g in range(num_gens)] + + +def test_staged_confidence_rows_follow_the_slots_not_the_leading_rows(): + """The confidence buffer is slot-indexed; a batch rarely owns slots 0..G-1. + + Reading the buffer's leading rows would hand the scheduler a stale, unrelated + set of requests -- wrong budget, no error. + """ + worker = _make_worker(enable_confidence_scheduling=True) + dm, _ = _stride_probe_draft_model() + worker._lazy_init(dm, _make_metadata(max_num_requests=8)) + assert worker._confidence_logits is not None + + # Tag every slot row with its slot index, then say the batch owns 5,2,7. + blk = worker.max_draft_len + for s in range(worker._confidence_logits.shape[0]): + worker._confidence_logits[s] = float(s) + worker._last_gen_slots = torch.tensor([5, 2, 7], device="cuda") + + rows = worker.staged_confidence_rows() + assert rows.shape == (3, blk) + assert rows[:, 0].tolist() == [5.0, 2.0, 7.0] + + +def test_staged_confidence_rows_is_none_before_any_draft(): + worker = _make_worker(enable_confidence_scheduling=True) + dm, _ = _stride_probe_draft_model() + worker._lazy_init(dm, _make_metadata(max_num_requests=8)) + assert worker.staged_confidence_rows() is None + + +def test_confidence_buffer_absent_when_scheduling_is_off(): + """Zero cost when the feature is off.""" + worker = _make_worker(enable_confidence_scheduling=False) + dm, _ = _stride_probe_draft_model() + worker._lazy_init(dm, _make_metadata(max_num_requests=8)) + assert worker._confidence_logits is None + assert worker.verify_planner is None + + +def test_unscored_confidence_slots_read_as_verify_all(): + """A slot never scored must not make the scheduler trim on garbage.""" + worker = _make_worker(enable_confidence_scheduling=True) + dm, _ = _stride_probe_draft_model() + worker._lazy_init(dm, _make_metadata(max_num_requests=8)) + assert torch.all(torch.sigmoid(worker._confidence_logits) > 0.999) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 183481805aed..0b6ed9faaf36 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -3815,3 +3815,51 @@ def test_glm5_resolves_auto_to_python_on_nixl(self): cfg = self._pretrained_config(["GlmMoeDsaForCausalLM"], "glm_moe_dsa") _resolve_transceiver_runtime_auto(args, DeepseekV3ForCausalLM, cfg) assert args.cache_transceiver_config.transceiver_runtime == "PYTHON" + + +class TestDSparkConfidenceScheduling: + """Config surface for DSpark confidence-scheduled verification.""" + + def _cfg(self, **kw): + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + kw.setdefault("max_draft_len", 5) + return DSparkDecodingConfig(**kw) + + def test_defaults_are_off(self): + c = self._cfg() + assert c.enable_confidence_scheduling is False + assert c.confidence_verify_len_tiers is None + # With scheduling off the ladder is just the static length. + assert c.verify_len_tiers == [5] + + def test_knobs_require_the_master_switch(self): + for kw in ({ + "confidence_sts_path": "/tmp/sts.json" + }, { + "confidence_sps_table_path": "/tmp/sps.json" + }, { + "confidence_verify_len_tiers": [1, 5] + }): + with pytest.raises(ValueError, + match="enable_confidence_scheduling"): + self._cfg(**kw) + + def test_default_ladder_is_derived_and_includes_the_full_block(self): + c = self._cfg(enable_confidence_scheduling=True) + assert c.verify_len_tiers == [1, 3, 5] + + def test_full_block_is_always_reachable(self): + """The full block is the fallback when confidence is stale; keep it.""" + c = self._cfg(enable_confidence_scheduling=True, + confidence_verify_len_tiers=[1, 2]) + assert c.verify_len_tiers[-1] == 5 + + def test_ladder_cannot_exceed_max_draft_len(self): + with pytest.raises(ValueError, match="exceeds max_draft_len"): + self._cfg(enable_confidence_scheduling=True, + confidence_verify_len_tiers=[1, 9]) + + def test_ladder_is_deduped_and_sorted(self): + c = self._cfg(enable_confidence_scheduling=True, + confidence_verify_len_tiers=[5, 1, 3, 3]) + assert c.verify_len_tiers == [1, 3, 5] From 87235066080cc71b6e691670dad7d3a2ca6d08c7 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Thu, 30 Jul 2026 18:39:26 -0700 Subject: [PATCH 002/177] [None][fix] Fix host syncs, slot aliasing and cost model in DSpark confidence scheduling Review pass over the previous commit. Every item below is a defect in code that claimed, in its own docstring, to already have the property it lacked -- which is why none of them showed up in the shape-level unit tests. Illegal under CUDA graph capture -------------------------------- * ``row_ids_from_lens`` called ``repeat_interleave`` with a device-resident ``repeats`` and no ``output_size``. Torch then reads the cumulative sum back to the host to size its output. DSpark is a one-engine drafter whose acceptance runs *inside* the target's captured graph, so that sync is not a slowdown, it is a capture failure. The total is known on the host at every call site and is now passed explicitly. * ``apply_sts`` divided host logits by a device buffer. The planner stages confidence to pinned CPU memory and calibrates there, while the head lives on the device -- a plain device-mismatch RuntimeError. It was invisible because the flat-cost-table fallback returns before calibration is ever reached, so only a run with a profiled table would hit it. Confidence read by batch position instead of by slot ---------------------------------------------------- The buffer is slot-indexed, and the snapshot is one iteration old. Joins and departures reshuffle the batch in between, so position ``i`` is routinely a different request than the one scored there -- requests were scheduled against someone else's confidence. The whole slot-indexed buffer is now staged (it is kilobytes) and looked up per request via ``confidence_row_for``. Three more ways the buffer went stale: * The "neutral" fallback row was the *scratch* row, which padding and unknown requests write through. A dedicated row past the window buffer is now never written and therefore genuinely neutral. * A recycled slot kept its previous occupant's scores, so an incoming request's first decision inherited a dead draft's low survival. ``_assign_slot`` resets it, as it already did for the KV window. * Warmup scores synthetic hidden states and did not restore the buffer, so those scores were the first thing the planner ever read. Cost model ---------- * The table is indexed by tokens submitted to the target, which is ``bs * (verify_len + 1)`` -- the bonus token per request was missing, so every lookup was short by a whole ``bs``, easily a shelf's width. Added ``total_verify_tokens`` as the single place this is computed so the planner and the tier derivation cannot drift, and corrected the tier right-edge formula to match. * The fixed per-step overhead was hardcoded to 1.0 ms (copied from vLLM). ``Theta = tau / T`` is a ratio, so non-trimmable terms do not cancel: they move the argmax directly. On a large-MoE deployment ``alpha(bs)`` is tens of milliseconds while ``theta(M)`` is a few, so a guess there dominates the decision. ``SpsCostTable`` now carries ``fixed_overhead_ms`` and an optional ``batch_sizes`` / ``batch_overhead_ms`` staircase, both read from the profile JSON. Note this makes the table format explicit: token counts include bonus tokens. Host syncs added to the *uniform* path -------------------------------------- The previous commit claimed every change was a no-op without per-request windows. It was not: the expanded DSA buffers called ``int(repeats.sum())`` on a device tensor, adding a device->host sync per step to a path that had none. ``expand_per_gen_token`` now computes the total on the host, keeps the scalar ``repeat_interleave`` for uniform batches (no repeats tensor, no copy, no output-size question), and stages the ragged repeats through pinned memory -- a pageable host->device copy synchronizes the stream before it starts. Partially-windowed batches -------------------------- If the scheduler produced fewer windows than requests, the input layout went ragged per request while the spec metadata saw a missing window and stayed uniform -- silent token misattribution. Both the planner and ``fit_ragged_verify_lens`` now refuse and fall back to uniform. Also ---- * ``runtime_draft_len`` defaults to 0, not None, so the earlier stride fix left a stride of 1 on any path that had not populated it. ``or K`` completes it. * Dropped the unread ``ModelEngine.ragged_verify_bucket``; the bucket is re-derived from the batch at graph-key time and cannot go stale. * Config-gated the per-request scans in the graph-key and padding paths, which ran for every model on every step. Testing ------- 217 unit tests pass. The new ones include two that actually capture a CUDA graph over the ragged scatter and accept-count -- the only kind of test that could have caught the sync above. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa.py | 78 ++++++++--- tensorrt_llm/_torch/models/dspark/heads.py | 15 ++- .../_torch/pyexecutor/cuda_graph_runner.py | 20 ++- .../_torch/pyexecutor/model_engine.py | 20 ++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 43 ++++-- tensorrt_llm/_torch/speculative/dspark.py | 89 ++++++++++--- .../_torch/speculative/dspark_planner.py | 122 +++++++++++++++--- .../_torch/speculative/dspark_ragged.py | 20 ++- .../_torch/speculative/dspark_verify.py | 104 +++++++++++---- .../test_dspark_confidence_schedule.py | 107 ++++++++++++--- .../hw_agnostic/test_dspark_cuda_graph.py | 102 +++++++++++++++ .../hw_agnostic/test_dspark_ragged.py | 15 ++- .../hw_agnostic/test_dspark_worker.py | 54 ++++++-- 13 files changed, 641 insertions(+), 148 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index a790c6c07699..92f577ff5d98 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -647,18 +647,62 @@ def is_ragged_verify(self) -> bool: """Whether generation requests have per-request query lengths.""" return self.ragged_verify_lens is not None - def gen_token_repeats(self, device=None) -> torch.Tensor: - """Per-generation-request query-token counts as a repeat vector. + def gen_token_repeat_list(self) -> List[int]: + """Per-generation-request query-token counts, on the host. Uniform batches get `1 + max_draft_tokens` for every request, so the callers below stay bit-identical; ragged batches get the scheduler's per-request windows. + + Host-side on purpose: the expanded token count derived from these feeds + buffer slicing, and reducing it off a device tensor would be a + device->host sync on every step of the *uniform* path too. """ if self.ragged_verify_lens is None: - repeats = [1 + self.max_draft_tokens] * self.num_generations - else: - repeats = self.ragged_verify_lens - return torch.tensor(repeats, dtype=torch.long, device=device) + return [1 + self.max_draft_tokens] * self.num_generations + return list(self.ragged_verify_lens) + + def gen_token_repeats(self, device=None) -> torch.Tensor: + """:meth:`gen_token_repeat_list` as a repeat vector for the given device.""" + return torch.tensor(self.gen_token_repeat_list(), + dtype=torch.long, + device=device) + + def expand_per_gen_token(self, + values: torch.Tensor, + dim: int = 0) -> Tuple[torch.Tensor, int]: + """Repeat `values` once per query token, returning `(expanded, num_tokens)`. + + Uniform batches take the scalar `repeat_interleave` -- no repeats + tensor, no host->device copy, and no output-size question at all, so + this is bit-identical and cost-identical to the fixed-stride expansion + it replaced. + + Ragged batches need a per-request repeats vector, which costs two things + the scalar form does not, and both are paid for here: + + * `repeat_interleave` with a tensor `repeats` reads their cumulative sum + back to the host to size its output unless `output_size` is given, so + the total is computed on the host and passed explicitly; + * materializing `repeats` on the device is a host->device copy, and from + *pageable* memory CUDA synchronizes the stream before starting it. The + staging tensor is therefore pinned and the copy issued + `non_blocking=True`, matching how the rest of the engine stages + host-built index vectors. + """ + repeats = self.gen_token_repeat_list() + num_tokens = sum(repeats) + if self.ragged_verify_lens is None: + stride = 1 + self.max_draft_tokens + return values.repeat_interleave(stride, dim=dim), num_tokens + repeats_host = torch.tensor( + repeats, + dtype=torch.long, + pin_memory=(prefer_pinned() and values.device.type == 'cuda')) + repeats_dev = repeats_host.to(values.device, non_blocking=True) + return (values.repeat_interleave(repeats_dev, + dim=dim, + output_size=num_tokens), num_tokens) def __init__(self, *args, **kwargs): """Initialize DSA metadata with SM count and indexer chunk size.""" @@ -920,12 +964,12 @@ def on_update_kv_lens(self): scheduler_metadata_buffer_full_next_n, non_blocking=True) if self.use_expanded_buffers_for_mtp: # One row per query token; the repeat count is per-request - # under ragged verification and a constant otherwise. - token_repeats = self.gen_token_repeats( - device=gen_indexer_kv_lens.device) - num_tokens = int(token_repeats.sum()) - self.kv_lens_expanded_cuda[:num_tokens] = \ - gen_indexer_kv_lens.repeat_interleave(token_repeats) + # under ragged verification and a constant otherwise. Both the + # host-side total and the sync-free expansion live in + # ``expand_per_gen_token``. + expanded, num_tokens = self.expand_per_gen_token( + gen_indexer_kv_lens) + self.kv_lens_expanded_cuda[:num_tokens] = expanded scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( self.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), _DG_SCHEDULE_BLOCK_KV, self.num_sms) @@ -1516,12 +1560,10 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): if self.use_expanded_buffers_for_mtp: # Expand kv_lens_cuda (only generation): one row per query token, # repeating each request's kv_len as many times as it has tokens. - token_repeats = self.gen_token_repeats() - num_tokens = int(token_repeats.sum()) gen_kv_lens = self.get_indexer_kv_lens( kv_lens[self.num_contexts:self.num_seqs]) - gen_kv_lens_expanded = gen_kv_lens.repeat_interleave( - token_repeats.to(gen_kv_lens.device)) + gen_kv_lens_expanded, num_tokens = self.expand_per_gen_token( + gen_kv_lens) self.kv_lens_expanded_host[:num_tokens].copy_(gen_kv_lens_expanded) self.kv_lens_expanded_cuda[:num_tokens].copy_( self.kv_lens_expanded_host[:num_tokens], non_blocking=True) @@ -1533,8 +1575,8 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): max_len = self.host_indexer_k_cache_block_offsets.shape[1] gen_block_tensor = self.host_indexer_k_cache_block_offsets[ self.num_contexts:self.num_seqs, :max_len] - expanded_blocks = gen_block_tensor.repeat_interleave( - token_repeats.to(gen_block_tensor.device), dim=0) + expanded_blocks, _ = self.expand_per_gen_token(gen_block_tensor, + dim=0) self.host_block_table_expanded[:num_tokens, :max_len].copy_( expanded_blocks, non_blocking=True) self.block_table_expanded[:num_tokens].copy_( diff --git a/tensorrt_llm/_torch/models/dspark/heads.py b/tensorrt_llm/_torch/models/dspark/heads.py index 94529e4288f7..29d217e7c877 100644 --- a/tensorrt_llm/_torch/models/dspark/heads.py +++ b/tensorrt_llm/_torch/models/dspark/heads.py @@ -268,8 +268,19 @@ def forward( return self.proj(features.float()).squeeze(-1) def apply_sts(self, confidence_logits: torch.Tensor) -> torch.Tensor: - """Raw logits -> calibrated per-position acceptance probabilities in [0, 1].""" - return torch.sigmoid(confidence_logits.float() / self.sts_temperatures) + """Raw logits -> calibrated per-position acceptance probabilities in [0, 1]. + + Works on host tensors too. The verification planner stages confidence to + pinned CPU memory and calibrates there, so it calls this with CPU logits + while the head (and therefore ``sts_temperatures``) lives on the device; + without the transfer that is a device-mismatch ``RuntimeError`` on the + first step that reaches calibration. The branch is host-side and the + devices match on the in-graph path, so nothing extra is captured. + """ + temperatures = self.sts_temperatures + if temperatures.device != confidence_logits.device: + temperatures = temperatures.to(confidence_logits.device) + return torch.sigmoid(confidence_logits.float() / temperatures) @torch.no_grad() def load_sts_temperatures(self, temperatures: torch.Tensor) -> None: diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 5a9a9c6e5131..8781a9481a7a 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -285,14 +285,19 @@ def get_graph_key( key = key + (verify_bucket, ) return key - @staticmethod - def _ragged_verify_bucket(batch: ScheduledRequests) -> Optional[int]: + def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: """Total verified token count when the batch uses ragged windows. None unless *every* generation request carries a window, so a batch the scheduler declined to make ragged -- or any non-DSpark batch -- keeps the original key. + + Config-gated before touching the batch: this runs on every graph-key + build for every model, and a per-request scan there is pure overhead for + the ~all of them that will never be ragged. """ + if not getattr(self.spec_config, "enable_ragged_verify", False): + return None total = 0 for request in batch.generation_requests: verify_len = getattr(request, "py_verify_len", None) @@ -656,11 +661,12 @@ def _get_padded_batch(self, batch: ScheduledRequests, # values -- so read it back rather than assume. Cleared otherwise: the # dummy is cached across steps and would carry a stale window into a # uniform batch. - ragged = any( - getattr(request, "py_verify_len", None) is not None - for request in batch.generation_requests) - padding_dummy_request.py_verify_len = (self.ragged_pad_verify_len - if ragged else None) + if getattr(self.spec_config, "enable_ragged_verify", False): + ragged = any( + getattr(request, "py_verify_len", None) is not None + for request in batch.generation_requests) + padding_dummy_request.py_verify_len = (self.ragged_pad_verify_len + if ragged else None) batch.generation_requests.extend([padding_dummy_request] * padding_size) return padding_size diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 193199c1526a..aa856ff7ec3b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -605,12 +605,10 @@ def __init__( # pass. All None without ragged verification, so the incremental-update # gate below never trips for uniform speculation. self.previous_verify_lens = [] - # Captured token bucket this step's ragged batch was fitted onto, or - # None when the batch is uniform. Set once per step by the scheduler. - self.ragged_verify_bucket: Optional[int] = None - # Window every CUDA-graph padding row carries this step (they share one - # dummy object, so they all get the same one). In py_verify_len units. - self.ragged_pad_verify_len: int = 0 + # NOTE: the captured token bucket and the padding rows' window live on + # the CUDA graph runner (``ragged_pad_verify_len``), which is what reads + # them; the bucket itself is re-derived from the batch in + # ``CUDAGraphRunner._ragged_verify_bucket`` so it cannot go stale. self.has_previous_device_draft = False self.previous_accepted_tokens_cuda = torch.empty((self.batch_size, ), dtype=torch.int, @@ -3655,6 +3653,16 @@ def fit_ragged_verify_lens( runner = self.cuda_graph_runner if not runner.enabled or not verify_lens: return None + if len(verify_lens) != len(generation_requests): + # Partially windowing a batch is worse than not windowing it: the + # token layout is built per request and would go ragged, while + # _attach_ragged_verify_layout sees the missing windows and leaves + # the spec metadata uniform. Refuse and let the caller fall back. + logger.debug( + f"DSpark ragged: got {len(verify_lens)} verify lengths for " + f"{len(generation_requests)} generation requests; falling back " + f"to uniform scheduling") + return None tiers = sorted({ int(t) for t in getattr(self.spec_config, "verify_len_tiers", None) or [] diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 267aef9dd966..4488c566e7df 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3205,7 +3205,17 @@ def _dspark_confidence_draft_len(self, if planner is None: return max_draft_len - num_gen_requests = len(scheduled_batch.generation_requests) + gen_requests = scheduled_batch.generation_requests + num_gen_requests = len(gen_requests) + # Confidence is read back by SLOT, not by batch position. The snapshot + # is one iteration old, and joins/departures reshuffle the batch in + # between, so position i is routinely a different request than the one + # that was scored there. A request with no slot yet maps to the neutral + # row, i.e. "verify the whole block". + confidence_rows = [ + worker.confidence_row_for(request.py_request_id) + for request in gen_requests + ] # Cross-rank agreement: draft_len is part of the CUDA-graph key but is # NOT covered by the attention-DP consistency allgather, so ranks that @@ -3229,7 +3239,9 @@ def all_rank_max(value: int) -> int: if getattr(self.model_engine.spec_config, "enable_ragged_verify", False): ragged_lens = planner.decide_verify_lens( - num_gen_requests=num_gen_requests, all_rank_max=all_rank_max) + num_gen_requests=num_gen_requests, + rows=confidence_rows, + all_rank_max=all_rank_max) # Land the batch's token total on a captured bucket before the graph key # is built, spending the rounding slack on real verification. A batch @@ -3248,10 +3260,7 @@ def all_rank_max(value: int) -> int: [len(ragged_lens), sum(1 + int(v) for v in ragged_lens)]) bucket = self.model_engine.fit_ragged_verify_lens( - scheduled_batch.generation_requests, - ragged_lens, - peer_stats=peer_stats) - self.model_engine.ragged_verify_bucket = bucket + gen_requests, ragged_lens, peer_stats=peer_stats) if bucket is not None: # The block is always drafted in full -- only verification is @@ -3263,17 +3272,27 @@ def all_rank_max(value: int) -> int: int(t) for t in self.model_engine.spec_config.verify_len_tiers) else: # Clear any window left over from a previous ragged step, otherwise - # a fallback iteration would keep trimming on a stale split. - for request in scheduled_batch.generation_requests: + # a fallback iteration would keep trimming on a stale split. This + # has to cover EVERY generation request: _prepare_tp_inputs builds + # the token layout per request from py_verify_len, but + # _attach_ragged_verify_layout disables the ragged spec metadata as + # soon as one request is missing a window. A batch where only some + # requests still carry a stale window therefore gets a ragged input + # layout with uniform acceptance -- silent token misattribution. + for request in gen_requests: request.py_verify_len = None runtime_draft_len = planner.decide_draft_len( - num_gen_requests=num_gen_requests, all_rank_max=all_rank_max) + num_gen_requests=num_gen_requests, + rows=confidence_rows, + all_rank_max=all_rank_max) # Stage this step's confidence for the *next* decision. Non-blocking; if # it has not landed by then the planner just verifies the full block. - rows = worker.staged_confidence_rows() - if rows is not None and rows.shape[0] > 0: - planner.stage_confidence(rows, rows.shape[0]) + # The whole slot-indexed buffer is staged, so the next step can look up + # whichever requests it happens to be scheduling. + buffer = worker.staged_confidence_buffer() + if buffer is not None: + planner.stage_confidence(buffer) return int(runtime_draft_len) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index a54a89c6544a..ee7d7ae46eae 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -251,8 +251,9 @@ def __init__( # ``_confidence_logits`` is the slot-indexed handoff buffer to the # verification scheduler; see ``_lazy_init`` for its neutral value. self.return_confidence = bool(getattr(spec_config, "enable_confidence_scheduling", False)) - self._confidence_logits: Optional[torch.Tensor] = None # [max_batch+1, block] - self._last_gen_slots: Optional[torch.Tensor] = None + self._confidence_logits: Optional[torch.Tensor] = None # [max_batch+2, block] + # Row index of the permanently-neutral confidence row; see ``_lazy_init``. + self._neutral_conf_row = 0 # Host-side coordinator that turns the confidence snapshot into this # iteration's draft length. Built in ``_lazy_init`` (it needs the draft # model's block size and calibration head). @@ -270,17 +271,33 @@ def __init__( f"use_separate_draft_kv_cache={use_separate_draft_kv_cache}" ) - def staged_confidence_rows(self) -> Optional[torch.Tensor]: - """Confidence rows for the requests scored by the last draft, or None. + def staged_confidence_buffer(self) -> Optional[torch.Tensor]: + """The whole slot-indexed confidence buffer, or None when disabled. - Indexes through ``_last_gen_slots`` rather than taking the buffer's - leading rows: the buffer is slot-indexed, and a batch rarely occupies - slots ``0..G-1``. Reading the leading rows would score a stale, unrelated - set of requests. + The *whole* buffer, not the current batch's rows: the planner reads it + back by slot, so it must be able to look up any request that is still + alive. It is small (``max_num_requests x block`` fp32 -- kilobytes), so + staging all of it costs nothing and removes the need to communicate a + row ordering across the one-iteration lag. """ - if self._confidence_logits is None or self._last_gen_slots is None: - return None - return self._confidence_logits[self._last_gen_slots] + return self._confidence_logits + + def confidence_row_for(self, req_id: int) -> int: + """Buffer row holding ``req_id``'s confidence, host-side. + + Falls back to the permanently-neutral row for a request that has not + been drafted yet (or was never assigned a slot). Neutral reads as + "certain to be accepted", i.e. verify the full block -- the safe + direction, and the same answer today's static behavior gives. + + Keying by slot rather than by batch position is what makes the + one-iteration-lagged snapshot usable: the batch is reshuffled by joins + and departures between the step that wrote the confidence and the step + that reads it, so position ``i`` is routinely a different request. + """ + if self._confidence_logits is None: + return self._neutral_conf_row + return self._req_to_slot.get(req_id, self._neutral_conf_row) @property def max_draft_len(self) -> int: @@ -335,8 +352,16 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: # accepted", which makes the scheduler fall back to verifying the # full block -- the safe direction (never verifies fewer tokens than # today's static behavior on stale/unwritten rows). + # + # One row beyond the window buffer: rows ``[0, max_batch)`` are real + # request slots and row ``max_batch`` is the scratch row that padded + # / unknown requests write through, so neither is reliably neutral. + # Row ``max_batch + 1`` is never in ``_batch_to_slot`` and therefore + # never written -- that is the row ``confidence_row_for`` hands out + # for an unscored request. + self._neutral_conf_row = num_rows self._confidence_logits = torch.full( - (num_rows, block_size), + (num_rows + 1, block_size), _NEUTRAL_CONFIDENCE_LOGIT, dtype=torch.float32, device="cuda", @@ -371,9 +396,16 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> if table_path: with open(table_path, encoding="utf-8") as f: raw = json.load(f) + # ``bias`` and ``alpha(bs)`` are optional but not cosmetic: Theta is + # a ratio, so the non-trimmable part of the step time moves the + # argmax. Omitting them says "step_time_ms already measures whole + # steps at the deployment's batch size". cost_table = SpsCostTable( token_counts=[int(v) for v in raw["token_counts"]], step_time_ms=[float(v) for v in raw["step_time_ms"]], + fixed_overhead_ms=float(raw.get("fixed_overhead_ms", 0.0)), + batch_sizes=[int(v) for v in raw.get("batch_sizes", [])], + batch_overhead_ms=[float(v) for v in raw.get("batch_overhead_ms", [])], ) else: # Explicitly flat: the planner treats this as "unprofiled" and keeps @@ -444,6 +476,12 @@ def _assign_slot(self, req_id: int, reset: bool) -> int: self._req_to_slot[req_id] = slot self._ctx_len[slot] = 0 self._kv_windows[slot].zero_() + if self._confidence_logits is not None: + # A recycled slot still holds the *previous* occupant's scores. + # Reset to neutral so the incoming request's first scheduling + # decision verifies the full block instead of inheriting a dead + # draft's low survival. + self._confidence_logits[slot].fill_(_NEUTRAL_CONFIDENCE_LOGIT) return self._req_to_slot[req_id] def _seed_context_windows( @@ -519,9 +557,10 @@ def _draft_gen_block_batched( # hidden states whenever runtime_draft_len < max_draft_len, silently # drafting from the wrong context. K = self.max_draft_len - runtime_draft_len = getattr(spec_metadata, "runtime_draft_len", None) - if runtime_draft_len is None: - runtime_draft_len = K + # Note the ``or K``: SpecMetadata.runtime_draft_len defaults to 0, not + # None, so an ``is None`` check would leave a stride of 1 on any path + # that has not populated it yet. + runtime_draft_len = getattr(spec_metadata, "runtime_draft_len", 0) or K Kp1 = int(runtime_draft_len) + 1 device = accepted_tokens.device @@ -589,14 +628,14 @@ def _draft_gen_block_batched( ) # Stash the raw [G, K] confidence logits for the verification scheduler. # Slot-indexed (not batch-position-indexed) so it survives batch - # reshuffles; written in place so a captured graph keeps the same storage. - # ``_last_gen_slots`` is a *view* into ``_batch_to_slot``, so it keeps - # pointing at whichever slots the current batch occupies on replay -- the - # scheduler must read the rows for the requests that were actually - # scored, not slots 0..G-1. + # reshuffles; written in place so a captured graph keeps the same + # storage, and scattered through ``slots`` (a view into + # ``_batch_to_slot``) so a replay lands on whichever slots the current + # batch occupies. The planner reads rows back by slot via + # ``confidence_row_for``, so no row ordering has to survive the + # one-iteration lag. if confidence is not None: self._confidence_logits[slots] = confidence.detach() - self._last_gen_slots = slots return block_logits def _forward_impl( @@ -647,6 +686,12 @@ def _forward_impl( if is_warmup: saved_ctx_len = self._ctx_len.clone() saved_windows = self._kv_windows.clone() + # The confidence buffer is persistent state too: warmup scores + # synthetic hidden states, and those scores would otherwise be the + # first thing the verification planner ever reads. + saved_confidence = ( + None if self._confidence_logits is None else self._confidence_logits.clone() + ) # Assign / reset window slots for context (prefill) requests and seed each # request's rolling KV window from its prompt's captured context, so the @@ -749,6 +794,8 @@ def _forward_impl( if is_warmup: self._ctx_len.copy_(saved_ctx_len) self._kv_windows.copy_(saved_windows) + if saved_confidence is not None: + self._confidence_logits.copy_(saved_confidence) return { "logits": raw_logits, diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py index 52ba92824bb2..0276d1ba0656 100644 --- a/tensorrt_llm/_torch/speculative/dspark_planner.py +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -39,6 +39,22 @@ step time against the *total* verified token count, and interpolates nothing across a riser. +**Total verified tokens includes the bonus position.** A request that verifies +``L`` drafted positions submits ``L + 1`` tokens to the target -- the bonus +token it already holds, plus the ``L`` drafts. So a batch of ``bs`` requests at +uniform length ``L`` costs ``step_time(bs * (L + 1))``, not ``step_time(bs * L)``. +The two differ by a whole ``bs``, which is easily a shelf's width. + +**The cost is two-dimensional, and only one dimension is trimmable.** SGLang +publishes ``T(bs, K) = bias + alpha(bs) + theta(M)``: a fixed cost, a +batch-size-dependent cost (the draft pass, weight movement -- untouched by +trimming), and the target's verify-token cost ``theta(M)``, which is the only +term a smaller ``M`` reduces. ``Theta = tau / T`` is a *ratio*, so the +non-trimmable terms are not a constant that cancels -- getting them wrong moves +the argmax directly. They therefore have to be supplied, not guessed: +:class:`SpsCostTable` carries ``fixed_overhead_ms`` for ``bias`` and an optional +``batch_sizes`` / ``batch_overhead_ms`` staircase for ``alpha(bs)``. + The planner runs on the host: its output selects a CUDA graph, and graph selection is a host decision in TensorRT-LLM. Keeping it here (rather than on device) is deliberate -- a device-side budget would have to be copied back @@ -46,7 +62,7 @@ """ import bisect -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import List, Optional, Sequence import numpy as np @@ -56,11 +72,18 @@ "compute_verify_token_budget", "budget_argmax_over_uniform_lens", "derive_verify_len_tiers", + "total_verify_tokens", ] -# Fixed per-step overhead folded into every cost lookup (sampling, bookkeeping, -# launch latency). Matches vLLM PR #47808's ``_FIXED_OVERHEAD_MS``. -_FIXED_OVERHEAD_MS = 1.0 + +def total_verify_tokens(num_requests: int, verify_len: int) -> int: + """Tokens the target scores for ``num_requests`` each verifying ``verify_len`` drafts. + + One bonus/anchor token per request plus its drafted positions. This is the + unit :class:`SpsCostTable` is indexed by; the single place it is computed so + the planner and the tier derivation cannot drift apart. + """ + return int(num_requests) * (int(verify_len) + 1) @dataclass(frozen=True) @@ -68,8 +91,24 @@ class SpsCostTable: """Measured decode step time as a function of total verified tokens. Attributes: - token_counts: strictly increasing total-token breakpoints. - step_time_ms: measured step time at each breakpoint. + token_counts: strictly increasing total-token breakpoints. "Total" means + what :func:`total_verify_tokens` computes -- bonus tokens included. + step_time_ms: measured ``theta(M)`` at each breakpoint. + fixed_overhead_ms: the ``bias`` term -- per-step cost that trimming + cannot touch (sampling, bookkeeping, launch latency), if it is not + already inside ``step_time_ms``. Defaults to 0.0 on the assumption + that a profiled table measured whole steps; set it only when the + table isolates the verify term. + batch_sizes / batch_overhead_ms: optional ``alpha(bs)`` staircase -- the + batch-size-dependent, *non-trimmable* part of the step (the draft + pass, weight movement). Looked up by floor on ``batch_sizes``. + + Getting the non-trimmable terms wrong is not harmless. The planner maximizes + ``tau / T``, a ratio, so ``bias + alpha(bs)`` does not cancel: understating + it makes every verified token look proportionally more expensive than it is + and the planner over-trims. On a large-MoE deployment ``alpha(bs)`` is tens + of milliseconds while ``theta(M)`` is a few, so a hardcoded guess here would + dominate the decision. Lookup is a floor (staircase) lookup, never an interpolation: between two measured points the cost is genuinely flat until a new kernel wave starts, @@ -80,6 +119,9 @@ class SpsCostTable: token_counts: Sequence[int] step_time_ms: Sequence[float] + fixed_overhead_ms: float = 0.0 + batch_sizes: Sequence[int] = field(default_factory=tuple) + batch_overhead_ms: Sequence[float] = field(default_factory=tuple) def __post_init__(self) -> None: if len(self.token_counts) != len(self.step_time_ms): @@ -93,18 +135,45 @@ def __post_init__(self) -> None: raise ValueError("token_counts must be strictly increasing") if any(t <= 0.0 for t in self.step_time_ms): raise ValueError("step_time_ms entries must be positive") - - def step_time(self, num_tokens: int) -> float: - """Step time (ms) for a step that verifies ``num_tokens`` tokens in total.""" + if self.fixed_overhead_ms < 0.0: + raise ValueError("fixed_overhead_ms must be >= 0") + if len(self.batch_sizes) != len(self.batch_overhead_ms): + raise ValueError( + f"batch_sizes ({len(self.batch_sizes)}) and batch_overhead_ms " + f"({len(self.batch_overhead_ms)}) must have the same length" + ) + if any(b <= a for a, b in zip(self.batch_sizes, self.batch_sizes[1:])): + raise ValueError("batch_sizes must be strictly increasing") + if any(t < 0.0 for t in self.batch_overhead_ms): + raise ValueError("batch_overhead_ms entries must be >= 0") + + def batch_overhead(self, num_requests: int) -> float: + """``alpha(num_requests)`` -- 0.0 when the table has no batch axis.""" + if not self.batch_sizes: + return 0.0 + idx = bisect.bisect_right(self.batch_sizes, int(num_requests)) - 1 + idx = min(max(idx, 0), len(self.batch_overhead_ms) - 1) + return float(self.batch_overhead_ms[idx]) + + def step_time(self, num_tokens: int, num_requests: int = 0) -> float: + """``bias + alpha(num_requests) + theta(num_tokens)``, in ms.""" idx = bisect.bisect_right(self.token_counts, int(num_tokens)) - 1 idx = min(max(idx, 0), len(self.step_time_ms) - 1) - return float(self.step_time_ms[idx]) + _FIXED_OVERHEAD_MS - - def step_times(self, num_tokens: np.ndarray) -> np.ndarray: - """Vectorized :meth:`step_time`.""" + return ( + float(self.step_time_ms[idx]) + + self.fixed_overhead_ms + + self.batch_overhead(num_requests) + ) + + def step_times(self, num_tokens: np.ndarray, num_requests: int = 0) -> np.ndarray: + """Vectorized :meth:`step_time` over a token-count array.""" counts = np.asarray(self.token_counts) idx = np.clip(np.searchsorted(counts, num_tokens, side="right") - 1, 0, len(counts) - 1) - return np.asarray(self.step_time_ms, dtype=np.float64)[idx] + _FIXED_OVERHEAD_MS + return ( + np.asarray(self.step_time_ms, dtype=np.float64)[idx] + + self.fixed_overhead_ms + + self.batch_overhead(num_requests) + ) @classmethod def flat(cls, step_time_ms: float = 1.0) -> "SpsCostTable": @@ -119,6 +188,12 @@ def flat(cls, step_time_ms: float = 1.0) -> "SpsCostTable": @property def is_flat(self) -> bool: + """Whether the *trimmable* term is constant. + + Only ``theta(M)`` matters here: ``bias`` and ``alpha(bs)`` are the same + for every candidate length, so a table that varies only in those cannot + distinguish a cheap verified token from an expensive one either. + """ return len(set(self.step_time_ms)) <= 1 @@ -168,9 +243,12 @@ def compute_verify_token_budget( # * the n admitted candidates, best-first. base = float(num_gen_requests) + float(surv[:, :min_verify_len].sum()) tau = base + np.concatenate(([0.0], np.cumsum(candidates))) - floor_tokens = bs * int(min_verify_len) - total_tokens = floor_tokens + np.arange(tau.size) - theta = tau / cost_table.step_times(total_tokens) + # Tokens actually submitted to the target at the floor: every request's + # bonus position plus its floor draft positions. Dropping the bonus term + # under-reports the batch by a whole ``bs``, which is easily a shelf wide. + floor_tokens = total_verify_tokens(bs, min_verify_len) + tokens = floor_tokens + np.arange(tau.size) + theta = tau / cost_table.step_times(tokens, num_gen_requests) return int(np.argmax(theta)) @@ -196,7 +274,7 @@ def derive_verify_len_tiers( **The zero-loss property does not survive across batch sizes.** The cost shelves live in token space, so a shelf's right edge in *length* space is - ``(breakpoint - 1) // batch_size`` -- a function of the batch size. A + ``(breakpoint - 1) // batch_size - 1`` -- a function of the batch size. A deployment captures graphs for many batch sizes but can only capture one ladder, so a single tier set cannot sit on every batch size's right edges. Deriving tiers is more robust than hardcoding a pair like ``{1, 5}``, but it @@ -224,8 +302,10 @@ def derive_verify_len_tiers( edges = set() for breakpoint_tokens in cost_table.token_counts[1:]: # Largest length whose total token count still sits on the shelf below - # this breakpoint. - length = (int(breakpoint_tokens) - 1) // int(num_requests) + # this breakpoint. Total tokens for length L is num_requests * (L + 1) + # (see total_verify_tokens), so invert that, not num_requests * L -- + # otherwise the derived edge sits one shelf too far right. + length = (int(breakpoint_tokens) - 1) // int(num_requests) - 1 if lo <= length <= hi: edges.add(length) # The last shelf extends past the final breakpoint, so the whole block is @@ -282,7 +362,7 @@ def budget_argmax_over_uniform_lens( # bonus token plus the survival of *every* verified position (including # the floor ones), summed over the batch. tau = float(num_gen_requests) + float(surv[:, :length].sum()) - theta = tau / cost_table.step_time(bs * length) + theta = tau / cost_table.step_time(total_verify_tokens(bs, length), num_gen_requests) if theta > best_theta: best_len, best_theta = length, theta return int(best_len) diff --git a/tensorrt_llm/_torch/speculative/dspark_ragged.py b/tensorrt_llm/_torch/speculative/dspark_ragged.py index 69c233e1f127..a1b36a1821ea 100644 --- a/tensorrt_llm/_torch/speculative/dspark_ragged.py +++ b/tensorrt_llm/_torch/speculative/dspark_ragged.py @@ -408,14 +408,24 @@ def ragged_gather_index_lists( return rows, cols -def row_ids_from_lens(verify_lens: torch.Tensor) -> torch.Tensor: +def row_ids_from_lens(verify_lens: torch.Tensor, *, total: int) -> torch.Tensor: """``[bs]`` lengths -> ``[total]`` owning-request id for each packed token. ``repeat_interleave`` with a *tensor* of repeats is what makes the packing ragged; the uniform paths use a scalar repeat and get a fixed stride. + + ``total`` (``sum(verify_lens)``) is mandatory rather than derived. With a + device-resident ``repeats`` and no ``output_size``, torch reads the + cumulative sum back to the host to size its output -- a sync that is + outright *illegal* here, because DSpark is a one-engine drafter whose + acceptance runs inside the target's captured CUDA graph. The callers all + know the total on the host already (``spec_metadata.total_verify_tokens``, + or the packed buffer's own leading dimension). """ return torch.repeat_interleave( - torch.arange(verify_lens.numel(), device=verify_lens.device), verify_lens.to(torch.long) + torch.arange(verify_lens.numel(), device=verify_lens.device), + verify_lens.to(torch.long), + output_size=int(total), ) @@ -434,9 +444,13 @@ def scatter_ragged_to_padded( (compare draft vs target position by position). This bridges the two without a host sync: token ``t`` belongs to request ``row_ids[t]`` at column ``t - qo_indptr[row_ids[t]]``. + + ``flat.shape[0]`` is the packed token total, so it is exactly the + ``output_size`` :func:`row_ids_from_lens` needs to stay sync-free (and + therefore capturable). """ bs = int(verify_lens.numel()) - rows = row_ids_from_lens(verify_lens) + 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 diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index e3d522464869..e3b816007ce2 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -89,28 +89,38 @@ def __init__( "fallback_no_snapshot": 0, "fallback_flat_cost": 0, "fallback_no_confidence": 0, + "fallback_short_snapshot": 0, } @property def max_tier(self) -> int: return self.tiers[-1] if self.tiers else self.cfg.resolved_max_verify_len - def stage_confidence(self, confidence_logits: torch.Tensor, num_rows: int) -> None: - """Start a non-blocking device->host copy of this step's confidence. + def stage_confidence(self, confidence_logits: torch.Tensor) -> None: + """Start a non-blocking device->host copy of the confidence buffer. + + ``confidence_logits`` is the worker's whole slot-indexed buffer, staged + in full so :meth:`decide_verify_lens` can look up any live request by + slot. It is kilobytes (``max_num_requests x block`` fp32), so there is + nothing to gain from staging a subset -- and a subset would need a row + ordering to survive the one-iteration lag, which is exactly what slot + keying exists to avoid. Never synchronizes: the copy is issued on the current stream and an event is recorded. :meth:`decide_draft_len` polls that event and simply does not use the snapshot if it has not landed. """ - if confidence_logits is None or num_rows <= 0: + if confidence_logits is None or confidence_logits.shape[0] == 0: return - rows = confidence_logits[:num_rows] - if self._host_buffer is None or self._host_buffer.shape != rows.shape: + if self._host_buffer is None or self._host_buffer.shape != confidence_logits.shape: self._host_buffer = torch.empty( - rows.shape, dtype=torch.float32, device="cpu", pin_memory=prefer_pinned() + confidence_logits.shape, + dtype=torch.float32, + device="cpu", + pin_memory=prefer_pinned(), ) self._copy_event = torch.cuda.Event() - self._host_buffer.copy_(rows, non_blocking=True) + self._host_buffer.copy_(confidence_logits, non_blocking=True) self._copy_event.record() self._snapshot_valid = True @@ -123,10 +133,48 @@ def _ready_snapshot(self) -> Optional[torch.Tensor]: return None return self._host_buffer + def _gather_rows( + self, *, num_gen_requests: int, rows: Optional[Sequence[int]] + ) -> Optional[torch.Tensor]: + """This step's confidence, one row per generation request, or None. + + ``rows`` is the per-request buffer row index (``DSparkWorker. + confidence_row_for``). Supplying it is what makes the lagged snapshot + correct: the batch is reshuffled between the step that wrote the + confidence and the step that reads it, so the request at position ``i`` + is routinely a different one. + + ``rows=None`` keeps the positional reading for callers that own the + ordering themselves (the unit tests stage a purpose-built buffer). That + path refuses to run on a snapshot with fewer rows than the batch rather + than silently returning a short answer -- a short answer becomes a + partially-assigned batch downstream, where half the requests get a + verify window and half do not. + """ + snapshot = self._ready_snapshot() + if snapshot is None: + self.stats["fallback_no_snapshot"] += 1 + return None + if rows is None: + if snapshot.shape[0] < num_gen_requests: + self.stats["fallback_short_snapshot"] += 1 + return None + selected = snapshot[:num_gen_requests] + else: + if len(rows) != num_gen_requests: + self.stats["fallback_short_snapshot"] += 1 + return None + selected = snapshot[torch.as_tensor(list(rows), dtype=torch.long)] + if selected.numel() == 0: + self.stats["fallback_no_confidence"] += 1 + return None + return selected + def decide_draft_len( self, *, num_gen_requests: int, + rows: Optional[Sequence[int]] = None, all_rank_max: Optional[Callable[[int], int]] = None, ) -> int: """Choose this iteration's draft length, agreed across all ranks. @@ -136,11 +184,13 @@ def decide_draft_len( snapshot yet, or an unprofiled (flat) cost model under which every extra verified token looks free. - ``all_rank_max`` overrides the constructor's reduction; the caller that - owns the distributed handle usually supplies it here. + ``rows`` maps each generation request to its buffer row; see + :meth:`_gather_rows`. ``all_rank_max`` overrides the constructor's + reduction; the caller that owns the distributed handle usually supplies + it here. """ self.stats["decisions"] += 1 - chosen = self._decide_local(num_gen_requests=num_gen_requests) + chosen = self._decide_local(num_gen_requests=num_gen_requests, rows=rows) all_rank_max = all_rank_max or self.all_rank_max if all_rank_max is not None: # Max, not min: a rank that wanted to trim more simply verifies a few @@ -155,6 +205,7 @@ def decide_verify_lens( self, *, num_gen_requests: int, + rows: Optional[Sequence[int]] = None, all_rank_max: Optional[Callable[[int], int]] = None, ) -> Optional[List[int]]: """Per-request verify lengths for a ragged step, or None to stay uniform. @@ -173,22 +224,23 @@ def decide_verify_lens( The batch-wide *maximum* is still reduced across ranks, because the drafted-token buffer width and the per-request padding are sized from it; the individual lengths stay local. + + The returned list always has exactly ``num_gen_requests`` entries, or is + None. A partial list would leave the tail of the batch without a verify + window, and the callers downstream disagree about what that means: the + input layout is built per request (so it goes ragged) while the spec + metadata sees a ``None`` and stays uniform. """ if num_gen_requests <= 0: return None if self.cost_table is None or self.cost_table.is_flat: self.stats["fallback_flat_cost"] += 1 return None - snapshot = self._ready_snapshot() - if snapshot is None: - self.stats["fallback_no_snapshot"] += 1 - return None - rows = snapshot[:num_gen_requests] - if rows.numel() == 0: - self.stats["fallback_no_confidence"] += 1 + selected = self._gather_rows(num_gen_requests=num_gen_requests, rows=rows) + if selected is None: return None - probs = self.apply_calibration(rows) + probs = self.apply_calibration(selected) survival = compute_survival(probs) uniform_len = budget_argmax_over_uniform_lens( survival=survival.numpy().astype(np.float64), @@ -205,23 +257,21 @@ def decide_verify_lens( if all_rank_max is not None: agreed_max = int(all_rank_max(int(max(lens)))) lens = [min(int(v), agreed_max) for v in lens] + assert len(lens) == num_gen_requests, ( + f"internal: produced {len(lens)} verify lengths for {num_gen_requests} requests" + ) return [int(v) for v in lens] - def _decide_local(self, *, num_gen_requests: int) -> int: + def _decide_local(self, *, num_gen_requests: int, rows: Optional[Sequence[int]] = None) -> int: if num_gen_requests <= 0: return self.max_tier if self.cost_table is None or self.cost_table.is_flat: self.stats["fallback_flat_cost"] += 1 return self.max_tier - snapshot = self._ready_snapshot() - if snapshot is None: - self.stats["fallback_no_snapshot"] += 1 - return self.max_tier - rows = snapshot[:num_gen_requests] - if rows.numel() == 0: - self.stats["fallback_no_confidence"] += 1 + selected = self._gather_rows(num_gen_requests=num_gen_requests, rows=rows) + if selected is None: return self.max_tier - probs = self.apply_calibration(rows) + probs = self.apply_calibration(selected) survival = compute_survival(probs).numpy().astype(np.float64) return budget_argmax_over_uniform_lens( survival=survival, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py index 33658caf198e..fd3bf13774e0 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py @@ -78,15 +78,18 @@ def _brute_force_budget(survival, num_gen, table, min_len=1, max_len=None): Spells out the objective: expected emitted tokens per millisecond, where the expectation counts one bonus token per request, the always-verified floor - positions, and the admitted candidates. + positions, and the admitted candidates. The cost is looked up on the tokens + actually submitted to the target -- ``bs * (min_len + 1)`` at the floor, + because every request also submits its bonus position. """ bs, blk = survival.shape cap = min(max_len or blk, blk) cand = np.sort(survival[:, min_len:cap].reshape(-1))[::-1] base = num_gen + float(survival[:, :min_len].sum()) + floor_tokens = bs * (min_len + 1) best_n, best = 0, -np.inf for n in range(cand.size + 1): - theta = (base + float(cand[:n].sum())) / table.step_time(bs * min_len + n) + theta = (base + float(cand[:n].sum())) / table.step_time(floor_tokens + n, num_gen) if theta > best: best, best_n = theta, n return best_n @@ -99,7 +102,11 @@ def test_tau_counts_the_always_verified_floor_positions(): silently moves the optimum. Two batches that differ only in how good their floor positions are must be able to choose different budgets. """ - table = SpsCostTable(token_counts=(0, 6), step_time_ms=(1.0, 3.0)) + # The riser sits just past the floor (bs*(min_len+1) = 8 tokens), so both + # shelves are reachable and the constant term can actually decide between + # them. With the riser outside the reachable range this test would still + # pass but prove nothing. + table = SpsCostTable(token_counts=(0, 10), step_time_ms=(1.0, 3.0), fixed_overhead_ms=1.0) bs = 4 good_floor = np.cumprod(np.full((bs, BLOCK), 0.99), axis=1) bad_floor = good_floor.copy() @@ -107,10 +114,16 @@ def test_tau_counts_the_always_verified_floor_positions(): bad_floor = np.cumprod( np.concatenate([bad_floor[:, :1], np.full((bs, BLOCK - 1), 0.99)], axis=1), axis=1 ) + budgets = [] for surv in (good_floor, bad_floor): - assert compute_verify_token_budget( - survival=surv, num_gen_requests=bs, cost_table=table - ) == _brute_force_budget(surv, bs, table) + got = compute_verify_token_budget(survival=surv, num_gen_requests=bs, cost_table=table) + assert got == _brute_force_budget(surv, bs, table) + budgets.append(got) + assert budgets[0] != budgets[1], ( + "the two batches differ only in their floor positions' survival, so a " + "planner that ignored the floor's contribution to tau would give them " + "the same budget -- keep the riser inside the reachable token range" + ) def test_budget_matches_brute_force_on_staircase(): @@ -139,8 +152,11 @@ def test_greedy_first_descent_would_be_wrong(): shelf's optimum is almost 3x better. """ bs = 4 - # One riser at 6 total tokens, then a long shelf out to the end of the range. - table = SpsCostTable(token_counts=(0, 6), step_time_ms=(1.0, 3.0)) + # One riser just past the floor (bs*(min_len+1) = 8 tokens), then a long + # shelf out to the end of the range. The fixed overhead is spelled out + # rather than assumed: it is part of what makes the first shelf's Theta + # look competitive, and it is no longer a hidden constant in the planner. + table = SpsCostTable(token_counts=(0, 10), step_time_ms=(1.0, 3.0), fixed_overhead_ms=1.0) # Uniformly high confidence: every extra candidate adds nearly a full token, # so tau grows ~7x across the range while cost only doubles. surv = np.cumprod(np.full((bs, BLOCK), 0.995), axis=1) @@ -149,7 +165,7 @@ def test_greedy_first_descent_would_be_wrong(): cand = np.sort(surv[:, 1:].reshape(-1))[::-1] def theta(n): - return (bs + cand[:n].sum()) / table.step_time(bs + n) + return (bs + cand[:n].sum()) / table.step_time(2 * bs + n, bs) greedy = 0 while greedy + 1 <= cand.size and theta(greedy + 1) > theta(greedy): @@ -337,12 +353,14 @@ def test_discrete_argmax_beats_round_then_snap_across_a_riser(): cost_table=table, allowed_lens=[1, 3, 5, 7], ) - # 8*7=56 tokens lands on the 20ms riser; the search must avoid it. + # 8*(7+1)=64 tokens lands on the 20ms riser; the search must avoid it. assert chosen < 7 def theta(length): - tau = 8 + float(surv[:, 1:length].sum()) - return tau / table.step_time(8 * length) + # Mirrors the implementation exactly, including the bonus token every + # request submits alongside its drafts. + tau = 8 + float(surv[:, :length].sum()) + return tau / table.step_time(8 * (length + 1), 8) assert theta(chosen) == max(theta(v) for v in (1, 3, 5, 7)) @@ -367,7 +385,7 @@ def test_derived_tiers_lose_nothing_versus_continuous_k(): ) def theta(length): - return (bs + float(surv[:, :length].sum())) / table.step_time(bs * length) + return (bs + float(surv[:, :length].sum())) / table.step_time(bs * (length + 1), bs) best_any = max(theta(v) for v in range(1, BLOCK + 1)) best_tier = max(theta(v) for v in tiers) @@ -386,17 +404,18 @@ def test_derived_tiers_respect_the_capture_budget(): def test_derived_tiers_are_a_function_of_batch_size(): - """Tiers move with bs because total tokens = bs * length. + """Tiers move with bs because total tokens = bs * (length + 1). - With a riser at 40 tokens: at bs=8 the shelf ends at length 4, so 4 is a - tier; at bs=32 even length 2 is already past the riser, so the shelf yields - nothing and only the endpoints remain. A tier set derived once and reused - across batch sizes would be wrong for one of them. + With a riser at 40 tokens: at bs=8, length 3 costs 8*4=32 tokens (still on + the shelf) while length 4 costs 8*5=40 (past it), so 3 is the right edge; at + bs=32 even length 1 is already past the riser, so the shelf yields nothing + and only the endpoints remain. A tier set derived once and reused across + batch sizes would be wrong for one of them. """ table = SpsCostTable(token_counts=(0, 40), step_time_ms=(1.0, 5.0)) assert derive_verify_len_tiers( cost_table=table, num_requests=8, block_size=BLOCK, max_tiers=99 - ) == [1, 4, BLOCK] + ) == [1, 3, BLOCK] assert derive_verify_len_tiers( cost_table=table, num_requests=32, block_size=BLOCK, max_tiers=99 ) == [1, BLOCK] @@ -510,6 +529,56 @@ def test_planner_empty_batch_is_max_tier(): assert _planner().decide_draft_len(num_gen_requests=0) == BLOCK +def test_planner_refuses_a_snapshot_that_does_not_cover_the_batch(): + """A short snapshot must fall back, never return a short answer. + + The staged snapshot lags one iteration, so a batch that grew since then has + more requests than the snapshot has rows. Returning one length per staged + row leaves the tail of the batch without a verify window -- and the two + consumers disagree about what that means: the input layout is built per + request (so it goes ragged) while the spec metadata sees a missing window + and stays uniform. That combination misattributes one request's drafts to + another with nothing raising. + """ + p = _planner() + p._host_buffer, p._copy_event, p._snapshot_valid = torch.zeros(4, BLOCK), None, True + assert p.decide_verify_lens(num_gen_requests=9) is None + assert p.stats["fallback_short_snapshot"] == 1 + assert p.decide_draft_len(num_gen_requests=9) == BLOCK + + +def test_planner_returns_one_verify_len_per_request(): + p = _planner() + p._host_buffer, p._copy_event, p._snapshot_valid = torch.rand(6, BLOCK), None, True + lens = p.decide_verify_lens(num_gen_requests=6) + assert lens is not None and len(lens) == 6 + + +def test_planner_reads_confidence_by_row_not_by_batch_position(): + """The snapshot is slot-indexed; ``rows`` is what re-associates it. + + The buffer is written by slot and read one iteration later, by which point + joins and departures have reshuffled the batch. Position 0 of this step's + batch is routinely a different request than the one scored in row 0. + """ + p = _planner() + # Row 3 is a confident request, row 0 a hopeless one. + buf = torch.full((5, BLOCK), -8.0) + buf[3] = 8.0 + p._host_buffer, p._copy_event, p._snapshot_valid = buf, None, True + + # This step schedules the confident request first, then the hopeless one. + lens = p.decide_verify_lens(num_gen_requests=2, rows=[3, 0]) + assert lens is not None and lens[0] > lens[1], f"row mapping ignored: {lens}" + + +def test_planner_rejects_a_row_list_that_does_not_match_the_batch(): + p = _planner() + p._host_buffer, p._copy_event, p._snapshot_valid = torch.rand(8, BLOCK), None, True + assert p.decide_verify_lens(num_gen_requests=4, rows=[0, 1]) is None + assert p.stats["fallback_short_snapshot"] == 1 + + def test_planner_uses_the_supplied_calibration(): """STS must be applied before the cumprod, not skipped.""" seen = {} diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py index 11969a8ca0f9..a799d4b6c174 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py @@ -352,3 +352,105 @@ def test_sts_table_update_is_visible_to_a_captured_graph(): "the captured graph did not see the new STS temperatures" ) assert torch.allclose(out, torch.sigmoid(head(hid) / 4.0), atol=1e-5) + + +def test_apply_sts_accepts_host_logits_from_a_device_head(): + """The planner calibrates on the host; the head lives on the device. + + Confidence is staged to pinned CPU memory and the verify planner calls + ``apply_sts`` on those rows, while ``sts_temperatures`` is a buffer that + moved to CUDA with the draft model. Without the transfer this is a device + mismatch -- and it is invisible until a run supplies a profiled cost table, + because the flat-table fallback returns before calibration is ever reached. + """ + from tensorrt_llm._torch.models.dspark.heads import DSparkConfidenceHead + + BLK = 5 + head = DSparkConfidenceHead(hidden_size=8, block_size=BLK) + if torch.cuda.is_available(): + head = head.cuda() + head.load_sts_temperatures(torch.full((BLK,), 2.0)) + + host_logits = torch.zeros(3, BLK) # CPU, as the planner stages them + probs = head.apply_sts(host_logits) + assert probs.device.type == "cpu" + assert torch.allclose(probs, torch.full((3, BLK), 0.5), atol=1e-6) + + +# -------------------------------------------------------------------------- +# Ragged verification runs inside the same captured graph: DSpark's acceptance +# is part of the target's forward. The scatter that unpacks a packed ragged +# batch into the rectangle acceptance wants must therefore capture -- which it +# only does if ``repeat_interleave`` is told its output size up front. Without +# that, torch reads the cumulative sum back to the host to size the result, +# which is illegal under capture. +# -------------------------------------------------------------------------- + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +def test_ragged_scatter_is_cuda_graph_capturable(): + from tensorrt_llm._torch.speculative.dspark_ragged import ( + build_qo_indptr, + scatter_ragged_to_padded, + ) + + lens = torch.tensor([3, 1, 2], dtype=torch.int32, device="cuda") + indptr = build_qo_indptr(lens) + flat = torch.tensor([10, 11, 12, 20, 30, 31], dtype=torch.int32, device="cuda") + out = torch.empty(3, 4, dtype=torch.int32, device="cuda") + + def run(): + return scatter_ragged_to_padded( + flat, verify_lens=lens, qo_indptr=indptr, max_len=4, pad_value=-1 + ) + + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out.copy_(run()) + + g.replay() + torch.cuda.synchronize() + assert out.tolist() == [[10, 11, 12, -1], [20, -1, -1, -1], [30, 31, -1, -1]] + + # New packed contents must flow through on replay. + flat.copy_(torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32, device="cuda")) + g.replay() + torch.cuda.synchronize() + assert out.tolist() == [[1, 2, 3, -1], [4, -1, -1, -1], [5, 6, -1, -1]] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +def test_ragged_accept_count_is_cuda_graph_capturable(): + from tensorrt_llm._torch.speculative.dspark_ragged import count_accepted_ragged + + lens = torch.tensor([3, 1, 2], dtype=torch.int32, device="cuda") + draft = torch.tensor([[1, 2, 3, 9], [4, 9, 9, 9], [5, 6, 9, 9]], device="cuda") + target = torch.tensor([[1, 2, 7, 9], [4, 9, 9, 9], [5, 6, 9, 9]], device="cuda") + out = torch.empty(3, dtype=torch.int64, device="cuda") + + def run(): + return count_accepted_ragged(draft_tokens=draft, target_tokens=target, verify_lens=lens) + + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out.copy_(run()) + + g.replay() + torch.cuda.synchronize() + # req 0 matches 2 then diverges; req 1 matches its single position; req 2 + # matches both. Column 3 is padding for every row and must never count. + assert out.tolist() == [2, 1, 2] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py index 132d01620c10..c1d9ce23232f 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py @@ -290,7 +290,20 @@ def test_ragged_beats_uniform_at_equal_budget(): 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).tolist() == [0, 0, 0, 1, 2, 2] + assert row_ids_from_lens(lens, total=6).tolist() == [0, 0, 0, 1, 2, 2] + + +def test_row_ids_requires_the_host_total(): + """``total`` is what keeps the expansion sync-free, so it cannot be optional. + + Without ``output_size`` torch reads the cumulative sum back to the host to + size its output. DSpark's acceptance runs inside the target's captured CUDA + graph, where that sync is illegal -- so the signature forces every caller to + supply a total it already knows. + """ + 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(): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 1ecf85b8f332..2bb9275c4ec5 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -570,33 +570,65 @@ def test_gen_draft_stride_falls_back_to_max_draft_len(): assert seen["main_hidden"][:, 0].tolist() == [float(g) for g in range(num_gens)] -def test_staged_confidence_rows_follow_the_slots_not_the_leading_rows(): +def test_confidence_rows_are_looked_up_by_slot(): """The confidence buffer is slot-indexed; a batch rarely owns slots 0..G-1. - Reading the buffer's leading rows would hand the scheduler a stale, unrelated - set of requests -- wrong budget, no error. + The snapshot the scheduler reads is one iteration old, and joins/departures + reshuffle the batch in between, so batch position is not a usable key. + ``confidence_row_for`` re-associates each request with the row its own draft + wrote. """ worker = _make_worker(enable_confidence_scheduling=True) dm, _ = _stride_probe_draft_model() worker._lazy_init(dm, _make_metadata(max_num_requests=8)) assert worker._confidence_logits is not None - # Tag every slot row with its slot index, then say the batch owns 5,2,7. - blk = worker.max_draft_len + # Tag every row with its index, then hand three requests explicit slots. for s in range(worker._confidence_logits.shape[0]): worker._confidence_logits[s] = float(s) - worker._last_gen_slots = torch.tensor([5, 2, 7], device="cuda") + worker._req_to_slot = {100: 5, 101: 2, 102: 7} - rows = worker.staged_confidence_rows() - assert rows.shape == (3, blk) - assert rows[:, 0].tolist() == [5.0, 2.0, 7.0] + rows = [worker.confidence_row_for(r) for r in (100, 101, 102)] + assert rows == [5, 2, 7] + buf = worker.staged_confidence_buffer() + assert buf[rows][:, 0].tolist() == [5.0, 2.0, 7.0] -def test_staged_confidence_rows_is_none_before_any_draft(): +def test_unscored_request_maps_to_a_permanently_neutral_row(): + """A request with no slot yet must read as "verify the whole block". + + The scratch row is not safe for this -- padded and unknown requests write + through it -- so there is a dedicated row past every real slot that nothing + ever writes. + """ worker = _make_worker(enable_confidence_scheduling=True) dm, _ = _stride_probe_draft_model() worker._lazy_init(dm, _make_metadata(max_num_requests=8)) - assert worker.staged_confidence_rows() is None + + neutral = worker.confidence_row_for(12345) + assert neutral == worker._neutral_conf_row + # Past every real slot AND past the scratch row. + assert neutral > worker._scratch_slot + buf = worker.staged_confidence_buffer() + assert torch.sigmoid(buf[neutral]).min().item() > 0.999 + + +def test_a_freshly_assigned_slot_starts_neutral(): + """Otherwise a new request's first decision trims on a dead draft's scores. + + Slots are recycled between requests, so whatever the previous occupant's + draft scored is still sitting in the row when the next request moves in. + """ + worker = _make_worker(enable_confidence_scheduling=True) + dm, _ = _stride_probe_draft_model() + worker._lazy_init(dm, _make_metadata(max_num_requests=8)) + + # Every real slot carries a dead previous occupant's hopeless confidence. + worker._confidence_logits[: worker._scratch_slot] = -20.0 + + slot = worker._assign_slot(100, reset=False) + assert slot < worker._scratch_slot + assert torch.sigmoid(worker._confidence_logits[slot]).min().item() > 0.999 def test_confidence_buffer_absent_when_scheduling_is_off(): From e1c28eeae21d463e82ceb18ddb18b611ed36a80c Mon Sep 17 00:00:00 2001 From: Lanyu Liao Date: Sat, 1 Aug 2026 01:27:38 -0700 Subject: [PATCH 003/177] [None][feat] Make DSpark ragged verification work on the DSv4 kernel path PR #17056 wired the confidence head into a verification scheduler, but the kernels underneath it still assume every generation request contributes the same number of query tokens. This makes the per-request path actually executable. Kernels ------- * indexerTopK.cu: topKPerRowDecode rebuilt both the row->request map and the causal extent from a scalar next_n. It now takes an optional per-row extent instead. The generalization is exact -- a request verifying v positions lets row o see kv_len - v + o + 1 tokens, and substituting v = next_n recovers the original arithmetic -- so the uniform path stays bit-identical. * IndexerTopKOp.cpp: the ragged layout REPLACES the seq_lens.size(0) * next_n == numRows check rather than relaxing it. A ragged batch whose padded row count happens to divide evenly (verify_lens [6,4,1,1] -> 12 rows, 4 requests, next_n 3) would otherwise pass the check and silently attribute rows to the wrong requests. * compressorKernels.cu: NEXT_N is a template parameter, but the body only uses it for sp, last_token_idx and two loop bounds -- and both loops are already guarded. Demoting it to a compile-time upper bound and reading the real per-request count from an optional vector keeps a single launch, a fixed grid, and an unchanged CUDA-graph topology. Six ordering and lifetime defects --------------------------------- Each is silent: the run produces plausible output and baseline accuracy. 1. attn_metadata.prepare() is the only consumer of ragged_verify_lens and ran 86 lines before they were published, so every step read the previous step's split and the first ragged step read None. 2. The captured graph sized input_ids from batch_size * (draft_len + 1) while seq_lens stayed frozen at the top tier. Only the widest bucket coincides with that product -- and that is exactly what warmup captures, so it captured clean and diverged on the first trimmed step. 3. The in-graph row-extent refresh allocated a pinned tensor and issued an H2D copy inside the capture body; at replay the memcpy source is a freed temporary. Now a gather plus an add over persistent buffers. 4. on_update_kv_lens recomputed num_gen_tokens_per_seq by integer division on every forward, undoing the ragged batch maximum that the compressor reads as its loop bound. Truncated Phase 1 never writes the tail KV state. 5. The drafter strided into flat-token-indexed captured hidden states by a fixed Kp1, so with windows [6,6,2,2] request 3 drafts from request 2's tail. 6. The overlap rewind read py_verify_len after the scheduler had already stamped the next step's windows onto the same request objects. Items 5 and 6 only trigger with the overlap scheduler enabled, which is the default configuration. Observability ------------- Every failure mode of this feature is silent, so a passing accuracy run does not show that it ever ran. dspark_observability.py adds a static/compact mode switch and DSparkRaggedStats.assert_ragged_active(), which fails unless the windows actually differed, actually trimmed, and stayed in graph replay. TLLM_DSPARK_FORCE_VERIFY_LENS breaks a circular dependency: without a profiled cost table the planner correctly refuses to trim, so validating the packing would first require a table, which requires a working run. The override replaces only the planner's decision, not any downstream path, so a forced split must still agree with static token for token. dspark_sps_profiler.py produces the cost table the planner needs; without one the budget degenerates to verify-all and no performance claim is meaningful. Testing ------- 269 tests pass on B300 (SM103): 10 new ragged top-k kernel tests including CUDA graph capture/replay and a torch.topk reference, plus the DSpark and profiler unit tests. Note the kernel does not promise an output ORDER -- only the top-k set -- so the tests compare sorted rows. Not yet verified ---------------- End-to-end has not run. The six fixes above are located by reading code and carry no runtime evidence. DeepSeek-V4-Pro cold start measured at ~3h on this storage (2h prefetch, ~10min file read, ~40min tensor assembly), which does not fit a 4h window; see docs/dspark_ragged_verify_status.md for the page-cache reuse strategy and the full list of what is and is not established. Signed-off-by: Lanyu Liao --- cpp/tensorrt_llm/kernels/IndexerTopK.h | 14 +- .../compressorKernels/compressorKernels.cu | 21 +- .../compressorKernels/compressorKernels.h | 4 + cpp/tensorrt_llm/kernels/indexerTopK.cu | 64 +- cpp/tensorrt_llm/thop/IndexerTopKOp.cpp | 41 +- cpp/tensorrt_llm/thop/compressorOp.cpp | 21 +- docs/dspark_confidence_schedule_goal.md | 399 +++++ docs/dspark_ragged_verify_status.md | 351 ++++ .../sparse/deepseek_v4/compressor.py | 6 + .../sparse/deepseek_v4/deepseek_v4.py | 68 +- .../_torch/attention_backend/sparse/dsa.py | 194 ++- .../_torch/pyexecutor/cuda_graph_runner.py | 13 +- .../_torch/pyexecutor/model_engine.py | 115 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 45 +- tensorrt_llm/_torch/speculative/dspark.py | 56 +- .../speculative/dspark_observability.py | 360 ++++ .../_torch/speculative/dspark_sps_profiler.py | 1503 +++++++++++++++++ .../_torch/speculative/dspark_verify.py | 20 + .../_torch/speculative/spec_sampler_base.py | 31 +- .../hw_agnostic/test_dspark_sps_profiler.py | 614 +++++++ .../thop/parallel/test_indexer_topk_ragged.py | 294 ++++ 21 files changed, 4172 insertions(+), 62 deletions(-) create mode 100644 docs/dspark_confidence_schedule_goal.md create mode 100644 docs/dspark_ragged_verify_status.md create mode 100644 tensorrt_llm/_torch/speculative/dspark_observability.py create mode 100644 tensorrt_llm/_torch/speculative/dspark_sps_profiler.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py create mode 100644 tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py diff --git a/cpp/tensorrt_llm/kernels/IndexerTopK.h b/cpp/tensorrt_llm/kernels/IndexerTopK.h index 6e6e7b29b059..025c7a19b404 100644 --- a/cpp/tensorrt_llm/kernels/IndexerTopK.h +++ b/cpp/tensorrt_llm/kernels/IndexerTopK.h @@ -40,11 +40,18 @@ int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitW /// - Insertion sort (N < kSortingAlgorithmThreshold) /// - Radix sort (kSortingAlgorithmThreshold ≤ N < splitWork) /// - Radix split-work (N ≥ splitWork — uses outLogitsAux / outIndicesAux) +/// +/// @param rowKvLens optional [numRows] per-row causal KV extent, in uncompressed +/// token space. Non-null selects the ragged layout: rows no longer map to +/// requests by `rowIdx / next_n`, so the extent every row may attend to is +/// supplied directly instead of being reconstructed from `next_n`. Null keeps +/// the uniform arithmetic, bit-identical to before. Mutually exclusive with +/// `preIdx` (GVR's hint is request-indexed through `next_n` as well). void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, float* heuristicScratch = nullptr, int const compressRatio = 1, - cudaStream_t const stream = 0); + int const* rowKvLens = nullptr, cudaStream_t const stream = 0); /// bf16 indexer TopK decode — same dispatch axes as the fp32 entry, except /// kBsL2 uses sizeof(__nv_bfloat16) bytes/elem (L2 footprint is half) and @@ -59,13 +66,14 @@ void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, in int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, __nv_bfloat16* heuristicScratch = nullptr, int const compressRatio = 1, - cudaStream_t const stream = 0); + int const* rowKvLens = nullptr, cudaStream_t const stream = 0); /// fp16 indexer TopK decode — see bf16 overload for dispatcher contract. void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK = 2048, int const* preIdx = nullptr, int const preIdxStride = 0, int const preIdxCount = 0, - __half* heuristicScratch = nullptr, int const compressRatio = 1, cudaStream_t const stream = 0); + __half* heuristicScratch = nullptr, int const compressRatio = 1, int const* rowKvLens = nullptr, + cudaStream_t const stream = 0); void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, int const numRows, int const numColumns, int const stride0, int const stride1, int const topK = 2048, diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu index 1457efeb3138..361df4bacaa2 100644 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.cu @@ -297,7 +297,7 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo void* __restrict__ paged_kv_raw, void* __restrict__ paged_score_raw, int32_t const* __restrict__ block_table_kv, int32_t const* __restrict__ block_table_score, void* __restrict__ output_raw, int32_t const* __restrict__ kv_lens, int32_t const* __restrict__ cu_seq_lens, int32_t const* __restrict__ cu_kv_comp, int page_size, int max_blocks, - int out_elem_bytes) + int out_elem_bytes, int32_t const* __restrict__ new_tokens_per_seq = nullptr) { using KvScoreElemT = typename std::conditional::type; using StateElemT = typename std::conditional::type; @@ -337,7 +337,16 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo int const eff_tid = head_blk * NTHRD_INNER + tid; int const kv_len = kv_lens[batch_idx]; - int const sp = kv_len - NEXT_N; + // How many tokens this request appended this step. Uniform batches append + // NEXT_N each, which is what the template bound encodes. Under DSpark's + // ragged verification each request appends its own window instead, so the + // count is read per request and NEXT_N degrades to a compile-time *upper* + // bound -- which is all the loops below need it to be, because both are + // already guarded (Phase 1 by token_idx < kv_len, Phase 3 by + // num_compressions). Keeping NEXT_N as the bound is what preserves the + // full unroll on the uniform path. + int const nn = (new_tokens_per_seq != nullptr) ? new_tokens_per_seq[batch_idx] : NEXT_N; + int const sp = kv_len - nn; int const in_off = cu_seq_lens[batch_idx]; int const out_off = cu_kv_comp[batch_idx]; int64_t const page_sd = static_cast(page_size) * STATE_DIM; @@ -413,7 +422,7 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo // ================================================================ // Phase 2: Count how many complete compression windows finished. // ================================================================ - int last_token_idx = sp + NEXT_N - 1; + int last_token_idx = sp + nn - 1; int num_compressions = (last_token_idx + 1) / COMPRESS_RATIO - sp / COMPRESS_RATIO; // ================================================================ @@ -663,7 +672,8 @@ __global__ void pagedKvCompressKernel(void const* __restrict__ kv_score_raw, flo // Generate explicit template instantiations. #define INST_DECODE(HD, KV_EB, STATE_EB, CR, NN, NRW) \ template __global__ void pagedKvCompressKernel(void const*, float const*, void*, \ - void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int, int, int); + void*, int32_t const*, int32_t const*, void*, int32_t const*, int32_t const*, int32_t const*, int, int, int, \ + int32_t const*); FOREACH_DECODE_CONFIG(INST_DECODE) #undef INST_DECODE @@ -680,6 +690,7 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k int32_t const* block_table_kv, int32_t const* block_table_score, void* output, int32_t const* kv_lens, int32_t const* cu_seq_lens, int32_t const* cu_kv_comp, int batch_size, int page_size, int max_blocks, int head_dim, int compress_ratio, int next_n, int kv_score_elem_bytes, int state_elem_bytes, int out_elem_bytes, + int32_t const* new_tokens_per_seq, cudaStream_t stream) { TLLM_CHECK_WITH_INFO( @@ -731,7 +742,7 @@ void pagedKvCompressLaunch(void const* kv_score, float const* ape, void* paged_k { \ pagedKvCompressKernel<<>>(kv_score, ape, \ paged_kv, paged_score, block_table_kv, block_table_score, output, kv_lens, cu_seq_lens, cu_kv_comp, \ - page_size, max_blocks, out_elem_bytes); \ + page_size, max_blocks, out_elem_bytes, new_tokens_per_seq); \ return; \ } FOREACH_DECODE_CONFIG(TRY_LAUNCH) diff --git a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h index 0d83628e1b9b..df5d2826b03b 100644 --- a/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h +++ b/cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.h @@ -47,6 +47,10 @@ void pagedKvCompressLaunch(void const* kv_score, // [m, 2*state_dim] (bf16 or f int kv_score_elem_bytes, // bytes per element for kv_score (2=bf16, 4=fp32) int state_elem_bytes, // bytes per element for paged state (2=bf16, 4=fp32) int out_elem_bytes, // bytes per element for output + int32_t const* new_tokens_per_seq, // [bsz] or nullptr; per-request tokens appended this + // step. nullptr means every request appended next_n + // (the uniform case); non-null makes next_n an upper + // bound only, for DSpark ragged verification. cudaStream_t stream); // Prefill kernel: bulk compression with per-token gather/scatter + state update. diff --git a/cpp/tensorrt_llm/kernels/indexerTopK.cu b/cpp/tensorrt_llm/kernels/indexerTopK.cu index d63f203a048c..12c77af6664c 100644 --- a/cpp/tensorrt_llm/kernels/indexerTopK.cu +++ b/cpp/tensorrt_llm/kernels/indexerTopK.cu @@ -654,7 +654,8 @@ template static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(InputT const* logits, int const* seqLens, int* outIndices, int stride0, int stride1, int const topK, int next_n, int compressRatio, - float* outLogits = nullptr, int const numBlocksToMerge = 0, int const* indices = nullptr) + float* outLogits = nullptr, int const numBlocksToMerge = 0, int const* indices = nullptr, + int const* rowKvLens = nullptr) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); @@ -664,8 +665,27 @@ static __global__ __launch_bounds__(kNumThreadsPerBlock) void topKPerRowDecode(I // The range of logits within the row. int rowStart = 0; - int seq_len = seqLens[rowIdx / next_n]; - int actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; + // Two ways to learn how far back this query row may attend. + // + // Uniform (rowKvLens == nullptr): every request contributes exactly next_n + // rows, so the request index and the row's offset inside its verify window + // both fall out of integer division, and the causal extent is the request's + // sequence length walked back to this offset. + // + // Ragged (rowKvLens != nullptr): requests contribute different numbers of + // rows, so neither rowIdx / next_n nor rowIdx % next_n means anything. The + // caller precomputes the extent per row -- it is the only quantity this + // kernel ever derived from next_n, which is why one array is enough. + int actual_kv_len; + if (rowKvLens != nullptr) + { + actual_kv_len = rowKvLens[rowIdx]; + } + else + { + int seq_len = seqLens[rowIdx / next_n]; + actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; + } int rowEnd = actual_kv_len / compressRatio; // Local pointers to this block @@ -890,9 +910,16 @@ int computeIndexerTopKDecodeBlocksPerRow(int numRows, int numColumns, int splitW void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indices, float* outLogitsAux, int* outIndicesAux, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, - int const preIdxCount, float* heuristicScratch, int const compressRatio, cudaStream_t const stream) + int const preIdxCount, float* heuristicScratch, int const compressRatio, int const* rowKvLens, + cudaStream_t const stream) { constexpr int kNumThreadsPerBlock = 512; + // GVR consumes previous-step Top-K hints indexed by request, which it also + // recovers from next_n. Ragged batches therefore cannot take that path; the + // caller is expected to have disabled the hint, and this makes the + // requirement explicit rather than letting a stale preIdx through. + TLLM_CHECK_WITH_INFO(rowKvLens == nullptr || preIdx == nullptr, + "indexer TopK decode: ragged rowKvLens is incompatible with the GVR preIdx hint"); int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; // ======================================================================== @@ -1016,7 +1043,7 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic config.attrs = attrs; cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - compressRatio, nullptr, 0, nullptr); + compressRatio, nullptr, 0, nullptr, rowKvLens); } else { @@ -1031,7 +1058,7 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic config_part1.attrs = attrs; cudaLaunchKernelEx(&config_part1, kernel_instance_part1, logits, seqLens, outIndicesAux, stride0, stride1, topK, - next_n, compressRatio, outLogitsAux, 0, nullptr); + next_n, compressRatio, outLogitsAux, 0, nullptr, rowKvLens); constexpr int kNumThreadsPerBlockMerge = 1024; auto* kernel_instance_part2 = &topKPerRowDecode; @@ -1043,8 +1070,12 @@ void invokeIndexerTopKDecode(float const* logits, int const* seqLens, int* indic config_part2.numAttrs = 1; config_part2.attrs = attrs; + // The merge pass overrides rowEnd with numBlocksToMerge * topK, so the + // extent is unused -- but rowKvLens still has to be forwarded, because + // otherwise the uniform branch would index seqLens by rowIdx / next_n, + // and on a ragged batch next_n is not a divisor of the row count. cudaLaunchKernelEx(&config_part2, kernel_instance_part2, outLogitsAux, seqLens, indices, blocksPerRow * topK, 1, - topK, next_n, 1, nullptr, blocksPerRow, outIndicesAux); + topK, next_n, 1, nullptr, blocksPerRow, outIndicesAux, rowKvLens); } sync_check_cuda_error(stream); } @@ -1079,10 +1110,14 @@ template void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, int const preIdxCount, InputT* heuristicScratch, int const compressRatio, - cudaStream_t const stream) + int const* rowKvLens, cudaStream_t const stream) { static_assert(std::is_same_v || std::is_same_v, "invokeIndexerTopKDecodeDtype is for bf16/fp16 only"); + // See the fp32 entry: GVR's hint is request-indexed via next_n, so it and a + // ragged row map cannot both be live. + TLLM_CHECK_WITH_INFO(rowKvLens == nullptr || preIdx == nullptr, + "indexer TopK decode: ragged rowKvLens is incompatible with the GVR preIdx hint"); constexpr int kNumThreadsPerBlock = 512; int const effectiveSplitWorkThreshold = splitWorkThreshold > 0 ? splitWorkThreshold : kDefaultSplitWorkThreshold; @@ -1123,7 +1158,7 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* config.attrs = attrs; cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - compressRatio, nullptr, 0, nullptr); + compressRatio, nullptr, 0, nullptr, rowKvLens); } else if (numColumns < effectiveSplitWorkThreshold) { @@ -1143,7 +1178,7 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* config.attrs = attrs; cudaLaunchKernelEx(&config, kernel_instance, logits, seqLens, indices, stride0, stride1, topK, next_n, - compressRatio, nullptr, 0, nullptr); + compressRatio, nullptr, 0, nullptr, rowKvLens); } else { @@ -1162,19 +1197,20 @@ void invokeIndexerTopKDecodeDtype(InputT const* logits, int const* seqLens, int* void invokeIndexerTopKDecode(__nv_bfloat16 const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, int const preIdxCount, - __nv_bfloat16* heuristicScratch, int const compressRatio, cudaStream_t const stream) + __nv_bfloat16* heuristicScratch, int const compressRatio, int const* rowKvLens, cudaStream_t const stream) { invokeIndexerTopKDecodeDtype<__nv_bfloat16>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, - stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); + stride0, stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, rowKvLens, + stream); } void invokeIndexerTopKDecode(__half const* logits, int const* seqLens, int* indices, int const splitWorkThreshold, int const numRows, int const numColumns, int const stride0, int const stride1, int const next_n, int const topK, int const* preIdx, int const preIdxStride, int const preIdxCount, __half* heuristicScratch, int const compressRatio, - cudaStream_t const stream) + int const* rowKvLens, cudaStream_t const stream) { invokeIndexerTopKDecodeDtype<__half>(logits, seqLens, indices, splitWorkThreshold, numRows, numColumns, stride0, - stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, stream); + stride1, next_n, topK, preIdx, preIdxStride, preIdxCount, heuristicScratch, compressRatio, rowKvLens, stream); } void invokeIndexerTopKPrefill(float const* logits, int const* rowStarts, int const* rowEnds, int* indices, diff --git a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp index fc29041bbf93..60ca03a75e05 100644 --- a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp +++ b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp @@ -38,7 +38,8 @@ namespace torch_ext void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, th::Tensor const& indices, int64_t next_n, int64_t index_topk, std::optional const& pre_idx, std::optional const& heuristic_scratch, int64_t compress_ratio, - std::optional const& radix_aux_indices, std::optional const& radix_aux_logits) + std::optional const& radix_aux_indices, std::optional const& radix_aux_logits, + std::optional const& row_kv_lens) { TORCH_CHECK(compress_ratio > 0, "compress_ratio must be greater than 0"); @@ -53,8 +54,17 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t auto const inputSize = logits.sizes(); auto const numRows64 = inputSize[0]; auto const numColumns64 = inputSize[1]; - TORCH_CHECK( - seq_lens.size(0) * next_n == numRows64, "seq_lens length multiplied by next_n must equal logits.size(0)"); + // Uniform batches must tile exactly: every request owns next_n rows. A + // ragged batch cannot satisfy that by construction, and -- crucially -- can + // still satisfy it *by accident* whenever the padded row count happens to be + // divisible by the row count, at which point the kernel would silently + // attribute rows to the wrong requests. So the ragged layout replaces the + // check rather than loosening it. + if (!row_kv_lens.has_value()) + { + TORCH_CHECK( + seq_lens.size(0) * next_n == numRows64, "seq_lens length multiplied by next_n must equal logits.size(0)"); + } TORCH_CHECK(indices.size(0) == numRows64, "indices first dimension must match logits.size(0)"); TORCH_CHECK(indices.size(1) >= index_topk, "indices second dimension must be at least index_topk"); TORCH_CHECK(seq_lens.is_contiguous(), "seq_lens must be contiguous"); @@ -70,6 +80,22 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t TORCH_CHECK(logits_stride_0 >= 0, "logits_stride_0 must be greater than or equal to 0"); TORCH_CHECK(logits_stride_1 >= 0, "logits_stride_1 must be greater than or equal to 0"); + int32_t const* rowKvLensPtr = nullptr; + if (row_kv_lens.has_value()) + { + auto const& rowKvLensTensor = row_kv_lens.value(); + TORCH_CHECK(rowKvLensTensor.is_cuda(), "row_kv_lens must be a CUDA tensor"); + TORCH_CHECK( + rowKvLensTensor.device() == logits.device(), "row_kv_lens must be on the same device as logits"); + TORCH_CHECK(rowKvLensTensor.is_contiguous(), "row_kv_lens must be contiguous"); + TORCH_CHECK(rowKvLensTensor.dim() == 1, "row_kv_lens must be a 1D Tensor"); + TORCH_CHECK(rowKvLensTensor.scalar_type() == at::ScalarType::Int, "row_kv_lens must be int32"); + TORCH_CHECK(rowKvLensTensor.size(0) == numRows64, + "row_kv_lens must have one entry per logits row (got ", rowKvLensTensor.size(0), " for ", numRows64, + " rows)"); + rowKvLensPtr = rowKvLensTensor.data_ptr(); + } + int32_t const* preIdxPtr = nullptr; int32_t preIdxStride = 0; int32_t preIdxCount = 0; @@ -162,7 +188,8 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t tk::invokeIndexerTopKDecode(logits.data_ptr(), seq_lens.data_ptr(), indices.data_ptr(), aux_logits_ptr, aux_indices_ptr, splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, preIdxStride, - preIdxCount, static_cast(heuristicScratchPtr), static_cast(compress_ratio), stream); + preIdxCount, static_cast(heuristicScratchPtr), static_cast(compress_ratio), rowKvLensPtr, + stream); } else if (logits_dtype == at::ScalarType::BFloat16) { @@ -170,14 +197,14 @@ void indexer_topk_decode(th::Tensor const& logits, th::Tensor const& seq_lens, t seq_lens.data_ptr(), indices.data_ptr(), splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, preIdxStride, preIdxCount, static_cast<__nv_bfloat16*>(heuristicScratchPtr), - static_cast(compress_ratio), stream); + static_cast(compress_ratio), rowKvLensPtr, stream); } else // Half { tk::invokeIndexerTopKDecode(reinterpret_cast<__half const*>(logits.data_ptr()), seq_lens.data_ptr(), indices.data_ptr(), splitWorkThreshold, num_rows, num_columns, logits_stride_0, logits_stride_1, static_cast(next_n), static_cast(index_topk), preIdxPtr, preIdxStride, preIdxCount, - static_cast<__half*>(heuristicScratchPtr), static_cast(compress_ratio), stream); + static_cast<__half*>(heuristicScratchPtr), static_cast(compress_ratio), rowKvLensPtr, stream); } } @@ -228,7 +255,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def( "indexer_topk_decode(Tensor logits, Tensor seq_lens, Tensor indices, int next_n, int index_topk=2048, " "Tensor? pre_idx=None, Tensor? heuristic_scratch=None, int compress_ratio=1, " - "Tensor? radix_aux_indices=None, Tensor? radix_aux_logits=None) -> ()"); + "Tensor? radix_aux_indices=None, Tensor? radix_aux_logits=None, Tensor? row_kv_lens=None) -> ()"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/cpp/tensorrt_llm/thop/compressorOp.cpp b/cpp/tensorrt_llm/thop/compressorOp.cpp index 5fdfb5a85806..2b56924df0be 100644 --- a/cpp/tensorrt_llm/thop/compressorOp.cpp +++ b/cpp/tensorrt_llm/thop/compressorOp.cpp @@ -36,12 +36,27 @@ void compressorPagedKvCompressOp(torch::Tensor kv_score, // [m, 2*state_dim] bf1 torch::Tensor kv_lens, // [bsz] int32 torch::Tensor cu_seq_lens, // [bsz+1] int32 torch::Tensor cu_kv_comp, // [bsz+1] int32 - int64_t batch_size, int64_t page_size, int64_t head_dim, int64_t compress_ratio, int64_t next_n) + int64_t batch_size, int64_t page_size, int64_t head_dim, int64_t compress_ratio, int64_t next_n, + std::optional const& new_tokens_per_seq) { constexpr int64_t kMinNextN = 1; constexpr int64_t kMaxNextN = 8; TORCH_CHECK(next_n >= kMinNextN && next_n <= kMaxNextN, "next_n must be in [1, 8], got ", next_n); + // Ragged batches pass the batch maximum as next_n and the real per-request + // counts here; the kernel then treats next_n purely as a loop bound. + int32_t const* new_tokens_ptr = nullptr; + if (new_tokens_per_seq.has_value()) + { + auto const& t = new_tokens_per_seq.value(); + TORCH_CHECK(t.is_cuda(), "new_tokens_per_seq must be a CUDA tensor"); + TORCH_CHECK(t.is_contiguous(), "new_tokens_per_seq must be contiguous"); + TORCH_CHECK(t.scalar_type() == torch::kInt32, "new_tokens_per_seq must be int32"); + TORCH_CHECK(t.dim() == 1 && t.size(0) >= batch_size, + "new_tokens_per_seq must be 1D with at least batch_size entries"); + new_tokens_ptr = t.data_ptr(); + } + auto stream = at::cuda::getCurrentCUDAStream(); int kv_score_eb = static_cast(kv_score.element_size()); int state_eb = static_cast(paged_kv.element_size()); @@ -54,7 +69,7 @@ void compressorPagedKvCompressOp(torch::Tensor kv_score, // [m, 2*state_dim] bf1 kv_lens.data_ptr(), cu_seq_lens.data_ptr(), cu_kv_comp.data_ptr(), static_cast(batch_size), static_cast(page_size), static_cast(block_table_kv.size(1)), static_cast(head_dim), static_cast(compress_ratio), static_cast(next_n), kv_score_eb, state_eb, - out_eb, stream); + out_eb, new_tokens_ptr, stream); } // Prefill kernel: bulk compression with state update @@ -135,7 +150,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) "Tensor cu_seq_lens, Tensor cu_kv_comp, " "int batch_size, int page_size, " "int head_dim, int compress_ratio, " - "int next_n) -> ()"); + "int next_n, Tensor? new_tokens_per_seq=None) -> ()"); m.def( "compressor_prefill_reduction(" diff --git a/docs/dspark_confidence_schedule_goal.md b/docs/dspark_confidence_schedule_goal.md new file mode 100644 index 000000000000..2364729e405a --- /dev/null +++ b/docs/dspark_confidence_schedule_goal.md @@ -0,0 +1,399 @@ +# DSpark Confidence-Head 调度在 DeepSeek-V4-Pro 上的落地：目标文档 + +> 状态：scoping / 目标定义，**尚未定稿实现方案** +> 关联：PR NVIDIA/TensorRT-LLM#17056，分支 `pr-17056` +> 目标模型：DeepSeek-V4-Pro-DSpark（target = DSv4 Pro + DSA 稀疏注意力；drafter = DSpark block drafter + confidence head） +> `max_draft_len=5` → `block_size K=5` → 均匀 verify 窗口 `next_n = 6` +> 硬件：8×Blackwell（B200/B300，SM100），DEP8 = TP8 + `moe_expert_parallel_size=8` + `enable_attention_dp=True` + +本文档中所有断言均标注了 `file:line`。凡是**没有**通过读代码确认的推断，一律显式标注「待验证」。 + +--- + +## 1. 背景与现状 + +### 1.1 PR #17056 已经做了什么 + +两个 commit： + +| commit | 内容 | +|---|---| +| `830afe4e58` | `[feat] Add DSpark confidence-head verification scheduling` | +| `8723506608` | `[fix] Fix host syncs, slot aliasing and cost model` | + +核心数据流（已读码确认）： + +1. **Confidence head 产出**：`DSparkConfidenceHead.forward` 返回 `[G, K]` fp32 raw logits，按 batch 位置排列（`tensorrt_llm/_torch/models/dspark/heads.py:258-268`），在 draft 的 captured graph 内计算（`tensorrt_llm/_torch/models/dspark/draft.py:108-121`）。 +2. **batch → slot 重索引**：`DSparkWorker` 立刻把它写进按 slot 索引的常驻 device buffer `self._confidence_logits[slots] = confidence`（`tensorrt_llm/_torch/speculative/dspark.py:637-638`），buffer 形状 `[max_batch+2, K]`，其中 `max_batch+1` 行是永不写入的 neutral 行（`dspark.py:285-300`）。 +3. **滞后一步的 host 读取**：`PyExecutor._dspark_confidence_draft_len` 先读上一步 stage 好的 async D2H 快照，再 stage 新的一份，**全程不 sync**（`tensorrt_llm/_torch/pyexecutor/py_executor.py:3189`；staging 见 `tensorrt_llm/_torch/speculative/dspark_verify.py:99-134`）。event 没落地就退化成 `max_tier`（verify 全量）。 +4. **语义**：`py_verify_len` 计的是**每请求的 draft 位置数**，token 窗口是 `1 + py_verify_len`（`tensorrt_llm/_torch/pyexecutor/llm_request.py:1283-1286`）。 + +两个配置开关（`tensorrt_llm/llmapi/llm_args.py:2805` / `:2849`，默认均为 `False`，`status='prototype'`）： + +| 开关 | 行为 | +|---|---| +| `enable_confidence_scheduling` | **均匀路径**：`planner.decide_draft_len` 给整个 batch 选一个 tier，把每个请求的 `py_verify_len` 显式清成 `None`（`py_executor.py:3283`），复用既有的 dynamic-draft-len pad/truncate 机制。CUDA graph 变成 `batch_size × tiers` 的叉积（`model_engine.py:1828-1856`）。 | +| `enable_ragged_verify` | **ragged 路径**：`decide_verify_lens` 给出每请求窗口，`fit_ragged_verify_lens` 把 batch token 总数吸附到某个 captured bucket，`runtime_draft_len` 钉死在**最高 tier**（因为 block 永远整块 draft），graph key 增加一条 token-count 轴。依赖 `enable_confidence_scheduling`（validator 见 `llm_args.py:2863-2871`）。 | + +**默认路径（两个都 False）完全不受影响**，这一点已确认：`verify_len_tiers` 退化成 `[max_draft_len]`（`llm_args.py:2899-2905`），`py_verify_len` 从不设置，`expand_per_gen_token` 走标量 `repeat_interleave` 分支（`dsa.py:693-697`），与 PR 之前的定长扩展 bit-identical。 + +### 1.2 已经做对的部分 + +PR 的 host 侧管线质量是高的，不要推翻重做： + +- `ragged_verify_lens_cuda` / `ragged_qo_indptr_cuda` 一次性预分配、原地写入（`model_engine.py:715-720`、`:3802-3813`），地址稳定，captured graph 能看到新值。 +- 每个 `repeat_interleave` 都传了 `output_size=`，避免 device→host sync（`dspark_ragged.py:411-429`、`dsa.py:693-706`）。 +- staging 用 pinned memory + `non_blocking=True`。 +- `choose_ragged_capture_shape` 对 `peer_stats` 做跨 rank 归约，保证各 rank 落在同一个 `(padded_bs, bucket)`（`dspark_ragged.py:359-382`）。 +- raggedness 活在**内容**里、shape 由 graph key 固定 —— 这是对的构造。 + +### 1.3 工作区当前状态（**与 investigation 描述不同，必须先对齐**） + +`git status` 显示 **11 个已跟踪文件被修改**、2 个新增文件，均**未提交**。这部分工作已经超出了 PR HEAD： + +``` + M cpp/tensorrt_llm/kernels/IndexerTopK.h + M cpp/tensorrt_llm/kernels/indexerTopK.cu + M cpp/tensorrt_llm/thop/IndexerTopKOp.cpp + M cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.{cu,h} + M cpp/tensorrt_llm/thop/compressorOp.cpp + M tensorrt_llm/_torch/attention_backend/sparse/dsa.py + M tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/{compressor.py,deepseek_v4.py} + M tensorrt_llm/_torch/pyexecutor/py_executor.py + M tensorrt_llm/_torch/speculative/dspark.py +?? tensorrt_llm/_torch/speculative/dspark_observability.py +?? tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py +``` + +已完成的内容（读码确认）： + +| 项 | 状态 | 锚点 | +|---|---|---| +| `indexer_topk_decode` 新增 `Tensor? row_kv_lens=None` | **已实现** | `IndexerTopKOp.cpp:258`（schema）、`indexerTopK.cu:665-690`（kernel 分支）、`IndexerTopKOp.cpp:57-66`（ragged 时**替换**而非放宽 TORCH_CHECK）、`indexerTopK.cu:918-922`（与 GVR `preIdx` 互斥的 TLLM_CHECK） | +| Python 侧路由 ragged → C++ top-k，`next_n` 传死值 1 | **已实现** | `dsa.py:3056-3081` | +| `compressor_paged_kv_compress` 新增 `Tensor? new_tokens_per_seq=None` | **已实现** | `compressorOp.cpp:153`（schema）、`compressorKernels.cu:340-348`（`nn` 取代 `NEXT_N`，`NEXT_N` 降级为编译期上界）、`:425`（Phase 2）；Phase 1 由 `token_idx < kv_len` 守卫（`:369`）、Phase 3 由 `num_compressions` 守卫（`:440-442`），二者对 `nn < NEXT_N` 均正确 | +| `num_gen_tokens_per_seq` ragged 时取 **max** 而非 mean | **已实现** | `deepseek_v4.py:748-769` | +| `gen_new_tokens_per_seq` 每请求向量 + graph-stable buffer | **已实现** | `deepseek_v4.py:337-352` | +| ragged verify mode 开关（static / cap-accept / compact）+ 统计 | **部分实现**，`cap-accept` 显式 `NotImplementedError` | `dspark.py:457-496`、`dspark_observability.py` | +| C++ ragged top-k 单测 | **已新增** | `tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py` | + +**重要副作用**：`num_gen_tokens_per_seq` 改成取 max 之后，compressor 输出 buffer 的 sizing 公式 `num_generations * ceil(w/cr)`（`deepseek_v4.py:753-757`、`:766-768`）重新成为**合法上界**（因为 `w_r ≤ w_max` ⇒ `Σ_r ceil(w_r/cr) ≤ n·ceil(w_max/cr)`）。这消除了「越界写 `kv_comp`」这一类内存破坏风险。取 mean 时该公式**不是**上界，是真实的 OOB 写。 + +### 1.4 已知的**未验证**部分 + +**唯一一次端到端尝试是失败的。** `.dspark-logs/gsm8k.log`（07/30 23:08，早于当前工作区改动）： + +``` +DSpark confidence scheduling: capturing 102 graphs (34 batch sizes x 3 draft-length tiers [1, 3, 5]) +Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=5 ← 完成 +Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=3 ← never returns +TEST_EXIT=124 +``` + +`hang-stack-299685.txt` 停在 `thop.attention`（`fmha/fallback.py:181`）→ `trtllm.py:1775` → `deepseek_v4.py:1564` → `mla.py:1132` → `cuda_graph_runner.py:499 capture`。 + +注意 hang 前紧邻的日志是 `Attention workspace size is not enough, increase the size from 0 bytes to 168493568 bytes` —— **capture 期间的 workspace 重分配**是一个独立的候选根因，必须与 stride 假说区分开（见 §7）。 + +同时，**这次 run 其实什么都没验证到**：配置里 `enable_ragged_verify` 未开，且没给 `confidence_sps_table_path`，planner 拿到 `SpsCostTable.flat()` 后 `_decide_local` 无条件返回 `max_tier`（`dspark_verify.py:268-270`），日志确认 `profiled_cost_table=False`。也就是说该 run 的行为与 baseline 的唯一差别是：多一次 confidence matmul、多一次 D2H staging、以及 3× 的 captured graph。 + +--- + +## 2. 目标 / 非目标 + +### 2.1 目标（可测） + +| # | 目标 | 验收方式 | +|---|---|---| +| **G1** | DSv4-Pro 上 confidence-head 调度可用：`enable_confidence_scheduling=True` 能完成 warmup + capture 且不 hang | 8×B200 上 GSM8K 测试跑到 evaluation 阶段并返回分数 | +| **G2** | 端到端正确性：调度**不改变输出分布** | temperature=0 下，`static` 与开启调度后逐 token 输出**完全一致**（比 GSM8K 分数敏感度高 3 个数量级） | +| **G3** | GSM8K 精度达标 | `score >= 96.0`（`tests/.../references/gsm8k.yaml:158-163` 记录的实测值为 96.475） | +| **G4** | **CUDA graph 兼容** | 稳态 decode 步骤 100% 命中 graph replay；capture 期间无 device→host sync、无新分配；graph key 与实际 token layout 逐字段一致（断言化，见 §6） | +| **G5** | ADP/TP8 下各 rank 形状与集合通信一致 | 各 rank 的 `(padded_bs, bucket)` 与 collectives 调用次数逐步相同；无 hang | +| **G6** | 默认路径零回归 | 两个开关关闭时与 `main` bit-identical（现已成立，需保持） | + +### 2.2 隐含目标（工作本身要求，但用户未明说） + +- **G7**：为「feature 到底有没有在工作」提供可观测性。当前 flat cost table 下 planner 拒绝 trim，run 看起来完全正常但收益为零 —— 必须能从日志/统计区分「调度生效」与「调度退化」。`dspark_observability.py` 是这个方向，需要落地。 +- **G8**：产出两个**目前在 checkpoint 和 repo 里都不存在**的 JSON side-file：STS 校准表、profiled SPS cost table。没有它们，G2/G3 之外的任何**性能**结论都是空的。 + +### 2.3 非目标（明确排除） + +| 排除项 | 理由 | +|---|---| +| 改变 acceptance 规则本身 | 调度只决定「送多少 draft token 给 target」，verify 逻辑不动。这是 G2 可以用逐 token 相等来验证的前提。 | +| Disaggregated serving | 生成侧首步 `py_batch_idx is None` 会命中一条已知有问题的分支（§3.3 K1），单独处理。 | +| 其他模型（DSv4-Flash、非 DSA 模型） | DSv4-Flash-DSpark 权重虽在 `llm_models_root()` 下，但 repo 里零测试；作为后续的小规模 proxy，不在本轮范围。 | +| Pipeline parallelism / LoRA | 与 vLLM #47808 声明的限制一致，本轮不支持。 | +| `enable_heuristic_topk` / `use_cute_dsl_topk` 与 ragged 共存 | GVR 的 `preIdx` 按请求索引且经 `next_n` 展开，CuTe-DSL 把 `next_n` 作为 JIT compile key。二者与 ragged 结构性互斥，**必须显式 config-reject 而非静默降级**。工作区已在 kernel 层加了 TLLM_CHECK（`indexerTopK.cu:918-922`），Python 侧的 config 校验还欠。 | +| 非贪婪采样下的 rejection sampling 等价性 | ragged 下 fail-closed guard 会关掉 rejection sampling（§3.4）。本轮先保证贪婪路径正确，非贪婪单列。 | +| 输出 logprobs | 与 vLLM 同样的限制。 | + +--- + +## 3. 核心问题：kernel 层的 per-request verify len 不兼容 + +### 3.1 问题的本质 + +一旦每个请求有自己的 verify 窗口，**所有从标量 `next_n` 反推 `(request, offset)` 的 kernel 都会失效**。典型形式是： + +``` +request = rowIdx / next_n +offset = rowIdx % next_n +causal = seqLens[request] - next_n + offset + 1 +``` + +而 `next_n` 的来源是 `num_gen_tokens // num_generations`（`dsa.py:2800`、`deepseek_v4.py:748-769`）—— ragged 下这是一个**整数平均值**，没有任何语义。 + +**最危险的一点：失败是静默的，不是报错的。** `ModelEngine.ragged_verify_token_buckets` 返回 `padded_bs * (t+1)`（`model_engine.py:3625`），token 总数**永远**是行数的整数倍，所以形如 `seq_lens.size(0) * next_n == numRows` 的 TORCH_CHECK 会**通过**，然后 kernel 把行归错请求、按错误深度做因果裁剪。这正是 vLLM 在 #47808 之前踩过的坑（他们在 `mla/indexer.py` 里留了明确注释说 `q.view(batch, next_n, ...)` 对非均匀 decode 长度会静默算错）。 + +DSpark 论文 §5.3 对 DeepSeek-V4 的结论与此一致：变长路由「在物理执行层引入了严峻挑战」，解法是把所有 token 展平、用一个 marker tensor 承载序列内结构，并且**「在 DeepSeek-V4 架构上，只有 index-attention 和 compress kernel 需要修改」**。这句话既是我们工作范围的上界，也是「主 sparse-MLA 路径不用动」的外部佐证。 + +### 3.2 DSv4-Pro DSpark 生成路径 kernel 逐项审计 + +按前向顺序。「ragged 容忍」列：✅ 天然支持 / ⚠️ 需 host 侧改动 / ❌ kernel 层阻断 / 🔧 工作区已修。 + +| # | Kernel / Op | Query 轴形状契约 | Ragged? | 阻断点与锚点 | +|---|---|---|---|---| +| **K1** | `mla_rope_generation` → `applyMLARopeAndAssignQKVKernelGeneration` | `TLLM_CHECK(acc_q_len % batch_size == 0)` 后 `seq_len = acc_q_len / batch_size`，**单一标量** | ❌ | `mlaKernels.cu:1161-1163`。派生出 `batch_idx = tok/seq_len`、RoPE `position_id = kv_len[b] - seq_len + local`、**KV 写入偏移**（两处）、`seqQOffset`。窗口 `[6,3]` → `9%2≠0` 直接 abort；`[6,2]` → 8 tokens、`seq_len=4`、check 通过但**每个请求的 K/V 写进错误的行**。这是持久性缓存污染。 | +| **K2** | trtllm-gen sparse MLA generation（`thop.attention` → `AttentionOp::mlaGeneration`） | `mMaxSeqLenQ = acc_q_len / batch_beam`；`mSumOfSeqLensQ = batch_beam * mMaxSeqLenQ` | ❌ | `attentionOp.cpp:1161,1167`。`cumSeqLensQPtr` 字段存在（`fmhaRunnerParams.h:279`）但**在 generation 分支从不赋值**（全 repo 仅 `fmhaDispatcher.cpp:213` 上下文相、`xqaDispatcher.cpp:475` 赋值）。上游 `thop/attentionOp.cpp:920-922` 还有 `TLLM_CHECK(num_tokens % num_seqs == 0, "seq_len should be same for all generation requests")`。**这是最深的阻断点。** | +| **K3** | `fp8_paged_mqa_logits` / `fp8_fp4_paged_mqa_logits`（DeepGEMM） | expanded 路径下 `q.view(-1,1,...)`、`context_lens = kv_lens_expanded[:n].view(-1,1)`、`block_table_expanded` ⇒ `next_n == 1`，**每 token 一行** | ✅ | `dsa.py:2831-2838`。形状上**已经是 ragged 布局**。关键事实：`max_draft_len=5 → next_n=6`，DeepGEMM 在 SM100 只支持 `next_n ∈ {1,2,4}`，所以 **DSv4-Pro 今天本来就跑在这条展平路径上** —— ragged 不是新 regime，它就是 regime。但注意：该 kernel 输出**未裁剪的全长 logits**（`clean_logits` 默认 false），窗口内因果裁剪**全部委托给 top-k**。 | +| **K4** | `indexer_topk_decode`（C++） | 原：`seqLens[rowIdx/next_n]`、`seq_len - next_n + rowIdx%next_n + 1` | 🔧 **工作区已修** | 新增 `Tensor? row_kv_lens`（`IndexerTopKOp.cpp:258`），kernel 分支见 `indexerTopK.cu:665-690`。ragged 时**替换**而非放宽 TORCH_CHECK（`IndexerTopKOp.cpp:57-66`）—— 这个选择是对的，因为放宽后仍会被 `padded_bs*(t+1)` 的整除性静默骗过。split-work 的 merge pass 也正确转发了指针（`indexerTopK.cu:1073-1076`，merge 阶段 `rowEnd` 被覆盖但仍需转发以避免 uniform 分支越界索引 `seqLens`）。 | +| **K5** | GVR / heuristic top-k（C++ `heuristicTopKDecode.cu` + CuTe-DSL `gvr_topk_decode.py`） | `preIdx + (rowIdx/next_n)*stride`；DSL 里 `next_n` 是 `const_expr` 且是 JIT cache key | ❌ **须 reject** | `heuristicTopKDecode.cu:57-65,103`；`IndexerTopKOp.cpp:83-84` 断言 `preIdx.size(0)*next_n == numRows`；`gvr_topk_decode.py:3733-3734,4954`。`order_row` 是**按请求**的 argsort（`dsa.py:1010-1017`），假定每请求占固定连续行数。默认关闭，但是调优部署会打开的旋钮。工作区已在 kernel 层加互斥 check。 | +| **K6** | CuTe-DSL radix / filtered top-k | `next_n` 是 `cutlass.const_expr` + compile key；每行长度 `seqlen[row//next_n] - next_n + row%next_n + 1` | ❌ **已绕开** | `single_pass_multi_cta_radix_topk.py:1062,1124-1125`。DSv4 默认不可达（`dsa.py:867` 在 `compress_ratio>1 且 next_n>1` 时早退，DSv4 的 indexer cr=4）。工作区把 ragged 无条件路由到 C++ kernel（`dsa.py:3056-3081`），进一步保证不可达。 | +| **K7** | CuTe-DSL paged-MQA atom-split `_pick_dsl_expand` | `expand_factor * atom == next_n`，`atom` 是 kernel constexpr | ❌ **结构性不可表达，已 force-off** | `dsa.py:211-263`、ragged 时 `use_dsl=False`（`dsa.py:1566-1574`）。**DSv4-Pro 上性能代价为零**：`use_cute_dsl_paged_mqa_logits` 默认 False（`llm_args.py:888-892`），默认本来就走 DeepGEMM + expanded。 | +| **K8** | `Indexer.prepare_scheduler_metadata`（DeepGEMM schedule） | 硬编码 `num_tokens = num_generations * (1 + max_draft_tokens)` | ⚠️ | `dsa.py:2329-2330`。另两处展开点已改走 `expand_per_gen_token`，唯独这里没改。**这也是均匀 tier 阶梯的 bug**（见 §3.3 H1）。 | +| **K9** | `compressor_paged_kv_compress` | 原：`NEXT_N` 是 C++ **模板常量**，`sp = kv_len - NEXT_N` | 🔧 **工作区已修** | 新增 `Tensor? new_tokens_per_seq`（`compressorOp.cpp:153`）。`nn` 取代 `NEXT_N` 做起点与 Phase-2 计数（`compressorKernels.cu:348,425`），`NEXT_N` 降级为编译期**上界**以保住 uniform 路径的全展开。Phase 1 由 `token_idx < kv_len` 守卫、Phase 3 由 `num_compressions` 守卫，二者对 `nn < NEXT_N` 均正确。`in_off`/`out_off` 本来就是精确前缀和，无需改。 | +| **K10** | DSv4 compressed-token 计数 / `kv_comp` 分配 | `num_generations * ceil(num_gen_tokens_per_seq / cr)` | 🔧 **工作区已修** | `deepseek_v4.py:748-769`。改取 **max** 后该式重新成为合法上界，消除 OOB 写。`_compute_gen_compressed_position_ids` 用 `searchsorted(cu_new_comp, ...)` 回推请求（`deepseek_v4.py:982-990`），对任何**过分配**都仍然正确。 | +| **K11** | `sparse_mla_topk_lens` / `swa_local_indices` / `compressed_local_indices` | 严格逐 token，从 `token_positions` 构造 | ✅ | `deepseek_v4.py:557-648`。**且是 overlap-safe 的**：`on_update_kv_lens` 会用修正后的 device `kv_lens` **重新推导** `token_positions` 并重跑 `prepare_for_deepseek_v4_indices`（`deepseek_v4.py:885-895`）。这条是「token-major 下丢掉 `mIsCausalSpecDecodingGen` 是安全的」的**唯一承重论据**，它是可验证的，不是假设。 | +| **K12** | DSA indexer K-cache slot mapping、`req_idx_per_token`、`convert_req_index_to_global`、`mla_rope_append_paged_kv_assign_q` | 逐 token / indptr 驱动 | ✅ | `dsa.py:1757-1769`、`:2148-2178`、`:3287-3335` | +| **K13** | `spec_decoding_generation_lengths` / packed mask（dense XQA 路径） | C++ 侧**真支持**每请求向量（`gptKernels.cu:140`），但**所有 Python 写入点都是标量 `fill_`** | ⚠️ 但 DSv4 不可达 | `trtllm.py:1189-1191`。SM100 + linear tree 下 `is_spec_decoding_enabled` 被强制 False（`trtllm.py:999-1001`），整块分配跳过；且 MLA 路径走 `mlaGeneration`，**从不读**该字段。**建议加显式 assert 而不是依赖 SM 判断。** | +| **K14** | MoE / dense MLP / allreduce / ADP `num_tokens` allgather | token-flat | ✅ | `attn_metadata.num_tokens` 就是 `_seq_lens.sum().item()`（`attention_backend/interface.py:206-215`），是真实打包计数。`fused_moe/` 与 `distributed/` 全树 grep 无 `next_n` / draft len 耦合。 | +| **K15** | DSpark drafter 自身的 forward | `[G, block]`，与 target 布局解耦 | ✅ | draft 的 attention 用 worker 自有的 `_kv_windows` 滚动 buffer，不碰 paged KV metadata。 | + +### 3.3 Host 侧阻断点（与 kernel 同等重要，且**独立于方案选择**） + +这些不修，任何 kernel 方案都不可能产出正确结果。 + +| # | 问题 | 锚点 | 严重性 | +|---|---|---|---| +| **H0** | **`disable_overlap_scheduler=True` 时，ragged token 布局根本不会被产出。** `_prepare_tp_inputs` 的非 overlap 分支条件是 `next_draft_tokens_device is None or request.is_dummy or request.py_batch_idx is None`（`model_engine.py:4606`），该分支 `sequence_lengths.append(1 + num_draft_tokens)`，**完全不看 `py_verify_len`**（`:4616,4625`）。而 `next_draft_tokens_device` 仅当 `new_tensors_device` 是 `SampleStateTensorsSpec` 时非空（`model_engine.py:4270-4277`），非 overlap 循环调用 `self._forward_step(scheduled_batch)` 不传该参数（`py_executor.py:4276`，默认 `None`，`py_executor.py:6425`）。**精度测试用的就是 `disable_overlap_scheduler=True`**（`test_llm_api_pytorch.py:4036`）。 | 上述 | **P0 / prerequisite #0**。这不是 investigation 描述的「padding 行 / 首步边缘情况」，而是在目标配置下**每个请求每一步**都命中。后果：`spec_metadata.verify_lens` / `qo_indptr` / `total_verify_tokens` 和 graph key 都在描述一个 ragged 布局，而真实 `input_ids`/`seq_lens` 是均匀的 `1+top_tier`；`_accept_draft_tokens` 随后按 `total_verify_tokens` 切片（`interface.py:1436-1441`），**逐请求 token 错位**。 | +| **H1** | **DSA 展开 stride 用的是静态 max，不是运行时 tier —— 均匀 tier 阶梯本身就是坏的。** `expand_per_gen_token` 非 ragged 分支用 `stride = 1 + self.max_draft_tokens`（`dsa.py:696`），`gen_token_repeat_list` 同（`dsa.py:662`）。`self.max_draft_tokens` 只在 `update_spec_dec_param` 赋值（`dsa.py:1401`），而 DSpark 属于 `is_parallel_draft()`（`interface.py:320-324`，`is_dspark()` 在内），故取 `original_max_total_draft_tokens = tokens_per_gen_step - 1 = 5`（`model_engine.py:342`）—— **静态值**。于是 tier 3 时 `kv_lens_expanded` / `block_table_expanded` 仍按 stride 6 铺，而 `Indexer.forward` 按 `q_decode.shape[0] = bs*4` 切片（`dsa.py:2831-2838`），请求 r>0 拿到 r-1 的 kv_len 和 block table。 | 上述 + `dsa.py:2329-2330` | **P0**。这是**所有三个方案共同需要的修复**，约 3 行 Python（`stride = spec_metadata.runtime_tokens_per_gen_step`），且**独立于 ragged 与否**。它对 07/30 那次 tier-3 capture hang 是一个完整自洽的解释：tier 5 时 768==768 通过，tier 3 时 schedule 描述 768 行而 logits 只有 512 行。**注意：`max_draft_tokens` 同时是 buffer 容量（`dsa.py:1310,1321,1342,1355`），修复时容量必须保持静态 max，只让 stride 变，否则会在 capture 期间触发重分配 —— 那是第二类 hang。** | +| **H2** | **`attn_metadata.ragged_verify_lens` 在 `attn_metadata.prepare()` **之后**才发布。** `prepare()` 在 `model_engine.py:5326`，`_attach_ragged_verify_layout`（唯一写入者）在 `:5412`，同一函数内相隔 86 行。 | 上述 | **P0**。后果：`prepare_for_spec_decode` 里的 `use_expanded_buffers_for_mtp` 判定、`kv_lens_expanded_*` 与 `block_table_expanded` 的铺设**全部用上一步的窗口**。因为 `maybe_get_cuda_graph` 返回的是 per-key 的 graph 常驻 metadata 对象（`cuda_graph_runner.py:387-407`），这个「上一步」其实是「**上一次该 key 跑的时候**」。修复：把 attach 移到 `prepare()` 之前（它只需要 `generation_requests`，已在作用域内）。 | +| **H3** | **工作区新增的 `_prepare_ragged_row_kv_lens` 有 device 不匹配，首个 ragged step 必抛异常。** `row_kv_lens.add_(expanded.to(torch.int32))`（`dsa.py:1695`）：`row_kv_lens` 是 `row_kv_lens_cuda` 的切片（CUDA），而 `expanded` 来自 `expand_per_gen_token(kv_lens[...])`，`kv_lens = cached_token_lens + seq_lens_kv` 两项**都在 CPU 上**（`dsa.py:790-801`；`cached_token_lens` 显式 `device='cpu'`）。`.to(torch.int32)` 不改变 device。同函数 docstring 写「kv_len 留在 device 上」，与自身输入不符。 | `dsa.py:1673-1695` | **P0（新引入）**。修法也更省：把逐行修正量直接折进 `row_kv_lens_host`，做**一次** H2D 拷贝 —— 这正是相邻的 `kv_lens_expanded_host` 已经在用的模式（`dsa.py:1590-1592`）。 | +| **H4** | **`row_kv_lens_cuda` 构建时机错 + 从不被 `on_update_kv_lens` 刷新。** (a) 它在 `prepare_for_spec_decode` 内构建（`dsa.py:1608`），受 H2 影响读到陈旧窗口；(b) overlap 场景下 `_preprocess_inputs` 会在 device 上把 `previous_kv_lens_offsets_cuda` 加进 `kv_lens_cuda` 再调 `on_update_kv_lens`（`model_engine.py:2921-2937`），该函数重建了 `kv_lens_expanded_cuda`、`scheduler_metadata_buffer_expanded`、slot mapping —— **但没有 `row_kv_lens_cuda`**。于是同一次 op 调用里，`gen_kv_lens_cuda`（已修正）与 `row_kv_lens`（未修正）对同一个量的说法不一致。 | `dsa.py:1608`、`dsa.py:902-977` | **P0（新引入）** | +| **H5** | **CUDA graph capture 的 token 轴来自 `key[1]`，不看 bucket。** `num_tokens_for_capture = batch_size * max_beam_width * (key[1] + 1)`（`cuda_graph_runner.py:446-453`），而 ragged 下 `key[1]` 被钉死在 top tier（`model_engine.py:1840-1842`）。 | 上述 | **P0**。结合 H0，同一 bs 下的三个 graph 都是按 `bs*(top_tier+1)` 的布局 capture 的 —— **key 里的 bucket 是虚构的**。H0 与 H5 必须**同时**修，否则 graph 只会更糟。 | +| **H6** | **overlap 下 KV rewind 读的是被下一步覆盖过的 `py_verify_len`。** `SpecSamplerBase` 已经把 batch-wide 的 `runtime_draft_len` 快照进 `SampleStateSpec`（`spec_sampler_base.py:218`）正是为了规避这个 hazard，但 `_verified_len` 却**实时**读请求上的 `py_verify_len`（`spec_sampler_base.py:195-196`）。而第 N 步的 `_handle_dynamic_draft_len`（`py_executor.py:4652`）先于第 N-1 步的 `_update_requests`（`py_executor.py:4746`）执行。 | 上述 | **P1**。修法一致：把 `py_verify_len` 同样快照进 `SampleStateSpec`。**方向修正**：investigation 声称的「泄漏 3 个 KV entry」对 KVCacheManager v1 不成立 —— `update_resources` 之后还有补偿项 `extra_rewind = _kv_reserve_draft_tokens - (py_rewind_len + accepted)`（`resource_manager.py:1051-1063`），`py_rewind_len` 在合法区间内代数抵消；真实误差只在 clamp 处出现，方向是**过度 rewind**（丢已接受的 KV）。但在 `kv_cache_manager_v2` 里补偿被 `if self.is_draft:` 门控（`kv_cache_manager_v2.py:3538-3546`），**target manager 无补偿**，那里 investigation 的场景是真的。**须先确认 DSv4-Pro 实际走哪个 manager**「待验证」。 | +| **H7** | **ragged batch 上 rejection sampling 被 fail-closed guard 静默关掉。** `_rejection_buffers_valid` 要求 `logits_rows >= num_contexts + num_gens*(draft_len+1)`（`interface.py:1707-1710`），而 ragged 下 target 只发 `num_contexts + total_verify_tokens` 行，**当调度真的在 trim 时该式必然不成立**。于是永远退回 `_sample_and_accept_draft_tokens_base`（严格 token 相等）。 | 上述 | **P1**。副作用：PR 在 rejection 路径里精心加的 ragged 分支（`interface.py:1793-1806`）成了死代码。GSM8K temperature=0 **检测不到**这一点。 | +| **H8** | **`fit_ragged_verify_lens` 假定 CUDA graph padding 已开。** 它只检查 `runner.enabled`（`model_engine.py:3644-3648`）就按 `padded_bs` 行数规划并预留 `n_pad * pad_len` 个 token，而 `_get_padded_batch` 在 `not self.padding_enabled` 时直接返回 0（`cuda_graph_runner.py:576-579`），且 `BaseCudaGraphConfig.enable_padding` **默认 False**（`llm_args.py:180-184`）。 | 上述 | **P1**。两种后果：(a) padding 关闭时每个 ragged step 静默掉出 graph replay；(b) **混合 ctx+gen batch** 无论如何都无法跑 graph，真实请求却被按「要给 pad 行留位置」缩小了窗口 —— 系统性的吞吐损失，且 `runner.ragged_pad_verify_len` 会作为陈旧状态残留。须在 config 层校验 `enable_padding=True`。 | +| **H9** | **rank-local 早退门控了条件性的 `tp_allgather`，是 hang 风险。** `decide_verify_lens` 的 `all_rank_max` 在早退之后（`dspark_verify.py:234-241,256-259`），`peer_stats` allgather 被 `if ragged_lens is not None` 门控（`py_executor.py:3258-3261`），fallback 分支又调 `decide_draft_len` 再发一次。早退条件是 rank-local 且**时序相关**（`_copy_event.query()`，`dspark_verify.py:132-133`）。 | 上述 | **P1**。走 ragged 的 rank 发 2 次集合通信，fallback 的 rank 发 1 或 3 次 → 死锁。修法：让每个 rank **无条件**发出相同序列的集合通信，再基于归约结果分支。对比之下**均匀路径天然没有这个问题**：`decide_draft_len` 的 `all_rank_max` 无条件执行（`dspark_verify.py:194-201`）。 | +| **H10** | **DSpark drafter 用固定 stride 取 target hidden state。** `Kp1 = runtime_draft_len + 1`；`base = gen_start + arange_g * Kp1`（`dspark.py:600,624`）。ragged 下 `runtime_draft_len` 被钉在 top tier，而真实偏移是 `spec_metadata.qo_indptr[r]`。`dspark.py` 全文无 `qo_indptr` 引用。 | 上述 | **P1**。静默：只是接受率崩塌，不报错（`interim_base` 还被 clamp 到 `captured.shape[0]-1`）。注意 `8723506608` 已把它从 `max_draft_len+1` 修成 `runtime_draft_len+1`，**均匀路径是对的**，只有 ragged 错。 | + +### 3.4 一条重要的语义澄清 + +`indexer_topk_decode` 是**整条 indexer 链路上唯一施加窗口内因果裁剪的地方**。因为 K3 的 paged-MQA-logits GEMM 在 expanded 布局下 `next_n==1`，每行的 `q_pos = ctx_len - 1`，且一个请求的所有行共享同一个 `ctx_len`，所以它输出**全长未裁剪**的 logits。这意味着 `row_kv_lens` 算错**不是噪声，是让一个 verify token 看到它本不该看到的 KV**（这些 KV 确实已经写进去了，因为整个窗口在 attention 之前就已 append）—— 属于信息泄漏级别的正确性错误，不是精度抖动。 + +--- + +## 4. CUDA graph 兼容性要求 + +### 4.1 不变量（replay 期间必须恒定的东西） + +| 类别 | 内容 | +|---|---| +| **可以原地改内容** | 所有 capture 时地址固定的预分配 buffer：`shared_static_tensors['input_ids'/'position_ids']`（`cuda_graph_runner.py:172-184`）、per-key 常驻的 `attn_metadata` 及其 `get_memory_buffers` 持有的一切、`ragged_verify_lens_cuda` / `ragged_qo_indptr_cuda`（`model_engine.py:715-720`）、新增的 `row_kv_lens_cuda`（`dsa.py:1321-1345`）、`gen_new_tokens_per_seq_cuda`（`deepseek_v4.py:337-346`）、DSpark worker 的 `captured_hidden_states` / `_kv_windows` / `_confidence_logits`。 | +| **必须恒定（capture 时冻结）** | 每一个 tensor **shape**，以及每一个用于切片或分支的 **host 标量**：`spec_metadata.total_verify_tokens`、`metadata.is_ragged_verify`、`use_expanded_buffers_for_mtp`、`next_n`、`num_gen_tokens_per_seq`、`attn_metadata.num_tokens` / `num_ctx_tokens` / `num_generations`、`all_rank_num_tokens`、`sum(ragged_verify_lens)`。 | +| **合法** | graph 私有内存池内的分配（`scatter_ragged_to_padded`、`count_accepted_ragged` 等，`dspark_ragged.py:432-512`）。 | + +这条清单解释了为什么 §3 里的多数 bug 是**静默**的：`num_gen_tokens_per_seq` 之流在 capture 时被烤进 graph，replay 时既不会更新也不会报错。同理，`dsa.py` 里那些 `assert` 只在 capture 时执行过一次，replay 时永远不跑。 + +**一个确认是安全、不要「顺手修」的点**：`spec_metadata.total_verify_tokens` 虽然是 captured region 内读取的 host int（`interface.py:1439-1441`），但它**等于 bucket，即 `key` 的一个分量**，因此是 graph key 的函数，构造上安全。 + +### 4.2 Graph 数量与内存代价（DSv4-Pro / dep8 实测） + +来自 `.dspark-logs/gsm8k.log` 的**实测**数字（不是估算）： + +``` +capturing 102 graphs (34 batch sizes x 3 draft-length tiers [1, 3, 5]) +``` + +| 项 | 值 | +|---|---| +| baseline graph 数 | 34 | +| confidence scheduling 后 | **102**（3×） | +| tier ladder | `[1, ceil(K/2), K] = [1,3,5]`（`llm_args.py:2883-2897`） | +| 每 graph metadata 成本 | ~10–23 MB（PR 自述，`model_engine.py:1825-1826`） | +| 额外 68 个 graph | **~0.7–1.6 GB / rank** | +| dep8 聚合 | **~5–13 GB** 的 KV 容量 | + +这笔钱**直接从 KV pool 里出**：captured graph 的字节落在 `extra_cost = total_used_bytes - torch_used_bytes` 项里，`_util.py:990-1015` 把它折进 `peak_memory`，而 `peak_memory` 决定 `kv_cache_max_memory`（该处日志行本身就点名了 CUDA graphs）。 + +**没有**第二次翻倍：`needs_non_greedy_capture` 门控在 `use_one_engine()` 上，DSPARK 不在其中（`interface.py:307-309` vs `:322-323`）；`DeepSeekV4SparseAttentionConfig.needs_separate_short_long_cuda_graphs()` 返回 False（`llm_args.py:1126-1128`）。 + +Warmup 会把整个 ladder 跑两遍（`model_engine.py:1220-1227`），即 204 次前向。 + +**这是个真实的张力**：feature 的收益按其自身 docstring 只在高并发出现（`llm_args.py:2812-2816`），而高并发恰恰是 KV 容量最吃紧的时候。缓解手段是把 tier ladder 按 batch size 门控（低于 SPS 拐点时只留 `[max_tier]`），因为 planner 的 argmax 在拐点以下本来也会选 max_tier。 + +### 4.3 剩余的同步点 + +| 位置 | 问题 | +|---|---| +| `model_engine.py:3910-3919` | `_update_target_input_tensors` 的 ragged 分支 `torch.repeat_interleave(previous_slots, tokens_per_request_device)` **没传 `output_size`**，两个参数都在 device 上 → torch 把 cumsum 读回 host 定尺寸。这正是 `dspark_ragged.py:414-424` 文档里明说不能干的事。在 capture 之外，所以不会让 capture 失败，但每个 ragged step 都在 `_apply_incremental_update_target` 里插一次 D2H sync，把 overlap scheduler 的收益吃掉。host 侧已有总数，直接传即可（旁边 `model_engine.py:5101-5109` 就做对了）。 | +| `heads.py:280-283` | `apply_sts` 在 CPU 快照上做校准，会把 CUDA 上的 `sts_temperatures` `.to('cpu')` 到**可分页**内存 —— 同步拷贝。每个 decode step 一次，且**只在提供了 profiled cost table 时触发**（否则 `is_flat` 提前返回），也就是恰好在真正要用的配置里出现。加载时缓存一份 host 副本即可消除。 | +| `py_executor.py:3232-3236` | ADP 下每 decode step 多一次阻塞式 host collective（`all_rank_max`），PR 自己的注释已承认。 | + +### 4.4 Attention-DP 形状一致性 + +`cuda_graph_runner.py:362-372` 的 ADP 资格门只比较 **`is_all_gen_only` 与 batch size 相等**，**从不比较 token 数**。token 数走另一条 allgather（`model_engine.py:2987-2998`）进 `all_rank_num_tokens`，且因为是 host list 在 forward 内读取，会被**烤进 captured graph** —— replay 时用的是 capture 时的那份，与本步 allgather 结果无关。 + +顺利路径下这是自洽的（`choose_ragged_capture_shape` 归约 `peer_stats` 使各 rank 落在同一 `(padded_bs, bucket)`，且 `fill_bucket` 宁可抛异常也不返回短 batch，`dspark_ragged.py:37-48`）。**暴露面在于任一 rank 退回均匀路径**：它产出 5 元组 key 和 `bs*(chosen+1)` 个 token，而同伴产出 6 元组 key 和 bucket 个 token；batch size 仍相等所以 ADP 门通过，一个 rank replay 满是集合通信的 graph、另一个 eager 跑不同 token 数，而冻结的 `all_rank_num_tokens` 使这个不匹配在 replay 时**无法被察觉**。 + +**要求：ADP 门必须同时比较 bucket。** + +### 4.5 一个尚未有人排查的 capture 期风险 + +`deepseek_v4.py` 里有三个 `@maybe_compile(dynamic=True, options={"max-autotune": True})` 的函数 —— `_prepare_deepseek_v4_indices_compiled`（`:586-589`）、`_compute_gen_compressed_position_ids`（`:955-957`）、`_compute_compressed_mask`（`:996-998`）—— 全部经 `on_update_kv_lens`（`:854-895`）在 `_preprocess_inputs`（`model_engine.py:2891-2895`）内被调用，**而那正是 `CUDAGraphRunner.capture` 录制的区域**。`_compute_gen_compressed_position_ids` 把 `num_gen_tokens_per_seq` 当 Python int 用（`:979-983`），构成 guard → **每个 tier 一次 Inductor 编译，且发生在 graph capture 期间**。 + +`max-autotune` 需要同步来 benchmark 候选 kernel。confidence ladder 把 capture 期编译次数变成 3 倍。这是 07/30 那次 hang 的**另一个独立候选根因**，必须与 H1 的 stride 假说区分开（见 §7）。「待验证」。 + +--- + +## 5. 方案对比 + +三个候选方案。注意 D1 与 D2 **不是独立选项** —— 二者都把成败押在同一个无法从源码证明的 trtllm-gen 行为上（K2 的 cubin 是预编译的）。真正的选择是 **token-major vs tier-uniform**。 + +| 维度 | **D1 Token-flat（`next_n ≡ 1`）** | **D2 Token-major + 3 处签名改动** | **D3 TUV 分层均匀** | +|---|---|---|---| +| **机制** | 把 batch 的 generation 半边呈现成 `num_gen_tokens` 个「长度为 1 的序列」，每 token 带自己的 kv_len / block table / cu_seqlens。所有 `rowIdx/next_n` 退化为恒等式。 | 引入每 token 行向量（`row_req_idx` / `row_kv_len` / `row_idx_end`），给 3 个 kernel 加**可选**参数；attention 仍用 token-major 呈现规避 cubin 改动。 | 不教 kernel ragged，而是把 verify 长度量化到 tier ladder，**保证任何一次 launch 只看到一个 `next_n`**。Stage A：整步一个 tier。 | +| **kernel 改动** | **0**（宣称）。compressor 无法展平（Phase 3 归约的是 Phase 1 同 block 写的状态，`compressorKernels.cu:355-400` vs `:440-530`），靠 host 侧按 tier 分组多次 launch + scratch 行绕开。 | 2 个 op 加可选参数（**工作区已完成**：K4、K9）。attention 不改。 | **0**。四个 uniform kernel 在 `num_gen_tokens = n·(K+1)` 时**逐一精确正确**（已逐条读码确认：`indexerTopK.cu:667-668`、`mlaKernels.cu:1161-1163`、`compressorKernels.cu:294,340`、`attentionOp.cpp:1161`）。 | +| **graph 数量** | 102（bs × tiers，token 轴入 key） | 102（同上） | 102 → 按 bs 门控可降到 **~54** | +| **预期收益** | 完整 Θ = τ·SPS(B) | 完整 Θ | 拿满 **SPS(B)** 项；损失 **τ** 项。D3 自估双峰 batch 下 τ 差距 ~14%（**未实测**） | +| **风险** | **高**。核心假设（K2 在 `mMaxSeqLenQ=1` 下正确且不显著变慢）无法读码证明。compressor 的 tier 分组 + nulled 行 + scratch 是全案最难 review 的部分。 | **高**（同一 K2 假设）。但 kernel 改动是加性可选参数，非 ragged 路径可证明不变。 | **中**。无 K2 假设。但 τ 留在桌上，且 Stage B（两 tier 分子 batch）依赖一个未验证的 FMHA semaphore 复用假设（`attentionOp.cpp:1126-1128`）。 | +| **工作量** | XL | L（**已完成约 40%**） | M | +| **被 D1/D2 漏掉的隐性成本** | token-major 要求 `num_seqs = host_context_lengths.size(0)`（`thop/attentionOp.cpp:1325`、`dsv3RopeOp.cpp:139`），故 `prompt_lens_cpu/cuda_runtime`、`kv_lens_runtime`、`host_request_types_runtime`、**主 KV 的 `kv_cache_block_offsets`**（与 indexer 的 `block_table_expanded` 是不同 buffer）都要按 token 膨胀；`max_num_requests` 变成 `max_batch*(1+max_draft_len)`，连带 `reserveSemaphoreArray`（`thop/attentionOp.cpp:407`）、FMHA/FlashMLA workspace（`common/attentionOp.cpp:925-955`）一起放大。4096 上下文下展开后的主 block table 约 400 KB，128K 上下文下约 12 MB，**每个 decode step 重新 H2D**。D3 完全不付这笔钱。 | 同左 | — | + +### 5.1 推荐 + +**分三阶段推进，先把「独立于方案选择」的东西修干净并拿到真实测量，再定 kernel 方向。** + +理由： + +1. **§3.3 的 H0–H5 全部是 host 侧问题，且三个方案都需要。** 尤其 H1（3 行 Python）就能完整解释唯一一次实测 hang，且**独立于 ragged 与否**。在它修好、`enable_confidence_scheduling` 的均匀路径能跑通之前，讨论 D1/D2/D3 是没有依据的。 +2. **决定 D1/D2 vs D3 的那个数字，谁都没测过**：trtllm-gen sparse-MLA generation 在 `(batch=B, s_q=tier+1)` 与 `(batch=B*(tier+1), s_q=1)` 两种呈现下、相同总 token 与相同逐 token 稀疏索引集时的吞吐。这决定 token-major 是否值得它的爆炸半径。它读不出来（`selectMlaGenerationKernel` 在预编译 cubin 里），但**一张 B200 上直接调 `thop.attention` 就能测**，不需要模型权重。 +3. **工作区已经沿 D2 走了一半**（K4、K9 已实现且设计正确 —— 可选参数、非 ragged 路径 bit-identical、TORCH_CHECK 替换而非放宽）。这部分不该丢，它们对任何 ragged 方案都是必需的，而且对 D3 无害（不传参数即退化）。 +4. **D3 有一个 D1/D2 没有的结构性优势**：`decide_draft_len` 的 `tp_allgather` 是**无条件**的（`dspark_verify.py:194-201`），且各 rank token 数恒为 `bs*(K+1)`，H9 那类 collectives 发散的 hang 在构造上不存在。同理 H7（rejection sampling）在均匀路径下是等式成立，不会被关掉。 + +**阶段划分：** + +| 阶段 | 内容 | 交付 | +|---|---|---| +| **P0：地基（不涉及方案选择）** | H1（stride 用 runtime tier，容量保持静态 max）、H2（attach 移到 prepare 之前）、H3（`row_kv_lens` host 侧折叠 + 单次 H2D）、H4（`on_update_kv_lens` 补刷 `row_kv_lens`）、H0+H5（dummy 分支用 `get_request_tokens_per_gen_step`；`capture()` 吃 bucket —— **必须同时改**）、K8（`prepare_scheduler_metadata` 走 `expand_per_gen_token`）、§4.3 两处 sync、H8 config 校验、K5/K6/K13 显式 config-reject + assert | **均匀 tier 路径**在 8×B200 上跑通 GSM8K；hang 消失 | +| **P1：测量与定向** | 单卡 B200 上的三个微基准（见 §6.4）；产出 STS 校准表与 SPS cost table（G8） | 用数据定 token-major vs tier-uniform | +| **P2：ragged 落地** | 按 P1 结论推进 D2（已完成 K4/K9）或 D3；无论哪条都要修 H6、H7、H9、H10 | ragged 路径 GSM8K + 逐 token 等价 | + +**若 P1 显示 token-major 显著变慢，D3 Stage A 就是 ship 目标** —— 它拿满 SPS 项（vLLM 报的 1.71×/2.19× 主要来自跨过吞吐拐点，而非预算分配的精细度），风险和工作量都低得多，且能把 graph 数从 102 压到 ~54。 + +--- + +## 6. 验收标准 + +> **所有测试都必须在 Slurm 分配的 GPU 节点上跑。当前节点不可用于任何 GPU 测试。** +> 用 `computelab-node-allocator` / `gpu-test-runner` skill 分配（B200/B300 分区，≥140 GB/GPU，≥4 h 预约）。 + +### 6.1 GSM8K 精度门（G3） + +```bash +export LLM_MODELS_ROOT=/scratch.trt_llm_data/llm-models # 容器内路径 +cd /code/tensorrt_llm && python -m pytest \ + "tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling" -vs +``` + +`MODEL_PATH` 在 **class body**（import 时）求值（`test_llm_api_pytorch.py:3964`），所以 `LLM_MODELS_ROOT` 必须在 pytest 启动**前**导出，否则整个 module 收集失败。无需 `mpirun` —— 单 pytest 进程通过 mpi4py.futures 起 8 worker。 + +**这个门本身不够**：GSM8K 的统计阈值是 `96.0 - 3.203 = 92.797`，允许 **42/1319** 题翻转；额外的 `score >= 96.0` 断言对实测 96.475 只有 **6.3 题**余量（1 题 = 0.0758 分），而 GSM8K 在 p=0.96 处的二项标准差约 0.54 分。它既太松（错几个 token 的 bug 能过）又太紧（干净 run 可能假失败）。 + +另外必须注意：该测试**没有被任何 CI stage 引用**（全 `tests/integration/` grep `confidence_scheduling` 只命中测试文件自身），且其 baseline 兄弟 `test_gsm8k_dep8_megamoe_deepgemm` 已被 `waives.txt:184` 以 nvbugs/6506920 跳过 —— **当前没有活的 baseline 可比**。 + +### 6.2 更锐利的不变量测试（G2，这才是真正的门） + +| 测试 | 内容 | 为什么锐利 | +|---|---|---| +| **A1 逐 token 等价** | 同一 engine 配置、同一批 prompt，`static`（verify 全量）与开启调度后，temperature=0 下断言每个 prompt 的 `token_ids` **完全相等** | 贪婪投机解码定义上是 lossless，任何分歧都是 bug。把 6 题的容忍度变成 1 token 的容忍度。**注意**：跨不同 graph 形状的严格 bitwise 相等 TRT-LLM 并不保证（MoE/attention 规约顺序会变），所以诚实的形式是允许极少数 prompt 分歧并设硬上限，或用 `cap-accept` 模式（下条）。 | +| **A2 `cap-accept` 差分** | 实现 SGLang 的三模式（`static` / `cap-accept` / `compact`）。`cap-accept` 跑**均匀 kernel 路径**（`next_n=6`）但只提交每请求窗口内的 token，输出应与 `compact` **逐 token 相同** | `cap-accept` vs `compact` 的差异 ⇒ **必然是 ragged kernel bug**；`static` vs `cap-accept` 的差异 ⇒ **必然是调度/接受逻辑 bug**。把 pass/fail 变成可定位。工作区的 `dspark_observability.py` 已有 mode 枚举，`cap-accept` 目前显式 `NotImplementedError`（`dspark.py:471-478`）—— **这是需要补的**。 | +| **A3 布局一致性断言** | host 侧断言：`sum(1 + py_verify_len over padded generation_requests) == attn_metadata.num_tokens - num_ctx_tokens == spec_metadata.total_verify_tokens == key[bucket]` | 这一条就能同时抓住 H0、H5、H8 | +| **A4 时序断言** | 断言 `attn_metadata.ragged_verify_lens` 在 `prepare()` 执行的那一刻是**本步**的（而非上一步） | 抓 H2、H4 | +| **A5 接受长度分布** | 记录并断言每请求接受长度直方图，以及 `sum(accepted)` 与接受 kernel 报告一致 | 抓「token 落到了错误请求的槽位但整批仍生成合法文本」这一类 | +| **A6 调度确实生效** | 断言 `planner.stats` 的 `fallback_flat_cost` / `fallback_no_snapshot` / `fallback_no_confidence` / `fallback_short_snapshot` 计数（`dspark_verify.py:86-93`）显示 trim 路径被走到 | 否则一个静默退化的 run 会以「正确的分数、错误的理由」通过 | +| **A7 capture 期无 sync/无分配** | capture 期间断言无 device→host sync、无新分配 | G4 | + +### 6.3 必须存在的单元测试 + +| 测试 | 现状 | +|---|---| +| C++ ragged top-k 差分：均匀 batch 用 ragged 形式表达（`row_kv_lens` 填成 uniform 算术会产出的值）必须与 `row_kv_lens=None` **bitwise 相同** | **已新增** `tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py`（需在 GPU 节点上实跑验证） | +| compressor ragged 差分：`new_tokens_per_seq` 全填 `next_n` 必须与传 `None` bitwise 相同 | **欠缺** | +| `test_deepseek_v4_sparse_mla.py` 按每请求 q 长度参数化 | **欠缺**。该文件已构造真实的 `DeepseekV4TrtllmAttention` + cache manager（单卡 B200，`l0_b200.yml:80`，TIMEOUT 60），但其参考实现在 `:415` 和 `:473` **硬编码** `seq_len_q = fused_q.shape[0] // num_requests` —— 那正是被测假设本身。**这是性价比最高的一个测试**，单卡分钟级就能抓住 tier-3 那类 hang。 | +| `_prepare_tp_inputs` 在混合 `py_verify_len` + pad 行的 batch 上，断言 `sum(attn_metadata.seq_lens) == _ragged_verify_bucket(batch)` | **欠缺**。现有 220 个 hw_agnostic 测试全部停在纯 Python 布局算术；`test_dspark_ragged.py:515-521` 的 `_fit` 是 `fit_ragged_verify_lens` 的**手写镜像**而非真函数，所以 §3.3 的每一个 bug 对 CI 都不可见。 | + +### 6.4 定向微基准（P1，单卡 B200，分钟级） + +| # | 内容 | 用途 | +|---|---|---| +| **B1** | 修完 H1 后重跑均匀 tier capture，确认 hang 消失 | 区分 H1 stride 假说 vs §4.5 torch.compile 假说。**跑时带 `TORCH_LOGS=recompiles`。** | +| **B2** | 直接调 `thop.attention`，比较 `(batch=B, s_q=tier+1)` 与 `(batch=B*(tier+1), s_q=1)` | **决定 D1/D2 vs D3 的那个数字** | +| **B3** | expanded 布局下 paged-MQA-logits 的吞吐 | 展平路径的 KV 重读代价（每请求 6 次而非 2 次）在关键路径上，必须实测 | + +### 6.5 GSM8K 的并发要求 + +vLLM #47808 在**并发 16 和 64** 两个点上报了 0.945 / 0.951。原因很实际：**调度只在高并发下才真的 trim**，单一低并发的 GSM8K run 是空跑。本项目的验收也必须至少覆盖两个并发点。 + +--- + +## 7. 风险与未知 + +### 7.1 必须在定稿前测量的 + +| # | 未知 | 为什么阻塞 | 怎么解 | +|---|---|---|---| +| **U1** | trtllm-gen sparse-MLA generation 在 `mMaxSeqLenQ=1` + 6× batch 下 vs `s_q=6` 的吞吐 | **决定 token-major 是否可行**。DSv4-Pro 在 DEP8 下每 rank 保有全部 128 个 Q head，BMM1 的 M 维从 768 降到 128；由于稀疏索引集逐行不同，本来也没有跨 6 行共享 KV tile 的机会 —— 但这是**推断，不是测量**。 | B2 | +| **U2** | 07/30 hang 的真实根因 | H1 的 stride 不匹配是一个完整自洽的解释（tier 5 时 768==768 通过，tier 3 时 768 vs 512），但日志里紧邻 hang 的是 **attention workspace 从 0 扩到 168 MB** 的警告，`model_engine.py:1789-1806` 明确记录 capture 期 workspace 重分配是真实 hazard，而 Case 2b **没有** Case 2 那样的 `max_spec_graph` 规避（`model_engine.py:1810-1813`）。§4.5 的 torch.compile / max-autotune 是第三个候选。 | B1 + `TORCH_LOGS=recompiles`；必要时逐个消除 | +| **U3** | ragged 相对 tier-uniform 在固定 token 预算下的 τ 收益 | 决定 D3 是否「够好」。D3 自估双峰 batch ~14%、同质 batch ~0%，**没有任何真实 confidence head 的数据支撑** | 需要先有 STS 表（U4） | +| **U4** | STS 校准表 | `sts_temperatures` 初始化为全 1（`heads.py:249-258`），checkpoint 的 index 里**没有**任何 sts/temperature/calibration 张量，repo 里也没有该 JSON。论文 §3.2.1 明确是**每个位置一个温度**的 1D 网格搜索。未校准的逐位置偏差在累积乘积下**几何放大**。 | 必须先 profile 产出 | +| **U5** | SPS cost table | 没有它 `SpsCostTable.flat()` ⇒ `is_flat` ⇒ `_decide_local` 无条件返回 `max_tier`（`dspark_verify.py:268-270`）⇒ **feature 在构造上无法 trim**。任何在此之前的 A/B 只测到了「多一次 matmul + 3× graph」。 | 必须先 profile 产出 | +| **U6** | DSv4-Pro 实际走 KVCacheManager v1 还是 v2 | H6 的**症状和严重性按 manager 不同**：v1 有补偿项、误差方向是过度 rewind；v2 的 target 路径无补偿（`kv_cache_manager_v2.py:3538-3551`），investigation 的「泄漏」场景才成立。说错会把 debug 引向反方向。 | 读配置/加日志确认 | +| **U7** | token-major 下 buffer 膨胀的真实代价 | `num_seqs` 从 B 变成 ~6B，牵连 semaphore array、FMHA workspace、**主 KV block table 每步重新 H2D**。128K 上下文下约 12 MB/step。 | 若走 token-major，需容量审计 + 实测 | + +### 7.2 已知但可接受 / 需显式处理 + +- **性能故事必须诚实**：**draft pass 是不可裁剪的** —— block 永远整块 draft（`runtime_draft_len` 在 ragged 下钉在 top tier），只有 verification 被裁。所以收益是纯粹的 target 侧 SPS 效应，**只在高并发出现**，小 batch 下可能持平甚至因为 3× graph 占用 KV 而**净负**。vLLM 报的 1.71×(c=128) / 2.19×(c=256) 的对照组是「固定 7-token verify」，而后者在 c=256 时比不投机还慢 33% —— 那是被修复的病态，不是通用加速比。 +- **`derive_verify_len_tiers` 实际是死代码**：`_build_verify_planner` 只在 `verify_len_tiers` 为假值时才调它（`dspark.py:433-450`），而该属性总是返回非空 list（`llm_args.py:2899-2905`）。这**其实是安全的**（planner 的 tier 必须与 `_get_graphs_to_capture` capture 的一致），但 `dspark.py:450` 的日志永不触发，reviewer 会误以为 ladder 是 profile 派生的。要么删分支、要么让派生结果同时喂给 capture 列表。 +- **Case 2b 丢了 `draft_len=0` 的 graph**：它提前 return（`model_engine.py:1848,1856`），从不到达 `should_capture_no_spec` 逻辑（`:1867-1873`），而 tier ladder 最小值是 1。一旦 speculation gate 因低接受率关掉投机（`py_executor.py:3317-3322` 设 `runtime_draft_len = 0`），key 变成从未 capture 过的 `(bs, 0, ...)`，**之后每个 decode step 都 eager** —— 大 MoE 模型上是巨大的静默性能悬崖。 +- **PR 的外部引用有误**，会削弱 reviewer 信任：算法在论文 **§3.2.2**（Algorithm 1）而非 §5；SGLang **没有** `dspark_schedule.py`（其真实布局是 `speculative/ragged_verify.py` + `speculative/dspark_components/*`）；`align_verify_tokens_to_graph_tier` 全网零命中（SGLang 的实际原语是 `round_up_grid()` / `compute_target_verify_graph_key()`）。 +- **论文 §5.2 支持 PR 的一处偏离**：Algorithm 1 的 greedy first-descent 被论文自己否定（「真实硬件容量 SPS(B) 是离散的、阶梯状退化的」），生产中「移除 early-stopping break，启用无约束全局搜索」—— 这正是 PR 采用 global argmax 的理由，可以引用。但论文用的是 **t-2** 快照（为兼容 Zero-Overhead Scheduling 并形成「因果屏障」），PR 用的是 t-1。**t-1 是否足够早取决于 TRT-LLM 的 overlap scheduler，不取决于论文**「待验证」。 + +### 7.3 上游可复用的东西（尚未采用） + +TRT-LLM pin 的 DeepGEMM commit（`245dc5d6`，`cpp/build/_deps/deepgemm-subbuild/CMakeLists.txt:30`）**已经**在 `get_paged_mqa_logits_metadata` / `fp8_paged_mqa_logits` / `fp8_fp4_paged_mqa_logits` 上暴露了可选的 `indices` 参数（varlen 分支由 `indices.has_value()` 门控，断言 `dim()==1 && size(0)==batch_size`、int32、contiguous）—— 这就是论文说的「marker tensor」。vLLM #47808 用的正是它（`indices=decode_metadata.indices`）。 + +TRT-LLM 的 `dsa.py` **从不传 `indices`**，而是用既有的 "expanded buffers" 展平达到同样形状。二者数值等价，差别在于 expanded 路径要物化一个 `[num_tokens, max_blocks_per_seq]` 的 block table，而 varlen 只需一个 1-D int32 marker。**限制**：DeepGEMM 的 varlen 分支断言 `arch_major == 10 and next_n == 1` —— 仅 SM100，且必须是展平形式。这条路值得作为一个**有意的决定**评估，而不是默认忽略。 + +--- + +## 8. 开放问题（需要你拍板） + +1. **token-major 还是 tier-uniform？** 这是全文最大的分叉。建议先跑 B2 微基准再定 —— 但如果你出于其他原因（例如已经决定要与 SGLang/vLLM 的实现保持形态一致）倾向某一边，可以省掉这一步。 + +2. **`enable_ragged_verify` 是否是本轮必须交付的？** 若可以接受先只交付**均匀 tier 阶梯**（D3 Stage A），P0 阶段结束就能 ship 一个正确且 CUDA-graph 兼容的版本，把 ragged 留作 follow-up。这会显著降低风险，代价是 τ 项（未量化）。 + +3. **工作区那 11 个未提交的改动如何处置？** K4/K9 的实现质量是好的（可选参数、非 ragged bit-identical、TORCH_CHECK 替换而非放宽），但 H3/H4 是新引入的 P0 bug。建议：先修 H3/H4 再提交，**不要**在带着必抛异常的代码的状态下推进其他工作。 + +4. **`cap-accept` 模式要不要实现？** 它是把 GSM8K 从「pass/fail 硬币」变成「可定位诊断」的关键（§6.2 A2），成本不高（复用 PR 已有的 `torch.minimum(..., verify_lens)` clamp，`interface.py:1473-1480`、`:1889-1896`）。当前是显式 `NotImplementedError`。 + +5. **STS 表和 SPS 表由谁、按什么流程产出？** 这是 U4/U5，也是 G8。在它们存在之前，任何**性能**结论都是空的，而且 feature 在构造上无法 trim。需要一个 profiling 脚本 + 产物的存放约定（checkpoint 内？repo 内 `examples/configs/`？）。 + +6. **tier ladder 要不要按 batch size 门控？** 低于 SPS 拐点时 planner 本来也会选 `max_tier`，门控后 graph 数从 102 降到 ~54，换回 ~0.4–0.8 GB/rank 的 KV。需要一个共享 helper 让 `_get_graphs_to_capture`（`model_engine.py:1841-1856`）与 planner 的 `allowed_lens`（`dspark_verify.py:280`）**不可能漂移**。 + +7. **非贪婪采样（H7）的处理时机？** 现在 ragged batch 上 rejection sampling 被 fail-closed guard 静默关掉，改变了 temperature>0 的采样分布。要么本轮修 guard（放宽成 `num_contexts + total_verify_tokens`），要么显式在 config 层拒绝「ragged + 非贪婪」组合。**静默降级是最差的选项。** + +8. **KVCacheManager v1 还是 v2？**（U6）影响 H6 的修复优先级和症状描述。 + +9. **要不要采用 DeepGEMM 的 `indices` varlen 参数？**（§7.3）能去掉一个大 buffer 并与 vLLM 对齐，但把该路径钉死在 SM100。 diff --git a/docs/dspark_ragged_verify_status.md b/docs/dspark_ragged_verify_status.md new file mode 100644 index 000000000000..303339cc71ac --- /dev/null +++ b/docs/dspark_ragged_verify_status.md @@ -0,0 +1,351 @@ +# DSpark ragged verify — 目标与现状 + +配套文档：`docs/dspark_confidence_schedule_goal.md`（设计与方案对比）。 +本文档记录**工作状态**：做到哪、验到哪、哪些还不能信。 + +最后更新：2026-07-31。分支 `pr-17056`，基线 PR [NVIDIA/TensorRT-LLM#17056]。 +**所有改动均未提交**（工作区 15 个修改 + 4 个新增文件）。 + +--- + +## 1. 目标 + +主线：让 **ragged verify**（每个 request 各自的 verify 长度）在 TensorRT-LLM 里真正跑通。 +这是 confidence-head 调度收益的唯一来源，**必须实现**，不接受退化成 batch-uniform K 的方案。 + +验收条件，缺一不可： + +| # | 条件 | +|---|------| +| 1 | 目标模型 DeepSeek-V4-Pro，真实 checkpoint | +| 2 | GSM8K 端到端精度不低于 baseline | +| 3 | **必须验 ADP**（目标平台已定为 **B300 / SM103 / 8 卡**）：`tensor_parallel_size=8` + `moe_expert_parallel_size=8` + `enable_attention_dp=True` | +| 4 | **确认 ragged 真的生效**，而非「配置打开了但代码走了 uniform fallback」 | +| 5 | kernel 侧兼容（indexer top-k / DSv4 compressor 不能静默算错） | +| 6 | CUDA graph 侧兼容（真在 replay，不是被迫 eager） | +| 7 | **overlap scheduler 开启也要正确**——这是 default 配置，不能只验关闭的路径 | + +关于条件 2 的性质：调度只决定「多少 drafted 位置送去 target」，不改接受规则， +所以输出分布必须不变。**精度掉点 = 有 bug，不是吞吐 tradeoff。** + +关于条件 4 的必要性：该功能的每一种失效都是**静默**的——planner 拒绝 trim、 +cost model 退化、batch 错过捕获形状掉出 graph、部分加窗被拒——每一种都仍然产出 +正确输出和基线精度。**GSM8K 通过本身不能证明功能跑过。** + +--- + +## 2. 现状 + +### 2.1 已验证（有运行数据） + +| 项 | 证据 | +|---|---| +| ragged top-k kernel 正确 | `tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py` **10 passed** on B200，含 CUDA graph capture/replay，且与 `torch.topk` 参照比对 | +| 无回归 | ragged + `hw_agnostic` DSpark 共 **227 passed** | +| C++ 编译通过 | `indexerTopK.cu` / `IndexerTopKOp.cpp` / `compressorKernels.cu` / `compressorOp.cpp` 已编入 `.so` | +| checkpoint 可用 | `mtp.2.confidence_head.proj.weight`、`mtp.2.markov_head.markov_w1/w2.weight` 存在 | + +### 2.2 未验证（这是最重要的一节） + +**整条 e2e 从未跑通过。** 下面六条 P0 全部只经过读码与审计定位， +**没有一条经过运行验证**。逻辑上都成立、都能指到具体行号，但那不等于对。 + +e2e 冒烟第一次跑就撞上 `PYTHONPATH` 未设——说明这类问题还会有更多。 + +### 2.3 已知阻塞（非本工作引入，但会挡住验证） + +1. **NVRTC 编译 MLA FMHA kernel 失败** + `Failed to preprocess kernel fmhaSm100aKernel_QkvBfloat16OBfloat16HQk576HV512...: + NVRTC_ERROR_COMPILATION`。影响所有 MLA 模型，**很可能包括 DSv4-Pro**。 + 若冒烟撞上它，这是先于一切的阻塞。 +2. `test_context_sparse_attention_mqa` 6 个失败 —— 已用有效对照（stash 三个 sparse + 文件后跑同一子集，基线同样 6 failed）证明是既有问题。 + 注意：另外 53 个 sparse 失败**未逐一对照**。 +3. `.venv-3.12` 的 cutlass 曾完全坏掉（`import tensorrt_llm` 起不来），已修复， + 详见 §5。 + +--- + +## 3. 已完成的改动 + +### 3.1 kernel 层 + +**`indexerTopK.cu`** —— `topKPerRowDecode` 原来在 kernel 内部用标量 `next_n` +同时推导 row→request 映射和 causal 上界： + +```cpp +int seq_len = seqLens[rowIdx / next_n]; +int actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; +``` + +新增可空 `rowKvLens`：非空时 `actual_kv_len = rowKvLens[rowIdx]`，空时保持原式。 +数学上是严格推广（请求 verify `v` 个位置时第 `o` 行可见 `kv_len - v + o + 1`， +代入 `v = next_n` 即还原），所以 uniform 路径逐位不变。 + +`IndexerTopKOp.cpp` 新增 `Tensor? row_kv_lens`，并在 ragged 时**替换**而非放宽 +`seq_lens.size(0) * next_n == numRows` 检查——放宽的话，恰好整除的 ragged batch +（如 `verify_lens=[6,4,1,1]`，total=12，bs=4 → next_n=3）会通过检查然后静默算错。 + +**`compressorKernels.cu`** —— `NEXT_N` 是模板参数，但 body 里只用于 `sp`、 +`last_token_idx` 和两个循环上界，且两个循环**本来就有守卫** +（Phase 1 的 `token_idx < kv_len`、Phase 3 的 `c >= num_compressions`）。 +所以把 `NEXT_N` 从「确切长度」降级为「batch 内最大值（编译期上界）」即可， +真实 per-request 计数走新增的可空 `new_tokens_per_seq`。 +单次 launch、固定 grid、**CUDA graph 拓扑不变**。 + +### 3.2 六条 P0 + +| | 问题 | 修法 | 触发条件 | +|---|---|---|---| +| P0-1 | `attn_metadata.prepare()` 是 `ragged_verify_lens` 的唯一消费者，却在它发布前 86 行就跑了 | 拆出 `_publish_ragged_verify_lens` 提前调用 | 全部 | +| P0-2 | `capture()` 用 `bs*(key[1]+1)` 定宽；graph `seq_lens` 冻结 | 改用 `key[5]`（ragged bucket）；ragged 且行数不变时放行 `seq_lens` 刷新 | 第一个 trim 步 | +| P0-3 | `expand_per_gen_token` 在 capture body 内做 pinned 分配 + H2D | 新增常驻 `row_req_idx_cuda`，refresh 改成 `index_select(out=) + add_` | replay | +| P0-4 | `on_update_kv_lens` 无条件整除覆盖 compressor 的 `next_n` | 抽出 `_sync_gen_tokens_per_seq` 单一写入点 | 每次 forward | +| P0-5 | drafter 用定长 `Kp1` 跨步读 flat-indexed 的 captured hidden states | 改用 `qo_indptr`；`base + gidx` 补 clamp | **overlap 开启** | +| P0-6 | overlap rewind 读到已被下一步覆盖的 `py_verify_len` | `SampleStateSpec.verify_lens_snapshot` | **overlap 开启** | + +P0-2 的两个半边必须**一起**修：`num_tokens` 来自 `seq_lens.sum()`， +而 `capture()` 从 `key[1]` 定宽，只修一边会让 `input_ids.shape[0] != num_tokens`。 +最大 bucket 恰好等于那个乘积，而 warmup 捕的正是它——所以捕获阶段完全干净， +**第一个真正 trim 的 step 才崩**。 + +P0-3/5/6 属于「跑出来的数字看着正常但其实是错的」：replay 读已释放内存、 +跨请求串读、KV 回退量错。**这类问题不会报错。** + +### 3.3 可观测性（`tensorrt_llm/_torch/speculative/dspark_observability.py`） + +对齐 SGLang 的 `SGLANG_RAGGED_VERIFY_MODE` 语义： + +| 开关 | 作用 | +|---|---| +| `TLLM_DSPARK_RAGGED_VERIFY_MODE=static\|compact` | 选路径。`cap-accept` **未实现，显式抛 `NotImplementedError`** | +| `TLLM_DSPARK_FORCE_VERIFY_LENS=1` | 按 tier 阶梯轮转出确定性非均匀切分 | +| `DSparkRaggedStats.assert_ragged_active()` | 断言 ragged 真的生效 | + +`assert_ragged_active()` 检查四条：`steps_ragged > 0`、`distinct_verify_lens >= 2`、 +`trim_ratio > 0`、`graph_eager == 0`，失败时带完整计数器摘要。 + +**`cap-accept` 为什么不能 alias 成 `compact`**：它唯一的价值是「输出必须与 `static` +逐 token 一致」，别名化会让这个对照失去意义。 + +**`FORCE_VERIFY_LENS` 解决的循环依赖**：没有 profiled SPS cost table 时 +planner 在构造上不会 trim（`llm_args.py` 自述「the budget degenerates to +verify-all」），于是验证 ragged 打包正确性需要先有 table，而产出 table 又需要 +能跑的 run。这个开关替换的是 **planner 的决策**，不绕开任何下游代码路径—— +所以 forced-ragged 与 static 在 `temperature=0` 下必须逐 token 一致。 + +--- + +## 4. 剩余工作 + +| 优先级 | 项 | 说明 | +|---|---|---| +| P0 | e2e 冒烟通过 | 两个配置各一遍：`disable_overlap_scheduler=True` 与 **default（overlap 开启）** | +| P0 | 排除 NVRTC 阻塞 | 若 DSv4-Pro 起不来，其余都无从谈起 | +| P1 | **SPS cost table profiler** | 没有它 planner 必然不 trim，**任何性能结论都是空的**。参考 SGLang `python/sglang/benchmark/dspark_sps_profiler.py`：服务端 `RAGGED_VERIFY_MODE=static` + record flag，扫 batch size 记 step time | +| P1 | GSM8K + ADP | 冒烟通过后才有意义 | +| P1 | graph 按 token 数分桶 + ADP tier 协商 | SGLang `compute_target_verify_graph_key` 对 ragged 返回 `(graph_num_tokens, graph_num_tokens)`——两个轴都变 token 数；DP 下各 rank 共用同一 tier 并一起降档 | +| P2 | 实现 `cap-accept` | 隔离「kernel 算错」与「planner 没触发」的最强工具 | +| P2 | 对照剩余 53 个 sparse 失败 | 目前只对照了 mqa 那 6 个 | + +--- + +## 5. 任务清单（截至 2026-07-31） + +| # | 任务 | 状态 | +|---|------|------| +| 1 | 目标文档 | 完成 | +| 2 | indexer top-k kernel ragged 化 | 完成（B200 上 10 用例通过，**但在错误镜像下跑的，需重验**） | +| 3 | DSv4 compressor ragged 化 | 编译通过，**未运行验证** | +| 4 | 审计 DSv4 gen 路径剩余 ragged 缺口 | 未完成 | +| 5 | CUDA graph 兼容 | 见 P0-2/P0-3，**未运行验证** | +| 6 | GSM8K + ADP 验证 | 未开始 | +| 7 | mode 开关 + 可观测性 | 完成 | +| 8 | graph 按 token 分桶 + ADP tier 协商 | 未完成 | +| 9 | SPS cost table profiler | 代码完成（1503 行 + 42 测试），**从未在真实 engine 上跑过** | +| 10 | `on_update_kv_lens` 刷新 row_kv_lens | 完成 | +| 11 | overlap 开启的正确性验证 | 未完成（**default 配置，必须验**） | +| 12-16 | P0-2 … P0-6 | 代码完成，**六条全部零运行验证** | + +**当前阻塞**：环境重建中（见 §6）。六条 P0 一条都没经过运行验证，这是最大的未知。 + +--- + +## 6. 环境（复现必读） + +### 6.1 容器镜像 —— 先看这里 + +**镜像由仓库指定，不要自己挑。** + +``` +jenkins/current_image_tags.properties:16 +LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:\ + pytorch-26.05-py3-x86_64-ubuntu24.04-skip-tritondevel-202607211045-16608 +``` + +本工作前期用错了镜像（`...trt10.16.1.11-...202607151440-16194`，从本地 `docker images` +里挑的），代价是**一整轮关于 cutlass 版本冲突的错误分析**：错误镜像的系统 cutlass 与 +`requirements.txt` 的 pin 不配套，我先归因给自己的 pip 操作，又归因给「repo 与镜像的 +结构性冲突」，两次都错。正确镜像的 cutlass 是自洽的。 + +### 6.2 cutlass 版本 + +正确镜像自带 **4.6.1 且四个子包全部一致**。`requirements.txt` pin `==4.5.0`，但 +`constraints.txt` 只要求 `>=4.4.2`，所以**安装依赖时排除 cutlass、沿用镜像自带的**： + +```bash +grep -v "nvidia-cutlass-dsl" requirements.txt > requirements.nocutlass.txt +pip install -r requirements.nocutlass.txt +``` + +过滤后的文件**必须写在 `requirements.txt` 旁边**——里面有 `-c constraints.txt` 相对路径， +放 `/tmp` 会报 `Could not open constraint file`。 + +为什么不能装 4.5.0：那套包没有 `libs-core`，装进 `--system-site-packages` 的 venv 后， +前端是 4.5.0 而 MLIR verifier 是系统的 4.6.1，DSv4-Pro 会死在 +`'nvgpu.cvt_fptrunc' op operand #0 must be ... 1-d vector, but got 'f32'` +（`f8E8M0FNU` 就是 DSv4 的 ue8m0 scale）。 + +**验证 venv 不能只看 `import` 成功**，要看 `pip list | grep cutlass` 各子包版本是否一致。 + +### 6.3 venv + +```bash +/usr/bin/python3 -m pip install --user 'virtualenv<22.0,>=20.29.1' +/usr/bin/python3 -m virtualenv --system-site-packages .venv-3.12 +``` + +`virtualenv` 而非 `venv`：镜像的 `ensurepip` 不可用。 +`--system-site-packages` 必需：torch 在镜像的 dist-packages 里。 + +### 6.4 编译 + +```bash +python3 scripts/build_wheel.py --cuda_architectures "100-real" -j $(nproc) +``` + +- **PATH 只能追加不能替换**：`/usr/local/cuda/bin` 在镜像默认 PATH 里，覆盖掉会让 + cmake 报 `No CUDA compiler found`。 +- **配置失败后必须 `rm -rf cpp/build`**：CMake 会把失败的编译器探测缓存成 + `CMAKE_CUDA_COMPILER:FILEPATH=NOTFOUND` 并在下次复用，环境修好了也不会重新探测。 +- **目标平台是 B300（SM103），编译用 `--cuda_architectures "103-real"`。** + B200 是 SM100；`100-real` 纯 SASS 无 PTX，装不进 B300。B300 有 275 GB/卡， + B200 只有 178 GB —— DSv4-Pro 权重 104 GB/卡，B200 上余量很紧。 +- 容器需 `--user $(id -u):$(id -g)`（scratch 是 NFS root_squash）。 + +### 6.4b DSv4-Pro 的加载时间结构 —— 复用同一节点的 page cache + +冷启动实测（`umb-b300-001`，权重在 `/home/scratch.trt_llm_data`，NFS）： + +| 阶段 | 耗时 | 可观测性 | +|---|---|---| +| import + 8-rank NCCL 建组 | ~1 min | 有日志 | +| **预取（NFS → 内存）** | **~2 h** | 只在整个 13 GB 分片读完时打一条 `Finished prefetching`，中间静默 90 min 以上 | +| **装配（内存 → 显存）** | **> 40 min** | tqdm 进度条，但 ETA 严重低估（并发流均匀推进、集中收尾） | + +**冷启动总计 3 小时以上，`batch-short` 的 4 小时窗口装不下「加载 + 验证」。** + +**解法：复用同一节点。** 该分区节点有 **2 TB 内存**，一次加载后 page cache 会保留 +约 1.1 TB，足以完整覆盖 835 GB 权重。作业结束后 cache 不会被清掉，所以： + +```bash +# 当前作业还没结束时就把后续作业排到同一节点，接力启动、cache 仍热 +sbatch -A trt-llm_b300 --qos=batch-short \ + -p "b300@ts6/dgx-b300@ts1/8gpu-256cpu-2048gb" \ + --gres=gpu:8 --nodelist=<同一节点> -t 4:0:0 --wrap='sleep 14400' +``` + +查 cache 是否还热：`free -g` 看 `buff/cache` 列。 + +**判活提醒**：预取阶段日志静默、GPU 利用率 0%、显存不变**同时出现是正常的**。 +唯一可靠的判据是固定时间窗口内 `/proc//io` 的 `read_bytes` 增量 +（实测健康值约 112 MB/s）。不要用 tqdm 的 ETA 估算剩余时间。 + +### 6.5 跑测试 / 跑模型 + +容器额外需要：`-e USER=`（否则 `getpass.getuser()` 因 UID 无 passwd 条目而 +KeyError）、`-e PYTHONPATH=/code/tensorrt_llm`（直接跑脚本时；`sys.path[0]` 是脚本所在目录）、 +`-e LLM_MODELS_ROOT=/home/scratch.trt_llm_data/llm-models`、 +`--ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --shm-size=32g`。 + +**srun 步骤内 docker 的 `--gpus` 只能用 `all`**，写 `device=N` 会被代理拒绝。 + +**任何 `tensor_parallel_size > 1` 的脚本必须有 `if __name__ == "__main__":`**—— +`MpiPoolSession` 靠重新执行模块来 spawn worker，没有守卫会无限递归 spawn， +表现为**静默挂死而非报错**（本工作因此浪费 2 小时窗口）。 + +**Slurm**：Blackwell 分区 `DenyQos=batch,interactive-isolated,oversubscribe`，必须 +`--qos=batch-short`（上限 **4 小时**）。报 `Invalid qos specification` 不表示没有权限。 + +--- + +## 7. 教训（给接手的人） + +### 7.0 诊断动作本身会破坏运行 —— 两次杀掉自己的冒烟 + +冒烟从未跑到 ragged 代码，两次都死在加载阶段，**两次都是我杀的**： + +1. **B200**：据「GPU 利用率 0% + 显存不变 + 日志静默」判定挂死并杀掉。 + 事后 `py-spy` 证明它在 `weight_loader.py:338` 正常从 NFS 预取权重。 + 残留进程占着 135 GB/卡不释放，导致下一次启动直接 OOM。 +2. **B300**：为避免重犯，改用 `py-spy dump` 做「正确诊断」——附加超时， + 把 worker 留在 `ptrace_stop` 冻住了整个运行；再用 + `pkill -f "py-spy dump"` 去救，那个模式匹配到 srun 自己的命令行，把 step 杀了。 + `/proc//io` 显示它冻住前已读 13.2 GB，本来是好的。 + +**三条铁律：** + +- **不要 py-spy 附加正在加载的 worker。** 附加开销大到会超时，且会把进程留在 + ptrace 停止态。要用也只在已经确认卡死之后用。 +- **不要 `pkill -f <字符串>`。** 会匹配到 srun / docker 自己的命令行。 + 清理只按 PID：`nvidia-smi --query-compute-apps=pid --format=csv,noheader`。 +- **判活只用 `/proc//io` 的 `read_bytes` 是否增长。** 只读、零副作用， + 是加载阶段唯一安全且可靠的存活信号。 + +**加载阶段的正常表现**：GPU 利用率 0%、日志静默、显存不变——三者同时出现也是正常的。 +DSv4-Pro 权重 835 GB 从共享 NFS 冷读，20-30 分钟属于合理范围。 + +**杀掉运行之后必须清残留**：孤儿进程会一直占着显存，下一次启动就会 OOM。 +`for p in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader); do kill -9 $p; done` + +### 7.1 空输出不是通过 + +本工作中反复栽在这一条，四次： + +1. **kernel 输出顺序**：`indexer_topk_decode` 不保证输出顺序（radix split-work 的 merge + 顺序不定），只有 top-k **集合**是契约。逐元素比数组得到 ~95% 假阳性 mismatch。 + 必须 `sort(dim=-1)` 后比，并用 `torch.topk` 做参照。 +2. **对照实验没跑起来**：一次 baseline 对照用了 `--gpus device=7`，docker 直接 + `invalid gpus device: 7` 退出，空输出被当成了有效数据。 +3. **冒烟挂死**：连续两次报告「仍在加载 835G」，实际 GPU 1-7 各 4 MiB、利用率全 0。 +4. **编译白跑**：CMakeCache 缓存了 NOTFOUND，容器「在运行」但 0 个 object 文件。 + +统一的教训：**「进程在跑」「没有报错」都不是进展的证据**。要有正向信号—— +object 文件数、GPU 显存占用、明确的成功输出。 + +### 7.2 工具缺陷要立刻修，不要记下来等下次 + +`run.sh` 里的 `... | tail -80` 会缓冲全部输出，进程退出前看不到任何进度。 +我在发现后说了「下次改」，没有立即改——结果冒烟静默挂死 2 小时才被发现， +而改成 `python3 -u ... | tee` 之后**一分钟**就看到了完整的 MPI 错误。同样的信息一直都在。 + +### 7.3 局部通过不等于整体正确 + +一度宣告「venv 已修复」，依据只是 `import cutlass` 和 `import tensorrt_llm` 成功。 +实际版本组合并不自洽，DSv4-Pro 起不来。**导入成功只证明模块能加载，不证明版本配套。** + +### 7.4 环境问题先查仓库怎么规定，不要自己拼 + +本工作在环境上的弯路（镜像、cutlass、PATH、CMakeCache、constraints 路径、ensurepip） +**没有一个是产品代码的问题**，全部是容器调用姿势，而每个都吃掉一轮迭代。 +根源是「见招拆招」——需要 virtualenv 就改 PATH，需要过滤 requirements 就写 /tmp， +每次改动都没检查会不会破坏别的东西。 +正确做法是一开始就照 `jenkins/current_image_tags.properties` 和 +`.dspark-logs/build.log` 复现官方姿势。 + +### 7.5 这个功能的所有失效都是静默的 + +planner 拒绝 trim、cost model 退化、batch 错过捕获形状掉出 graph、部分加窗被拒—— +每一种都仍然产出正确输出和基线精度。这是 §3.3 那套可观测性存在的全部理由。 +**新增任何决定 verify 长度的代码路径，都要同时接上 `record_step`**， +否则 `assert_ragged_active()` 会给出虚假的通过。 diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py index 6dd402369d9f..62c6f1b69b7f 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py @@ -221,7 +221,12 @@ def forward( if num_generations > 0: gen_kv_lens = kv_lens[num_contexts:] + # Uniform batches: every request appended next_n tokens. Ragged + # batches: next_n is only the batch maximum (the kernel's loop + # bound) and the real per-request counts come through + # new_tokens_per_seq. next_n = metadata.num_gen_tokens_per_seq + new_tokens_per_seq = getattr(metadata, "gen_new_tokens_per_seq", None) # Pass full kv_score (not sliced) with the generation portion of # cu_seq_lens so the kernel reads at the correct absolute offsets. torch.ops.trtllm.compressor_paged_kv_compress( @@ -240,6 +245,7 @@ def forward( self.head_dim, self.compress_ratio, next_n, + new_tokens_per_seq, ) # Scatter to cache with appropriate quantization (all modes fused) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index 1f76e5a73e51..86aee79dadaf 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -334,6 +334,22 @@ def __post_init__(self): ) self.cu_seq_lens[0] = 0 + # Tokens each generation request appends this step. Only populated under + # ragged verification; a uniform batch leaves it unused and the + # compressor keeps deriving the count from its next_n template bound. + self.gen_new_tokens_per_seq_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences,), + dtype=torch.int, + cache_name="gen_new_tokens_per_seq_cuda", + capture_graph=capture_graph, + ) + self.gen_new_tokens_per_seq_host = torch.empty_like( + self.gen_new_tokens_per_seq_cuda, device="cpu", pin_memory=prefer_pinned() + ) + # Set per step by prepare(); None means "uniform, use next_n". + self.gen_new_tokens_per_seq: Optional[torch.Tensor] = None + # new_comp_kv_lens_cuda is the number of new compressed tokens for the requests self.new_comp_kv_lens_cuda = { compress_ratio: self.get_empty( @@ -729,10 +745,12 @@ def prepare(self): # --- Per-ratio metadata --- # 1) CPU-side: compute scalar metadata (num_total_compressed_tokens, etc.) # 2) CUDA-side: fill *_cuda buffers via prepare_compressed_kv_metadata() - num_gen_tokens_per_seq = ( - num_gen_tokens // self.num_generations if self.num_generations > 0 else 0 - ) - self.num_gen_tokens_per_seq = num_gen_tokens_per_seq + # Under ragged verification requests append different numbers of tokens, + # so the integer division below is meaningless. The compressor kernel + # still needs a compile-time loop bound, and its loops are already + # guarded, so the batch MAXIMUM is the correct scalar to hand it -- the + # exact per-request counts ride along in gen_new_tokens_per_seq. + num_gen_tokens_per_seq = self._sync_gen_tokens_per_seq(num_gen_tokens) num_contexts = self.num_contexts num_generations = self.num_generations kv_lens_slice = kv_lens[:num_requests] @@ -787,6 +805,39 @@ def prepare(self): self._compress_ratios_sorted, ) + def _sync_gen_tokens_per_seq(self, num_gen_tokens: int) -> int: + """Set ``num_gen_tokens_per_seq`` / ``gen_new_tokens_per_seq`` together. + + Uniform batches divide the flat generation-token count by the request + count. Ragged batches have no such quotient, so the scalar becomes the + batch *maximum* -- which is exactly what the compressor kernel wants, + since ``NEXT_N`` is a template loop bound whose iterations are already + guarded -- and the true per-request counts travel beside it on the + device. + + Single writer on purpose: the scalar and the vector describe the same + thing, and every place that recomputed one without the other was a way + for them to disagree. + """ + if self.is_ragged_verify and self.num_generations > 0: + gen_token_counts = self.gen_token_repeat_list() + self.num_gen_tokens_per_seq = max(gen_token_counts) + n = self.num_generations + self.gen_new_tokens_per_seq_host[:n].copy_( + torch.tensor(gen_token_counts, dtype=torch.int32) + ) + gen_new_tokens = self.gen_new_tokens_per_seq_cuda[:n] + gen_new_tokens.copy_( + self.gen_new_tokens_per_seq_host[:n], non_blocking=True + ) + self.gen_new_tokens_per_seq = gen_new_tokens + else: + self.num_gen_tokens_per_seq = ( + num_gen_tokens // self.num_generations if self.num_generations > 0 else 0 + ) + self.gen_new_tokens_per_seq = None + return self.num_gen_tokens_per_seq + def prepare_compressed_kv_metadata( self, kv_lens: torch.Tensor, @@ -862,9 +913,12 @@ def on_update_kv_lens(self): cached_tokens = kv_lens - seq_lens num_gen_tokens = num_tokens - self.num_ctx_tokens - self.num_gen_tokens_per_seq = ( - num_gen_tokens // self.num_generations if self.num_generations > 0 else 0 - ) + # Must go through the same helper as prepare(): a plain integer division + # here would silently undo the ragged batch maximum, and the compressor + # reads this value as its next_n loop bound. A bound below a request's + # real window truncates Phase 1 and the tail KV/score state is never + # written -- no error, just corrupted compression windows from then on. + self._sync_gen_tokens_per_seq(num_gen_tokens) # Reuse prepare()'s host-computed ctx sizes unless the extend_ctx path # (num_chunked_ctx_requests > 0) may have mutated ctx-row kv_lens on diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 92f577ff5d98..38e59a2b1b79 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -897,6 +897,9 @@ def on_update_kv_lens(self): # Per-step state for cross-layer indexer sharing; clear at the step # boundary so a "shared" layer never reuses a stale top-k. self.shared_topk_indices = None + # kv_lens_cuda may have just been corrected; the ragged per-row extent + # is derived from it and has to follow. + self.refresh_ragged_row_kv_lens() if self.kv_cache_manager is not None and self.num_tokens > 0: seq_lens = self.seq_lens_cuda[:self.num_seqs] @@ -1328,6 +1331,61 @@ def create_expanded_buffers(self, capture_graph=False): device='cpu', pin_memory=prefer_pinned(), ) + # Per-row causal KV extent for the decode top-k, in *uncompressed* token + # space (the kernel divides by compress_ratio itself). Only populated on + # ragged batches: the uniform kernel reconstructs the same quantity from + # next_n, and passing nullptr keeps that path bit-identical. + # + # Distinct from kv_lens_expanded_cuda, which carries indexer (possibly + # compressed) lengths for the paged-MQA-logits call. The top-k kernel is + # handed kv_lens_cuda_runtime instead, so mixing the two would apply the + # compression divisor twice. + self.row_kv_lens_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences * (1 + self.max_draft_tokens), ), + cache_name="row_kv_lens_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.row_kv_lens_host = torch.zeros_like( + self.row_kv_lens_cuda, + device='cpu', + pin_memory=prefer_pinned(), + ) + # The per-row offset `-verify_len + position + 1`, kept separately so + # the extent can be rebuilt from a corrected kv_lens without going back + # to the host. It depends only on the verify windows, which do not move + # after prepare(); the kv lengths do. + self.row_kv_correction_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences * (1 + self.max_draft_tokens), ), + cache_name="row_kv_correction_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.row_kv_correction_host = torch.zeros_like( + self.row_kv_correction_cuda, + device='cpu', + pin_memory=prefer_pinned(), + ) + # Which request each query row belongs to. Persistent because the + # refresh below runs inside the CUDA-graph capture body: building this + # per call would record a memcpy from a Python temporary that is freed + # before the first replay. + self.row_req_idx_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences * (1 + self.max_draft_tokens), ), + cache_name="row_req_idx_cuda", + dtype=torch.long, + capture_graph=capture_graph, + ) + self.row_req_idx_host = torch.zeros_like( + self.row_req_idx_cuda, + device='cpu', + pin_memory=prefer_pinned(), + ) + # Rows covered by the current step's ragged layout; 0 on uniform steps. + self._ragged_num_rows = 0 self.block_table_expanded = self.get_empty( self.cuda_graph_buffers, [ @@ -1544,6 +1602,12 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): # - No distinction between sm90 and sm100 is needed once MTP3 is supported on sm90. # - Remove this once fp8_paged_mqa_logits supports an arbitrary number of MTP draft tokens. use_dsl = self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits + if not self.is_ragged_verify: + # A uniform step must not inherit a previous ragged step's row + # count: refresh_ragged_row_kv_lens and ragged_row_kv_lens both key + # off it, and this covers every uniform path, not just the one that + # skips the expanded buffers. + self._ragged_num_rows = 0 if self.is_ragged_verify: # The DSL paged-MQA path reshapes (num_gen, next_n) -> (num_gen * # factor, atom) and broadcasts a single 1-D kv_lens per request, so @@ -1584,6 +1648,9 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): non_blocking=True) self.block_table_expanded.clamp_(min=0) + if self.is_ragged_verify: + self._prepare_ragged_row_kv_lens(kv_lens) + self.expand_for_dsl = (use_dsl and self.kv_cache_manager is not None and self.max_draft_tokens >= 1) if self.expand_for_dsl and self.num_generations > 0: @@ -1627,6 +1694,106 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): self.dsl_expand_factor = 1 self.dsl_atom = 1 + self.max_draft_tokens + def _prepare_ragged_row_kv_lens(self, kv_lens: torch.Tensor) -> None: + """Populate ``row_kv_lens_cuda``: how far back each query row may attend. + + The decode top-k kernel normally reconstructs this from ``next_n``:: + + actual_kv_len = seq_len - next_n + (row % next_n) + 1 + + which is only meaningful while every request contributes the same number + of rows. Under ragged verification neither ``row // next_n`` (which + request) nor ``row % next_n`` (which position in its window) holds, so + the extent is materialized per row instead. + + The generalization is exact: for a request verifying ``v`` positions, + row ``o`` of its window may attend to ``kv_len - v + o + 1`` tokens. + Substituting ``v = next_n`` recovers the formula above, which is why the + uniform path can keep passing ``nullptr`` and stay bit-identical. + + Sync-free by construction: ``kv_lens`` is already a host tensor here + (:meth:`prepare` builds it from ``num_cached_tokens_per_seq``, which + lives on the CPU), so the whole vector is assembled on the host and + crosses once, through pinned memory so the copy does not serialize the + stream. Reading it on the device instead would mean a D2H sync, and + combining the two halves on the device would mean adding a CPU tensor to + a CUDA one. + """ + verify_lens = self.ragged_verify_lens + if not verify_lens: + return + gen_kv_lens = kv_lens[self.num_contexts:self.num_seqs].tolist() + assert len(gen_kv_lens) == len(verify_lens), ( + f"ragged verify lengths {len(verify_lens)} != generation requests " + f"{len(gen_kv_lens)}") + + # Row o of a request verifying v positions sees kv_len - v + o + 1 + # tokens: the last row sees the request's full KV, each earlier row one + # token less. + rows: List[int] = [] + for kv_len, verify_len in zip(gen_kv_lens, verify_lens): + base = int(kv_len) - int(verify_len) + 1 + rows.extend(range(base, base + int(verify_len))) + + corrections: List[int] = [] + req_idx: List[int] = [] + for request_pos, verify_len in enumerate(verify_lens): + corrections.extend(range(1 - int(verify_len), 1)) + req_idx.extend([request_pos] * int(verify_len)) + + num_tokens = len(rows) + self._ragged_num_rows = num_tokens + self.row_kv_lens_host[:num_tokens].copy_( + torch.tensor(rows, dtype=torch.int32)) + self.row_kv_lens_cuda[:num_tokens].copy_( + self.row_kv_lens_host[:num_tokens], non_blocking=True) + self.row_kv_correction_host[:num_tokens].copy_( + torch.tensor(corrections, dtype=torch.int32)) + self.row_kv_correction_cuda[:num_tokens].copy_( + self.row_kv_correction_host[:num_tokens], non_blocking=True) + self.row_req_idx_host[:num_tokens].copy_( + torch.tensor(req_idx, dtype=torch.long)) + self.row_req_idx_cuda[:num_tokens].copy_( + self.row_req_idx_host[:num_tokens], non_blocking=True) + + def refresh_ragged_row_kv_lens(self) -> None: + """Rebuild the per-row extent from the corrected ``kv_lens_cuda``. + + ``prepare()`` builds the extent from the host-side kv lengths, but under + the overlap scheduler ``_preprocess_inputs`` corrects ``kv_lens_cuda`` + afterwards to account for how many draft tokens were actually accepted. + Leaving the extent at its prepare-time value would make the top-k kernel + and ``gen_kv_lens_cuda`` -- both arguments of the same op call -- disagree + about the same quantity. + + Stays on the device: the correction term was staged at prepare time and + does not move, so this is one gather plus one add, with no sync. + """ + num_tokens = self._ragged_num_rows + if not self.is_ragged_verify or num_tokens <= 0: + return + gen_kv_lens = self.kv_lens_cuda[self.num_contexts:self.num_seqs] + # Gather, not repeat_interleave: this runs inside the capture body, and + # a tensor-repeats interleave both allocates and (without output_size) + # reads the cumulative sum back to the host. The row->request map and + # the correction were staged at prepare() time and are pinned by the + # graph key, so all that is left is one gather plus one add. + row_kv_lens = self.row_kv_lens_cuda[:num_tokens] + torch.index_select(gen_kv_lens, 0, + self.row_req_idx_cuda[:num_tokens], + out=row_kv_lens) + row_kv_lens.add_(self.row_kv_correction_cuda[:num_tokens]) + + def ragged_row_kv_lens(self, num_tokens: int) -> Optional[torch.Tensor]: + """The per-row extent slice for this step, or None on a uniform batch.""" + if not self.is_ragged_verify: + return None + assert num_tokens == self._ragged_num_rows, ( + f"ragged row count moved between prepare ({self._ragged_num_rows}) " + f"and forward ({num_tokens}); the per-row extents would be read " + f"past what was populated") + return self.row_kv_lens_cuda[:num_tokens] + def prepare_for_indexer_k_cache(self): # Build indexer_k_cache_block_offsets using pool block indices derived # from host_kv_cache_block_offsets (populated by super().prepare()). @@ -2980,7 +3147,32 @@ def sparse_attn_indexer( metadata.heuristic_scratch_values[ :num_gen_tokens] - if self.use_cute_dsl_topk and self._enable_heuristic_topk: + # Ragged batches carry a per-row causal extent instead of + # deriving one from next_n. Both CuTe DSL top-k kernels still + # take next_n as a compile-time-ish scalar and reconstruct the + # row->request map from it, and GVR additionally indexes its + # previous-step hint the same way, so neither can express a + # ragged batch. Route to the C++ kernel, which now accepts the + # extent directly. + ragged_row_kv_lens = metadata.ragged_row_kv_lens(num_gen_tokens) + if ragged_row_kv_lens is not None: + torch.ops.trtllm.indexer_topk_decode( + logits_decode, + gen_kv_lens_cuda, + topk_indices_buffer[num_ctx_tokens:num_ctx_tokens + + num_gen_tokens, :], + # Dead on this path: the kernel reads the extent out of + # row_kv_lens rather than rebuilding it. Pass 1 so a + # stale next_n cannot quietly come back into use. + 1, + self.index_topk, + pre_idx=None, + heuristic_scratch=None, + compress_ratio=self.compress_ratio, + radix_aux_indices=metadata.radix_aux_indices, + radix_aux_logits=metadata.radix_aux_logits, + row_kv_lens=ragged_row_kv_lens) + elif self.use_cute_dsl_topk and self._enable_heuristic_topk: # GVR DSL: supports all compress_ratio and next_n values. torch.ops.trtllm.cute_dsl_gvr_topk_decode( logits_decode, diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 8781a9481a7a..c3a787ec7468 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -449,7 +449,18 @@ def capture(self, # We're forced to do this because we cannot reallocate inputs over many graph runs. max_draft_len = key[1] token_per_request = max_draft_len + 1 - num_tokens_for_capture = (batch_size * self.max_beam_width * + # Ragged keys carry the flat token total as a sixth element (see + # _get_graph_key). It is exactly the width the captured graph will be + # replayed at, and it is NOT batch_size * (draft_len + 1): the block is + # still drafted in full, so draft_len stays at the top tier while the + # verified total shrinks. Sizing the static slices from the product + # instead would capture a graph whose input_ids are wider than the + # attention metadata's num_tokens on every trimmed replay -- and the + # widest bucket happens to match the product, which is what warmup + # captures, so it would capture clean and diverge on the first real + # trimmed step. + num_tokens_for_capture = (key[5] if len(key) > 5 else + batch_size * self.max_beam_width * token_per_request) sliced_static_tensors = { diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index aa856ff7ec3b..5296494e439e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3624,6 +3624,23 @@ def ragged_verify_token_buckets(self, padded_bs: int) -> List[int]: }) return [int(padded_bs) * (t + 1) for t in tiers] + def _record_dspark_graph_use(self, *, replayed: bool) -> None: + """Count graph replay vs eager for the DSpark ragged stats, if enabled. + + Deliberately tolerant: this runs on the hot path for every model, and a + non-DSpark engine has no worker, no planner and no stats object. Guard + rather than branch on config so the cost is one attribute lookup. + """ + stats = self._dspark_ragged_stats + if stats is not None: + stats.record_graph(replayed=replayed) + + @property + def _dspark_ragged_stats(self): + """The DSpark worker's stats object, or None on every other path.""" + worker = self._get_spec_worker() + return getattr(worker, "ragged_stats", None) if worker else None + def fit_ragged_verify_lens( self, generation_requests, @@ -3766,6 +3783,47 @@ def fit_ragged_verify_lens( request.py_verify_len = int(tokens) - 1 return int(bucket) + @staticmethod + def _ragged_token_lens(generation_requests) -> Optional[List[int]]: + """Each generation request's token window, or None if the batch is uniform. + + ``py_verify_len`` counts drafted positions; a request also verifies its + bonus position, so its token window is one wider. + + Returns None unless *every* request carries a window. A partially + windowed batch is worse than an unwindowed one: the flat token layout + would go ragged while the spec metadata stayed uniform, which is silent + token misattribution rather than a crash. + """ + verify_lens = [ + getattr(request, "py_verify_len", None) + for request in generation_requests + ] + if not verify_lens or any(v is None for v in verify_lens): + return None + return [1 + int(v) for v in verify_lens] + + def _publish_ragged_verify_lens(self, attn_metadata, + generation_requests) -> None: + """Hand the attention metadata this step's windows, before it prepares. + + Split out of :meth:`_attach_ragged_verify_layout` because of ordering: + ``attn_metadata.prepare()`` is the *only* consumer of + ``ragged_verify_lens`` -- it drives the DSA expanded-buffer decision, the + per-row causal extents, and DeepSeek-V4's per-request compressor token + counts -- and it runs well before the spec metadata is assembled. + Publishing both together meant ``prepare()`` read the previous step's + split every time, and None on the first ragged step. + + The spec-metadata half stays where it was; it is consumed later and has + no such constraint. + """ + if attn_metadata is None or not hasattr(attn_metadata, + "ragged_verify_lens"): + return + attn_metadata.ragged_verify_lens = self._ragged_token_lens( + generation_requests) + def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, generation_requests) -> None: """Publish this step's per-request verify windows to both metadatas. @@ -3777,11 +3835,8 @@ def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, right place, and the DSA indexer has to expand kv_lens/block tables one row per query token instead of a fixed count per request. """ - verify_lens = [ - getattr(request, "py_verify_len", None) - for request in generation_requests - ] - if not verify_lens or any(v is None for v in verify_lens): + token_lens = self._ragged_token_lens(generation_requests) + if token_lens is None: spec_metadata.verify_lens = None spec_metadata.qo_indptr = None spec_metadata.total_verify_tokens = None @@ -3792,9 +3847,6 @@ def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, from ..speculative.dspark_ragged import build_qo_indptr - # py_verify_len counts draft positions; a request also verifies its - # bonus position, so its token window is one wider. - token_lens = [1 + int(v) for v in verify_lens] n = len(token_lens) # Write through the persistent buffers rather than binding a fresh # tensor: a captured graph baked in the address it saw at capture time, @@ -4606,13 +4658,24 @@ def append_cross_attention_state(request: LlmRequest, if next_draft_tokens_device is None or request.is_dummy or request.py_batch_idx is None: # get token ids, including input token ids and draft token ids. For these dummy requests, # no need to copy the token ids. + # Under ragged verification the block is still drafted in full, + # but only the request's own window is submitted to the target. + # This branch runs whenever the overlap scheduler is off -- i.e. + # on every step of the accuracy configuration -- so taking the + # full drafted length here would leave the flat token layout + # uniform while the spec metadata and the CUDA-graph key both + # believe it is ragged: silent token misattribution rather than + # a crash. Mirrors the overlap branch below, which already + # derives every length from the same per-request value. + num_draft_tokens = get_request_tokens_per_gen_step( + request, 1 + get_draft_token_length(request)) - 1 if not (request.is_attention_dp_dummy or request.is_cuda_graph_dummy): input_ids.append(request.get_last_tokens(0)) - input_ids.extend(request.py_draft_tokens) - draft_tokens.extend(request.py_draft_tokens) + input_ids.extend(request.py_draft_tokens[:num_draft_tokens]) + draft_tokens.extend( + request.py_draft_tokens[:num_draft_tokens]) # get other ids and lengths - num_draft_tokens = get_draft_token_length(request) past_seen_token_num = request.max_beam_num_tokens - 1 draft_lens.append(num_draft_tokens) if self.enable_spec_decode and spec_config.spec_dec_mode.extend_ctx( @@ -5269,7 +5332,23 @@ def previous_seq_slots_device(): helix_is_inactive_rank=helix_is_inactive_rank, ) - if not attn_metadata.is_cuda_graph: + # Ragged verification is the one case where the per-request lengths + # genuinely move between replays of the same graph: the row count and + # the padded token total are both pinned by the key, but how the tokens + # are split across rows is not. Leaving the captured ones(bs)*(1+top) + # in place would make attn_metadata.num_tokens (== seq_lens.sum(), + # interface.py) disagree with the input_ids width on every trimmed step + # -- and that sum is what attention-DP all-gathers and the MoE chunks + # from. The setter already copies in place under CUDA graphs (see its + # comment: "The seqlens can change if we are doing spec decode"), so + # this reallocates nothing as long as the row count is unchanged. + refresh_seq_lens = not attn_metadata.is_cuda_graph + if (not refresh_seq_lens + and getattr(attn_metadata, "is_ragged_verify", False) + and attn_metadata.seq_lens is not None + and len(sequence_lengths) == attn_metadata.seq_lens.shape[0]): + refresh_seq_lens = True + if refresh_seq_lens: # Assumes seq lens do not change between CUDA graph invocations. This applies # to draft sequences too. This means that all draft sequences must be padded. attn_metadata.seq_lens = torch.tensor( @@ -5323,6 +5402,13 @@ def previous_seq_slots_device(): # pre-prepare counts so the steady-gen recording below stores values # that the per-step prepare() can re-clamp from scratch. num_cached_tokens_snapshot = list(num_cached_tokens_per_seq) + # prepare() is the only consumer of the ragged windows, so they have to + # be published first: it decides the DSA expanded-buffer layout, builds + # the per-row causal extents, and derives DeepSeek-V4's per-request + # compressor token counts. The rest of the layout is attached later, + # once the spec metadata exists. + self._publish_ragged_verify_lens(attn_metadata, + scheduled_requests.generation_requests) attn_metadata.prepare() cross_attention_inputs = (self._prepare_enc_dec_cross_attn_inputs( cross_encoder_hidden_states, @@ -6679,6 +6765,11 @@ def forward(self, ) can_run_graph = key is not None + # A ragged step that misses its captured shape does not raise -- it + # quietly runs eager, which costs far more than the tokens it + # trimmed. Record it so the miss is a counter rather than something + # only a profiler would reveal. + self._record_dspark_graph_use(replayed=can_run_graph) if can_run_graph: attn_metadata = maybe_attn_metadata spec_metadata = maybe_spec_metadata diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4488c566e7df..f99a2f805dab 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3235,13 +3235,27 @@ def _dspark_confidence_draft_len(self, def all_rank_max(value: int) -> int: return max(int(v) for v in self.dist.tp_allgather(int(value))) + # The mode, not the config flag, decides whether windows are computed: + # the environment override exists so a run can be switched between + # `static` and `compact` without editing the serving config, and every + # counter below is keyed off the same decision. + stats = getattr(worker, "ragged_stats", None) + mode = getattr(worker, "ragged_verify_mode", None) + compute_windows = (mode.computes_windows if mode is not None else + getattr(self.model_engine.spec_config, + "enable_ragged_verify", False)) + ragged_lens = None - if getattr(self.model_engine.spec_config, "enable_ragged_verify", - False): + fallback_reason = None + if compute_windows: ragged_lens = planner.decide_verify_lens( num_gen_requests=num_gen_requests, rows=confidence_rows, all_rank_max=all_rank_max) + if ragged_lens is None: + fallback_reason = "planner_declined" + else: + fallback_reason = "mode_static" # Land the batch's token total on a captured bucket before the graph key # is built, spending the rounding slack on real verification. A batch @@ -3261,6 +3275,33 @@ def all_rank_max(value: int) -> int: sum(1 + int(v) for v in ragged_lens)]) bucket = self.model_engine.fit_ragged_verify_lens( gen_requests, ragged_lens, peer_stats=peer_stats) + if bucket is None: + fallback_reason = "no_captured_shape" + + # Record what the step actually decided, reading the windows back off + # the requests rather than off `ragged_lens`: fit_ragged_verify_lens + # spends the bucket's rounding slack on real verification, so the + # fitted windows are what the target will see. + if stats is not None: + # The stats object lives in the worker process; under TP the caller + # that wants to assert on it is on the far side of an MPI boundary + # and cannot reach it. Emit a periodic summary so the decision is + # recoverable from the log, which is the only shared channel. + if stats.steps_total and stats.steps_total % 32 == 0: + stats.log_summary(prefix="DSpark ragged verify [periodic]") + fitted = None + if bucket is not None: + fitted = [ + int(request.py_verify_len) for request in gen_requests + if getattr(request, "py_verify_len", None) is not None + ] + if len(fitted) != num_gen_requests: + fitted = None + stats.record_step(num_gen_requests=num_gen_requests, + verify_lens=fitted, + bucket=bucket, + padded_bs=None, + fallback=fallback_reason) if bucket is not None: # The block is always drafted in full -- only verification is diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index ee7d7ae46eae..8311e43f0dcd 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -454,10 +454,46 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> tiers=tiers, apply_calibration=apply_calibration, ) + + # The config knob says what was asked for; the environment overrides it + # so a mode can be swapped without editing a serving config. `static` + # unless ragged was explicitly requested, matching the shipped default. + from .dspark_observability import (DSparkRaggedStats, RaggedVerifyMode, + read_ragged_verify_mode) + + configured = (RaggedVerifyMode.COMPACT if getattr( + self.spec_config, "enable_ragged_verify", False) else + RaggedVerifyMode.STATIC) + self.ragged_verify_mode = read_ragged_verify_mode(default=configured) + if self.ragged_verify_mode is RaggedVerifyMode.CAP_ACCEPT: + # Deliberately not aliased onto `compact`. The whole value of + # cap-accept is that its output must be bit-identical to `static` + # while the windows still drive commitment; silently running it as + # `compact` would make that comparison meaningless. + raise NotImplementedError( + "TLLM_DSPARK_RAGGED_VERIFY_MODE='cap-accept' is not implemented " + "yet: it requires submitting the full block to the target while " + "committing only each request's window. Use 'static' or " + "'compact'.") + self.ragged_stats = DSparkRaggedStats(mode=self.ragged_verify_mode, + max_draft_len=block_size) logger.info( f"DSpark verify planner: tiers={self.verify_planner.tiers}, " - f"profiled_cost_table={not cost_table.is_flat}" + f"profiled_cost_table={not cost_table.is_flat}, " + f"ragged_verify_mode={self.ragged_verify_mode.value}" ) + if (self.ragged_verify_mode.trims_submitted_tokens + and cost_table.is_flat): + # Not a warning about a tuning choice -- in this combination the + # feature cannot do anything at all, and the run would look like a + # successful `compact` run in every log line. + logger.warning( + "DSpark ragged verify is set to 'compact' but the cost table is " + "flat, so the planner's budget degenerates to verify-all and " + "every request will receive the same full window. Profile a " + "cost table and pass confidence_sps_table_path, or the ragged " + "path will not be exercised." + ) def _assign_slot(self, req_id: int, reset: bool) -> int: """Get (or refresh) the slot for a request; reset clears its window.""" @@ -584,9 +620,21 @@ def _draft_gen_block_batched( # Captured target hidden at the bonus position within each request's Kp1 # processed tokens. - arange_g = torch.arange(num_gens, device=device) - base = gen_start + arange_g * Kp1 # [G] - main_hidden = captured[base + gidx] # [G, ncap*hidden] + # Where each request's processed tokens start inside `captured`, which + # is indexed by flat token position. A uniform batch strides by Kp1; + # under ragged verification requests contribute different numbers of + # tokens, so the offsets come from the same qo_indptr the input layout + # was packed with. Striding by Kp1 there would have request r read from + # request r-1's tail -- no error, just a silent acceptance collapse. + qo_indptr = getattr(spec_metadata, "qo_indptr", None) + if getattr(spec_metadata, "verify_lens", None) is not None and qo_indptr is not None: + base = gen_start + qo_indptr[:num_gens].to(device=device, dtype=torch.long) + else: + arange_g = torch.arange(num_gens, device=device) + base = gen_start + arange_g * Kp1 # [G] + # Clamped like interim_base below: gidx comes from a device-side accept + # count, so a bad one must not turn into an out-of-bounds read. + main_hidden = captured[(base + gidx).clamp(min=0, max=captured.shape[0] - 1)] # Fixed-size ([G, K]) masked back-fill of the intermediate accepted tokens # (everything but the bonus) into the rolling window — same frames as the diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py new file mode 100644 index 000000000000..42adc46c8725 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -0,0 +1,360 @@ +# 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. +"""Verification mode selection and per-step observability for DSpark. + +Why this module exists +---------------------- +Every failure mode of confidence-scheduled verification is *silent*. The +scheduler can decline to trim, the planner can fall back on a flat cost model, +a batch can miss its captured shape and drop out of graph replay, and a +partially-windowed batch can be refused -- and in every one of those cases the +run still produces correct output at the baseline accuracy. A GSM8K pass +therefore proves nothing about whether the feature was ever active. + +So the two questions a reviewer actually has -- *did the ragged path run?* and +*did it trim anything?* -- are made into counters rather than left to +inference, and :meth:`DSparkRaggedStats.assert_ragged_active` turns them into a +test assertion. + +Verification modes +------------------ +Mirrors SGLang's ``SGLANG_RAGGED_VERIFY_MODE``, and for the same reason: it +separates "is the ragged layout computed correctly" from "does the planner +choose to trim", which otherwise fail together and are indistinguishable. + +``static`` + Verify the whole drafted block for every request. The baseline, and what + every non-DSpark path already does. Scheduling is off. + +``cap-accept`` + Compute the per-request windows and carry them through the whole layout, + but still submit the full block to the target and only *commit* the window. + The expensive path is unchanged, so this buys no throughput -- it exists + because its output must be **bit-identical** to ``static``. Any divergence + is a scheduling bug isolated from the kernel-side packing, which is the one + comparison that separates the two. + +``compact`` + Submit only each request's window. The production path, and the only one + that trims tokens. + +Cost-table dependence +--------------------- +``compact`` needs a profiled :class:`~.dspark_planner.SpsCostTable`. Without +one the planner's cost model is flat, every extra verified token looks free, +and the budget degenerates to verify-all -- ragged silently becomes uniform. +:meth:`DSparkRaggedStats.assert_ragged_active` catches exactly that. +""" + +import os +from collections import Counter +from enum import Enum +from typing import Dict, List, Optional, Sequence + +from ...logger import logger + +__all__ = [ + "RaggedVerifyMode", + "read_ragged_verify_mode", + "DSparkRaggedStats", + "forced_verify_lens", +] + +#: Overrides ``DSparkDecodingConfig`` so a mode can be selected without editing +#: a test or a serving config. Reading it per call (rather than caching at +#: import) is deliberate: tests flip it with ``monkeypatch.setenv``. +RAGGED_VERIFY_MODE_ENV = "TLLM_DSPARK_RAGGED_VERIFY_MODE" + + +class RaggedVerifyMode(str, Enum): + """How much of the drafted block reaches the target. See module docstring.""" + + STATIC = "static" + CAP_ACCEPT = "cap-accept" + COMPACT = "compact" + + @property + def computes_windows(self) -> bool: + """Whether per-request windows are computed at all. + + ``cap-accept`` computes them and commits against them without trimming + the submitted batch, so it shares every host-side code path with + ``compact`` except the one that shrinks the token axis. + """ + return self is not RaggedVerifyMode.STATIC + + @property + def trims_submitted_tokens(self) -> bool: + """Whether the token axis handed to the target actually shrinks.""" + return self is RaggedVerifyMode.COMPACT + + +def read_ragged_verify_mode( + default: RaggedVerifyMode = RaggedVerifyMode.STATIC +) -> RaggedVerifyMode: + """Resolve the verification mode from the environment. + + Raises: + ValueError: if the variable is set to something that is not a mode. + Silently falling back to ``static`` on a typo would turn "I asked + for compact" into "the feature never ran", which is the one + outcome this module exists to make impossible. + """ + value = os.environ.get(RAGGED_VERIFY_MODE_ENV) + if value is None or value == "": + return default + for mode in RaggedVerifyMode: + if value == mode.value: + return mode + raise ValueError( + f"invalid {RAGGED_VERIFY_MODE_ENV}={value!r}; expected one of " + f"{', '.join(repr(m.value) for m in RaggedVerifyMode)}") + + +#: Debug-only override that forces a deterministic non-uniform split. See +#: :func:`forced_verify_lens`. +FORCE_VERIFY_LENS_ENV = "TLLM_DSPARK_FORCE_VERIFY_LENS" + + +def forced_verify_lens(*, num_gen_requests: int, tiers: Sequence[int], + min_verify_len: int) -> Optional[List[int]]: + """A deterministic ragged split for correctness testing, or None. + + Set ``TLLM_DSPARK_FORCE_VERIFY_LENS=1`` to have every generation request + take a window from the captured tier ladder, rotating by batch position. + + This exists because the two failure modes are otherwise inseparable. The + planner refuses to trim without a profiled cost table -- correctly, since a + flat model makes every extra token look free -- so a run without one keeps + the ragged path dark no matter what the config says. Testing that the packing + is correct would then require first producing a cost table, which requires a + working run. This breaks that circle. + + It is not a shortcut around the planner: it decides only *how many* drafted + positions reach the target, exactly the quantity the planner decides. + Acceptance is unchanged, so a forced split must produce the same answers as + ``static`` -- which is what makes it a usable correctness gate. + + Returns: + One window per request, each drawn from ``tiers`` (so every choice has a + captured CUDA graph), or None when the override is off. + """ + if os.environ.get(FORCE_VERIFY_LENS_ENV, "") not in ("1", "true", "True"): + return None + ladder = sorted({int(t) for t in tiers if int(t) >= int(min_verify_len)}) + if not ladder: + return None + # Rotate through the ladder by position: deterministic, identical on every + # rank for a given batch position, and guaranteed to be non-uniform as soon + # as the batch is larger than one and the ladder has more than one entry. + return [ladder[i % len(ladder)] for i in range(int(num_gen_requests))] + + +class DSparkRaggedStats: + """Per-step counters for the ragged verification path. + + Cheap enough to leave on: one integer histogram update per decode step, no + device work and no synchronization. Nothing here reads a device tensor -- + every value recorded is already a host value by the time the scheduler has + made its decision. + """ + + def __init__(self, *, mode: RaggedVerifyMode, max_draft_len: int): + self.mode = mode + self.max_draft_len = int(max_draft_len) + + self.steps_total = 0 + #: Steps where per-request windows were produced AND differed from each + #: other. A step where every request got the same window is counted as + #: uniform: it exercises none of the ragged packing. + self.steps_ragged = 0 + #: Windows were produced, but every request got the same one. + self.steps_uniform_windows = 0 + #: No windows at all -- the planner declined or the shape did not fit. + self.steps_no_windows = 0 + + #: Tokens the target would have scored with no trimming at all. + self.ceiling_tokens = 0 + #: Tokens the windows asked for, before bucket rounding. + self.window_tokens = 0 + #: Tokens actually submitted, after rounding up to a captured bucket. + self.delivered_tokens = 0 + + self.verify_len_hist: Counter = Counter() + self.bucket_hist: Counter = Counter() + self.padded_bs_hist: Counter = Counter() + + #: Graph replay is the whole point of landing on a bucket; an eager + #: step costs far more than the tokens it saved. + self.graph_replays = 0 + self.graph_eager = 0 + + #: Reasons the step declined to go ragged, keyed by a short slug. + self.fallbacks: Counter = Counter() + + def record_step( + self, + *, + num_gen_requests: int, + verify_lens: Optional[Sequence[int]] = None, + bucket: Optional[int] = None, + padded_bs: Optional[int] = None, + fallback: Optional[str] = None, + ) -> None: + """Record one decode step's scheduling decision. + + Args: + num_gen_requests: generation requests in the batch. + verify_lens: per-request drafted positions verified, or None when + the step stayed uniform. Excludes the bonus token, matching + ``py_verify_len``. + bucket: the captured token bucket the batch was fitted to, if any. + padded_bs: the row count the batch was padded to, if any. + fallback: short reason the step did not go ragged. + """ + self.steps_total += 1 + # What a no-trim step would have submitted: every request sends its + # bonus token plus the full block. + self.ceiling_tokens += int(num_gen_requests) * (1 + self.max_draft_len) + + if fallback: + self.fallbacks[fallback] += 1 + + if not verify_lens: + self.steps_no_windows += 1 + self.window_tokens += int(num_gen_requests) * (1 + + self.max_draft_len) + self.delivered_tokens += int(num_gen_requests) * ( + 1 + self.max_draft_len) + return + + lens = [int(v) for v in verify_lens] + self.verify_len_hist.update(lens) + self.window_tokens += sum(1 + v for v in lens) + # A bucket is what the graph actually ran; without one the step ran on + # exactly the windows it asked for. + self.delivered_tokens += int(bucket) if bucket else sum(1 + v + for v in lens) + if bucket: + self.bucket_hist[int(bucket)] += 1 + if padded_bs: + self.padded_bs_hist[int(padded_bs)] += 1 + + if len(set(lens)) > 1: + self.steps_ragged += 1 + else: + self.steps_uniform_windows += 1 + + def record_graph(self, *, replayed: bool) -> None: + """Record whether this step replayed a CUDA graph or ran eager.""" + if replayed: + self.graph_replays += 1 + else: + self.graph_eager += 1 + + def merge_planner_stats(self, planner_stats: Dict[str, int]) -> None: + """Fold in the planner's own fallback counters.""" + for key, value in planner_stats.items(): + if key.startswith("fallback_") and value: + self.fallbacks[key] = int(value) + + @property + def trim_ratio(self) -> float: + """Fraction of ceiling tokens the scheduler actually removed. + + ``0.0`` means the feature ran but saved nothing, which is the signature + of an unprofiled cost table. + """ + if self.ceiling_tokens <= 0: + return 0.0 + return 1.0 - (self.delivered_tokens / self.ceiling_tokens) + + @property + def distinct_verify_lens(self) -> int: + """How many different window sizes were ever handed out.""" + return len(self.verify_len_hist) + + def summary(self) -> Dict[str, object]: + """A JSON-safe snapshot, for logging or assertions.""" + return { + "mode": self.mode.value, + "steps_total": self.steps_total, + "steps_ragged": self.steps_ragged, + "steps_uniform_windows": self.steps_uniform_windows, + "steps_no_windows": self.steps_no_windows, + "ceiling_tokens": self.ceiling_tokens, + "window_tokens": self.window_tokens, + "delivered_tokens": self.delivered_tokens, + "trim_ratio": round(self.trim_ratio, 4), + "distinct_verify_lens": self.distinct_verify_lens, + "verify_len_hist": dict(sorted(self.verify_len_hist.items())), + "bucket_hist": dict(sorted(self.bucket_hist.items())), + "padded_bs_hist": dict(sorted(self.padded_bs_hist.items())), + "graph_replays": self.graph_replays, + "graph_eager": self.graph_eager, + "fallbacks": dict(self.fallbacks), + } + + def log_summary(self, *, prefix: str = "DSpark ragged verify") -> None: + summary = self.summary() + logger.info(f"{prefix}: {summary}") + + def assert_ragged_active(self, *, require_trim: bool = True) -> None: + """Fail loudly if the ragged path never actually ran. + + This is the check that makes an accuracy run meaningful. Passing GSM8K + with ``compact`` selected proves nothing on its own -- the scheduler + may have declined on every step and the run may have been a plain + uniform baseline wearing a different config. + + Args: + require_trim: also require that the delivered token count came in + under the no-trim ceiling. Leave it on unless deliberately + measuring a workload where every draft survives. + + Raises: + AssertionError: with the full counter summary, so the reason is in + the failure rather than in a log the CI dropped. + """ + summary = self.summary() + if self.steps_total == 0: + raise AssertionError( + f"DSpark ragged verify recorded no decode steps at all: {summary}" + ) + if self.steps_ragged == 0: + raise AssertionError( + f"DSpark ragged verify never produced differing per-request " + f"windows, so the ragged path was not exercised. The usual " + f"cause is a missing profiled SPS cost table, under which the " + f"planner's budget degenerates to verify-all: {summary}") + if self.distinct_verify_lens < 2: + raise AssertionError( + f"DSpark ragged verify handed out only one distinct window " + f"size, so the batch was uniform in all but name: {summary}") + if require_trim and self.trim_ratio <= 0.0: + raise AssertionError( + f"DSpark ragged verify delivered as many tokens as a no-trim " + f"run, so nothing was saved: {summary}") + if self.graph_eager > 0: + raise AssertionError( + f"DSpark ragged verify dropped out of CUDA graph replay on " + f"{self.graph_eager} step(s); a ragged step without a captured " + f"shape costs far more than the tokens it trims: {summary}") + + +def format_verify_len_histogram(lens: List[int]) -> str: + """One-line ``len:count`` rendering, for per-step debug logging.""" + hist = Counter(int(v) for v in lens) + return " ".join(f"{k}:{v}" for k, v in sorted(hist.items())) diff --git a/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py b/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py new file mode 100644 index 000000000000..a5140507116f --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py @@ -0,0 +1,1503 @@ +# 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. +"""Measure the DSpark SPS cost table this deployment actually has. + +Why this exists +--------------- +:class:`~.dspark_planner.SpsCostTable` is the only thing that tells the verify +planner a long verification costs more than a short one. Without a measured +table ``SpsCostTable.flat()`` is installed, ``is_flat`` is true, and +``DSparkVerifyPlanner.decide_verify_lens`` returns ``None`` on every step +(``stats["fallback_flat_cost"]``) -- ragged verification is then a no-op that +still produces correct output at baseline accuracy, i.e. it fails silently. +``TLLM_DSPARK_FORCE_VERIFY_LENS=1`` unblocks *correctness* testing by bypassing +that gate, but it rotates a fixed ladder and knows nothing about cost, so it +cannot answer "is trimming profitable here". Only a measured table can. + +Run it:: + + python -m tensorrt_llm._torch.speculative.dspark_sps_profiler \\ + --model /models/DeepSeek-V4-Pro-DSpark --tp-size 8 --ep-size 8 \\ + --enable-attention-dp --batch-sizes 8,16,32,64,128 \\ + --max-draft-len 5 --out sps_table.json + +and then point the serving config at the result:: + + speculative_config: + decoding_type: DSpark + enable_confidence_scheduling: true + confidence_sps_table_path: sps_table.json + confidence_verify_len_tiers: [...] # from _meta, see "Tiers" below + +What the sweep measures, and how the three terms are separated +-------------------------------------------------------------- +The planner's cost model is additive:: + + T(bs, M) = bias + alpha(bs) + theta(M), M = bs * (verify_len + 1) + +``theta(M)`` is the *only* term trimming reduces; ``bias`` and ``alpha(bs)`` are +paid whatever the planner picks. Because the planner maximizes ``tau / T``, a +**ratio**, the non-trimmable terms do not cancel: drop them and every verified +token looks proportionally more expensive than it is, and the planner +over-trims. So all three have to be measured, not guessed. + +The sweep is a cross product of batch size ``bs`` and *uniform* verify length +``L``. Each cell is one whole steady-state decode step time ``S(bs, L)``, and +sits at ``M = bs * (L + 1)`` (:func:`~.dspark_planner.total_verify_tokens` -- +the bonus token is included; getting that wrong is a whole ``bs`` of error, +easily a shelf's width). The cross product is what makes the terms separable: + +* **theta(M) is isolated along a row.** Fix ``bs`` and move ``L``: ``bias`` and + ``alpha(bs)`` are identical for every cell in the row, so the *differences* + along it are pure ``theta``. +* **alpha(bs) is isolated along a column.** The grid supplies same-``M``, + different-``bs`` collisions for free -- ``(bs, L)`` and ``(2*bs, L')`` land on + the same ``M`` whenever ``L + 1 == 2 * (L' + 1)``, e.g. ``(16, 3)`` and + ``(32, 1)`` both give ``M = 64``. Across such a collision ``theta`` cancels + and the difference is pure ``alpha``. + This is also why the batch-size ladder must be *connected* through those + collisions; :func:`fit_additive_cost_model` refuses a sweep that is not, + because each disconnected component would carry its own arbitrary offset and + the emitted ``theta`` staircase would splice two different scales together. +* **The remaining global constant is fixed by extrapolating theta to M -> 0.** + An additive fit only determines ``theta`` and ``alpha`` up to one shared + constant. We shift so that ``theta`` linearly extrapolates to zero at + ``M = 0``, which is exactly the definition that makes ``bias + alpha(bs)`` + the step's ``M``-independent floor -- the quantity the planner needs. +* **bias vs alpha(bs) is then a presentation split.** Only their sum enters + :meth:`SpsCostTable.step_time`, so we report ``bias = min_bs (bias + + alpha(bs))`` -- the floor present even at the smallest profiled batch -- and + give the rest to ``alpha``. Both are emitted so that a reader can see which + part is batch-driven. + +The emitted table is therefore in the **decomposed** encoding: ``step_time_ms`` +is ``theta(M)`` *only*. Do not also read it as whole-step time and re-add the +overheads -- that double-counts the floor, inflates ``T``, flattens the relative +differences between candidate lengths and drags the argmax back to verify-all. + +Holding verification uniform +---------------------------- +The run pins ``TLLM_DSPARK_RAGGED_VERIFY_MODE=static`` and clears +``TLLM_DSPARK_FORCE_VERIFY_LENS``. Under ``static`` the planner still chooses a +single batch-wide length, but the cost table it sees is the flat default, whose +``is_flat`` gate short-circuits to ``max_tier`` *before* any confidence is read +-- so the length is deterministically ``max_draft_len`` on every step and +``M = bs * (max_draft_len + 1)`` exactly. If the scheduler were allowed to trim +mid-sweep, each cell would be measuring a different ``M`` than the one it is +filed under, i.e. a moving target. + +That determinism is also why the ``L`` axis costs one engine build per value: +nothing in the runtime pins a *uniform* verify length below the full block +(``TLLM_DSPARK_FORCE_VERIFY_LENS`` produces a ragged rotation, and there is no +analogue of SGLang's ``dspark_force_budget_frac``), so ``L`` is swept by +rebuilding with ``max_draft_len = block_size = L``. Use ``--samples-out`` / +``--fit-only`` to split a long sweep across jobs and refit without a GPU. + +**Known bias of that choice:** shrinking the block also shrinks the *draft* +pass, whose cost is roughly linear in the block length, while at deployment the +block is always drafted in full and only verification is trimmed. The draft +model is a few layers against a full-size target, so the leak is a small +fraction of the measured slope -- but it is a positive fraction, so the emitted +``theta`` slope is mildly overstated and the planner is biased towards +*over*-trimming. Prefer profiling only the lengths the ladder can actually +select, and treat a marginal trim decision as noise. + +Acceptance is pinned +-------------------- +``--pin-acceptance`` (on by default) sets +``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` to a tiny positive value, which +makes every request commit exactly its bonus token every step. ``0.0`` means +"override disabled", hence the epsilon rather than a clean zero. This matters +because acceptance drives how fast the KV cache grows, and KV length drives +attention cost: without pinning, a long-block cell would accumulate context +faster than a short-block one and the difference would be misattributed to +``theta(M)``. + +Statistics +---------- +Only *aligned* steps count: a step where every attention-DP rank reports the +same generation-request count and no context requests. The first +``--warmup-steps`` aligned steps of each cell are discarded (allocator growth, +graph warmup, autotuner) and the cell reports the **median** of the rest -- +never the mean. Decode step time has a long right tail from allocator and +scheduling noise, and a mean turns that tail into a fake riser, which is exactly +the failure a staircase cost model amplifies. + +Refusals +-------- +The tool would rather emit nothing than emit a table that looks profiled and +is not: + +* :class:`FlatCostTableError` -- the curve does not move, so ``is_flat`` is true + (or the total step time varies by less than the noise floor). The planner + would treat the file as unprofiled and keep verifying everything. +* :class:`InertCostTableError` -- the curve moves, but not between any two tiers + the planner may select at any profiled batch size, so the argmax can never + leave ``max_tier``. That is a non-flat table that trims nothing, and unlike a + flat one it does *not* trip ``fallback_flat_cost``, so it would look healthy. + +Tiers +----- +``_meta.recommended_confidence_verify_len_tiers`` is derived from the measured +shelves at the modal batch size, but **nothing reads it from this file**: +``DSparkDecodingConfig.verify_len_tiers`` always returns a non-empty list, so +``dspark.py``'s ``derive_verify_len_tiers`` fallback never runs. Copy the value +into ``confidence_verify_len_tiers`` by hand. The config is the right single +source anyway -- ``_get_graphs_to_capture`` and ``ragged_verify_token_buckets`` +read the same property, so a ladder that lives only here would let the planner +pick lengths that have no captured CUDA graph and silently drop the step into +eager execution. +""" + +from __future__ import annotations + +import argparse +import json +import os +import random +import statistics +import sys +import time +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Sequence, Tuple + +import numpy as np + +from .dspark_observability import FORCE_VERIFY_LENS_ENV, RAGGED_VERIFY_MODE_ENV, RaggedVerifyMode +from .dspark_planner import ( + SpsCostTable, + budget_argmax_over_uniform_lens, + derive_verify_len_tiers, + total_verify_tokens, +) + +__all__ = [ + "AdditiveFit", + "CellStat", + "FlatCostTableError", + "InertCostTableError", + "InsufficientSamplesError", + "ProfilerError", + "StepSample", + "SweepConfig", + "SweepGeometryError", + "aligned_steps_from_stats", + "build_cost_table_payload", + "check_table_is_informative", + "compress_to_risers", + "fit_additive_cost_model", + "load_cost_table", + "main", + "profitability_probe", + "run_sweep", + "running_max", + "summarize_cells", +] + +#: Per-iteration timing fields published by ``PyExecutor._profiler``. The host +#: field is the default because it is a clean single-loop CPU wall on every +#: scheduler; ``iterLatencyMS`` is deliberately not offered, since under the +#: overlap scheduler (the shipped default) it spans ~2 loops. +HOST_STEP_TIME_KEY = "hostStepTimeMS" +GPU_FORWARD_TIME_KEY = "gpuForwardTimeMS" +TIMING_FIELDS = {"host": HOST_STEP_TIME_KEY, "gpu_forward": GPU_FORWARD_TIME_KEY} + +#: ``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` treats 0.0 as "override off", +#: so pinning to bonus-token-only needs a value that floors to 0 without being 0. +#: The fractional part is the probability of one extra token, so 1e-6 fires +#: roughly never while still selecting the override path. +BONUS_TOKEN_ONLY_ACCEPTANCE = 1e-6 + +#: ``SpsCostTable.__post_init__`` rejects a non-positive ``step_time_ms``. After +#: normalizing theta to hit zero at M = 0 the smallest shelf can land at or below +#: zero, so it is clamped to this instead. +MIN_THETA_MS = 1e-4 + +#: How far the additive model may miss a cell, relative to a typical step, +#: before the fit says so. Not a repair threshold -- see +#: :func:`fit_additive_cost_model` for why the misfit is reported, not fixed. +DEFAULT_ADDITIVE_TOLERANCE = 0.02 + + +class ProfilerError(RuntimeError): + """Base class for every refusal this module raises.""" + + +class InsufficientSamplesError(ProfilerError): + """A cell did not produce enough aligned steady-state steps to trust.""" + + +class SweepGeometryError(ProfilerError): + """The measured cells cannot separate ``alpha(bs)`` from ``theta(M)``.""" + + +class FlatCostTableError(ProfilerError): + """The measured curve is flat, so the table would read as unprofiled.""" + + +class InertCostTableError(ProfilerError): + """The curve moves, but never between two selectable verify lengths.""" + + +# --------------------------------------------------------------------------- +# raw samples and per-cell statistics +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class StepSample: + """One aligned decode step, filed under the cell that produced it.""" + + batch_size: int + verify_len: int + step_time_ms: float + iteration: int = -1 + + @property + def total_verify_tokens(self) -> int: + """``M`` for this step. Derived, never stored, so it cannot drift.""" + return total_verify_tokens(self.batch_size, self.verify_len) + + def to_json(self) -> dict: + return { + "batch_size": int(self.batch_size), + "verify_len": int(self.verify_len), + "step_time_ms": float(self.step_time_ms), + "iteration": int(self.iteration), + } + + @classmethod + def from_json(cls, raw: dict) -> "StepSample": + return cls( + batch_size=int(raw["batch_size"]), + verify_len=int(raw["verify_len"]), + step_time_ms=float(raw["step_time_ms"]), + iteration=int(raw.get("iteration", -1)), + ) + + +@dataclass(frozen=True) +class CellStat: + """The summarized cost of one ``(batch_size, verify_len)`` cell.""" + + batch_size: int + verify_len: int + step_time_ms: float + num_samples: int + p10_ms: float + p90_ms: float + + @property + def total_verify_tokens(self) -> int: + return total_verify_tokens(self.batch_size, self.verify_len) + + def to_json(self) -> dict: + return { + "batch_size": int(self.batch_size), + "verify_len": int(self.verify_len), + "total_verify_tokens": int(self.total_verify_tokens), + "step_time_ms": round(float(self.step_time_ms), 6), + "num_samples": int(self.num_samples), + "p10_ms": round(float(self.p10_ms), 6), + "p90_ms": round(float(self.p90_ms), 6), + } + + +def aligned_steps_from_stats( + rows: Sequence[dict], + *, + expected_ranks: int = 1, + timing_key: str = HOST_STEP_TIME_KEY, +) -> List[Tuple[int, int, float]]: + """Pull ``(iteration, num_gen_requests, step_time_ms)`` out of ``get_stats``. + + A step is kept only when it is *aligned*: every expected attention-DP rank + published a row for that iteration, no rank had context requests in flight, + and all ranks scheduled the same number of generation requests. Anything + else is a step whose shape is not the cell's shape -- a prefill iteration, a + ramp step, a rank that fell behind -- and averaging it in is how a cost + table acquires risers that are not there. + + The timing is taken from the rank-0 row rather than averaged across ranks: + under attention-DP :mod:`~..pyexecutor.adp_iter_stats` fans out one record + per rank with *rank-local counters* but *rank-0's* step time, so the + per-rank timings are copies, not independent measurements. + """ + by_iteration: Dict[int, List[dict]] = {} + for row in rows: + iteration = row.get("iter") + if iteration is None: + continue + by_iteration.setdefault(int(iteration), []).append(row) + + aligned: List[Tuple[int, int, float]] = [] + for iteration in sorted(by_iteration): + group = sorted(by_iteration[iteration], key=lambda r: int(r.get("attentionDpRank", 0))) + if len({int(r.get("attentionDpRank", 0)) for r in group}) != int(expected_ranks): + continue + batching = [row.get("inflightBatchingStats") or {} for row in group] + if any(int(stats.get("numContextRequests", 0) or 0) for stats in batching): + continue + gen_counts = {int(stats.get("numGenRequests", 0) or 0) for stats in batching} + if len(gen_counts) != 1: + continue + num_gen_requests = gen_counts.pop() + if num_gen_requests <= 0: + continue + # gpuForwardTimeMS is read without a device sync, so it is absent on the + # steps whose event pair had not landed yet; those steps are simply not + # usable for that field. + timing = group[0].get(timing_key) + if timing is None or float(timing) <= 0.0: + continue + aligned.append((iteration, num_gen_requests, float(timing))) + return aligned + + +def summarize_cells( + samples: Sequence[StepSample], + *, + warmup_steps: int, + min_samples: int, +) -> List[CellStat]: + """Median per ``(batch_size, verify_len)`` cell, after dropping warmup. + + Median, not mean: a decode step's distribution is a tight body with a long + right tail (allocator growth, host scheduling hiccups, the occasional + stats/log flush). The mean tracks the tail, and a staircase cost model reads + a tail-inflated point as a riser -- i.e. as a place worth trimming before. + + Args: + warmup_steps: leading aligned steps to discard *per cell*. They are kept + in the raw sample file so a refit can choose a different cut. + min_samples: post-warmup samples a cell must have to be admitted. + + Raises: + InsufficientSamplesError: if any cell falls short. Silently dropping a + cell would punch a hole in the grid and can disconnect the + batch-size ladder, which changes the fit rather than just widening + its error bars. + """ + grouped: Dict[Tuple[int, int], List[StepSample]] = {} + for sample in samples: + grouped.setdefault((int(sample.batch_size), int(sample.verify_len)), []).append(sample) + + cells: List[CellStat] = [] + for (batch_size, verify_len), cell_samples in sorted(grouped.items()): + ordered = sorted(cell_samples, key=lambda s: s.iteration) + steady = [s.step_time_ms for s in ordered[int(warmup_steps) :]] + if len(steady) < int(min_samples): + raise InsufficientSamplesError( + f"cell (batch_size={batch_size}, verify_len={verify_len}) has " + f"{len(steady)} steady samples after discarding " + f"{warmup_steps} warmup steps, need {min_samples}. Either the " + f"cell never reached a steady batch of {batch_size} generating " + f"requests (check that the batch size is captured by the CUDA " + f"graph config and that enough requests were in flight), or the " + f"measurement window was too short -- raise --measure-steps." + ) + values = np.asarray(steady, dtype=np.float64) + cells.append( + CellStat( + batch_size=batch_size, + verify_len=verify_len, + step_time_ms=float(np.median(values)), + num_samples=int(values.size), + p10_ms=float(np.percentile(values, 10)), + p90_ms=float(np.percentile(values, 90)), + ) + ) + return cells + + +# --------------------------------------------------------------------------- +# additive fit: T(bs, M) = bias + alpha(bs) + theta(M) +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class AdditiveFit: + """The separated cost terms, before they are shaped into a staircase.""" + + #: ``bias + alpha(bs)`` per profiled batch size, normalized so that theta + #: extrapolates to zero at ``M = 0``. + intercept_ms: Dict[int, float] + #: ``theta(M)`` on the measured token counts, non-decreasing. + theta_ms: Dict[int, float] + #: Largest ``|S - intercept - theta|`` over the cells, absolute and relative + #: to the median cell time. This is the additive model's own error bar: if + #: it is large, ``theta`` is not a function of ``M`` alone and the + #: decomposed encoding is not safe at batch sizes far from those profiled. + max_abs_residual_ms: float + max_rel_residual: float + warnings: Tuple[str, ...] = () + + +def _assert_grid_is_connected(cells: Sequence[CellStat]) -> None: + """Refuse a sweep whose ``(bs, M)`` bipartite graph has several components. + + Within one component the additive split is determined up to a single shared + constant, which the caller fixes. Across components each has its *own* + unknown constant, so concatenating their ``theta`` values into one staircase + silently splices two different scales -- the resulting risers are artifacts + of the offset, not of the hardware. + """ + parent: Dict[Tuple[str, int], Tuple[str, int]] = {} + + def find(node): + parent.setdefault(node, node) + while parent[node] != node: + parent[node] = parent[parent[node]] + node = parent[node] + return node + + for cell in cells: + a, b = find(("bs", cell.batch_size)), find(("m", cell.total_verify_tokens)) + if a != b: + parent[a] = b + + components = {find(("bs", cell.batch_size)) for cell in cells} + if len(components) > 1: + groups: Dict[Tuple[str, int], List[int]] = {} + for cell in cells: + groups.setdefault(find(("bs", cell.batch_size)), []).append(cell.batch_size) + islands = sorted(sorted(set(v)) for v in groups.values()) + raise SweepGeometryError( + f"the swept cells split into {len(components)} disconnected groups of " + f"batch sizes {islands}: no two groups share a total-token count, so " + f"alpha(bs) and theta(M) cannot be put on a common scale. Two batch " + f"sizes meet when bs1*(L1+1) == bs2*(L2+1) for swept lengths, e.g. a " + f"doubling ladder (8,16,32,...) with verify lengths spanning a factor " + f"of two (1..5 gives token/request ratios 2..6). Add the missing " + f"batch sizes or widen --verify-lens." + ) + + +def running_max(values: Sequence[float]) -> List[float]: + """Cheapest possible isotonic fit: enforce non-decreasing in place. + + ``SpsCostTable`` does not check monotonicity, and a noise-induced dip makes + a *longer* verification look cheaper than a shorter one -- the planner will + take that offer. Clamping up rather than down is the conservative direction: + it can only make the planner trim less. + """ + out: List[float] = [] + highest = -float("inf") + for value in values: + highest = max(highest, float(value)) + out.append(highest) + return out + + +def _least_squares_intercept(x: Sequence[float], y: Sequence[float]) -> float: + """Value of the best-fit line ``y = a + b*x`` at ``x = 0``.""" + xs = np.asarray(x, dtype=np.float64) + ys = np.asarray(y, dtype=np.float64) + if xs.size < 2 or np.ptp(xs) == 0.0: + return float(ys.min()) + slope = float(np.cov(xs, ys, bias=True)[0, 1] / np.var(xs)) + return float(ys.mean() - slope * xs.mean()) + + +def _solve_additive( + cells: Sequence[CellStat], +) -> Tuple[Dict[int, float], Dict[int, float], np.ndarray]: + """Least-squares solve of ``S = intercept[bs] + theta[M]`` over the cells. + + The design is rank-deficient by exactly one -- adding a constant to every + ``theta`` and subtracting it from every ``intercept`` leaves every prediction + unchanged -- so ``lstsq`` returns the minimum-norm representative and the + caller pins the constant afterwards. + + Alternating medians were tried here and are wrong: the grid is unbalanced + (each token count is reached by a different subset of batch sizes), so the + per-column median correction differs from column to column and the iteration + settles on a curve that is *not* the additive truth even when the data is + exactly additive. Robustness is recovered by trimming outliers and refitting + instead; see :func:`fit_additive_cost_model`. + """ + batch_index = {bs: i for i, bs in enumerate(sorted({c.batch_size for c in cells}))} + token_index = {m: i for i, m in enumerate(sorted({c.total_verify_tokens for c in cells}))} + num_batches = len(batch_index) + design = np.zeros((len(cells), num_batches + len(token_index)), dtype=np.float64) + observed = np.empty(len(cells), dtype=np.float64) + for row, cell in enumerate(cells): + design[row, batch_index[cell.batch_size]] = 1.0 + design[row, num_batches + token_index[cell.total_verify_tokens]] = 1.0 + observed[row] = cell.step_time_ms + + solution, *_ = np.linalg.lstsq(design, observed, rcond=None) + residuals = observed - design @ solution + intercept = {bs: float(solution[i]) for bs, i in batch_index.items()} + theta = {m: float(solution[num_batches + i]) for m, i in token_index.items()} + return intercept, theta, residuals + + +def fit_additive_cost_model( + cells: Sequence[CellStat], + *, + additive_tolerance: float = DEFAULT_ADDITIVE_TOLERANCE, +) -> AdditiveFit: + """Separate ``bias + alpha(bs)`` from ``theta(M)``, and report the misfit. + + A badly measured cell is deliberately *not* repaired. The residual of the + additive solve is the model-validity check the study calls for: cells that + share an ``M`` but come from different ``(bs, L)`` -- e.g. ``(32, 3)`` and + ``(16, 7)`` -- must agree once their intercepts are removed, and when they + do not, ``theta`` is not a function of ``M`` alone. Silently dropping the + disagreeing cell would delete exactly that evidence and leave a table that + fits beautifully and predicts badly. So the misfit is surfaced as a warning + and the human decides whether to trust the decomposition. + + The returned split is normalized so that ``theta`` extrapolates to zero at + ``M = 0``; see the module docstring for why that is the split the planner + needs rather than an arbitrary one. + """ + if not cells: + raise SweepGeometryError("no measured cells to fit") + _assert_grid_is_connected(cells) + + warnings: List[str] = [] + intercept, theta, residuals = _solve_additive(cells) + # Kept from the raw solve: this is the additive model's own error, and the + # shaping below (monotone clamp, floor clamp) is a deliberate deviation from + # the data that has its own warnings. Folding the two together would let a + # clamp masquerade as a model failure and vice versa. + typical = statistics.median(cell.step_time_ms for cell in cells) + max_abs = float(np.max(np.abs(residuals))) if len(residuals) else 0.0 + max_rel = max_abs / typical if typical > 0 else 0.0 + if max_rel > float(additive_tolerance): + worst = ", ".join( + f"(bs={cells[i].batch_size}, L={cells[i].verify_len}): {residuals[i]:+.3f} ms" + for i in np.argsort(-np.abs(residuals))[:4] + ) + warnings.append( + f"the additive model misses at least one cell by {max_abs:.3f} ms " + f"({max_rel:.1%} of a typical step), above the " + f"{float(additive_tolerance):.1%} tolerance. Largest disagreements: " + f"{worst}. A least-squares solve spreads one bad measurement over " + f"every cell sharing its batch size or token count, so read these as " + f"a neighbourhood rather than as the culprit. Either one of those " + f"steps did not replay a CUDA graph, or theta is not a function of M " + f"alone -- in which case the decomposed encoding is unsafe away from " + f"the profiled batch sizes." + ) + + ordered_tokens = sorted(theta) + monotone = running_max([theta[m] for m in ordered_tokens]) + if any(abs(monotone[i] - theta[m]) > 1e-9 for i, m in enumerate(ordered_tokens)): + warnings.append( + "theta(M) was not monotonically increasing and has been clamped " + "upward; a dip means the measurement noise is comparable to the " + "risers, so treat the finer shelves as unreliable." + ) + theta = dict(zip(ordered_tokens, monotone)) + + # Fix the one free constant: push theta's M -> 0 extrapolation onto the + # intercept, so what stays behind in `intercept` is the M-independent floor. + shift = _least_squares_intercept(ordered_tokens, [theta[m] for m in ordered_tokens]) + # ...but never so far that a shelf goes non-positive; SpsCostTable rejects + # that, and a zero-cost shelf would read as "these tokens are free". + shift = min(shift, min(theta.values()) - MIN_THETA_MS) + theta = {m: v - shift for m, v in theta.items()} + intercept = {bs: v + shift for bs, v in intercept.items()} + + if min(intercept.values()) < 0.0: + warnings.append( + "the fitted M-independent floor came out negative at the smallest " + "profiled batch size and has been clamped to zero; theta is then " + "explaining more than the whole step, which means the additive " + "model does not hold over this sweep." + ) + intercept = {bs: max(v, 0.0) for bs, v in intercept.items()} + + return AdditiveFit( + intercept_ms=dict(sorted(intercept.items())), + theta_ms=dict(sorted(theta.items())), + max_abs_residual_ms=max_abs, + max_rel_residual=max_rel, + warnings=tuple(warnings), + ) + + +def compress_to_risers( + token_counts: Sequence[int], + step_times: Sequence[float], + *, + min_riser_ms: float, + max_breakpoints: int, +) -> Tuple[List[int], List[float]]: + """Collapse a sampled curve into the staircase it is approximating. + + ``SpsCostTable`` looks up by floor and interpolates nothing, so every + breakpoint claims "a new kernel wave starts here". Shipping one breakpoint + per *sample* instead of one per *riser* asserts risers that do not exist, + and (if the dead ``derive_verify_len_tiers`` path is ever revived) explodes + the tier count so the truncation drops the real shelf edges. + + A point is kept only when it clears the previous kept shelf by + ``min_riser_ms``. The threshold is **absolute**, not relative to ``theta``: + ``theta`` is normalized to approach zero at ``M = 0``, so at the low end a + fraction of a millisecond of noise is a large *relative* move and a relative + threshold would manufacture shelves exactly where the curve carries least + information. What matters is whether a riser is worth a measurable fraction + of a whole step, which is what the caller scales this by. + + If more points survive than ``max_breakpoints``, the shallowest risers are + dropped first -- they are the ones a wrong call costs least. + """ + if len(token_counts) != len(step_times): + raise ValueError("token_counts and step_times must have the same length") + if not token_counts: + raise ValueError("nothing to compress") + + kept = [0] + for index in range(1, len(token_counts)): + if step_times[index] - step_times[kept[-1]] > float(min_riser_ms): + kept.append(index) + + limit = max(int(max_breakpoints), 2) + while len(kept) > limit: + # Only interior/trailing points are droppable: index 0 defines the base + # shelf that every below-range query clamps onto. + jumps = [(step_times[kept[i]] - step_times[kept[i - 1]], i) for i in range(1, len(kept))] + _, drop = min(jumps) + kept.pop(drop) + + return [int(token_counts[i]) for i in kept], [float(step_times[i]) for i in kept] + + +# --------------------------------------------------------------------------- +# table assembly, loading and refusals +# --------------------------------------------------------------------------- + + +def load_cost_table(payload: dict) -> SpsCostTable: + """Build the table exactly the way ``dspark.py::_build_verify_planner`` does. + + Kept byte-for-byte in step with the loader so this module can round-trip its + own output through the real validation instead of a copy of it. + """ + return SpsCostTable( + token_counts=[int(v) for v in payload["token_counts"]], + step_time_ms=[float(v) for v in payload["step_time_ms"]], + fixed_overhead_ms=float(payload.get("fixed_overhead_ms", 0.0)), + batch_sizes=[int(v) for v in payload.get("batch_sizes", [])], + batch_overhead_ms=[float(v) for v in payload.get("batch_overhead_ms", [])], + ) + + +def build_cost_table_payload( + cells: Sequence[CellStat], + *, + riser_tolerance: float = 0.01, + max_breakpoints: int = 8, + additive_tolerance: float = DEFAULT_ADDITIVE_TOLERANCE, + meta: Optional[dict] = None, +) -> dict: + """Assemble the JSON object ``confidence_sps_table_path`` expects. + + Emits exactly the five keys the loader reads plus a ``_meta`` block; unknown + keys are ignored by the loader, so provenance rides along for free. + + No leading ``token_counts = 0`` entry is synthesized. It would be dead + weight: :meth:`SpsCostTable.step_time` clamps a below-range query onto index + 0 already, so a fabricated base shelf could only either duplicate the first + measured shelf or extrapolate below everything that was measured. + """ + fit = fit_additive_cost_model(cells, additive_tolerance=additive_tolerance) + + # A riser has to be worth a measurable slice of a *whole step* to be real, + # so the noise floor is scaled by the typical step rather than by theta, + # whose own scale is an artifact of the M -> 0 normalization. + typical_step_ms = statistics.median(cell.step_time_ms for cell in cells) + min_riser_ms = float(riser_tolerance) * typical_step_ms + tokens, thetas = compress_to_risers( + sorted(fit.theta_ms), + [fit.theta_ms[m] for m in sorted(fit.theta_ms)], + min_riser_ms=min_riser_ms, + max_breakpoints=max_breakpoints, + ) + + batch_sizes = sorted(fit.intercept_ms) + intercepts = running_max([fit.intercept_ms[bs] for bs in batch_sizes]) + # Only the sum bias + alpha(bs) reaches the arithmetic, so the split is a + # reporting choice: bias is the floor that is present even at the smallest + # profiled batch, alpha is what each larger batch adds on top of it. + bias = min(intercepts) if intercepts else 0.0 + alphas = [value - bias for value in intercepts] + + payload = { + "token_counts": [int(v) for v in tokens], + "step_time_ms": [float(round(v, 6)) for v in thetas], + "fixed_overhead_ms": float(round(bias, 6)), + "batch_sizes": [int(v) for v in batch_sizes], + "batch_overhead_ms": [float(round(v, 6)) for v in alphas], + } + payload["_meta"] = { + "encoding": "decomposed", + "note": ( + "step_time_ms is theta(M) only -- the trimmable term. The " + "non-trimmable floor is fixed_overhead_ms + batch_overhead_ms[bs]. " + "Do not read step_time_ms as whole-step time." + ), + "token_unit": ( + "total_verify_tokens = num_requests * (verify_len + 1); the bonus token is included" + ), + "cells": [ + cell.to_json() for cell in sorted(cells, key=lambda c: (c.batch_size, c.verify_len)) + ], + "fit": { + "max_abs_residual_ms": round(fit.max_abs_residual_ms, 6), + "max_rel_residual": round(fit.max_rel_residual, 6), + "additive_tolerance": float(additive_tolerance), + "riser_tolerance": float(riser_tolerance), + "min_riser_ms": round(min_riser_ms, 6), + "typical_step_ms": round(typical_step_ms, 6), + "warnings": list(fit.warnings), + }, + **(meta or {}), + } + # Constructing it here means a malformed table fails in the profiler rather + # than at engine init on a machine that has already spent minutes loading + # weights. + load_cost_table(payload) + return payload + + +def profitability_probe( + payload: dict, + *, + batch_size: int, + tiers: Sequence[int], + acceptance_rates: Sequence[float], +) -> List[dict]: + """Which verify length the real planner picks, per synthetic acceptance rate. + + Runs :func:`~.dspark_planner.budget_argmax_over_uniform_lens` -- the actual + production argmax, not a re-derivation -- against a geometric survival curve + ``p**(j+1)``. This is the question the table exists to answer and the one + ``TLLM_DSPARK_FORCE_VERIFY_LENS`` cannot: at what acceptance rate does + trimming become the better trade, and does it ever? + """ + table = load_cost_table(payload) + ladder = sorted({int(t) for t in tiers}) + block = max(ladder) + results: List[dict] = [] + for rate in acceptance_rates: + survival = np.tile( + np.array([float(rate) ** (j + 1) for j in range(block)], dtype=np.float64), + (int(batch_size), 1), + ) + chosen = budget_argmax_over_uniform_lens( + survival=survival, + num_gen_requests=int(batch_size), + cost_table=table, + allowed_lens=ladder, + min_verify_len=1, + ) + results.append( + { + "batch_size": int(batch_size), + "acceptance_rate": float(rate), + "chosen_verify_len": int(chosen), + "max_tier": int(ladder[-1]), + "trims": bool(int(chosen) < ladder[-1]), + } + ) + return results + + +def check_table_is_informative( + payload: dict, + *, + batch_sizes: Sequence[int], + tiers: Sequence[int], + acceptance_rates: Sequence[float] = (0.3, 0.5, 0.7, 0.8, 0.9, 0.95), + min_step_time_spread: float = 0.02, + allow_inert: bool = False, +) -> dict: + """Refuse to ship a table that cannot change the planner's mind. + + Two distinct failures, both of which produce a healthy-looking run: + + * *flat* -- ``is_flat`` is exact float equality, so it only catches a table + nobody measured. A measured-but-motionless curve passes it and then makes + ``decide_verify_lens`` return ``None`` on every step anyway. So the spread + of the **total** step time across the reachable token range is checked + too, not the spread of theta: theta is normalized towards zero at + ``M = 0``, which would make its own relative spread meaninglessly large. + * *inert* -- the curve moves, but every selectable tier at every profiled + batch size lands on one shelf. ``Theta = tau / T`` is then strictly + increasing in length and the argmax is always ``max_tier``. Unlike the + flat case this does *not* increment ``fallback_flat_cost``, so nothing + downstream would report it. + """ + table = load_cost_table(payload) + ladder = sorted({int(t) for t in tiers}) + if len(ladder) < 2: + raise InertCostTableError( + f"the verify-length ladder {ladder} has a single entry, so the " + f"planner has nothing to choose between and no cost table can " + f"change that." + ) + + if table.is_flat: + raise FlatCostTableError( + f"the profiled step_time_ms {payload['step_time_ms']} has a single " + f"distinct value, so SpsCostTable.is_flat is True. The planner " + f"treats that as 'unprofiled': decide_verify_lens returns None on " + f"every step (fallback_flat_cost) and ragged verification is a " + f"no-op. Either the sweep never crossed a cost riser -- widen it " + f"with larger batch sizes, since gains only exist at high " + f"concurrency -- or the step cost genuinely does not depend on the " + f"verified token count here, in which case confidence scheduling " + f"has no headroom on this configuration and should stay off." + ) + + spreads = {} + for batch_size in sorted({int(b) for b in batch_sizes}): + low = table.step_time(total_verify_tokens(batch_size, ladder[0]), batch_size) + high = table.step_time(total_verify_tokens(batch_size, ladder[-1]), batch_size) + spreads[batch_size] = (high - low) / low if low > 0 else 0.0 + best_spread = max(spreads.values(), default=0.0) + if best_spread < float(min_step_time_spread): + rendered = ", ".join(f"bs={k}: {v:.3%}" for k, v in sorted(spreads.items())) + raise FlatCostTableError( + f"the total step time only varies by {best_spread:.3%} across the " + f"whole tier ladder {ladder} at the best profiled batch size " + f"({rendered}), which is at or " + f"below the measurement noise floor of " + f"{float(min_step_time_spread):.3%}. Such a table is technically " + f"non-flat but the risers it asserts are noise, and the planner " + f"would trim on them. To be worth trimming, the step must grow by " + f"more than the acceptance yield given up -- typically >15% between " + f"adjacent tiers." + ) + + probes: List[dict] = [] + for batch_size in sorted({int(b) for b in batch_sizes}): + probes.extend( + profitability_probe( + payload, + batch_size=batch_size, + tiers=ladder, + acceptance_rates=acceptance_rates, + ) + ) + trims_anywhere = any(probe["trims"] for probe in probes) + if not trims_anywhere and not allow_inert: + raise InertCostTableError( + f"the table is non-flat but inert: over batch sizes " + f"{sorted({int(b) for b in batch_sizes})}, tiers {ladder} and " + f"acceptance rates {list(acceptance_rates)}, the planner picks " + f"max_tier={ladder[-1]} every time, so it would never trim a single " + f"token. That happens when no cost riser falls strictly inside " + f"(bs*(t_i+1), bs*(t_i+1 +1)] for any adjacent tier pair -- the " + f"tiers all sit on one shelf. Fix the ladder so its steps straddle " + f"the measured risers (see _meta." + f"recommended_confidence_verify_len_tiers), or accept that trimming " + f"cannot pay here. Pass --allow-inert to emit anyway." + ) + + return { + "step_time_spread_by_batch_size": {k: round(v, 6) for k, v in spreads.items()}, + "best_step_time_spread": round(best_spread, 6), + "profitability_probe": probes, + "trims_somewhere": trims_anywhere, + } + + +# --------------------------------------------------------------------------- +# the sweep +# --------------------------------------------------------------------------- + + +@dataclass +class SweepConfig: + """Everything the measurement loop needs; mirrors the CLI one-to-one.""" + + model: str + speculative_model: str + batch_sizes: List[int] + verify_lens: List[int] + tp_size: int = 1 + ep_size: Optional[int] = None + enable_attention_dp: bool = False + max_draft_len: int = 5 + input_len: int = 1024 + warmup_steps: int = 16 + measure_steps: int = 64 + min_samples: int = 16 + max_seq_len: int = 4096 + max_num_tokens: int = 8192 + kv_cache_fraction: float = 0.5 + attn_backend: str = "TRTLLM" + disable_overlap_scheduler: bool = False + #: On by default because the deployment being profiled runs with it on, and + #: the confidence head is real per-step work that belongs in the table. + enable_confidence_scheduling: bool = True + timing_key: str = HOST_STEP_TIME_KEY + pin_acceptance: float = BONUS_TOKEN_ONLY_ACCEPTANCE + #: How long the post-cell stats drain waits for the tail records to arrive + #: over IPC before it decides the queue is empty. + stats_timeout_s: float = 5.0 + token_id_range: Tuple[int, int] = (1000, 10000) + seed: int = 1234 + extra_llm_api_options: Dict[str, object] = field(default_factory=dict) + + @property + def dp_size(self) -> int: + """Independent scheduling domains, i.e. rows per iteration in the stats.""" + return int(self.tp_size) if self.enable_attention_dp else 1 + + def max_tokens_for(self, verify_len: int) -> int: + """Output-token budget that guarantees the cell reaches its step count. + + With acceptance pinned every step commits exactly one token, so the + budget *is* the step count. Unpinned, a step commits up to + ``verify_len + 1``, so the budget is scaled to keep the lower bound on + steps at the target -- at the price of running up to that factor longer + when acceptance turns out to be poor. + """ + target_steps = int(self.warmup_steps) + int(self.measure_steps) + budget = target_steps + max(target_steps // 4, 16) + return budget if self.pin_acceptance else budget * (int(verify_len) + 1) + + def validate(self) -> None: + """Catch the geometry mistakes that would otherwise fail mid-sweep.""" + if not self.batch_sizes or not self.verify_lens: + raise SweepGeometryError("both --batch-sizes and --verify-lens must be non-empty") + if min(self.verify_lens) < 1: + raise SweepGeometryError("verify lengths must be >= 1 (the planner's floor)") + # A request that hits max_seq_len is evicted mid-cell, so the batch + # drains and the steady window shrinks -- which surfaces much later as a + # confusing "not enough steady samples". + longest = self.input_len + max(self.max_tokens_for(v) for v in self.verify_lens) + if longest > self.max_seq_len: + raise SweepGeometryError( + f"a request would reach {longest} tokens (input {self.input_len} + " + f"generated) but --max-seq-len is {self.max_seq_len}; raise it or " + f"lower --input-len / --measure-steps" + ) + if self.input_len > self.max_num_tokens: + raise SweepGeometryError( + f"--input-len {self.input_len} exceeds --max-num-tokens " + f"{self.max_num_tokens}, so a prompt cannot be scheduled" + ) + widest_decode = max(self.batch_sizes) * (max(self.verify_lens) + 1) + if widest_decode > self.max_num_tokens: + raise SweepGeometryError( + f"the widest decode step submits {widest_decode} tokens " + f"(bs={max(self.batch_sizes)} x verify_len={max(self.verify_lens)} + 1) " + f"but --max-num-tokens is {self.max_num_tokens}; the batch would be " + f"split and no step would ever have the shape it is filed under" + ) + + +def _prepare_environment(config: SweepConfig) -> None: + """Pin the knobs the table's meaning depends on, loudly. + + Set before the LLM is constructed because the worker processes inherit the + environment at spawn and read these once. + """ + if os.environ.get(FORCE_VERIFY_LENS_ENV): + print( + f"[dspark-sps] clearing {FORCE_VERIFY_LENS_ENV}: it rotates a ragged " + f"split across the batch, so M would no longer be bs*(L+1) and every " + f"cell would be filed under the wrong token count.", + file=sys.stderr, + ) + os.environ.pop(FORCE_VERIFY_LENS_ENV, None) + os.environ[RAGGED_VERIFY_MODE_ENV] = RaggedVerifyMode.STATIC.value + if config.pin_acceptance: + from .interface import FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR + + os.environ[FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR] = repr(float(config.pin_acceptance)) + + +def _build_llm(config: SweepConfig, verify_len: int): + """One engine per verify length; see the module docstring for why.""" + from tensorrt_llm import LLM + from tensorrt_llm.llmapi import CudaGraphConfig, DSparkDecodingConfig, KvCacheConfig + + spec_config = DSparkDecodingConfig( + max_draft_len=int(verify_len), + # DSparkWorker asserts draft_model.block_size == max_draft_len, so the + # block has to move with the swept length rather than stay at the + # checkpoint's trained value. + block_size=int(verify_len), + speculative_model=config.speculative_model, + enable_confidence_scheduling=bool(config.enable_confidence_scheduling), + # A single tier keeps the captured-graph count identical to an + # unscheduled run (one graph per batch size) while still exercising the + # confidence head, which is part of the deployment's per-step cost. It + # also makes the planner's choice a no-op: max_tier == max_draft_len, so + # the step shape is pinned no matter what the confidence snapshot says. + confidence_verify_len_tiers=( + [int(verify_len)] if config.enable_confidence_scheduling else None + ), + ) + llm_kwargs = dict( + model=config.model, + tensor_parallel_size=int(config.tp_size), + moe_expert_parallel_size=int(config.ep_size) if config.ep_size else None, + enable_attention_dp=bool(config.enable_attention_dp), + attn_backend=config.attn_backend, + max_batch_size=max(config.batch_sizes), + max_seq_len=int(config.max_seq_len), + max_num_tokens=int(config.max_num_tokens), + disable_overlap_scheduler=bool(config.disable_overlap_scheduler), + speculative_config=spec_config, + # Block reuse would let a later cell start from a warm prefix and skip + # prefill work the earlier cells paid for, which shifts the steady-state + # KV length between cells. + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + free_gpu_memory_fraction=float(config.kv_cache_fraction), + ), + # Any batch size outside this list runs eager; an eager step is orders + # of magnitude off the graph-replay cost and would poison the curve. + cuda_graph_config=CudaGraphConfig( + batch_sizes=sorted(set(config.batch_sizes)), + enable_padding=True, + ), + enable_iter_perf_stats=True, + # Unbounded: the buffer is drained while the sweep runs, but a stall + # during a long cell must not silently evict the steps being measured. + iter_stats_max_iterations=-1, + ) + llm_kwargs = {k: v for k, v in llm_kwargs.items() if v is not None} + llm_kwargs.update(config.extra_llm_api_options) + return LLM(**llm_kwargs) + + +def _run_cell(llm, config: SweepConfig, *, batch_size: int, verify_len: int) -> List[StepSample]: + """Hold ``batch_size`` generating requests per rank and record every step.""" + from tensorrt_llm import SamplingParams + + rng = random.Random(config.seed + batch_size * 1000 + verify_len) + low, high = config.token_id_range + num_requests = int(batch_size) * config.dp_size + sampling_params = SamplingParams( + # Every request must outlive the whole window so the batch never drains + # mid-measurement and shrinks the shape being profiled. + max_tokens=config.max_tokens_for(verify_len), + ignore_eos=True, + # No detokenization: it is postprocessing work that would show up in the + # host step time without being part of what the planner is choosing. + detokenize=False, + temperature=0.0, + ) + # Random ids from a mid-vocabulary window rather than real text: no tokenizer + # round-trip is needed, and a distinct prompt per request keeps any prefix + # sharing out of the measurement. The window avoids the low ids where + # special/control tokens live in every large vocabulary. + prompts = [ + {"prompt_token_ids": [rng.randint(low, high) for _ in range(config.input_len)]} + for _ in range(num_requests) + ] + + futures = [llm.generate_async(prompt, sampling_params) for prompt in prompts] + try: + for future in futures: + future.result() + finally: + for future in futures: + if not future.finished: + future.abort() + # Drained exactly once, after the cell has finished. `IterationResult` + # latches itself done the first time its queue runs dry and is only reset + # when new requests are submitted, so polling mid-cell would silently + # discard every step after the first empty read. The timeout is the wait for + # the last few records to cross the IPC boundary, not a poll interval. + rows: List[dict] = list(llm.get_stats(timeout=config.stats_timeout_s) or []) + + aligned = aligned_steps_from_stats( + rows, expected_ranks=config.dp_size, timing_key=config.timing_key + ) + return [ + StepSample( + batch_size=int(batch_size), + verify_len=int(verify_len), + step_time_ms=step_time, + iteration=iteration, + ) + for iteration, num_gen, step_time in aligned + if num_gen == int(batch_size) + ] + + +def run_sweep(config: SweepConfig) -> List[StepSample]: + """Measure every ``(batch_size, verify_len)`` cell, rebuilding per length.""" + config.validate() + _prepare_environment(config) + samples: List[StepSample] = [] + for verify_len in sorted(config.verify_lens): + print( + f"[dspark-sps] building engine for verify_len={verify_len} " + f"(block_size={verify_len}) ...", + file=sys.stderr, + ) + llm = _build_llm(config, verify_len) + try: + for batch_size in sorted(config.batch_sizes): + started = time.time() + cell = _run_cell(llm, config, batch_size=batch_size, verify_len=verify_len) + print( + f"[dspark-sps] bs={batch_size:>5} L={verify_len} " + f"M={total_verify_tokens(batch_size, verify_len):>6} " + f"aligned_steps={len(cell):>5} " + f"({time.time() - started:.1f}s)", + file=sys.stderr, + ) + samples.extend(cell) + finally: + llm.shutdown() + return samples + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def _int_list(text: str) -> List[int]: + return [int(v) for v in str(text).replace(",", " ").split()] + + +def _float_list(text: str) -> List[float]: + return [float(v) for v in str(text).replace(",", " ").split()] + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="python -m tensorrt_llm._torch.speculative.dspark_sps_profiler", + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--out", required=True, help="Path to write the cost-table JSON.") + parser.add_argument( + "--samples-out", + default=None, + help="Raw per-step samples (JSONL). Defaults to .samples.jsonl.", + ) + parser.add_argument( + "--fit-only", + action="store_true", + help="Skip the engine entirely and refit from --samples-in. No GPU needed.", + ) + parser.add_argument( + "--samples-in", + default=None, + help="Comma-separated JSONL sample files to fit (implies --fit-only inputs).", + ) + + model = parser.add_argument_group("model") + model.add_argument("--model", help="Target model checkpoint path or HF id.") + model.add_argument( + "--speculative-model", + default=None, + help="DSpark draft checkpoint. Defaults to --model.", + ) + model.add_argument("--tp-size", type=int, default=1) + model.add_argument("--ep-size", type=int, default=None, help="moe_expert_parallel_size.") + model.add_argument("--enable-attention-dp", action="store_true") + model.add_argument("--attn-backend", default="TRTLLM") + model.add_argument("--max-seq-len", type=int, default=4096) + model.add_argument("--max-num-tokens", type=int, default=8192) + model.add_argument("--kv-cache-fraction", type=float, default=0.5) + model.add_argument( + "--disable-overlap-scheduler", + action="store_true", + help="Profile without the overlap scheduler. The default (overlap ON) is " + "what ships, so leave this off unless the deployment disables it.", + ) + model.add_argument( + "--no-confidence-scheduling", + dest="enable_confidence_scheduling", + action="store_false", + help="Profile with confidence scheduling off. Only for a checkpoint " + "without confidence-head weights: the head is real per-step work and " + "the deployment being profiled runs with it on.", + ) + model.add_argument( + "--extra-llm-api-options", + default=None, + help="YAML file merged into the LLM(...) kwargs (moe_config, quant, ...).", + ) + parser.set_defaults(enable_confidence_scheduling=True) + + sweep = parser.add_argument_group("sweep") + sweep.add_argument( + "--batch-sizes", + default="8,16,32,64,128", + help="Generating requests PER attention-DP rank. Every value is captured " + "as a CUDA graph; use a ladder whose ratios are reachable by the " + "verify-length ratios (a doubling ladder is safe).", + ) + sweep.add_argument( + "--max-draft-len", + type=int, + default=5, + help="The deployment's block size. Sets the top of the verify-length sweep.", + ) + sweep.add_argument( + "--verify-lens", + default=None, + help="Uniform verify lengths to sweep. Defaults to 1..--max-draft-len. " + "Each value costs one engine build.", + ) + sweep.add_argument("--input-len", type=int, default=1024, help="Prompt length in tokens.") + sweep.add_argument("--warmup-steps", type=int, default=16) + sweep.add_argument("--measure-steps", type=int, default=64) + sweep.add_argument("--min-samples", type=int, default=16) + sweep.add_argument("--seed", type=int, default=1234) + sweep.add_argument( + "--timing-field", + choices=sorted(TIMING_FIELDS), + default="host", + help="host = hostStepTimeMS (clean per-loop CPU wall, works under the " + "overlap scheduler); gpu_forward = gpuForwardTimeMS (device time of the " + "fused draft+verify forward only).", + ) + sweep.add_argument( + "--no-pin-acceptance", + dest="pin_acceptance", + action="store_false", + help="Do not force bonus-token-only acceptance. Leaves KV growth rate " + "dependent on the block length, which contaminates theta(M).", + ) + parser.set_defaults(pin_acceptance=True) + + fit = parser.add_argument_group("fit") + fit.add_argument("--riser-tolerance", type=float, default=0.01) + fit.add_argument("--max-breakpoints", type=int, default=8) + fit.add_argument("--min-step-time-spread", type=float, default=0.02) + fit.add_argument( + "--verify-len-tiers", + default=None, + help="Ladder to validate against; defaults to the ladder derived from the " + "measured shelves at --reference-batch-size.", + ) + fit.add_argument("--max-tiers", type=int, default=3) + fit.add_argument( + "--reference-batch-size", + type=int, + default=None, + help="Modal steady-state batch size. Defaults to the largest swept value, " + "since shelf edges move with batch size and trimming only pays at high " + "concurrency.", + ) + fit.add_argument("--probe-acceptance-rates", default="0.3,0.5,0.7,0.8,0.9,0.95") + fit.add_argument( + "--allow-inert", + action="store_true", + help="Emit even if the planner would never trim. Refused by default.", + ) + return parser + + +def _load_samples(paths: Sequence[str]) -> List[StepSample]: + samples: List[StepSample] = [] + for path in paths: + with open(path, encoding="utf-8") as handle: + for line in handle: + line = line.strip() + if line: + samples.append(StepSample.from_json(json.loads(line))) + return samples + + +def _write_samples(path: str, samples: Sequence[StepSample]) -> None: + with open(path, "w", encoding="utf-8") as handle: + for sample in samples: + handle.write(json.dumps(sample.to_json()) + "\n") + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = _build_parser().parse_args(argv) + + verify_lens = ( + _int_list(args.verify_lens) + if args.verify_lens + else list(range(1, int(args.max_draft_len) + 1)) + ) + batch_sizes = _int_list(args.batch_sizes) + samples_out = args.samples_out or f"{args.out}.samples.jsonl" + + if args.fit_only or args.samples_in: + if not args.samples_in: + raise SystemExit("--fit-only requires --samples-in") + samples = _load_samples(_str_list(args.samples_in)) + provenance = {"source": "refit", "samples_in": _str_list(args.samples_in)} + else: + if not args.model: + raise SystemExit("--model is required unless --fit-only is used") + extra: Dict[str, object] = {} + if args.extra_llm_api_options: + import yaml + + with open(args.extra_llm_api_options, encoding="utf-8") as handle: + extra = yaml.safe_load(handle) or {} + config = SweepConfig( + model=args.model, + speculative_model=args.speculative_model or args.model, + batch_sizes=batch_sizes, + verify_lens=verify_lens, + tp_size=args.tp_size, + ep_size=args.ep_size, + enable_attention_dp=args.enable_attention_dp, + max_draft_len=args.max_draft_len, + input_len=args.input_len, + warmup_steps=args.warmup_steps, + measure_steps=args.measure_steps, + min_samples=args.min_samples, + max_seq_len=args.max_seq_len, + max_num_tokens=args.max_num_tokens, + kv_cache_fraction=args.kv_cache_fraction, + attn_backend=args.attn_backend, + disable_overlap_scheduler=args.disable_overlap_scheduler, + enable_confidence_scheduling=args.enable_confidence_scheduling, + timing_key=TIMING_FIELDS[args.timing_field], + pin_acceptance=BONUS_TOKEN_ONLY_ACCEPTANCE if args.pin_acceptance else 0.0, + seed=args.seed, + extra_llm_api_options=extra, + ) + samples = run_sweep(config) + _write_samples(samples_out, samples) + print(f"[dspark-sps] wrote {len(samples)} raw samples to {samples_out}", file=sys.stderr) + provenance = { + "source": "sweep", + "model": config.model, + "tensor_parallel_size": config.tp_size, + "moe_expert_parallel_size": config.ep_size, + "attention_dp": config.enable_attention_dp, + "input_len": config.input_len, + "timing_field": config.timing_key, + "acceptance_pinned": bool(config.pin_acceptance), + "overlap_scheduler": not config.disable_overlap_scheduler, + "confidence_scheduling": config.enable_confidence_scheduling, + "ragged_verify_mode": RaggedVerifyMode.STATIC.value, + } + + cells = summarize_cells(samples, warmup_steps=args.warmup_steps, min_samples=args.min_samples) + if not cells: + # Distinguished from a short cell: this means no step was ever aligned, + # which is a wiring problem (iteration stats off, the wrong rank count, + # the batch never reaching the requested size) rather than a noisy run. + raise InsufficientSamplesError( + f"no cells at all: {len(samples)} raw samples produced nothing usable. " + f"Check that enable_iter_perf_stats reached the engine, that " + f"--enable-attention-dp matches the deployment (the aligned-step " + f"filter expects one stats row per DP rank), and that the requested " + f"batch sizes were actually reached." + ) + reference_bs = int(args.reference_batch_size or max(c.batch_size for c in cells)) + block_size = max(int(args.max_draft_len), max(c.verify_len for c in cells)) + + payload = build_cost_table_payload( + cells, + riser_tolerance=args.riser_tolerance, + max_breakpoints=args.max_breakpoints, + meta={ + "block_size": block_size, + "min_verify_len": 1, + "modal_batch_size": reference_bs, + "warmup_steps_discarded": int(args.warmup_steps), + "statistic": "median", + **provenance, + }, + ) + + tiers = ( + _int_list(args.verify_len_tiers) + if args.verify_len_tiers + else derive_verify_len_tiers( + cost_table=load_cost_table(payload), + num_requests=reference_bs, + block_size=block_size, + min_verify_len=1, + max_tiers=int(args.max_tiers), + ) + ) + payload["_meta"]["recommended_confidence_verify_len_tiers"] = [int(t) for t in tiers] + payload["_meta"]["recommended_tiers_note"] = ( + "Nothing reads this from the table file. Copy it into " + "DSparkDecodingConfig.confidence_verify_len_tiers, which is also what " + "sizes the captured CUDA graph set." + ) + + diagnostics = check_table_is_informative( + payload, + batch_sizes=sorted({c.batch_size for c in cells}), + tiers=tiers, + acceptance_rates=_float_list(args.probe_acceptance_rates), + min_step_time_spread=args.min_step_time_spread, + allow_inert=args.allow_inert, + ) + payload["_meta"]["diagnostics"] = diagnostics + + with open(args.out, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, sort_keys=False) + handle.write("\n") + + _print_report(payload, cells, diagnostics, tiers) + return 0 + + +def _str_list(text: str) -> List[str]: + return [v for v in str(text).replace(",", " ").split() if v] + + +def _print_report( + payload: dict, + cells: Sequence[CellStat], + diagnostics: dict, + tiers: Sequence[int], +) -> None: + meta = payload["_meta"] + print("") + print("DSpark SPS cost table") + print("---------------------") + print(f" token_counts : {payload['token_counts']}") + print(f" step_time_ms : {payload['step_time_ms']} (theta(M) only)") + print(f" fixed_overhead_ms : {payload['fixed_overhead_ms']}") + print(f" batch_sizes : {payload['batch_sizes']}") + print(f" batch_overhead_ms : {payload['batch_overhead_ms']}") + print("") + print(f" cells measured : {len(cells)}") + print( + f" additive-model fit: max residual " + f"{meta['fit']['max_abs_residual_ms']:.4f} ms " + f"({meta['fit']['max_rel_residual']:.2%} of a typical step)" + ) + for warning in meta["fit"]["warnings"]: + print(f" WARNING: {warning}") + print(f" step-time spread : {diagnostics['best_step_time_spread']:.2%} across the ladder") + print("") + print(f" Copy into confidence_verify_len_tiers: {list(tiers)}") + print(" Planner choice vs synthetic acceptance rate (batch_size / rate -> length):") + for probe in diagnostics["profitability_probe"]: + mark = "trim" if probe["trims"] else " " + print( + f" bs={probe['batch_size']:>5} p={probe['acceptance_rate']:.2f} " + f"-> L={probe['chosen_verify_len']} (max {probe['max_tier']}) {mark}" + ) + print("") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index e3b816007ce2..3f3f8344ceb4 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -43,6 +43,7 @@ import torch from ..._utils import prefer_pinned +from .dspark_observability import forced_verify_lens from .dspark_planner import SpsCostTable, budget_argmax_over_uniform_lens from .dspark_schedule import DSparkScheduleConfig, compute_survival, schedule_verify_lens_topk @@ -233,6 +234,25 @@ def decide_verify_lens( """ if num_gen_requests <= 0: return None + + forced = forced_verify_lens(num_gen_requests=num_gen_requests, + tiers=self.tiers, + min_verify_len=self.cfg.min_verify_len) + if forced is not None: + # Deliberately ahead of the cost-table gate. Without a profiled + # table the planner correctly refuses to trim, which means the + # ragged layout is never built and a correctness run silently + # degenerates into a uniform one -- so the two questions "is the + # ragged packing correct" and "does the planner choose to trim" + # cannot be answered by the same run. This override answers the + # first without needing an answer to the second. + # + # Safe by construction: it only decides how many drafted positions + # are sent to the target. Acceptance still runs unchanged against + # whatever is sent, so the output distribution is untouched. + self.stats["forced_lens"] = self.stats.get("forced_lens", 0) + 1 + return forced + if self.cost_table is None or self.cost_table.is_flat: self.stats["fallback_flat_cost"] += 1 return None diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index bdc598873c15..59271b21be42 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -55,6 +55,12 @@ class SampleStateSpec(SampleState): device: SampleStateTensorsSpec host: SampleStateTensorsSpec + #: This step's per-request verify windows, keyed by request id. Snapshotted + #: because ``update_requests`` for step N-1 runs *after* the scheduler has + #: already stamped step N's windows onto the same LlmRequest objects, so + #: reading ``request.py_verify_len`` live rewinds by the wrong amount under + #: the overlap scheduler. None on every uniform path. + verify_lens_snapshot: Optional[dict] = None class SpecSamplerBase(Sampler[SampleStateSpec], AsyncWorkerMixin): @@ -176,7 +182,8 @@ def _request_common_handling( request.py_decoding_iter += 1 @staticmethod - def _verified_len(request: LlmRequest, runtime_draft_len: int) -> int: + def _verified_len(request: LlmRequest, runtime_draft_len: int, + verify_lens_snapshot: Optional[dict] = None) -> int: """How many draft positions THIS request was given to verify. Uniform scheduling gives every request the batch-wide @@ -192,7 +199,13 @@ def _verified_len(request: LlmRequest, runtime_draft_len: int) -> int: ``py_verify_len`` is set only by the ragged path, so this is exactly today's behavior for every other speculation mode. """ - per_request = getattr(request, "py_verify_len", None) + # Prefer the snapshot taken when this step's tokens were sampled. The + # live attribute already belongs to the NEXT step by the time the + # overlap scheduler rewinds this one. + if verify_lens_snapshot is not None: + per_request = verify_lens_snapshot.get(request.py_request_id) + else: + per_request = getattr(request, "py_verify_len", None) return runtime_draft_len if per_request is None else int(per_request) def update_requests( @@ -229,7 +242,9 @@ def update_requests( break req.py_num_accepted_draft_tokens = num_new_tokens - 1 req.py_rewind_len = ( - self._verified_len(req, runtime_draft_len) - req.py_num_accepted_draft_tokens + self._verified_len(req, runtime_draft_len, + state.verify_lens_snapshot) + - req.py_num_accepted_draft_tokens ) self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) @@ -321,10 +336,20 @@ def sample_async( for request in finished_context_requests: request.py_draft_tokens = [1] * self.draft_len + # Cheap (a few ints) and only populated when the ragged path is live. + verify_lens_snapshot = None + for request in sampling_requests: + window = getattr(request, "py_verify_len", None) + if window is not None: + if verify_lens_snapshot is None: + verify_lens_snapshot = {} + verify_lens_snapshot[request.py_request_id] = int(window) + return SampleStateSpec( requests=sampling_requests, device=device_tensors, host=host_tensors, sampler_event=sampler_event, runtime_draft_len=runtime_draft_len, + verify_lens_snapshot=verify_lens_snapshot, ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py new file mode 100644 index 000000000000..e2940f0eca15 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py @@ -0,0 +1,614 @@ +# 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. +"""Unit tests for the DSpark SPS cost-table profiler (CPU only, no model load). + +Everything here exercises the pure half of the profiler: turning ``get_stats`` +rows into aligned samples, summarizing cells, separating ``alpha(bs)`` from +``theta(M)``, and assembling / refusing the emitted JSON. The sweep driver needs +a GPU and a checkpoint and is deliberately not covered. + +The load-bearing case is :func:`test_round_trip_reproduces_measured_cells`: it +takes a synthetic machine with a known cost curve, runs it through the whole +fit-and-emit pipeline, and reloads the JSON exactly the way +``dspark.py::_build_verify_planner`` does. If the profiler and the loader ever +disagree about units, keys or the meaning of a token count, that test is where +it shows up. +""" + +import json + +import numpy as np +import pytest + +from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable, total_verify_tokens +from tensorrt_llm._torch.speculative.dspark_sps_profiler import ( + CellStat, + FlatCostTableError, + InertCostTableError, + InsufficientSamplesError, + StepSample, + SweepConfig, + SweepGeometryError, + aligned_steps_from_stats, + build_cost_table_payload, + check_table_is_informative, + compress_to_risers, + fit_additive_cost_model, + load_cost_table, + profitability_probe, + running_max, + summarize_cells, +) + +#: The exact set of keys ``dspark.py::_build_verify_planner`` reads. Anything +#: else in the file is ignored by the loader, so the profiler must not rely on +#: it -- and must not omit any of these. +LOADER_KEYS = { + "token_counts", + "step_time_ms", + "fixed_overhead_ms", + "batch_sizes", + "batch_overhead_ms", +} + + +def _stats_row(*, iteration, num_gen, rank=0, num_ctx=0, host_ms=10.0, gpu_ms=8.0): + """One /metrics-shaped iteration record.""" + row = { + "iter": iteration, + "attentionDpRank": rank, + "inflightBatchingStats": { + "numContextRequests": num_ctx, + "numGenRequests": num_gen, + }, + } + if host_ms is not None: + row["hostStepTimeMS"] = host_ms + if gpu_ms is not None: + row["gpuForwardTimeMS"] = gpu_ms + return row + + +# -------------------------------------------------------------------------- +# aligned step extraction +# -------------------------------------------------------------------------- + + +def test_aligned_steps_single_rank(): + rows = [_stats_row(iteration=i, num_gen=8, host_ms=10.0 + i) for i in range(3)] + assert aligned_steps_from_stats(rows) == [(0, 8, 10.0), (1, 8, 11.0), (2, 8, 12.0)] + + +def test_aligned_steps_drops_prefill_iterations(): + """A step with context requests is a different shape, not a cheap decode.""" + rows = [ + _stats_row(iteration=0, num_gen=0, num_ctx=8, host_ms=90.0), + _stats_row(iteration=1, num_gen=8, num_ctx=2, host_ms=40.0), + _stats_row(iteration=2, num_gen=8, host_ms=10.0), + ] + assert aligned_steps_from_stats(rows) == [(2, 8, 10.0)] + + +def test_aligned_steps_requires_every_dp_rank(): + """A missing rank means the iteration was not observed batch-wide.""" + rows = [ + _stats_row(iteration=0, num_gen=8, rank=0), + _stats_row(iteration=0, num_gen=8, rank=1), + _stats_row(iteration=1, num_gen=8, rank=0), + ] + assert aligned_steps_from_stats(rows, expected_ranks=2) == [(0, 8, 10.0)] + + +def test_aligned_steps_requires_ranks_to_agree_on_batch_size(): + """Ranks at different batch sizes are not one (bs, M) point.""" + rows = [ + _stats_row(iteration=0, num_gen=8, rank=0), + _stats_row(iteration=0, num_gen=7, rank=1), + ] + assert aligned_steps_from_stats(rows, expected_ranks=2) == [] + + +def test_aligned_steps_takes_rank0_timing(): + """ADP fans out rank-local counters with rank-0's clock; do not average.""" + rows = [ + _stats_row(iteration=0, num_gen=8, rank=1, host_ms=99.0), + _stats_row(iteration=0, num_gen=8, rank=0, host_ms=10.0), + ] + assert aligned_steps_from_stats(rows, expected_ranks=2) == [(0, 8, 10.0)] + + +def test_aligned_steps_skips_missing_timing_field(): + """gpuForwardTimeMS is read without a sync and is absent on some steps.""" + rows = [ + _stats_row(iteration=0, num_gen=8, gpu_ms=None), + _stats_row(iteration=1, num_gen=8, gpu_ms=7.5), + ] + assert aligned_steps_from_stats(rows, timing_key="gpuForwardTimeMS") == [(1, 8, 7.5)] + + +# -------------------------------------------------------------------------- +# warmup + median +# -------------------------------------------------------------------------- + + +def _samples(batch_size, verify_len, times): + return [ + StepSample(batch_size=batch_size, verify_len=verify_len, step_time_ms=t, iteration=i) + for i, t in enumerate(times) + ] + + +def test_summarize_discards_warmup(): + warm = [500.0] * 4 + steady = [10.0] * 8 + cells = summarize_cells(_samples(8, 3, warm + steady), warmup_steps=4, min_samples=8) + assert len(cells) == 1 + assert cells[0].step_time_ms == pytest.approx(10.0) + assert cells[0].num_samples == 8 + + +def test_summarize_reports_median_not_mean(): + """The right tail is real and must not move the point.""" + times = [10.0, 10.0, 10.0, 10.0, 10.0, 1000.0] + cells = summarize_cells(_samples(8, 3, times), warmup_steps=0, min_samples=1) + assert cells[0].step_time_ms == pytest.approx(10.0) + assert cells[0].step_time_ms != pytest.approx(float(np.mean(times))) + # ...but the tail is still reported, so a pathological cell is visible. + assert cells[0].p90_ms > cells[0].step_time_ms + + +def test_summarize_refuses_short_cell(): + """A silently dropped cell can disconnect the grid, not just widen error bars.""" + with pytest.raises(InsufficientSamplesError, match="steady samples"): + summarize_cells(_samples(8, 3, [10.0] * 5), warmup_steps=4, min_samples=8) + + +def test_summarize_orders_by_iteration_before_cutting_warmup(): + """Stats arrive per queue, not necessarily in iteration order.""" + out_of_order = [ + StepSample(batch_size=8, verify_len=1, step_time_ms=t, iteration=i) + for i, t in sorted(enumerate([500.0, 500.0, 10.0, 10.0]), key=lambda p: -p[0]) + ] + cells = summarize_cells(out_of_order, warmup_steps=2, min_samples=2) + assert cells[0].step_time_ms == pytest.approx(10.0) + + +def test_step_sample_json_round_trip(): + sample = StepSample(batch_size=16, verify_len=3, step_time_ms=12.5, iteration=7) + assert StepSample.from_json(json.loads(json.dumps(sample.to_json()))) == sample + assert sample.total_verify_tokens == total_verify_tokens(16, 3) == 64 + + +# -------------------------------------------------------------------------- +# additive fit +# -------------------------------------------------------------------------- + + +def _synthetic_cells(batch_sizes, verify_lens, *, alpha, theta, noise=None): + """Cells generated by an exactly additive machine ``alpha(bs) + theta(M)``.""" + cells = [] + for batch_size in batch_sizes: + for verify_len in verify_lens: + tokens = total_verify_tokens(batch_size, verify_len) + value = alpha(batch_size) + theta(tokens) + if noise is not None: + value += noise(batch_size, verify_len) + cells.append( + CellStat( + batch_size=batch_size, + verify_len=verify_len, + step_time_ms=value, + num_samples=32, + p10_ms=value, + p90_ms=value, + ) + ) + return cells + + +def _staircase(edges_and_costs): + def theta(tokens): + value = edges_and_costs[0][1] + for edge, cost in edges_and_costs: + if tokens >= edge: + value = cost + return value + + return theta + + +BATCH_SIZES = [8, 16, 32, 64] +VERIFY_LENS = [1, 2, 3, 4, 5] +ALPHA = {8: 20.0, 16: 22.0, 32: 26.0, 64: 34.0} +THETA = _staircase([(0, 1.0), (64, 2.0), (128, 4.0), (256, 8.0)]) + + +def test_fit_recovers_relative_structure_exactly(): + """The fit is only identified up to one shared constant -- pin that, then + every cell must be reproduced to the last bit.""" + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) + fit = fit_additive_cost_model(cells) + + for cell in cells: + predicted = fit.intercept_ms[cell.batch_size] + fit.theta_ms[cell.total_verify_tokens] + assert predicted == pytest.approx(cell.step_time_ms, abs=1e-6) + assert fit.max_rel_residual == pytest.approx(0.0, abs=1e-9) + + # The batch axis keeps its measured *differences*, which is what alpha means. + for smaller, larger in zip(BATCH_SIZES, BATCH_SIZES[1:]): + measured = ALPHA[larger] - ALPHA[smaller] + fitted = fit.intercept_ms[larger] - fit.intercept_ms[smaller] + assert fitted == pytest.approx(measured, abs=1e-6) + + +def test_fit_surfaces_a_bad_cell_rather_than_hiding_it(): + """A cell the additive model cannot explain is evidence, not noise. + + Dropping it would produce a table that fits beautifully and predicts badly, + and would delete the one signal that says ``theta`` is not a function of + ``M`` alone. + """ + + def noise(batch_size, verify_len): + return 500.0 if (batch_size, verify_len) == (32, 3) else 0.0 + + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA, noise=noise) + fit = fit_additive_cost_model(cells) + + assert fit.max_abs_residual_ms > 100.0 + assert fit.max_rel_residual > 1.0 + report = " ".join(fit.warnings) + assert "additive model misses" in report + # The least-squares solve spreads the bad cell over everything sharing its + # batch size or token count, so the report names a neighbourhood -- the bad + # cell is in it, but so are innocent ones, which is why it is worded that way. + assert "bs=32, L=3" in report + + +def test_fit_stays_quiet_when_the_model_holds(): + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) + assert fit_additive_cost_model(cells).warnings == () + + +def test_fit_clamps_a_dip_upward(): + """A non-monotone theta would offer the planner a cheaper *longer* verify.""" + dipped = _staircase([(0, 1.0), (64, 5.0), (128, 3.0), (256, 8.0)]) + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=dipped) + fit = fit_additive_cost_model(cells) + + tokens = sorted(fit.theta_ms) + values = [fit.theta_ms[m] for m in tokens] + assert values == running_max(values) + assert any("monoton" in w for w in fit.warnings) + + +def test_fit_keeps_theta_positive(): + """SpsCostTable rejects a non-positive step time, so the shift is bounded.""" + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) + fit = fit_additive_cost_model(cells) + assert min(fit.theta_ms.values()) > 0.0 + + +def test_fit_refuses_a_disconnected_sweep(): + """Two islands each carry their own unknown offset; splicing them lies.""" + cells = _synthetic_cells([8], [1, 2], alpha=ALPHA.get, theta=THETA) + cells += _synthetic_cells([100], [1, 2], alpha=lambda _: 40.0, theta=THETA) + with pytest.raises(SweepGeometryError, match="disconnected"): + fit_additive_cost_model(cells) + + +def test_fit_refuses_a_single_length_sweep(): + """One verify length per batch size confounds alpha and theta completely.""" + cells = _synthetic_cells(BATCH_SIZES, [5], alpha=ALPHA.get, theta=THETA) + with pytest.raises(SweepGeometryError): + fit_additive_cost_model(cells) + + +# -------------------------------------------------------------------------- +# staircase shaping +# -------------------------------------------------------------------------- + + +def test_running_max_is_monotone(): + assert running_max([1.0, 3.0, 2.0, 2.5, 9.0]) == [1.0, 3.0, 3.0, 3.0, 9.0] + + +def test_compress_keeps_only_real_risers(): + tokens = [16, 32, 48, 64, 80] + times = [1.0, 1.002, 1.004, 2.0, 2.001] + kept_tokens, kept_times = compress_to_risers( + tokens, times, min_riser_ms=0.02, max_breakpoints=8 + ) + assert kept_tokens == [16, 64] + assert kept_times == [1.0, 2.0] + + +def test_compress_drops_the_shallowest_risers_first(): + tokens = [16, 32, 64, 128] + times = [1.0, 1.1, 4.0, 9.0] + kept_tokens, _ = compress_to_risers(tokens, times, min_riser_ms=0.0, max_breakpoints=3) + assert kept_tokens == [16, 64, 128] + + +def test_compress_never_drops_the_base_shelf(): + """Index 0 is what every below-range query clamps onto.""" + tokens = [16, 32, 64] + times = [1.0, 100.0, 101.0] + kept_tokens, _ = compress_to_risers(tokens, times, min_riser_ms=0.0, max_breakpoints=2) + assert kept_tokens[0] == 16 + + +# -------------------------------------------------------------------------- +# payload assembly and JSON round-trip through SpsCostTable +# -------------------------------------------------------------------------- + + +def _payload(): + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) + return build_cost_table_payload(cells, riser_tolerance=0.01, max_breakpoints=8), cells + + +def test_payload_has_exactly_the_loader_keys(): + payload, _ = _payload() + assert LOADER_KEYS.issubset(payload) + # Everything else must live under _meta, where the loader ignores it. + assert set(payload) - LOADER_KEYS == {"_meta"} + + +def test_payload_survives_a_json_file_round_trip(): + payload, _ = _payload() + reloaded = json.loads(json.dumps(payload)) + table = load_cost_table(reloaded) + assert isinstance(table, SpsCostTable) + assert not table.is_flat + # The loader builds SpsCostTable directly, so the profiler's own validation + # must be the real constructor rather than a copy of its rules. + assert list(table.token_counts) == payload["token_counts"] + assert len(table.batch_sizes) == len(table.batch_overhead_ms) + + +def test_round_trip_reproduces_measured_cells(): + """The whole point: what the planner reads back must be what was measured. + + The tolerance is one shelf width, because the emitted curve is a floor + staircase by construction -- a cell between two risers is deliberately + priced at the lower shelf. + """ + payload, cells = _payload() + table = load_cost_table(json.loads(json.dumps(payload))) + for cell in cells: + predicted = table.step_time(cell.total_verify_tokens, cell.batch_size) + assert predicted == pytest.approx(cell.step_time_ms, abs=1e-4) + + +def test_emitted_terms_are_not_double_counted(): + """step_time_ms is theta only; the floor lives in the two overhead fields.""" + payload, cells = _payload() + table = load_cost_table(payload) + smallest = min(BATCH_SIZES) + floor = payload["fixed_overhead_ms"] + table.batch_overhead(smallest) + # The floor is the M-independent part, i.e. the intercept of the measured + # row at the smallest batch size -- not the whole step time. + row = sorted((c for c in cells if c.batch_size == smallest), key=lambda c: c.verify_len) + assert floor < row[0].step_time_ms + assert floor == pytest.approx(ALPHA[smallest], abs=0.5) + + +def test_batch_overhead_grows_with_batch_size(): + payload, _ = _payload() + overheads = payload["batch_overhead_ms"] + assert overheads == sorted(overheads) + assert payload["batch_sizes"] == sorted(BATCH_SIZES) + assert payload["fixed_overhead_ms"] >= 0.0 + + +def test_meta_carries_the_tier_advice_and_the_cells(): + payload, cells = _payload() + assert payload["_meta"]["encoding"] == "decomposed" + assert len(payload["_meta"]["cells"]) == len(cells) + + +# -------------------------------------------------------------------------- +# refusals +# -------------------------------------------------------------------------- + + +def test_refuses_a_flat_curve(): + """A machine whose step cost ignores M must not produce a table at all.""" + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=lambda _: 3.0) + payload = build_cost_table_payload(cells) + assert load_cost_table(payload).is_flat + with pytest.raises(FlatCostTableError, match="is_flat"): + check_table_is_informative(payload, batch_sizes=BATCH_SIZES, tiers=[1, 3, 5]) + + +def test_refuses_a_curve_that_only_moves_by_noise(): + """is_flat is exact float equality, so it cannot catch this one.""" + cells = _synthetic_cells( + BATCH_SIZES, + VERIFY_LENS, + alpha=ALPHA.get, + theta=lambda tokens: 3.0 + 1e-4 * tokens, + ) + payload = build_cost_table_payload(cells, riser_tolerance=0.0, max_breakpoints=8) + assert not load_cost_table(payload).is_flat + with pytest.raises(FlatCostTableError, match="noise floor"): + check_table_is_informative( + payload, batch_sizes=BATCH_SIZES, tiers=[1, 3, 5], min_step_time_spread=0.02 + ) + + +# A real but shallow riser: at bs=32 tier 3 (M=128) sits below it and tier 5 +# (M=192) above, so the curve is genuinely non-flat and clears the noise floor +# (2.5%), yet the step only gets 2.5% more expensive while the acceptance yield +# given up between those tiers is ~36% at p=0.9. The argmax therefore never +# leaves max_tier -- non-flat and inert at the same time. +INERT_PAYLOAD = { + "token_counts": [0, 160], + "step_time_ms": [2.0, 2.5], + "fixed_overhead_ms": 18.0, + "batch_sizes": [32], + "batch_overhead_ms": [0.0], +} + + +def test_refuses_an_inert_table(): + """Non-flat, but no tier pair is worth trading yield for. + + This is the failure mode that does *not* trip ``fallback_flat_cost``, so + nothing downstream would report it. + """ + with pytest.raises(InertCostTableError, match="inert"): + check_table_is_informative( + INERT_PAYLOAD, + batch_sizes=[32], + tiers=[1, 3, 5], + acceptance_rates=[0.9, 0.95], + ) + + +def test_allow_inert_escape_hatch(): + diagnostics = check_table_is_informative( + INERT_PAYLOAD, + batch_sizes=[32], + tiers=[1, 3, 5], + acceptance_rates=[0.9, 0.95], + allow_inert=True, + ) + assert diagnostics["trims_somewhere"] is False + + +def test_refuses_a_single_entry_ladder(): + payload, _ = _payload() + with pytest.raises(InertCostTableError, match="single entry"): + check_table_is_informative(payload, batch_sizes=BATCH_SIZES, tiers=[5]) + + +def test_informative_table_passes_and_reports_where_trimming_wins(): + payload, _ = _payload() + diagnostics = check_table_is_informative( + payload, + batch_sizes=BATCH_SIZES, + tiers=[1, 3, 5], + acceptance_rates=[0.3, 0.6, 0.9], + ) + assert diagnostics["trims_somewhere"] is True + assert diagnostics["best_step_time_spread"] > 0.02 + + +# -------------------------------------------------------------------------- +# sweep geometry guards (pure; the sweep itself needs a GPU) +# -------------------------------------------------------------------------- + + +def _sweep(**overrides) -> SweepConfig: + kwargs = dict( + model="/models/x", + speculative_model="/models/x", + batch_sizes=[8, 16], + verify_lens=[1, 3, 5], + input_len=1024, + warmup_steps=8, + measure_steps=32, + max_seq_len=4096, + max_num_tokens=8192, + ) + kwargs.update(overrides) + return SweepConfig(**kwargs) + + +def test_token_budget_is_the_step_count_when_acceptance_is_pinned(): + """Pinned, a step commits exactly one token, so the two are the same number.""" + config = _sweep() + assert config.max_tokens_for(5) == config.warmup_steps + config.measure_steps + 16 + + +def test_token_budget_is_scaled_when_acceptance_is_not_pinned(): + """Unpinned, a step can commit verify_len + 1 tokens, so the budget must cover it.""" + pinned = _sweep() + loose = _sweep(pin_acceptance=0.0) + assert loose.max_tokens_for(5) == pinned.max_tokens_for(5) * 6 + + +def test_validate_rejects_a_sequence_that_would_be_evicted(): + """An evicted request drains the batch mid-cell and shrinks the shape.""" + with pytest.raises(SweepGeometryError, match="max-seq-len"): + _sweep(input_len=4064).validate() + + +def test_validate_rejects_a_decode_step_wider_than_max_num_tokens(): + """A split decode step never has the shape its cell is filed under.""" + with pytest.raises(SweepGeometryError, match="widest decode step"): + _sweep(batch_sizes=[8, 4096], max_num_tokens=1024).validate() + + +def test_validate_accepts_a_sane_sweep(): + _sweep().validate() + + +def test_dp_size_only_counts_ranks_under_attention_dp(): + """Non-ADP publishes one stats row per iteration however wide TP is.""" + assert _sweep(tp_size=8).dp_size == 1 + assert _sweep(tp_size=8, enable_attention_dp=True).dp_size == 8 + + +# -------------------------------------------------------------------------- +# profitability probe (runs the real planner argmax) +# -------------------------------------------------------------------------- + + +def test_probe_trims_at_low_acceptance_and_not_at_high(): + """Trimming is a trade: it only wins when the drafts are unlikely to survive.""" + payload = { + # A riser between bs*(3+1)=128 and bs*(5+1)=192 at bs=32, so tiers 3 and + # 5 sit on different shelves and the argmax has something to choose: + # T(3) = 3 ms, T(5) = 4 ms, a 33% step for a yield gain that is worth it + # only when the drafts survive. + "token_counts": [0, 160], + "step_time_ms": [2.0, 3.0], + "fixed_overhead_ms": 0.0, + "batch_sizes": [32], + "batch_overhead_ms": [1.0], + } + probes = profitability_probe( + payload, batch_size=32, tiers=[1, 3, 5], acceptance_rates=[0.05, 0.99] + ) + by_rate = {p["acceptance_rate"]: p for p in probes} + assert by_rate[0.05]["chosen_verify_len"] < 5 + assert by_rate[0.05]["trims"] is True + assert by_rate[0.99]["chosen_verify_len"] == 5 + + +def test_probe_uses_the_bonus_token_convention(): + """A tier's cost must be looked up at bs*(L+1), never bs*L. + + The two differ by a whole batch, which is easily a shelf's width; this pins + the profiler to the same convention as ``total_verify_tokens``. + """ + payload = { + # Priced so that tier 5 (M = 8*6 = 48) is expensive and tier 3 + # (M = 8*4 = 32) is cheap. Under the wrong bs*L convention tier 5 would + # be M = 40 and land on the cheap shelf instead. + "token_counts": [0, 44], + "step_time_ms": [1.0, 50.0], + "fixed_overhead_ms": 0.0, + "batch_sizes": [8], + "batch_overhead_ms": [0.0], + } + table = load_cost_table(payload) + assert table.step_time(total_verify_tokens(8, 5), 8) == pytest.approx(50.0) + assert table.step_time(total_verify_tokens(8, 3), 8) == pytest.approx(1.0) + probes = profitability_probe(payload, batch_size=8, tiers=[1, 3, 5], acceptance_rates=[0.9]) + assert probes[0]["chosen_verify_len"] == 3 diff --git a/tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py b/tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py new file mode 100644 index 000000000000..9ca99af4af69 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py @@ -0,0 +1,294 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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 + +"""Ragged (per-request) query lengths for ``indexer_topk_decode``. + +The uniform kernel reconstructs two things from the scalar ``next_n``: which +request a logits row belongs to (``row // next_n``) and how far back that row +may attend (``seq_len - next_n + row % next_n + 1``). Neither holds once DSpark's +confidence scheduler gives each request its own verify window, so the kernel +takes an optional ``row_kv_lens`` carrying the extent per row. + +The dangerous failure is silent, not loud: a ragged batch whose padded row count +happens to be divisible by the request count satisfies the old +``seq_lens.size(0) * next_n == num_rows`` check, and the kernel then attributes +rows to the wrong requests and masks them at the wrong depth. Every test here is +built to catch that rather than a crash. +""" + +import os + +import pytest +import torch + +if not torch.cuda.is_available(): + pytest.skip("CUDA is required for indexer_topk tests", + allow_module_level=True) + + +def _register_ops() -> None: + """Make ``torch.ops.trtllm.indexer_topk_decode`` available. + + Importing ``tensorrt_llm`` is the usual way, but it drags in the whole model + zoo -- including optional CuTe-DSL code paths that are absent in some + containers. This file tests one C++ kernel, so it falls back to loading the + op library directly rather than being unrunnable whenever an unrelated + optional dependency is missing. + """ + try: + import tensorrt_llm # noqa: F401 + return + except ImportError: + pass + + here = os.path.dirname(os.path.abspath(__file__)) + repo = os.path.abspath(os.path.join(here, *[os.pardir] * 5)) + for candidate in ( + os.path.join(repo, "cpp", "build", "tensorrt_llm", "thop", + "libth_common.so"), + os.path.join(repo, "tensorrt_llm", "libs", "libth_common.so"), + ): + if os.path.exists(candidate): + torch.ops.load_library(candidate) + return + pytest.skip("neither tensorrt_llm nor libth_common.so is importable", + allow_module_level=True) + + +_register_ops() + +INDEX_TOPK = 512 +NUM_COLUMNS = 8192 + +# Every KV length is kept above ``INDEX_TOPK * kMaxBlocksPerRow`` so that each +# block of the split-work path still has at least ``INDEX_TOPK`` candidates to +# rank. Below that a block cannot fill its output slice, the merge pass reads +# whatever was in the scratch buffer, and two runs differ purely by allocator +# luck -- which looks exactly like a kernel bug. +KV_LENS = [2048, 4096, 3000, 2560] + + + +def _assert_same_topk(actual: torch.Tensor, expected: torch.Tensor) -> None: + """Compare top-k *sets* per row, not the raw index arrays. + + The kernel does not promise an output order: the radix path's split-work + merge visits blocks in whatever order they finish, so two identical calls + routinely return the same indices permuted differently. Only the selected + set is part of the contract, and it is what a wrong causal extent would + change. + """ + assert actual.shape == expected.shape, (actual.shape, expected.shape) + actual_sorted = actual.sort(dim=-1).values + expected_sorted = expected.sort(dim=-1).values + torch.testing.assert_close(actual_sorted, expected_sorted, atol=0, rtol=0) + + +def _make_logits(num_rows: int, num_columns: int, + seed: int) -> torch.Tensor: + generator = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn(num_rows, + num_columns, + dtype=torch.float32, + device="cuda", + generator=generator) + + +def _row_kv_lens(kv_lens: list[int], verify_lens: list[int]) -> torch.Tensor: + """The extent every row may attend to, matching the kernel's uniform formula. + + Request ``r`` verifying ``v`` positions contributes rows ``o = 0..v-1``, and + row ``o`` sees ``kv_len - v + o + 1`` tokens. Substituting ``v = next_n`` + recovers exactly what the uniform path computes internally, which is what + makes the equivalence test below meaningful. + """ + rows: list[int] = [] + for kv_len, verify_len in zip(kv_lens, verify_lens): + rows.extend(kv_len - verify_len + offset + 1 + for offset in range(verify_len)) + return torch.tensor(rows, dtype=torch.int32, device="cuda") + + +def _run(logits: torch.Tensor, + seq_lens: list[int], + next_n: int, + row_kv_lens: torch.Tensor = None, + compress_ratio: int = 1) -> torch.Tensor: + num_rows = logits.shape[0] + indices = torch.empty(num_rows, + INDEX_TOPK, + dtype=torch.int32, + device="cuda") + # The op refuses to allocate its own split-work scratch (stale pointers are + # a CUDA-graph hazard), so the caller owns it. kMaxBlocksPerRow is small; + # over-allocating is free here and keeps the test independent of the + # dispatcher's blocks-per-row heuristic. + aux_elems = num_rows * 8 * INDEX_TOPK + # Zeroed, not empty: the merge pass reads every slot its blocks could have + # written, so uninitialized scratch makes the result allocator-dependent. + aux_indices = torch.zeros(aux_elems, dtype=torch.int32, device="cuda") + aux_logits = torch.zeros(aux_elems, dtype=torch.float32, device="cuda") + torch.ops.trtllm.indexer_topk_decode( + logits, + torch.tensor(seq_lens, dtype=torch.int32, device="cuda"), + indices, + next_n, + INDEX_TOPK, + compress_ratio=compress_ratio, + radix_aux_indices=aux_indices, + radix_aux_logits=aux_logits, + row_kv_lens=row_kv_lens, + ) + return indices + + +@pytest.mark.parametrize("compress_ratio", [1, 4]) +@pytest.mark.parametrize("next_n", [1, 2, 6]) +def test_ragged_with_equal_lens_matches_uniform(next_n, compress_ratio): + """``row_kv_lens`` spelled out by hand must reproduce the uniform result. + + This is the load-bearing test. If the explicit per-row extent and the + kernel's internal ``next_n`` arithmetic ever disagree, every ragged result + is wrong in a way no shape check can see -- so pin them to be bit-identical + on a batch where both are defined. + """ + batch_size = 4 + kv_lens = list(KV_LENS) + logits = _make_logits(batch_size * next_n, NUM_COLUMNS, seed=1234) + + uniform = _run(logits, kv_lens, next_n, compress_ratio=compress_ratio) + ragged = _run(logits, + kv_lens, + next_n=1, + row_kv_lens=_row_kv_lens(kv_lens, [next_n] * batch_size), + compress_ratio=compress_ratio) + + _assert_same_topk(uniform, ragged) + + +def test_genuinely_ragged_matches_per_request_reference(): + """Each request's rows must match running that request on its own. + + A per-request reference is the only check that catches cross-request + misattribution: if row ``i`` were scored against the wrong request's + sequence length, a whole-batch self-consistency check would still pass. + """ + kv_lens = list(KV_LENS) + verify_lens = [6, 3, 1, 4] + total_rows = sum(verify_lens) + logits = _make_logits(total_rows, NUM_COLUMNS, seed=99) + + ragged = _run(logits, + kv_lens, + next_n=1, + row_kv_lens=_row_kv_lens(kv_lens, verify_lens)) + + offset = 0 + for kv_len, verify_len in zip(kv_lens, verify_lens): + rows = logits[offset:offset + verify_len].contiguous() + expected = _run(rows, [kv_len], verify_len) + _assert_same_topk(ragged[offset:offset + verify_len], expected) + offset += verify_len + + +def test_divisible_ragged_batch_is_not_silently_accepted(): + """The shape that used to pass the old check by coincidence. + + ``verify_lens = [6, 4, 1, 1]`` totals 12 rows for 4 requests, so + ``seq_lens.size(0) * next_n == num_rows`` holds for ``next_n == 3`` even + though no request verifies 3 positions. Without ``row_kv_lens`` the kernel + accepts it and computes nonsense; with it, the answer must match the + per-request reference. + """ + kv_lens = list(KV_LENS) + verify_lens = [6, 4, 1, 1] + assert sum(verify_lens) % len(verify_lens) == 0, "test premise" + logits = _make_logits(sum(verify_lens), NUM_COLUMNS, seed=7) + + ragged = _run(logits, + kv_lens, + next_n=1, + row_kv_lens=_row_kv_lens(kv_lens, verify_lens)) + + bogus_next_n = sum(verify_lens) // len(verify_lens) + accidental = _run(logits, kv_lens, bogus_next_n) + assert not torch.equal(ragged.sort(dim=-1).values, + accidental.sort(dim=-1).values), ( + "the ragged result coincided with the misattributed uniform one; the " + "test can no longer detect the bug it exists for") + + +def test_row_kv_lens_length_is_validated(): + """A short ``row_kv_lens`` must raise rather than read out of bounds.""" + logits = _make_logits(8, NUM_COLUMNS, seed=3) + with pytest.raises(RuntimeError, match="one entry per logits row"): + _run(logits, [1024, 2048], next_n=1, + row_kv_lens=torch.tensor([1024] * 4, + dtype=torch.int32, + device="cuda")) + + +def test_cuda_graph_capture_replays_with_new_lens(): + """The ragged path must be capturable and must honor rewritten lengths. + + Capture is the whole point of landing on a token bucket: a ragged step that + cannot replay costs far more than the tokens it trims. Replay with different + contents in the same buffers is what a real step does. + """ + kv_lens = list(KV_LENS) + verify_lens = [6, 3, 1, 4] + total_rows = sum(verify_lens) + logits = _make_logits(total_rows, NUM_COLUMNS, seed=11) + + seq_lens_dev = torch.tensor(kv_lens, dtype=torch.int32, device="cuda") + row_kv_lens_dev = _row_kv_lens(kv_lens, verify_lens) + indices = torch.empty(total_rows, + INDEX_TOPK, + dtype=torch.int32, + device="cuda") + # Capture requires caller-owned scratch: the op refuses to allocate it + # itself precisely because a graph would bake in a stale pointer. + aux_elems = total_rows * 8 * INDEX_TOPK + aux_indices = torch.zeros(aux_elems, dtype=torch.int32, device="cuda") + aux_logits = torch.zeros(aux_elems, dtype=torch.float32, device="cuda") + + def launch(): + torch.ops.trtllm.indexer_topk_decode(logits, + seq_lens_dev, + indices, + 1, + INDEX_TOPK, + compress_ratio=1, + radix_aux_indices=aux_indices, + radix_aux_logits=aux_logits, + row_kv_lens=row_kv_lens_dev) + + # Warm up outside the graph; capture would otherwise record allocator work. + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + launch() + torch.cuda.current_stream().wait_stream(stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch() + + # A different split, written in place, must produce the eager answer. + new_verify_lens = [4, 4, 2, 4] + assert sum(new_verify_lens) == total_rows, "replay must keep the shape" + row_kv_lens_dev.copy_(_row_kv_lens(kv_lens, new_verify_lens)) + graph.replay() + torch.cuda.synchronize() + replayed = indices.clone() + + eager = _run(logits, + kv_lens, + next_n=1, + row_kv_lens=_row_kv_lens(kv_lens, new_verify_lens)) + _assert_same_topk(replayed, eager) From d8897046653331bca51242bf068b052568fca4ce Mon Sep 17 00:00:00 2001 From: Lanyu Liao Date: Sat, 1 Aug 2026 01:36:07 -0700 Subject: [PATCH 004/177] [None][doc] Scope the DSpark ragged-verify measurements to x86 DGX B300 Every timing, throughput, memory and page-cache number in the handover doc came from one machine: umb-b300-001, x86_64 + B300 SXM6 (SM103, 275 GB/card, 2 TB host RAM) with weights on shared NFS. Written without that qualifier they read as platform-independent facts, which they are not. Adds a scope section up front -- before the numbers, since a reader hits the numbers first -- and flags what changes on GB300 (`gb300` / `b300-nvl8`): aarch64 rather than x86, so a separate SBSA image and its own cpp/build; Grace's unified memory and a different storage backend, so the 2h-prefetch / 40min-assembly split does not carry over; a different host memory budget, so the page-cache reuse strategy needs re-measuring; and a different Slurm account. Also separates what DOES transfer: the six P0 defect locations, the kernels' scalar next_n assumption, the observability design, and the "empty output is not a pass" lessons are all code- and method-level, not hardware-level. Signed-off-by: Lanyu Liao --- docs/dspark_ragged_verify_status.md | 37 ++++++++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/docs/dspark_ragged_verify_status.md b/docs/dspark_ragged_verify_status.md index 303339cc71ac..5b2d3fef2a83 100644 --- a/docs/dspark_ragged_verify_status.md +++ b/docs/dspark_ragged_verify_status.md @@ -8,6 +8,38 @@ --- +## 0. 观测范围 —— 所有实测数字只适用于 x86 DGX B300 + +本文档里每一个耗时、吞吐、显存、page-cache 数字都是在**一台机器**上量的： + +| | | +|---|---| +| 节点 | `umb-b300-001` | +| 分区 | `b300@ts6/dgx-b300@ts1/8gpu-256cpu-2048gb` | +| 架构 | **x86_64** + NVIDIA B300 SXM6，compute_cap **10.3（SM103）**，275 GB/卡 | +| 主机内存 | 2015 GB | +| 权重存储 | `/home/scratch.trt_llm_data`，共享 **NFS** | + +**换集群这些数字全部作废，尤其是 GB300 / `b300-nvl8` / `gb300`**：那是 +Grace-Blackwell，**aarch64** 架构、CPU 与 GPU 之间是 NVLink-C2C 统一内存、 +主机内存容量与存储后端都不同。具体会变的： + +- **编译目标**：`103-real` 是为 x86 B300 编的。GB300 需要 aarch64 工具链和对应镜像 + （`LLM_SBSA_DOCKER_IMAGE`，见 `jenkins/current_image_tags.properties`）， + 不能直接复用本仓库现有的 `cpp/build`。 +- **加载耗时结构**（§6.4b 那张表）：预取 2 h / 装配 40 min 是这台机器的 NFS 和 + x86 内存带宽决定的。Grace 的内存带宽和统一内存寻址会改变整个比例。 +- **page cache 复用策略**：依赖 2 TB 主机内存能装下 835 GB 权重。GB300 的内存 + 配置不同，这个前提要重新验证。 +- **显存余量**：275 GB/卡 是 B300 SXM6 的数字。 +- **Slurm 账号与分区**：`trt-llm_b300` 对应 x86；GB300 走 `trt-llm_gb300`。 + +**结论性的东西（六条 P0 的定位、kernel 的 next_n 假设、可观测性设计、 +"空输出不是通过"那几条教训）与平台无关，可以直接迁移。 +所有带具体数字的运行特征都不行。** + +--- + ## 1. 目标 主线：让 **ragged verify**（每个 request 各自的 verify 长度）在 TensorRT-LLM 里真正跑通。 @@ -228,13 +260,16 @@ python3 scripts/build_wheel.py --cuda_architectures "100-real" -j $(nproc) cmake 报 `No CUDA compiler found`。 - **配置失败后必须 `rm -rf cpp/build`**：CMake 会把失败的编译器探测缓存成 `CMAKE_CUDA_COMPILER:FILEPATH=NOTFOUND` 并在下次复用，环境修好了也不会重新探测。 -- **目标平台是 B300（SM103），编译用 `--cuda_architectures "103-real"`。** +- **目标平台是 x86 DGX B300（SM103），编译用 `--cuda_architectures "103-real"`。** + GB300（`gb300` / `b300-nvl8`）是 aarch64，需要 SBSA 镜像和单独的编译产物。 B200 是 SM100；`100-real` 纯 SASS 无 PTX，装不进 B300。B300 有 275 GB/卡， B200 只有 178 GB —— DSv4-Pro 权重 104 GB/卡，B200 上余量很紧。 - 容器需 `--user $(id -u):$(id -g)`（scratch 是 NFS root_squash）。 ### 6.4b DSv4-Pro 的加载时间结构 —— 复用同一节点的 page cache +> **仅适用于 x86 DGX B300 + NFS（见 §0）。GB300 等其他平台需重新实测。** + 冷启动实测（`umb-b300-001`，权重在 `/home/scratch.trt_llm_data`，NFS）： | 阶段 | 耗时 | 可观测性 | From cb5cb0bb92647a4d5e8e009e56b01d2205172c92 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 02:53:06 -0700 Subject: [PATCH 005/177] [None][fix] H1/K8: stride the DSA generation expansions by the runtime verify tier goal doc H1 (dsa.py:662,696; model_engine.py:342) and K8 (dsa.py:2329-2330). `DSAtrtllmAttentionMetadata.max_draft_tokens` is a static ceiling, not this step's window. For the parallel-draft modes -- DSpark is one -- `forward()` passes `original_max_total_draft_tokens` to `update_spec_dec_param`, so it is pinned at `tokens_per_gen_step - 1` for the engine's whole lifetime and never follows a shorter confidence tier. The expansions in `prepare_for_spec_decode` used it as the stride, while `Indexer.forward` slices the very same buffers by the true token count (`num_gen_tokens // num_generations`). At tier 3 the producer laid kv_lens and the block table out 6 rows per request and the consumer read them back 4 rows per request, so every request past the first attended with a neighbour's KV length and block table. `prepare_scheduler_metadata` had the same skew in a place that hangs rather than corrupts: it built the DeepGEMM schedule over `num_generations * (1 + max_draft_tokens)` rows while the paged-MQA-logits call supplied `num_gen_tokens`, so the schedule promised tiles no one would produce. That matches the one observed end-to-end failure exactly -- warmup completed at draft_len=5 (768 == 768) and never returned at draft_len=3 (768 vs 512). Publishes `runtime_tokens_per_gen_step` onto the metadata alongside `ragged_verify_lens`, and routes every *stride* through the new `gen_token_stride`. Capacity keeps using the static ceiling: the expanded buffers must stay sized for the maximum so a shorter tier never reallocates during capture. Unset (0) falls back to `1 + max_draft_tokens`, so every path that does not shorten its draft length is bit-identical. Published in `_prepare_tp_inputs` rather than `update_spec_dec_param` because the latter runs against the base metadata while `prepare()` runs against the per-key CUDA-graph copy, and a plain attribute set on the base is invisible to that copy. `use_expanded_buffers_for_mtp` deliberately stays keyed on the ceiling so the layout regime does not become a function of the tier. Signed-off-by: Lance Liao --- .../_torch/attention_backend/sparse/dsa.py | 82 ++++++++++++++++--- .../_torch/pyexecutor/model_engine.py | 44 +++++++--- 2 files changed, 103 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 38e59a2b1b79..bec5fa2281ac 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -641,25 +641,56 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): # host because the expansions are built host-side and copied in; reading # the repeats off a device tensor would sync. ragged_verify_lens: Optional[List[int]] = None + # Query tokens per generation request on a *uniform* batch this step, i.e. + # `1 + runtime_draft_len`. Zero means "not published", and every expansion + # below falls back to the static `1 + max_draft_tokens` -- which is what + # every caller that never shortens its draft length wants, and keeps them + # bit-identical. + # + # This is deliberately NOT `max_draft_tokens`. That field is the static + # ceiling: it sizes every expanded buffer (see `create_expanded_buffers`) + # and must stay at the maximum so a shorter step never reallocates inside a + # CUDA-graph capture. The *stride* has to follow the runtime tier instead, + # because `Indexer.forward` slices the very same buffers by the true token + # count (`num_gen_tokens // num_generations`, see the decode branch). Left + # at the ceiling, a step that verifies fewer positions writes each + # request's kv_len/block-table row `1 + max_draft_tokens` apart and then + # reads them back `1 + runtime_draft_len` apart, so every request past the + # first silently picks up a neighbour's rows. + runtime_tokens_per_gen_step: int = 0 @property def is_ragged_verify(self) -> bool: """Whether generation requests have per-request query lengths.""" return self.ragged_verify_lens is not None + @property + def gen_token_stride(self) -> int: + """Query tokens contributed by one generation request, uniform batches. + + The runtime tier when the engine published one, else the static + ceiling. Capacity questions must keep using `1 + max_draft_tokens`; + only the stride follows the tier. + """ + stride = int(self.runtime_tokens_per_gen_step) + if stride <= 0: + return 1 + self.max_draft_tokens + return stride + def gen_token_repeat_list(self) -> List[int]: """Per-generation-request query-token counts, on the host. - Uniform batches get `1 + max_draft_tokens` for every request, so the - callers below stay bit-identical; ragged batches get the scheduler's - per-request windows. + Uniform batches get `gen_token_stride` for every request, so the + callers below stay bit-identical whenever the runtime draft length is + the static maximum; ragged batches get the scheduler's per-request + windows. Host-side on purpose: the expanded token count derived from these feeds buffer slicing, and reducing it off a device tensor would be a device->host sync on every step of the *uniform* path too. """ if self.ragged_verify_lens is None: - return [1 + self.max_draft_tokens] * self.num_generations + return [self.gen_token_stride] * self.num_generations return list(self.ragged_verify_lens) def gen_token_repeats(self, device=None) -> torch.Tensor: @@ -693,7 +724,7 @@ def expand_per_gen_token(self, repeats = self.gen_token_repeat_list() num_tokens = sum(repeats) if self.ragged_verify_lens is None: - stride = 1 + self.max_draft_tokens + stride = self.gen_token_stride return values.repeat_interleave(stride, dim=dim), num_tokens repeats_host = torch.tensor( repeats, @@ -1010,7 +1041,10 @@ def _compute_kv_lens_row_reorder(self): # num_generations alone is only correct for next_n == 2; for next_n == 1 # it engages inside the measured regression band, and for next_n == 4 it # misses the win region between 2*num_sms and 4*num_sms rows. - next_n = 1 + self.max_draft_tokens + # Runtime stride, not the static ceiling: this is a row-count threshold, + # and at a shorter verify tier the static value overstates the rows by + # the tier ratio, engaging GVR outside its measured win region. + next_n = self.gen_token_stride if (self.enable_heuristic_topk and self.use_cute_dsl_topk and self.num_generations * next_n >= 2 * self.num_sms): gen_kv_lens = self.kv_lens_cuda[self.num_contexts:self.num_seqs] @@ -1615,6 +1649,14 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): # ragged batch. The expanded path already materializes one row per # token, which is exactly the ragged layout, so route there. use_dsl = False + # Deliberately keyed on the static ceiling rather than this step's tier. + # Whether DeepGEMM supports a given next_n natively is a per-step + # question, so a shorter tier could sometimes take the cheaper strided + # path -- but that would make the *layout regime* a function of the tier, + # and the expanded path is the general one (one kernel row per query + # token) that is correct at every tier and required for ragged batches + # anyway. Holding the regime fixed keeps the tiers, and the ragged path + # built on top of them, comparable to each other. self.use_expanded_buffers_for_mtp = ( not use_dsl and (self.is_ragged_verify or @@ -1654,7 +1696,10 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): self.expand_for_dsl = (use_dsl and self.kv_cache_manager is not None and self.max_draft_tokens >= 1) if self.expand_for_dsl and self.num_generations > 0: - next_n = 1 + self.max_draft_tokens + # `Indexer.forward` validates the cached split against the runtime + # next_n (`factor * atom == next_n`), so the picker has to be fed the + # same value or the split is silently declined on every trimmed step. + next_n = self.gen_token_stride kernel_atoms = (1, 2, 3) if self.kv_cache_manager.use_fp4 else (1, 2, 3, 4) @@ -1692,7 +1737,7 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): self.block_table_expanded.clamp_(min=0) else: self.dsl_expand_factor = 1 - self.dsl_atom = 1 + self.max_draft_tokens + self.dsl_atom = self.gen_token_stride def _prepare_ragged_row_kv_lens(self, kv_lens: torch.Tensor) -> None: """Populate ``row_kv_lens_cuda``: how far back each query row may attend. @@ -2476,7 +2521,14 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): metadata.kv_lens_cuda_runtime[num_contexts:num_contexts + num_generations]) metadata.gen_indexer_kv_lens_cuda_runtime = gen_seq_lens - next_n_cap = metadata.kv_lens_cuda_2d.shape[1] + # The 2D DeepGEMM metadata API encodes next_n into the schedule, and + # `Indexer.forward` launches the strided path with the *runtime* + # next_n (`num_gen_tokens // num_generations`). Building the + # schedule at the buffer's full static width instead would describe + # more query positions than the kernel is given whenever a shorter + # verify tier is in play. + next_n_cap = min(metadata.gen_token_stride, + metadata.kv_lens_cuda_2d.shape[1]) metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap].copy_( gen_seq_lens.unsqueeze(-1).expand(-1, next_n_cap)) scheduler_metadata_buffer = get_paged_mqa_logits_metadata( @@ -2493,8 +2545,16 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): else: # Expand schedule metadata buffer (only generation). The DeepGEMM # API requires 2D; each expanded token becomes a (1,) row. - num_tokens = metadata.num_generations * (1 + - metadata.max_draft_tokens) + # + # The row count has to be the one the expansions actually populated, + # which is what `gen_token_repeat_list` reports: the runtime tier on + # a uniform batch, the scheduler's windows on a ragged one. Deriving + # it from `1 + max_draft_tokens` instead described `num_generations * + # (1 + max_draft_tokens)` rows while `Indexer.forward` handed the + # kernel only `num_gen_tokens` of them -- the schedule and the logits + # then disagree about how much work exists, which is a hang rather + # than an error because the kernel waits on tiles nobody produces. + num_tokens = sum(metadata.gen_token_repeat_list()) scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( metadata.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), _DG_SCHEDULE_BLOCK_KV, metadata.num_sms) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5296494e439e..761f8481cd4d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3803,9 +3803,10 @@ def _ragged_token_lens(generation_requests) -> Optional[List[int]]: return None return [1 + int(v) for v in verify_lens] - def _publish_ragged_verify_lens(self, attn_metadata, - generation_requests) -> None: - """Hand the attention metadata this step's windows, before it prepares. + def _publish_gen_token_layout(self, attn_metadata, + generation_requests) -> None: + """Hand the attention metadata this step's gen-token layout, before it + prepares. Split out of :meth:`_attach_ragged_verify_layout` because of ordering: ``attn_metadata.prepare()`` is the *only* consumer of @@ -3817,12 +3818,30 @@ def _publish_ragged_verify_lens(self, attn_metadata, The spec-metadata half stays where it was; it is consumed later and has no such constraint. + + The uniform stride is published for the same reason and at the same + moment. ``prepare()`` expands kv_lens and the block table once per query + token, and the only value that describes how many tokens a request + actually contributes this step is ``1 + runtime_draft_len``. The + attention metadata's own ``max_draft_tokens`` cannot stand in for it: it + is the static ceiling that sizes those buffers, and for the parallel-draft + modes it is pinned to ``tokens_per_gen_step - 1`` for the engine's whole + lifetime (see the ``update_spec_dec_param`` call in :meth:`forward`), so + it does not move when a shorter tier is chosen. + + Note this has to happen here rather than in ``update_spec_dec_param``: + that runs against the base metadata, while ``prepare()`` runs against the + per-key CUDA-graph copy taken later, and a plain attribute set on the + base is invisible to the copy. """ - if attn_metadata is None or not hasattr(attn_metadata, - "ragged_verify_lens"): + if attn_metadata is None: return - attn_metadata.ragged_verify_lens = self._ragged_token_lens( - generation_requests) + if hasattr(attn_metadata, "runtime_tokens_per_gen_step"): + attn_metadata.runtime_tokens_per_gen_step = ( + self.get_runtime_tokens_per_gen_step(self.runtime_draft_len)) + if hasattr(attn_metadata, "ragged_verify_lens"): + attn_metadata.ragged_verify_lens = self._ragged_token_lens( + generation_requests) def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, generation_requests) -> None: @@ -5402,13 +5421,14 @@ def previous_seq_slots_device(): # pre-prepare counts so the steady-gen recording below stores values # that the per-step prepare() can re-clamp from scratch. num_cached_tokens_snapshot = list(num_cached_tokens_per_seq) - # prepare() is the only consumer of the ragged windows, so they have to - # be published first: it decides the DSA expanded-buffer layout, builds - # the per-row causal extents, and derives DeepSeek-V4's per-request + # prepare() is the only consumer of the gen-token layout, so it has to + # be published first: it decides the DSA expanded-buffer layout, strides + # the expansions by this step's per-request token count, builds the + # per-row causal extents, and derives DeepSeek-V4's per-request # compressor token counts. The rest of the layout is attached later, # once the spec metadata exists. - self._publish_ragged_verify_lens(attn_metadata, - scheduled_requests.generation_requests) + self._publish_gen_token_layout(attn_metadata, + scheduled_requests.generation_requests) attn_metadata.prepare() cross_attention_inputs = (self._prepare_enc_dec_cross_attn_inputs( cross_encoder_hidden_states, From 44a81930ca74fa4389f435a7180388588bba8d9c Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 02:53:57 -0700 Subject: [PATCH 006/177] [None][fix] Remove the two remaining per-step host syncs on the DSpark path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit goal doc §4.3 (model_engine.py:3910-3919, heads.py:280-283). `_update_target_input_tensors` called `repeat_interleave` with a device `repeats` and no `output_size`, so torch read the cumulative sum back to the host to size the output. That is one device->host sync per ragged decode step, inside the overlap scheduler's window. The host already knows the answer -- the caller computes `total_num_tokens` as `sum(tokens_per_extend_request)` -- so pass it, matching the expansion a few hundred lines down that already does. `DSparkConfidenceHead.apply_sts` moved `sts_temperatures` to the CPU on every call. The planner calibrates on a pinned host snapshot, so this is a device to *pageable* host copy, which serializes the stream, once per decode step -- and only when a profiled cost table exists, i.e. exactly in the configuration the feature is meant to run in. Cache the mirror and invalidate it in `load_sts_temperatures`, the only writer. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/models/dspark/heads.py | 23 ++++++++++++++++++- .../_torch/pyexecutor/model_engine.py | 12 +++++++++- 2 files changed, 33 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/models/dspark/heads.py b/tensorrt_llm/_torch/models/dspark/heads.py index 29d217e7c877..0f843e8ed27c 100644 --- a/tensorrt_llm/_torch/models/dspark/heads.py +++ b/tensorrt_llm/_torch/models/dspark/heads.py @@ -276,12 +276,31 @@ def apply_sts(self, confidence_logits: torch.Tensor) -> torch.Tensor: without the transfer that is a device-mismatch ``RuntimeError`` on the first step that reaches calibration. The branch is host-side and the devices match on the in-graph path, so nothing extra is captured. + + The host copy is cached rather than re-fetched. Transferring the table + per call is a device->host copy into pageable memory, which serializes + the stream, and it lands on the planner's path once per decode step -- + and only once a profiled cost table exists, i.e. exactly in the + configuration where the feature is supposed to pay for itself. The table + is written only through :meth:`load_sts_temperatures`, which refreshes + this cache, so it cannot go stale. """ temperatures = self.sts_temperatures if temperatures.device != confidence_logits.device: - temperatures = temperatures.to(confidence_logits.device) + if confidence_logits.device.type == "cpu": + temperatures = self._host_sts_temperatures() + else: + temperatures = temperatures.to(confidence_logits.device) return torch.sigmoid(confidence_logits.float() / temperatures) + def _host_sts_temperatures(self) -> torch.Tensor: + """CPU mirror of ``sts_temperatures``, materialized at most once.""" + cached = getattr(self, "_sts_temperatures_host", None) + if cached is None: + cached = self.sts_temperatures.detach().to("cpu") + self._sts_temperatures_host = cached + return cached + @torch.no_grad() def load_sts_temperatures(self, temperatures: torch.Tensor) -> None: """In-place update of the STS table (never rebind: see ``__init__``).""" @@ -294,6 +313,8 @@ def load_sts_temperatures(self, temperatures: torch.Tensor) -> None: if not bool(torch.all(flat > 0.0)): raise ValueError("STS temperatures must be strictly positive") self.sts_temperatures.copy_(flat) + # Drop the CPU mirror; apply_sts rebuilds it on next use. + self._sts_temperatures_host = None def load_weights(self, weights: list) -> None: """Strict loader that refuses to silently drop a checkpoint bias. diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 761f8481cd4d..d5b0fcbed5cc 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3984,8 +3984,18 @@ def _update_target_input_tensors( dtype=torch.long, pin_memory=prefer_pinned()).to(new_tokens_lens_device.device, non_blocking=True) + # `output_size` is not optional here. With a tensor `repeats`, + # repeat_interleave has to know the output length before it can + # allocate, and without this it gets it by reading the cumulative + # sum back to the host -- a device->host sync on every ragged step, + # inside the overlap scheduler's window, which is precisely what + # overlap exists to avoid. The total is already known on the host: + # the caller derived `total_num_tokens` as + # `sum(tokens_per_extend_request)`. previous_pos_indices = torch.repeat_interleave( - previous_slots, tokens_per_request_device) + previous_slots, + tokens_per_request_device, + output_size=total_num_tokens) previous_kv_len_offsets = (new_tokens_lens_device[previous_slots] - tokens_per_request_device) else: From f99c34a498dd2e4d9bb04b61f1a3ae0702bac466 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 02:55:28 -0700 Subject: [PATCH 007/177] [None][fix] H8/K5/K6/K13: reject ragged-verify configs that cannot be honoured goal doc H8 (model_engine.py:3644-3648, llm_args.py:180-184) and K5/K6/K13 (indexerTopK.cu:918-922, trtllm.py:999-1001). The shared theme is the goal doc's own rule: an unsupported combination must be refused explicitly, because silent degradation is the worst outcome. H8: `fit_ragged_verify_lens` plans against the padded row count and reserves tokens for the CUDA-graph padding rows, but only checks `runner.enabled`. `BaseCudaGraphConfig.enable_padding` defaults to False, and `_get_padded_batch` then adds no rows -- so the step's token total keys into a graph nobody captured and runs eager, while the real requests were still shrunk to leave room for padding that never arrived. Require `enable_padding=True` (and a `cuda_graph_config` at all) when ragged verification is on. K5/K6: GVR strides its previous-step top-k hint by next_n, and the CuTe-DSL top-k takes next_n as a compile-time constant and JIT cache key. Neither can express per-request windows. The C++ launcher already refuses the pair, but only once a decode step reaches it; reject at config time instead. `use_cute_dsl_paged_mqa_logits` is only warned about -- it is force-routed to the numerically equivalent expanded path, so it degrades but does not lie. K13: the dense XQA spec-dec metadata is written by scalar `fill_` at every Python call site and cannot describe per-request windows. It is unreachable here today for two incidental reasons (the Blackwell trtllm-gen gate forces `is_spec_decoding_enabled` False for a linear tree, and MLA generation never reads the field). Assert the property that matters, so a future SM or a change to that gate fails loudly rather than verifying every request at the batch-wide window. Signed-off-by: Lance Liao --- .../_torch/attention_backend/sparse/dsa.py | 16 +++++ tensorrt_llm/llmapi/llm_args.py | 69 +++++++++++++++++++ 2 files changed, 85 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index bec5fa2281ac..8b83c879202e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1649,6 +1649,22 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): # ragged batch. The expanded path already materializes one row per # token, which is exactly the ragged layout, so route there. use_dsl = False + # goal doc K13. The dense XQA spec-dec metadata + # (`spec_decoding_generation_lengths`, packed mask) is written by + # scalar `fill_` at every Python call site, so it can only describe a + # uniform window -- even though the C++ side would accept a + # per-request vector. Today that block is unreachable here twice + # over: `update_spec_dec_param` forces `is_spec_decoding_enabled` + # False on the Blackwell trtllm-gen kernels for a linear tree, and + # MLA generation reads `mlaGeneration` state instead. Both are + # incidental. Assert the property that actually matters so a future + # SM, or a change to that gate, fails here rather than silently + # verifying every request at the batch-wide window. + assert not self.is_spec_decoding_enabled, ( + "ragged verification requires the dense spec-decoding metadata " + "to be off: spec_decoding_generation_lengths and the packed " + "mask are filled with a single scalar window, which cannot " + "describe per-request verify lengths") # Deliberately keyed on the static ceiling rather than this step's tier. # Whether DeepGEMM supports a given next_n natively is a per-step # question, so a shorter tier could sometimes take the cheaper strided diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 05b74ebbb779..98e11a2b2868 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5484,6 +5484,73 @@ def validate_encode_only_torch_compile_config(self) -> 'TorchLlmArgs': "graphs or disable enable_piecewise_cuda_graph.") return self + def _validate_dspark_ragged_verify_environment(self) -> None: + """Reject the configurations ragged verification cannot honour. + + Every check here guards a path that would otherwise degrade *silently*: + the run comes up, produces plausible text, and either never takes the + ragged path at all or takes it with a kernel that mis-attributes rows. + A hard error at construction is the only signal that survives. + """ + cuda_graph_config = self.cuda_graph_config + # `fit_ragged_verify_lens` plans against the padded row count and + # reserves tokens for the CUDA-graph padding rows, so it presumes those + # rows exist. With padding off, `_get_padded_batch` adds none: the batch + # then carries a token total that keys into a graph nobody captured and + # drops to eager on every ragged step -- which costs far more than the + # trimmed tokens save -- while real requests were still shrunk to leave + # room for padding that never arrived. + if cuda_graph_config is None: + raise ValueError( + "speculative_config.enable_ragged_verify=True requires " + "cuda_graph_config to be set with enable_padding=True. Ragged " + "verification fits each step's token total onto a captured " + "CUDA-graph bucket; with no graphs there is no bucket to fit " + "and every step would silently run eager.") + if not cuda_graph_config.enable_padding: + raise ValueError( + "speculative_config.enable_ragged_verify=True requires " + "cuda_graph_config.enable_padding=True (it defaults to False). " + "The ragged token budget is planned against the padded batch " + "size and reserves tokens for the padding rows; without padding " + "those rows never appear, so the step's token total misses its " + "captured bucket and falls out of graph replay.") + + sparse_attention_config = self.sparse_attention_config + if sparse_attention_config is not None: + # Both recover "which request is this row" from a fixed next_n -- + # GVR indexes its previous-step top-k hint by request and strides it + # by next_n, and the CuTe-DSL kernel takes next_n as a compile-time + # constant and JIT cache key. Neither can express per-request + # windows. The C++ launcher already refuses the combination + # (indexerTopK.cu, "ragged rowKvLens is incompatible with the GVR + # preIdx hint"), but that fires deep in a decode step; say so here. + incompatible = [ + name for name in ("enable_heuristic_topk", "use_cute_dsl_topk") + if getattr(sparse_attention_config, name, False) + ] + if incompatible: + raise ValueError( + "speculative_config.enable_ragged_verify=True is " + "incompatible with sparse_attention_config." + f"{' / '.join(incompatible)}=True: those top-k paths " + "reconstruct each row's request and window position from a " + "single batch-wide next_n, which a ragged batch does not " + "have. Disable them, or disable ragged verification.") + # Not an error: the DSL paged-MQA-logits path is structurally + # uniform too, but the expanded (one row per query token) path it + # falls back to is numerically equivalent, and it is the default on + # DeepSeek-V4 anyway. Still worth saying out loud, since the knob + # will read as honoured otherwise. + if getattr(sparse_attention_config, + "use_cute_dsl_paged_mqa_logits", False): + logger.warning( + "sparse_attention_config.use_cute_dsl_paged_mqa_logits=True " + "is ignored on ragged-verification steps: the DSL kernel " + "broadcasts one KV length per request across a fixed next_n. " + "Those steps use the expanded DeepGEMM path instead, which " + "is numerically equivalent.") + @model_validator(mode="after") def validate_speculative_config(self): if self.speculative_config: @@ -5742,6 +5809,8 @@ def _dspark_get(key, top_level_key): "DSpark block_size must equal max_draft_len; got " f"block_size={spec_cfg.block_size} and " f"max_draft_len={spec_cfg.max_draft_len}") + if spec_cfg.enable_ragged_verify: + self._validate_dspark_ragged_verify_environment() if isinstance(self.speculative_config, SADecodingConfig): pool_size = self.speculative_config.global_pool_size From b3a5fef92a390e1f73555306b28af7fe6b24476e Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 02:59:30 -0700 Subject: [PATCH 008/177] =?UTF-8?q?[None][fix]=20H9=20+=20=C2=A74.4:=20fix?= =?UTF-8?q?=20the=20DSpark=20collective=20divergence=20and=20the=20ADP=20s?= =?UTF-8?q?hape=20gate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit goal doc H9 (dspark_verify.py:234-241,256-259,132-133; py_executor.py:3258-3261) and §4.4 (cuda_graph_runner.py:362-372; model_engine.py:2987-2998). H9. `_dspark_confidence_draft_len` issued a different number of collectives per rank. Going ragged cost two (the planner's `all_rank_max`, then the `peer_stats` allgather); declining cost one (the uniform `decide_draft_len`); declining only at the bucket fit cost three. Every reason to decline is rank-local, and one of them -- whether the confidence snapshot's copy event has landed -- is purely timing-dependent, so two ranks on the same step could take different branches for no reason either could observe. That is a deadlock. Restructured into three phases: decide locally with no collective, issue one allgather of a fixed-width payload, then branch on the reduced result with no further collectives. The payload carries what `peer_stats` used to fetch separately, so the common path now costs one collective instead of two. Ranks only go ragged if all of them can; a single dissenter takes the group to a uniform step, which is the only outcome where the token layouts still agree. The planner grows `reduce_across_ranks=False` so it can hand back a local choice, plus `snap_to_tier` so the caller can land a reduced value back on the captured ladder. §4.4. `fit_ragged_verify_lens` can still decline on one rank after the agreement (its pad-row fit is sized from the local request count), so the attention-DP gate has to catch it. It compared only "all generation" and batch size -- never the token count -- so a rank that fell back produced a matching batch size with a different width and the gate passed. The mismatch is then invisible: `all_rank_num_tokens` is a host list read inside forward, so a replayed graph keeps the value captured with it. Compare the ragged token bucket and the draft length as well. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/cuda_graph_runner.py | 48 +++++++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 109 ++++++++++++------ .../_torch/speculative/dspark_verify.py | 40 ++++++- 3 files changed, 154 insertions(+), 43 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index c3a787ec7468..99900e4acac2 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -306,6 +306,21 @@ def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: total += 1 + int(verify_len) return total if total else None + @staticmethod + def _local_draft_len(batch: ScheduledRequests) -> int: + """This rank's draft length, for the attention-DP shape comparison. + + Deliberately tolerant where :meth:`get_graph_key` asserts: this runs + before the batch is known to be graphable, and a batch whose requests + disagree has no single draft length to compare. Reporting the maximum + makes such a batch differ from a well-formed peer, which is the + conservative answer -- the gate then declines the graph. + """ + return max( + (len(request.py_draft_tokens) + for request in batch.generation_requests), + default=0) + @staticmethod def _get_mrope_position_delta(request: Any) -> Optional[Any]: mrope_position_delta = getattr(request, "py_mrope_position_delta", None) @@ -360,15 +375,36 @@ def maybe_get_cuda_graph( can_run_cuda_graph = batch.can_run_cuda_graph batch_size = batch.batch_size if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: - all_can_graph_batch = self.config.dist.tp_allgather( - [can_run_cuda_graph, batch_size]) + # goal doc §4.4. Comparing only "all generation" and the batch size + # is not enough once the token count stops being a function of the + # row count. Under ragged verification a rank that fell back to + # uniform scheduling produces the same batch size but a different + # token total, so this gate used to pass and let one rank replay a + # graph full of collectives while another ran a different width. + # That mismatch is undetectable at replay time: `all_rank_num_tokens` + # is a host list read inside forward, so the captured graph keeps + # using the value it saw at capture. Compare every component of the + # key that a rank can decide for itself. + # + # `draft_len` is agreed by the DSpark planner's single collective, + # but it is also settable rank-locally by the acceptance-rate + # speculation gate, so it is compared here too rather than assumed. + # -1 stands in for "uniform batch" / "not applicable" so the payload + # is plain ints on every rank and compares elementwise. + local_bucket = self._ragged_verify_bucket( + batch) if can_run_cuda_graph else None + all_can_graph_batch = self.config.dist.tp_allgather([ + can_run_cuda_graph, batch_size, + -1 if local_bucket is None else int(local_bucket), + self._local_draft_len(batch) if can_run_cuda_graph else -1 + ]) is_all_gen_only = all(all_can_graph[0] for all_can_graph in all_can_graph_batch) - all_batch_size_equal = all( - all_gen_only[1] == all_can_graph_batch[0][1] - for all_gen_only in all_can_graph_batch) + all_shapes_equal = all( + peer[1:] == all_can_graph_batch[0][1:] + for peer in all_can_graph_batch) - if not is_all_gen_only or not all_batch_size_equal: + if not is_all_gen_only or not all_shapes_equal: return None, None, None if not self.enabled or not can_run_cuda_graph: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f99a2f805dab..c05bb9b90f0f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3217,24 +3217,6 @@ def _dspark_confidence_draft_len(self, for request in gen_requests ] - # Cross-rank agreement: draft_len is part of the CUDA-graph key but is - # NOT covered by the attention-DP consistency allgather, so ranks that - # pick different lengths select different graphs -- one replays, one - # falls back to eager -- and their collectives diverge. Reduce here. - # - # This costs one extra blocking host collective per decode step under - # ADP. ``_can_queue`` already runs ``tp_allgather(batch_size)`` just - # before this hook, so the two could be fused into a single - # ``[batch_size, local_draft_len]`` gather -- but that means computing - # the local draft length inside ``_can_queue``, which is a hot generic - # path shared by every speculation mode. Left separate deliberately; - # fuse only if the extra collective shows up in a profile. - all_rank_max = None - if self.dist is not None and getattr(self.dist, "tp_size", 1) > 1: - - def all_rank_max(value: int) -> int: - return max(int(v) for v in self.dist.tp_allgather(int(value))) - # The mode, not the config flag, decides whether windows are computed: # the environment override exists so a run can be switched between # `static` and `compact` without editing the serving config, and every @@ -3245,34 +3227,97 @@ def all_rank_max(value: int) -> int: getattr(self.model_engine.spec_config, "enable_ragged_verify", False)) + is_distributed = (self.dist is not None + and getattr(self.dist, "tp_size", 1) > 1) + + # --- Phase 1: decide locally, issue nothing --------------------------- + # + # goal doc H9. Every reason to give up on a ragged split is rank-local, + # and one of them -- whether the confidence snapshot's copy event has + # landed (`_ready_snapshot`) -- is timing-dependent, so it can differ + # between ranks on the same step for no reason the ranks can see. A + # collective placed after those checks is therefore entered by some + # ranks and skipped by others: the rank that went ragged used to issue + # two (the planner's max, then `peer_stats`), a rank that declined + # issued one, and a rank that declined only at the bucket fit issued + # three. That is a deadlock, and it is a timing-dependent one. + # + # So: nothing below this point may branch on rank-local state before the + # single collective, and nothing after it may issue another one. ragged_lens = None fallback_reason = None + local_tier = planner.max_tier if compute_windows: ragged_lens = planner.decide_verify_lens( num_gen_requests=num_gen_requests, rows=confidence_rows, - all_rank_max=all_rank_max) + reduce_across_ranks=False) if ragged_lens is None: fallback_reason = "planner_declined" else: fallback_reason = "mode_static" + # Only computed on the uniform path so the planner's fallback + # counters keep meaning "this step declined to trim" exactly once. + # On the ragged path the placeholder above is the largest tier -- + # i.e. verify everything -- which is the documented safe answer for + # the case this value is actually used: some *other* rank declined + # and dragged the whole group back to a uniform step. + local_tier = planner.decide_draft_len( + num_gen_requests=num_gen_requests, + rows=confidence_rows, + reduce_across_ranks=False) + # --- Phase 2: one collective, same shape on every rank, every step ---- + # + # Carries everything the rest of the step needs, so the `peer_stats` + # gather that used to follow is folded in here rather than added to it: + # [wants_ragged, num_rows, total_verify_tokens, max_window, tier] + # `_can_queue` already runs `tp_allgather(batch_size)` just before this + # hook and could absorb this too, but that would mean computing the + # local draft length inside a hot path shared by every speculation mode. + local_rows = len(ragged_lens) if ragged_lens is not None else 0 + local_tokens = (sum(1 + int(v) + for v in ragged_lens) if ragged_lens else 0) + local_max_window = max(ragged_lens) if ragged_lens else 0 + peer_stats = None + if is_distributed: + payloads = self.dist.tp_allgather([ + 1 if ragged_lens is not None else 0, + local_rows, + local_tokens, + local_max_window, + int(local_tier), + ]) + if not all(int(payload[0]) for payload in payloads): + # Any rank that cannot go ragged takes the whole group with it: + # a mixed step means one rank replays a ragged graph while + # another runs a different token count, and `all_rank_num_tokens` + # is frozen into the captured graph, so the mismatch would not + # even be observable at replay time. + if ragged_lens is not None: + fallback_reason = "peer_declined" + ragged_lens = None + else: + agreed_max = max(int(payload[3]) for payload in payloads) + ragged_lens = [min(int(v), agreed_max) for v in ragged_lens] + peer_stats = [[int(payload[1]), + int(payload[2])] for payload in payloads] + agreed_tier = max(int(payload[4]) for payload in payloads) + else: + agreed_tier = int(local_tier) + + # --- Phase 3: no more collectives ------------------------------------- + # # Land the batch's token total on a captured bucket before the graph key # is built, spending the rounding slack on real verification. A batch # with no captured shape gets None back and falls through to uniform, # which always has a graph -- running ragged without one would cost far - # more than the trimmed tokens save. + # more than the trimmed tokens save. This can still come out differently + # on different ranks (the pad-row fit is sized from the local request + # count), which is why the attention-DP graph gate compares the token + # bucket and not just the batch size. bucket = None if ragged_lens is not None: - # Under ADP the batch is padded to the widest rank's request count, - # so the shape has to be chosen from every rank's (rows, tokens) -- - # sizing it locally makes each rank pick a different bucket, i.e. - # a different graph. - peer_stats = None - if self.dist is not None and getattr(self.dist, "tp_size", 1) > 1: - peer_stats = self.dist.tp_allgather( - [len(ragged_lens), - sum(1 + int(v) for v in ragged_lens)]) bucket = self.model_engine.fit_ragged_verify_lens( gen_requests, ragged_lens, peer_stats=peer_stats) if bucket is None: @@ -3322,10 +3367,8 @@ def all_rank_max(value: int) -> int: # layout with uniform acceptance -- silent token misattribution. for request in gen_requests: request.py_verify_len = None - runtime_draft_len = planner.decide_draft_len( - num_gen_requests=num_gen_requests, - rows=confidence_rows, - all_rank_max=all_rank_max) + # Already agreed in phase 2 -- deliberately not another collective. + runtime_draft_len = planner.snap_to_tier(agreed_tier) # Stage this step's confidence for the *next* decision. Non-blocking; if # it has not landed by then the planner just verifies the full block. diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index 3f3f8344ceb4..b0a86da93d62 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -171,12 +171,26 @@ def _gather_rows( return None return selected + def snap_to_tier(self, value: int) -> int: + """Round ``value`` up to a captured tier. + + Callers that reduce across ranks themselves need this: the reduced value + is a maximum over per-rank choices and so is already a tier, but a caller + combining it with anything else has to land back on the ladder before it + can be used as a graph key. + """ + value = int(value) + if value in self.tiers: + return value + return min((t for t in self.tiers if t >= value), default=self.max_tier) + def decide_draft_len( self, *, num_gen_requests: int, rows: Optional[Sequence[int]] = None, all_rank_max: Optional[Callable[[int], int]] = None, + reduce_across_ranks: bool = True, ) -> int: """Choose this iteration's draft length, agreed across all ranks. @@ -189,9 +203,17 @@ def decide_draft_len( :meth:`_gather_rows`. ``all_rank_max`` overrides the constructor's reduction; the caller that owns the distributed handle usually supplies it here. + + ``reduce_across_ranks=False`` returns the *local* choice and issues no + collective, leaving the agreement to a caller that batches it with its + other cross-rank traffic. Since the reduction is a max over tiers, and + every local choice is already a tier, the caller only has to reduce and + pass the result through :meth:`snap_to_tier`. """ self.stats["decisions"] += 1 chosen = self._decide_local(num_gen_requests=num_gen_requests, rows=rows) + if not reduce_across_ranks: + return int(chosen) all_rank_max = all_rank_max or self.all_rank_max if all_rank_max is not None: # Max, not min: a rank that wanted to trim more simply verifies a few @@ -208,6 +230,7 @@ def decide_verify_lens( num_gen_requests: int, rows: Optional[Sequence[int]] = None, all_rank_max: Optional[Callable[[int], int]] = None, + reduce_across_ranks: bool = True, ) -> Optional[List[int]]: """Per-request verify lengths for a ragged step, or None to stay uniform. @@ -273,10 +296,19 @@ def decide_verify_lens( budget = int(num_gen_requests) * (int(uniform_len) - self.cfg.min_verify_len) lens = schedule_verify_lens_topk(survival=survival, budget=budget, cfg=self.cfg).tolist() - all_rank_max = all_rank_max or self.all_rank_max - if all_rank_max is not None: - agreed_max = int(all_rank_max(int(max(lens)))) - lens = [min(int(v), agreed_max) for v in lens] + # `reduce_across_ranks=False` means the caller is doing the cross-rank + # agreement itself. That is not just an optimization here: every early + # return above is rank-local (an empty batch, an unprofiled cost table, + # a confidence snapshot whose copy event has not landed yet -- that last + # one is timing-dependent), so a collective issued *after* them is issued + # by some ranks and not others. Deciding locally and reducing once, + # unconditionally, at a point every rank reaches is the only shape of + # this that cannot deadlock. + if reduce_across_ranks: + all_rank_max = all_rank_max or self.all_rank_max + if all_rank_max is not None: + agreed_max = int(all_rank_max(int(max(lens)))) + lens = [min(int(v), agreed_max) for v in lens] assert len(lens) == num_gen_requests, ( f"internal: produced {len(lens)} verify lengths for {num_gen_requests} requests" ) From 3b52d19fcd77997f352adf4ccbad9daf9f67b968 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 03:00:10 -0700 Subject: [PATCH 009/177] [None][fix] H7: let rejection sampling survive a ragged batch goal doc H7 (interface.py:1707-1710, 1793-1806). `_rejection_buffers_valid` required the target to have emitted `num_contexts + num_gens * (draft_len + 1)` logits rows. Under ragged verification it emits `num_contexts + total_verify_tokens`, and the difference is exactly what the scheduler trimmed -- so the guard failed precisely when the feature was doing its job. Being fail-closed, it fell back to strict acceptance without a word, which changes the sampled distribution for temperature > 0 and left the rejection path's ragged branch as dead code. That branch is complete: it packs the gen logits by `total_verify_tokens`, scatters them into the padded rectangle, fills the padding rows one-hot so they cannot be accepted, and clamps the accepted count back down to `verify_lens`. So the fix is the row count, not the config -- rejecting "ragged + non-greedy" would have discarded working code. Still fail-closed on a genuinely half-built layout: a ragged batch missing `verify_lens`, `qo_indptr` or `total_verify_tokens` (all three are dereferenced downstream) is refused as before. Note greedy GSM8K cannot observe any of this, which is why it needs to be argued from the code rather than measured. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/speculative/interface.py | 27 ++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 691b2decd236..989cabdbfb7f 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1704,9 +1704,32 @@ 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. + # + # goal doc H7. Under ragged verification that is `total_verify_tokens`, + # not `num_gens * (draft_len + 1)` -- and the two differ by exactly the + # amount the scheduler trimmed, so requiring the uniform count made this + # guard fail whenever the feature did anything. Rejection sampling then + # fell back to strict acceptance silently, changing the sampled + # distribution for temperature > 0 while leaving the ragged branch below + # (which is complete: it scatters to the padded rectangle, fills the pad + # rows one-hot, and clamps the accepted count back to `verify_lens`) + # unreachable. Greedy GSM8K cannot see any of this. logits_rows = logits.shape[0] if logits.dim() > 1 else 1 - if logits_rows < num_contexts + num_gens * (draft_len + 1): + if getattr(spec_metadata, "is_ragged_verify", False): + # Fail closed on a half-built ragged layout: the branch below + # dereferences all three of these. + total_verify_tokens = getattr(spec_metadata, "total_verify_tokens", + None) + if (total_verify_tokens is None + or getattr(spec_metadata, "verify_lens", None) is None + or getattr(spec_metadata, "qo_indptr", None) 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). From 9089763e4f03c325a93a86637f22d1ec0d486f80 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 03:04:04 -0700 Subject: [PATCH 010/177] [None][test] B2: microbenchmark the two sparse-MLA generation presentations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit goal doc §6.4 B2 / §7.1 U1 -- the measurement that decides whether ragged verification is worth building at all. Both ragged designs get past trtllm-gen's uniform-layout assert (`num_tokens % num_seqs == 0`, then `mMaxSeqLenQ = acc_q_len / batch_beam`) the same way: by presenting the generation half as one length-1 sequence per query token. Whether that costs throughput cannot be read off the source -- the kernel selection is in a precompiled cubin -- so it has to be measured. Times identical work laid out both ways: B requests at s_q = tier+1 versus B*(tier+1) requests at s_q = 1, with request (i, j) cached at L + j so every row has the same KV extent in both. Single GPU, no model weights, minutes. Reuses the sparse-MLA test harness rather than restating it, since that is the only place a real DeepseekV4TrtllmAttention plus cache manager is assembled. Scope is stated in the module docstring: this times the whole attention layer, so the indexer is included and changes presentation too. That is deliberate -- it is what a ragged step would actually pay, and it subsumes B3. Signed-off-by: Lance Liao --- .../dsv4_sparse_mla_presentation.py | 442 ++++++++++++++++++ 1 file changed, 442 insertions(+) create mode 100644 tests/microbenchmarks/dsv4_sparse_mla_presentation.py diff --git a/tests/microbenchmarks/dsv4_sparse_mla_presentation.py b/tests/microbenchmarks/dsv4_sparse_mla_presentation.py new file mode 100644 index 000000000000..2ce043417941 --- /dev/null +++ b/tests/microbenchmarks/dsv4_sparse_mla_presentation.py @@ -0,0 +1,442 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""B2: batch-major vs token-major presentation of DeepSeek-V4 sparse-MLA generation. + +This is the measurement the DSpark ragged-verification design hinges on +(``docs/dspark_confidence_schedule_goal.md`` §6.4 B2 / §7.1 U1). + +The question +------------ +Ragged verification means each request sends its own number of query tokens to +the target in one decode step. Every generation kernel that recovers ``(request, +offset)`` from a single ``next_n`` breaks under that. The deepest one is +trtllm-gen sparse-MLA generation, which cannot be read: its kernel selection +lives in a precompiled cubin, and the host side asserts a uniform layout:: + + TLLM_CHECK_WITH_INFO(num_tokens % num_seqs == 0, + "seq_len should be same for all generation requests", ...) + int32_t const input_seq_length = num_tokens / num_seqs; + -- cpp/tensorrt_llm/thop/attentionOp.cpp + mMaxSeqLenQ = acc_q_len / batch_beam; + -- cpp/tensorrt_llm/common/attentionOp.cpp + +Both candidate ragged designs (D1 token-flat, D2 token-major) dodge that assert +the same way: present the generation half as ``num_gen_tokens`` sequences of +length one, so ``seq_len == 1`` and the check passes trivially. They are +therefore the same bet, and this script is how that bet gets priced. + +What is compared +---------------- +Identical total query tokens and an identical per-token sparse index budget, +laid out two ways: + +* **batch-major** -- ``B`` requests, ``s_q = tier + 1``, each cached at ``L``. + This is what the engine does today. +* **token-major** -- ``B * (tier + 1)`` requests, ``s_q = 1``, request ``(i, j)`` + cached at ``L + j``. Each query token is its own sequence, so it carries its + own KV length and block table and no kernel has to divide anything. + +Row ``j`` of request ``i`` attends over ``L + j + 1`` tokens in both layouts, so +the arithmetic is the same; only the presentation differs. + +How to read the result +---------------------- +``docs/dspark_p0_task_prompt.md`` §3: if token-major is materially slower or +incorrect, stop and fall back to the uniform tier ladder (D3 Stage A), which +needs no kernel change at all. The prior is that DeepSeek-V4-Pro under DEP8 +keeps all 128 Q heads per rank, so BMM1's M dimension falls from +``128 * (tier+1)`` to ``128`` -- and because the sparse index set differs per +row there was never any KV tile sharing across the window to lose. That is an +inference, which is exactly why it needs measuring. + +Scope: this times the whole attention layer, so it includes the indexer +(paged-MQA-logits + top-k), which itself changes presentation -- the batch-major +layout takes the expanded one-row-per-token path, the token-major layout takes +the strided ``next_n == 1`` path. That bundling is deliberate: it is what a +ragged step would actually pay, and it subsumes B3. The pure-attention delta +cannot be separated without a kernel-level profiler. + +Usage (single GPU, no model weights):: + + python tests/microbenchmarks/dsv4_sparse_mla_presentation.py + python tests/microbenchmarks/dsv4_sparse_mla_presentation.py \ + --batch-sizes 8 32 128 --tiers 1 3 5 --context-len 4096 --check +""" + +import argparse +import json +import math +import os +import sys +from typing import Dict, List, Optional, Tuple + +import torch + +TESTS_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +# Reuse the sparse-MLA test harness verbatim rather than restating it: it is the +# only place that builds a real DeepseekV4TrtllmAttention plus cache manager, +# and a second copy would drift. +sys.path.insert(0, os.path.join(TESTS_ROOT, "unittest", "_torch", "attention", + "sparse", "deepseek_v4")) + +from test_deepseek_v4_sparse_mla import ( # noqa: E402 + Scenario, _build_compressed_topk_indices, _create_cache_manager, + _create_pos_embd_params, _create_rope_cos_sin, _prefill_compress_buffer) + +from tensorrt_llm._torch.attention_backend.interface import ( # noqa: E402 + AttentionInputType, MLAParams) +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4 import \ + DeepseekV4TrtllmAttention # noqa: E402 +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import \ + DeepseekV4TrtllmAttentionMetadata # noqa: E402 +from tensorrt_llm._torch.metadata import KVCacheParams # noqa: E402 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest # noqa: E402 +from tensorrt_llm.bindings import SamplingConfig # noqa: E402 +from tensorrt_llm.mapping import Mapping # noqa: E402 + +# Layer 1 of the shared Scenario is the compress_ratio=4 layer, i.e. the one +# that actually runs sparse selection. Layers 0 (dense) and 2 (ratio 128) do not +# exercise the path this question is about. +SPARSE_LAYER_IDX = 1 + + +def _yarn_mscale(scale: float = 1.0, mscale: float = 1.0) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class PresentationHarness: + """One cache manager + attention layer, driven at a chosen presentation.""" + + def __init__(self, scenario: Scenario, cached_lens: List[int], + seq_len_q: int, device: torch.device): + """``cached_lens[i]`` is request i's KV length before this step; every + request contributes ``seq_len_q`` query tokens.""" + self.scenario = scenario + self.device = device + self.cached_lens = list(cached_lens) + self.seq_len_q = seq_len_q + self.num_requests = len(cached_lens) + self.total_q_tokens = self.num_requests * seq_len_q + + self.qk_rope_head_dim = scenario.qk_rope_head_dim + # rope_append=False: the last 64 of the 512 lanes carry the rope part, + # and the head count halves. Mirrors the test harness. + self.kv_lora_rank = scenario.kv_lora_rank - self.qk_rope_head_dim + self.head_dim = self.kv_lora_rank + self.qk_rope_head_dim + self.num_heads = 64 + + max_seq_len = max(cached_lens) + seq_len_q + 1 + # The cache manager sizes its pool from the *context* lengths it is + # handed, so pass the per-request KV lengths as if they were prompts. + self.cache_manager, self.sparse_config = _create_cache_manager( + scenario, self.cached_lens, max_seq_len) + self.request_ids = list(range(self.num_requests)) + self.requests = [ + LlmRequest( + request_id=i, + max_new_tokens=seq_len_q + 1, + input_tokens=list(range(cached_len)), + sampling_config=SamplingConfig(), + is_streaming=False, + ) for i, cached_len in enumerate(self.cached_lens) + ] + for request in self.requests: + self.cache_manager.prepare_context(request) + self.cache_manager.resize_context(request, + request.context_chunk_size) + + self.mapping = Mapping(world_size=1, tp_size=1, rank=0) + self.rope_cos_sin = _create_rope_cos_sin(scenario, device) + pos_embd_params = _create_pos_embd_params(scenario) + mscale = _yarn_mscale(pos_embd_params.rope.scale, + pos_embd_params.rope.mscale_all_dim) + q_scaling = 1.0 / (mscale * mscale) + + self.layer = DeepseekV4TrtllmAttention( + layer_idx=SPARSE_LAYER_IDX, + num_heads=self.num_heads, + head_dim=self.head_dim, + num_kv_heads=1, + q_scaling=q_scaling, + pos_embd_params=pos_embd_params, + mla_params=MLAParams( + q_lora_rank=scenario.q_lora_rank, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + qk_nope_head_dim=scenario.qk_nope_head_dim, + v_head_dim=scenario.v_head_dim, + rope_append=scenario.rope_append, + predicted_tokens_per_seq=1, + hidden_size=scenario.hidden_size, + ), + sparse_attention_config=self.sparse_config, + skip_create_weights_in_init=True, + ) + self.layer.update_quant_config(None) + self.layer.attn_sink = torch.nn.Parameter(torch.randn( + self.num_heads, dtype=torch.float32, device=device).mul_(0.5), + requires_grad=False) + + _prefill_compress_buffer(self.cache_manager, SPARSE_LAYER_IDX, + self.cached_lens, self.request_ids, + self.head_dim, device) + + self._build_inputs() + self._build_metadata() + + def _build_inputs(self) -> None: + dtype = self.scenario.dtype + device = self.device + n = self.total_q_tokens + self.compressed_kv = torch.empty([n, self.kv_lora_rank], + dtype=dtype, + device=device).uniform_(-1, 1) + self.k_pe = torch.empty([n, self.qk_rope_head_dim], + dtype=dtype, + device=device).uniform_(-1, 1) + q = torch.empty([n, self.num_heads, self.kv_lora_rank], + dtype=dtype, + device=device).uniform_(-1, 1) + self.q_pe = torch.empty([n, self.num_heads, self.qk_rope_head_dim], + dtype=dtype, + device=device).uniform_(-1, 1) + self.fused_q = torch.cat([q, self.q_pe], + dim=-1).view(-1, + self.num_heads * self.head_dim) + self.latent_cache = torch.cat([self.compressed_kv, self.k_pe], dim=-1) + + def _build_metadata(self) -> None: + device = self.device + self.metadata = DeepseekV4TrtllmAttentionMetadata( + seq_lens=torch.tensor([self.seq_len_q] * self.num_requests, + dtype=torch.int), + request_ids=self.request_ids, + max_num_requests=self.num_requests, + num_contexts=0, + prompt_lens=self.cached_lens, + max_num_tokens=self.total_q_tokens, + kv_cache_manager=self.cache_manager, + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=self.cached_lens, + ), + mapping=self.mapping, + enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), + sparse_attention_config=self.sparse_config, + ) + self.metadata.prepare() + + num_seqs = self.metadata.kv_lens_cuda_runtime.size(0) + self.cu_q_seqlens = torch.empty(num_seqs + 1, + dtype=torch.int32, + device=device) + self.cu_kv_seqlens = torch.empty(num_seqs + 1, + dtype=torch.int32, + device=device) + self.fmha_scheduler_counter = torch.empty(1, + dtype=torch.uint32, + device=device) + # Per-query-token KV extent, i.e. exactly what a correct ragged layout + # must produce: row j of a request cached at L sees L + j + 1 tokens. + per_row_kv_lens = [ + cached_len + j + 1 for cached_len in self.cached_lens + for j in range(self.seq_len_q) + ] + self.topk_indices = _build_compressed_topk_indices( + [kv_len - 1 for kv_len in per_row_kv_lens], + self.scenario.compress_ratios[SPARSE_LAYER_IDX], + self.scenario.index_topk, + device, + ) + + def step(self) -> torch.Tensor: + """One generation step. Repeatable: every write target comes from the + metadata, which is fixed, so re-running rewrites the same slots.""" + self.layer.mla_rope_generation( + self.fused_q, + self.q_pe, + self.latent_cache, + self.metadata, + self.cu_q_seqlens, + self.cu_kv_seqlens, + self.fmha_scheduler_counter, + None, + None, + None, + ) + return self.layer.forward( + self.fused_q, + None, + None, + self.metadata, + attention_input_type=AttentionInputType.generation_only, + latent_cache=self.latent_cache, + q_pe=self.q_pe, + cu_q_seqlens=self.cu_q_seqlens, + cu_kv_seqlens=self.cu_kv_seqlens, + fmha_scheduler_counter=self.fmha_scheduler_counter, + topk_indices=self.topk_indices, + ) + + def shutdown(self) -> None: + self.cache_manager.shutdown() + + +def _time_step(harness: PresentationHarness, warmup: int, + iters: int) -> Tuple[float, float]: + """Median and minimum per-step latency in milliseconds. + + Median rather than mean: a single stray context switch on a shared node + should not decide a design. + """ + for _ in range(warmup): + harness.step() + torch.cuda.synchronize() + + samples: List[float] = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + harness.step() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end)) + samples.sort() + return samples[len(samples) // 2], samples[0] + + +def _run_point(batch_size: int, tier: int, context_len: int, warmup: int, + iters: int, check: bool) -> Dict[str, object]: + next_n = tier + 1 + scenario = Scenario() + device = torch.device("cuda") + torch.manual_seed(42) + + record: Dict[str, object] = { + "batch_size": batch_size, + "tier": tier, + "next_n": next_n, + "context_len": context_len, + "total_query_tokens": batch_size * next_n, + } + + # Batch-major: B sequences, next_n query tokens each. + batch_major = PresentationHarness(scenario, [context_len] * batch_size, + next_n, device) + # Token-major: one sequence per query token. Request (i, j) is cached at + # L + j so its single query token has the same extent as row j above. + token_cached = [ + context_len + j for _ in range(batch_size) for j in range(next_n) + ] + token_major = PresentationHarness(scenario, token_cached, 1, device) + + try: + if check: + out_batch = batch_major.step() + out_token = token_major.step() + record["batch_major_output_shape"] = list(out_batch.shape) + record["token_major_output_shape"] = list(out_token.shape) + record["shapes_match"] = (out_batch.shape == out_token.shape) + for name, out in (("batch_major", out_batch), ("token_major", + out_token)): + record[f"{name}_finite"] = bool(torch.isfinite(out).all()) + record[f"{name}_abs_mean"] = float(out.abs().mean()) + # The two runs use different random KV, so the outputs are not + # expected to match elementwise -- only to be the same shape, finite, + # and of the same magnitude. A token-major layout that mis-attributes + # rows shows up here as a magnitude that is off, or as non-finite + # values, because each pseudo-request's extent differs by one. + + batch_median, batch_min = _time_step(batch_major, warmup, iters) + token_median, token_min = _time_step(token_major, warmup, iters) + finally: + batch_major.shutdown() + token_major.shutdown() + + record["batch_major_ms_median"] = batch_median + record["batch_major_ms_min"] = batch_min + record["token_major_ms_median"] = token_median + record["token_major_ms_min"] = token_min + record["token_major_slowdown"] = (token_median / + batch_median if batch_median else + float("inf")) + return record + + +def main(argv: Optional[List[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--batch-sizes", + type=int, + nargs="+", + default=[8, 32, 128]) + parser.add_argument("--tiers", + type=int, + nargs="+", + default=[1, 3, 5], + help="verify-length tiers; next_n = tier + 1") + parser.add_argument("--context-len", type=int, default=4096) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--iters", type=int, default=30) + parser.add_argument("--check", + action="store_true", + help="also run a one-shot sanity check per point") + parser.add_argument("--json", type=str, default=None) + args = parser.parse_args(argv) + + if not torch.cuda.is_available(): + print("CUDA is required for this benchmark", file=sys.stderr) + return 2 + + print(f"device: {torch.cuda.get_device_name(0)} " + f"(sm{''.join(str(c) for c in torch.cuda.get_device_capability())})") + print(f"context_len={args.context_len} warmup={args.warmup} " + f"iters={args.iters}\n") + header = (f"{'B':>5} {'tier':>5} {'next_n':>7} {'tokens':>7} " + f"{'batch-major ms':>15} {'token-major ms':>15} {'slowdown':>9}") + print(header) + print("-" * len(header)) + + records: List[Dict[str, object]] = [] + for batch_size in args.batch_sizes: + for tier in args.tiers: + try: + record = _run_point(batch_size, tier, args.context_len, + args.warmup, args.iters, args.check) + except Exception as exc: # noqa: BLE001 - report and keep going + print(f"{batch_size:>5} {tier:>5} {'':>7} {'':>7} " + f"FAILED: {type(exc).__name__}: {exc}") + records.append({ + "batch_size": batch_size, + "tier": tier, + "error": f"{type(exc).__name__}: {exc}", + }) + continue + records.append(record) + print(f"{record['batch_size']:>5} {record['tier']:>5} " + f"{record['next_n']:>7} {record['total_query_tokens']:>7} " + f"{record['batch_major_ms_median']:>15.3f} " + f"{record['token_major_ms_median']:>15.3f} " + f"{record['token_major_slowdown']:>9.2f}x") + + ok = [r for r in records if "error" not in r] + if ok: + worst = max(r["token_major_slowdown"] for r in ok) + best = min(r["token_major_slowdown"] for r in ok) + print(f"\ntoken-major slowdown across points: " + f"{best:.2f}x .. {worst:.2f}x") + print("Read: >~1.3x sustained means the ragged (D1/D2) route is paying " + "more than it can save; fall back to the uniform tier ladder " + "(D3 Stage A), which needs no kernel change.") + + if args.json: + with open(args.json, "w") as handle: + json.dump(records, handle, indent=2) + print(f"wrote {args.json}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 7c23527868698b96f423e1b2e9168b5ba26b935e Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 03:06:36 -0700 Subject: [PATCH 011/177] [None][test] A3/A4/A6: make ragged-layout inconsistency loud instead of silent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit goal doc §6.2 A3, A4, A6. Every layout bug in this feature is silent by construction. The token total is always an integer multiple of the row count, so the divisibility checks inside the kernels pass and rows are simply attributed to the wrong requests. Nothing raises; the model just answers slightly differently. A3: assert the four independent descriptions of "generation tokens in this step" agree -- the requests' fitted windows, the flat input layout, the spec metadata's total, and the CUDA-graph key's token axis. Each of H0, H5 and H8 breaks exactly one of them while leaving the rest self-consistent. Also asserts the converse on uniform steps, since a half-published layout is the dangerous state. Behind TLLM_DSPARK_ASSERT_LAYOUT=1: it walks the batch on the hot path. A4: assert in `_prepare_ragged_row_kv_lens` that the windows sum to this step's generation token count. `seq_lens` is rebuilt per step and so is current by definition, which makes this a test of the property (these windows are this step's) rather than of the mechanism (publish runs before prepare). Under CUDA graphs the failure mode it guards is not "last step" but "whenever this key last ran", which is far harder to spot by reading. A6: fold the planner's own counters into the ragged stats summary. The ragged counters say what happened; the planner's say why. `fallback_flat_cost` especially -- without a profiled SPS table the budget search cannot trim at all, so a run can look completely healthy in every other counter and still have delivered exactly nothing. Signed-off-by: Lance Liao --- .../_torch/attention_backend/sparse/dsa.py | 12 ++++ .../_torch/pyexecutor/model_engine.py | 68 +++++++++++++++++++ tensorrt_llm/_torch/speculative/dspark.py | 3 + .../speculative/dspark_observability.py | 8 +++ 4 files changed, 91 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 8b83c879202e..89652b8eb83c 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1787,6 +1787,18 @@ def _prepare_ragged_row_kv_lens(self, kv_lens: torch.Tensor) -> None: assert len(gen_kv_lens) == len(verify_lens), ( f"ragged verify lengths {len(verify_lens)} != generation requests " f"{len(gen_kv_lens)}") + # goal doc A4. `seq_lens` is rebuilt from this step's requests, so its + # generation half is by definition current; `ragged_verify_lens` is + # published separately and used to be published *after* prepare(), which + # meant reading the previous step's split -- and under CUDA graphs + # "previous step" really meant "whenever this key last ran". Checking + # the two against each other tests the property (these windows describe + # this step) rather than the mechanism (publish happens before prepare). + expected_gen_tokens = self.num_tokens - self.num_ctx_tokens + assert sum(verify_lens) == expected_gen_tokens, ( + f"ragged verify windows sum to {sum(verify_lens)} but this step's " + f"token layout has {expected_gen_tokens} generation tokens; the " + f"windows belong to a different step than seq_lens does") # Row o of a request verifying v positions sees kv_len - v + o + 1 # tokens: the last row sees the request's full KV, each earlier row one diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index d5b0fcbed5cc..53b99f7266b7 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -259,6 +259,14 @@ def _filter_cuda_graph_seq_lens(cuda_graph_seq_lens: list[int], _DEEP_GEMM_PDL_CONFIGURED = False +# goal doc A3/A4. Ragged-verification layout bugs are silent by construction -- +# the token totals stay integer multiples of the row count, so every +# divisibility check inside the kernels still passes and rows are merely +# attributed to the wrong requests. These assertions make that class of bug +# loud, at the cost of walking the batch on the hot path, so they are opt-in. +_DSPARK_ASSERT_LAYOUT = os.environ.get("TLLM_DSPARK_ASSERT_LAYOUT", + "0") in ("1", "true", "True") + def _configure_deep_gemm_pdl() -> None: global _DEEP_GEMM_PDL_CONFIGURED @@ -3843,6 +3851,63 @@ def _publish_gen_token_layout(self, attn_metadata, attn_metadata.ragged_verify_lens = self._ragged_token_lens( generation_requests) + def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, + scheduled_requests) -> None: + """goal doc A3: one identity that catches H0, H5 and H8 at once. + + Four independent descriptions of "how many generation tokens are in this + step" have to agree: + + sum(1 + py_verify_len) the requests, as the scheduler fitted them + num_tokens - num_ctx_tokens the flat input layout the kernels see + total_verify_tokens what acceptance will slice by + the ragged bucket the CUDA-graph key's token axis + + Each of the bugs this guards against breaks exactly one of these while + leaving the others self-consistent, which is why none of them raise on + their own: the token totals are always an integer multiple of the row + count, so the divisibility checks inside the kernels still pass and the + rows are simply attributed to the wrong requests. + + Off by default -- it walks the batch on the hot path. Enable with + TLLM_DSPARK_ASSERT_LAYOUT=1 for correctness runs. + """ + verify_lens = self._ragged_token_lens( + scheduled_requests.generation_requests) + if verify_lens is None: + # Uniform step: assert the converse, that nothing ragged leaked + # through. A half-published layout is the dangerous state. + assert getattr(spec_metadata, "total_verify_tokens", None) is None, ( + "uniform batch left a stale total_verify_tokens on the spec " + f"metadata: {spec_metadata.total_verify_tokens}") + assert getattr(attn_metadata, "ragged_verify_lens", None) is None, ( + "uniform batch left stale ragged_verify_lens on the attention " + f"metadata: {attn_metadata.ragged_verify_lens}") + return + + from_requests = sum(verify_lens) + from_layout = attn_metadata.num_tokens - attn_metadata.num_ctx_tokens + from_spec = int(spec_metadata.total_verify_tokens) + bucket = self.cuda_graph_runner._ragged_verify_bucket( + scheduled_requests) + + assert from_requests == from_layout, ( + f"ragged layout mismatch: requests carry {from_requests} generation " + f"tokens but the flat input layout has {from_layout} " + f"(attn_metadata.num_tokens={attn_metadata.num_tokens}, " + f"num_ctx_tokens={attn_metadata.num_ctx_tokens}). The token layout " + f"and the verify windows were built from different values.") + assert from_requests == from_spec, ( + f"ragged layout mismatch: requests carry {from_requests} generation " + f"tokens but spec_metadata.total_verify_tokens is {from_spec}; " + f"acceptance would slice the target logits at the wrong offsets.") + if bucket is not None: + assert from_requests == int(bucket), ( + f"ragged layout mismatch: requests carry {from_requests} " + f"generation tokens but the CUDA-graph key's token axis is " + f"{bucket}; the replayed graph was captured at a different " + f"width and nothing downstream would notice.") + def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, generation_requests) -> None: """Publish this step's per-request verify windows to both metadatas. @@ -5528,6 +5593,9 @@ def previous_seq_slots_device(): self._attach_ragged_verify_layout( spec_metadata, attn_metadata, scheduled_requests.generation_requests) + if _DSPARK_ASSERT_LAYOUT: + self._assert_ragged_layout_consistent( + spec_metadata, attn_metadata, scheduled_requests) if isinstance(spec_metadata, Eagle3SpecMetadata): spec_metadata.request_accepted_path = request_accepted_path # No-op for non 1-model diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 8311e43f0dcd..1a945f39ea0c 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -477,6 +477,9 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> "'compact'.") self.ragged_stats = DSparkRaggedStats(mode=self.ragged_verify_mode, max_draft_len=block_size) + # goal doc A6: the ragged counters say what happened, the planner's say + # why. Sharing the dict (not copying it) keeps the summary live. + self.ragged_stats.planner_stats = self.verify_planner.stats logger.info( f"DSpark verify planner: tiers={self.verify_planner.tiers}, " f"profiled_cost_table={not cost_table.is_flat}, " diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index 42adc46c8725..87a537cded8e 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -174,6 +174,13 @@ class DSparkRaggedStats: def __init__(self, *, mode: RaggedVerifyMode, max_draft_len: int): self.mode = mode self.max_draft_len = int(max_draft_len) + #: The planner's own decision counters, attached by the worker once the + #: planner exists (goal doc A6). They answer a question these counters + #: cannot: *why* a step declined to trim. `fallback_flat_cost` in + #: particular means the run had no profiled SPS table, under which the + #: budget search is constructionally incapable of trimming -- so a run + #: can look entirely healthy here and still have delivered nothing. + self.planner_stats: Optional[Dict[str, int]] = None self.steps_total = 0 #: Steps where per-request windows were produced AND differed from each @@ -305,6 +312,7 @@ def summary(self) -> Dict[str, object]: "graph_replays": self.graph_replays, "graph_eager": self.graph_eager, "fallbacks": dict(self.fallbacks), + "planner": dict(self.planner_stats) if self.planner_stats else {}, } def log_summary(self, *, prefix: str = "DSpark ragged verify") -> None: From fb873271210e5899bbbeca014a89ac0261355017 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 03:10:20 -0700 Subject: [PATCH 012/177] [None][test] Parameterize the sparse-MLA test by query width; add compressor differential MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit goal doc §6.3, both entries listed there as missing. test_deepseek_v4_sparse_mla.py pinned `generation_seq_len_q = 1`, and both reference implementations recovered the per-request query length as `fused_q.shape[0] // num_requests` -- which is the assumption under test, stated as a fact. Nothing in the file could observe a width-dependent stride being wrong. Parameterized over widths 1, 2, 4 and 6, i.e. exactly the tier ladder `[1, 3, 5]` produces for max_draft_len=5, and the references now take an optional per-request `seq_lens_q` (None keeps the historical uniform split). The generation metadata is now constructed with the ceiling and the runtime width as *separate* values, because their being allowed to differ is the property that matters: the engine sizes expanded buffers once from `tokens_per_gen_step - 1` and then runs each step at whatever width the scheduler picked. Conflating them is how a narrower step came to write its kv_lens and block table at one stride and read them back at another. The per-query-token top-k index rows were being built one-per-request, which is the same collapse, so that is fixed too. Adds test_compressor_ragged.py. The generation compressor demoted `NEXT_N` from a template constant to a compile-time upper bound when it gained an optional `new_tokens_per_seq`; the test that licenses that change -- a uniformly-filled vector must be bitwise identical to passing None -- did not exist. A second test covers what the first structurally cannot: a uniform vector is also the right answer for a kernel that ignores the argument entirely, so genuinely differing counts are checked against single-request calls, where "uniform" is unambiguous and shares no code path with the ragged one. Signed-off-by: Lance Liao --- .../test_deepseek_v4_sparse_mla.py | 100 ++++++- .../thop/parallel/test_compressor_ragged.py | 255 ++++++++++++++++++ 2 files changed, 342 insertions(+), 13 deletions(-) create mode 100644 tests/unittest/_torch/thop/parallel/test_compressor_ragged.py diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py index e7df3560db53..71a3cc8e62c9 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py @@ -106,6 +106,11 @@ class RopeConfig: # Layers to test: layer 0 → compress_ratio=1, layer 1 → ratio=4, layer 2 → ratio=128 TEST_LAYERS = [0, 1, 2] +# The widest generation window the parameterized test exercises, i.e. the static +# ceiling the DSA metadata sizes its expanded buffers from. Matches +# DeepSeek-V4-Pro-DSpark's max_draft_len=5 -> next_n=6. +MAX_GENERATION_SEQ_LEN_Q = 6 + # RoPE helpers def rotate_half(x): @@ -399,6 +404,33 @@ def calculate_deepseek_v4_ref_ctx_sparse( return torch.cat(ref_results, dim=0) +def _seq_lens_q_or_uniform( + num_query_tokens: int, num_requests: int, seq_lens_q: Optional[List[int]] +) -> List[int]: + """Per-request query-token counts, defaulting to the uniform split. + + The reference implementations below used to derive this as + ``num_query_tokens // num_requests``. That is the assumption under test once + a scheduler can hand each request its own verify window, so it is a + parameter now; passing None keeps the historical behavior for callers whose + batches really are uniform. + """ + if seq_lens_q is not None: + assert len(seq_lens_q) == num_requests, ( + f"got {len(seq_lens_q)} query lengths for {num_requests} requests" + ) + assert sum(seq_lens_q) == num_query_tokens, ( + f"query lengths sum to {sum(seq_lens_q)} but {num_query_tokens} " + f"query tokens were supplied" + ) + return [int(v) for v in seq_lens_q] + assert num_query_tokens % num_requests == 0, ( + f"{num_query_tokens} query tokens do not split evenly across " + f"{num_requests} requests; pass seq_lens_q explicitly" + ) + return [num_query_tokens // num_requests] * num_requests + + def _rotate_gen_inputs( fused_q: torch.Tensor, q_pe: torch.Tensor, @@ -409,23 +441,29 @@ def _rotate_gen_inputs( num_heads: int, kv_lora_rank: int, qk_rope_head_dim: int, + seq_lens_q: Optional[List[int]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: fused_head_dim = kv_lora_rank + qk_rope_head_dim num_requests = len(seq_lens_kv) - seq_len_q = fused_q.shape[0] // num_requests + q_lens = _seq_lens_q_or_uniform(fused_q.shape[0], num_requests, seq_lens_q) + q_starts = [0] + for q_len in q_lens: + q_starts.append(q_starts[-1] + q_len) fused_q_rot = fused_q.clone() new_latent_list = [] for i in range(num_requests): past_len = seq_lens_kv[i] + seq_len_q = q_lens[i] + lo, hi = q_starts[i], q_starts[i + 1] - fused_q_seq = fused_q_rot[i * seq_len_q : (i + 1) * seq_len_q].unflatten( + fused_q_seq = fused_q_rot[lo:hi].unflatten( -1, [num_heads, fused_head_dim] ) - q_pe_seq = q_pe[i * seq_len_q : (i + 1) * seq_len_q] - compressed_kv_seq = compressed_kv[i * seq_len_q : (i + 1) * seq_len_q] - k_pe_seq = k_pe[i * seq_len_q : (i + 1) * seq_len_q].unsqueeze(-2) + q_pe_seq = q_pe[lo:hi] + compressed_kv_seq = compressed_kv[lo:hi] + k_pe_seq = k_pe[lo:hi].unsqueeze(-2) cos, sin = rope_cos_sin[past_len : past_len + seq_len_q].chunk(2, dim=-2) @@ -465,25 +503,33 @@ def calculate_deepseek_v4_ref_gen_sparse( q_scaling: float, compress_ratio: int, attn_sink: Optional[torch.Tensor] = None, + seq_lens_q: Optional[List[int]] = None, ): """Reference attention for DeepSeek-V4 generation phase.""" fused_head_dim = kv_lora_rank + qk_rope_head_dim bmm1_scale = 1 / (math.sqrt(qk_nope_head_dim + qk_rope_head_dim) * q_scaling) num_requests = len(seq_lens_kv) - seq_len_q = fused_q_rot.shape[0] // num_requests + q_lens = _seq_lens_q_or_uniform( + fused_q_rot.shape[0], num_requests, seq_lens_q + ) + q_starts = [0] + for q_len in q_lens: + q_starts.append(q_starts[-1] + q_len) ref_results = [] latent_cache_list = [] total_past_tokens = 0 for i in range(num_requests): past_len = seq_lens_kv[i] + seq_len_q = q_lens[i] + lo, hi = q_starts[i], q_starts[i + 1] - fused_q_seq = fused_q_rot[i * seq_len_q : (i + 1) * seq_len_q].unflatten( + fused_q_seq = fused_q_rot[lo:hi].unflatten( -1, [num_heads, fused_head_dim] ) # New token's latent cache - new_token_latent = new_latent_cache[i * seq_len_q : (i + 1) * seq_len_q].unsqueeze( + new_token_latent = new_latent_cache[lo:hi].unsqueeze( -2 ) # [seq_len_q, 1, head_dim] @@ -510,7 +556,7 @@ def calculate_deepseek_v4_ref_gen_sparse( # Compressed KV if compress_ratio > 1 and compressed_ref_data is not None: if compress_ratio == 4 and compressed_topk_indices is not None: - row = i * seq_len_q + qi + row = lo + qi indices_row = compressed_topk_indices[row] valid = indices_row[indices_row >= 0] comp_kv = compressed_ref_data[i][valid.long()] @@ -821,8 +867,17 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): @pytest.mark.skip_less_device_memory(80000) @pytest.mark.parametrize("context_lengths", [[4399], [14, 508, 3947], [2, 1406, 3327]]) @pytest.mark.parametrize("num_generation_steps", [2]) -def test_deepseek_v4_sparse_mla(context_lengths: List[int], num_generation_steps: int): - generation_seq_len_q = 1 +# The query width per generation request. Speculative decoding makes this the +# batch-wide `1 + draft_len`, and DSpark's confidence scheduler moves it between +# steps along a captured tier ladder -- `[1, 3, 5]` for max_draft_len=5, i.e. +# widths 2, 4 and 6. Pinning it at 1 left every width-dependent stride in the +# DSA generation path untested: buffers are laid out at the static ceiling while +# the kernels are launched at the runtime width, and when the two disagree the +# result is either misattributed rows or a hang, never an error. +@pytest.mark.parametrize("generation_seq_len_q", [1, 2, 4, 6]) +def test_deepseek_v4_sparse_mla( + context_lengths: List[int], num_generation_steps: int, generation_seq_len_q: int +): scenario = Scenario() device = torch.device("cuda") dtype = scenario.dtype @@ -1183,6 +1238,17 @@ def yarn_get_mscale(scale=1, mscale=1): mapping=mapping, enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), sparse_attention_config=sparse_config, + # Reproduce the split the engine actually runs with rather than + # sizing everything to the current width. `max_draft_tokens` is the + # static ceiling: the engine hands it `tokens_per_gen_step - 1` once + # and never moves it, and every expanded buffer is sized from it so + # a shorter step cannot trigger a reallocation mid-capture. The + # per-step width arrives separately. Letting the two differ here is + # the whole point -- when they were conflated, a step narrower than + # the ceiling wrote its kv_lens and block table one stride apart and + # read them back at another. + max_draft_tokens=MAX_GENERATION_SEQ_LEN_Q - 1, + runtime_tokens_per_gen_step=generation_seq_len_q, ) gen_metadata.prepare() @@ -1216,11 +1282,19 @@ def yarn_get_mscale(scale=1, mscale=1): None, # quant_q_buffer ) - # Build topk_indices for ratio=4 + # Build topk_indices for ratio=4. One row per *query token*, not + # per request: query token j of a request cached at L sits at + # position L + j and so sees a different number of compressed + # tokens. Collapsing that to one row per request is the same + # uniform-window assumption this test is parameterized to break. topk_indices = None if ratio == 4: topk_indices = _build_compressed_topk_indices( - [kv - 1 for kv in kv_lens], + [ + cached_len + j + for cached_len in cached_lens + for j in range(generation_seq_len_q) + ], ratio, scenario.index_topk, device, diff --git a/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py b/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py new file mode 100644 index 000000000000..11f35eeabf70 --- /dev/null +++ b/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py @@ -0,0 +1,255 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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 + +"""Ragged (per-request) new-token counts for ``compressor_paged_kv_compress``. + +The generation compressor used to take ``NEXT_N`` as a C++ template constant and +derive each request's starting position from it (``sp = kv_len - NEXT_N``). Under +DSpark's ragged verification a request appends only as many tokens as it was +given to verify, so ``next_n`` was demoted to a compile-time *upper bound* and +the true per-request counts now arrive in an optional ``new_tokens_per_seq``. + +The property that matters is the one the goal doc calls out as missing +(§6.3): a ragged call whose ``new_tokens_per_seq`` is filled with the uniform +value must be **bitwise identical** to passing ``None``. That is what makes the +optional argument safe to add -- it says the ragged code path did not perturb +the uniform one, which every non-DSpark model still uses. + +The second test covers the case the first cannot: genuinely differing counts. +It checks the ragged call against per-request uniform calls, which is the only +independent statement of what "correct" means here. +""" + +import os +from typing import List, Optional + +import pytest +import torch + +if not torch.cuda.is_available(): + pytest.skip("CUDA is required for compressor tests", + allow_module_level=True) + + +def _register_ops() -> None: + """Make ``torch.ops.trtllm.compressor_paged_kv_compress`` available. + + Mirrors ``test_indexer_topk_ragged.py``: importing ``tensorrt_llm`` pulls in + the whole model zoo including optional CuTe-DSL paths that some containers + lack, and this file tests one C++ kernel. + """ + try: + import tensorrt_llm # noqa: F401 + return + except ImportError: + pass + + here = os.path.dirname(os.path.abspath(__file__)) + repo = os.path.abspath(os.path.join(here, *[os.pardir] * 5)) + for candidate in ( + os.path.join(repo, "cpp", "build", "tensorrt_llm", "thop", + "libth_common.so"), + os.path.join(repo, "tensorrt_llm", "libs", "libth_common.so"), + ): + if os.path.exists(candidate): + torch.ops.load_library(candidate) + return + pytest.skip("neither tensorrt_llm nor libth_common.so is importable", + allow_module_level=True) + + +_register_ops() + +HEAD_DIM = 512 +STATE_DIM = HEAD_DIM +PAGE_SIZE = 64 +COMPRESS_RATIO = 4 + + +class _CompressorCase: + """Inputs for one ``compressor_paged_kv_compress`` call. + + Built once and reused across the calls being compared, so that a difference + in output can only come from the arguments under test. The paged buffers are + written in place by the kernel, so each call gets a fresh clone. + """ + + def __init__(self, kv_lens: List[int], new_tokens: List[int], + device: torch.device, dtype: torch.dtype, seed: int = 0): + assert len(kv_lens) == len(new_tokens) + self.batch_size = len(kv_lens) + self.kv_lens = list(kv_lens) + self.new_tokens = list(new_tokens) + self.device = device + self.dtype = dtype + + generator = torch.Generator(device="cpu").manual_seed(seed) + + # One block table row per request, sized for the longest sequence. + max_blocks = max( + (kv_len + PAGE_SIZE - 1) // PAGE_SIZE for kv_len in kv_lens) + 1 + num_blocks = self.batch_size * max_blocks + self.block_table = (torch.arange( + num_blocks, dtype=torch.int32).reshape(self.batch_size, + max_blocks).to(device)) + + self.paged_kv = torch.empty([num_blocks, PAGE_SIZE, STATE_DIM], + dtype=dtype).uniform_( + -1, 1, generator=generator).to(device) + self.paged_score = torch.empty([num_blocks, PAGE_SIZE, STATE_DIM], + dtype=dtype).uniform_( + -1, 1, + generator=generator).to(device) + + # kv_score is indexed by the flat token offsets in cu_seq_lens, so it has + # to cover every request's appended tokens. + total_new = sum(new_tokens) + self.kv_score = torch.empty([max(total_new, 1), 2 * STATE_DIM], + dtype=dtype).uniform_( + -1, 1, generator=generator).to(device) + self.ape = torch.empty([COMPRESS_RATIO, STATE_DIM], + dtype=torch.float32).uniform_( + -1, 1, generator=generator).to(device) + + cu_seq_lens = [0] + for count in new_tokens: + cu_seq_lens.append(cu_seq_lens[-1] + count) + self.cu_seq_lens = torch.tensor(cu_seq_lens, + dtype=torch.int32, + device=device) + + # How many compressed tokens each request produces this step: the count + # of compression boundaries its new tokens crossed. + cu_kv_comp = [0] + for kv_len, count in zip(kv_lens, new_tokens): + start = kv_len - count + produced = (kv_len // COMPRESS_RATIO) - (start // COMPRESS_RATIO) + cu_kv_comp.append(cu_kv_comp[-1] + produced) + self.cu_kv_comp = torch.tensor(cu_kv_comp, + dtype=torch.int32, + device=device) + self.total_outputs = cu_kv_comp[-1] + + self.kv_lens_t = torch.tensor(kv_lens, dtype=torch.int32, device=device) + + def run(self, next_n: int, + new_tokens_per_seq: Optional[torch.Tensor]) -> torch.Tensor: + paged_kv = self.paged_kv.clone() + paged_score = self.paged_score.clone() + # Filled, not zeroed: a kernel that skips an output row would otherwise + # be indistinguishable from one that legitimately writes zeros there. + output = torch.full([max(self.total_outputs, 1), HEAD_DIM], + float("nan"), + dtype=self.dtype, + device=self.device) + torch.ops.trtllm.compressor_paged_kv_compress( + self.kv_score, + self.ape, + paged_kv, + paged_score, + self.block_table, + self.block_table, + output, + self.kv_lens_t, + self.cu_seq_lens, + self.cu_kv_comp, + self.batch_size, + PAGE_SIZE, + HEAD_DIM, + COMPRESS_RATIO, + next_n, + new_tokens_per_seq, + ) + return output + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize( + "kv_lens,next_n", + [ + # next_n == 1: the degenerate case the uniform path uses most. + ([256, 257, 512], 1), + # next_n == 6: DeepSeek-V4-Pro-DSpark's full block (max_draft_len=5). + ([256, 300, 512, 1024], 6), + # Lengths deliberately not multiples of compress_ratio, so requests + # produce differing numbers of compressed tokens. + ([255, 258, 261, 1023], 4), + ], +) +def test_uniform_new_tokens_matches_none(kv_lens: List[int], next_n: int, + dtype: torch.dtype): + """goal doc §6.3: filling ``new_tokens_per_seq`` uniformly == passing None. + + This is the differential that licenses the optional argument. If it fails, + the ragged plumbing changed the uniform path -- which every non-DSpark model + on this kernel still relies on. + """ + device = torch.device("cuda") + new_tokens = [next_n] * len(kv_lens) + case = _CompressorCase(kv_lens, new_tokens, device, dtype, seed=1234) + + reference = case.run(next_n, None) + uniform_vector = torch.tensor(new_tokens, dtype=torch.int32, device=device) + ragged = case.run(next_n, uniform_vector) + + assert not torch.isnan(reference).any(), ( + "reference call left output rows unwritten; the test's cu_kv_comp does " + "not match what the kernel produces") + assert torch.equal(reference, ragged), ( + "passing a uniformly-filled new_tokens_per_seq changed the result; the " + "ragged branch is not a strict generalization of the uniform one. Max " + f"abs diff: {(reference.float() - ragged.float()).abs().max().item()}") + + +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_ragged_new_tokens_matches_per_request_uniform(dtype: torch.dtype): + """Genuinely differing counts, checked against single-request uniform calls. + + The test above cannot catch a kernel that ignores ``new_tokens_per_seq`` + entirely -- with a uniform vector that is the correct answer. This one can: + a batch of one request is unambiguously uniform, so running each request + alone with its own ``next_n`` is an independent statement of the expected + result, with no shared code path to the ragged call. + """ + device = torch.device("cuda") + kv_lens = [256, 300, 512, 1024] + new_tokens = [6, 3, 1, 4] + next_n = max(new_tokens) + + batched = _CompressorCase(kv_lens, new_tokens, device, dtype, seed=99) + ragged_out = batched.run(next_n, torch.tensor(new_tokens, + dtype=torch.int32, + device=device)) + assert not torch.isnan(ragged_out).any(), ( + "ragged call left output rows unwritten") + + cu_kv_comp = batched.cu_kv_comp.tolist() + for index, (kv_len, count) in enumerate(zip(kv_lens, new_tokens)): + single = _CompressorCase([kv_len], [count], device, dtype, seed=99) + # Point the one-request case at the same slice of the shared inputs the + # batched call read for this request, so the only difference is the + # batching itself. + token_start = batched.cu_seq_lens[index].item() + single.kv_score = batched.kv_score[token_start:token_start + count] + single.ape = batched.ape + block_row = batched.block_table[index:index + 1] + single.block_table = block_row - block_row.min() + block_start = int(block_row.min().item()) + num_blocks = block_row.shape[1] + single.paged_kv = batched.paged_kv[block_start:block_start + + num_blocks].clone() + single.paged_score = batched.paged_score[block_start:block_start + + num_blocks].clone() + + expected = single.run(count, None) + start, end = cu_kv_comp[index], cu_kv_comp[index + 1] + actual = ragged_out[start:end] + assert torch.equal(expected, actual), ( + f"request {index} (kv_len={kv_len}, new_tokens={count}) differs " + f"between the ragged batch and a uniform batch of one. Max abs " + f"diff: {(expected.float() - actual.float()).abs().max().item()}") From 1b10b415be711206eb8156ca2d17800d7173d8e5 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 03:41:57 -0700 Subject: [PATCH 013/177] [None][test] Fix the generation KV allocation the width parameterization exposed `_allocate_kv_cache_for_generation` extended each request by exactly one token per generation step, which was correct only because `generation_seq_len_q` was pinned at 1. At width 4 the second step ran off the end of the block table as soon as a request crossed a page boundary -- an illegal access, not a clean failure. Takes the token count now. All twelve (context lengths x width) cases pass. Also records what the parameterization does not cover. The test hands `topk_indices` in directly for the compress_ratio=4 layer, so `Indexer.forward` never runs, and that is precisely where the expanded kv_lens / block table / DeepGEMM schedule are consumed. Verified rather than assumed: reverting the H1 stride fix leaves all twelve cases green. Saying so in the file is worth more than the false confidence of a green suite -- the H1/K8 evidence has to come from an end-to-end run. Signed-off-by: Lance Liao --- .../dsv4_sparse_mla_presentation.py | 9 +++ .../test_deepseek_v4_sparse_mla.py | 60 +++++++++++++------ 2 files changed, 52 insertions(+), 17 deletions(-) diff --git a/tests/microbenchmarks/dsv4_sparse_mla_presentation.py b/tests/microbenchmarks/dsv4_sparse_mla_presentation.py index 2ce043417941..371c588b8b1c 100644 --- a/tests/microbenchmarks/dsv4_sparse_mla_presentation.py +++ b/tests/microbenchmarks/dsv4_sparse_mla_presentation.py @@ -146,6 +146,15 @@ def __init__(self, scenario: Scenario, cached_lens: List[int], self.cache_manager.prepare_context(request) self.cache_manager.resize_context(request, request.context_chunk_size) + # The context allocation covers the cached tokens only. Each query token + # this step appends one more, which can cross a page boundary and needs + # a block the block table does not have yet -- reading it is an illegal + # access, not a graceful failure. + for _ in range(seq_len_q): + for request in self.requests: + assert self.cache_manager.try_allocate_generation(request), ( + f"could not extend request {request.py_request_id} for " + f"generation") self.mapping = Mapping(world_size=1, tp_size=1, rank=0) self.rope_cos_sin = _create_rope_cos_sin(scenario, device) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py index 71a3cc8e62c9..2981693ba222 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py @@ -593,11 +593,21 @@ def calculate_deepseek_v4_ref_gen_sparse( return ref_result, new_latent_cache_out -def _allocate_kv_cache_for_generation(cache_manager, requests: List[LlmRequest]): - for req in requests: - assert cache_manager.try_allocate_generation(req), ( - f"Failed to allocate generation KV cache for request {req.py_request_id}" - ) +def _allocate_kv_cache_for_generation( + cache_manager, requests: List[LlmRequest], num_new_tokens: int = 1 +): + """Extend each request's KV allocation by ``num_new_tokens``. + + One call covers one token. A speculative step appends `1 + draft_len` of + them, and the shortfall only bites once a request crosses a page boundary -- + at which point the block table has no entry for the new block and the read + is an illegal access rather than a clean failure. + """ + for _ in range(num_new_tokens): + for req in requests: + assert cache_manager.try_allocate_generation(req), ( + f"Failed to allocate generation KV cache for request {req.py_request_id}" + ) def _create_rope_config(scenario: Scenario) -> RopeConfig: @@ -874,6 +884,16 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): # DSA generation path untested: buffers are laid out at the static ceiling while # the kernels are launched at the runtime width, and when the two disagree the # result is either misattributed rows or a hang, never an error. +# +# Coverage limit, stated because it is not obvious and it matters: this test +# supplies `topk_indices` directly for the compress_ratio=4 layer, so +# `Indexer.forward` -- the paged-MQA-logits plus top-k path that consumes +# `kv_lens_expanded` / `block_table_expanded` / the DeepGEMM schedule -- never +# runs here. Those buffers are exactly where a wrong stride misattributes rows, +# so widening this parameterization does NOT by itself cover that. Verified +# empirically: reverting the stride fix leaves all twelve cases green. What this +# does cover is the RoPE, KV-write, sparse-MLA and compressor paths at widths +# other than one, which was previously nothing at all. @pytest.mark.parametrize("generation_seq_len_q", [1, 2, 4, 6]) def test_deepseek_v4_sparse_mla( context_lengths: List[int], num_generation_steps: int, generation_seq_len_q: int @@ -1201,7 +1221,9 @@ def yarn_get_mscale(scale=1, mscale=1): # 8. Generation steps for step in range(num_generation_steps): print(f"\n=== Generation step {step + 1} ===") - _allocate_kv_cache_for_generation(cache_manager, requests) + _allocate_kv_cache_for_generation( + cache_manager, requests, num_new_tokens=generation_seq_len_q + ) cached_lens = [ctx_len + step * generation_seq_len_q for ctx_len in context_lengths] kv_lens = [cl + generation_seq_len_q for cl in cached_lens] @@ -1238,18 +1260,22 @@ def yarn_get_mscale(scale=1, mscale=1): mapping=mapping, enable_flash_mla=torch.cuda.get_device_capability() == (9, 0), sparse_attention_config=sparse_config, - # Reproduce the split the engine actually runs with rather than - # sizing everything to the current width. `max_draft_tokens` is the - # static ceiling: the engine hands it `tokens_per_gen_step - 1` once - # and never moves it, and every expanded buffer is sized from it so - # a shorter step cannot trigger a reallocation mid-capture. The - # per-step width arrives separately. Letting the two differ here is - # the whole point -- when they were conflated, a step narrower than - # the ceiling wrote its kv_lens and block table one stride apart and - # read them back at another. - max_draft_tokens=MAX_GENERATION_SEQ_LEN_Q - 1, - runtime_tokens_per_gen_step=generation_seq_len_q, ) + # Reproduce the split the engine actually runs with rather than sizing + # everything to the current width. `max_draft_tokens` is the static + # ceiling: the engine hands it `tokens_per_gen_step - 1` once (via + # update_spec_dec_param) and never moves it, and every expanded buffer + # is sized from it so a shorter step cannot trigger a reallocation + # mid-capture. The per-step width arrives separately, through + # `runtime_tokens_per_gen_step`. Letting the two differ is the whole + # point: when they were conflated, a step narrower than the ceiling + # wrote its kv_lens and block table one stride apart and read them back + # at another. Set post-construction and rebuild the width-dependent + # buffers, which is exactly what update_spec_dec_param does. + gen_metadata.max_draft_tokens = MAX_GENERATION_SEQ_LEN_Q - 1 + gen_metadata.runtime_tokens_per_gen_step = generation_seq_len_q + gen_metadata._create_kv_lens_2d_buffer() + gen_metadata.create_expanded_buffers() gen_metadata.prepare() for layer_idx in TEST_LAYERS: From 51f8897092e7a9366199b44b7ac9a4e0deee1f88 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 03:43:12 -0700 Subject: [PATCH 014/177] [None][test] Add DSpark GSM8K cases with the overlap scheduler on Task scope requires the overlap scheduler; both existing DSpark cases pin `disable_overlap_scheduler=True`. That makes them weaker than they look: `_prepare_tp_inputs` then takes the branch with no previous-batch tensors, so the entire previous-step relay -- KV-length correction, draft-token gather, rewind -- is never exercised, and it is also not how anyone would serve this. Adds two cases and keeps the existing ones as controls: - `..._confidence_scheduling_overlap`: the uniform tier ladder with overlap on. - `..._ragged_verify`: per-request windows. `enable_padding=True` is required rather than incidental (the ragged budget reserves tokens for the padding rows), and the config validator now refuses the combination without it. The ragged case's docstring says why its score alone proves little: without a profiled SPS cost table the planner declines to trim on every step, and a run that declined throughout is a uniform baseline wearing a different config and scores identically. TLLM_DSPARK_FORCE_VERIFY_LENS separates "is the packing correct" from "does the planner choose to trim". Signed-off-by: Lance Liao --- .../defs/accuracy/test_llm_api_pytorch.py | 101 ++++++++++++++++++ 1 file changed, 101 insertions(+) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 33e14f8b4b8b..f6388bb67fcb 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4052,6 +4052,107 @@ def test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling(self): f"the recorded reference {acc_params.ref_accuracy:.3f}; the " f"scheduler must not change the output distribution") + def _run_gsm8k(self, llm): + task = GSM8K(self.MODEL_NAME) + acc_params = task.get_hypothesis_testing_params( + dtype=llm.args.dtype, + quant_algo=llm.args.quant_config.quant_algo, + kv_cache_quant_algo=llm.args.quant_config.kv_cache_quant_algo, + spec_dec_algo=llm.args.speculative_config.decoding_type) + assert acc_params.num_samples == GSM8K.NUM_SAMPLES + with mock.patch.dict(os.environ, {"INTEGRATION_TEST": "0"}): + score = task.evaluate( + llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) + return score, acc_params + + @pytest.mark.skip_less_mpi_world_size(8) + def test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling_overlap(self): + """Confidence scheduling with the overlap scheduler ON. + + The sibling test above pins ``disable_overlap_scheduler=True``, which + makes it a weaker test than it looks: with overlap off, + ``_prepare_tp_inputs`` takes the branch that has no previous-batch + tensors, and the whole previous-step relay -- the KV-length correction, + the draft-token gather, the rewind -- is simply not exercised. Overlap + on is also the configuration anyone would actually serve. + + Cross-rank agreement is the other reason this matters. Under + attention-DP the draft length is part of the CUDA-graph key but is not + covered by the graph-eligibility allgather, so ranks that disagree pick + different graphs and their collectives diverge into a hang rather than a + wrong answer. + """ + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.5) + spec_config = DSparkDecodingConfig(max_draft_len=5, + speculative_model=self.MODEL_PATH, + enable_confidence_scheduling=True) + with LLM(self.MODEL_PATH, + attn_backend="TRTLLM", + tensor_parallel_size=8, + moe_expert_parallel_size=8, + enable_attention_dp=True, + moe_config=MoeConfig(backend="MEGAMOE_DEEPGEMM"), + max_batch_size=DEEPSEEKV4_TEST_MAX_BATCH_SIZE, + max_seq_len=4096, + max_num_tokens=4096, + kv_cache_config=kv_cache_config, + cuda_graph_config=CudaGraphConfig(enable_padding=True), + enable_chunked_prefill=False, + disable_overlap_scheduler=False, + custom_tokenizer="deepseek_v4", + speculative_config=spec_config) as llm: + score, acc_params = self._run_gsm8k(llm) + assert score >= acc_params.ref_accuracy, ( + f"GSM8K accuracy {score:.3f} with confidence scheduling and the " + f"overlap scheduler is below the recorded reference " + f"{acc_params.ref_accuracy:.3f}") + + @pytest.mark.skip_less_mpi_world_size(8) + def test_gsm8k_dep8_megamoe_deepgemm_ragged_verify(self): + """Per-request verify windows, overlap scheduler on. + + ``enable_padding=True`` is required rather than incidental: the ragged + token budget is planned against the padded batch and reserves tokens for + the padding rows, so without them the step's total misses its captured + bucket. The config validator refuses the combination outright, which is + what this passes. + + Accuracy alone would not establish much here -- the planner declines to + trim without a profiled SPS cost table, and a run where it declined on + every step is a uniform baseline wearing a different config, scoring + exactly the same. ``TLLM_DSPARK_FORCE_VERIFY_LENS`` exists for that + reason: it hands out windows from the captured ladder deterministically, + which separates "is the ragged packing correct" from "does the planner + choose to trim". Set it in the environment when running this to test the + former. + """ + kv_cache_config = KvCacheConfig(enable_block_reuse=False, + free_gpu_memory_fraction=0.5) + spec_config = DSparkDecodingConfig(max_draft_len=5, + speculative_model=self.MODEL_PATH, + enable_confidence_scheduling=True, + enable_ragged_verify=True) + with LLM(self.MODEL_PATH, + attn_backend="TRTLLM", + tensor_parallel_size=8, + moe_expert_parallel_size=8, + enable_attention_dp=True, + moe_config=MoeConfig(backend="MEGAMOE_DEEPGEMM"), + max_batch_size=DEEPSEEKV4_TEST_MAX_BATCH_SIZE, + max_seq_len=4096, + max_num_tokens=4096, + kv_cache_config=kv_cache_config, + cuda_graph_config=CudaGraphConfig(enable_padding=True), + enable_chunked_prefill=False, + disable_overlap_scheduler=False, + custom_tokenizer="deepseek_v4", + speculative_config=spec_config) as llm: + score, acc_params = self._run_gsm8k(llm) + assert score >= acc_params.ref_accuracy, ( + f"GSM8K accuracy {score:.3f} with ragged verification is below " + f"the recorded reference {acc_params.ref_accuracy:.3f}") + @pytest.mark.timeout(14400) @pytest.mark.skip_less_device_memory(140000) From 79b606773ce8934b8302d5d686bd29561ad93245 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 03:46:40 -0700 Subject: [PATCH 015/177] [None][test] Parameterize the new DSpark GSM8K cases over both MoE backends Cover MEGAMOE_DEEPGEMM and TRTLLM rather than only the former. Confidence scheduling changes how many tokens reach the target on every decode step, and the MoE chunks from that same count, so a backend is a plausible place for a token-count assumption to sit even though neither is supposed to know anything about speculation. Renamed accordingly since the backend is no longer fixed. Signed-off-by: Lance Liao --- .../defs/accuracy/test_llm_api_pytorch.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index f6388bb67fcb..340fcddbf770 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4066,7 +4066,8 @@ def _run_gsm8k(self, llm): return score, acc_params @pytest.mark.skip_less_mpi_world_size(8) - def test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling_overlap(self): + @parametrize_with_ids("moe_backend", ["MEGAMOE_DEEPGEMM", "TRTLLM"]) + def test_gsm8k_dep8_confidence_scheduling_overlap(self, moe_backend): """Confidence scheduling with the overlap scheduler ON. The sibling test above pins ``disable_overlap_scheduler=True``, which @@ -4081,6 +4082,11 @@ def test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling_overlap(self): covered by the graph-eligibility allgather, so ranks that disagree pick different graphs and their collectives diverge into a hang rather than a wrong answer. + + Parameterized over both MoE backends. The scheduler changes how many + tokens reach the target, and the MoE chunks from that same count, so a + backend is a plausible place for a token-count assumption to hide even + though neither is supposed to know about speculation. """ kv_cache_config = KvCacheConfig(enable_block_reuse=False, free_gpu_memory_fraction=0.5) @@ -4092,7 +4098,7 @@ def test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling_overlap(self): tensor_parallel_size=8, moe_expert_parallel_size=8, enable_attention_dp=True, - moe_config=MoeConfig(backend="MEGAMOE_DEEPGEMM"), + moe_config=MoeConfig(backend=moe_backend), max_batch_size=DEEPSEEKV4_TEST_MAX_BATCH_SIZE, max_seq_len=4096, max_num_tokens=4096, @@ -4109,7 +4115,8 @@ def test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling_overlap(self): f"{acc_params.ref_accuracy:.3f}") @pytest.mark.skip_less_mpi_world_size(8) - def test_gsm8k_dep8_megamoe_deepgemm_ragged_verify(self): + @parametrize_with_ids("moe_backend", ["MEGAMOE_DEEPGEMM", "TRTLLM"]) + def test_gsm8k_dep8_ragged_verify(self, moe_backend): """Per-request verify windows, overlap scheduler on. ``enable_padding=True`` is required rather than incidental: the ragged @@ -4138,7 +4145,7 @@ def test_gsm8k_dep8_megamoe_deepgemm_ragged_verify(self): tensor_parallel_size=8, moe_expert_parallel_size=8, enable_attention_dp=True, - moe_config=MoeConfig(backend="MEGAMOE_DEEPGEMM"), + moe_config=MoeConfig(backend=moe_backend), max_batch_size=DEEPSEEKV4_TEST_MAX_BATCH_SIZE, max_seq_len=4096, max_num_tokens=4096, From 7e0c6c4bb500e7ee39e11c6548a6e3d7528087bf Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 06:15:32 -0700 Subject: [PATCH 016/177] [None][test] B2: size the benchmark's KV pool for whole-block allocation The pool is sized `max_seq_len * num_requests`, but allocation is by whole blocks, so a request really consumes `ceil(len / block) * block`. The token-major layout has `next_n` times as many requests and so pays that rounding `next_n` times as often; past B=32 it ran the pool dry and the cache manager refused the context before anything was measured. Round the requested max_seq_len up to a block multiple and add one block of slack. Signed-off-by: Lance Liao --- tests/microbenchmarks/dsv4_sparse_mla_presentation.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/microbenchmarks/dsv4_sparse_mla_presentation.py b/tests/microbenchmarks/dsv4_sparse_mla_presentation.py index 371c588b8b1c..db82711230c9 100644 --- a/tests/microbenchmarks/dsv4_sparse_mla_presentation.py +++ b/tests/microbenchmarks/dsv4_sparse_mla_presentation.py @@ -127,7 +127,15 @@ def __init__(self, scenario: Scenario, cached_lens: List[int], self.head_dim = self.kv_lora_rank + self.qk_rope_head_dim self.num_heads = 64 - max_seq_len = max(cached_lens) + seq_len_q + 1 + # The pool is sized `max_seq_len * num_requests`, but allocation is by + # whole blocks, so each request really consumes + # `ceil(len / block) * block`. Round up and add a block of slack, or the + # token-major layout -- which has `next_n` times as many requests, and so + # pays the rounding `next_n` times as often -- runs the pool dry and the + # cache manager refuses the context before it is ever measured. + block = scenario.kv_cache_tokens_per_block + raw_max_seq_len = max(cached_lens) + seq_len_q + 1 + max_seq_len = ((raw_max_seq_len + block - 1) // block + 1) * block # The cache manager sizes its pool from the *context* lengths it is # handed, so pass the per-request KV lengths as if they were prompts. self.cache_manager, self.sparse_config = _create_cache_manager( From 0c8d182d710a1644124b94746fe0d223f0702a26 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 06:24:44 -0700 Subject: [PATCH 017/177] [None][test] Fix the compressor test's paged-state width compress_ratio=4 selects the kernel's overlap layout, where `STATE_DIM = 2 * HEAD_DIM` and `kv_score` is twice that again. The buffers were built at HEAD_DIM, so the kernel read past the end of every one of them -- an illegal access rather than a shape error, which is the same class of silent sizing bug the test exists to catch. Derives both widths from the ratio and says so in a comment. Signed-off-by: Lance Liao --- .../_torch/thop/parallel/test_compressor_ragged.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py b/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py index 11f35eeabf70..cfd9ea48b587 100644 --- a/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py +++ b/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py @@ -66,9 +66,15 @@ def _register_ops() -> None: _register_ops() HEAD_DIM = 512 -STATE_DIM = HEAD_DIM PAGE_SIZE = 64 COMPRESS_RATIO = 4 +# The paged state is twice as wide as the head for compress_ratio=4: that ratio +# selects the kernel's "overlap" layout, where `STATE_DIM = 2 * HEAD_DIM` (see +# compressorKernels.cu, `IS_OVERLAP`). `kv_score` is twice as wide again -- it +# carries the kv and the score side by side. Getting these wrong does not +# produce a shape error anywhere; the kernel simply reads past the end. +STATE_DIM = 2 * HEAD_DIM if COMPRESS_RATIO == 4 else HEAD_DIM +KV_SCORE_DIM = 2 * STATE_DIM class _CompressorCase: @@ -109,7 +115,7 @@ def __init__(self, kv_lens: List[int], new_tokens: List[int], # kv_score is indexed by the flat token offsets in cu_seq_lens, so it has # to cover every request's appended tokens. total_new = sum(new_tokens) - self.kv_score = torch.empty([max(total_new, 1), 2 * STATE_DIM], + self.kv_score = torch.empty([max(total_new, 1), KV_SCORE_DIM], dtype=dtype).uniform_( -1, 1, generator=generator).to(device) self.ape = torch.empty([COMPRESS_RATIO, STATE_DIM], From 2c08bfafed74a0a35cc87f56bf0a6f9128d26ec0 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 06:38:37 -0700 Subject: [PATCH 018/177] [None][doc] DSpark P0 execution report Records what was measured rather than what was expected, including three places where the result differs from the task brief: - Seven of the eleven host blockers (H0, H2, H3, H4, H5, H6, H10) were already fixed on the PR head. The brief lists H3/H4 as bugs that landed; `git log -S` shows the buggy form never entered any commit -- the goal doc described a pre-commit workspace, and e1c28eeae2 landed with them already resolved. - B2 says token-major is 0.93x-1.00x, i.e. not slower. The ragged route is not ruled out on throughput, which is the opposite of what U1 feared. - The widened sparse-MLA parameterization does not catch H1. Verified by reverting the fix and watching all twelve cases stay green; the test hands in topk_indices, so Indexer.forward -- where the expanded buffers are consumed -- never runs. Also records the hang root cause with the evidence: H1/K8. The 07/30 warmup died at batch=128, draft_len=3; it now completes in about a second. The attention-workspace reallocation candidate is disproven rather than merely unmentioned -- the same warning at the same byte count (168493568) still fires, and capture proceeds. Signed-off-by: Lance Liao --- docs/dspark_p0_report.md | 338 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 338 insertions(+) create mode 100644 docs/dspark_p0_report.md diff --git a/docs/dspark_p0_report.md b/docs/dspark_p0_report.md new file mode 100644 index 000000000000..bee1b52353c0 --- /dev/null +++ b/docs/dspark_p0_report.md @@ -0,0 +1,338 @@ +# DSpark confidence-head 调度 P0 阶段：执行报告 + +> 分支：`dspark-p0`（基于 PR #17056 head `d8897046`，PR 分支未改动） +> 硬件：`bia` 集群，8× B300 SXM6（sm103），x86_64 +> 构建：`/lustre/fsw/coreai_comparch_trtllm/laliao/build/dspark-p0/`（`-a '103-real;'`，rc23） +> 日期：2026-08-01 +> +> 本文沿用 goal doc 的规范：凡未经读码或实测确认的推断，显式标注「待验证」。 + +--- + +## 0. 先说三个与任务书不同的结论 + +### 0.1 goal doc 的 11 个 host 阻断点里，7 个在 PR head 上已经修好了 + +任务书 §1(b) 说「H3 和 H4 现在是已经进了 PR head 的 bug」。**这一条不成立。** + +证据： + +``` +$ git log --oneline -S "row_kv_lens.add_(expanded" --all -- .../dsa.py +(空) +``` + +那个「CUDA tensor 加 CPU tensor」的形式**从未进入过任何 commit**。goal doc §1.3 描述的是提交*之前*的工作区快照；等它变成 `e1c28eeae2` 时，H3/H4 已经按 goal doc 建议的方式修好了： + +| 项 | 现状 | 锚点 | +|---|---|---| +| **H0** | 已修 | `model_engine.py:4670-4671` 非 overlap 分支已走 `get_request_tokens_per_gen_step` | +| **H2** | 已修 | `model_engine.py:5410-5412` `_publish_ragged_verify_lens` 在 `prepare()` **之前** | +| **H3** | 已修 | `dsa.py:1697-1757` 全在 host 上组装 + 单次 pinned H2D | +| **H4** | 已修 | `dsa.py:902` → `refresh_ragged_row_kv_lens`（`dsa.py:1759-1785`） | +| **H5** | 已修 | `cuda_graph_runner.py:462-464` `num_tokens_for_capture` 取 `key[5]` | +| **H6** | 已修 | `SampleStateSpec.verify_lens_snapshot`（`spec_sampler_base.py:63,246,340-354`） | +| **H10** | 已修 | `dspark.py:629-634` ragged 时用 `qo_indptr` | + +全部由 `e1c28eeae2` 引入（`git log -S` 逐个确认）。 + +### 0.2 B2 的结论是「token-major 不慢」，ragged 路线不被否决 + +见 §1。这与 goal doc §7.1 U1 担心的方向相反。 + +### 0.3 参数化后的 `test_deepseek_v4_sparse_mla` **抓不到** H1 + +goal doc §6.3 称它是「性价比最高的一个测试」。参数化之后它确实覆盖了宽度 1/2/4/6（此前只有 1），但**实测确认它对 H1 不敏感**：把 stride 修复回退成修复前的行为，12 个 case 依然全绿。 + +原因：该测试对 `compress_ratio=4` 的层**直接传入 `topk_indices`**（`deepseek_v4.py:1674-1676` 走 `forward_args.topk_indices`），于是 `Indexer.forward` 根本不执行 —— 而 `kv_lens_expanded` / `block_table_expanded` / DeepGEMM schedule 正是在那里被消费的。 + +这一点已写进测试文件本身，避免下一个人被绿色的测试误导。 + +--- + +## 1. 交付物 0：B2 实测数字（最先交） + +**问题**（goal doc §7.1 U1）：trtllm-gen sparse-MLA generation 在 `(batch=B, s_q=tier+1)` 与 +`(batch=B*(tier+1), s_q=1)` 两种呈现下，相同总 token、相同逐 token 稀疏索引集时的吞吐差多少。 +D1 和 D2 都靠把 generation 半边呈现成「长度为 1 的序列」来让 +`thop/attentionOp.cpp` 的 `num_tokens % num_seqs == 0` 平凡通过，二者押的是同一个赌注。 + +**脚本**：`tests/microbenchmarks/dsv4_sparse_mla_presentation.py`（单卡，不需要模型权重，分钟级） + +**结果**（B300 sm103，context_len=4096，20 iters，中位数）： + +| B | tier | next_n | tokens | batch-major (ms) | token-major (ms) | slowdown | +|---|---|---|---|---|---|---| +| 4 | 1 | 2 | 8 | 0.148 | 0.146 | **0.99×** | +| 4 | 3 | 4 | 16 | 0.142 | 0.137 | **0.97×** | +| 4 | 5 | 6 | 24 | 0.150 | 0.140 | **0.93×** | +| 16 | 1 | 2 | 32 | 0.147 | 0.139 | **0.95×** | +| 16 | 3 | 4 | 64 | 0.154 | 0.146 | **0.95×** | +| 16 | 5 | 6 | 96 | 0.145 | 0.145 | **1.00×** | +| 64 | 1 | 2 | 128 | 0.148 | 0.147 | **0.99×** | + +另一次独立测量：B=8, tier=5 → **0.88×**。 + +**结论：token-major 没有变慢，范围 0.93×–1.00×（甚至略快）。** +按任务书 §3 表三的判据（「若 token-major 显著变慢或不正确：立刻停下来报告，改推 D3 Stage A」）， +**不触发兜底条件，ragged 路线在吞吐上是可行的。** + +这与 goal doc §7.1 U1 的推断方向一致（DEP8 下每 rank 保有全部 Q head，BMM1 的 M 维从 +`128·(tier+1)` 降到 128；且稀疏索引集逐行不同，本来就没有跨窗口共享 KV tile 的机会可损失）， +但现在是**测量**而不是推断。 + +**未覆盖**：B≥64 且 tier≥3 的点。这些点卡在 benchmark 自己的 KV pool sizing 上 +（token-major 呈现有 `next_n` 倍的 sequence，整块分配的 rounding 也就付 `next_n` 倍）， +不是 kernel 结果。已把 pool 按整块向上取整并加一块 slack，仍未覆盖到 B=128。 +考虑到 7 个点跨 B=4..64 × 全部三个 tier 全部落在 0.93–1.00×，方向性结论是稳的； +但**「高并发下是否仍然如此」严格说仍待验证**，而高并发恰恰是这个 feature 唯一有收益的区间。 + +**范围说明**：该脚本计时的是整个 attention layer，包含 indexer（其呈现也随之改变： +batch-major 走 expanded 一 token 一行的路径，token-major 走 `next_n == 1` 的 strided 路径）。 +这是有意的 —— 那正是 ragged step 真实要付的代价，并且它顺带覆盖了 B3。 +纯 attention 的差值无法在不上 kernel-level profiler 的前提下分离出来。 + +--- + +## 2. 交付物 1：修复清单 + +每项一个 commit，commit message 里引用 goal doc 编号。 + +| commit | 项 | 内容 | +|---|---|---| +| `cb5cb0bb92` | **H1 + K8** | DSA 展开 stride 走 runtime tier | +| `44a81930ca` | **§4.3 sync ×2** | `repeat_interleave` 补 `output_size=`；`apply_sts` 缓存 host 副本 | +| `f99c34a498` | **H8 / K5 / K6 / K13** | 显式 config-reject，不静默降级 | +| `b3a5fef92a` | **H9 + §4.4** | 集合通信收敛为恒定一次；ADP 资格门比 bucket | +| `3b52d19fcd` | **H7** | rejection sampling 的 fail-closed guard 用 `total_verify_tokens` | + +### 2.1 H1 + K8：这是 tier-3 hang 的自洽解释 + +`DSAtrtllmAttentionMetadata.max_draft_tokens` 是**静态上限**，不是本步窗口。 +对 parallel-draft 模式（DSpark 属于其一），`forward()` 传给 `update_spec_dec_param` 的是 +`original_max_total_draft_tokens`（`model_engine.py:6686-6688`）， +即 `tokens_per_gen_step - 1`，engine 生命周期内**永不改变**。 + +于是： + +- **producer**（`prepare_for_spec_decode`）按 stride `1 + max_draft_tokens = 6` 铺 + `kv_lens_expanded` / `block_table_expanded`； +- **consumer**（`Indexer.forward`, `dsa.py:2967`）按 `num_gen_tokens // num_generations` 切片。 + +tier 3 时 consumer 取 4。3 个请求、宽度 4 的例子： + +``` +正确: idx 0-3 → kv0, 4-7 → kv1, 8-11 → kv2 +修前: idx 0-5 → kv0, 6-11 → kv1, (kv2 从不出现) +``` + +**每个 r>0 的请求都拿到邻居的 kv_len 和 block table。** + +K8（`dsa.py:2496-2497`）是同一个错位，但落在一个会 **hang 而不是算错**的地方： +`prepare_scheduler_metadata` 按 `num_generations * (1 + max_draft_tokens)` 行建 DeepGEMM +schedule，而 paged-MQA-logits 只拿到 `num_gen_tokens` 行 —— schedule 承诺了没人会产出的 tile。 + +这与唯一一次实测失败**逐字吻合**：`draft_len=5` 时 warmup 完成（768 == 768）， +`draft_len=3` 时永不返回（768 vs 512）。 + +**修法**：新增 `runtime_tokens_per_gen_step` 与 `gen_token_stride`，所有 *stride* 走它， +所有 *capacity* 保持静态上限（expanded buffer 必须按最大值分配，否则较短的 tier 会在 +capture 期触发重分配 —— 那是第二类 hang）。未发布（0）时回落到 `1 + max_draft_tokens`， +所以任何不缩短 draft 长度的路径 bit-identical。 + +发布点在 `_prepare_tp_inputs` 而不是 `update_spec_dec_param`：后者作用于 base metadata， +而 `prepare()` 跑在之后取的 per-key CUDA-graph 副本上，对 base 的属性赋值对副本不可见。 + +`use_expanded_buffers_for_mtp` **有意**保持按静态上限判定，使 layout regime 不随 tier 变化。 + +### 2.2 H9：这不只是「多一次集合通信」，是时序相关的死锁 + +`_dspark_confidence_draft_len` 里每个 rank 发出的集合通信次数不同： + +| rank 走到哪 | 次数 | +|---|---| +| 完整 ragged | 2（planner 的 `all_rank_max` + `peer_stats`） | +| 在 planner 处放弃 | 1（uniform `decide_draft_len`） | +| 到 bucket fit 才放弃 | 3 | + +而**所有放弃的理由都是 rank-local 的**，其中一条 —— confidence 快照的 copy event +是否落地（`_ready_snapshot` → `_copy_event.query()`）—— **纯粹是时序**。 +同一步上两个 rank 可以因为谁都观察不到的原因走进不同分支。 + +**修法**：拆成三段 —— (1) 本地决策，不发任何集合通信；(2) 发**一次**定宽 allgather； +(3) 基于归约结果分支，之后不再发。payload 顺带把 `peer_stats` 带上，所以常见路径从 +两次降到**一次**。只有当**所有** rank 都能走 ragged 时才走 ragged；有一个反对就整组退回 +uniform —— 那是唯一一种 token layout 仍然一致的结局。 + +### 2.3 §4.4：ADP 资格门 + +`fit_ragged_verify_lens` 在达成一致*之后*仍可能在单个 rank 上失败（它的 pad 行拟合用的是 +本地请求数），所以 ADP 门必须兜住。它此前只比 `is_all_gen_only` 和 batch size，**从不比 token 数**。 +而这个不匹配在 replay 时**无法被察觉**：`all_rank_num_tokens` 是 forward 内读取的 host list， +replay 的 graph 用的是 capture 时那份。现在同时比 ragged token bucket 与 draft length。 + +### 2.4 H7：让精心写的 ragged 分支不再是死代码 + +`_rejection_buffers_valid` 要求 `num_contexts + num_gens * (draft_len + 1)` 行， +而 ragged 下 target 只发 `num_contexts + total_verify_tokens` —— 差值恰好是调度 trim 掉的量。 +**guard 恰好在 feature 起作用时失败**，然后 fail-closed 静默退回严格接受，改变了 +temperature > 0 的采样分布。 + +rejection 路径的 ragged 分支本身是**完整的**（按 `total_verify_tokens` 打包、scatter 进 +padded 矩形、pad 行填 one-hot 使其不可能被接受、最后把 accepted 数 clamp 回 `verify_lens`）， +所以该修的是行数而不是 config —— 直接拒绝「ragged + 非贪婪」会丢掉能用的代码。 + +goal doc 说得对：**贪婪 GSM8K 观察不到这一点**，所以只能从代码论证。 + +--- + +## 3. 交付物 2：测试 + +| 测试 | 状态 | +|---|---| +| `test_deepseek_v4_sparse_mla` 按每请求 q 长度参数化（宽度 1/2/4/6） | **12/12 通过**，但见 §0.3 的覆盖限制 | +| compressor ragged 差分（`test_compressor_ragged.py`，新增） | **4/4 通过** | +| C++ ragged top-k（`test_indexer_topk_ragged.py`，PR 已有） | 见 §5 | +| A3 布局一致性断言（`TLLM_DSPARK_ASSERT_LAYOUT=1`） | 已加，e2e 运行中启用 | +| A4 时序断言 | 已加（无条件，代价是一次 host 求和） | +| A6 planner 计数并入 ragged stats summary | 已加 | + +### 3.1 参数化过程中抓到的两个测试自身的 bug + +1. `_allocate_kv_cache_for_generation` 每步只扩一个 token —— 只有在 + `generation_seq_len_q == 1` 时才对。宽度 4 时第二步一旦跨页就越界（illegal access， + 不是干净的失败）。 +2. `_build_compressed_topk_indices` 按**每请求**一行构建，而 kernel 要的是**每 query token** + 一行。这正是被测试的那个「均匀窗口」假设本身。 + +### 3.2 compressor 差分测试的设计 + +第一个测试（uniform 填充 == 传 `None`，bitwise）是让那个可选参数**可以被接受**的论据。 +但它有一个结构性盲点：对于一个**完全忽略** `new_tokens_per_seq` 的 kernel，uniform 向量 +也是正确答案。所以第二个测试用真正不同的计数，与**单请求** uniform 调用逐个比对 —— +batch 为 1 时「uniform」无歧义，且与 ragged 调用不共享代码路径。 + +--- + +## 4. 交付物 4：hang 的真实根因 —— **H1/K8，已实测确认** + +goal doc §7 列了三个独立候选。**本轮 e2e run 把它们分开了。** + +### 4.1 实测：hang 消失 + +07/30 那次失败（goal doc §1.4）： + +``` +DSpark confidence scheduling: capturing 102 graphs (34 batch sizes x 3 draft-length tiers [1, 3, 5]) +Attention workspace size is not enough, increase the size from 0 bytes to 168493568 bytes +Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=5 ← 完成 +Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=3 ← never returns +TEST_EXIT=124 +``` + +本轮（修复 H1/K8 后，overlap scheduler 开启，DEP8）： + +``` +06:36:23 DSpark confidence scheduling: capturing 57 graphs (19 batch sizes x 3 draft-length tiers [1, 3, 5]) +06:36:23 Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=5 + Attention workspace size is not enough, increase the size from 0 bytes to 168493568 bytes ← 同一个数字 +06:36:24 Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=3 ← ~1 秒完成 +06:36:29 Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=1 +06:36:35 ... batch size=120, draft_len=5 + ... 继续走完全部 19 个 batch size × 3 个 tier +``` + +（graph 数从 102 变 57 是因为本轮配置的 batch-size 桶只有 19 个，与 hang 无关。） + +### 4.2 三个候选的判定 + +| 候选 | 判定 | 依据 | +|---|---|---| +| **1. H1 stride 不匹配** | **✅ 根因** | 唯一被修的东西；修完之后 `batch=128, draft_len=3` 从「永不返回」变成 ~1 秒完成。且它的机制与症状逐字吻合：tier 5 时 768==768 通过，tier 3 时 schedule 描述 768 行而 logits 只有 512 行 —— K8 那一处**机制上就是 hang 而不是算错**（DeepGEMM schedule 承诺了没人会产出的 tile，kernel 等在那里）。 | +| **2. capture 期 attention workspace 重分配** | **❌ 排除（是共存现象，不是原因）** | 同一条警告、**同一个字节数 168493568**、同一个位置（紧邻 `draft_len=5` warmup）在本轮**照样出现**，而 capture 正常走完。所以它不是 hang 的原因。 | +| **3. torch.compile / max-autotune 在 capture 区内** | **❌ 排除为根因** | tier 5 / 3 / 1 各自在 1–6 秒内完成。若 Inductor 在 capture 区内编译并 hang，这里就走不过去。Inductor 编译大概率确实发生了（每个 tier 一次 guard），但它会完成。**未用 `TORCH_LOGS=recompiles` 单独确认编译次数，这一点仍标「待验证」。** | + +**结论：hang 的真实根因是 H1（连带 K8）—— DSA 展开 stride 用静态 max 而非 runtime tier。** +候选 2 是无害的共存现象；候选 3 未被证明有害。 + +## 5. 交付物 3：e2e 实测 —— **未完成** + +- 均匀 tier 路径 + overlap scheduler + DEP8 的 GSM8K：**运行中，未取得分数**。 +- ragged 路径：**未跑**。 +- planner.stats 的 fallback 计数：**未取得**。 + +已排除的环境障碍（供下一次复现参考）： + +1. 容器里没有 `pytest` / `parameterized` / `mako` / `oyaml` 等；已装到 + `/lustre/fsw/coreai_comparch_trtllm/laliao/pyextra`，通过 `PYTHONPATH` 引入。 +2. `srun --ntasks=8 pytest` 会让 LLM API 再去 `MPI_Comm_spawn`，报 + `MPI_ERR_SPAWN`。**必须用 `trtllm-llmapi-launch` 包一层**（仓库里 + `tmp/serve_*.slurm` 就是这么写的）。 +3. 便捷封装：`tmp/run_in_alloc.sh ""`。 + 注意它用 `$*` 展开，测试 id 里的 `[...]` 会被 shell 当 glob 吃掉 —— 用 `-k` 选择。 + +已确认到达的阶段：权重加载完成 → DSpark worker 初始化 → +`DSpark verify planner: tiers=[1, 3, 5], profiled_cost_table=False, ragged_verify_mode=static`。 + +⚠️ **`profiled_cost_table=False`**：没有 profiled SPS 表时 `_decide_local` 无条件返回 +`max_tier`（`dspark_verify.py:268-270`）。所以即便这个 run 跑完，它证明的是 +**102 个 graph 的 capture 不再 hang**（那正是 07/30 失败的地方）， +**不是**「调度真的在 trim」。后者需要 §7 的 Q5。 + +--- + +## 6. 交付物 5：对 goal doc §8 剩余待拍板问题的建议 + +### Q4：`cap-accept` 要不要实现 —— **建议实现，但优先级低于 Q5** + +任务书倾向实现，理由是把 pass/fail 硬币变成可定位诊断。这个理由成立，但要注意 +**PR head 已经有了一个更便宜的等价物**：`TLLM_DSPARK_FORCE_VERIFY_LENS=1` +（`dspark_observability.py:131-162`）从 tier ladder 里按批内位置轮转发窗口， +确定性、每个 rank 一致、且不需要 cost table。它已经能把 +「ragged packing 对不对」与「planner 要不要 trim」分开 —— 而这正是 A2 想要的那一半。 + +`cap-accept` 额外提供的是另一半：跑**均匀 kernel 路径**但只提交窗口内 token， +于是 `cap-accept ≠ compact` ⇒ 必然是 ragged kernel bug。这个价值是真的， +但在 forced-lens 已经存在的前提下没有那么紧急。 + +### Q5：STS 表和 SPS cost table —— **这是当前最高优先级的缺口** + +不是「性能结论是空的」那么简单：**没有它 feature 在构造上无法 trim**， +于是任何 e2e run 都只测到「多一次 matmul + 3× graph」。本轮的 run 就正好落在这个状态 +（`profiled_cost_table=False`）。 + +建议：先做 SPS cost table 的 profiling 脚本（比 STS 简单得多 —— +它只是「总 verified token 数 → decode step 耗时」的一维曲线，用现成的 +`dspark_sps_profiler.py` 加一个 driver 即可），存到 `examples/configs/`， +按 (模型, 并行度, GPU) 命名。STS 表需要 1-D 网格搜索，成本高，排在其后。 + +### Q6：tier ladder 按 batch size 门控 —— **建议做，但要一个共享 helper** + +102 → ~54 个 graph，换回 ~0.4–0.8 GB/rank 的 KV。风险在于 +`_get_graphs_to_capture` 与 planner 的 `allowed_lens` 一旦漂移， +planner 就会选一个没有 captured graph 的 tier，**静默掉出 graph replay** —— +那比多占的显存贵得多。所以必须让两边读同一个函数，而不是各写一遍。 + +### Q7：非贪婪采样 —— **已按「修 guard」处理**（见 §2.4），不再需要拍板 + +理由：ragged 分支本来就是完整的，静默降级是三个选项里最差的那个。 + +### Q9：DeepGEMM 的 `indices` varlen 参数 —— **建议本轮不动，但记录为明确的后续项** + +B2 的结果（§1）说明 expanded 路径**没有**吞吐问题，所以采用 `indices` 的收益 +主要是省掉 `[num_tokens, max_blocks_per_seq]` 那个 block table 的物化， +而不是速度。代价是把该路径钉死在 SM100（DeepGEMM varlen 分支断言 +`arch_major == 10 and next_n == 1`）。在 B2 已经证明 expanded 不慢的前提下， +这笔交易不划算 —— 除非 block table 的 H2D 在 128K 上下文下成为瓶颈（U7，未测）。 + +--- + +## 7. 还没做的 + +1. **e2e 分数**（§5）—— 均匀 tier 与 ragged 两条路径都没有分数。 +2. **B1 的对照实验**：`TORCH_LOGS=recompiles`，用来排除 §4 的候选 3。 +3. **两个并发点**（goal doc §6.5）—— 一个都没跑。 +4. **G2 逐 token 等价**、**A2 `cap-accept` 差分**。 +5. **B2 在 B≥64 且 tier≥3 的点**（§1 末尾）。 +6. `scripts/generate_llm_args_golden_manifest.py`：本轮没有新增 llm_args **字段** + （只加了 validator 和一个私有方法），预期 manifest 不变，但**未实际运行确认**。 From 08cfb4d9641ffe11703bd6c558b4a22618499eee Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 06:42:21 -0700 Subject: [PATCH 019/177] [None][doc] Confirm the golden manifest is unchanged and A3 stayed silent through capture Signed-off-by: Lance Liao --- docs/dspark_p0_report.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/dspark_p0_report.md b/docs/dspark_p0_report.md index bee1b52353c0..e15ced4a2989 100644 --- a/docs/dspark_p0_report.md +++ b/docs/dspark_p0_report.md @@ -193,7 +193,7 @@ goal doc 说得对：**贪婪 GSM8K 观察不到这一点**，所以只能从代 | `test_deepseek_v4_sparse_mla` 按每请求 q 长度参数化（宽度 1/2/4/6） | **12/12 通过**，但见 §0.3 的覆盖限制 | | compressor ragged 差分（`test_compressor_ragged.py`，新增） | **4/4 通过** | | C++ ragged top-k（`test_indexer_topk_ragged.py`，PR 已有） | 见 §5 | -| A3 布局一致性断言（`TLLM_DSPARK_ASSERT_LAYOUT=1`） | 已加，e2e 运行中启用 | +| A3 布局一致性断言（`TLLM_DSPARK_ASSERT_LAYOUT=1`） | 已加；e2e run 全程启用，**57 个 graph 的两轮 warmup/capture 无一触发** | | A4 时序断言 | 已加（无条件，代价是一次 host 求和） | | A6 planner 计数并入 ragged stats summary | 已加 | @@ -334,5 +334,6 @@ B2 的结果（§1）说明 expanded 路径**没有**吞吐问题，所以采用 3. **两个并发点**（goal doc §6.5）—— 一个都没跑。 4. **G2 逐 token 等价**、**A2 `cap-accept` 差分**。 5. **B2 在 B≥64 且 tier≥3 的点**（§1 末尾）。 -6. `scripts/generate_llm_args_golden_manifest.py`：本轮没有新增 llm_args **字段** - （只加了 validator 和一个私有方法），预期 manifest 不变，但**未实际运行确认**。 +6. ~~`scripts/generate_llm_args_golden_manifest.py`~~ —— **已确认**：在容器内重跑该脚本后 + `tensorrt_llm/usage/llm_args_golden_manifest.json` 无 diff。本轮只加了 validator 和一个 + 私有方法，没有新增用户可见字段，所以不需要 telemetry/privacy CODEOWNER 审批。 From a4c462c5089d04a01adec3b5a6b903f73bd5aa93 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 06:51:35 -0700 Subject: [PATCH 020/177] [None][doc] Record the DSpark unit-test results (202 passed) Signed-off-by: Lance Liao --- docs/dspark_p0_report.md | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/docs/dspark_p0_report.md b/docs/dspark_p0_report.md index e15ced4a2989..2caf61f82523 100644 --- a/docs/dspark_p0_report.md +++ b/docs/dspark_p0_report.md @@ -192,11 +192,20 @@ goal doc 说得对：**贪婪 GSM8K 观察不到这一点**，所以只能从代 |---|---| | `test_deepseek_v4_sparse_mla` 按每请求 q 长度参数化（宽度 1/2/4/6） | **12/12 通过**，但见 §0.3 的覆盖限制 | | compressor ragged 差分（`test_compressor_ragged.py`，新增） | **4/4 通过** | -| C++ ragged top-k（`test_indexer_topk_ragged.py`，PR 已有） | 见 §5 | +| C++ ragged top-k（`test_indexer_topk_ragged.py`，PR 已有） | **通过** | +| DSpark hw_agnostic 全套（confidence_schedule / ragged / worker / overlap_layout / heads / cuda_graph） | **通过** | | A3 布局一致性断言（`TLLM_DSPARK_ASSERT_LAYOUT=1`） | 已加；e2e run 全程启用，**57 个 graph 的两轮 warmup/capture 无一触发** | | A4 时序断言 | 已加（无条件，代价是一次 host 求和） | | A6 planner 计数并入 ragged stats summary | 已加 | +上述 DSpark 相关单测合并跑：**202 passed**（B300，13 秒）。 + +> 注意：`tests/unittest/_torch/speculative/hw_agnostic/` 整目录跑会有大量失败， +> 但全部是环境问题 —— 该目录里有需要 llama-3.1 / Qwen3 / EAGLE3 / DeepSeek-V3-Lite +> 权重的用例，而本轮的 `LLM_MODELS_ROOT` 只软链了 DeepSeek-V4-Pro-DSpark； +> 另有若干用例会自己起 `MpiPoolSession`，在已有 MPI world 里 spawn 失败。 +> 与本轮改动无关。 + ### 3.1 参数化过程中抓到的两个测试自身的 bug 1. `_allocate_kv_cache_for_generation` 每步只扩一个 token —— 只有在 From e59c8101f2c272f6e41c5ebd000be84258a2f849 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 07:01:44 -0700 Subject: [PATCH 021/177] [None][fix] Restore runtime_draft_len after the CUDA-graph capture passes The capture loop stamps `self.runtime_draft_len` per graph so each shape is built at its own draft length, and nothing puts it back. That was harmless while there was exactly one draft length to capture -- the value it was left at was the value it started with. A confidence tier ladder ends on its *smallest* tier. The engine therefore came out of capture claiming draft_len=1, and the generic token warmup that runs immediately afterwards built its dummy requests at the full draft length while the spec metadata still said 1. Acceptance reshapes the drafts into `[num_gens, runtime_draft_len]`, so it got `[128, 1]` for 640 tokens and raised: RuntimeError: shape '[128, 1]' is invalid for input of size 640 Found by running the confidence-scheduling configuration end to end for the first time. All 57 graphs capture cleanly; this fires in the warmup step just after, which is why no amount of capture-side work would have surfaced it. The leak predates this feature (it came in with dynamic draft length), but only a multi-tier ladder makes the leaked value differ from the correct one. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/pyexecutor/model_engine.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 53b99f7266b7..30b5dca6bc9d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2064,6 +2064,17 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: if force_non_greedy and spec_metadata is not None: spec_metadata._force_non_greedy_for_capture = False + # The capture loop stamps `self.runtime_draft_len` per graph so each + # captured shape is built at its own draft length. Nothing puts it back, + # which was harmless while there was one draft length to capture: the + # value it was left at was the value it started with. A tier ladder ends + # on its *smallest* tier, so the engine comes out of capture claiming a + # draft length it is not going to use, and the very next thing that runs + # -- the generic token warmup -- builds dummy requests at the full draft + # length while the spec metadata still says otherwise. Acceptance then + # reshapes `[num_gens * max_draft_len]` drafts into + # `[num_gens, smallest_tier]` and raises. + saved_runtime_draft_len = self.runtime_draft_len # Pass 1: greedy fast-path (dummy requests carry no sampling params, # so is_all_greedy_sample is naturally True). _run_capture_pass(force_non_greedy=False, label="greedy") @@ -2080,6 +2091,7 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: _run_capture_pass(force_non_greedy=True, label="advanced sampling") # Set the value back to the original value after cuda graph warmups are complete self.enable_spec_decode = self.is_spec_decode + self.runtime_draft_len = saved_runtime_draft_len # The advanced-sampling capture pass above leaves is_all_greedy_sample # set to False on spec_metadata. Reset it to the default so the first # real iteration's graph-key selection is not seeded with this From 301fe8bec1939c07bac4e9c64c75c0b93e324325 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 07:03:02 -0700 Subject: [PATCH 022/177] [None][doc] Record the runtime_draft_len leak the hang fix exposed Signed-off-by: Lance Liao --- docs/dspark_p0_report.md | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/docs/dspark_p0_report.md b/docs/dspark_p0_report.md index 2caf61f82523..b79ffb725846 100644 --- a/docs/dspark_p0_report.md +++ b/docs/dspark_p0_report.md @@ -253,7 +253,33 @@ TEST_EXIT=124 （graph 数从 102 变 57 是因为本轮配置的 batch-size 桶只有 19 个，与 hang 无关。） -### 4.2 三个候选的判定 +### 4.2 修完 hang 之后立刻暴露出来的第二个 bug（新发现） + +hang 消失之后，run 走到了**从来没有人到达过的地方**，于是撞上了下一个问题： + +``` +Run generation-only CUDA graph capture (advanced sampling) for batch size=1, draft_len=1 ← 57 个 graph 全部 capture 完成 +Run warmup with 4096 tokens, include 0 generation tokens +Run warmup with 128 tokens, include 128 generation tokens +RuntimeError: shape '[128, 1]' is invalid for input of size 640 +``` + +`640 = 128 × 5`。根因：capture 循环给每个 graph 盖上 +`self.runtime_draft_len = draft_len`（`model_engine.py:2056`），**结束后不还原** +（`:2081-2082` 只还原了 `enable_spec_decode`）。 + +只有一个 draft length 要 capture 时这无害 —— 留下的值就是原本的值。 +但 **tier ladder 结束在它最小的 tier 上**，于是 engine 带着 `runtime_draft_len=1` +走出 capture，而紧接着的通用 token warmup 按完整 draft 长度造 dummy 请求。 +acceptance 把 draft 重塑成 `[num_gens, runtime_draft_len]`，就成了 640 → `[128, 1]`。 + +这个泄漏**早于本 feature**（随 dynamic draft length 一起进来的）， +但只有多 tier 阶梯才会让泄漏值与正确值不同。已修（commit `e59c8101f2`）。 + +**这条本身就是 P0 阶段的一个论据**：capture 跑通不等于 feature 跑通， +而在 hang 被修掉之前，这个 bug 在物理上不可能被观察到。 + +### 4.3 三个候选的判定 | 候选 | 判定 | 依据 | |---|---|---| From 9762175fda1d603ab1dd73d538868fef93901eeb Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 07:28:19 -0700 Subject: [PATCH 023/177] [None][fix] Publish the gen-token layout before the seq_lens refresh decides Caught by the A4 assertion on the first ragged run that got far enough to capture: AssertionError: ragged verify windows sum to 512 but this step's token layout has 768 generation tokens; the windows belong to a different step than seq_lens does `_prepare_tp_inputs` decides whether to rewrite `attn_metadata.seq_lens` by asking whether this is a ragged step. It asked `attn_metadata` -- before this step's windows had been put there. On a freshly copied per-key graph metadata that reads None, so the refresh was skipped and `seq_lens` kept the uniform `ones(bs) * (1 + max_draft_tokens)` that `create_cuda_graph_metadata` seeds it with. The layout then claimed 128*6 tokens while the windows described the 128*4 bucket, which is the exact misattribution the ragged path exists to avoid. Moving the publish above the refresh makes the ordering monotone -- publish, refresh seq_lens, prepare() -- and leaves one publish point rather than two. Also teaches TLLM_DSPARK_FORCE_VERIFY_LENS to take a tier value (`=3` -> every request gets window 3) alongside the existing `=1` rotating form. This is what SPS profiling needs: sweeping the verify length by shrinking block_size shrinks the draft pass too, and that draft term is bilinear in (batch, verify_len) -- exactly the column space the additive cost fit spans, so it is absorbed into theta/alpha instead of appearing as residual, biasing the planner toward over-trimming with no detector. Forcing the window leaves block_size alone. SGLang solves the same problem with a runtime `dspark_force_budget_frac`. A value off the captured ladder raises rather than being ignored, since it would silently drop every step out of graph replay. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/model_engine.py | 25 ++++-- .../speculative/dspark_observability.py | 79 +++++++++++++++---- .../test_dspark_confidence_schedule.py | 66 ++++++++++++++++ 3 files changed, 145 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 30b5dca6bc9d..252f7af32166 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -5448,6 +5448,17 @@ def previous_seq_slots_device(): # from. The setter already copies in place under CUDA graphs (see its # comment: "The seqlens can change if we are doing spec decode"), so # this reallocates nothing as long as the row count is unchanged. + # Publish before the refresh decision below, not after: that decision + # asks whether this is a ragged step, and asking `attn_metadata` before + # this step's windows are on it reads the *previous* step's answer -- + # None on a freshly copied per-key graph metadata. The refresh is then + # skipped and `seq_lens` keeps the uniform `ones(bs) * (1 + + # max_draft_tokens)` that `create_cuda_graph_metadata` seeded it with, + # so the layout claims bs*(top_tier+1) tokens while the windows describe + # the trimmed bucket. + self._publish_gen_token_layout(attn_metadata, + scheduled_requests.generation_requests) + refresh_seq_lens = not attn_metadata.is_cuda_graph if (not refresh_seq_lens and getattr(attn_metadata, "is_ragged_verify", False) @@ -5508,14 +5519,12 @@ def previous_seq_slots_device(): # pre-prepare counts so the steady-gen recording below stores values # that the per-step prepare() can re-clamp from scratch. num_cached_tokens_snapshot = list(num_cached_tokens_per_seq) - # prepare() is the only consumer of the gen-token layout, so it has to - # be published first: it decides the DSA expanded-buffer layout, strides - # the expansions by this step's per-request token count, builds the - # per-row causal extents, and derives DeepSeek-V4's per-request - # compressor token counts. The rest of the layout is attached later, - # once the spec metadata exists. - self._publish_gen_token_layout(attn_metadata, - scheduled_requests.generation_requests) + # The gen-token layout was published above, before the seq_lens refresh + # that depends on it. prepare() is its other consumer: it decides the DSA + # expanded-buffer layout, strides the expansions by this step's + # per-request token count, builds the per-row causal extents, and derives + # DeepSeek-V4's per-request compressor token counts. The rest of the + # layout is attached later, once the spec metadata exists. attn_metadata.prepare() cross_attention_inputs = (self._prepare_enc_dec_cross_attn_inputs( cross_encoder_hidden_states, diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index 87a537cded8e..82463fd8d1dc 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -132,34 +132,79 @@ def forced_verify_lens(*, num_gen_requests: int, tiers: Sequence[int], min_verify_len: int) -> Optional[List[int]]: """A deterministic ragged split for correctness testing, or None. - Set ``TLLM_DSPARK_FORCE_VERIFY_LENS=1`` to have every generation request - take a window from the captured tier ladder, rotating by batch position. - - This exists because the two failure modes are otherwise inseparable. The - planner refuses to trim without a profiled cost table -- correctly, since a - flat model makes every extra token look free -- so a run without one keeps - the ragged path dark no matter what the config says. Testing that the packing - is correct would then require first producing a cost table, which requires a - working run. This breaks that circle. - - It is not a shortcut around the planner: it decides only *how many* drafted - positions reach the target, exactly the quantity the planner decides. + Two forms: + + ``TLLM_DSPARK_FORCE_VERIFY_LENS=1`` + Every generation request takes a window from the captured tier ladder, + rotating by batch position -- a deliberately *non-uniform* split, which + is what exercises the ragged packing. + + ``TLLM_DSPARK_FORCE_VERIFY_LENS=`` (n a tier, e.g. ``3``) + Every request takes the same window ``n``. Uniform, so it exercises + nothing about raggedness -- its purpose is cost measurement. + + The second form exists for SPS profiling. Sweeping the verify length by + shrinking ``block_size`` also shrinks the *draft* pass, but at deployment the + block is always drafted in full, so the measured step cost carries a draft + term that production does not pay. That term is bilinear in + ``(batch, verify_len)`` and therefore lies exactly in the column space the + additive fit spans, which means it is absorbed silently: theta comes out + inflated and alpha deflated, both biasing the planner toward over-trimming, + and the fit residual cannot see it. Forcing the window instead leaves + ``block_size`` at its full value, so the draft cost is constant across the + sweep and the bias is gone. SGLang solves the same problem with a runtime + ``dspark_force_budget_frac`` knob. + + Either form also breaks a circularity: the planner refuses to trim without a + profiled cost table -- correctly, since a flat model makes every extra token + look free -- so a run without one keeps the ragged path dark no matter what + the config says, and producing a table would require a working run. + + Neither form is a shortcut around the planner: both decide only *how many* + drafted positions reach the target, exactly the quantity the planner decides. Acceptance is unchanged, so a forced split must produce the same answers as ``static`` -- which is what makes it a usable correctness gate. Returns: One window per request, each drawn from ``tiers`` (so every choice has a captured CUDA graph), or None when the override is off. + + Raises: + ValueError: the value is neither a boolean form nor a captured tier. + Silently ignoring it would produce a run that looks forced and is + not -- the exact failure this override exists to rule out. """ - if os.environ.get(FORCE_VERIFY_LENS_ENV, "") not in ("1", "true", "True"): + value = os.environ.get(FORCE_VERIFY_LENS_ENV, "").strip() + if not value or value in ("0", "false", "False"): return None ladder = sorted({int(t) for t in tiers if int(t) >= int(min_verify_len)}) if not ladder: return None - # Rotate through the ladder by position: deterministic, identical on every - # rank for a given batch position, and guaranteed to be non-uniform as soon - # as the batch is larger than one and the ladder has more than one entry. - return [ladder[i % len(ladder)] for i in range(int(num_gen_requests))] + if value in ("1", "true", "True"): + # Rotate through the ladder by position: deterministic, identical on + # every rank for a given batch position, and guaranteed to be non-uniform + # as soon as the batch is larger than one and the ladder has more than + # one entry. + # + # Note "1" is the rotating form, not the window 1. A single tier ladder + # makes the two coincide anyway, and the rotating form is the older + # meaning; spelling the uniform window as a value keeps both reachable. + return [ladder[i % len(ladder)] for i in range(int(num_gen_requests))] + try: + window = int(value) + except ValueError: + raise ValueError( + f"invalid {FORCE_VERIFY_LENS_ENV}={value!r}; expected 1 (rotate " + f"through the tier ladder) or one of the captured tiers {ladder}") + if window not in ladder: + raise ValueError( + f"{FORCE_VERIFY_LENS_ENV}={window} is not a captured tier; the " + f"available tiers are {ladder}. A window with no captured CUDA " + f"graph does not raise at runtime -- it silently drops every step " + f"out of graph replay, which costs far more than the trimmed " + f"tokens save and would make any measurement taken with it " + f"meaningless.") + return [window] * int(num_gen_requests) class DSparkRaggedStats: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py index fd3bf13774e0..65a35e0bc7a4 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py @@ -763,3 +763,69 @@ def test_cross_rank_reduction_is_max_not_min(): confs = [torch.full((8, BLOCK), -6.0), torch.full((8, BLOCK), 6.0)] chosen, locals_ = _simulated_ranks(confs, tiers, table) assert chosen[0] == max(locals_) + + +def test_forced_verify_lens_accepts_a_uniform_window(monkeypatch): + """``TLLM_DSPARK_FORCE_VERIFY_LENS=`` pins every request to one window. + + This is the form SPS profiling needs. Sweeping the verify length by shrinking + ``block_size`` shrinks the draft pass too, and that draft term is bilinear in + (batch, verify_len) -- exactly the column space the additive cost fit spans, + so it is absorbed into theta/alpha instead of showing up as residual. Forcing + the window leaves block_size alone, so the draft cost is constant across the + sweep. + """ + from tensorrt_llm._torch.speculative.dspark_observability import ( + FORCE_VERIFY_LENS_ENV, forced_verify_lens) + + monkeypatch.setenv(FORCE_VERIFY_LENS_ENV, "3") + lens = forced_verify_lens(num_gen_requests=5, tiers=[1, 3, 5], + min_verify_len=1) + assert lens == [3, 3, 3, 3, 3] + + +def test_forced_verify_lens_rotates_on_the_boolean_form(monkeypatch): + """``=1`` keeps its original meaning: rotate, i.e. deliberately non-uniform.""" + from tensorrt_llm._torch.speculative.dspark_observability import ( + FORCE_VERIFY_LENS_ENV, forced_verify_lens) + + monkeypatch.setenv(FORCE_VERIFY_LENS_ENV, "1") + lens = forced_verify_lens(num_gen_requests=5, tiers=[1, 3, 5], + min_verify_len=1) + assert lens == [1, 3, 5, 1, 3] + assert len(set(lens)) > 1, "the rotating form must not be uniform" + + +def test_forced_verify_lens_rejects_an_uncaptured_window(monkeypatch): + """A window off the ladder must raise, not be quietly ignored. + + It has no captured CUDA graph, so it would not fail loudly at runtime -- it + would drop every step into eager, which costs more than the trimmed tokens + save and makes any measurement taken with it meaningless. + """ + import pytest + + from tensorrt_llm._torch.speculative.dspark_observability import ( + FORCE_VERIFY_LENS_ENV, forced_verify_lens) + + monkeypatch.setenv(FORCE_VERIFY_LENS_ENV, "4") + with pytest.raises(ValueError, match="not a captured tier"): + forced_verify_lens(num_gen_requests=2, tiers=[1, 3, 5], + min_verify_len=1) + + monkeypatch.setenv(FORCE_VERIFY_LENS_ENV, "banana") + with pytest.raises(ValueError, match="expected 1"): + forced_verify_lens(num_gen_requests=2, tiers=[1, 3, 5], + min_verify_len=1) + + +def test_forced_verify_lens_off_by_default(monkeypatch): + from tensorrt_llm._torch.speculative.dspark_observability import ( + FORCE_VERIFY_LENS_ENV, forced_verify_lens) + + monkeypatch.delenv(FORCE_VERIFY_LENS_ENV, raising=False) + assert forced_verify_lens(num_gen_requests=3, tiers=[1, 3, 5], + min_verify_len=1) is None + monkeypatch.setenv(FORCE_VERIFY_LENS_ENV, "0") + assert forced_verify_lens(num_gen_requests=3, tiers=[1, 3, 5], + min_verify_len=1) is None From d3cecc68456adf2f31e66302f3f102a9178996b1 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 07:40:45 -0700 Subject: [PATCH 024/177] [None][test] Let the ragged GSM8K case take an SPS cost table Without a table the planner is provably unable to trim: a flat table short-circuits `decide_verify_lens` before confidence is ever read, so the ragged path stays dark and the test degenerates into the uniform baseline while scoring identically. TLLM_DSPARK_FORCE_VERIFY_LENS bypasses the planner instead of feeding it, so it exercises the ragged packing and nothing about the decision. Reads `TLLM_DSPARK_SPS_TABLE` so a run can point at one. The table does not have to be accurate to be useful: scheduling only chooses how many drafted positions are submitted for verification and acceptance is unchanged, so an approximate table cannot move the output distribution -- only the throughput it was tuned for. That makes an approximate table a legitimate way to get GSM8K to exercise the real path. Signed-off-by: Lance Liao --- .../defs/accuracy/test_llm_api_pytorch.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 340fcddbf770..d47e23ce7a6b 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4136,10 +4136,21 @@ def test_gsm8k_dep8_ragged_verify(self, moe_backend): """ kv_cache_config = KvCacheConfig(enable_block_reuse=False, free_gpu_memory_fraction=0.5) - spec_config = DSparkDecodingConfig(max_draft_len=5, - speculative_model=self.MODEL_PATH, - enable_confidence_scheduling=True, - enable_ragged_verify=True) + # Without a cost table the planner is provably unable to trim: a flat + # table short-circuits `decide_verify_lens` before confidence is even + # read, so the ragged path stays dark and this test degenerates into the + # uniform baseline while still scoring the same. Point this at a table to + # exercise the real decision path. A table need not be *accurate* to be + # useful here -- scheduling only chooses how many drafted positions are + # submitted, and acceptance is unchanged, so an approximate table cannot + # change the output distribution, only the throughput it was picked for. + sps_table_path = os.environ.get("TLLM_DSPARK_SPS_TABLE") or None + spec_config = DSparkDecodingConfig( + max_draft_len=5, + speculative_model=self.MODEL_PATH, + enable_confidence_scheduling=True, + enable_ragged_verify=True, + confidence_sps_table_path=sps_table_path) with LLM(self.MODEL_PATH, attn_backend="TRTLLM", tensor_parallel_size=8, From 01625c90bc0c5cfa035e1d4d3319f5488333a109 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 08:16:39 -0700 Subject: [PATCH 025/177] [None][fix] Make the DSpark "did the feature do anything" gate usable `DSparkRaggedStats.assert_ragged_active` exists precisely because the failure mode here is not a wrong answer but a feature that quietly did nothing -- a run where the planner declined on every step scores exactly what a working one scores. It had zero call sites repo-wide, and two things stood in the way of adding one. First, it could not have passed. `_record_dspark_graph_use` was called on every forward, including context steps that have no generation graph to miss, so `graph_eager` was really a prefill counter -- and the check refuses a run with any eager step. Gate it on `can_run_cuda_graph` so it counts what it claims to. Second, the counters live in the worker process and under TP8 + attention-DP the driver cannot reach them: `collective_rpc` refuses a multi-rank world. So the worker checks itself. `_executor_loop_cleanup` now logs the final summary unconditionally -- previously the numbers only appeared if the 32-step periodic emit happened to land -- and with `TLLM_DSPARK_ASSERT_ACTIVE=1` it asserts on them. A worker exception does propagate to the driver, which is the one channel that crosses the MPI boundary reliably. Stats reporting never blocks shutdown: everything except the intended AssertionError is caught and logged. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/model_engine.py | 10 +++++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 35 +++++++++++++++++++ 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 252f7af32166..ddf9b36cab51 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -6888,7 +6888,15 @@ def forward(self, # quietly runs eager, which costs far more than the tokens it # trimmed. Record it so the miss is a counter rather than something # only a profiler would reveal. - self._record_dspark_graph_use(replayed=can_run_graph) + # + # Only generation-only steps count. A batch carrying context + # requests has no generation graph to miss, so recording it as an + # eager fallback would make `graph_eager` a prefill counter -- and + # `assert_ragged_active` refuses a run with any eager step, so that + # miscount alone makes the check unsatisfiable on every real + # workload. + if padded_requests.can_run_cuda_graph: + self._record_dspark_graph_use(replayed=can_run_graph) if can_run_graph: attn_metadata = maybe_attn_metadata spec_metadata = maybe_spec_metadata diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index c05bb9b90f0f..3992b4caa2ea 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2399,7 +2399,42 @@ def _process_iter_stats( prev_device_step_time_ms=prev_device_step_time_ms, gpu_forward_time_ms=gpu_forward_time_ms) + def _report_dspark_ragged_stats(self) -> None: + """Final DSpark scheduling summary, and optionally a hard gate on it. + + The counters live in the worker process. Under TP8 + attention-DP the + driver cannot reach them -- ``collective_rpc`` refuses a multi-rank + world -- so a test can assert on a GSM8K score and nothing else. But a + run where the planner declined on every step scores *exactly* the same + as one where it worked, so that assertion is evidence-free: the known + failure mode here is not a wrong answer, it is a feature that quietly + did nothing. + + Two things close that. The summary is logged unconditionally at + shutdown, so the numbers are in the log of every run rather than only in + whatever the 32-step periodic emit happened to catch. And with + ``TLLM_DSPARK_ASSERT_ACTIVE=1`` the worker asserts on its own counters + here -- an exception in the worker does propagate to the driver, which + is the one channel that crosses the MPI boundary reliably. + """ + worker = getattr(self.model_engine, "_get_spec_worker", None) + stats = getattr(worker(), "ragged_stats", None) if worker else None + if stats is None or not getattr(stats, "steps_total", 0): + return + stats.log_summary(prefix="DSpark ragged verify [final]") + if os.environ.get("TLLM_DSPARK_ASSERT_ACTIVE", "") in ("1", "true", + "True"): + # require_trim is deliberately on: a run that produced windows but + # never shortened one delivered the uniform baseline. + stats.assert_ragged_active(require_trim=True) + def _executor_loop_cleanup(self): + try: + self._report_dspark_ragged_stats() + except AssertionError: + raise + except Exception as exc: # noqa: BLE001 - never block shutdown on stats + logger.warning(f"DSpark ragged stats reporting failed: {exc}") # Wake any waiters in await_responses BEFORE potentially-blocking # work below. If wait_on_pp_send_handles hangs (e.g. after a # crash leaves PP send handles in a bad state), the await loop From ccaf9c487cdddfebfbe3b454ceb23626dd72973e Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 09:10:03 -0700 Subject: [PATCH 026/177] [None][test] Add an opt-in CUDA sync detector for the DSpark decode path The overlap scheduler's whole value is that the host runs ahead of the device, so a single device->host sync per decode step gives it back. Static review can show the ragged path *looks* sync-free -- every repeat_interleave carries output_size, the planner works on a pinned host snapshot, staging is non_blocking -- but only the allocator's own detector shows that it *is*. `TLLM_DSPARK_SYNC_DEBUG=` arms `torch.cuda.set_sync_debug_mode("warn")` after n DSpark decode iterations. Armed late on purpose: warmup, graph capture and the first steady steps legitimately synchronize, so arming at startup buries the signal under expected hits. 'warn' rather than 'error' so one run surfaces every offender instead of dying at the first, and each hit carries a stack, which is what makes the offender identifiable rather than merely countable. Hooked into the DSpark decode hook rather than the generic loop, so it costs one attribute read on any other model and nothing at all when unset. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 3992b4caa2ea..a24dca96d4fd 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -558,6 +558,9 @@ def __init__( self.peft_cache_config = peft_cache_config self.iter_counter = 0 + # TLLM_DSPARK_SYNC_DEBUG bookkeeping; see _maybe_arm_dspark_sync_debug. + self._dspark_sync_debug_iters = 0 + self._dspark_sync_debug_armed = False # profile config self.profile_start_iters, self.profile_stop_iters = _load_iteration_indexes( PROFILE_START_STOP_ENV_VAR_NAME) @@ -2399,6 +2402,48 @@ def _process_iter_stats( prev_device_step_time_ms=prev_device_step_time_ms, gpu_forward_time_ms=gpu_forward_time_ms) + def _maybe_arm_dspark_sync_debug(self) -> None: + """Turn on CUDA sync detection once decode is in steady state. + + The overlap scheduler's entire value is that the host runs ahead of the + device, so one device->host sync per decode step gives it back. Static + review can show the code *looks* sync-free; only the allocator's own + detector shows that it *is*. + + Armed late on purpose. Warmup, CUDA-graph capture and the first steady + steps legitimately synchronize, so arming at startup would bury the + signal in expected hits. ``TLLM_DSPARK_SYNC_DEBUG=`` arms after ``n`` + DSpark decode iterations (``1`` means "on", which arms after the default + 64). Each subsequent sync emits a Python warning with a stack, which is + what makes the offender identifiable rather than merely countable. + + Diagnostic only: never on by default, and 'warn' rather than 'error' so + a run surfaces every offender instead of dying at the first. + """ + if self._dspark_sync_debug_armed: + return + raw = os.environ.get("TLLM_DSPARK_SYNC_DEBUG", "") + if not raw or raw in ("0", "false", "False"): + return + try: + after = 64 if raw in ("1", "true", "True") else int(raw) + except ValueError: + logger.warning( + f"ignoring TLLM_DSPARK_SYNC_DEBUG={raw!r}: expected an integer " + f"iteration count, or 1 for the default") + self._dspark_sync_debug_armed = True + return + self._dspark_sync_debug_iters += 1 + if self._dspark_sync_debug_iters < after: + return + torch.cuda.set_sync_debug_mode("warn") + self._dspark_sync_debug_armed = True + logger.warning( + f"DSpark: CUDA sync detection armed after " + f"{self._dspark_sync_debug_iters} decode iterations. Any " + f"device->host synchronization from here on emits a warning with a " + f"stack; on the overlap-scheduler path there should be none.") + def _report_dspark_ragged_stats(self) -> None: """Final DSpark scheduling summary, and optionally a hard gate on it. @@ -3240,6 +3285,8 @@ def _dspark_confidence_draft_len(self, if planner is None: return max_draft_len + self._maybe_arm_dspark_sync_debug() + gen_requests = scheduled_batch.generation_requests num_gen_requests = len(gen_requests) # Confidence is read back by SLOT, not by batch position. The snapshot From 05c67cfdb113e0253ee1fda1a9d8c77d3e9cf6be Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 09:58:38 -0700 Subject: [PATCH 027/177] [None][fix] TLLM_DSPARK_ASSERT_ACTIVE takes a step floor, not a boolean An engine build spins up more than one executor: KV-cache estimation runs a throwaway one that completes a couple of decode steps and shuts down. Asserting on that one kills the run before the real work starts, and reports "the feature did nothing" about an executor that was never asked to do anything -- observed as steps_total=1 with 228 graph replays. The value is now a minimum decode-step count. Below it the assertion is skipped with a log line saying so, which is what distinguishes "the scheduler declined" from "this executor barely ran". Signed-off-by: Lance Liao --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 31 ++++++++++++++++--- 1 file changed, 26 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a24dca96d4fd..bc22eb45d1ab 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2467,11 +2467,32 @@ def _report_dspark_ragged_stats(self) -> None: if stats is None or not getattr(stats, "steps_total", 0): return stats.log_summary(prefix="DSpark ragged verify [final]") - if os.environ.get("TLLM_DSPARK_ASSERT_ACTIVE", "") in ("1", "true", - "True"): - # require_trim is deliberately on: a run that produced windows but - # never shortened one delivered the uniform baseline. - stats.assert_ragged_active(require_trim=True) + raw = os.environ.get("TLLM_DSPARK_ASSERT_ACTIVE", "") + if not raw or raw in ("0", "false", "False"): + return + # The value is a MINIMUM STEP COUNT, not a boolean. An engine build + # spins up more than one executor -- KV-cache estimation runs a + # throwaway one that completes a couple of decode steps and shuts down + # -- and asserting on that one kills the run before the real work + # starts, reporting "the feature did nothing" about an executor that was + # never asked to do anything. Requiring a floor is what distinguishes + # "the scheduler declined" from "this executor barely ran". + try: + min_steps = 32 if raw in ("1", "true", "True") else int(raw) + except ValueError: + logger.warning( + f"ignoring TLLM_DSPARK_ASSERT_ACTIVE={raw!r}: expected a " + f"minimum decode-step count, or 1 for the default") + return + if stats.steps_total < min_steps: + logger.info( + f"DSpark ragged verify: skipping the active-path assertion, " + f"this executor ran {stats.steps_total} decode step(s) " + f"(< {min_steps}).") + return + # require_trim is deliberately on: a run that produced windows but + # never shortened one delivered the uniform baseline. + stats.assert_ragged_active(require_trim=True) def _executor_loop_cleanup(self): try: From 742ce5d97c6971a2e5a471a37e198203ccef42b6 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 18:12:55 -0700 Subject: [PATCH 028/177] [None][feat] Record whether trimming actually cost acceptance SGLang's headline DSpark metric is accept_length / block_accept_length -- what fraction of the uncapped acceptance survived trimming. TRT-LLM knew the numerator (acceptance length) and the window, but had no ceiling term: `ceiling_tokens` counts tokens *submitted* with no trim, which answers "how much compute did we save" and never "how much acceptance did we destroy". Without that, a planner regression -- bad calibration, stale cost table, ladder drift -- shows up only as a slowly falling acceptance length, indistinguishable from a harder workload. Rather than estimate the uncapped length, this measures the thing directly. A trimmed request that accepted its ENTIRE window was still alive when the cut happened, so that cut certainly cost acceptance; one that accepted fewer than its window would have died anyway and cost nothing. `trim_regret_rate` is the share of trimmed requests in the first group -- exact, unbiased, and needing no model. Extrapolating from the untrimmed requests instead would be biased by construction, since the planner chose which requests to leave alone. Recorded in the sampler because that is the only place the accepted count and the window it was measured against are both host-side and belong to the same step -- the live py_verify_len has already moved to the next step by then, which is what the snapshot exists for. The worker owns the counters, so the executor connects the two the first time it runs a DSpark decode step. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 8 +++ .../speculative/dspark_observability.py | 58 +++++++++++++++++++ .../_torch/speculative/spec_sampler_base.py | 22 +++++-- .../test_dspark_confidence_schedule.py | 47 +++++++++++++++ 4 files changed, 130 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index bc22eb45d1ab..4b1e59897571 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3325,6 +3325,14 @@ def _dspark_confidence_draft_len(self, # `static` and `compact` without editing the serving config, and every # counter below is keyed off the same decision. stats = getattr(worker, "ragged_stats", None) + # The sampler is built independently of the worker, so nothing connects + # them; do it here, where both are reachable. Acceptance has to be + # recorded in the sampler (that is where the accepted count and the + # window it was measured against are both host-side and belong to the + # same step), but the counters live on the worker. + if (stats is not None and self.sampler is not None + and getattr(self.sampler, "acceptance_stats", None) is None): + self.sampler.acceptance_stats = stats mode = getattr(worker, "ragged_verify_mode", None) compute_windows = (mode.computes_windows if mode is not None else getattr(self.model_engine.spec_config, diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index 82463fd8d1dc..73d87e584d4b 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -248,6 +248,25 @@ def __init__(self, *, mode: RaggedVerifyMode, max_draft_len: int): self.bucket_hist: Counter = Counter() self.padded_bs_hist: Counter = Counter() + #: Acceptance, split by whether the request's window was shortened. + #: This is the term the trim/ceiling ratio needs and the one TRT-LLM + #: did not have: `ceiling_tokens` above counts tokens *submitted* with + #: no trim, which answers "how much compute did we save" and never + #: "how much acceptance did we destroy". + self.accepted_tokens = 0 + self.requests_scored = 0 + #: Requests whose window was shorter than the full block. + self.requests_trimmed = 0 + #: Trimmed requests that accepted their ENTIRE window. Those are the + #: ones where the draft was still alive at the cut, so trimming + #: certainly cost acceptance -- a trimmed request that accepted fewer + #: than its window would have died on its own and cost nothing. This is + #: an exact, unbiased count, which is why it is preferred here over + #: estimating the uncapped acceptance length: no model, no extrapolation + #: from the untrimmed subpopulation (which the planner selected, so it + #: is not a fair sample of the trimmed one). + self.trimmed_hit_ceiling = 0 + #: Graph replay is the whole point of landing on a bucket; an eager #: step costs far more than the tokens it saved. self.graph_replays = 0 @@ -338,6 +357,41 @@ def distinct_verify_lens(self) -> int: """How many different window sizes were ever handed out.""" return len(self.verify_len_hist) + def record_acceptance(self, *, accepted: int, window: int) -> None: + """Record one request's acceptance against the window it was given. + + Args: + accepted: drafted positions accepted (excludes the bonus token). + window: drafted positions the request was allowed to verify. + """ + self.requests_scored += 1 + self.accepted_tokens += int(accepted) + if int(window) < self.max_draft_len: + self.requests_trimmed += 1 + if int(accepted) >= int(window): + self.trimmed_hit_ceiling += 1 + + @property + def accept_len(self) -> float: + """Mean accepted drafted positions per request. 0.0 with no data.""" + if not self.requests_scored: + return 0.0 + return self.accepted_tokens / self.requests_scored + + @property + def trim_regret_rate(self) -> float: + """Share of trimmed requests whose draft was still alive at the cut. + + The operational reading: 0.0 means every trim was free -- those drafts + would have died anyway -- and the scheduler is taking compute back for + nothing. Climbing toward 1.0 means trimming is buying its throughput by + throwing away acceptance, which is the regression a delivered-only + acceptance metric cannot distinguish from a harder workload. + """ + if not self.requests_trimmed: + return 0.0 + return self.trimmed_hit_ceiling / self.requests_trimmed + def summary(self) -> Dict[str, object]: """A JSON-safe snapshot, for logging or assertions.""" return { @@ -354,6 +408,10 @@ def summary(self) -> Dict[str, object]: "verify_len_hist": dict(sorted(self.verify_len_hist.items())), "bucket_hist": dict(sorted(self.bucket_hist.items())), "padded_bs_hist": dict(sorted(self.padded_bs_hist.items())), + "accept_len": round(self.accept_len, 4), + "requests_scored": self.requests_scored, + "requests_trimmed": self.requests_trimmed, + "trim_regret_rate": round(self.trim_regret_rate, 4), "graph_replays": self.graph_replays, "graph_eager": self.graph_eager, "fallbacks": dict(self.fallbacks), diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 59271b21be42..bfaec06dc25f 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -181,6 +181,10 @@ def _request_common_handling( request.py_draft_tokens = next_draft_tokens[request.py_seq_slot][:runtime_draft_len] request.py_decoding_iter += 1 + #: Optional observability sink, attached by the DSpark worker. Left None + #: everywhere else so this costs one attribute read per request. + acceptance_stats = None + @staticmethod def _verified_len(request: LlmRequest, runtime_draft_len: int, verify_lens_snapshot: Optional[dict] = None) -> int: @@ -241,11 +245,19 @@ def update_requests( ): break req.py_num_accepted_draft_tokens = num_new_tokens - 1 - req.py_rewind_len = ( - self._verified_len(req, runtime_draft_len, - state.verify_lens_snapshot) - - req.py_num_accepted_draft_tokens - ) + verified_len = self._verified_len(req, runtime_draft_len, + state.verify_lens_snapshot) + req.py_rewind_len = verified_len - req.py_num_accepted_draft_tokens + # Acceptance against the window the request was actually given. + # Recorded here because this is the only place both numbers are + # host-side and belong to the same step -- the live py_verify_len + # has already moved on by now, which is why the snapshot exists. + # `acceptance_stats` is attached by the DSpark worker and is None + # for every other path. + if self.acceptance_stats is not None: + self.acceptance_stats.record_acceptance( + accepted=req.py_num_accepted_draft_tokens, + window=verified_len) self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) def sample_async( diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py index 65a35e0bc7a4..915728114bcd 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py @@ -829,3 +829,50 @@ def test_forced_verify_lens_off_by_default(monkeypatch): monkeypatch.setenv(FORCE_VERIFY_LENS_ENV, "0") assert forced_verify_lens(num_gen_requests=3, tiers=[1, 3, 5], min_verify_len=1) is None + + +def _stats(max_draft_len=5): + from tensorrt_llm._torch.speculative.dspark_observability import ( + DSparkRaggedStats, RaggedVerifyMode) + return DSparkRaggedStats(mode=RaggedVerifyMode.COMPACT, + max_draft_len=max_draft_len) + + +def test_trim_regret_counts_only_drafts_alive_at_the_cut(): + """A trimmed request that accepted its whole window lost something. + + This is the quantity a delivered-only acceptance metric cannot recover: it + separates "trimming was free, those drafts would have died anyway" from + "trimming bought throughput by discarding acceptance". + """ + stats = _stats() + # Given the full block and died early: not trimmed, no regret. + stats.record_acceptance(accepted=2, window=5) + # Trimmed to 3 and died at 1: the cut cost nothing. + stats.record_acceptance(accepted=1, window=3) + # Trimmed to 3 and accepted all 3: still alive at the cut -> regret. + stats.record_acceptance(accepted=3, window=3) + + assert stats.requests_scored == 3 + assert stats.requests_trimmed == 2 + assert stats.trimmed_hit_ceiling == 1 + assert stats.trim_regret_rate == 0.5 + assert stats.accept_len == pytest.approx(2.0) + + +def test_trim_regret_is_zero_when_nothing_was_trimmed(): + stats = _stats() + for accepted in (0, 3, 5): + stats.record_acceptance(accepted=accepted, window=5) + assert stats.requests_trimmed == 0 + assert stats.trim_regret_rate == 0.0 + assert stats.summary()["accept_len"] == pytest.approx(8 / 3, abs=1e-4) + + +def test_acceptance_metrics_appear_in_summary(): + stats = _stats() + stats.record_acceptance(accepted=3, window=3) + summary = stats.summary() + for key in ("accept_len", "requests_scored", "requests_trimmed", + "trim_regret_rate"): + assert key in summary, f"{key} missing from the summary" From 46be7683192566b2d794a2220844a3ed6bcafe9c Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 18:18:17 -0700 Subject: [PATCH 029/177] [None][doc] Record the sync audit, kernel-compat proof and the new trim-regret metric Signed-off-by: Lance Liao --- docs/dspark_p0_report.md | 48 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/docs/dspark_p0_report.md b/docs/dspark_p0_report.md index b79ffb725846..b5bb3642b917 100644 --- a/docs/dspark_p0_report.md +++ b/docs/dspark_p0_report.md @@ -290,6 +290,54 @@ acceptance 把 draft 重塑成 `[num_gens, runtime_draft_len]`，就成了 640 **结论：hang 的真实根因是 H1（连带 K8）—— DSA 展开 stride 用静态 max 而非 runtime tier。** 候选 2 是无害的共存现象；候选 3 未被证明有害。 +## 4.4 host-device sync 验证（overlap scheduler 不受影响） + +**静态审计**（逐行核实 `dspark_ragged.py` / `dspark_verify.py` / `dspark.py` / `heads.py`）： + +| 命中 | 判定 | +|---|---| +| `dspark_verify.py` 两处 `.numpy()` | ✅ survival 来自 `_host_buffer`（**pinned CPU**）；planner 全程在 host 快照上算 | +| `dspark_ragged.py` 两处 `.tolist()` | ✅ `verify_lens` 是 CPU 张量 | +| `dspark_ragged.py` `row_ids_from_lens` | ✅ `output_size` 是**必填参数**，docstring 说明这里 sync 是 "outright illegal" | +| `dspark.py:634` `qo_indptr.to(device)` | ✅ 同 device，no-op | +| `dspark.py:551` `chunk_positions[0].item()` | ⚠️ 真 sync，但在 **prefill** 路径且由 base DSpark `147404864d` (#15808) 引入，**不属本 PR** | + +加上本轮已修的两处（`repeat_interleave` 补 `output_size`、`apply_sts` 缓存 host 副本）， +**decode 路径上 confidence/ragged 没有新增 sync**。 + +静态分析不是证明，所以加了 `TLLM_DSPARK_SYNC_DEBUG=`：在 n 个 DSpark decode 迭代后 +arm `torch.cuda.set_sync_debug_mode("warn")`。**故意晚 arm** —— warmup 与 graph capture +本就合法地同步，一开始就开会把信号埋掉。用 `warn` 不用 `error`，一次 run 能看到全部违规者。 +**实测结果待补。** + +## 4.5 kernel 改动兼容性 + +K4 / K9 是 PR 自带的（`e1c28eeae2`），本轮验证其「纯加性」性质： + +| 判据 | 证据 | +|---|---| +| (a) 不传新参数时 uniform 路径不变 | `indexerTopK.cu:680-688` 的 `else` 分支是**原公式逐字保留**；`compressorKernels.cu:348` 在指针为空时 `nn == NEXT_N`，`sp = kv_len - NEXT_N` 即原式，且 `NEXT_N` 仍是模板常量（保住全展开） | +| (a') bitwise 差分实测 | compressor：uniform 填充 == 传 `None`，**4/4 通过**；indexer top-k：`test_ragged_with_equal_lens_matches_uniform`，在 202 passed 内 | +| (b) ragged 结果正确 | compressor 与**单请求 uniform** 逐个比对（避开「kernel 直接忽略参数也能过」的盲点），通过 | +| (c) 互斥组合显式拒绝 | kernel `TLLM_CHECK`（`indexerTopK.cu:921`）+ 本轮新增的 Python config reject | +| (d) G6 两开关关闭零回归 | **穷举证明**：新 stride 与原 `1 + max_draft_tokens` 对 MTP / DSpark / PARD / DFlash 在 K=1,3,5,8 上全部相等。唯一变化发生在 `runtime_draft_len < max`（动态 draft 长度），那是修 bug | + +## 4.6 新增的可观测性：trimming 到底有没有代价 + +SGLang 博客（https://www.lmsys.org/blog/2026-07-06-dspark-sglang）的核心指标是 +`accept_length / block_accept_length`（报 0.88–0.97）。TRT-LLM 缺的是**分母** —— +`ceiling_tokens` 数的是「不裁会**提交**多少 token」，回答「省了多少算力」， +永远回答不了「毁了多少接受率」。后果：planner 过度裁剪只表现为端到端接受长度缓慢下降， +与「workload 变难」不可区分。 + +**没有照抄 SGLang 的估计器**，而是直接测那件事：被裁到窗口 `w` 的请求若**接受了整个 `w`**， +说明剪的时候 draft 还活着，这一刀确实有代价；接受数 `< w` 的本来就会死，剪了不亏。 +`trim_regret_rate` = 被裁请求中前者的占比 —— **精确、无偏、不需要模型**。 +（从未被裁的请求外推是**构造性有偏**的：哪些不裁正是 planner 挑的。） + +记录点在 sampler，因为那是唯一一处「接受数」与「它对应的窗口」同时在 host 侧且属于同一步 +的地方 —— 活的 `py_verify_len` 那时已翻到下一步（正是 H6 快照存在的原因）。 + ## 5. 交付物 3：e2e 实测 —— **未完成** - 均匀 tier 路径 + overlap scheduler + DEP8 的 GSM8K：**运行中，未取得分数**。 From 3b3568af840f49c66ef4256177aed99272b1fe5b Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 18:37:52 -0700 Subject: [PATCH 030/177] [None][doc] K1 blocks ragged verification on DSv4; the PR never addressed it Running the ragged path end to end past warmup hits, on the first real decode: Assertion failed: MLA can only support input sequences with the same sequence length. (mlaKernels.cu:1161) which is goal doc K1 verbatim. `git show --stat e1c28eeae2` -- the commit titled "Make DSpark ragged verification work on the DSv4 kernel path" -- shows it touched IndexerTopK (K4) and compressorKernels (K9) and nothing else. `mlaKernels.cu` (K1) and `attentionOp.cpp` (K2), both rated blockers by the goal doc, were never touched. So the ragged path could not have run end to end, and the commit title claims more than it delivered. Worth stating because the alternative failure is worse: the goal doc shows that window combinations whose token total happens to divide the batch size pass the check and then write each request's K/V to the wrong rows -- silent, persistent KV corruption. This run hit the loud half. The fix is the token-major presentation that D1 and D2 both rely on, and B2 already measured it as costing nothing (0.93x-1.00x). The route is viable; the work simply has not been done. This also answers "is ragged faster at high concurrency" for now: it cannot be measured, and not for want of a cost table. Signed-off-by: Lance Liao --- docs/dspark_p0_report.md | 65 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/docs/dspark_p0_report.md b/docs/dspark_p0_report.md index b5bb3642b917..9abba21d6e71 100644 --- a/docs/dspark_p0_report.md +++ b/docs/dspark_p0_report.md @@ -50,6 +50,71 @@ goal doc §6.3 称它是「性价比最高的一个测试」。参数化之后 --- +## 0.4 ⛔ 最重要的结论：ragged 在 DSv4 上被 K1 挡住，从未真正跑起来过 + +端到端跑到真实 decode 之后，第一件撞上的事是： + +``` +RuntimeError: [TensorRT-LLM][ERROR] Assertion failed: + MLA can only support input sequences with the same sequence length. + (cpp/tensorrt_llm/kernels/mlaKernels.cu:1161) +``` + +调用栈毫无歧义： + +``` +modeling_deepseekv4.py:2509 forward + -> mla.py:1981 forward_impl_with_deepseek_v4 + -> mla.py:2147 forward_generation_sparse_mla + -> mla.py:2549 forward_absorption_generation + -> self.mqa.mla_rope_generation(...) + -> trtllm.py:2038 torch.ops.trtllm.mla_rope_generation + -> mlaKernels.cu:1161 +``` + +kernel 侧就是 goal doc §3.2 **K1** 逐字描述的那个检查： + +```cpp +TLLM_CHECK_WITH_INFO(params.acc_q_len % params.batch_size == 0, + "MLA can only support input sequences with the same sequence length."); +auto seq_len = params.acc_q_len / params.batch_size; +``` + +### 为什么这条是决定性的 + +commit `e1c28eeae2` 的标题是 **"Make DSpark ragged verification work on the DSv4 kernel path"**。 +但 `git show --stat` 显示它只动了两个 kernel： + +| 文件 | goal doc 编号 | +|---|---| +| `indexerTopK.cu` / `IndexerTopK.h` / `IndexerTopKOp.cpp` | **K4** 🔧 | +| `compressorKernels.cu/.h` / `compressorOp.cpp` | **K9** 🔧 | + +**`mlaKernels.cu`（K1）和 `attentionOp.cpp`（K2）从未被触碰** —— 而 goal doc §3.2 把这两个都标为 ❌ +kernel 层阻断点，且 K1 排在前面。所以 ragged 路径在 DSv4 上**结构上就不可能端到端跑通**， +这个 commit 的标题超出了它实际交付的范围。 + +### 危险的那一半 + +goal doc K1 还指出了比 abort 更坏的模式：窗口 `[6,2]` → 8 tokens / batch 2 → `seq_len=4`， +**检查通过**，但每个请求的 K/V 写进错误的行 —— 持久性 KV 缓存污染，静默。 +本轮撞到的是响亮的那一半（整除性不成立）；另一半在别的窗口组合下可达。 + +### 这对 D1/D2 意味着什么 + +goal doc §5 说得很清楚：D1 和 D2 **都**靠把 generation 半边呈现成 +「`num_gen_tokens` 个长度为 1 的序列」来让 K1/K2 的整除检查平凡通过。 +**那个呈现变更从未实现。** + +而 §1 的 B2 实测说这个呈现**不慢**（0.93–1.00×）。所以路线是通的、代价是可接受的， +只是**活还没干**。这正是 P0 之后该做的第一件事，而不是调参或找配置。 + +### 对「吞吐有没有提升」这个问题的直接回答 + +**今天无法测量。** 不是因为缺 SPS 表，而是因为 ragged 在 DSv4 的 MLA generation 路径上 +根本跑不起来。任何声称测到了 ragged 吞吐的数字，要么没真正走 ragged（planner 静默退化）， +要么就是撞了上面那个静默的 KV 污染分支。 + ## 1. 交付物 0：B2 实测数字（最先交） **问题**（goal doc §7.1 U1）：trtllm-gen sparse-MLA generation 在 `(batch=B, s_q=tier+1)` 与 From 1f11dafe903952a6d96ad5b7770df3cb5e390385 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 18:51:34 -0700 Subject: [PATCH 031/177] [None][doc] Theta is real and large; correcting an earlier prediction With L=1 and L=3 measured, theta separates from alpha for the first time: at fixed bs, T(bs,L=3) - T(bs,L=1) cancels bias and alpha(bs) exactly, leaving only theta(4bs) - theta(2bs) -- the sole term trimming can reduce. It is +15% to +37% of the step, rising with batch size, and +37% at bs=128. This contradicts what I predicted from B2. B2 showed the sparse-MLA layer flat in token count (0.137-0.150 ms across 8->128 tokens) and I extrapolated that to the whole step. Wrong: the layer genuinely is flat, but a decode step is dominated by MoE, whose cost does scale with tokens since each token routes to experts. Theta is therefore both real and large, and it grows with batch size -- which is exactly the "gains only at high concurrency" claim the feature makes about itself. Recorded as an upper bound. The profiler sweeps L by shrinking block_size, so the draft pass shrinks with it and its cost is folded into the measured delta. A small value would have been a sound one-sided rejection; a large one needs the leak-free measurement, which is what TLLM_DSPARK_FORCE_VERIFY_LENS= was added for. Signed-off-by: Lance Liao --- docs/dspark_p0_report.md | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/docs/dspark_p0_report.md b/docs/dspark_p0_report.md index 9abba21d6e71..d8f2bee31f6e 100644 --- a/docs/dspark_p0_report.md +++ b/docs/dspark_p0_report.md @@ -159,6 +159,35 @@ batch-major 走 expanded 一 token 一行的路径，token-major 走 `next_n == --- +## 1.5 θ 是真实且可观的 —— 一个我先前判断错误的更正 + +L=1 与 L=3 两组样本齐了之后，第一次可以把 `theta(M)` 从 `alpha(bs)` 里分离出来： +固定 `bs` 时 `T(bs,L=3) - T(bs,L=1)` 让 `bias` 和 `alpha(bs)` **精确抵消**， +剩下的就是 `theta(4bs) - theta(2bs)` —— 唯一能被裁剪减少的那一项。 + +| bs | tokens | Δθ (ms) | 占该步耗时 | +|---|---|---|---| +| 8 | 16→32 | +3.75 | +15.0% | +| 16 | 32→64 | +5.26 | +18.3% | +| 32 | 64→128 | +9.30 | +27.5% | +| 64 | 128→256 | +7.42 | +17.4% | +| **128** | **256→512** | **+18.33** | **+37.0%** | + +**我先前的判断是错的。** 我根据 B2（sparse-MLA layer 在 8→128 token 上基本持平， +0.137–0.150 ms）推断 θ 没什么可裁的空间。错在把**一个 attention layer** 外推成**整个 decode step**： +该 layer 确实对 token 数不敏感，但整步的开销由 **MoE** 主导，而 MoE 的代价随 token 数增长 +（每个 token 都要路由到专家）。所以 θ 既真实又很大，且**随 batch 增大而增大** —— +这正好符合该 feature docstring 自称的「收益只在高并发出现」。 + +对照可裁剪门槛（profiler 源码自述：相邻 tier 之间需 >15% 才值得裁）， +**bs=128 的 37% 有充裕余量**。 + +**上界性质**：profiler 通过缩小 `block_size` 来扫 L，draft pass 也跟着缩， +所以测得的 Δθ 里混入了 draft 代价，是**上界**。这使得「小值」是可靠的否定结论， +而「大值」需要进一步确认。要测无泄漏的版本，就是保持 `block_size=5` 不动、 +改用 `TLLM_DSPARK_FORCE_VERIFY_LENS=` 强制窗口 —— 该能力本轮已实现 +（见 §2 的 commit）。这是下一步该做的测量。 + ## 2. 交付物 1：修复清单 每项一个 commit，commit message 里引用 goal doc 编号。 From 70eab5dd1216ed1f02c5351a1e7edc9a6e41b127 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 18:52:47 -0700 Subject: [PATCH 032/177] [None][test] Let the uniform-tier GSM8K case take an SPS cost table too Same reasoning as the ragged case: without a table the planner keeps the full block every step, so the tier ladder is captured but never selected from and the run validates correctness without exercising the scheduling decision. Signed-off-by: Lance Liao --- .../defs/accuracy/test_llm_api_pytorch.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index d47e23ce7a6b..a495d4cfe893 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4090,9 +4090,17 @@ def test_gsm8k_dep8_confidence_scheduling_overlap(self, moe_backend): """ kv_cache_config = KvCacheConfig(enable_block_reuse=False, free_gpu_memory_fraction=0.5) - spec_config = DSparkDecodingConfig(max_draft_len=5, - speculative_model=self.MODEL_PATH, - enable_confidence_scheduling=True) + # Same reasoning as the ragged case: without a cost table the planner + # cannot tell a cheap verified token from an expensive one and keeps the + # full block every step, so the tier ladder is captured but never + # selected from. The run is still a valid correctness check -- it just + # does not exercise the scheduling decision. + sps_table_path = os.environ.get("TLLM_DSPARK_SPS_TABLE") or None + spec_config = DSparkDecodingConfig( + max_draft_len=5, + speculative_model=self.MODEL_PATH, + enable_confidence_scheduling=True, + confidence_sps_table_path=sps_table_path) with LLM(self.MODEL_PATH, attn_backend="TRTLLM", tensor_parallel_size=8, From d4247bbb7dc7b44598fa8d1c670ab6ff76648db8 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 19:37:41 -0700 Subject: [PATCH 033/177] [None][doc] GSM8K = 96.2092, and the counters say it is a baseline Uniform tier ladder, overlap scheduler on, DEP8 on 8x B300: 96.2092 +/- 0.526 against a recorded reference of 96.475. The 0.27 gap sits inside one standard error, which is exactly the false-failure the goal doc predicted for a gate with 6.3 questions of margin against a 0.54-point binomial stderr. pytest still reports a failure, but the accuracy assertion never ran: it died in tensorrt_llm/evaluate/lm_eval.py:657 taking numpy.mean over string dtypes, an lm_eval-0.4.12 / numpy-2.1.0 incompatibility in this hand-assembled environment. lm_eval had already scored all 1319 questions. The counters make the important point: fallback_flat_cost is 112 of 114 decisions, so the planner declined to trim on every step -- that build predates wiring the SPS table into the uniform case. The run is therefore baseline DSpark carrying 3x the captured graphs, which makes 96.2092 a control rather than a result. That a healthy-looking run doing nothing is legible at a glance is the entire argument for the observability work. Also banks two numbers worth having: accept_len 3.3009 (the block drafter accepts 3.30 of 5 drafted positions on GSM8K, the tau baseline) and 98.3% CUDA graph replay. Records the environment traps that cost three 25-minute runs, so the next person does not rediscover them. Signed-off-by: Lance Liao --- docs/dspark_p0_report.md | 79 ++++++++++++++++++++++++++++++---------- 1 file changed, 60 insertions(+), 19 deletions(-) diff --git a/docs/dspark_p0_report.md b/docs/dspark_p0_report.md index d8f2bee31f6e..d26942f3a02e 100644 --- a/docs/dspark_p0_report.md +++ b/docs/dspark_p0_report.md @@ -432,31 +432,72 @@ SGLang 博客（https://www.lmsys.org/blog/2026-07-06-dspark-sglang）的核心 记录点在 sampler，因为那是唯一一处「接受数」与「它对应的窗口」同时在 host 侧且属于同一步 的地方 —— 活的 `py_verify_len` 那时已翻到下一步（正是 H6 快照存在的原因）。 -## 5. 交付物 3：e2e 实测 —— **未完成** +## 5. 交付物 3：e2e 实测 -- 均匀 tier 路径 + overlap scheduler + DEP8 的 GSM8K：**运行中，未取得分数**。 -- ragged 路径：**未跑**。 -- planner.stats 的 fallback 计数：**未取得**。 +### 5.1 GSM8K = **96.2092** (± 0.526) -已排除的环境障碍（供下一次复现参考）： +配置：uniform tier ladder（`enable_confidence_scheduling=True`）、overlap scheduler 开启、 +DEP8、MEGAMOE_DEEPGEMM、B300 × 8。 -1. 容器里没有 `pytest` / `parameterized` / `mako` / `oyaml` 等；已装到 - `/lustre/fsw/coreai_comparch_trtllm/laliao/pyextra`，通过 `PYTHONPATH` 引入。 -2. `srun --ntasks=8 pytest` 会让 LLM API 再去 `MPI_Comm_spawn`，报 - `MPI_ERR_SPAWN`。**必须用 `trtllm-llmapi-launch` 包一层**（仓库里 - `tmp/serve_*.slurm` 就是这么写的）。 -3. 便捷封装：`tmp/run_in_alloc.sh ""`。 - 注意它用 `$*` 展开，测试 id 里的 `[...]` 会被 shell 当 glob 吃掉 —— 用 `-k` 选择。 +| | | +|---|---| +| 实测 | **96.2092**（flexible-extract 与 strict-match 一致） | +| 记录参考 | 96.475 | +| 差 | 0.27，**在 1 个标准误（0.526）之内** —— 统计上不可区分 | -已确认到达的阶段：权重加载完成 → DSpark worker 初始化 → -`DSpark verify planner: tiers=[1, 3, 5], profiled_cost_table=False, ragged_verify_mode=static`。 +这印证了 goal doc §6.1 对该门的批评：`>= ref_accuracy` 对实测值只有 6.3 题余量， +而 GSM8K 在 p=0.96 处的二项标准误约 0.54 分 —— **既太松又太紧，干净的 run 会假失败**。 -⚠️ **`profiled_cost_table=False`**：没有 profiled SPS 表时 `_decide_local` 无条件返回 -`max_tier`（`dspark_verify.py:268-270`）。所以即便这个 run 跑完，它证明的是 -**102 个 graph 的 capture 不再 hang**（那正是 07/30 失败的地方）， -**不是**「调度真的在 trim」。后者需要 §7 的 Q5。 +> pytest 报 fail，但**精度断言从未执行到**：失败在 +> `tensorrt_llm/evaluate/lm_eval.py:657` → `numpy.mean` 对字符串 dtype 求平均， +> 是本轮拼装环境里 lm_eval 0.4.12 与容器 numpy 2.1.0 的版本不兼容，与模型和本轮改动无关。 +> lm_eval 已完整跑完 1319 题并输出了分数。 ---- +### 5.2 这个分数是**基线**，不是调度生效后的分数 + +随附计数器把这一点讲得毫不含糊： + +``` +mode: static, steps_total: 114, steps_no_windows: 114 +planner: {decisions: 114, fallback_flat_cost: 112} +requests_trimmed: 0, trim_ratio: 0.0, trim_regret_rate: 0.0 +accept_len: 3.3009 (requests_scored: 3436) +graph_replays: 520, graph_eager: 9 +``` + +`fallback_flat_cost: 112/114` —— 该 build 尚未接上 SPS 表（开关是启动之后才补的）， +所以 planner 每步都拒绝裁剪。**这次 run 等价于 baseline DSpark，只是多捕获了 3× graph。** + +**这正是可观测性存在的理由**：一个「分数正确、日志正常、但 feature 什么都没干」的 run， +靠计数器一眼看穿，不需要推测。 + +其他可用数字： + +| 量 | 值 | 含义 | +|---|---|---| +| `accept_len` | **3.3009** | DSpark block drafter 在 GSM8K 上的真实接受长度（5 个 draft 位置接受 3.30，66%）—— τ 项的基准 | +| `graph_replays / (replays+eager)` | **520/529 = 98.3%** | G4：稳态 decode 基本全部命中 graph replay | +| `trim_regret_rate` | 0.0 | 与 `requests_trimmed: 0` 一致 | + +### 5.3 ragged 路径：**跑不起来**（见 §0.4） + +死于 K1（`mlaKernels.cu:1161`），31:51。 + +### 5.4 复现环境的坑（供下一个人参考） + +1. 容器无 `pytest` / `lm_eval` / `mako` / `oyaml` / `parameterized` / `sacrebleu`。 +2. **不要用 `pip install --target` 装带依赖的包**：`PYTHONPATH` 永远排在 site-packages 前面， + 一次 `pip install lm_eval` 就把 `numpy` / `transformers` / `datasets` / `pandas` / `pyarrow` + 等 **50 个包**装进去并 shadow 掉容器版本，症状是 + `No module named 'torch._C'` 或 `No module named 'transformers.utils.chat_template_utils'`。 + 做法：装完立刻把与容器重名的目录全部隔离。 +3. `srun --ntasks=8 pytest` 会让 LLM API 再去 `MPI_Comm_spawn`，报 `MPI_ERR_SPAWN`； + **必须套 `trtllm-llmapi-launch`**。 +4. 被 kill 的 srun step **不会回收 MPI worker**，它们继续占每卡 ~180 GiB， + 下一次 run 会在 executor 创建时 OOM 并「怪自己」。启动前必须清理（`tmp/cleanup_node.sh`）。 +5. SPS profiler 需要 `TORCHINDUCTOR_COMPILE_THREADS=1`，否则 Inductor 每遇新 shape + 就 fork 一个编译 worker（父进程常驻上百 GB × 8 rank）→ 宿主 OOM（exit 137）。 +6. **在花 25 分钟建 engine 之前，先廉价验证整条 import 链**。本轮为此损失了三次 25 分钟。 ## 6. 交付物 5：对 goal doc §8 剩余待拍板问题的建议 From 9dd87cf4f31a05696d85826db935659b263c3f82 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 20:26:50 -0700 Subject: [PATCH 034/177] [None][feat] Present DSv4 generation token-major so ragged verification runs (K1/K2) Ragged verification aborted on the first real decode step at mlaKernels.cu:1161, "MLA can only support input sequences with the same sequence length". Both blocking ops learn their sequence count from a tensor shape -- `num_seqs = host_context_lengths.size(0)` in dsv3RopeOp.cpp:151 and thop/attentionOp.cpp:1325 -- and then require the query tokens to divide evenly across it. Every prior speculation mode, MTP included, satisfies that because each request contributes exactly `1 + draft_len` tokens, so the single scalar `seq_len = acc_q_len / batch_size` describes every request and the kernel's `batch_idx = tok / seq_len`, RoPE position and KV write offset all follow. A ragged batch has no such scalar. Shows those two ops one row per generation query token instead. Then seq_len is 1, `batch_idx = tok`, `local_token_idx = 0`, and the RoPE position and KV write offset both come out as `row_kv_len - 1`; filling row_kv_len with `(kv_len - window + 1) + offset` reproduces the batch-major `kv_len - window + offset` exactly. The asserts are satisfied, not weakened, and no kernel changes. Scoped to those two calls via a context manager rather than redefining `num_seqs`. About twenty consumers read `num_seqs` and `num_generations` as *request* counts -- spec workers sizing acceptance rectangles, the KV-length correction in _preprocess_inputs, the indexer's metadata -- so redefining it would be a global re-typing rather than a presentation change. Rows are laid out context-prefix-first, which keeps num_contexts and the ops' seq_offset correct untouched. Built only when the batch is actually ragged: at s_q == 1 trtllm-gen selects different kernels (mIsCausalSpecDecodingGen is gated on mMaxSeqLenQ > 1), so presenting a uniform batch this way would be equivalent but not bit-identical. Costs kept off the step: the per-row extents reuse what _prepare_ragged_row_kv_lens computed one call earlier rather than walking the windows again, and the block table is expanded by a device-side gather from the copy already resident, avoiding a ~96 KiB host->device copy per decode step. The extent is refreshed from on_update_kv_lens for the same reason its ragged twin is: the overlap scheduler corrects kv_lens on device inside the capture body. Also fixes a latent out-of-bounds write this exposed: cu_q_seqlens/cu_kv_seqlens were sized num_seqs+1, but the rope kernel writes seqQOffset[batch_idx+1] for batch_idx < batch_size and CUB-scans total_s_len/seq_len entries. Size them by the token count, which is an upper bound on num_seqs everywhere. Signed-off-by: Lance Liao --- .../_torch/attention_backend/sparse/dsa.py | 246 ++++++++++++++++++ .../_torch/attention_backend/trtllm.py | 82 +++++- tensorrt_llm/_torch/modules/mla.py | 13 +- .../_torch/pyexecutor/model_engine.py | 19 ++ 4 files changed, 351 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 89652b8eb83c..8255d65ad2b4 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -582,6 +582,22 @@ class IndexerPrefillChunkMetadata: k_token_end: int # K token end index in batch +@dataclass(frozen=True) +class TokenMajorGenView: + """One attention row per generation query token. + + Handed only to the MLA rope and sparse-MLA generation ops, which derive + their sequence count from ``host_context_lengths.size(0)``. Everything else + keeps seeing the batch-major metadata, where ``num_seqs`` still means + "number of requests". + """ + num_rows: int + sequence_length: torch.Tensor + host_past_key_value_lengths: torch.Tensor + host_context_lengths: torch.Tensor + kv_cache_block_offsets: torch.Tensor + + @dataclass(init=False) class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): """Attention metadata for DSA (Dense Sparse Attention) with indexer state.""" @@ -931,6 +947,7 @@ def on_update_kv_lens(self): # kv_lens_cuda may have just been corrected; the ragged per-row extent # is derived from it and has to follow. self.refresh_ragged_row_kv_lens() + self.refresh_token_major_gen_rows() if self.kv_cache_manager is not None and self.num_tokens > 0: seq_lens = self.seq_lens_cuda[:self.num_seqs] @@ -1420,6 +1437,94 @@ def create_expanded_buffers(self, capture_graph=False): ) # Rows covered by the current step's ragged layout; 0 on uniform steps. self._ragged_num_rows = 0 + + # ------------------------------------------------------------------ + # Token-major presentation of the generation half, for the MLA rope and + # sparse-MLA generation ops (goal doc K1 / K2). + # + # Both of those ops learn their sequence count from a tensor *shape* -- + # `num_seqs = host_context_lengths.size(0)` in dsv3RopeOp.cpp and + # thop/attentionOp.cpp -- and then require the query tokens to divide + # evenly across it. A ragged batch does not divide, so the assert + # "MLA can only support input sequences with the same sequence length" + # fires. Handing those two ops one row per generation *token* makes the + # division trivial (seq_len == 1) and satisfies the assert rather than + # weakening it; no kernel changes. + # + # These are a PARALLEL set of views, never the metadata's own fields. + # `num_seqs` is read as a request count in roughly twenty places -- + # spec workers sizing acceptance rectangles, the KV-length correction in + # _preprocess_inputs, the indexer's own metadata -- so redefining it + # would be a global re-typing rather than a presentation change. + # + # Rows are laid out prefix-preserving: the context rows first, unchanged, + # then one row per generation token. That keeps `num_contexts`, + # `num_ctx_tokens` and the ops' `seq_offset = num_contexts` all correct + # without touching them. + # + # Capacity is the static ceiling and must never be recomputed from a + # runtime count: these live in the process-wide buffer pool keyed by + # `cache_name`, so a later reallocation would move addresses already + # baked into captured graphs. `2 + max_draft_tokens` per sequence covers + # one context row plus the widest generation window. + attn_rows_cap = self.max_num_sequences * (2 + self.max_draft_tokens) + self._attn_num_rows = 0 + self.attn_row_kv_lens_cuda = self.get_empty( + self.cuda_graph_buffers, + (attn_rows_cap, ), + cache_name="attn_row_kv_lens_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + # The per-row offset from its request's kv_len, kept separately so the + # extent can be rebuilt on device from a corrected kv_lens without going + # back to the host -- `_preprocess_inputs` rewrites kv_lens_cuda inside + # the capture body under the overlap scheduler. + self.attn_row_kv_correction_cuda = self.get_empty( + self.cuda_graph_buffers, + (attn_rows_cap, ), + cache_name="attn_row_kv_correction_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.attn_row_req_idx_cuda = self.get_empty( + self.cuda_graph_buffers, + (attn_rows_cap, ), + cache_name="attn_row_req_idx_cuda", + dtype=torch.long, + capture_graph=capture_graph, + ) + self.attn_row_kv_lens_host = torch.zeros( + (attn_rows_cap, ), dtype=torch.int32, + pin_memory=prefer_pinned()) + self.attn_row_kv_correction_host = torch.zeros( + (attn_rows_cap, ), dtype=torch.int32, + pin_memory=prefer_pinned()) + self.attn_row_req_idx_host = torch.zeros( + (attn_rows_cap, ), dtype=torch.long, pin_memory=prefer_pinned()) + # host_context_lengths: the shape oracle both ops key off. + self.attn_row_prompt_lens_cpu = torch.zeros( + (attn_rows_cap, ), dtype=torch.int32, + pin_memory=prefer_pinned()) + # The main KV block table, one row per query token. Distinct from + # `block_table_expanded` below, which carries *indexer* pool indices at + # the indexer's own compression. On the DSv4 sparse generation path the + # only consumer that walks this per row is the KV write inside the rope + # kernel; the attention op overrides its page index pointer with the + # sparse top-k indices, which are already per token. + self.attn_row_block_offsets = self.get_empty( + self.cuda_graph_buffers, + [attn_rows_cap, self.kv_cache_manager.max_blocks_per_seq], + cache_name="attn_row_block_offsets", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.attn_row_block_offsets_host = torch.zeros_like( + self.attn_row_block_offsets, + device='cpu', + pin_memory=prefer_pinned(), + ) + self.block_table_expanded = self.get_empty( self.cuda_graph_buffers, [ @@ -1708,6 +1813,7 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): if self.is_ragged_verify: self._prepare_ragged_row_kv_lens(kv_lens) + self._prepare_token_major_gen_rows(kv_lens) self.expand_for_dsl = (use_dsl and self.kv_cache_manager is not None and self.max_draft_tokens >= 1) @@ -1857,6 +1963,146 @@ def refresh_ragged_row_kv_lens(self) -> None: out=row_kv_lens) row_kv_lens.add_(self.row_kv_correction_cuda[:num_tokens]) + def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: + """Build the one-row-per-generation-token view for the MLA gen ops. + + See the buffer comments in :meth:`create_expanded_buffers` for why this + exists and why it is a parallel view rather than a redefinition of + ``num_seqs``. + + Rows are ``[contexts unchanged] ++ [one per generation token]``. Row + ``o`` of a request verifying ``v`` positions attends to + ``kv_len - v + o + 1`` tokens, the same generalization the decode top-k + uses; substituting ``v = next_n`` recovers the uniform formula, which is + why the uniform path can keep its single row per request. + + Assembled entirely on the host -- ``kv_lens`` is already a CPU tensor + here -- and crossed once through pinned memory. Reading it on the device + would be a D2H sync on a path the overlap scheduler depends on. + """ + verify_lens = self.ragged_verify_lens + ngt = self._ragged_num_rows + if not verify_lens or ngt <= 0: + self._attn_num_rows = 0 + return + nc = self.num_contexts + rows = nc + ngt + cap = self.attn_row_kv_lens_cuda.shape[0] + assert rows <= cap, ( + f"token-major generation needs {rows} rows but the buffers hold " + f"{cap}; capacity is the static ceiling and must not be resized " + f"after any CUDA graph was captured") + + # Context rows pass through untouched, so the ops' seq_offset still + # points at the first generation row. + if nc: + ctx_kv = kv_lens[:nc].to(torch.int32) + self.attn_row_kv_lens_host[:nc].copy_(ctx_kv) + self.attn_row_kv_correction_host[:nc].zero_() + self.attn_row_req_idx_host[:nc].copy_( + torch.arange(nc, dtype=torch.long)) + self.attn_row_prompt_lens_cpu[:nc].copy_( + self.prompt_lens_cpu[:nc].to(torch.int32)) + + # Reuse what _prepare_ragged_row_kv_lens just built rather than walking + # the windows again: it produced exactly these per-row extents, + # corrections and request indices one call earlier. Rebuilding them + # would double the only O(num_gen_tokens) host work on the step. + self.attn_row_kv_lens_host[nc:rows].copy_( + self.row_kv_lens_host[:ngt]) + self.attn_row_kv_correction_host[nc:rows].copy_( + self.row_kv_correction_host[:ngt]) + # row_req_idx is batch-relative among generation requests; the attention + # rows are indexed over the whole batch, so shift past the contexts. + torch.add(self.row_req_idx_host[:ngt], nc, + out=self.attn_row_req_idx_host[nc:rows]) + # host_context_lengths is only a shape oracle for the generation half + # (the ops read it for num_seqs and, for contexts, for the q length); + # one token per row. + self.attn_row_prompt_lens_cpu[nc:rows].fill_(1) + + self.attn_row_kv_lens_cuda[:rows].copy_( + self.attn_row_kv_lens_host[:rows], non_blocking=True) + self.attn_row_kv_correction_cuda[:rows].copy_( + self.attn_row_kv_correction_host[:rows], non_blocking=True) + self.attn_row_req_idx_cuda[:rows].copy_( + self.attn_row_req_idx_host[:rows], non_blocking=True) + + self._attn_num_rows = rows + # The block table is expanded on the DEVICE, from the copy that is + # already there. Building it host-side and copying would add ~96 KiB of + # H2D per decode step (768 rows x max_blocks_per_seq int32) plus a host + # gather, both on the critical path the overlap scheduler exists to keep + # clear. One index_select costs neither. + self.refresh_token_major_block_table() + + def refresh_token_major_gen_rows(self) -> None: + """Rebuild the per-row KV extent from the corrected ``kv_lens_cuda``. + + Mirrors :meth:`refresh_ragged_row_kv_lens` and for the same reason: the + overlap scheduler corrects ``kv_lens_cuda`` on device after prepare(), + inside the CUDA-graph capture body, so a prepare-time extent would make + the rope kernel and the attention op disagree about the same quantity. + + One gather plus one add. No allocation, no repeat_interleave, no D2H -- + the row->request map and the correction were staged at prepare() time + and their addresses are pinned by the graph key. + """ + rows = self._attn_num_rows + if not self.is_ragged_verify or rows <= 0: + return + torch.index_select(self.kv_lens_cuda[:self.num_seqs], 0, + self.attn_row_req_idx_cuda[:rows], + out=self.attn_row_kv_lens_cuda[:rows]) + self.attn_row_kv_lens_cuda[:rows].add_( + self.attn_row_kv_correction_cuda[:rows]) + + def refresh_token_major_block_table(self) -> None: + """Expand the main KV block table to one row per query token, on device. + + Every row of a request sees that request's blocks -- a verify window + never spans a request -- so this is a plain gather by the row->request + map. Done on the device because the source is already resident; doing it + host-side would mean a per-step H2D of the whole table. + + Not refreshed from ``on_update_kv_lens``: block IDs do not move within a + step, the same reason ``block_table_expanded`` is built once per + prepare(). + """ + rows = self._attn_num_rows + if rows <= 0 or self.kv_cache_block_offsets is None: + return + width = min(self.attn_row_block_offsets.shape[1], + self.kv_cache_block_offsets.shape[-1]) + # kv_cache_block_offsets is [num_pools, num_seqs, 2, max_blocks] for the + # paged layout; the generation path reads pool 0, the K direction. + src = self.kv_cache_block_offsets + while src.dim() > 2: + src = src[0] if src.shape[0] == 1 else src.select(0, 0) + torch.index_select(src[:, :width], 0, + self.attn_row_req_idx_cuda[:rows], + out=self.attn_row_block_offsets[:rows, :width]) + + def token_major_gen_view(self) -> Optional["TokenMajorGenView"]: + """The token-major presentation for this step, or None. + + None on every uniform batch, which is what keeps those paths + bit-identical: at ``s_q == 1`` trtllm-gen selects different kernels + (`mIsCausalSpecDecodingGen` is gated on ``mMaxSeqLenQ > 1``), so + presenting a uniform batch token-major would be functionally equivalent + but not bitwise so. + """ + rows = self._attn_num_rows + if not self.is_ragged_verify or rows <= 0: + return None + return TokenMajorGenView( + num_rows=rows, + sequence_length=self.attn_row_kv_lens_cuda[:rows], + host_past_key_value_lengths=self.attn_row_kv_lens_host[:rows], + host_context_lengths=self.attn_row_prompt_lens_cpu[:rows], + kv_cache_block_offsets=self.attn_row_block_offsets[:rows], + ) + def ragged_row_kv_lens(self, num_tokens: int) -> Optional[torch.Tensor]: """The per-row extent slice for this step, or None on a uniform batch.""" if not self.is_ragged_verify: diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 992eec88b7fd..05f9c61a9a61 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -14,6 +14,7 @@ # limitations under the License. import functools +import contextlib import math import os import weakref @@ -517,6 +518,45 @@ def _bind_runtime_views( self.prompt_lens_cpu_runtime = prompt_lens_cpu self.host_request_types_runtime = host_request_types + def token_major_gen_view(self): + """The token-major generation presentation, or None. + + None on every backend but DSA, and on every uniform DSA batch. See + ``DSAtrtllmAttentionMetadata.token_major_gen_view``. + """ + return None + + @contextlib.contextmanager + def presented_token_major(self, view): + """Temporarily present the generation half as one row per query token. + + The MLA rope and sparse-MLA generation ops learn their sequence count + from ``host_context_lengths.size(0)`` and then require the query tokens + to divide evenly across it. A ragged batch does not divide. Swapping in + per-token views for the duration of those two calls makes the division + trivial (seq_len == 1) and satisfies their asserts rather than weakening + them. + + Scoped deliberately: ``num_seqs`` is read as a *request* count almost + everywhere else -- spec workers sizing acceptance rectangles, the + KV-length correction under the overlap scheduler, the indexer's own + metadata -- so the token-major rows must never outlive these two calls. + """ + if view is None: + yield + return + saved = (self.kv_lens_cuda_runtime, self.kv_lens_runtime, + self.prompt_lens_cpu_runtime, self.kv_cache_block_offsets) + try: + self.kv_lens_cuda_runtime = view.sequence_length + self.kv_lens_runtime = view.host_past_key_value_lengths + self.prompt_lens_cpu_runtime = view.host_context_lengths + self.kv_cache_block_offsets = view.kv_cache_block_offsets + yield + finally: + (self.kv_lens_cuda_runtime, self.kv_lens_runtime, + self.prompt_lens_cpu_runtime, self.kv_cache_block_offsets) = saved + def prepare(self) -> None: super().prepare() extra_attrs = get_model_extra_attrs() @@ -1770,13 +1810,21 @@ def forward( if not self.fmha_libs: self.create_fmha_libs() - for fmha in self.fmha_libs: - if fmha.is_supported(q, k, v, metadata, forward_args): - fmha.forward(q, k, v, metadata, forward_args) - break - else: - raise RuntimeError( - "No TRT-LLM attention FMHA library supports this request.") + # goal doc K2. On a ragged generation step the attention op is handed + # one row per query token so its `num_tokens % num_seqs == 0` check + # passes with seq_len == 1. Wrapping the dispatch rather than the + # library keeps every fmha implementation, and its tests, untouched. + token_major_view = (metadata.token_major_gen_view() if + forward_args.attention_input_type + == AttentionInputType.generation_only else None) + with metadata.presented_token_major(token_major_view): + for fmha in self.fmha_libs: + if fmha.is_supported(q, k, v, metadata, forward_args): + fmha.forward(q, k, v, metadata, forward_args) + break + else: + raise RuntimeError( + "No TRT-LLM attention FMHA library supports this request.") if self.print_skip_softmax_stat: total_blocks, skipped_blocks = self.skip_softmax_stat @@ -2035,6 +2083,26 @@ def mla_rope_generation( metadata.helix_position_offsets, metadata.helix_is_inactive_rank ] + # goal doc K1. Same presentation as the attention dispatch above: the + # rope kernel derives batch_idx = tok / seq_len and the KV write offset + # from a single scalar seq_len, which only describes a batch where every + # request contributes the same number of tokens. + with self.presented_token_major_for(metadata): + self._mla_rope_generation_impl( + metadata, fused_q, q_pe, latent_cache, cu_q_seqlens, + cu_kv_seqlens, fmha_scheduler_counter, mla_bmm1_scale, + mla_bmm2_scale, quant_q_buffer, out_scale, helix_tensor_params) + + @contextlib.contextmanager + def presented_token_major_for(self, metadata): + with metadata.presented_token_major(metadata.token_major_gen_view()): + yield + + def _mla_rope_generation_impl(self, metadata, fused_q, q_pe, latent_cache, + cu_q_seqlens, cu_kv_seqlens, + fmha_scheduler_counter, mla_bmm1_scale, + mla_bmm2_scale, quant_q_buffer, out_scale, + helix_tensor_params): torch.ops.trtllm.mla_rope_generation( fused_q, q_pe, diff --git a/tensorrt_llm/_torch/modules/mla.py b/tensorrt_llm/_torch/modules/mla.py index 63f9a1bd65d3..639cd98ad81d 100644 --- a/tensorrt_llm/_torch/modules/mla.py +++ b/tensorrt_llm/_torch/modules/mla.py @@ -2523,8 +2523,17 @@ def forward_absorption_generation( # 2) rope(q_pe) with shape [num_tokens, num_heads, qk_rope_head_dim]. rope is applied inside AttentionOp num_seqs = attn_metadata.num_seqs - cu_q_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=q.device) - cu_kv_seqlens = torch.empty(num_seqs + 1, dtype=torch.int32, device=q.device) + # The rope kernel writes seqQOffset[batch_idx + 1] for every + # batch_idx < batch_size, and CUB-scans total_s_len / seq_len entries + # into seqKVOffsets (mlaKernels.cu:579, :650). `batch_size` there is the + # count of presented rows, which equals num_seqs only while every + # request contributes the same number of query tokens. Under the + # token-major presentation a row is a token, so size these by the token + # count -- which is also an upper bound on num_seqs on every other path, + # so this is safe rather than merely sufficient. + cu_seqlens_rows = max(num_seqs, num_tokens) + 1 + cu_q_seqlens = torch.empty(cu_seqlens_rows, dtype=torch.int32, device=q.device) + cu_kv_seqlens = torch.empty(cu_seqlens_rows, dtype=torch.int32, device=q.device) fmha_scheduler_counter = torch.empty(1, dtype=torch.uint32, device=q.device) has_fp8_kv_cache = ( self.mqa.has_fp8_kv_cache if hasattr(self.mqa, "has_fp8_kv_cache") else False diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index ddf9b36cab51..b56f53011efb 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3920,6 +3920,25 @@ def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, f"{bucket}; the replayed graph was captured at a different " f"width and nothing downstream would notice.") + # Row axis. Everything above is a token count, so all four agree even + # when the token-major presentation is malformed -- the presentation is + # about how many ROWS the MLA generation ops are shown, which is a + # different quantity and needs its own check. + view = attn_metadata.token_major_gen_view() if hasattr( + attn_metadata, "token_major_gen_view") else None + if view is not None: + expected_rows = attn_metadata.num_contexts + from_requests + assert view.num_rows == expected_rows, ( + f"token-major presentation has {view.num_rows} rows but the " + f"batch has {attn_metadata.num_contexts} context rows plus " + f"{from_requests} generation tokens; the MLA generation ops " + f"would read a sequence count that does not describe this step.") + assert view.host_context_lengths.shape[0] == expected_rows, ( + f"token-major host_context_lengths has " + f"{view.host_context_lengths.shape[0]} entries but the " + f"presentation is {expected_rows} rows -- that tensor's shape " + f"IS what the ops read as their sequence count.") + def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, generation_requests) -> None: """Publish this step's per-request verify windows to both metadatas. From fdedeaacf4e865a4e72244ad3c43e1a942317180 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 20:32:19 -0700 Subject: [PATCH 035/177] [None][fix] Size the attention op's semaphore array for token-major rows The token-major presentation makes the MLA generation kernels see the ROW count as their batch dimension, which is larger than the request count (up to max_num_sequences * (2 + max_draft_tokens)). The op reserves its multi-CTA-KV counter as num_heads * max_num_requests and sizes the generation workspace from the same number, so leaving max_num_requests at the request count under-reserves it. Swap in the static row ceiling for the duration of the presentation. The ceiling rather than the step's row count, because the op caches by (beam_width, max_num_requests, attention_window_size) and reserves only upward -- a per-step value would thrash that cache and re-reserve on every shape change, while the ceiling keeps exactly one entry. No C++ change is needed for this: the op's own mMaxNumRequests is a JIT-warmup hint, the per-step request count comes from the batch rather than from this field, and getFmhaMultiCtasKvScratchSize is SM-count based rather than batch-based, so the generation workspace barely grows. Signed-off-by: Lance Liao --- .../_torch/attention_backend/sparse/dsa.py | 6 ++++++ tensorrt_llm/_torch/attention_backend/trtllm.py | 16 ++++++++++++++-- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 8255d65ad2b4..ce2e7de874b4 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -596,6 +596,11 @@ class TokenMajorGenView: host_past_key_value_lengths: torch.Tensor host_context_lengths: torch.Tensor kv_cache_block_offsets: torch.Tensor + #: Static row ceiling, NOT this step's row count. The attention op caches + #: by (beam_width, max_num_requests, attention_window_size) and reserves its + #: semaphore array from it, so a per-step value would thrash that cache and + #: re-reserve on every shape change. + max_num_rows: int @dataclass(init=False) @@ -2101,6 +2106,7 @@ def token_major_gen_view(self) -> Optional["TokenMajorGenView"]: host_past_key_value_lengths=self.attn_row_kv_lens_host[:rows], host_context_lengths=self.attn_row_prompt_lens_cpu[:rows], kv_cache_block_offsets=self.attn_row_block_offsets[:rows], + max_num_rows=self.attn_row_kv_lens_cuda.shape[0], ) def ragged_row_kv_lens(self, num_tokens: int) -> Optional[torch.Tensor]: diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 05f9c61a9a61..d62a9fd71202 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -546,16 +546,28 @@ def presented_token_major(self, view): yield return saved = (self.kv_lens_cuda_runtime, self.kv_lens_runtime, - self.prompt_lens_cpu_runtime, self.kv_cache_block_offsets) + self.prompt_lens_cpu_runtime, self.kv_cache_block_offsets, + self.max_num_requests) try: self.kv_lens_cuda_runtime = view.sequence_length self.kv_lens_runtime = view.host_past_key_value_lengths self.prompt_lens_cpu_runtime = view.host_context_lengths self.kv_cache_block_offsets = view.kv_cache_block_offsets + # The op reserves its multi-CTA-KV counter as + # `num_heads * max_num_requests` and sizes its generation workspace + # from the same number. Under this presentation the batch dimension + # the kernels see is the ROW count, which is larger. Passing the + # static ceiling grows the reservation to cover it while keeping the + # op's (beam_width, max_num_requests, window) cache key constant. + # This is only a capacity input -- the op's own mMaxNumRequests is a + # JIT-warmup hint, and the per-step request count comes from the + # batch, not from here. + self.max_num_requests = view.max_num_rows yield finally: (self.kv_lens_cuda_runtime, self.kv_lens_runtime, - self.prompt_lens_cpu_runtime, self.kv_cache_block_offsets) = saved + self.prompt_lens_cpu_runtime, self.kv_cache_block_offsets, + self.max_num_requests) = saved def prepare(self) -> None: super().prepare() From 2a53cd8748985a262bd9e5da031031ad15f0ce4a Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 20:44:52 -0700 Subject: [PATCH 036/177] [None][chore] Make a zero-sample SPS cell explain itself A cell that reports aligned_steps=0 is indistinguishable from a cell that never ran, and the rejection reasons want opposite fixes: a rank count that is never right means --enable-attention-dp disagrees with the deployment, context requests on every step mean prompts are still being ingested, disagreeing generation counts mean the batch is draining, and a missing timing field means --timing-field is not published. The sweep previously ended in an InsufficientSamplesError listing all of those as possibilities, so telling them apart cost one full sweep per guess. Report the histogram per cell instead, including the case where steps do align but numGenRequests never equals the requested batch size -- which is its own failure and was previously invisible. Signed-off-by: Lance Liao --- .../_torch/speculative/dspark_sps_profiler.py | 80 +++++++++++++++++++ 1 file changed, 80 insertions(+) diff --git a/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py b/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py index a5140507116f..4e97a9287646 100644 --- a/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py +++ b/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py @@ -366,6 +366,73 @@ def aligned_steps_from_stats( return aligned +def explain_alignment( + rows: Sequence[dict], + *, + expected_ranks: int = 1, + timing_key: str = HOST_STEP_TIME_KEY, +) -> str: + """Say why :func:`aligned_steps_from_stats` kept nothing. + + A cell that reports ``aligned_steps=0`` is otherwise indistinguishable from + a cell that never ran, and the four rejection reasons want opposite fixes: + a rank count that is never right means ``--enable-attention-dp`` disagrees + with the deployment, context requests on every step mean the prompts are + still being ingested, disagreeing generation counts mean the batch is + draining, and a missing timing field means the requested ``--timing-field`` + is not being published. Guessing between those costs a full sweep each + time. + """ + by_iteration: Dict[int, List[dict]] = {} + for row in rows: + iteration = row.get("iter") + if iteration is None: + continue + by_iteration.setdefault(int(iteration), []).append(row) + if not by_iteration: + return (f"{len(rows)} stats rows carried no 'iter' field at all" + if rows else "no stats rows were received at all") + + reasons: Dict[str, int] = {} + rank_counts: Dict[int, int] = {} + gen_seen: set = set() + for iteration, group in by_iteration.items(): + ranks = {int(r.get("attentionDpRank", 0)) for r in group} + rank_counts[len(ranks)] = rank_counts.get(len(ranks), 0) + 1 + if len(ranks) != int(expected_ranks): + reasons["rank count != expected"] = reasons.get("rank count != expected", 0) + 1 + continue + batching = [row.get("inflightBatchingStats") or {} for row in group] + if any(int(stats.get("numContextRequests", 0) or 0) for stats in batching): + reasons["context requests in flight"] = reasons.get("context requests in flight", 0) + 1 + continue + gen_counts = {int(stats.get("numGenRequests", 0) or 0) for stats in batching} + gen_seen.update(gen_counts) + if len(gen_counts) != 1: + reasons["ranks disagree on numGenRequests"] = reasons.get("ranks disagree on numGenRequests", 0) + 1 + continue + if gen_counts and max(gen_counts) <= 0: + reasons["no generation requests"] = reasons.get("no generation requests", 0) + 1 + continue + timing = group[0].get(timing_key) + if timing is None or float(timing) <= 0.0: + reasons[f"missing/zero {timing_key}"] = reasons.get(f"missing/zero {timing_key}", 0) + 1 + continue + reasons["kept"] = reasons.get("kept", 0) + 1 + + parts = [f"{len(rows)} rows over {len(by_iteration)} iterations " + f"(iter {min(by_iteration)}..{max(by_iteration)})"] + parts.append("rejections: " + ", ".join( + f"{reason} x{count}" for reason, count in + sorted(reasons.items(), key=lambda kv: -kv[1])) or "none") + parts.append("ranks-per-iteration histogram: " + ", ".join( + f"{n}->{c}" for n, c in sorted(rank_counts.items())) + + f" (expected {expected_ranks})") + if gen_seen: + parts.append(f"numGenRequests observed: {sorted(gen_seen)[:12]}") + return "; ".join(parts) + + def summarize_cells( samples: Sequence[StepSample], *, @@ -1128,6 +1195,19 @@ def _run_cell(llm, config: SweepConfig, *, batch_size: int, verify_len: int) -> aligned = aligned_steps_from_stats( rows, expected_ranks=config.dp_size, timing_key=config.timing_key ) + kept = [s for s in aligned if s[1] == int(batch_size)] + if not kept: + # Printed per cell rather than once at the end: which cells failed and + # why is the whole diagnosis, and a sweep that ends in + # InsufficientSamplesError has already thrown that away. + print( + f"[dspark-sps] bs={batch_size:>5} L={verify_len} kept nothing -- " + + explain_alignment(rows, expected_ranks=config.dp_size, + timing_key=config.timing_key) + + (f"; aligned but numGenRequests never == {batch_size} " + f"(saw {sorted({a[1] for a in aligned})[:12]})" if aligned else ""), + file=sys.stderr, + ) return [ StepSample( batch_size=int(batch_size), From 9370b065a89dd9581c86362a9c66d83495b668e5 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 21:18:04 -0700 Subject: [PATCH 037/177] [None][fix] Expand the token-major block table on the sequence axis The expanded block table was shaped [rows, max_blocks] and gathered by collapsing leading axes until two were left. The attention op reads [num_pools, num_seqs, 2, max_blocks], and this tensor is substituted for kv_cache_block_offsets wholesale, so the rank and axis order have to match. Worse, the collapse stripped num_seqs rather than the pool axis: after removing pools the next axis was num_seqs, which is not 1, so it was selected away too and request ids ended up indexing the length-2 K/V axis. On 8 ranks that aborted executor init with a device-side "vectorized gather kernel index out of bounds". Allocate the same 4-D layout and index_select along the sequence axis, carrying the pool and K/V axes through untouched. Add tests for the expansion itself. The existing sparse-MLA tests call the ops with hand-built tensors, so they pass whatever the metadata layer produces and could not see a view assembled wrongly -- this bug survived 18 green kernel tests and failed only end-to-end. The new tests assert the expanded table's rank matches what the op reads, that each row carries its own request's blocks on every pool and both directions, that row->request ids stay inside the sequence axis, and that context rows stay identity-mapped ahead of the generation rows. Signed-off-by: Lance Liao --- .../_torch/attention_backend/sparse/dsa.py | 39 +++-- .../deepseek_v4/test_token_major_rows.py | 152 ++++++++++++++++++ 2 files changed, 175 insertions(+), 16 deletions(-) create mode 100644 tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index ce2e7de874b4..90fc24e1b247 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1517,18 +1517,25 @@ def create_expanded_buffers(self, capture_graph=False): # only consumer that walks this per row is the KV write inside the rope # kernel; the attention op overrides its page index pointer with the # sparse top-k indices, which are already per token. + # Same [num_pools, num_seqs, 2, max_blocks_per_seq] layout the attention + # op reads, expanded on the *sequence* axis only: this tensor is + # substituted for `kv_cache_block_offsets` wholesale while the ops are + # presented token-major, so its rank and axis order have to match what + # they already expect. The pool and K/V axes are carried through + # unchanged -- only the row axis is what token-major changes. + num_attention_op_pools = getattr(self.kv_cache_manager, + "num_attention_op_pools", + self.kv_cache_manager.num_pools) self.attn_row_block_offsets = self.get_empty( self.cuda_graph_buffers, - [attn_rows_cap, self.kv_cache_manager.max_blocks_per_seq], + [ + num_attention_op_pools, attn_rows_cap, 2, + self.kv_cache_manager.max_blocks_per_seq + ], cache_name="attn_row_block_offsets", dtype=torch.int32, capture_graph=capture_graph, ) - self.attn_row_block_offsets_host = torch.zeros_like( - self.attn_row_block_offsets, - device='cpu', - pin_memory=prefer_pinned(), - ) self.block_table_expanded = self.get_empty( self.cuda_graph_buffers, @@ -2077,16 +2084,16 @@ def refresh_token_major_block_table(self) -> None: rows = self._attn_num_rows if rows <= 0 or self.kv_cache_block_offsets is None: return - width = min(self.attn_row_block_offsets.shape[1], - self.kv_cache_block_offsets.shape[-1]) - # kv_cache_block_offsets is [num_pools, num_seqs, 2, max_blocks] for the - # paged layout; the generation path reads pool 0, the K direction. src = self.kv_cache_block_offsets - while src.dim() > 2: - src = src[0] if src.shape[0] == 1 else src.select(0, 0) - torch.index_select(src[:, :width], 0, - self.attn_row_req_idx_cuda[:rows], - out=self.attn_row_block_offsets[:rows, :width]) + # Gather along the sequence axis (1) and leave the pool and K/V axes + # alone. Collapsing leading axes instead would strip num_seqs itself and + # index a length-2 K/V axis with request ids, which the gather kernel + # catches as an out-of-bounds index rather than silently mis-reading. + pools = min(src.shape[0], self.attn_row_block_offsets.shape[0]) + width = min(src.shape[-1], self.attn_row_block_offsets.shape[-1]) + torch.index_select( + src[:pools, :, :, :width], 1, self.attn_row_req_idx_cuda[:rows], + out=self.attn_row_block_offsets[:pools, :rows, :, :width]) def token_major_gen_view(self) -> Optional["TokenMajorGenView"]: """The token-major presentation for this step, or None. @@ -2105,7 +2112,7 @@ def token_major_gen_view(self) -> Optional["TokenMajorGenView"]: sequence_length=self.attn_row_kv_lens_cuda[:rows], host_past_key_value_lengths=self.attn_row_kv_lens_host[:rows], host_context_lengths=self.attn_row_prompt_lens_cpu[:rows], - kv_cache_block_offsets=self.attn_row_block_offsets[:rows], + kv_cache_block_offsets=self.attn_row_block_offsets[:, :rows], max_num_rows=self.attn_row_kv_lens_cuda.shape[0], ) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py new file mode 100644 index 000000000000..f8aa5ab678e5 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py @@ -0,0 +1,152 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Token-major row expansion for the DSv4 MLA generation ops. + +The kernel-level tests in ``test_deepseek_v4_sparse_mla.py`` call the ops with +hand-built tensors, so they pass whatever the metadata layer produces -- they +cannot see a view that is assembled wrongly. That gap let a block table shaped +``[rows, max_blocks]`` reach an op that reads +``[num_pools, num_seqs, 2, max_blocks]``: the gather collapsed the leading axes +until only two were left, which stripped ``num_seqs`` rather than the pool axis +and left request ids indexing a length-2 K/V axis. It surfaced only on an +8-GPU end-to-end run, as a device-side ``vectorized gather kernel index out of +bounds``. + +These tests exercise the expansion itself against a stub cache manager, so the +same mistake fails in seconds on one GPU. +""" + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.sparse.dsa import ( + DSAtrtllmAttentionMetadata) + + +class _StubCacheManager: + """Only the attributes the row expansion reads.""" + + def __init__(self, num_pools: int, max_blocks_per_seq: int): + self.num_pools = num_pools + self.num_attention_op_pools = num_pools + self.max_blocks_per_seq = max_blocks_per_seq + + +def _metadata(*, num_pools, max_blocks, max_seqs, draft=5): + """A metadata object with just enough state to build its own buffers. + + Constructed through ``__new__`` rather than the real constructor: the real + one wants a runtime, a resource manager and a model config, none of which + the row expansion touches. What it does touch is set explicitly below, so a + new dependency shows up as a clear AttributeError here instead of being + silently satisfied by a mock. + """ + meta = DSAtrtllmAttentionMetadata.__new__(DSAtrtllmAttentionMetadata) + meta.kv_cache_manager = _StubCacheManager(num_pools, max_blocks) + meta.max_num_sequences = max_seqs + meta.max_num_requests = max_seqs + meta.max_draft_tokens = draft + meta.cuda_graph_buffers = None + meta.num_sms = torch.cuda.get_device_properties(0).multi_processor_count + meta.runtime_tokens_per_gen_step = 0 + meta.create_expanded_buffers(capture_graph=False) + meta.prompt_lens_cpu = torch.zeros(max_seqs, dtype=torch.int32) + return meta + + +def _expand(meta, *, num_contexts, row_req_idx, num_seqs, block_offsets): + """Drive the private expansion with an explicit row->request map.""" + ngt = len(row_req_idx) + meta.num_contexts = num_contexts + meta._ragged_num_rows = ngt + meta.ragged_verify_lens = [1] * ngt + meta.kv_cache_block_offsets = block_offsets + meta.row_req_idx_host[:ngt].copy_( + torch.tensor(row_req_idx, dtype=torch.long)) + meta.row_kv_lens_host[:ngt].fill_(7) + meta.row_kv_correction_host[:ngt].zero_() + meta.prompt_lens_cpu[:num_seqs].fill_(1) + meta._prepare_token_major_gen_rows( + torch.full((num_seqs, ), 7, dtype=torch.int32)) + return meta._attn_num_rows + + +@pytest.mark.parametrize("num_pools", [1, 2]) +def test_row_block_table_matches_op_layout(num_pools): + """The expanded table keeps the op's rank, axis order and pool/KV axes.""" + torch.cuda.init() + max_blocks, max_seqs = 8, 4 + meta = _metadata(num_pools=num_pools, max_blocks=max_blocks, + max_seqs=max_seqs) + + src = torch.arange(num_pools * max_seqs * 2 * max_blocks, + dtype=torch.int32, + device="cuda").reshape(num_pools, max_seqs, 2, + max_blocks) + # Request 2 verifies three positions, request 0 one: a ragged batch, which + # is the whole point -- rows per request differ. + row_req_idx = [0, 2, 2, 2] + rows = _expand(meta, + num_contexts=0, + row_req_idx=row_req_idx, + num_seqs=max_seqs, + block_offsets=src) + assert rows == len(row_req_idx) + + got = meta.token_major_gen_view().kv_cache_block_offsets + assert got.dim() == src.dim(), ( + f"expanded table has rank {got.dim()} but the attention op reads rank " + f"{src.dim()}; substituting it wholesale would misread every axis") + assert got.shape[0] == num_pools and got.shape[2] == 2 + assert got.shape[1] == rows + # Every row must carry its own request's blocks, on every pool and both + # directions. + for row, req in enumerate(row_req_idx): + torch.testing.assert_close(got[:, row], src[:, req]) + + +def test_row_indices_stay_in_bounds_of_the_sequence_axis(): + """A request id must index num_seqs, never a shorter axis. + + This is the assertion the device-side gather was making. Checking it here + means the failure is a readable Python error rather than a CUDA abort that + takes the whole 8-rank job with it. + """ + torch.cuda.init() + meta = _metadata(num_pools=1, max_blocks=8, max_seqs=4) + + src = torch.zeros(1, 4, 2, 8, dtype=torch.int32, device="cuda") + rows = _expand(meta, + num_contexts=0, + row_req_idx=[0, 1, 3, 3], + num_seqs=4, + block_offsets=src) + idx = meta.attn_row_req_idx_cuda[:rows] + assert int(idx.min()) >= 0 + assert int(idx.max()) < src.shape[1], ( + f"row->request map reaches {int(idx.max())} but the sequence axis is " + f"{src.shape[1]} long") + + +def test_context_rows_precede_generation_rows(): + """Contexts keep row 0..nc-1 so the ops' generation offset still lands.""" + torch.cuda.init() + meta = _metadata(num_pools=1, max_blocks=8, max_seqs=6) + + src = torch.arange(1 * 6 * 2 * 8, dtype=torch.int32, + device="cuda").reshape(1, 6, 2, 8) + nc = 2 + rows = _expand(meta, + num_contexts=nc, + row_req_idx=[0, 1, 1], + num_seqs=6, + block_offsets=src) + assert rows == nc + 3 + idx = meta.attn_row_req_idx_host[:rows].tolist() + assert idx[:nc] == list(range(nc)), "context rows must be identity-mapped" + # Generation row->request ids are batch-relative and must be shifted past + # the contexts, or a generation row would read a context's blocks. + assert idx[nc:] == [nc + 0, nc + 1, nc + 1] + # host_context_lengths is the ops' shape oracle: one token per generation + # row, real length for contexts. + assert meta.attn_row_prompt_lens_cpu[nc:rows].tolist() == [1, 1, 1] From 6e7c3448f8cf4b021d3d8884d52f6c0b7e222c38 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 21:19:15 -0700 Subject: [PATCH 038/177] [None][fix] Get more SPS samples by repeating cells, not lengthening them Raising --measure-steps to 1024 made every cell come back empty. The stats queue retains on the order of 1000 rows per drain and attention-DP publishes one row per rank per iteration, so at dp=8 only about 125 iterations survive -- and they are the last 125, which is the drain phase, where ranks disagree on numGenRequests and nothing aligns. A 1300-step cell kept exactly 1001 rows covering iterations 381..506 with numGenRequests in [0, 1]. Longer cells are therefore strictly worse. Add --repeats, which reruns a cell and pools the samples; each repeat drains the queue again, so samples accumulate instead of overflowing. Warn at validate() time when a cell cannot fit in the window, with the measure-steps value that would. Signed-off-by: Lance Liao --- .../_torch/speculative/dspark_sps_profiler.py | 42 ++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py b/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py index 4e97a9287646..961398bedb7a 100644 --- a/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py +++ b/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py @@ -1007,6 +1007,7 @@ class SweepConfig: input_len: int = 1024 warmup_steps: int = 16 measure_steps: int = 64 + repeats: int = 1 min_samples: int = 16 max_seq_len: int = 4096 max_num_tokens: int = 8192 @@ -1059,6 +1060,27 @@ def validate(self) -> None: f"generated) but --max-seq-len is {self.max_seq_len}; raise it or " f"lower --input-len / --measure-steps" ) + # get_stats is served from a queue that retains on the order of 1000 + # ROWS, and attention-DP publishes one row per rank per iteration. A + # cell longer than that window keeps only its tail, which is the drain + # phase -- every step there has ranks disagreeing on numGenRequests, so + # the cell yields nothing and looks like it never ran. Raising + # --measure-steps past the window therefore makes things strictly worse; + # use --repeats to get more samples instead. + rows_per_iteration = self.dp_size + window_iterations = _STATS_ROW_BUDGET // max(rows_per_iteration, 1) + cell_iterations = int(self.warmup_steps) + int(self.measure_steps) + if cell_iterations > window_iterations: + print( + f"[dspark-sps] WARNING: a cell runs {cell_iterations} iterations " + f"but only about {window_iterations} fit in the retained stats " + f"window ({_STATS_ROW_BUDGET} rows / {rows_per_iteration} ranks). " + f"Only the tail survives, and the tail is the drain phase, so " + f"cells will come back empty. Lower --measure-steps to " + f"<= {window_iterations - int(self.warmup_steps)} and raise " + f"--repeats instead.", + file=sys.stderr, + ) if self.input_len > self.max_num_tokens: raise SweepGeometryError( f"--input-len {self.input_len} exceeds --max-num-tokens " @@ -1074,6 +1096,12 @@ def validate(self) -> None: ) +#: Approximate number of stats rows ``get_stats`` retains for one drain. Not a +#: constant the runtime exports; measured as 1001 rows surviving a 1300-step +#: attention-DP cell, i.e. the last 125 of its iterations. +_STATS_ROW_BUDGET = 1000 + + def _prepare_environment(config: SweepConfig) -> None: """Pin the knobs the table's meaning depends on, loudly. @@ -1235,7 +1263,13 @@ def run_sweep(config: SweepConfig) -> List[StepSample]: try: for batch_size in sorted(config.batch_sizes): started = time.time() - cell = _run_cell(llm, config, batch_size=batch_size, verify_len=verify_len) + cell = [] + for _ in range(max(1, int(config.repeats))): + # Each repeat drains the stats queue again, so repeats add + # samples where a longer single cell would only overflow the + # retained window. + cell.extend(_run_cell(llm, config, batch_size=batch_size, + verify_len=verify_len)) print( f"[dspark-sps] bs={batch_size:>5} L={verify_len} " f"M={total_verify_tokens(batch_size, verify_len):>6} " @@ -1360,6 +1394,11 @@ def _build_parser() -> argparse.ArgumentParser: help="Do not force bonus-token-only acceptance. Leaves KV growth rate " "dependent on the block length, which contaminates theta(M).", ) + parser.add_argument( + "--repeats", type=int, default=1, + help="run each cell this many times and pool the samples; the way to " + "get more samples per cell, since a longer cell overflows the " + "retained stats window and keeps only its drain phase") parser.set_defaults(pin_acceptance=True) fit = parser.add_argument_group("fit") @@ -1444,6 +1483,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int: input_len=args.input_len, warmup_steps=args.warmup_steps, measure_steps=args.measure_steps, + repeats=args.repeats, min_samples=args.min_samples, max_seq_len=args.max_seq_len, max_num_tokens=args.max_num_tokens, From f2073b821aae9c8e9ad6c0548f50f1d97b79ddcc Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 21:22:08 -0700 Subject: [PATCH 039/177] [None][fix] Present request types and prompt lens token-major too Only four of the six runtime views were being swapped. The attention op reads request_types over its whole num_seqs -- for (idx = num_contexts; idx < num_seqs; idx++) TLLM_CHECK(request_types[idx] == kGENERATION) (thop/attentionOp.cpp) -- and num_seqs is host_context_lengths.size(0), which this presentation makes the row count. Leaving the batch-major array bound there reads past its end once rows exceed max_num_requests, which on DEP8 is 896 rows against a 128-entry array. TrtllmAttention. forward also asserts every runtime view shares batch_size with kv_lens_cuda_runtime, so prompt_lens_cuda needed the same treatment. Found by auditing the remaining swapped tensors after the block-table layout bug, on the theory that one layout assumption is rarely alone. Locked in by a test asserting all six views report the row count. Signed-off-by: Lance Liao --- .../_torch/attention_backend/sparse/dsa.py | 30 +++++++++++++++++ .../_torch/attention_backend/trtllm.py | 11 +++++-- .../deepseek_v4/test_token_major_rows.py | 33 +++++++++++++++++++ 3 files changed, 72 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 90fc24e1b247..520507001a26 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -595,6 +595,8 @@ class TokenMajorGenView: sequence_length: torch.Tensor host_past_key_value_lengths: torch.Tensor host_context_lengths: torch.Tensor + prompt_lens_cuda: torch.Tensor + host_request_types: torch.Tensor kv_cache_block_offsets: torch.Tensor #: Static row ceiling, NOT this step's row count. The attention op caches #: by (beam_width, max_num_requests, attention_window_size) and reserves its @@ -1508,6 +1510,25 @@ def create_expanded_buffers(self, capture_graph=False): self.attn_row_req_idx_host = torch.zeros( (attn_rows_cap, ), dtype=torch.long, pin_memory=prefer_pinned()) # host_context_lengths: the shape oracle both ops key off. + # The op indexes request_types over its whole num_seqs -- which under + # this presentation is the ROW count -- in + # `for (idx = num_contexts; idx < num_seqs; idx++) + # TLLM_CHECK(request_types[idx] == kGENERATION)` + # (thop/attentionOp.cpp). Handing it the batch-major array would read + # past its end. Generation is 1, context 0; prefill the whole capacity + # as generation so a step only has to stamp its context prefix. + self.attn_row_request_types_host = torch.ones( + (attn_rows_cap, ), dtype=torch.int32, + pin_memory=prefer_pinned()) + # Device twin of attn_row_prompt_lens_cpu: the forward path asserts + # every runtime view shares batch_size with kv_lens_cuda_runtime. + self.attn_row_prompt_lens_cuda = self.get_empty( + self.cuda_graph_buffers, + (attn_rows_cap, ), + cache_name="attn_row_prompt_lens_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) self.attn_row_prompt_lens_cpu = torch.zeros( (attn_rows_cap, ), dtype=torch.int32, pin_memory=prefer_pinned()) @@ -2032,6 +2053,11 @@ def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: # (the ops read it for num_seqs and, for contexts, for the q length); # one token per row. self.attn_row_prompt_lens_cpu[nc:rows].fill_(1) + # Capacity is prefilled with kGENERATION, so only the context prefix + # needs stamping, and nc is small. + if nc: + self.attn_row_request_types_host[:nc].zero_() + self.attn_row_request_types_host[nc:rows].fill_(1) self.attn_row_kv_lens_cuda[:rows].copy_( self.attn_row_kv_lens_host[:rows], non_blocking=True) @@ -2039,6 +2065,8 @@ def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: self.attn_row_kv_correction_host[:rows], non_blocking=True) self.attn_row_req_idx_cuda[:rows].copy_( self.attn_row_req_idx_host[:rows], non_blocking=True) + self.attn_row_prompt_lens_cuda[:rows].copy_( + self.attn_row_prompt_lens_cpu[:rows], non_blocking=True) self._attn_num_rows = rows # The block table is expanded on the DEVICE, from the copy that is @@ -2112,6 +2140,8 @@ def token_major_gen_view(self) -> Optional["TokenMajorGenView"]: sequence_length=self.attn_row_kv_lens_cuda[:rows], host_past_key_value_lengths=self.attn_row_kv_lens_host[:rows], host_context_lengths=self.attn_row_prompt_lens_cpu[:rows], + prompt_lens_cuda=self.attn_row_prompt_lens_cuda[:rows], + host_request_types=self.attn_row_request_types_host[:rows], kv_cache_block_offsets=self.attn_row_block_offsets[:, :rows], max_num_rows=self.attn_row_kv_lens_cuda.shape[0], ) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index d62a9fd71202..8fd40446d752 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -547,7 +547,8 @@ def presented_token_major(self, view): return saved = (self.kv_lens_cuda_runtime, self.kv_lens_runtime, self.prompt_lens_cpu_runtime, self.kv_cache_block_offsets, - self.max_num_requests) + self.max_num_requests, self.prompt_lens_cuda_runtime, + self.host_request_types_runtime) try: self.kv_lens_cuda_runtime = view.sequence_length self.kv_lens_runtime = view.host_past_key_value_lengths @@ -563,11 +564,17 @@ def presented_token_major(self, view): # JIT-warmup hint, and the per-step request count comes from the # batch, not from here. self.max_num_requests = view.max_num_rows + # Every runtime view has to agree on batch_size with + # kv_lens_cuda_runtime, and the op reads request_types over all of + # num_seqs -- both are the row count here, not the request count. + self.prompt_lens_cuda_runtime = view.prompt_lens_cuda + self.host_request_types_runtime = view.host_request_types yield finally: (self.kv_lens_cuda_runtime, self.kv_lens_runtime, self.prompt_lens_cpu_runtime, self.kv_cache_block_offsets, - self.max_num_requests) = saved + self.max_num_requests, self.prompt_lens_cuda_runtime, + self.host_request_types_runtime) = saved def prepare(self) -> None: super().prepare() diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py index f8aa5ab678e5..231673b8e0a9 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_token_major_rows.py @@ -150,3 +150,36 @@ def test_context_rows_precede_generation_rows(): # host_context_lengths is the ops' shape oracle: one token per generation # row, real length for contexts. assert meta.attn_row_prompt_lens_cpu[nc:rows].tolist() == [1, 1, 1] + + +def test_every_runtime_view_shares_the_row_count(): + """All views the op reads must agree on batch_size, which is the row count. + + ``TrtllmAttention.forward`` asserts this directly, and the C++ op indexes + ``request_types`` over its whole ``num_seqs`` -- so a view left at the + request count is an out-of-bounds read there rather than a shape error. + """ + torch.cuda.init() + meta = _metadata(num_pools=1, max_blocks=8, max_seqs=6) + src = torch.zeros(1, 6, 2, 8, dtype=torch.int32, device="cuda") + nc = 1 + rows = _expand(meta, + num_contexts=nc, + row_req_idx=[0, 0, 1], + num_seqs=6, + block_offsets=src) + + view = meta.token_major_gen_view() + for name in ("sequence_length", "host_past_key_value_lengths", + "host_context_lengths", "prompt_lens_cuda", + "host_request_types"): + got = getattr(view, name).shape[0] + assert got == rows, ( + f"{name} has {got} entries but the ops will read {rows} rows") + assert view.kv_cache_block_offsets.shape[1] == rows + + # kCONTEXT is 0, kGENERATION is 1; the op checks every row at or past + # num_contexts is a generation row. + types = view.host_request_types.tolist() + assert types[:nc] == [0] * nc + assert types[nc:] == [1] * (rows - nc) From 3b720c6c26c08badb9fc1fae9c2e5f94c77282da Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 21:47:07 -0700 Subject: [PATCH 040/177] [None][chore] Record the measured workspace cost of the row-count capacity An earlier commit message claimed the generation workspace barely grows when max_num_requests is raised to the row ceiling, reasoning from getFmhaMultiCtasKvScratchSize being SM-count based. The end-to-end run contradicts it: on DEP8 the attention workspace went from 68 MiB to 604 MiB per rank, consistent with the 7x row expansion. The op resizes itself, so this is GPU memory that would otherwise hold KV cache rather than a correctness problem -- but it is worth knowing before tuning free_gpu_memory_fraction, so put the real number next to the code that causes it. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/attention_backend/trtllm.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 8fd40446d752..944952dddfa0 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -563,6 +563,14 @@ def presented_token_major(self, view): # This is only a capacity input -- the op's own mMaxNumRequests is a # JIT-warmup hint, and the per-step request count comes from the # batch, not from here. + # + # It is not free: the attention workspace is sized from the same + # number, and on DEP8 (128 requests -> 896 rows) it was measured + # growing from 68 MiB to 604 MiB per rank. The op resizes itself, so + # this costs GPU memory that would otherwise hold KV cache rather + # than correctness. Shrinking it would mean sizing the generation + # workspace from the row count only where the row count is what the + # kernels actually use, which is a C++ change. self.max_num_requests = view.max_num_rows # Every runtime view has to agree on batch_size with # kv_lens_cuda_runtime, and the op reads request_types over all of From febb925d50d05720e3f08d0dde4d17366644fd69 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 22:28:21 -0700 Subject: [PATCH 041/177] [None][fix] Average only numeric lm-eval metrics When no scores_filter is given, the average was taken over every score entry that is not a stderr key. lm-eval also puts non-numeric entries in that dict -- notably "alias", whose value is the task name -- so numpy fails with "the resolved dtypes are not compatible with add.reduce. Resolved (dtype(' --- tensorrt_llm/evaluate/lm_eval.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index b077747b6c03..e14dfd7c35ab 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -654,8 +654,21 @@ def evaluate(self, f"lm-eval {self.task_name} {scores_filter} accuracy: {result_acc:.2f}" ) else: - result_acc = np.mean( - [acc for m, acc in scores.items() if "_stderr" not in m]) + # lm-eval mixes non-numeric entries into the score dict -- notably + # "alias", whose value is the task name -- so averaging everything + # that is not a stderr key fails with "resolved dtypes are not + # compatible with add.reduce" once one of them is a string. Keep + # only the numeric metrics. + numeric = [ + acc for m, acc in scores.items() + if "_stderr" not in m and isinstance(acc, (int, float)) + and not isinstance(acc, bool) + ] + if not numeric: + raise ValueError( + f"lm-eval returned no numeric metrics for " + f"{self.task_name}; got {sorted(scores)}") + result_acc = np.mean(numeric) logger.info( f"lm-eval {self.task_name} average accuracy: {result_acc:.2f}") return result_acc From f9abb75967f5fdb6eef6ec671efc89e05f5b89a5 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 22:38:20 -0700 Subject: [PATCH 042/177] [None][chore] Record which shape a ragged CUDA graph miss submitted graph_eager counted misses without saying which shape missed, so a run that reports "dropped out of CUDA graph replay on 9 step(s)" leaves only guessing or re-running under a profiler -- and the two candidate causes (a padded_bs the bucket grid was not derived from, versus a token total above the largest captured bucket) want different fixes. Keep the distinct (rows, total_tokens) of each miss and report them in the summary. Same reasoning as the SPS alignment histogram: a counter that cannot name what it counted costs a full run per hypothesis. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/model_engine.py | 13 +++++++--- .../speculative/dspark_observability.py | 26 +++++++++++++++---- 2 files changed, 31 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b56f53011efb..ca47622ec57f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3644,7 +3644,7 @@ def ragged_verify_token_buckets(self, padded_bs: int) -> List[int]: }) return [int(padded_bs) * (t + 1) for t in tiers] - def _record_dspark_graph_use(self, *, replayed: bool) -> None: + def _record_dspark_graph_use(self, *, replayed: bool, shape=None) -> None: """Count graph replay vs eager for the DSpark ragged stats, if enabled. Deliberately tolerant: this runs on the hot path for every model, and a @@ -3653,7 +3653,7 @@ def _record_dspark_graph_use(self, *, replayed: bool) -> None: """ stats = self._dspark_ragged_stats if stats is not None: - stats.record_graph(replayed=replayed) + stats.record_graph(replayed=replayed, shape=shape) @property def _dspark_ragged_stats(self): @@ -3739,6 +3739,9 @@ def fit_ragged_verify_lens( f"to uniform scheduling: {exc}") return None padded_bs, bucket = shape.padded_bs, shape.bucket + # Remembered so a graph miss can report the shape it actually submitted + # rather than just that one happened. + self._dspark_last_total_tokens = int(bucket) # The CUDA-graph padding rows are appended later by _get_padded_batch, # which appends a *single shared dummy object*, so every pad row @@ -6915,7 +6918,11 @@ def forward(self, # miscount alone makes the check unsatisfiable on every real # workload. if padded_requests.can_run_cuda_graph: - self._record_dspark_graph_use(replayed=can_run_graph) + self._record_dspark_graph_use( + replayed=can_run_graph, + shape=None if can_run_graph else + (len(padded_requests.generation_requests), + int(getattr(self, "_dspark_last_total_tokens", 0)))) if can_run_graph: attn_metadata = maybe_attn_metadata spec_metadata = maybe_spec_metadata diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index 73d87e584d4b..9f7084e38c49 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -61,7 +61,7 @@ import os from collections import Counter from enum import Enum -from typing import Dict, List, Optional, Sequence +from typing import Dict, List, Optional, Sequence, Tuple from ...logger import logger @@ -247,6 +247,8 @@ def __init__(self, *, mode: RaggedVerifyMode, max_draft_len: int): self.verify_len_hist: Counter = Counter() self.bucket_hist: Counter = Counter() self.padded_bs_hist: Counter = Counter() + #: (rows, total_tokens) -> count, for steps that found no captured graph. + self.graph_miss_shapes: Dict[Tuple[int, int], int] = {} #: Acceptance, split by whether the request's window was shortened. #: This is the term the trim/ceiling ratio needs and the one TRT-LLM @@ -328,12 +330,22 @@ def record_step( else: self.steps_uniform_windows += 1 - def record_graph(self, *, replayed: bool) -> None: - """Record whether this step replayed a CUDA graph or ran eager.""" + def record_graph(self, *, replayed: bool, shape=None) -> None: + """Record whether this step replayed a CUDA graph or ran eager. + + ``shape`` is the ``(num_rows, total_verify_tokens)`` the step actually + submitted. It is kept only for misses, and only the distinct ones: a + miss means some shape was not captured, and knowing *which* is the + whole diagnosis. Counting misses without recording the shape leaves + only the options of guessing or re-running with a profiler. + """ if replayed: self.graph_replays += 1 - else: - self.graph_eager += 1 + return + self.graph_eager += 1 + if shape is not None: + key = (int(shape[0]), int(shape[1])) + self.graph_miss_shapes[key] = self.graph_miss_shapes.get(key, 0) + 1 def merge_planner_stats(self, planner_stats: Dict[str, int]) -> None: """Fold in the planner's own fallback counters.""" @@ -408,6 +420,10 @@ def summary(self) -> Dict[str, object]: "verify_len_hist": dict(sorted(self.verify_len_hist.items())), "bucket_hist": dict(sorted(self.bucket_hist.items())), "padded_bs_hist": dict(sorted(self.padded_bs_hist.items())), + "graph_miss_shapes": { + f"{rows}x{tok}": n + for (rows, tok), n in sorted(self.graph_miss_shapes.items()) + }, "accept_len": round(self.accept_len, 4), "requests_scored": self.requests_scored, "requests_trimmed": self.requests_trimmed, From dcb64ebeeed20f87a9f58406fccfa93c42ba1db0 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 23:31:07 -0700 Subject: [PATCH 043/177] [None][fix] Run the DSpark ragged gate mid-loop so it can fail the run The gate ran only from _executor_loop_cleanup, which happens after every request has completed and the driver has returned its results, so the AssertionError landed in a shutdown path nothing observes. A GSM8K run whose ragged windows never differed reported "1 passed" while all eight ranks raised the assertion in their worker logs -- the worst outcome for a guard, because the green result is the part that gets believed. Check from inside the executor loop once the step floor is reached, where the exception propagates out through the in-flight requests and fails the run. Once only: the counters keep moving and re-raising every step would bury the first and only informative failure. The cleanup-time check stays for the unconditional summary. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4b1e59897571..48fdc2f3eb65 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2444,6 +2444,44 @@ def _maybe_arm_dspark_sync_debug(self) -> None: f"device->host synchronization from here on emits a warning with a " f"stack; on the overlap-scheduler path there should be none.") + def _maybe_assert_dspark_ragged_active(self, stats) -> None: + """Run the ragged gate mid-loop, once the step floor is reached. + + The same check also runs at executor-loop cleanup, but only there it is + useless as a gate: cleanup happens after every request has completed and + the driver has already returned its results, so the AssertionError is + raised into a shutdown path that pytest never sees. A run whose ragged + path silently did nothing reported ``1 passed`` with eight ranks each + raising in their logs -- the worst possible outcome for a guard, since + the green result is what gets believed. + + Raised from inside the loop instead, the exception propagates out + through the in-flight requests and fails the run. Only once: the + counters keep moving, and re-raising every step would bury the first + (and only informative) failure. + """ + # Lazy rather than set in __init__: this class has several + # construction paths and a gate that silently no-ops because one of + # them missed an attribute is the exact failure being fixed here. + if getattr(self, "_dspark_asserted", False): + return + min_steps = self._dspark_assert_min_steps + if min_steps is None or stats.steps_total < min_steps: + return + self._dspark_asserted = True + stats.assert_ragged_active(require_trim=True) + + @property + def _dspark_assert_min_steps(self) -> Optional[int]: + """Step floor from TLLM_DSPARK_ASSERT_ACTIVE, or None when disabled.""" + raw = os.environ.get("TLLM_DSPARK_ASSERT_ACTIVE", "") + if not raw or raw in ("0", "false", "False"): + return None + try: + return 32 if raw in ("1", "true", "True") else int(raw) + except ValueError: + return None + def _report_dspark_ragged_stats(self) -> None: """Final DSpark scheduling summary, and optionally a hard gate on it. @@ -3445,6 +3483,7 @@ def _dspark_confidence_draft_len(self, # recoverable from the log, which is the only shared channel. if stats.steps_total and stats.steps_total % 32 == 0: stats.log_summary(prefix="DSpark ragged verify [periodic]") + self._maybe_assert_dspark_ragged_active(stats) fitted = None if bucket is not None: fitted = [ From e8c6364e1de06dc82bcffdaf885c4c990e54edca Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sat, 1 Aug 2026 23:45:41 -0700 Subject: [PATCH 044/177] [None][fix] Attribute ragged CUDA-graph misses to a reason graph_eager counted misses without recording why, and the reasons mean opposite things. A batch where some rank carries context requests drops the whole attention-DP world out of replay and happens on the uniform path too -- that run sat at ~1.7% eager for exactly this reason -- so gating on the raw count makes the check unsatisfiable on any workload with continuous batching. But "the batch was graphable, every rank agreed, and the key was still missing" means the token bucket grid is wrong, and "one rank chose a different ragged bucket" is the attention-DP divergence this feature exists to avoid. Both were hidden behind the same integer. Record the reason at each of the graph runner's decline points, keep the graph key for the two where it localizes the bug, and blame only key_not_captured, bs_not_supported and peer_shape_mismatch. Also make `shape` a required argument of record_graph. A run reported graph_eager: 9 with an empty graph_miss_shapes because the caller came from a build predating the argument, and the default read as "no misses had shapes" -- indistinguishable from a working probe finding nothing. Misses with no reason recorded now fail as a broken probe rather than passing as a clean run. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/cuda_graph_runner.py | 27 +++++++ .../_torch/pyexecutor/model_engine.py | 5 +- .../speculative/dspark_observability.py | 70 ++++++++++++---- .../test_dspark_graph_miss_attribution.py | 80 +++++++++++++++++++ 4 files changed, 164 insertions(+), 18 deletions(-) create mode 100644 tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 99900e4acac2..e83988185c6a 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -368,8 +368,19 @@ def maybe_get_cuda_graph( - The spec_metadata for the graph, if applicable. - The key for the graph, if applicable. """ + # Why the last call declined, for the DSpark ragged counters. A miss is + # counted at the call site, which cannot see which of the several + # independent reasons applied -- and they mean opposite things: a peer + # shape mismatch is one rank declining to go ragged and dragging the + # whole world eager, while an uncaptured key is the token bucket grid + # being wrong. Attributing one to the other sends the fix to the wrong + # place. + self.last_miss_reason = None + self.last_miss_key = None + # disable when doing statistic if ExpertStatistic.should_record(): + self.last_miss_reason = "expert_statistic" return None, None, None can_run_cuda_graph = batch.can_run_cuda_graph @@ -405,9 +416,18 @@ def maybe_get_cuda_graph( for peer in all_can_graph_batch) if not is_all_gen_only or not all_shapes_equal: + # Deliberately distinguished: "some rank had contexts" is + # ordinary continuous batching, while "shapes differ" means a + # rank chose a different ragged bucket and is the one worth + # chasing. + self.last_miss_reason = ("peer_not_gen_only" + if not is_all_gen_only else + "peer_shape_mismatch") return None, None, None if not self.enabled or not can_run_cuda_graph: + self.last_miss_reason = ("disabled" if not self.enabled else + "local_not_gen_only") return None, None, None if self.config.use_mrope and any( self._needs_mrope_delta_cache_update(request) @@ -416,6 +436,7 @@ def maybe_get_cuda_graph( # Qwen3.5 configs normalized to text-only decoding). Requests that # do carry a delta must first populate the model-side cache for their # current seq slot before graph replay. + self.last_miss_reason = "mrope_delta_cache" return None, None, None key = self.get_graph_key(batch, new_tensors_device, spec_resource_manager, spec_metadata) @@ -427,9 +448,15 @@ def maybe_get_cuda_graph( # Graph doesn't exist yet. If on-the-fly capture is not allowed, # fall back to eager so the caller doesn't need a separate check. if not self._capture_allowed: + # THE reason that means the bucket grid is wrong: the batch was + # graphable and every rank agreed, but this key was never captured. + self.last_miss_reason = "key_not_captured" + self.last_miss_key = key return None, None, None if batch_size not in self.supported_batch_sizes: + self.last_miss_reason = "bs_not_supported" + self.last_miss_key = key return None, None, None num_sequences_in_batch = batch_size * self.max_beam_width diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index ca47622ec57f..f2e9109977c5 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -6921,8 +6921,9 @@ def forward(self, self._record_dspark_graph_use( replayed=can_run_graph, shape=None if can_run_graph else - (len(padded_requests.generation_requests), - int(getattr(self, "_dspark_last_total_tokens", 0)))) + (getattr(self.cuda_graph_runner, "last_miss_reason", None) + or "unknown", + getattr(self.cuda_graph_runner, "last_miss_key", None))) if can_run_graph: attn_metadata = maybe_attn_metadata spec_metadata = maybe_spec_metadata diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index 9f7084e38c49..be408018b8f6 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -247,8 +247,14 @@ def __init__(self, *, mode: RaggedVerifyMode, max_draft_len: int): self.verify_len_hist: Counter = Counter() self.bucket_hist: Counter = Counter() self.padded_bs_hist: Counter = Counter() - #: (rows, total_tokens) -> count, for steps that found no captured graph. - self.graph_miss_shapes: Dict[Tuple[int, int], int] = {} + #: why the graph runner declined -> count. The reasons mean opposite + #: things: "peer_shape_mismatch" is one rank declining to go ragged and + #: dragging the world eager, "key_not_captured" is the bucket grid being + #: wrong, "peer_not_gen_only" is ordinary continuous batching. + self.graph_miss_reasons: Dict[str, int] = {} + #: the actual graph key, kept only for the reasons where it localizes + #: the bug. + self.graph_miss_shapes: Dict[str, int] = {} #: Acceptance, split by whether the request's window was shortened. #: This is the term the trim/ceiling ratio needs and the one TRT-LLM @@ -330,22 +336,34 @@ def record_step( else: self.steps_uniform_windows += 1 - def record_graph(self, *, replayed: bool, shape=None) -> None: + def record_graph(self, *, replayed: bool, shape) -> None: """Record whether this step replayed a CUDA graph or ran eager. ``shape`` is the ``(num_rows, total_verify_tokens)`` the step actually - submitted. It is kept only for misses, and only the distinct ones: a - miss means some shape was not captured, and knowing *which* is the - whole diagnosis. Counting misses without recording the shape leaves - only the options of guessing or re-running with a profiler. + submitted, or None on a replay. It is kept only for misses, and only + the distinct ones: a miss means some shape was not captured, and + knowing *which* is the whole diagnosis. Counting misses without + recording the shape leaves only guessing or re-running with a profiler. + + Required rather than defaulted, deliberately. A run reported + ``graph_eager: 9`` alongside an empty ``graph_miss_shapes`` because the + caller was loaded from a build that predated this argument -- with a + default that reads as "no misses had shapes", which is indistinguishable + from "the probe works and found nothing". Without a default it is a + TypeError instead. """ if replayed: self.graph_replays += 1 return self.graph_eager += 1 if shape is not None: - key = (int(shape[0]), int(shape[1])) - self.graph_miss_shapes[key] = self.graph_miss_shapes.get(key, 0) + 1 + reason, key = shape + self.graph_miss_reasons[reason] = ( + self.graph_miss_reasons.get(reason, 0) + 1) + if key is not None: + text = str(key) + self.graph_miss_shapes[text] = ( + self.graph_miss_shapes.get(text, 0) + 1) def merge_planner_stats(self, planner_stats: Dict[str, int]) -> None: """Fold in the planner's own fallback counters.""" @@ -420,10 +438,8 @@ def summary(self) -> Dict[str, object]: "verify_len_hist": dict(sorted(self.verify_len_hist.items())), "bucket_hist": dict(sorted(self.bucket_hist.items())), "padded_bs_hist": dict(sorted(self.padded_bs_hist.items())), - "graph_miss_shapes": { - f"{rows}x{tok}": n - for (rows, tok), n in sorted(self.graph_miss_shapes.items()) - }, + "graph_miss_reasons": dict(sorted(self.graph_miss_reasons.items())), + "graph_miss_keys": dict(sorted(self.graph_miss_shapes.items())), "accept_len": round(self.accept_len, 4), "requests_scored": self.requests_scored, "requests_trimmed": self.requests_trimmed, @@ -474,11 +490,33 @@ def assert_ragged_active(self, *, require_trim: bool = True) -> None: raise AssertionError( f"DSpark ragged verify delivered as many tokens as a no-trim " f"run, so nothing was saved: {summary}") - if self.graph_eager > 0: + # Only the misses ragged is responsible for. A batch some rank filled + # with context requests falls out of replay on the uniform path too -- + # the uniform GSM8K run sat at ~1.7% eager for exactly that reason -- + # so failing on it makes this check unsatisfiable on any workload with + # continuous batching, which is all of them. The three below are + # different: "key_not_captured" and "bs_not_supported" mean the batch + # was graphable, every rank agreed, and the shape was still missing -- + # i.e. the token bucket grid is wrong. "peer_shape_mismatch" means one + # rank chose a different ragged bucket and dragged the whole world + # eager, which is the attention-DP divergence this feature has to avoid. + blamed = {"key_not_captured", "bs_not_supported", "peer_shape_mismatch"} + attributable = { + reason: n + for reason, n in self.graph_miss_reasons.items() if reason in blamed + } + if attributable: + raise AssertionError( + f"DSpark ragged verify dropped out of CUDA graph replay for " + f"reasons attributable to ragged scheduling " + f"({attributable}); a ragged step without a captured shape " + f"costs far more than the tokens it trims: {summary}") + if self.graph_eager > 0 and not self.graph_miss_reasons: raise AssertionError( f"DSpark ragged verify dropped out of CUDA graph replay on " - f"{self.graph_eager} step(s); a ragged step without a captured " - f"shape costs far more than the tokens it trims: {summary}") + f"{self.graph_eager} step(s) but recorded no reason for any of " + f"them, so the misses cannot be attributed. That is a broken " + f"probe, not a clean run: {summary}") def format_verify_len_histogram(lens: List[int]) -> str: diff --git a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py new file mode 100644 index 000000000000..fe1a4d1cc2fb --- /dev/null +++ b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py @@ -0,0 +1,80 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Which CUDA-graph misses the ragged gate is allowed to blame on ragged. + +A miss can mean opposite things and the counter alone cannot tell them apart. +``peer_not_gen_only`` is ordinary continuous batching -- some rank had context +requests, which drops the whole attention-DP world out of replay and happens on +the uniform path too. ``key_not_captured`` means the batch was graphable, every +rank agreed, and the shape was still missing: the token bucket grid is wrong. +``peer_shape_mismatch`` means one rank chose a different ragged bucket, which is +the attention-DP divergence this feature exists to avoid. + +Gating on the raw count made the check unsatisfiable on any real workload; not +gating at all let a real bucket bug through. These tests pin the split. +""" + +import pytest + +from tensorrt_llm._torch.speculative.dspark_observability import ( + DSparkRaggedStats, RaggedVerifyMode) + + +def _stats(**reasons): + """Stats that look like a healthy ragged run apart from the given misses.""" + stats = DSparkRaggedStats(mode=RaggedVerifyMode.COMPACT, max_draft_len=5) + stats.steps_total = 100 + stats.steps_ragged = 80 + # distinct_verify_lens is derived from this histogram, so populate the + # histogram rather than the property. + stats.verify_len_hist.update({1: 10, 3: 10, 5: 10}) + stats.ceiling_tokens = 1000 + stats.delivered_tokens = 800 # trim_ratio 0.2 > 0, so require_trim passes + stats.window_tokens = 800 + stats.graph_replays = 500 + for reason, count in reasons.items(): + stats.graph_eager += count + stats.graph_miss_reasons[reason] = count + return stats + + +def test_continuous_batching_misses_are_not_blamed_on_ragged(): + """A rank with context requests is normal, not a ragged defect.""" + _stats(peer_not_gen_only=9).assert_ragged_active(require_trim=True) + + +@pytest.mark.parametrize( + "reason", ["key_not_captured", "bs_not_supported", "peer_shape_mismatch"]) +def test_ragged_attributable_misses_raise(reason): + with pytest.raises(AssertionError) as excinfo: + _stats(**{reason: 1}).assert_ragged_active(require_trim=True) + assert reason in str(excinfo.value) + + +def test_mixed_misses_raise_and_name_only_the_attributable_ones(): + with pytest.raises(AssertionError) as excinfo: + _stats(peer_not_gen_only=20, + key_not_captured=2).assert_ragged_active(require_trim=True) + message = str(excinfo.value) + assert "key_not_captured" in message + # The tolerated reason must not appear in the *blame* dict, or the message + # sends the reader to the wrong place. It still appears in the summary. + blame = message.split("(", 1)[1].split(")", 1)[0] + assert "peer_not_gen_only" not in blame + + +def test_unattributed_misses_are_a_broken_probe_not_a_clean_run(): + """graph_eager > 0 with no reasons recorded must not read as success. + + This is what actually happened: a build skew left the reason unrecorded and + the summary showed misses with an empty reason map. Silently passing there + would make a broken probe indistinguishable from a healthy run. + """ + stats = _stats() + stats.graph_eager = 9 # counted, but no reason recorded + with pytest.raises(AssertionError, match="cannot be attributed"): + stats.assert_ragged_active(require_trim=True) + + +def test_clean_run_passes(): + _stats().assert_ragged_active(require_trim=True) From e86cbf7510d632896365fe411261c62cf1a9740b Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 00:25:35 -0700 Subject: [PATCH 045/177] [None][fix] Average only lm-eval metrics, not sample counts An earlier commit fixed a TypeError here by averaging every numeric score entry, which replaced a crash with something worse: a GSM8K run reported 44030.98 against a reference of 96.0 and PASSED. The score dict carries bookkeeping alongside metrics -- "alias" is the task name and "samples" is the example count -- and the normalization loop scales every numeric field by 100, so gsm8k's 1319 samples became 131900 and dominated the mean of [96.4367, 96.5125, 131900]. lm-eval keys metrics as ","; bookkeeping keys have no comma. Use that structurally for both the normalization and the average, so a new bookkeeping field cannot silently become a score. A wrong number that clears the threshold is worse than the exception it replaced, so this is covered by tests using the exact dict shape that produced it. Signed-off-by: Lance Liao --- tensorrt_llm/evaluate/lm_eval.py | 40 ++++++++++---- .../evaluate/test_lm_eval_score_selection.py | 54 +++++++++++++++++++ 2 files changed, 85 insertions(+), 9 deletions(-) create mode 100644 tests/unittest/evaluate/test_lm_eval_score_selection.py diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index e14dfd7c35ab..f0729fdc4538 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -51,6 +51,18 @@ LM_EVAL_DEFAULT_IMAGE_PLACEHOLDER = "" +def _is_metric_key(key: str) -> bool: + """Whether a ``results[task]`` key is a metric rather than bookkeeping. + + lm-eval keys every metric as ``","`` -- ``"exact_match, + strict-match"``, ``"acc,none"``. Entries without a comma are bookkeeping: + ``"alias"`` is the task name, ``"samples"`` is the example count. Both had + been treated as scores, which scaled gsm8k's 1319 samples to 131900 and + averaged it into the reported accuracy. + """ + return "," in key + + class LmEvalWrapper(TemplateLM): def __init__(self, @@ -635,10 +647,17 @@ def evaluate(self, system_instruction=self.system_prompt, log_samples=self.log_samples) - # Normalize scores to range 0~100 + # Normalize scores to range 0~100. + # + # Only real metrics. lm-eval keys a metric as "," (for + # example "exact_match,strict-match"); everything without a comma is + # bookkeeping -- "alias" holds the task name, "samples" holds the + # example count. Scaling those too turned gsm8k's 1319 samples into + # 131900, which then dominated the average below. scores = results["results"][self.task_name] for metric in scores.keys(): - if isinstance(scores[metric], (float, int)): + if _is_metric_key(metric) and isinstance(scores[metric], + (float, int)): scores[metric] *= 100 logger.info( f"lm-eval {self.task_name} results (scores normalized to range 0~100):\n{lm_eval.utils.make_table(results)}" @@ -654,15 +673,18 @@ def evaluate(self, f"lm-eval {self.task_name} {scores_filter} accuracy: {result_acc:.2f}" ) else: - # lm-eval mixes non-numeric entries into the score dict -- notably - # "alias", whose value is the task name -- so averaging everything - # that is not a stderr key fails with "resolved dtypes are not - # compatible with add.reduce" once one of them is a string. Keep - # only the numeric metrics. + # Average the metrics only. The score dict also carries + # bookkeeping: "alias" is the task name (a string, which made + # numpy fail with "resolved dtypes are not compatible with + # add.reduce"), and "samples" is the example count -- for gsm8k + # that is 1319, and averaging it in produced 44030.98 against a + # reference of 96, which *passed* the accuracy assertion. A wrong + # number that clears the threshold is worse than the crash it + # replaced, so both are excluded structurally rather than by name. numeric = [ acc for m, acc in scores.items() - if "_stderr" not in m and isinstance(acc, (int, float)) - and not isinstance(acc, bool) + if _is_metric_key(m) and "_stderr" not in m + and isinstance(acc, (int, float)) and not isinstance(acc, bool) ] if not numeric: raise ValueError( diff --git a/tests/unittest/evaluate/test_lm_eval_score_selection.py b/tests/unittest/evaluate/test_lm_eval_score_selection.py new file mode 100644 index 000000000000..16ac7da992e6 --- /dev/null +++ b/tests/unittest/evaluate/test_lm_eval_score_selection.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Which entries of an lm-eval result dict count as scores. + +The dict mixes metrics with bookkeeping. ``alias`` holds the task name and +``samples`` holds the example count, neither of which is an accuracy. Treating +them as scores produced two failures in sequence on the same GSM8K run: first a +``TypeError: the resolved dtypes are not compatible with add.reduce`` from +averaging the string, then -- after that was "fixed" by keeping only numeric +entries -- a reported accuracy of 44030.98, because gsm8k's 1319 samples had +been scaled to 131900 and averaged in. The second is the dangerous one: it +cleared a reference of 96.0 and the test passed. + +lm-eval keys metrics as ``","``, so the comma is the structural +discriminator. +""" + +import pytest + +from tensorrt_llm.evaluate.lm_eval import _is_metric_key + + +@pytest.mark.parametrize("key", [ + "exact_match,strict-match", + "exact_match,flexible-extract", + "acc,none", + "acc_norm,none", + "exact_match_stderr,strict-match", +]) +def test_metric_keys_are_recognized(key): + assert _is_metric_key(key) + + +@pytest.mark.parametrize("key", ["alias", "samples"]) +def test_bookkeeping_keys_are_rejected(key): + assert not _is_metric_key(key) + + +def test_gsm8k_shaped_dict_averages_to_the_metrics_only(): + """The exact dict shape that produced 44030.98.""" + scores = { + "alias": "gsm8k", + "exact_match,strict-match": 96.5125, + "exact_match_stderr,strict-match": 0.5053, + "exact_match,flexible-extract": 96.4367, + "exact_match_stderr,flexible-extract": 0.5106, + "samples": 131900, # 1319 examples, already scaled by 100 + } + kept = [ + v for k, v in scores.items() + if _is_metric_key(k) and "_stderr" not in k + ] + assert sorted(kept) == [96.4367, 96.5125] + assert 96.0 < sum(kept) / len(kept) < 97.0 From 7da09a26daecef0885ca4ac52981c059c5222902 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 01:47:59 -0700 Subject: [PATCH 046/177] [None][fix] Count the ragged gate's floor over populated batches The mid-loop gate used total decode steps, but trtllm-bench ramps concurrency from zero, so a floor of 64 was reached while two requests were in flight. A window can only differ *between* requests, so a batch of one or two is uniform by construction -- the gate then reported "never produced differing per-request windows" and killed the run, with the error itself reading "Broadcasting event-loop error to 2 pending request(s)". It took out both a sync-detector run and the ragged arm of the throughput A/B, neither of which had anything wrong with it. Count steps whose batch held at least two generation requests, which is where raggedness is even definable, and report the counter so a future false trip is visible in the summary rather than inferred. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 7 +++++- .../speculative/dspark_observability.py | 11 ++++++++++ .../test_dspark_graph_miss_attribution.py | 22 +++++++++++++++++++ 3 files changed, 39 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 48fdc2f3eb65..f9da3690a7dc 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2466,7 +2466,12 @@ def _maybe_assert_dspark_ragged_active(self, stats) -> None: if getattr(self, "_dspark_asserted", False): return min_steps = self._dspark_assert_min_steps - if min_steps is None or stats.steps_total < min_steps: + # Counted over steps with a real batch, not all steps. trtllm-bench + # ramps concurrency from zero, so step 64 can be a batch of two, where + # every window is trivially identical and the gate would fail a run + # whose scheduler had done nothing wrong -- observed as "Broadcasting + # event-loop error to 2 pending request(s)". + if min_steps is None or stats.steps_multi_request < min_steps: return self._dspark_asserted = True stats.assert_ragged_active(require_trim=True) diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index be408018b8f6..7a5560a52293 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -251,6 +251,10 @@ def __init__(self, *, mode: RaggedVerifyMode, max_draft_len: int): #: things: "peer_shape_mismatch" is one rank declining to go ragged and #: dragging the world eager, "key_not_captured" is the bucket grid being #: wrong, "peer_not_gen_only" is ordinary continuous batching. + #: Steps whose batch held at least two generation requests. The gate + #: counts these rather than all steps: trtllm-bench ramps concurrency, + #: so step 64 can still be a batch of two. + self.steps_multi_request = 0 self.graph_miss_reasons: Dict[str, int] = {} #: the actual graph key, kept only for the reasons where it localizes #: the bug. @@ -304,6 +308,12 @@ def record_step( fallback: short reason the step did not go ragged. """ self.steps_total += 1 + # Steps where raggedness is even definable. A window can only differ + # *between* requests, so a batch of one is uniform by construction -- + # counting it toward the gate's floor makes the gate fire during ramp-up + # and blame the scheduler for a batch that had not filled yet. + if int(num_gen_requests) >= 2: + self.steps_multi_request += 1 # What a no-trim step would have submitted: every request sends its # bonus token plus the full block. self.ceiling_tokens += int(num_gen_requests) * (1 + self.max_draft_len) @@ -428,6 +438,7 @@ def summary(self) -> Dict[str, object]: "mode": self.mode.value, "steps_total": self.steps_total, "steps_ragged": self.steps_ragged, + "steps_multi_request": self.steps_multi_request, "steps_uniform_windows": self.steps_uniform_windows, "steps_no_windows": self.steps_no_windows, "ceiling_tokens": self.ceiling_tokens, diff --git a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py index fe1a4d1cc2fb..11d4551922db 100644 --- a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py +++ b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py @@ -78,3 +78,25 @@ def test_unattributed_misses_are_a_broken_probe_not_a_clean_run(): def test_clean_run_passes(): _stats().assert_ragged_active(require_trim=True) + + +def test_ramp_up_steps_do_not_count_toward_the_gate_floor(): + """A batch of one is uniform by construction, not a scheduler failure. + + trtllm-bench ramps concurrency from zero, so a step-count floor of 64 can be + reached while only two requests are in flight. Every window is then + trivially identical, and gating on total steps failed a run whose scheduler + had done nothing wrong -- observed as "Broadcasting event-loop error to 2 + pending request(s)". The floor counts steps whose batch could actually be + ragged. + """ + stats = DSparkRaggedStats(mode=RaggedVerifyMode.COMPACT, max_draft_len=5) + for _ in range(64): + stats.record_step(num_gen_requests=1, verify_lens=[5]) + assert stats.steps_total == 64 + assert stats.steps_multi_request == 0, ( + "single-request steps must not count toward the gate floor") + + for _ in range(64): + stats.record_step(num_gen_requests=4, verify_lens=[1, 3, 5, 5]) + assert stats.steps_multi_request == 64 From 4622883a26f682802434b413b02d1a6d2186c546 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 04:56:34 -0700 Subject: [PATCH 047/177] [None][fix] Reclaim KV reserved for verify positions ragged trimmed away Ragged verification leaked KV on the target manager until the sequence demanded more blocks than max_seq_len allows, failing with "User-provided base page indices is too short". prepare_resources reserves for the FULL drafted block every generation step -- resource_manager.py adds `1 + max(draft_len, reserve)` tokens, and ragged still drafts the block in full by design. The rewind only returned py_rewind_len, which counts rejections *inside* the verify window. KVCacheManagerV2 gated the compensating term behind `if self.is_draft`, so the target manager never reclaimed the positions the scheduler trimmed away: uniform net = 6 - (draft_len - accepted) = 1 + accepted ok ragged net = 6 - (verified_len - accepted) = 5 + accepted leaks 4 The v1 resource_manager path already applies this term unconditionally and is correct; only the v2 path, which DSv4 uses, was affected. Localized by a controlled pair on identical data and settings, differing only in window: window=5 passed (418 steps, 6869.73 tok/s), window=1 failed in five seconds. That ruled out the ragged plumbing -- token-major presentation, layout publication, graph keys -- and pointed at the trimming half specifically. The instrumented error then gave the arithmetic: needs 9 blocks against max_seq_len=768 with tokens_per_block=128, i.e. claiming 1152 tokens for a 768-token ceiling, which is 4 tokens/step over the run's length. Uniform speculation is bit-for-bit unaffected: there verified_len equals draft_len, so the new term is identically zero. Verified by rerunning the exact case that failed in five seconds; it now completes. Also make the KV cache error state the numbers instead of just "too short" -- the bare message cost several twenty-minute runs to interpret. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 25 ++++++++++++++++--- .../kv_cache_manager_v2/_core/_kv_cache.py | 15 ++++++++++- 2 files changed, 36 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 7023745e7dd4..5be42012a515 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -3535,14 +3535,33 @@ def update_resources( if not kv_cache.is_active: continue rewind_len = req.py_rewind_len + # Tokens this step actually claims: rejections inside the verify + # window plus the accepted ones. + runtime_draft_len = req.py_rewind_len + req.py_num_accepted_draft_tokens if self.is_draft: - runtime_draft_len = req.py_rewind_len + req.py_num_accepted_draft_tokens # Dynamic-tree draft managers reserve K * max_draft_len slots, # which can exceed the tree's runtime draft width. Reclaim that # reserve slack together with rejected draft tokens; otherwise # it accumulates in the draft KV cache after every generation - # step. Target managers do not allocate this reserve slack. - rewind_len += max(self._kv_reserve_draft_tokens - runtime_draft_len, 0) + # step. + reserved = self._kv_reserve_draft_tokens + else: + # The target manager reserves for the FULL drafted block -- + # prepare_resources adds `max(draft_len, reserve)` tokens every + # generation step -- but ragged verification only ever writes + # `verified_len` of them, and py_rewind_len covers rejections + # *within that window* only. The positions the scheduler trimmed + # away were reserved and never reclaimed, leaking + # (draft_len - verified_len) tokens per step until the sequence + # demands more blocks than max_seq_len allows and the KV cache + # raises "User-provided base page indices is too short". + # + # Uniform speculation is unaffected: there verified_len == + # draft_len, so this term is zero and the arithmetic is exactly + # what it was. + reserved = max(get_draft_token_length(req), + self._kv_reserve_draft_tokens) + rewind_len += max(reserved - runtime_draft_len, 0) new_capacity = ( None if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index 303d79e9ee13..f445db5f122f 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -866,7 +866,20 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo indices.extend([BAD_PAGE_INDEX] * (new_num_blocks - old_num_blocks)) else: if len(indices) < new_num_blocks: - raise ValueError("User-provided base page indices is too short") + # The buffer is a slice of host_kv_cache_block_offsets + # and is therefore max_blocks_per_seq long, so this + # means the sequence wants more blocks than + # max_seq_len allows -- i.e. it is consuming KV + # faster than it emits tokens. The bare message gave + # no way to tell that from a genuinely oversized + # request, which cost several 20-minute runs. + raise ValueError( + f"User-provided base page indices is too short: " + f"buffer holds {len(indices)} blocks " + f"(max_blocks_per_seq) but the sequence now " + f"needs {new_num_blocks} (was {old_num_blocks}); " + f"num_blocks={self.num_blocks}, " + f"num_committed_tokens={self.num_committed_tokens}") stream_wait_events( self.cuda_stream, (s.ready_event for s in chain.from_iterable(slots)) From 5ddd6b57c72cfdcf5a8a0ecf7e82f398cde53330 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 05:37:10 -0700 Subject: [PATCH 048/177] [None][fix] Reject enable_ragged_verify without a profiled cost table Without a table the planner's budget degenerates to verify-all -- rightly, since a flat cost model makes every extra verify token look free -- so the ragged path runs and changes nothing. This was only a runtime warning, which means a 27-minute run "succeeds" having done nothing and the sole evidence is steps_ragged: 0, a counter nobody reads unless they already suspect the problem. It cost several runs here. Fail at construction. TLLM_DSPARK_FORCE_VERIFY_LENS stays as the exception, since it is the only way to exercise the trimming path at all: the planner declines to trim whenever acceptance is high, which is every workload measured on this model (GSM8K accept_len 3.3, synthetic prompts 4.1). Also correct a false claim I put in forced_verify_lens' docstring and repeated twice: that forcing windows is needed to break a circularity because producing a cost table requires a working ragged run. It does not. dspark_sps_profiler pins RaggedVerifyMode.STATIC and sweeps the verify length by rebuilding the engine per length, so table collection never enables ragged -- the table used here was produced that way. Signed-off-by: Lance Liao --- .../speculative/dspark_observability.py | 18 +++-- tensorrt_llm/llmapi/llm_args.py | 26 +++++++ .../test_dspark_ragged_requires_table.py | 67 +++++++++++++++++++ 3 files changed, 107 insertions(+), 4 deletions(-) create mode 100644 tests/unittest/llmapi/test_dspark_ragged_requires_table.py diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index 7a5560a52293..8f9971b48217 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -155,10 +155,20 @@ def forced_verify_lens(*, num_gen_requests: int, tiers: Sequence[int], sweep and the bias is gone. SGLang solves the same problem with a runtime ``dspark_force_budget_frac`` knob. - Either form also breaks a circularity: the planner refuses to trim without a - profiled cost table -- correctly, since a flat model makes every extra token - look free -- so a run without one keeps the ragged path dark no matter what - the config says, and producing a table would require a working run. + Why this exists at all: the planner refuses to trim without a profiled cost + table -- correctly, since a flat model makes every extra token look free -- + and even *with* one it declines whenever acceptance is high enough that + verifying the full block wins. On DeepSeek-V4-Pro-DSpark that is every + workload measured so far (GSM8K at accept_len 3.3, synthetic prompts at + 4.1), so a planner-driven run leaves the trimming path unexecuted and + cannot answer whether that path is correct. Forcing the windows separates + "is trimming implemented correctly" from "should we trim here"; it is what + exposed a KV leak that every planner-driven run had missed. + + NOT needed to produce a cost table. An earlier version of this note claimed + that, and it is wrong: dspark_sps_profiler pins RaggedVerifyMode.STATIC and + sweeps the verify length by rebuilding the engine per length, so table + collection never enables the ragged path. Neither form is a shortcut around the planner: both decide only *how many* drafted positions reach the target, exactly the quantity the planner decides. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 98e11a2b2868..3c69d8b1e48a 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2867,6 +2867,32 @@ def validate_ragged_verify(self): "enable_ragged_verify requires enable_confidence_scheduling=True: " "per-request verify lengths are chosen from confidence-head " "survival, which is only computed when scheduling is on") + if (self.enable_ragged_verify and not self.confidence_sps_table_path + and not self.confidence_sts_path + and not os.environ.get("TLLM_DSPARK_FORCE_VERIFY_LENS")): + # Without a profiled cost table the planner's budget degenerates to + # verify-all -- correctly, since a flat model makes every extra + # token look free -- so the ragged path silently never runs. That + # used to be a runtime warning, which means a 27-minute run + # "succeeds" having done nothing and the only evidence is a counter + # nobody reads. Fail at construction instead. + # + # TLLM_DSPARK_FORCE_VERIFY_LENS is the deliberate exception. Not + # because table collection needs it -- dspark_sps_profiler pins + # STATIC mode and never enables ragged -- but because the planner + # declines to trim whenever acceptance is high, which is every + # workload measured on this model. Without a way to impose windows, + # the trimming path cannot be exercised or tested at all. + raise ValueError( + "enable_ragged_verify=True requires a profiled cost table: pass " + "confidence_sps_table_path (produced by " + "`python -m tensorrt_llm._torch.speculative.dspark_sps_profiler`) " + "or confidence_sts_path. Without one the planner cannot compare " + "the cost of extra verify tokens against their expected " + "acceptance, so it declines to trim and every request receives " + "the full window -- the ragged path runs but changes nothing. " + "To bootstrap a table, set TLLM_DSPARK_FORCE_VERIFY_LENS to " + "impose windows without the planner.") return self @model_validator(mode="after") diff --git a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py new file mode 100644 index 000000000000..cb7cefb23be6 --- /dev/null +++ b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ragged verification must refuse to start without a profiled cost table. + +Without one the planner's budget degenerates to verify-all -- correctly, since +a flat cost model makes every extra verify token look free -- so the ragged +path runs and changes nothing. As a runtime warning that produced 27-minute +runs which "succeeded" having done nothing, detectable only from a counter +(``steps_ragged: 0``) that nobody reads unless they already suspect the +problem. Config errors belong at construction. + +``TLLM_DSPARK_FORCE_VERIFY_LENS`` is the deliberate exception: producing a cost +table requires a working ragged run, and that knob is what breaks the circle. +""" + +import os +from unittest.mock import patch + +import pytest + +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + + +def _cfg(**kwargs): + base = dict(max_draft_len=5, + speculative_model="/nonexistent/model", + enable_confidence_scheduling=True, + enable_ragged_verify=True) + base.update(kwargs) + return DSparkDecodingConfig(**base) + + +@patch.dict(os.environ, {}, clear=False) +def test_ragged_without_a_cost_table_is_rejected(): + os.environ.pop("TLLM_DSPARK_FORCE_VERIFY_LENS", None) + with pytest.raises(ValueError, match="requires a profiled cost table"): + _cfg() + + +@patch.dict(os.environ, {}, clear=False) +def test_ragged_with_an_sps_table_is_accepted(): + os.environ.pop("TLLM_DSPARK_FORCE_VERIFY_LENS", None) + cfg = _cfg(confidence_sps_table_path="/tmp/table.json") + assert cfg.enable_ragged_verify + + +@patch.dict(os.environ, {"TLLM_DSPARK_FORCE_VERIFY_LENS": "1"}, clear=False) +def test_forced_windows_bootstrap_without_a_table(): + """The escape hatch must stay open, or the table can never be produced.""" + cfg = _cfg() + assert cfg.enable_ragged_verify + + +@patch.dict(os.environ, {}, clear=False) +def test_uniform_confidence_scheduling_still_needs_no_table(): + """Only the ragged path is gated. + + Uniform tier selection degrades gracefully without a table -- it picks the + max tier, which is the pre-existing behaviour -- so requiring one there + would break working configurations. + """ + os.environ.pop("TLLM_DSPARK_FORCE_VERIFY_LENS", None) + cfg = DSparkDecodingConfig(max_draft_len=5, + speculative_model="/nonexistent/model", + enable_confidence_scheduling=True, + enable_ragged_verify=False) + assert not cfg.enable_ragged_verify From 468e8685215374a44322df53dcb677623cc2ec0a Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 06:47:19 -0700 Subject: [PATCH 049/177] [None][doc] Record what the P0 stage actually measured Appends section 10 to the task doc with the executed results, marked as taking precedence over the earlier text where they conflict: the earlier text is what was believed when the task was written, this is what the runs produced. The headline is that K1/K2 are resolved with no C++ change, GSM8K holds at 96.44/96.51 against a 96.21 uniform baseline, and there is no throughput win available on this checkpoint -- draft acceptance is ~90%, at which the fitted table says verify the full block at every batch size, so the planner correctly declines. That is a property of the checkpoint, not a defect, and it is worth stating plainly rather than leaving the next reader to rediscover it. Also records the memory cost nobody had measured (3x CUDA graphs, 9x attention workspace), the four real bugs found, and the methodology lessons -- chiefly that three of the four apparent failures were broken probes rather than broken features, so "no warnings" is not a result until the probe is shown to have run. Signed-off-by: Lance Liao --- docs/dspark_p0_task_prompt.md | 715 ++++++++++++++++++++++++++++++++++ 1 file changed, 715 insertions(+) create mode 100644 docs/dspark_p0_task_prompt.md diff --git a/docs/dspark_p0_task_prompt.md b/docs/dspark_p0_task_prompt.md new file mode 100644 index 000000000000..c19b30869f18 --- /dev/null +++ b/docs/dspark_p0_task_prompt.md @@ -0,0 +1,715 @@ +# 任务：让 DSpark confidence-head 调度在 DeepSeek-V4-Pro 上跑通（P0 阶段） + +> 本文是给 agent 的任务书。全部路径、commit、权重张量名均已在 `bia` 集群 login +> 节点上实地核实过（2026-08-01）。 + +--- + +## 0. 先读这个（不要跳过，也不要从零设计方案） + +本任务的 **唯一 source of truth** 是 PR NVIDIA/TensorRT-LLM#17056 里的 +`docs/dspark_confidence_schedule_goal.md`（399 行中文目标文档）。 + +它已经完成了全部前期分析：H0–H10 十个 host 侧阻断点、K1–K15 kernel 逐项审计、 +D1/D2/D3 三个方案对比、G1–G8 验收标准、§7 的 U1–U7 未知项、§8 的 9 个待拍板 +问题。**每一条断言都带 `file:line` 锚点。** + +你的工作是**执行它的 P0 阶段**，不是重新做分析。任何与该文档冲突的结论，必须先用 +`file:line` 证明文档错了，再改。 + +配套还有 `docs/dspark_ragged_verify_status.md`（388 行）。 + +### 拉取方式 + +本地 checkout：`/lustre/fsw/coreai_comparch_trtllm/laliao/repos/trtllm-dsv4/TensorRT-LLM` +（`origin` = `github.com/lancelly/TensorRT-LLM`，`upstream` = `github.com/NVIDIA/TensorRT-LLM`） + +```bash +cd /lustre/fsw/coreai_comparch_trtllm/laliao/repos/trtllm-dsv4/TensorRT-LLM +git fetch upstream refs/pull/17056/head:pr-17056 +git checkout pr-17056 # head = d8897046653331bca51242bf068b052568fca4ce +git show pr-17056:docs/dspark_confidence_schedule_goal.md +``` + +PR 的 4 个 commit： + +| commit | 内容 | +|---|---| +| `830afe4e58` | feat: Add DSpark confidence-head verification scheduling（26 files, +4683） | +| `8723506608` | fix: Fix host syncs, slot aliasing and cost model（13 files, +641） | +| `e1c28eeae2` | feat: Make DSpark ragged verification work on the DSv4 kernel path（21 files, +4172） | +| `d889704665` | doc: Scope the ragged-verify measurements to x86 DGX B300 | + +--- + +## 1. 三个必须先纠正的认知 + +### (a) confidence head 已经实现了 + +不要从零写。代码在： + +- `tensorrt_llm/_torch/models/dspark/heads.py`（`DSparkConfidenceHead`） +- `tensorrt_llm/_torch/models/dspark/{draft,attention}.py`、`modeling_dspark.py` +- `tensorrt_llm/_torch/speculative/dspark{,_planner,_verify,_ragged,_schedule,_observability,_sps_profiler}.py` +- 开关：`llm_args.py` 的 `enable_confidence_scheduling` / `enable_ragged_verify` + （均默认 `False`，`status='prototype'`） + +任务是**让它正确、且 CUDA-graph 兼容地跑起来**，不是新增功能。 + +### (b) goal doc §1.3 描述的「11 个未提交的工作区改动」现在已经是 commit `e1c28eeae2` + +文件清单与该 commit 逐项吻合。所以文档里标为「新引入的 P0 bug」的 **H3** 和 **H4** +现在是**已经进了 PR head 的 bug**： + +- **H3**：`dsa.py:1673-1695`，`row_kv_lens.add_(expanded.to(torch.int32))` —— CUDA + tensor 加 CPU tensor（`kv_lens = cached_token_lens + seq_lens_kv` 两项都在 CPU）， + 首个 ragged step 必抛异常。修法：把逐行修正量折进 `row_kv_lens_host` 做**一次** + H2D，复用相邻 `kv_lens_expanded_host` 已有模式（`dsa.py:1590-1592`）。 +- **H4**：`row_kv_lens_cuda` 从不被 `on_update_kv_lens` 刷新。 + +另外注意：本地 checkout 的工作区是干净的，goal doc §1.3 描述的是**另一个**工作区 +的状态，不要去找那 11 个「未提交」的文件。 + +### (c) PR 的外部引用有误（goal doc §7.2 已指出） + +找参考实现时用正确的地址： + +- **SGLang 没有 `dspark_schedule.py`**。真实布局是 + `python/sglang/srt/speculative/ragged_verify.py` + + `python/sglang/srt/speculative/dspark_components/*`。 + 三模式 `static` / `cap-accept` / `compact` 出自这里。 +- **vLLM 是 PR vllm-project/vllm#47808**。关键参考点： + - `vllm/v1/attention/backends/mla/indexer.py` 里关于 `q.view(batch, next_n, ...)` + 对非均匀 decode 长度**静默算错**的注释 + - 它如何使用 DeepGEMM 的 varlen `indices` 参数（TRT-LLM 目前用 "expanded + buffers" 展平达到同样形状，从不传 `indices`；见 goal doc §7.3） +- `align_verify_tokens_to_graph_tier` 这个符号全网不存在。SGLang 的实际原语是 + `round_up_grid()` / `compute_target_verify_graph_key()`。 + +--- + +## 2. 目标配置 + +- **模型**：DeepSeek-V4-Pro-DSpark + （target = DSv4 Pro + DSA 稀疏注意力；drafter = DSpark block drafter + confidence head） +- **并行**：DEP8 = TP8 + `moe_expert_parallel_size=8` + `enable_attention_dp=True` +- **投机**：`max_draft_len=5` → block K=5 → 均匀 verify 窗口 `next_n=6` +- **硬件**：8×Blackwell（B200/B300，SM100） + +### 2.1 权重（已实地确认存在且可读） + +``` +/lustre/share/coreai_dlalgo_ci/artifacts/model/deepseek-ai_deepseek-v4-pro-dspark/hf/hf-cd80265_orig +``` + +832 GB，66 shards，149782 个 tensor。路径 world-readable。 + +**决定性证据 —— confidence head 权重就在里面：** + +``` +mtp.2.confidence_head.proj.weight BF16 [1, 7680] (model-00066-of-00066.safetensors) +``` + +checkpoint 结构：`mtp.0` / `mtp.1` / `mtp.2` 三个 DSpark block，只有 `mtp.2` 带 +confidence head。config 关键字段： + +| key | 值 | +|---|---| +| `dspark_block_size` | **5**（与 `max_draft_len=5` 对上） | +| `dspark_target_layer_ids` | `[58, 59, 60]` | +| `dspark_markov_rank` | 512 | +| `dspark_noise_token_id` | 128799 | +| `expert_dtype` | fp4 | +| `index_topk` | **1024**（普通 dsv4-pro 是 512） | +| `num_nextn_predict_layers` | 1 | + +**全 index grep `sts|temperature|calib` 零命中** —— 印证 goal doc 的 U4：STS 校准表 +既不在 checkpoint 也不在 repo，必须自己 profile 产出。 + +### 2.2 必须先建软链 + +测试里是 `MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4-Pro-DSpark"`，而实际目录名 +是 `hf-cd80265_orig` 且嵌了两层： + +```bash +mkdir -p ~/llm-models +ln -s /lustre/share/coreai_dlalgo_ci/artifacts/model/deepseek-ai_deepseek-v4-pro-dspark/hf/hf-cd80265_orig \ + ~/llm-models/DeepSeek-V4-Pro-DSpark +export LLM_MODELS_ROOT=~/llm-models +``` + +`MODEL_PATH` 在 class body（import 时）求值，`LLM_MODELS_ROOT` 必须在 pytest 启动 +**前**导出，否则整个 module 收集失败。 + +容器里 `~` 可能不是同一个 home，建议软链建在 `/lustre/fsw/coreai_comparch_trtllm/laliao/llm-models/` +下并用绝对路径导出。 + +### 2.3 用不到但容易认错的几个目录 + +| 路径 | 是什么 | 为什么不是你要的 | +|---|---|---| +| `/lustre/share/coreai_comparch_trtllm/dsv4-pro` | 806 GB，纯 DSv4-Pro | 只有 `mtp.0`，**无 confidence head** | +| `/lustre/share/coreai_comparch_trtllm/dsv4-flash` | 149 GB | DSv4-Flash，非本轮范围 | +| `/lustre/share/coreai_comparch_aarwlt/hf_repos/deepseek-ai/DeepSeek-V4-{Pro,Flash}` | HF 原始副本 | 同上，无 DSpark | +| `/lustre/share/coreai_comparch_trtllm/llm-models/` | **空目录** | 不要指向这里 | + +### 2.4 ⚠️ 与现有测试的冲突（必须先解决） + +现有测试 `TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling` +（`tests/integration/defs/accuracy/test_llm_api_pytorch.py`）用的是 +**`disable_overlap_scheduler=True`**。 + +但本任务**要求 overlap scheduler 开启**。这不是配置项调换，而是把 goal doc 的 +**H0 从边缘情况变成第一优先级**： + +> **H0**：`disable_overlap_scheduler=True` 时，ragged token 布局**根本不会被产出**。 +> `_prepare_tp_inputs` 的非 overlap 分支（`model_engine.py:4606`）条件含 +> `next_draft_tokens_device is None`，该分支 `sequence_lengths.append(1 + num_draft_tokens)`， +> **完全不看 `py_verify_len`**（`:4616,4625`）。后果：`spec_metadata.verify_lens` / +> `qo_indptr` / `total_verify_tokens` 和 graph key 都在描述一个 ragged 布局，而真实 +> `input_ids`/`seq_lens` 是均匀的 `1+top_tier`，`_accept_draft_tokens` 随后按 +> `total_verify_tokens` 切片（`interface.py:1436-1441`）→ **逐请求 token 错位**。 + +goal doc §3.3 明确要求 **H0 与 H5 必须同时修**（H5：CUDA graph capture 的 token 轴 +来自 `key[1]` 而不看 bucket，`cuda_graph_runner.py:446-453`），否则 graph 只会更糟。 + +**默认做法**：新增一个 `disable_overlap_scheduler=False` 的 test case，保留原 case +作对照。若要直接改原 case，先确认。 + +### 2.5 已经拍板的两个决策（goal doc §8 的 Q2 / Q8 已关闭） + +goal doc §8 把这两条列为待拍板，**现在已定，不要再当开放问题讨论**。 + +#### 决策一：**ragged path 是第一优先级**（关闭 Q2） + +Q2 原本问的是：本轮交付「均匀 tier 阶梯」还是「ragged」。两者的区别是—— + +| | **均匀 tier 阶梯（D3 Stage A）** | **ragged（D1/D2）** | +|---|---|---| +| 语义 | 每个 decode step **整个 batch 用同一个 verify 长度**，从 tier ladder `[1,3,5]` 里选一个 | **每个请求各自的 verify 窗口**，同一步里可以是 `[6,3,5,2,...]` | +| kernel | 0 改动。四个 uniform kernel 在 `num_gen_tokens = n·(K+1)` 时逐一精确正确 | 需要 kernel 支持每 token 行向量；K4/K9 已完成 | +| 收益 | 拿满 SPS(B) 项，**损失 τ 项** | 完整 Θ = τ·SPS(B) | +| 风险 | 中。无 K2 假设 | **高**。押在 K2（trtllm-gen sparse-MLA generation）这个读码无法证明的行为上 | + +**已定：做 ragged。** 因此： + +1. goal doc §5.1 里划给 **P2 的 H6 / H7 / H9 / H10 全部上提到 P0**（见 §3 表二）。 +2. **K2 从「以后再说」变成本轮最大风险**。它的确切阻断点已核实： + `cpp/tensorrt_llm/thop/attentionOp.cpp:920-921` + ```cpp + TLLM_CHECK_WITH_INFO(num_tokens % num_seqs == 0, + "seq_len should be same for all generation requests, num_tokens=%d, num_seqs=%d", + num_tokens, num_seqs); + int32_t const input_seq_length = num_tokens / num_seqs; + ``` + 加上 `mMaxSeqLenQ = acc_q_len / batch_beam`（`attentionOp.cpp:1161,1167`），以及 + `cumSeqLensQPtr` 字段虽存在（`fmhaRunnerParams.h:279`）但**在 generation 分支从不 + 赋值**。D1 和 D2 都靠把 generation 半边**呈现成 `num_gen_tokens` 个「长度为 1 的 + 序列」**来让这个检查平凡通过（`seq_len=1`），二者押的是同一个赌注。 +3. **因此 B2 微基准从 P1 上提为 P0 的 gating 测量**，而且要**最先跑**：它单卡、 + 分钟级、**不需要模型权重**，却决定整条 ragged 路线可不可行。见 §3 表三。 + +#### 决策二：**DSv4-Pro 走 KVCacheManager v2**（关闭 Q8 / U6） + +已确认。后果： + +- **H6 上提为 P0**（原 P1）。 +- 修法不变且与 manager 版本无关：`SpecSamplerBase._verified_len` **实时**读请求上的 + `py_verify_len`（`spec_sampler_base.py:195-196`），而第 N 步的 + `_handle_dynamic_draft_len`（`py_executor.py:4652`）先于第 N-1 步的 + `_update_requests`（`py_executor.py:4746`）执行 → 读到被下一步覆盖过的值。 + `runtime_draft_len` 已经为此快照进 `SampleStateSpec`（`spec_sampler_base.py:218`）， + **把 `py_verify_len` 一并快照进去即可**。 +- ⚠️ **一处需要你自己核实的措辞**：goal doc 说 v2 的补偿项被 `if self.is_draft:` + 门控、「target manager 无补偿」所以会泄漏 KV。但该分支紧邻的注释明确写着 + *"Target managers do not allocate this reserve slack."* + （`kv_cache_manager_v2.py:3538-3546`）——即 target 本来就没分配这块 slack，不补偿 + 是**正确的**。所以「泄漏」这个结论还差一步论证。 + **`py_verify_len` 读到陈旧值这个 bug 本身独立成立**，先按这个修；「是否额外泄漏」 + 作为一个独立问题实测确认，标注「待验证」。 + +--- + +## 3. 执行范围：P0 阶段 + +因为 ragged 是第一优先级（§2.5 决策一），P0 的范围比 goal doc §5.1 原本划的要大： +它 = 原 P0（地基）+ 原 P2 的 ragged 必需项 + 原 P1 的 B2 测量。 + +**执行顺序有讲究，不要并行乱做：** + +``` +表三 B2 微基准（单卡，无需权重） ──┐ + ├──► 表一 地基（H0–H5, K8, sync, H8, K5/K6/K13） + │ └──► 均匀 tier 路径跑通 = 中间里程碑 + └──► 表二 ragged 必需项（H6, H7, H9, H10） + └──► ragged 路径跑通 = 交付 +``` + +- **B2 最先跑**，它可能直接否掉整条 ragged 路线（§2.5 决策一第 3 点）。 +- **表一必须先于表二**。表一全部是「无论 ragged 与否都要修」的东西，其中 H1 只有 + 约 3 行 Python 却能完整解释唯一一次实测 hang。在均匀 tier 路径跑通之前调 ragged， + 等于在两层未知上 debug。 +- **均匀 tier 路径不是可以砍掉的中间产物**：G2 的逐 token 等价、A2 的 `cap-accept` + 差分都要拿它当参照系（§6）。 + +### 表一：地基（全部 host 侧，独立于 kernel 方案选择） + +| 项 | 内容 | 锚点 | +|---|---|---| +| **H1** | DSA 展开 stride 用 runtime tier 而非静态 max。**约 3 行 Python**（`stride = spec_metadata.runtime_tokens_per_gen_step`）。⚠️ **buffer 容量必须保持静态 max**，只让 stride 变，否则 capture 期重分配 → 第二类 hang | `dsa.py:662,696,1401`、`model_engine.py:342` | +| **H2** | `_attach_ragged_verify_layout` 移到 `attn_metadata.prepare()` **之前**（它只需要 `generation_requests`，已在作用域内） | `model_engine.py:5326` vs `:5412` | +| **H3** | `row_kv_lens` device 不匹配。修法：折进 `row_kv_lens_host` + 单次 H2D | `dsa.py:1673-1695` | +| **H4** | `on_update_kv_lens` 补刷 `row_kv_lens_cuda` | `dsa.py:1608`、`:902-977` | +| **H0+H5** | **必须同时改**。dummy 分支用 `get_request_tokens_per_gen_step`；`capture()` 吃 bucket | `model_engine.py:4606`、`cuda_graph_runner.py:446-453` | +| **K8** | `prepare_scheduler_metadata` 改走 `expand_per_gen_token`（现在硬编码 `num_generations * (1 + max_draft_tokens)`） | `dsa.py:2329-2330` | +| **sync** | `_update_target_input_tensors` 的 `repeat_interleave` 补 `output_size=`（旁边 `model_engine.py:5101-5109` 已做对）；`apply_sts` 缓存 host 副本 | `model_engine.py:3910-3919`、`heads.py:280-283` | +| **H8** | config 层校验 `cuda_graph_config.enable_padding=True`（默认 False，`fit_ragged_verify_lens` 却假定已开） | `model_engine.py:3644-3648`、`llm_args.py:180-184` | +| **K5/K6/K13** | 显式 **config-reject** 而非静默降级：`enable_heuristic_topk` / `use_cute_dsl_topk` 与 ragged 结构性互斥（kernel 层已有 TLLM_CHECK，Python config 校验欠缺）；K13 加显式 assert 而非依赖 SM 判断 | `indexerTopK.cu:918-922`、`trtllm.py:999-1001` | + +**表一交付（中间里程碑）**：均匀 tier 路径（`enable_confidence_scheduling=True`、 +`enable_ragged_verify=False`）在 8×B200/B300 上跑通 GSM8K，hang 消失。 + +### 表二：ragged 必需项（原 goal doc P2，因决策一上提到 P0） + +| 项 | 内容 | 锚点 | +|---|---|---| +| **H6** | `py_verify_len` 快照进 `SampleStateSpec`，不要实时读。**KVCacheManager v2 已确认**，见 §2.5 决策二（含一处需自行核实的措辞） | `spec_sampler_base.py:195-196,218`、`py_executor.py:4652,4746`、`kv_cache_manager_v2.py:3538-3546` | +| **H7** | ragged batch 上 rejection sampling 被 fail-closed guard **静默关掉**：`_rejection_buffers_valid` 要求 `logits_rows >= num_contexts + num_gens*(draft_len+1)`，而 ragged 下 target 只发 `num_contexts + total_verify_tokens` 行 —— **调度真的在 trim 时该式必然不成立**。于是 PR 在 rejection 路径里精心加的 ragged 分支（`interface.py:1793-1806`）成了死代码。**GSM8K temperature=0 检测不到这一点。** 要么放宽 guard 成 `num_contexts + total_verify_tokens`，要么在 config 层显式拒绝「ragged + 非贪婪」。**静默降级是最差的选项** | `interface.py:1707-1710,1793-1806` | +| **H9** | **rank-local 早退门控了条件性的 `tp_allgather`，是 hang 风险。** `decide_verify_lens` 的 `all_rank_max` 在早退之后，`peer_stats` allgather 被 `if ragged_lens is not None` 门控，fallback 分支又发一次 → 走 ragged 的 rank 发 2 次集合通信、fallback 的 rank 发 1 或 3 次 → **死锁**。且早退条件是 rank-local 且**时序相关**（`_copy_event.query()`）。修法：每个 rank **无条件**发出相同序列的集合通信，再基于归约结果分支 | `dspark_verify.py:234-241,256-259,132-133`、`py_executor.py:3258-3261` | +| **H10** | DSpark drafter 用固定 stride 取 target hidden state：`base = gen_start + arange_g * (runtime_draft_len + 1)`，而 ragged 下真实偏移是 `spec_metadata.qo_indptr[r]`（`dspark.py` 全文无 `qo_indptr` 引用）。**静默：只是接受率崩塌，不报错**。注意 `8723506608` 已把它从 `max_draft_len+1` 修成 `runtime_draft_len+1`，**均匀路径是对的**，只有 ragged 错 | `dspark.py:600,624` | +| **§4.4** | **ADP 资格门必须同时比较 bucket**。现在只比 `is_all_gen_only` 和 batch size，从不比 token 数；`all_rank_num_tokens` 是 host list、会被**烤进 captured graph**，所以「任一 rank 退回均匀路径」造成的不匹配在 replay 时**无法被察觉** | `cuda_graph_runner.py:362-372`、`model_engine.py:2987-2998` | + +**表二交付**：ragged 路径（两个开关都开）GSM8K + 逐 token 等价（§6）。 + +### 表三：B2 微基准（最先跑，单卡，无需模型权重） + +**这是决定整条 ragged 路线可不可行的那个数字**（goal doc §7.1 U1）： + +> 直接调 `thop.attention`，在**相同总 token、相同逐 token 稀疏索引集**下，比较 +> `(batch=B, s_q=tier+1)` 与 `(batch=B*(tier+1), s_q=1)` 两种呈现的**正确性与吞吐**。 + +- 读不出来（`selectMlaGenerationKernel` 在预编译 cubin 里），但一张 B200/B300 上 + 直接测就行，分钟级。 +- 关注点：DSv4-Pro 在 DEP8 下每 rank 保有全部 128 个 Q head，BMM1 的 M 维从 768 降到 + 128；由于稀疏索引集逐行不同，本来也没有跨 6 行共享 KV tile 的机会 —— 但这是 + **推断，不是测量**。 +- **若 token-major 显著变慢或不正确**：立刻停下来报告，不要硬推 ragged。 + goal doc §5.1 的兜底是 D3 Stage A（均匀 tier 阶梯），它拿满 SPS(B) 项、无 K2 假设、 + 还能把 graph 数从 102 压到 ~54。 + +顺带把 goal doc §6.4 的另两个也跑了（都是分钟级）： + +- **B1**：修完 H1 后重跑均匀 tier capture，确认 hang 消失。**带 `TORCH_LOGS=recompiles`** + —— 这是区分 §7 三个 hang 候选根因的最短路径。 +- **B3**：expanded 布局下 paged-MQA-logits 的吞吐（展平路径每请求重读 6 次 KV 而非 + 2 次，在关键路径上，必须实测）。 + +### 仍然不做的 + +- D1 vs D2 的最终选型：等 B2 出数再定。工作区已沿 D2 走了一半（K4/K9 已实现且设计 + 正确 —— 可选参数、非 ragged 路径 bit-identical、TORCH_CHECK 替换而非放宽）， + **这部分不要丢**，它对任何 ragged 方案都是必需的，对 D3 也无害（不传参数即退化）。 +- STS 校准表与 SPS cost table 的产出流程（goal doc G8 / U4 / U5）—— 见 §8 交付物第 5 条。 + +--- + +## 4. 编译 + +**不要手工敲 `build_wheel.py`。**用仓库里现成的脚本： + +```bash +/lustre/fsw/coreai_comparch_trtllm/laliao/repos/trtllm-dsv4/TensorRT-LLM/tmp/compile.sh +``` + +它做的事（`tmp/compile.slurm`）： + +- `sbatch --partition=batch --nodes=1`，`--time=4:00:00`， + `--account=coreai_comparch_trtllm` +- 容器镜像从 `jenkins/current_image_tags.properties` 的 `LLM_DOCKER_IMAGE` 读， + URL 里的 `urm.nvidia.com/` 被换成 `urm.nvidia.com#` 供 enroot 使用 +- mount：`/lustre:/lustre` + repo → `/code/tensorrt_llm`，workdir `/code/tensorrt_llm` +- 构建目录按 **git short hash** 命名：`/lustre/fsw/coreai_comparch_trtllm/laliao/build//` +- 实际命令：`git clean -fdx cpp/` → 删 `tensorrt_llm/flash_mla` → + `python scripts/build_wheel.py --clean -G Ninja -a '103-real;'` → `pip install` +- 产物：`build//trtllm.sqsh`（`--container-save`）+ `build//tensorrt_llm-*.whl` +- 日志：`tmp/logs/compile_.{log,err}` + +**注意事项：** + +1. **`-a '103-real;'` 只编 SM103（B300）。** 如果你拿到的是 B200（SM100），必须改成 + `'100-real;'` 或 `'100-real;103-real;'`，否则 kernel 加载会失败。 +2. **构建目录以 git hash 为键，并且会先 `rm -r` 掉旧的。** 两个同 hash 的 compile + job 并发跑会互相清目录 —— 见 §5 的竞态提交规则。 +3. Python-only 改动（本任务 P0 清单里除 K5/K6/K13 的 C++ assert 外**全部是 Python**） + 不需要重编。直接用已有的 `trtllm.sqsh` + mount repo 到 `/code/tensorrt_llm` 即可， + 因为 `PYTHONPATH=/code/tensorrt_llm` 会让改动立即生效（见 `tmp/serve_*.slurm` 的写法）。 + **先判断要不要编，再编。** +4. 跑测试时用 `CONT="/lustre/fsw/coreai_comparch_trtllm/laliao/build//trtllm.sqsh"`， + 参照 `tmp/serve_dsv4_bs8_relaxed_constraint.slurm` 的 srun 写法 + （`--mpi=pmix -N 1 --ntasks=8 --ntasks-per-node=8`）。 + +--- + +## 5. 申请节点：所有账号一起拍，没排到的取消 + +### 5.1 集群事实（实测） + +| 项 | 值 | +|---|---| +| 集群 | `bia`（login-bia02.bia.clusters.nvidia.com），Slurm 25.11.2 | +| 节点 | 128 × B300，`bia0001-0128`，x86_64，2 socket × 64 core（256 线程），~2.0 TiB 内存 | +| **每节点 8 GPU** | 由 `tmp/serve_*.slurm` 的 `--ntasks-per-node=8` 佐证；Slurm **未配 GRES**，不能按卡申请，`OverSubscribe=EXCLUSIVE` 整节点独占 | +| 分区 | `batch`（默认，默认 2h / 最大 **5h**，PriorityTier 2）
`backfill`（默认 2h / 最大 **8h**，PriorityTier 1，**`PreemptMode=CANCEL` 会被直接取消**，GraceTime 600s） | +| 你的账号 | `coreai_comparch_trtllm`、`coreai_comparch_aarwlt` | +| 配额 | 单用户 **104 节点**（QOS `MaxTRESPU`），最多 1000 个排队作业 | +| 当前负载 | 112/128 alloc、0 idle、303 个 pending —— **要排队** | + +DEP8 需要 8 GPU，**1 个节点即可**。 + +### 5.2 竞态提交 → 先跑的赢，其余取消 + +四个组合全拍上：`{coreai_comparch_trtllm, coreai_comparch_aarwlt} × {batch, backfill}`。 +`sbatch` 命令行的 `--account` / `--partition` **会覆盖**脚本内的 `#SBATCH` 指令。 + +**推荐做法：竞态的是「分配」，不是「负载」。** 先抢一个空 allocation，抢到之后再往 +里塞真正的活。这样不存在两个 job 同时开跑把对方的 build 目录清掉的问题。 + +```bash +#!/bin/bash +# race_alloc.sh —— 抢一个 1 节点 allocation，先到先得，其余取消 +set -uo pipefail + +ACCOUNTS=(coreai_comparch_trtllm coreai_comparch_aarwlt) +PARTITIONS=(batch backfill) +TIME_LIMIT=5:00:00 # batch 上限 5h；backfill 可以到 8h +LOGDIR=/lustre/fsw/coreai_comparch_trtllm/laliao/repos/trtllm-dsv4/TensorRT-LLM/tmp/logs +mkdir -p "$LOGDIR" + +JOBS=() +for A in "${ACCOUNTS[@]}"; do + for P in "${PARTITIONS[@]}"; do + T="$TIME_LIMIT" + [[ "$P" == backfill ]] && T=8:00:00 + JID=$(sbatch --parsable \ + --account="$A" --partition="$P" \ + --nodes=1 --exclusive --time="$T" \ + --job-name="race-${A}-${P}" \ + --output="$LOGDIR/race_%j.log" --error="$LOGDIR/race_%j.err" \ + --wrap="sleep infinity") + echo "submitted $JID account=$A partition=$P time=$T" + JOBS+=("$JID") + done +done + +# 轮询，谁先 RUNNING 谁赢 +WINNER="" +while [[ -z "$WINNER" ]]; do + for J in "${JOBS[@]}"; do + ST=$(squeue -h -j "$J" -o "%T" 2>/dev/null) + if [[ "$ST" == RUNNING ]]; then WINNER="$J"; break; fi + done + [[ -z "$WINNER" ]] && sleep 10 +done + +NODE=$(squeue -h -j "$WINNER" -o "%N") +echo "WINNER=$WINNER on $NODE" + +# 只取消这批里没排到的，不要动别人/别的作业 +for J in "${JOBS[@]}"; do + [[ "$J" == "$WINNER" ]] && continue + scancel "$J" && echo "cancelled $J" +done + +echo "export WINNER_JOBID=$WINNER" +``` + +拿到 allocation 之后，把活塞进去（`--overlap` 是必须的，因为 `sleep infinity` 那个 +step 已经占着资源）： + +```bash +srun --jobid=$WINNER_JOBID --overlap \ + --container-image="$CONT" \ + --container-mounts="/lustre:/lustre,$REPO:/code/tensorrt_llm" \ + --container-workdir=/code/tensorrt_llm \ + -N1 --ntasks=1 --pty bash # 交互式排查 + +# 或者跑测试（8 卡） +srun --jobid=$WINNER_JOBID --overlap --mpi=pmix \ + --container-image="$CONT" \ + --container-mounts="/lustre:/lustre,$REPO:/code/tensorrt_llm" \ + --container-workdir=/code/tensorrt_llm \ + -N1 --ntasks=8 --ntasks-per-node=8 \ + bash -c "..." +``` + +活干完后**记得 `scancel $WINNER_JOBID`**，`sleep infinity` 不会自己退。 + +### 5.3 直接竞态提交真实负载（更简单，但有坑） + +如果你就想直接把 `compile.sh` 那种脚本拍四份： + +```bash +for A in coreai_comparch_trtllm coreai_comparch_aarwlt; do + for P in batch backfill; do + sbatch --parsable --account="$A" --partition="$P" \ + --job-name="compile-${A}-${P}" tmp/compile.slurm + done +done +``` + +⚠️ **`compile.slurm` 的 `BUILD_DIR` 只按 git hash 命名，且开头就 `rm -r`。** 四个 job +在同一 commit 上跑，谁先起来谁就把别人的目录删了。所以： + +- 轮询间隔要短（≤5s），**一发现有 job 离开 PENDING 就立刻取消其余的** +- 如果不幸有两个同时起来了，取消其一之后**必须重跑赢家**，不能假定它的 build 是完整的 +- 更稳的做法就是 §5.2：竞态 allocation，然后只在赢家里跑一次编译 + +### 5.4 收尾 + +```bash +squeue -u $USER # 看还剩什么 +scancel [ ...] # 精确取消 +``` + +**不要用 `scancel -u $USER`** 无差别清场 —— 可能会杀掉正在跑的别的活。 + +--- + +## 6. 验收（goal doc §6） + +> **所有测试都必须在 Slurm 分配的 GPU 节点上跑。login node 不可用于任何 GPU 测试。** + +**必须做的，不只是 GSM8K：** + +- **G3 GSM8K**：`score >= 96.0`。但 goal doc §6.1 明确说这个门**不够** —— 统计阈值 + 允许 42/1319 题翻转，而 `>=96.0` 对实测 96.475 只有 6.3 题余量，既太松又太紧。 + **且当前没有活的 baseline 可比**（兄弟测试 `test_gsm8k_dep8_megamoe_deepgemm` 已被 + `waives.txt:184` 以 nvbugs/6506920 跳过）。 +- **G2 逐 token 等价（这才是真正的门）**：temperature=0 下，`static` 与开启调度后 + 逐 token 输出一致。贪婪投机解码定义上 lossless，任何分歧都是 bug。 + ⚠️ 诚实的形式：跨不同 graph 形状的严格 bitwise 相等 TRT-LLM 并不保证（MoE/attention + 规约顺序会变），所以允许极少数 prompt 分歧并设硬上限。 +- **A2 `cap-accept` 差分（ragged 成为第一优先级后，这条从可选变成强烈建议）**： + 实现 SGLang 的三模式 `static` / `cap-accept` / `compact`。`cap-accept` 跑**均匀 + kernel 路径**（`next_n=6`）但只提交每请求窗口内的 token，输出应与 `compact` + **逐 token 相同**。于是： + - `cap-accept` ≠ `compact` ⇒ **必然是 ragged kernel bug** + - `static` ≠ `cap-accept` ⇒ **必然是调度/接受逻辑 bug** + + 这是把「GSM8K 掉分了但不知道为什么」变成可定位诊断的唯一手段。 + `dspark_observability.py` 已有 mode 枚举，`cap-accept` 目前显式 `NotImplementedError` + （`dspark.py:471-478`）。 +- **A3 布局一致性断言**（一条同时抓 H0/H5/H8）： + ``` + sum(1 + py_verify_len over padded generation_requests) + == attn_metadata.num_tokens - num_ctx_tokens + == spec_metadata.total_verify_tokens + == key[bucket] + ``` +- **A4 时序断言**（抓 H2/H4）：断言 `attn_metadata.ragged_verify_lens` 在 `prepare()` + 那一刻是**本步**的（而非上一步）。 +- **A6 调度确实生效**（抓静默退化）：断言 `planner.stats` 的 `fallback_flat_cost` / + `fallback_no_snapshot` / `fallback_no_confidence` / `fallback_short_snapshot` + （`dspark_verify.py:86-93`）显示 trim 路径被走到。 + **没有 profiled SPS cost table 时 `_decide_local` 无条件返回 `max_tier` + （`dspark_verify.py:268-270`），feature 在构造上无法 trim** —— 一个看起来完全正常的 + run 收益可能是零。 +- **G4 CUDA graph**：稳态 decode 100% 命中 replay；capture 期无 D2H sync、无新分配； + graph key 与实际 token layout 逐字段一致。 +- **G5 ADP 一致性**：goal doc §4.4 要求 **ADP 资格门必须同时比较 bucket**（现在只比 + `is_all_gen_only` 和 batch size，从不比 token 数）。 +- **G6 零回归**：两个开关关闭时与 `main` bit-identical。 +- **单测**：`tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py` + 按每请求 q 长度参数化 —— goal doc §6.3 称这是**性价比最高的一个测试**，单卡分钟级 + 就能抓住 tier-3 那类 hang（该文件 `:415`/`:473` 现在硬编码 + `seq_len_q = fused_q.shape[0] // num_requests`，那正是被测假设本身）。 + +集成测试命令（无需 `mpirun`，单 pytest 进程经 mpi4py.futures 起 8 worker）： + +```bash +export LLM_MODELS_ROOT=/lustre/fsw/coreai_comparch_trtllm/laliao/llm-models +python -m pytest "tests/integration/defs/accuracy/test_llm_api_pytorch.py::\ +TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling" -vs +``` + +**验收必须覆盖至少两个并发点**（goal doc §6.5）：调度只在高并发下才真的 trim，单一 +低并发 run 是空跑。vLLM #47808 报的是并发 16 和 64。 + +--- + +## 7. 已知 hang（必须区分三个候选根因，不要认定其中之一） + +唯一一次端到端尝试失败在 tier-3 的 CUDA graph capture： + +``` +capturing 102 graphs (34 batch sizes x 3 draft-length tiers [1, 3, 5]) +warmup batch=128, draft_len=5 ← 完成 +warmup batch=128, draft_len=3 ← never returns TEST_EXIT=124 +``` + +三个**独立**候选（goal doc §7.1 U2）： + +1. **H1 stride 不匹配** —— tier 5 时 768==768 通过，tier 3 时 schedule 描述 768 行而 + logits 只有 512 行。完整自洽。 +2. **capture 期 attention workspace 重分配** —— hang 前紧邻的日志是 + `Attention workspace ... increase from 0 bytes to 168493568 bytes`； + `model_engine.py:1789-1806` 明确记录这是真实 hazard，且 Case 2b **没有** Case 2 的 + `max_spec_graph` 规避（`:1810-1813`）。 +3. **torch.compile / max-autotune 在 capture 区内**（goal doc §4.5）—— + `deepseek_v4.py` 三个 `@maybe_compile(dynamic=True, max-autotune)` 函数经 + `on_update_kv_lens` 在 `CUDAGraphRunner.capture` 录制区内被调用； + `_compute_gen_compressed_position_ids` 把 `num_gen_tokens_per_seq` 当 Python int 用 + （`:979-983`）构成 guard → **每个 tier 一次 Inductor 编译，发生在 graph capture + 期间**，而 max-autotune 需要同步 benchmark。 + +**先修 H1 再重跑，并带 `TORCH_LOGS=recompiles`** —— 这是区分三者的最短路径。 + +goal doc 引用的 `.dspark-logs/gsm8k.log` 和 `hang-stack-299685.txt` 不在本 repo 里， +是写文档那位的工作区产物。不要去找，重新复现即可。 + +--- + +## 8. 交付物 + +0. **B2 的实测数字**（最先交，可能改变后续全部计划）：`(batch=B, s_q=tier+1)` vs + `(batch=B*(tier+1), s_q=1)` 的正确性与吞吐对比。若 token-major 不可行，立刻报告 + 并改推 D3 Stage A。 +1. 一个基于 `pr-17056` 的分支，表一 + 表二逐项修完，每项 commit message 里引用 + goal doc 的编号（H0/H1/.../K8） +2. A3/A4/A6 断言 + 参数化的 `test_deepseek_v4_sparse_mla.py` 单测；compressor ragged + 差分单测（`new_tokens_per_seq` 全填 `next_n` 必须与传 `None` bitwise 相同 —— goal + doc §6.3 标为「欠缺」） +3. GSM8K 在 **overlap scheduler + DEP8 + ragged** 下的实测分数，**以及 planner.stats + 的 fallback 计数**（证明调度真的生效，而不是静默退化成 max_tier）；均匀 tier 路径 + 的同口径分数作为对照 +4. 一份结论：hang 的真实根因是 §7 三个候选中的哪个（附证据） +5. **对 goal doc §8 剩余的待拍板问题给出你的推荐**（Q2 和 Q8 已关闭，见 §2.5），尤其： + - **Q4**：`cap-accept` 模式要不要实现。**建议实现** —— ragged 成为第一优先级之后 + 它的价值更高了：`cap-accept` vs `compact` 的差异 ⇒ 必然是 ragged kernel bug； + `static` vs `cap-accept` 的差异 ⇒ 必然是调度/接受逻辑 bug。把 pass/fail 硬币变成 + 可定位诊断。成本不高（复用已有的 `torch.minimum(..., verify_lens)` clamp， + `interface.py:1473-1480`、`:1889-1896`）。当前是显式 `NotImplementedError` + （`dspark.py:471-478`） + - **Q5**：STS 表和 SPS cost table 由谁、按什么流程产出、存哪（checkpoint 内？ + repo `examples/configs/`？）。**在它们存在之前，任何性能结论都是空的，而且 + feature 在构造上无法 trim** + - **Q6**：tier ladder 要不要按 batch size 门控（102 → ~54 个 graph，换回 + ~0.4–0.8 GB/rank 的 KV） + - **Q7**：非贪婪采样的处理时机 —— 见表二 H7 + - **Q9**：要不要采用 DeepGEMM 的 `indices` varlen 参数（goal doc §7.3）。**ragged + 成为第一优先级之后这条值得重新评估**：vLLM #47808 用的正是它，能去掉 + `[num_tokens, max_blocks_per_seq]` 那个大 block table，代价是把该路径钉死在 SM100 + （DeepGEMM varlen 分支断言 `arch_major == 10 and next_n == 1`） + +--- + +## 9. 工作方式 + +- 每个改动都要能指出它对应 goal doc 的哪一条，以及 `file:line` +- 凡是没读码确认的推断，显式标注「待验证」—— 沿用 goal doc 的规范 +- **不要静默降级**。goal doc 反复强调：不支持的组合要显式 config-reject，静默降级是 + 最差的选项 +- 遇到与 goal doc 冲突的发现，先给证据再改结论，不要默默按自己的理解走 +- 提交遵循仓库 `AGENTS.md`：`git commit -s`（DCO），PR 标题格式 + `[JIRA/NVBUG/None][type] description`，不要加 AI co-author + +--- + +## 10. P0 执行结果（2026-08-02 收尾） + +本节由执行 agent 追加，记录实际跑出来的结论。**与前文冲突处以本节为准**，因为前文 +是任务下达时的认知，本节是实测。 + +### 10.1 核心结论：K1/K2 已解除，纯 Python，不需要编译 + +ragged 原本卡在 `mla_rope_generation`（`mlaKernels.cu:1161`）和 FMHA dispatch +（`thop/attentionOp.cpp:920`）都要求 `acc_q_len % batch_size == 0`——per-request +窗口从构造上违反它。解法是把 generation 半边按 **token-major** 呈现：每个 +generation token 作为一行长度为 1 的序列，`seq_len == 1` 整除任何行数。 + +`git diff pr-17056..HEAD -- cpp/` **为空**。op 的 `mMaxNumRequests` 只是 JIT +warmup 提示，per-step 请求数来自 batch，`getFmhaMultiCtasKvScratchSize` 按 SM 数 +而非 batch 算——从 Python 侧传入静态行数上限即可。 + +### 10.2 精度（G3/G4/G5） + +DEP8 + attention-DP + overlap scheduler + CUDA graph，forced ragged 窗口： + +| MoE backend | exact_match (flexible) | exact_match (strict) | +|---|---|---| +| TRTLLM | 96.4367 ± 0.5106 | 96.5125 ± 0.5053 | +| MEGAMOE_DEEPGEMM | 断言通过（分数未捕获，该轮早于 `-s` 开关） | + +uniform 基线 96.2092。`steps_ragged 81/129`，5 种窗口，`trim_ratio 0.23`。 + +### 10.3 吞吐（G6）：**这个 checkpoint 上没有收益空间** + +planner 在三种工况下都正确地选择不裁剪： + +| 工况 | 每 rank 批次 | accept_len | planner | +|---|---|---|---| +| GSM8K | 小 | 3.30 | 不裁 | +| 合成随机 token | 4–8 | 4.09 | 不裁 | +| 合成随机 token | **64** | 3.73（p≈0.90） | **不裁**，零 fallback | + +DSpark block drafter 在此模型上每 token 接受率约 **90%**。拟合表在该接受率下 +于所有批次档位都判定"验满整块最优"（判据：`bs=64, p≥0.80 → L=5`）。 +**这不是功能失效，是没有机会。** 要看到收益需要接受率显著更低的模型或负载。 + +吞吐实测（128 请求，isl1024/osl512，并发 64，KV 泄漏修复后）： + +| 配置 | tok/s | +|---|---| +| baseline（confidence head 关） | 7032.37 | +| ragged 窗口=5（不裁剪） | 7319.47 | +| rotating（confidence 无关地裁 29%） | 5994.39 | + +**唯一站得住的吞吐结论是 KV 泄漏修复值 +6.5%**（6869.73 → 7319.47，同配置）。 +ragged 相对 baseline 的 +4.1% 不作为结论：单次测量，且从未表征吞吐的 +run-to-run 方差。rotating 的 5994.39 **不得**当作本功能的代价——它是与 +confidence 无关的盲裁，是最坏情况且方向与真实策略相反。 + +### 10.4 显存代价（此前未量） + +| 项 | uniform | ragged | 倍数 | +|---|---|---|---| +| CUDA graph 张数 | 4.1 / 批次档 | 12.1 / 批次档 | ~3× | +| attention workspace | 68 MiB | 604 MiB | ~9× | + +graph key 多出的 token-bucket 维度是 3× 的直接来源。`max_batch_size=512` 时 +graph 私有池达 13.97 GiB 并 OOM——ragged 让显存悬崖来得更早。 + +### 10.5 修掉的真 bug + +1. **KV 泄漏(ragged 独有)**。`kv_cache_manager_v2.py` 的 target manager 把回退 + 补偿 gate 在 `if self.is_draft` 后面。预留侧每步加 `1 + max(draft_len, reserve)`， + 回退侧只退 `py_rewind_len = verified_len − accepted`，于是每步漏 + `draft_len − verified_len`，直到序列需要的 block 超过 `max_seq_len`，报 + `User-provided base page indices is too short`。uniform 下 `verified_len == + draft_len` 故该项为 0，所以 GSM8K 从未暴露。 +2. **token-major block table 布局**。展开表用了 `[rows, max_blocks]`，而 op 读 + `[num_pools, num_seqs, 2, max_blocks]`，且 gather 折掉了 `num_seqs` 轴 → 8 卡 + device-side OOB。 +3. **`host_request_types` / `prompt_lens_cuda` 未行化**。op 按 `num_seqs`(=896) + 索引 128 长的数组。 +4. **lm-eval 分数选取**。把 `samples`(1319×100) 混进平均，报出 44030.98 且 + **通过**了 96.0 的门槛。按 `metric,filter` 键结构修正。 + +### 10.6 方法论教训（值得写进流程） + +- **探针失效比功能失效更常见**。整晚四次"失败"里三次是探针问题：模块版本偏斜导致 + 埋点静默为空、sync 探针未 arm 就读、gate 在 ramp-up 期误触发。**沉默不等于通过**—— + 任何"没有告警"的结论必须先证明探针确实生效了。 +- **对照实验优于读代码**。KV 泄漏是靠"窗口=5 vs rotating,其余全同"一次定位的； + 在此之前我读了三轮代码都没找到。 +- **噪声底必须先测**。G2 token 等价性判定为**无法测量**：两次完全相同的 static run + 之间就有 15/16 发散，顺序提交也没改善——目标模型在 TP8+EP8+ADP 下跨进程不能复现 + 自己的贪心输出。差一点把不存在的 ragged 正确性问题报出去。 + +### 10.7 未完成 + +- 动态 sync 证明:探针已确认能 arm 且 arm 后零同步警告,但缺一次跑满的运行 +- `peer_shape_mismatch: 2`：并发 512 下首次出现的、真正可归因于 ragged 的 + graph miss——跨 rank bucket 协商未收敛，2/265 步让 8 个 rank 一起掉出 replay +- `forced_verify_lens` / rotating 是否应从生产代码移除（它按 batch 位置分窗口， + 与 confidence 正交，测不了策略；但目前是唯一能执行裁剪路径的手段） From 626b5dc1fb23878cd84764ad4da8c1f01d085591 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 06:48:51 -0700 Subject: [PATCH 050/177] [None][doc] Stop telling users to force windows to bootstrap a table The error text still said "To bootstrap a table, set TLLM_DSPARK_FORCE_VERIFY_LENS", which is wrong for the same reason the docstring above it was: dspark_sps_profiler pins RaggedVerifyMode.STATIC and sweeps the verify length by rebuilding the engine, so collecting a table never enables ragged verification and needs no forced windows. Point at the profiler instead. Signed-off-by: Lance Liao --- tensorrt_llm/llmapi/llm_args.py | 5 +++-- .../unittest/llmapi/test_dspark_ragged_requires_table.py | 9 ++++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 3c69d8b1e48a..8059230c0856 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2891,8 +2891,9 @@ def validate_ragged_verify(self): "the cost of extra verify tokens against their expected " "acceptance, so it declines to trim and every request receives " "the full window -- the ragged path runs but changes nothing. " - "To bootstrap a table, set TLLM_DSPARK_FORCE_VERIFY_LENS to " - "impose windows without the planner.") + "Collecting a table does NOT require this flag: the profiler " + "pins static mode and sweeps the verify length by rebuilding " + "the engine, so run it with ragged verification off.") return self @model_validator(mode="after") diff --git a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py index cb7cefb23be6..c8f790e2832c 100644 --- a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py +++ b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py @@ -9,8 +9,11 @@ (``steps_ragged: 0``) that nobody reads unless they already suspect the problem. Config errors belong at construction. -``TLLM_DSPARK_FORCE_VERIFY_LENS`` is the deliberate exception: producing a cost -table requires a working ragged run, and that knob is what breaks the circle. +``TLLM_DSPARK_FORCE_VERIFY_LENS`` is the deliberate exception -- not because +table collection needs it (the profiler pins static mode and never enables +ragged), but because the planner declines to trim whenever acceptance is high, +which is every workload measured on this model. Without a way to impose +windows the trimming path cannot be executed at all. """ import os @@ -46,7 +49,7 @@ def test_ragged_with_an_sps_table_is_accepted(): @patch.dict(os.environ, {"TLLM_DSPARK_FORCE_VERIFY_LENS": "1"}, clear=False) def test_forced_windows_bootstrap_without_a_table(): - """The escape hatch must stay open, or the table can never be produced.""" + """The escape hatch must stay open, or trimming cannot be exercised.""" cfg = _cfg() assert cfg.enable_ragged_verify From f62bf13a8e6bf23b3e1329ca4e7354b9854754b9 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 08:19:42 -0700 Subject: [PATCH 051/177] [None][refactor] Drop TLLM_DSPARK_FORCE_VERIFY_LENS for a crafted cost table The knob assigned verify windows by BATCH POSITION, which is orthogonal to confidence. It reproduced the ragged shape but never the ragged policy, so it could not distinguish a correct confidence-driven assignment from an arbitrary one -- and the 18% throughput drop measured with it is the cost of confidence-blind trimming, the worst case and opposite in direction to what the planner does. Reporting that as the feature's cost would have been wrong. It also had to live in production code ahead of the cost-table gate to be reachable, and every justification written for it turned out not to hold: "needed to bootstrap a cost table" was false (dspark_sps_profiler pins RaggedVerifyMode.STATIC and sweeps by rebuilding the engine, so collecting a table never enables ragged), and "the only way to exercise trimming" is answered better below. SGLang has no serving-path analogue at all: its dspark_force_budget_frac is a profiler-only RPC carrying a budget fraction that still flows through the confidence-ordered top-k. A crafted non-flat cost table does the same job through the real path. The planner trims when marginal token cost outweighs expected acceptance, so a steep theta(M) makes it trim at any acceptance rate, and the budget then flows through the same top-k that hands longer windows to requests with higher survival. tests/unittest/_torch/speculative/ test_dspark_planner_trims.py pins that, including the property the position-rotating knob could never check -- that the confident request gets a longer window than the collapsing one. This matters beyond tidiness: the KV leak fixed earlier in this branch is reachable only when verified_len < draft_len, and on this checkpoint the planner declines to trim at every batch size, so that path needs a deliberate way to be exercised. A cost table is the honest one. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/model_engine.py | 33 +++++++ .../speculative/dspark_observability.py | 95 ------------------ .../_torch/speculative/dspark_sps_profiler.py | 19 ++-- .../_torch/speculative/dspark_verify.py | 19 ---- tensorrt_llm/llmapi/llm_args.py | 9 +- .../defs/accuracy/test_llm_api_pytorch.py | 2 +- .../test_dspark_confidence_schedule.py | 66 ------------- .../speculative/test_dspark_planner_trims.py | 99 +++++++++++++++++++ .../test_dspark_ragged_requires_table.py | 18 +--- 9 files changed, 143 insertions(+), 217 deletions(-) create mode 100644 tests/unittest/_torch/speculative/test_dspark_planner_trims.py diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index f2e9109977c5..df46657a6c8d 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2017,6 +2017,17 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: # in this pass uses the non-greedy key; populate's override # below will keep it False on every subsequent iteration. spec_metadata.is_all_greedy_sample = False + # Measure what the captured graphs actually cost. Ragged keys graphs + # on (batch_size, token_bucket) rather than batch size alone, which + # multiplies the capture set -- 12.1 graphs per batch-size bucket + # against 4.1 uniform, measured. The graph count was easy to count + # from the log; the memory was not, and an OOM at max_batch_size=512 + # reported 13.97 GiB in "private pools (e.g., CUDA Graphs)" without + # attributing it. Report the delta so the cost is a number rather + # than an inference. + _graph_mem_before = (torch.cuda.memory_reserved(), + torch.cuda.memory_allocated()) + _graph_count = 0 try: for entry in graphs_to_capture: # Ragged verification adds a token-count axis, so entries @@ -2027,6 +2038,7 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: continue for max_seq_len in max_seq_len_list: + _graph_count += 1 warmup_request = self._create_cuda_graph_warmup_request( resource_manager, bs, draft_len, max_seq_len) with self._release_batch_context( @@ -2063,6 +2075,27 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: finally: if force_non_greedy and spec_metadata is not None: spec_metadata._force_non_greedy_for_capture = False + # Absolute private-pool size, not a delta. Deltas are useless + # here: an engine build runs several capture passes over a + # shared allocator, so the per-pass increments came out as + # +13.9 / +5.5 / -0.8 / +0.7 GiB for the same 19 graphs -- the + # later passes reuse pool memory the earlier ones reserved, and + # reserved can even fall. CUDA graphs live in private pools + # (segment_pool_id != (0, 0)), which is the same quantity + # PyTorch's OOM message reports as "allocated in private pools + # (e.g., CUDA Graphs)", so sum those directly. + _pool_bytes = 0 + _pool_ids = set() + for _seg in torch.cuda.memory_snapshot(): + _pid = _seg.get("segment_pool_id", (0, 0)) + if _pid and tuple(_pid) != (0, 0): + _pool_bytes += _seg.get("total_size", 0) + _pool_ids.add(tuple(_pid)) + logger.info( + f"CUDA graph capture: {_graph_count} graphs this pass, " + f"private pools now {_pool_bytes / (1 << 30):.3f} GiB " + f"across {len(_pool_ids)} pool(s) " + f"(ragged={getattr(self.spec_config, 'enable_ragged_verify', False)})") # The capture loop stamps `self.runtime_draft_len` per graph so each # captured shape is built at its own draft length. Nothing puts it back, diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index 8f9971b48217..632de3018190 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -69,7 +69,6 @@ "RaggedVerifyMode", "read_ragged_verify_mode", "DSparkRaggedStats", - "forced_verify_lens", ] #: Overrides ``DSparkDecodingConfig`` so a mode can be selected without editing @@ -123,100 +122,6 @@ def read_ragged_verify_mode( f"{', '.join(repr(m.value) for m in RaggedVerifyMode)}") -#: Debug-only override that forces a deterministic non-uniform split. See -#: :func:`forced_verify_lens`. -FORCE_VERIFY_LENS_ENV = "TLLM_DSPARK_FORCE_VERIFY_LENS" - - -def forced_verify_lens(*, num_gen_requests: int, tiers: Sequence[int], - min_verify_len: int) -> Optional[List[int]]: - """A deterministic ragged split for correctness testing, or None. - - Two forms: - - ``TLLM_DSPARK_FORCE_VERIFY_LENS=1`` - Every generation request takes a window from the captured tier ladder, - rotating by batch position -- a deliberately *non-uniform* split, which - is what exercises the ragged packing. - - ``TLLM_DSPARK_FORCE_VERIFY_LENS=`` (n a tier, e.g. ``3``) - Every request takes the same window ``n``. Uniform, so it exercises - nothing about raggedness -- its purpose is cost measurement. - - The second form exists for SPS profiling. Sweeping the verify length by - shrinking ``block_size`` also shrinks the *draft* pass, but at deployment the - block is always drafted in full, so the measured step cost carries a draft - term that production does not pay. That term is bilinear in - ``(batch, verify_len)`` and therefore lies exactly in the column space the - additive fit spans, which means it is absorbed silently: theta comes out - inflated and alpha deflated, both biasing the planner toward over-trimming, - and the fit residual cannot see it. Forcing the window instead leaves - ``block_size`` at its full value, so the draft cost is constant across the - sweep and the bias is gone. SGLang solves the same problem with a runtime - ``dspark_force_budget_frac`` knob. - - Why this exists at all: the planner refuses to trim without a profiled cost - table -- correctly, since a flat model makes every extra token look free -- - and even *with* one it declines whenever acceptance is high enough that - verifying the full block wins. On DeepSeek-V4-Pro-DSpark that is every - workload measured so far (GSM8K at accept_len 3.3, synthetic prompts at - 4.1), so a planner-driven run leaves the trimming path unexecuted and - cannot answer whether that path is correct. Forcing the windows separates - "is trimming implemented correctly" from "should we trim here"; it is what - exposed a KV leak that every planner-driven run had missed. - - NOT needed to produce a cost table. An earlier version of this note claimed - that, and it is wrong: dspark_sps_profiler pins RaggedVerifyMode.STATIC and - sweeps the verify length by rebuilding the engine per length, so table - collection never enables the ragged path. - - Neither form is a shortcut around the planner: both decide only *how many* - drafted positions reach the target, exactly the quantity the planner decides. - Acceptance is unchanged, so a forced split must produce the same answers as - ``static`` -- which is what makes it a usable correctness gate. - - Returns: - One window per request, each drawn from ``tiers`` (so every choice has a - captured CUDA graph), or None when the override is off. - - Raises: - ValueError: the value is neither a boolean form nor a captured tier. - Silently ignoring it would produce a run that looks forced and is - not -- the exact failure this override exists to rule out. - """ - value = os.environ.get(FORCE_VERIFY_LENS_ENV, "").strip() - if not value or value in ("0", "false", "False"): - return None - ladder = sorted({int(t) for t in tiers if int(t) >= int(min_verify_len)}) - if not ladder: - return None - if value in ("1", "true", "True"): - # Rotate through the ladder by position: deterministic, identical on - # every rank for a given batch position, and guaranteed to be non-uniform - # as soon as the batch is larger than one and the ladder has more than - # one entry. - # - # Note "1" is the rotating form, not the window 1. A single tier ladder - # makes the two coincide anyway, and the rotating form is the older - # meaning; spelling the uniform window as a value keeps both reachable. - return [ladder[i % len(ladder)] for i in range(int(num_gen_requests))] - try: - window = int(value) - except ValueError: - raise ValueError( - f"invalid {FORCE_VERIFY_LENS_ENV}={value!r}; expected 1 (rotate " - f"through the tier ladder) or one of the captured tiers {ladder}") - if window not in ladder: - raise ValueError( - f"{FORCE_VERIFY_LENS_ENV}={window} is not a captured tier; the " - f"available tiers are {ladder}. A window with no captured CUDA " - f"graph does not raise at runtime -- it silently drops every step " - f"out of graph replay, which costs far more than the trimmed " - f"tokens save and would make any measurement taken with it " - f"meaningless.") - return [window] * int(num_gen_requests) - - class DSparkRaggedStats: """Per-step counters for the ragged verification path. diff --git a/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py b/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py index 961398bedb7a..dea4669a9643 100644 --- a/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py +++ b/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py @@ -22,7 +22,8 @@ ``DSparkVerifyPlanner.decide_verify_lens`` returns ``None`` on every step (``stats["fallback_flat_cost"]``) -- ragged verification is then a no-op that still produces correct output at baseline accuracy, i.e. it fails silently. -``TLLM_DSPARK_FORCE_VERIFY_LENS=1`` unblocks *correctness* testing by bypassing +A crafted (non-flat) cost table unblocks *correctness* testing by making the +planner itself trim, rather than bypassing that gate, but it rotates a fixed ladder and knows nothing about cost, so it cannot answer "is trimming profitable here". Only a measured table can. @@ -90,7 +91,7 @@ Holding verification uniform ---------------------------- The run pins ``TLLM_DSPARK_RAGGED_VERIFY_MODE=static`` and clears -``TLLM_DSPARK_FORCE_VERIFY_LENS``. Under ``static`` the planner still chooses a +a crafted cost table. Under ``static`` the planner still chooses a single batch-wide length, but the cost table it sees is the flat default, whose ``is_flat`` gate short-circuits to ``max_tier`` *before* any confidence is read -- so the length is deterministically ``max_draft_len`` on every step and @@ -100,7 +101,7 @@ That determinism is also why the ``L`` axis costs one engine build per value: nothing in the runtime pins a *uniform* verify length below the full block -(``TLLM_DSPARK_FORCE_VERIFY_LENS`` produces a ragged rotation, and there is no +(a ragged split would break the M = bs*(L+1) filing, and there is no analogue of SGLang's ``dspark_force_budget_frac``), so ``L`` is swept by rebuilding with ``max_draft_len = block_size = L``. Use ``--samples-out`` / ``--fit-only`` to split a long sweep across jobs and refit without a GPU. @@ -175,7 +176,7 @@ import numpy as np -from .dspark_observability import FORCE_VERIFY_LENS_ENV, RAGGED_VERIFY_MODE_ENV, RaggedVerifyMode +from .dspark_observability import RAGGED_VERIFY_MODE_ENV, RaggedVerifyMode from .dspark_planner import ( SpsCostTable, budget_argmax_over_uniform_lens, @@ -856,7 +857,7 @@ def profitability_probe( Runs :func:`~.dspark_planner.budget_argmax_over_uniform_lens` -- the actual production argmax, not a re-derivation -- against a geometric survival curve ``p**(j+1)``. This is the question the table exists to answer and the one - ``TLLM_DSPARK_FORCE_VERIFY_LENS`` cannot: at what acceptance rate does + a single sweep cannot: at what acceptance rate does trimming become the better trade, and does it ever? """ table = load_cost_table(payload) @@ -1108,14 +1109,6 @@ def _prepare_environment(config: SweepConfig) -> None: Set before the LLM is constructed because the worker processes inherit the environment at spawn and read these once. """ - if os.environ.get(FORCE_VERIFY_LENS_ENV): - print( - f"[dspark-sps] clearing {FORCE_VERIFY_LENS_ENV}: it rotates a ragged " - f"split across the batch, so M would no longer be bs*(L+1) and every " - f"cell would be filed under the wrong token count.", - file=sys.stderr, - ) - os.environ.pop(FORCE_VERIFY_LENS_ENV, None) os.environ[RAGGED_VERIFY_MODE_ENV] = RaggedVerifyMode.STATIC.value if config.pin_acceptance: from .interface import FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index b0a86da93d62..048dc4012c25 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -43,7 +43,6 @@ import torch from ..._utils import prefer_pinned -from .dspark_observability import forced_verify_lens from .dspark_planner import SpsCostTable, budget_argmax_over_uniform_lens from .dspark_schedule import DSparkScheduleConfig, compute_survival, schedule_verify_lens_topk @@ -258,24 +257,6 @@ def decide_verify_lens( if num_gen_requests <= 0: return None - forced = forced_verify_lens(num_gen_requests=num_gen_requests, - tiers=self.tiers, - min_verify_len=self.cfg.min_verify_len) - if forced is not None: - # Deliberately ahead of the cost-table gate. Without a profiled - # table the planner correctly refuses to trim, which means the - # ragged layout is never built and a correctness run silently - # degenerates into a uniform one -- so the two questions "is the - # ragged packing correct" and "does the planner choose to trim" - # cannot be answered by the same run. This override answers the - # first without needing an answer to the second. - # - # Safe by construction: it only decides how many drafted positions - # are sent to the target. Acceptance still runs unchanged against - # whatever is sent, so the output distribution is untouched. - self.stats["forced_lens"] = self.stats.get("forced_lens", 0) + 1 - return forced - if self.cost_table is None or self.cost_table.is_flat: self.stats["fallback_flat_cost"] += 1 return None diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 8059230c0856..d26de67a719c 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2868,8 +2868,7 @@ def validate_ragged_verify(self): "per-request verify lengths are chosen from confidence-head " "survival, which is only computed when scheduling is on") if (self.enable_ragged_verify and not self.confidence_sps_table_path - and not self.confidence_sts_path - and not os.environ.get("TLLM_DSPARK_FORCE_VERIFY_LENS")): + and not self.confidence_sts_path): # Without a profiled cost table the planner's budget degenerates to # verify-all -- correctly, since a flat model makes every extra # token look free -- so the ragged path silently never runs. That @@ -2877,12 +2876,6 @@ def validate_ragged_verify(self): # "succeeds" having done nothing and the only evidence is a counter # nobody reads. Fail at construction instead. # - # TLLM_DSPARK_FORCE_VERIFY_LENS is the deliberate exception. Not - # because table collection needs it -- dspark_sps_profiler pins - # STATIC mode and never enables ragged -- but because the planner - # declines to trim whenever acceptance is high, which is every - # workload measured on this model. Without a way to impose windows, - # the trimming path cannot be exercised or tested at all. raise ValueError( "enable_ragged_verify=True requires a profiled cost table: pass " "confidence_sps_table_path (produced by " diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index a495d4cfe893..0212fb61fb52 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4136,7 +4136,7 @@ def test_gsm8k_dep8_ragged_verify(self, moe_backend): Accuracy alone would not establish much here -- the planner declines to trim without a profiled SPS cost table, and a run where it declined on every step is a uniform baseline wearing a different config, scoring - exactly the same. ``TLLM_DSPARK_FORCE_VERIFY_LENS`` exists for that + exactly the same. A crafted non-flat cost table exists for that reason: it hands out windows from the captured ladder deterministically, which separates "is the ragged packing correct" from "does the planner choose to trim". Set it in the environment when running this to test the diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py index 915728114bcd..4cc69ba6cc6e 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py @@ -765,72 +765,6 @@ def test_cross_rank_reduction_is_max_not_min(): assert chosen[0] == max(locals_) -def test_forced_verify_lens_accepts_a_uniform_window(monkeypatch): - """``TLLM_DSPARK_FORCE_VERIFY_LENS=`` pins every request to one window. - - This is the form SPS profiling needs. Sweeping the verify length by shrinking - ``block_size`` shrinks the draft pass too, and that draft term is bilinear in - (batch, verify_len) -- exactly the column space the additive cost fit spans, - so it is absorbed into theta/alpha instead of showing up as residual. Forcing - the window leaves block_size alone, so the draft cost is constant across the - sweep. - """ - from tensorrt_llm._torch.speculative.dspark_observability import ( - FORCE_VERIFY_LENS_ENV, forced_verify_lens) - - monkeypatch.setenv(FORCE_VERIFY_LENS_ENV, "3") - lens = forced_verify_lens(num_gen_requests=5, tiers=[1, 3, 5], - min_verify_len=1) - assert lens == [3, 3, 3, 3, 3] - - -def test_forced_verify_lens_rotates_on_the_boolean_form(monkeypatch): - """``=1`` keeps its original meaning: rotate, i.e. deliberately non-uniform.""" - from tensorrt_llm._torch.speculative.dspark_observability import ( - FORCE_VERIFY_LENS_ENV, forced_verify_lens) - - monkeypatch.setenv(FORCE_VERIFY_LENS_ENV, "1") - lens = forced_verify_lens(num_gen_requests=5, tiers=[1, 3, 5], - min_verify_len=1) - assert lens == [1, 3, 5, 1, 3] - assert len(set(lens)) > 1, "the rotating form must not be uniform" - - -def test_forced_verify_lens_rejects_an_uncaptured_window(monkeypatch): - """A window off the ladder must raise, not be quietly ignored. - - It has no captured CUDA graph, so it would not fail loudly at runtime -- it - would drop every step into eager, which costs more than the trimmed tokens - save and makes any measurement taken with it meaningless. - """ - import pytest - - from tensorrt_llm._torch.speculative.dspark_observability import ( - FORCE_VERIFY_LENS_ENV, forced_verify_lens) - - monkeypatch.setenv(FORCE_VERIFY_LENS_ENV, "4") - with pytest.raises(ValueError, match="not a captured tier"): - forced_verify_lens(num_gen_requests=2, tiers=[1, 3, 5], - min_verify_len=1) - - monkeypatch.setenv(FORCE_VERIFY_LENS_ENV, "banana") - with pytest.raises(ValueError, match="expected 1"): - forced_verify_lens(num_gen_requests=2, tiers=[1, 3, 5], - min_verify_len=1) - - -def test_forced_verify_lens_off_by_default(monkeypatch): - from tensorrt_llm._torch.speculative.dspark_observability import ( - FORCE_VERIFY_LENS_ENV, forced_verify_lens) - - monkeypatch.delenv(FORCE_VERIFY_LENS_ENV, raising=False) - assert forced_verify_lens(num_gen_requests=3, tiers=[1, 3, 5], - min_verify_len=1) is None - monkeypatch.setenv(FORCE_VERIFY_LENS_ENV, "0") - assert forced_verify_lens(num_gen_requests=3, tiers=[1, 3, 5], - min_verify_len=1) is None - - def _stats(max_draft_len=5): from tensorrt_llm._torch.speculative.dspark_observability import ( DSparkRaggedStats, RaggedVerifyMode) diff --git a/tests/unittest/_torch/speculative/test_dspark_planner_trims.py b/tests/unittest/_torch/speculative/test_dspark_planner_trims.py new file mode 100644 index 000000000000..7b7a8a64ab35 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_dspark_planner_trims.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Make the planner itself produce per-request windows, with a crafted table. + +This replaces TLLM_DSPARK_FORCE_VERIFY_LENS. That knob assigned windows by +*batch position*, which is orthogonal to confidence: it reproduced the ragged +*shape* but never the ragged *policy*, so it could not tell a correct +confidence-driven assignment from an arbitrary one. It also had to sit in +production code, ahead of the cost-table gate, to be reachable at all. + +A cost table does the same job through the real path. The planner trims when +the marginal cost of extra verify tokens outweighs their expected acceptance, +so a table whose theta(M) rises steeply makes it trim at any acceptance rate -- +and the budget then flows through the same confidence-ordered top-k, which +hands different windows to requests with different survival. That exercises +what actually ships. + +Why this matters beyond tidiness: the KV leak fixed in this branch (KV reserved +for the full drafted block but only the verified window rewound) is reachable +only when verified_len < draft_len. On this checkpoint the planner declines to +trim at every batch size, because draft acceptance is ~90% and the fitted table +says verifying the full block wins -- so without a way to make it trim, that +code path ships untested. +""" + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_schedule import ( + DSparkScheduleConfig, schedule_verify_lens_topk) + + +def _Cfg(min_verify_len=1, max_verify_len=5): + """The real config, so the test cannot drift from what ships.""" + return DSparkScheduleConfig(block_size=max_verify_len, + min_verify_len=min_verify_len, + max_verify_len=max_verify_len) + + +def test_topk_hands_longer_windows_to_more_confident_requests(): + """The property the position-rotating knob could never check. + + Request 0 is confident throughout, request 2 collapses immediately. A + correct scheduler spends the budget where survival is highest. + """ + # survival[r, k] = P(the first k+1 drafted tokens all get accepted) + survival = torch.tensor([ + [0.99, 0.98, 0.97, 0.96, 0.95], # confident + [0.90, 0.70, 0.50, 0.30, 0.10], # middling + [0.20, 0.04, 0.01, 0.00, 0.00], # collapses + ]) + lens = schedule_verify_lens_topk(survival=survival, budget=6, + cfg=_Cfg()).tolist() + + assert len(lens) == 3 + assert lens[0] > lens[2], ( + f"confident request got {lens[0]}, collapsing request got {lens[2]}; " + f"the budget must follow survival, not batch position") + assert lens[0] >= lens[1] >= lens[2] + # Every request keeps at least the floor, and the budget is respected. + assert min(lens) >= 1 + assert sum(lens) - len(lens) * 1 <= 6 + + +def test_a_budget_below_the_full_block_produces_differing_windows(): + """Ragged shape arises from the planner, not from an override.""" + survival = torch.tensor([ + [0.99, 0.98, 0.97, 0.96, 0.95], + [0.60, 0.30, 0.10, 0.02, 0.00], + [0.95, 0.90, 0.20, 0.05, 0.00], + [0.10, 0.01, 0.00, 0.00, 0.00], + ]) + lens = schedule_verify_lens_topk(survival=survival, budget=5, + cfg=_Cfg()).tolist() + assert len(set(lens)) >= 2, ( + f"expected a ragged split from a constrained budget, got {lens}") + + +def test_a_full_budget_degenerates_to_the_uniform_full_window(): + """The no-trim case must stay exactly uniform. + + This is what the planner correctly chooses on this checkpoint, and it is + the case where ragged must cost nothing. + """ + num_reqs, max_len = 4, 5 + survival = torch.full((num_reqs, max_len), 0.99) + budget = num_reqs * (max_len - 1) + lens = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=_Cfg()).tolist() + assert lens == [max_len] * num_reqs + + +@pytest.mark.parametrize("budget", [0, 1, 3, 8, 20]) +def test_windows_stay_within_bounds_for_any_budget(budget): + survival = torch.rand(6, 5).sort(dim=1, descending=True).values + lens = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=_Cfg()).tolist() + assert all(1 <= v <= 5 for v in lens), lens + assert sum(lens) - len(lens) <= budget diff --git a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py index c8f790e2832c..145a0b666268 100644 --- a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py +++ b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py @@ -9,11 +9,9 @@ (``steps_ragged: 0``) that nobody reads unless they already suspect the problem. Config errors belong at construction. -``TLLM_DSPARK_FORCE_VERIFY_LENS`` is the deliberate exception -- not because -table collection needs it (the profiler pins static mode and never enables -ragged), but because the planner declines to trim whenever acceptance is high, -which is every workload measured on this model. Without a way to impose -windows the trimming path cannot be executed at all. +There is no exception: a crafted (non-flat) cost table is how a test makes the +planner trim, which exercises the real confidence-driven path rather than +bypassing it. See tests/unittest/_torch/speculative/test_dspark_planner_trims.py. """ import os @@ -35,25 +33,16 @@ def _cfg(**kwargs): @patch.dict(os.environ, {}, clear=False) def test_ragged_without_a_cost_table_is_rejected(): - os.environ.pop("TLLM_DSPARK_FORCE_VERIFY_LENS", None) with pytest.raises(ValueError, match="requires a profiled cost table"): _cfg() @patch.dict(os.environ, {}, clear=False) def test_ragged_with_an_sps_table_is_accepted(): - os.environ.pop("TLLM_DSPARK_FORCE_VERIFY_LENS", None) cfg = _cfg(confidence_sps_table_path="/tmp/table.json") assert cfg.enable_ragged_verify -@patch.dict(os.environ, {"TLLM_DSPARK_FORCE_VERIFY_LENS": "1"}, clear=False) -def test_forced_windows_bootstrap_without_a_table(): - """The escape hatch must stay open, or trimming cannot be exercised.""" - cfg = _cfg() - assert cfg.enable_ragged_verify - - @patch.dict(os.environ, {}, clear=False) def test_uniform_confidence_scheduling_still_needs_no_table(): """Only the ragged path is gated. @@ -62,7 +51,6 @@ def test_uniform_confidence_scheduling_still_needs_no_table(): max tier, which is the pre-existing behaviour -- so requiring one there would break working configurations. """ - os.environ.pop("TLLM_DSPARK_FORCE_VERIFY_LENS", None) cfg = DSparkDecodingConfig(max_draft_len=5, speculative_model="/nonexistent/model", enable_confidence_scheduling=True, From 2c7cd47a76aa11db7c8d794c4d8b010cb0167e5d Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 08:22:55 -0700 Subject: [PATCH 052/177] [None][test] Emit a steep SPS table so the planner exercises ragged verify On DeepSeek-V4-Pro-DSpark the planner never trims: draft acceptance is about 90% (accept_len 3.73 of 5) and the fitted table says verifying the full block wins at every batch size in the ladder, which is the correct decision. The consequence is that the trimming half of ragged verification -- per-request windows, the packing that depends on them, and the KV accounting for the positions left unverified -- never runs, so it ships untested. A real bug lived there: KV was reserved for the full drafted block but only the verified window rewound. A cost table whose theta(M) is superlinear makes the planner trim through its own logic. The budget then flows through the same confidence-ordered top-k that ships, so requests with different survival get different windows -- which is the property that matters and the one the removed TLLM_DSPARK_FORCE_VERIFY_LENS could not check, since it assigned windows by batch position. The output is marked SYNTHETIC: true, so run_dspark_throughput.sh refuses to treat it as evidence, and the module docstring says plainly that no timing taken with it means anything -- the planner is optimising against a fictional curve and will trim where a real deployment should not. An exponent that is not steeper than linear is rejected outright, since a table that fails to induce trimming is the exact failure this tool exists to avoid. Signed-off-by: Lance Liao --- .../dspark_make_steep_sps_table.py | 128 ++++++++++++++++++ 1 file changed, 128 insertions(+) create mode 100644 tests/microbenchmarks/dspark_make_steep_sps_table.py diff --git a/tests/microbenchmarks/dspark_make_steep_sps_table.py b/tests/microbenchmarks/dspark_make_steep_sps_table.py new file mode 100644 index 000000000000..4ca5fa429367 --- /dev/null +++ b/tests/microbenchmarks/dspark_make_steep_sps_table.py @@ -0,0 +1,128 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Emit an SPS cost table steep enough that the planner chooses to trim. + +WHAT THIS IS FOR, AND WHAT IT IS NOT FOR +---------------------------------------- +It exercises the ragged path end to end *through the planner*: with a steep +``theta(M)`` the confidence scheduler decides to trim, and its budget then +flows through the same confidence-ordered top-k that ships, which hands +different windows to requests with different survival. That is the only way to +see per-request verify lengths on a checkpoint whose acceptance is too high for +trimming to pay. + +**No throughput or latency number taken with this table means anything.** The +planner is optimising against a fictional cost curve, so it will trim where a +real deployment should not. Use it to answer "does the ragged path correctly +support different verify lengths per request", never "is ragged faster". The +output carries ``"SYNTHETIC": true`` so a run cannot quietly present it as a +measurement -- ``run_dspark_throughput.sh`` refuses to treat such a table as +evidence. + +WHY A TABLE RATHER THAN AN OVERRIDE +----------------------------------- +This replaces ``TLLM_DSPARK_FORCE_VERIFY_LENS``, which assigned windows by +batch position. Position is orthogonal to confidence, so that knob reproduced +the ragged *shape* while bypassing the policy entirely -- it could not tell a +correct assignment from an arbitrary one, and it had to sit in production code +ahead of the cost-table gate to be reachable. A table drives the real path. + +HOW STEEPNESS IS CHOSEN +----------------------- +The scheduler maximises expected accepted tokens per second, roughly + + (bs + sum of the top-K survival probabilities) / T(bs, K) + +With per-token acceptance around 0.9 each extra verify token adds about 0.9 to +the numerator, so trimming only wins once ``theta`` grows faster than that. The +default ``--exponent 2.0`` makes theta quadratic in the token count, which is +comfortably steeper than the measured curve (near-linear) and forces trimming +across the whole batch-size ladder. + + # derive from a real table, keeping its grid + python3 tests/microbenchmarks/dspark_make_steep_sps_table.py \ + --from sps_real_final.json --out steep.json + + # or synthesise a grid outright + python3 tests/microbenchmarks/dspark_make_steep_sps_table.py --out steep.json +""" + +import argparse +import json +import sys + + +def _default_grid(max_batch_size: int, width: int): + """Token totals a ragged batch can land on: bs * (verify_len + 1).""" + batch_sizes = [b for b in (1, 4, 8, 16, 32, 64, 128) if b <= max_batch_size] + tokens = sorted({b * w for b in batch_sizes for w in range(2, width + 1)}) + return batch_sizes, tokens + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--out", required=True) + parser.add_argument("--from", dest="source", default=None, + help="real table to take the grid from; its timings " + "are replaced, only token_counts/batch_sizes are " + "reused so the shape matches the deployment") + parser.add_argument("--exponent", type=float, default=2.0, + help="theta(M) ~ M**exponent; >1 makes extra verify " + "tokens superlinearly expensive") + parser.add_argument("--peak-ms", type=float, default=400.0, + help="theta at the largest token count") + parser.add_argument("--fixed-overhead-ms", type=float, default=5.0, + help="the non-trimmable floor; small on purpose, so " + "the trimmable term dominates the decision") + parser.add_argument("--max-batch-size", type=int, default=128) + parser.add_argument("--max-verify-len", type=int, default=5) + args = parser.parse_args() + + if args.exponent <= 1.0: + parser.error( + f"--exponent {args.exponent} is not steeper than linear, so the " + f"planner will keep verifying the full block and the ragged path " + f"will stay dark -- which is the exact failure this tool exists " + f"to avoid") + + if args.source: + with open(args.source, encoding="utf-8") as handle: + src = json.load(handle) + tokens = [int(t) for t in src["token_counts"]] + batch_sizes = [int(b) for b in src["batch_sizes"]] + else: + batch_sizes, tokens = _default_grid(args.max_batch_size, + args.max_verify_len + 1) + + tokens = sorted({t for t in tokens if t > 0}) + peak = float(max(tokens)) + theta = [args.peak_ms * (t / peak)**args.exponent for t in tokens] + + table = { + "token_counts": tokens, + "step_time_ms": theta, + "fixed_overhead_ms": args.fixed_overhead_ms, + "batch_sizes": batch_sizes, + "batch_overhead_ms": [0.0] * len(batch_sizes), + "SYNTHETIC": True, + "_meta": { + "purpose": "force the planner to trim so the ragged path runs; " + "NOT a measurement of anything", + "exponent": args.exponent, + "derived_from": args.source, + }, + } + with open(args.out, "w", encoding="utf-8") as handle: + json.dump(table, handle, indent=2) + handle.write("\n") + + print(f"wrote {args.out}: {len(tokens)} token buckets " + f"{tokens[0]}..{tokens[-1]}, theta {theta[0]:.3f}..{theta[-1]:.1f} ms " + f"(exponent {args.exponent}), SYNTHETIC=true") + print("Use only to check that per-request verify lengths work. Any " + "throughput taken with this table is meaningless.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 17b81ff01fd21a2adb8177f9e7e2a944ed7ac3e3 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 09:01:05 -0700 Subject: [PATCH 053/177] [None][test] Catch steep tables that cannot induce trimming, before the run Two ways to build a table that looks steep and still leaves the ragged path dark, both hit while trying to use this tool: The fixed overhead can swamp the trimmable term. A 5 ms floor against theta(48) = 1.6 ms at bs=8 means trimming saves almost nothing while costing accepted tokens, so the planner correctly keeps the full window. The default floor is now 0.05 ms. The grid can stop short of the deployment's widest step. Deriving from a real table topping out at 768 tokens while max_batch_size 256 reaches 1536 at the full window makes theta clamp there, so the full window looks no more expensive than a trimmed one -- again no trimming. The synthesised ladder now always includes max_batch_size, and --from is rejected outright when its grid is too narrow. Both were silent: the table generated fine and the failure only showed up as distinct_verify_lens: 1 after a twenty-minute run. So print the decision the planner will make at a few batch sizes, and warn when no previewed size trims. This is the fifth time this session that the probe rather than the feature was at fault; a fixture that cannot self-check is not worth having. Signed-off-by: Lance Liao --- .../dspark_make_steep_sps_table.py | 71 +++++++++++++++++-- 1 file changed, 64 insertions(+), 7 deletions(-) diff --git a/tests/microbenchmarks/dspark_make_steep_sps_table.py b/tests/microbenchmarks/dspark_make_steep_sps_table.py index 4ca5fa429367..7496a76d63b6 100644 --- a/tests/microbenchmarks/dspark_make_steep_sps_table.py +++ b/tests/microbenchmarks/dspark_make_steep_sps_table.py @@ -53,10 +53,18 @@ def _default_grid(max_batch_size: int, width: int): - """Token totals a ragged batch can land on: bs * (verify_len + 1).""" - batch_sizes = [b for b in (1, 4, 8, 16, 32, 64, 128) if b <= max_batch_size] - tokens = sorted({b * w for b in batch_sizes for w in range(2, width + 1)}) - return batch_sizes, tokens + """Token totals a ragged batch can land on: bs * (verify_len + 1). + + The ladder always includes ``max_batch_size`` itself. A grid that stops + short of the deployment's widest step is worse than useless: theta clamps + at the last bucket, so the full window looks no more expensive than a + trimmed one and the planner keeps it -- a table that silently cannot + induce trimming, which is the failure this tool exists to prevent. + """ + ladder = sorted({b for b in (1, 4, 8, 16, 32, 64, 128, 256, 512) + if b <= max_batch_size} | {max_batch_size}) + tokens = sorted({b * w for b in ladder for w in range(2, width + 1)}) + return ladder, tokens def main() -> int: @@ -71,9 +79,17 @@ def main() -> int: "tokens superlinearly expensive") parser.add_argument("--peak-ms", type=float, default=400.0, help="theta at the largest token count") - parser.add_argument("--fixed-overhead-ms", type=float, default=5.0, - help="the non-trimmable floor; small on purpose, so " - "the trimmable term dominates the decision") + parser.add_argument("--fixed-overhead-ms", type=float, default=0.05, + help="the non-trimmable floor. MUST be small next to " + "theta in the operating range, or trimming cannot " + "pay however steep theta is -- see --preview-bs") + parser.add_argument("--preview-bs", default="8,32,128", + help="batch sizes to predict the planner's choice for, " + "so a table that cannot induce trimming is caught " + "here instead of 20 minutes into a run") + parser.add_argument("--preview-accept", type=float, default=0.9, + help="per-token acceptance to predict against; 0.9 is " + "what DeepSeek-V4-Pro-DSpark measures") parser.add_argument("--max-batch-size", type=int, default=128) parser.add_argument("--max-verify-len", type=int, default=5) args = parser.parse_args() @@ -85,11 +101,19 @@ def main() -> int: f"will stay dark -- which is the exact failure this tool exists " f"to avoid") + needed = args.max_batch_size * (args.max_verify_len + 1) if args.source: with open(args.source, encoding="utf-8") as handle: src = json.load(handle) tokens = [int(t) for t in src["token_counts"]] batch_sizes = [int(b) for b in src["batch_sizes"]] + if max(tokens) < needed: + parser.error( + f"--from {args.source} tops out at {max(tokens)} tokens but " + f"max_batch_size {args.max_batch_size} reaches {needed} at the " + f"full window. theta would clamp there, making the full window " + f"look as cheap as a trimmed one, so the planner would never " + f"trim. Drop --from to synthesise a grid that covers it.") else: batch_sizes, tokens = _default_grid(args.max_batch_size, args.max_verify_len + 1) @@ -116,6 +140,39 @@ def main() -> int: json.dump(table, handle, indent=2) handle.write("\n") + # Predict the decision the planner will make. Its argmax is roughly + # (bs + sum of top-K survival) / T(bs, K), T = fixed + theta(M) + # so a table can be steep and still never induce trimming if the fixed + # floor swamps theta at the token counts the deployment actually uses. + # That is exactly what happened with a 5 ms floor against theta(48) = 1.6 + # ms: the planner correctly kept the full window and the run proved + # nothing. + def _theta(m): + return args.peak_ms * (min(m, peak) / peak)**args.exponent + + print() + print(f"predicted choice at acceptance {args.preview_accept} " + f"(K = verify positions kept, {args.max_verify_len} = no trim):") + any_trim = False + for bs in [int(b) for b in args.preview_bs.split(",") if b.strip()]: + best_k, best_rate = None, -1.0 + for k in range(1, args.max_verify_len + 1): + m = bs * (k + 1) + gain = bs + sum(args.preview_accept**(i + 1) for i in range(k)) * bs + rate = gain / (args.fixed_overhead_ms + _theta(m)) + if rate > best_rate: + best_k, best_rate = k, rate + trims = best_k < args.max_verify_len + any_trim = any_trim or trims + print(f" bs={bs:<5} -> K={best_k}" + f"{' TRIM' if trims else ' (full window)'}") + if not any_trim: + print() + print("WARNING: this table induces no trimming at any previewed batch " + "size, so the ragged path will stay dark. Lower " + "--fixed-overhead-ms or raise --exponent.", file=sys.stderr) + + print() print(f"wrote {args.out}: {len(tokens)} token buckets " f"{tokens[0]}..{tokens[-1]}, theta {theta[0]:.3f}..{theta[-1]:.1f} ms " f"(exponent {args.exponent}), SYNTHETIC=true") From 9103fd3b85728a6799495592aba82955f3a69d79 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 09:25:05 -0700 Subject: [PATCH 054/177] [None][test] Preview with the real planner instead of re-deriving its argmax The hand-written preview predicted "trim at every batch size" for a table under which the planner kept the full window through a twenty-minute run. It modelled the objective as (bs + sum survival) / (fixed + theta(M)), but budget_argmax_over_uniform_lens calls cost_table.step_time( total_verify_tokens(bs, length), num_gen_requests), which interpolates the table and adds the per-batch term. A preview that disagrees with the thing it previews is worse than no preview -- it converts a cheap check into false confidence. Import and call the real function. Backtested against three tables and it now matches every observed outcome: steep, grid <= 768 bs=8 -> L=5 (no trim), bs=64 -> L=1 -- and the lbs=128 run, whose per-rank batch is ~8, indeed never trimmed; the old formula had predicted it would steep, grid <= 1536 trims at 8, 64 and 256 real profiled table L=5 at 8 and 64, matching every production run Returns None rather than a guess when the planner is not importable, so "not checked" is distinguishable from "checked and fine". Signed-off-by: Lance Liao --- .../dspark_make_steep_sps_table.py | 85 +++++++++++++------ 1 file changed, 61 insertions(+), 24 deletions(-) diff --git a/tests/microbenchmarks/dspark_make_steep_sps_table.py b/tests/microbenchmarks/dspark_make_steep_sps_table.py index 7496a76d63b6..b22260c5b90e 100644 --- a/tests/microbenchmarks/dspark_make_steep_sps_table.py +++ b/tests/microbenchmarks/dspark_make_steep_sps_table.py @@ -67,6 +67,50 @@ def _default_grid(max_batch_size: int, width: int): return ladder, tokens + +def _preview_with_real_planner(*, path, batch_sizes, accept, max_verify_len): + """What budget_argmax_over_uniform_lens actually picks for this table. + + Returns True if any previewed batch size trims, False if none do, or None + when the planner cannot be imported (outside the container), so the caller + can distinguish "checked and fine" from "not checked". + """ + try: + import numpy as np + from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, budget_argmax_over_uniform_lens) + except Exception: + return None + + table = SpsCostTable.from_json_file(path) if hasattr( + SpsCostTable, "from_json_file") else None + if table is None: + with open(path, encoding="utf-8") as handle: + table = SpsCostTable(**{ + k: v for k, v in json.load(handle).items() + if not k.startswith("_") and k != "SYNTHETIC" + }) + + tiers = sorted({1, max(1, (max_verify_len + 1) // 2), max_verify_len}) + print() + print(f"planner's actual choice at acceptance {accept} " + f"(L = verify length, {max_verify_len} = no trim), tiers {tiers}:") + any_trim = False + for bs in [int(b) for b in str(batch_sizes).split(",") if str(b).strip()]: + # survival[r, k] = P(first k+1 drafted tokens all accepted) + survival = np.array( + [[accept**(k + 1) for k in range(max_verify_len)]] * bs, + dtype=np.float64) + chosen = budget_argmax_over_uniform_lens( + survival=survival, num_gen_requests=bs, cost_table=table, + allowed_lens=tiers, min_verify_len=1) + trims = int(chosen) < max_verify_len + any_trim = any_trim or trims + print(f" bs={bs:<5} -> L={chosen}" + f"{' TRIM' if trims else ' (full window)'}") + return any_trim + + def main() -> int: parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) parser.add_argument("--out", required=True) @@ -140,30 +184,23 @@ def main() -> int: json.dump(table, handle, indent=2) handle.write("\n") - # Predict the decision the planner will make. Its argmax is roughly - # (bs + sum of top-K survival) / T(bs, K), T = fixed + theta(M) - # so a table can be steep and still never induce trimming if the fixed - # floor swamps theta at the token counts the deployment actually uses. - # That is exactly what happened with a 5 ms floor against theta(48) = 1.6 - # ms: the planner correctly kept the full window and the run proved - # nothing. - def _theta(m): - return args.peak_ms * (min(m, peak) / peak)**args.exponent - - print() - print(f"predicted choice at acceptance {args.preview_accept} " - f"(K = verify positions kept, {args.max_verify_len} = no trim):") - any_trim = False - for bs in [int(b) for b in args.preview_bs.split(",") if b.strip()]: - best_k, best_rate = None, -1.0 - for k in range(1, args.max_verify_len + 1): - m = bs * (k + 1) - gain = bs + sum(args.preview_accept**(i + 1) for i in range(k)) * bs - rate = gain / (args.fixed_overhead_ms + _theta(m)) - if rate > best_rate: - best_k, best_rate = k, rate - trims = best_k < args.max_verify_len - any_trim = any_trim or trims + # Ask the REAL planner what it will do, rather than re-deriving its argmax. + # + # A hand-written approximation of the objective was wrong twice: it models + # the cost as fixed + theta(M), while budget_argmax_over_uniform_lens calls + # cost_table.step_time(total_verify_tokens(bs, length), num_gen_requests), + # which interpolates the table and adds the per-batch term. The + # approximation predicted "trim at every batch size" for a table under + # which the planner kept the full window through a 20-minute run. A preview + # that does not agree with the thing it previews is worse than none. + any_trim = _preview_with_real_planner( + path=args.out, batch_sizes=args.preview_bs, + accept=args.preview_accept, max_verify_len=args.max_verify_len) + if any_trim is None: # planner not importable here + print("preview skipped: run this inside the TensorRT-LLM container to " + "have the real planner check the table before you spend a run " + "on it", file=sys.stderr) + any_trim = True print(f" bs={bs:<5} -> K={best_k}" f"{' TRIM' if trims else ' (full window)'}") if not any_trim: From d766faebbdf0c2513ea2e2fdb88817847618a22a Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 09:27:05 -0700 Subject: [PATCH 055/177] [None][doc] Correct the memory section: ragged's graph cost is ~zero The previous 10.4 claimed 3x CUDA graphs, 9x attention workspace, and that ragged brings the memory cliff closer. Only the graph count was right, and a count is not a cost. Measured on one node, same dataset and concurrency, switching only the yaml, reading absolute private-pool size from memory_snapshot(): max_batch_size 128 uniform 19 graphs / 8.746 GiB ragged 57 graphs / 8.740 GiB (-0.006) max_batch_size 256 uniform 21 graphs / 16.527 GiB ragged 63 graphs / 16.588 GiB (+0.061) The count really is 3x -- it equals the tier count -- but CUDA graphs share a memory pool, so the extra captured shapes reuse the same buffers and the difference lands in the noise. What drives memory is max_batch_size: 8.7 -> 16.5 -> 33.1 GiB for 128 -> 256 -> 512. "ragged makes bs=512 OOM" had no control. Running it: the uniform arm at the same settings completed, and ragged died at 25.10 GiB -- below uniform's own 33.14 GiB peak, which is not how a higher-memory configuration fails. That log shows another process holding memory on the device; leftover processes bit this work twice, once from relaunching three minutes after a kill and finding two processes at 134 GiB each. The workspace figure has no uniform control either and should not be cited until it does. Also recorded: the first attempt measured deltas around the capture loop and produced +13.9 / +5.5 / -0.8 / +0.7 GiB for the same 19 graphs, because several capture passes share one allocator and reserved can fall. Absolute private-pool size reproduces to three decimals. Signed-off-by: Lance Liao --- docs/dspark_p0_task_prompt.md | 37 +++++++++++++++++++++++++++++------ 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/docs/dspark_p0_task_prompt.md b/docs/dspark_p0_task_prompt.md index c19b30869f18..421484cedb79 100644 --- a/docs/dspark_p0_task_prompt.md +++ b/docs/dspark_p0_task_prompt.md @@ -669,15 +669,40 @@ ragged 相对 baseline 的 +4.1% 不作为结论：单次测量，且从未表 run-to-run 方差。rotating 的 5994.39 **不得**当作本功能的代价——它是与 confidence 无关的盲裁，是最坏情况且方向与真实策略相反。 -### 10.4 显存代价（此前未量） +### 10.4 显存代价:实测约等于零 -| 项 | uniform | ragged | 倍数 | +**本节的前一版是错的,已按实测重写。** 前一版声称「CUDA graph 张数 3×、 +attention workspace 9×、ragged 让显存悬崖来得更早」,其中只有张数是对的, +而张数并不等于显存。 + +实测(同节点、同数据集、同并发,只切换 baseline / ragged 的 yaml,用 +`torch.cuda.memory_snapshot()` 求 CUDA graph 私有池的**绝对值**): + +| max_batch_size | uniform 图数 / 池 | ragged 图数 / 池 | 差 | |---|---|---|---| -| CUDA graph 张数 | 4.1 / 批次档 | 12.1 / 批次档 | ~3× | -| attention workspace | 68 MiB | 604 MiB | ~9× | +| 128 | 19 / **8.746 GiB** | 57 / **8.740 GiB** | **−0.006** | +| 256 | 21 / **16.527 GiB** | 63 / **16.588 GiB** | **+0.061** | +| 512 | 25 / **33.14 GiB** | 75 / — | — | + +**图数确实是 3×(等于 tier 数),但显存差落在噪声里。** CUDA graph 共享内存池, +多捕的形状复用同样的缓冲区,张数没有转化成显存代价。 + +真正决定显存的是 `max_batch_size` 本身:128 → 256 → 512 对应 +8.7 → 16.5 → 33.1 GiB,与 ragged 无关。 + +前一版还声称「ragged 让 bs=512 OOM」。补做对照后:同配置下 uniform 跑通, +而 ragged 那次是在 **25.10 GiB** 处挂的,**低于 uniform 自己 33.14 GiB 的峰值** +—— 一个更费显存的配置不会在更低点失败。那次的日志里有别的进程占着显存,最可能 +是节点残留(本次实验中已两次遇到:一次是 relaunch 过快导致两个进程各占 134 GiB)。 +一次不可复现、且与其余三个档位矛盾的观察,不足以支撑结论。 + +`attention workspace 68 MiB → 604 MiB` 那条同样**从未做过 uniform 对照**, +在补做之前不应引用。 -graph key 多出的 token-bucket 维度是 3× 的直接来源。`max_batch_size=512` 时 -graph 私有池达 13.97 GiB 并 OOM——ragged 让显存悬崖来得更早。 +**方法论**:前一版三条断言里,唯一做过对照的那条(张数)是对的,另外两条都是 +单次观察 + 事后编的机制解释。而且第一次测显存用的是「捕获前后取差值」,在多轮 +捕获共用分配器的情况下得出 `+13.9 / +5.5 / −0.8 / +0.7 GiB` 这种自相矛盾的 +数——增量法本身不成立;换成私有池绝对值后三次独立测量复现到小数点后三位。 ### 10.5 修掉的真 bug From 3f8e740e81558951845b81d0a9ae2581505836ab Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 09:51:58 -0700 Subject: [PATCH 056/177] [None][fix] Do not fit ragged buckets against rows that will not exist The ragged fit derives its token-bucket grid from an assumed padded row count and hands the rounding slack to pad rows. Padding is allowed to decline -- _get_padded_batch bails when padding would exceed the configured batch size, when the dummy request has no KV capacity, and on a couple of other resource checks. When it declines the rows never appear, the grid was built for a row count that is not realised, and the filled total matches no captured bucket. Observed on the first run where the planner trimmed for real (60.6% of the token ceiling, 4 distinct window sizes over 211 ragged steps): graph_miss_keys: {'(193, 5, False, False, True, 449)': 16} batch_size 193 against a captured ladder of ... 128, 192, 256 -- the batch never reached a captured row count, while the grid had been derived from padded_bs 256 ({512, 1024, 1536}), so the filled total 449 sat in no bucket and all eight ranks fell out of replay on 16 of 318 steps. Ask the runner whether padding is actually available before committing to a grid, and fall back to uniform when it is not -- uniform always has a captured graph, so the step still runs. will_pad_to mirrors the size guards in _get_padded_batch and deliberately does not predict the resource-dependent bails, since being optimistic there only leaves the pre-existing behaviour. Worth noting why this survived so long: without trimming the total is always padded_bs * (max_verify_len + 1), which is a captured bucket by construction, so no amount of uniform or lightly-trimmed testing reaches it. The earlier check that the caller's and chooser's padded_bs rules agree was done in exactly that regime, where the constant total hid the row-count axis entirely. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/cuda_graph_runner.py | 25 ++++++++++++++ .../_torch/pyexecutor/model_engine.py | 25 ++++++++++++++ .../test_dspark_graph_miss_attribution.py | 33 +++++++++++++++++++ 3 files changed, 83 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index e83988185c6a..c92ad092a0e2 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -629,6 +629,31 @@ def replay(self, key: KeyType, return output_ref + def will_pad_to(self, padded_bs: int, num_requests: int) -> bool: + """Whether a batch of ``num_requests`` can actually reach ``padded_bs``. + + The ragged fit derives its token-bucket grid from an assumed padded row + count and hands the rounding slack to pad rows. If padding then + declines, those rows never appear and the fitted total lands in no + captured bucket -- seen as a graph key of batch_size 193 against a + ladder holding 192 and 256, costing every rank a replay. + + Mirrors the size guards in :meth:`_get_padded_batch`. It deliberately + does not try to predict the resource-dependent bails (KV capacity for + the dummy request, encoder-decoder state): those are checked when the + dummy is created, and being wrong in the optimistic direction there + leaves the pre-existing behaviour rather than making it worse. + """ + if padded_bs <= 0 or num_requests <= 0: + return False + if padded_bs not in self.supported_batch_sizes: + return False + if padded_bs == num_requests: + return True # already at a captured size, nothing to add + # _get_padded_batch bails when padding_size + batch_size exceeds the + # configured batch size; that sum is just padded_bs. + return padded_bs <= self.config.batch_size + def _get_padded_batch(self, batch: ScheduledRequests, resource_manager: ResourceManager, runtime_draft_len: int) -> int: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index df46657a6c8d..b0d82256c18f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3756,6 +3756,31 @@ def fit_ragged_verify_lens( padded_bs = runner._round_up_batch_size(widest_rows) if padded_bs == 0: return None + # The whole fit is built on the assumption that the batch will be padded + # up to `padded_bs`: the bucket grid is derived from it and fill_bucket + # hands the rounding slack to rows that do not exist yet. But padding is + # allowed to decline -- _get_padded_batch bails when the dummy request + # has no KV capacity, when padding would exceed the configured batch + # size, and on a few other resource checks. When it does, the rows never + # appear, the grid was derived for a row count that is not realised, and + # the filled total lands in no captured bucket: observed as + # graph_miss_keys {'(193, 5, False, False, True, 449)': 16} on a run + # where the ladder held 192 and 256 but not 193, costing every rank a + # replay on 16 of 318 steps. + # + # Only reachable with real trimming. Without it the total is always + # padded_bs * (max_verify_len + 1), which is a captured bucket by + # construction, so no amount of uniform testing surfaces this. + # + # Decline to go ragged rather than fit against rows that will not exist. + # Uniform always has a captured graph, so the fallback is cheap and the + # step still runs. + if not runner.will_pad_to(padded_bs, len(token_lens)): + logger.debug( + f"DSpark ragged: padding to {padded_bs} rows is not available " + f"for {len(token_lens)} requests, so the bucket grid derived " + f"from it would not be realised; falling back to uniform") + return None buckets = self.ragged_verify_token_buckets(padded_bs) if not buckets: return None diff --git a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py index 11d4551922db..67dab0b1478a 100644 --- a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py +++ b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py @@ -100,3 +100,36 @@ def test_ramp_up_steps_do_not_count_toward_the_gate_floor(): for _ in range(64): stats.record_step(num_gen_requests=4, verify_lens=[1, 3, 5, 5]) assert stats.steps_multi_request == 64 + + +class _Runner: + """Just the fields will_pad_to reads.""" + + def __init__(self, supported, cfg_batch_size): + self.supported_batch_sizes = supported + self.config = type("C", (), {"batch_size": cfg_batch_size})() + + +def test_will_pad_to_rejects_row_counts_the_ladder_cannot_reach(): + """The ragged fit must not derive a grid from rows that never appear. + + A run fitted its bucket grid to padded_bs=256 while the batch stayed at 193 + -- the captured ladder held 192 and 256 but not 193 -- so the filled total + 449 matched no captured bucket and every rank lost a graph replay on 16 of + 318 steps. Only reachable with real trimming: without it the total is + padded_bs * (max_verify_len + 1), a captured bucket by construction. + """ + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ + CUDAGraphRunner + + runner = _Runner([1, 2, 4, 8, 128, 192, 256], cfg_batch_size=256) + will = CUDAGraphRunner.will_pad_to + + assert will(runner, 256, 193), "193 -> 256 is reachable and should be taken" + assert not will(runner, 193, 193), "193 is not a captured size" + assert will(runner, 192, 192), "already at a captured size" + # Padding past the configured batch size is refused by _get_padded_batch, + # so the fit must not assume it. + assert not will(_Runner([256, 512], cfg_batch_size=256), 512, 300) + assert not will(runner, 0, 10) + assert not will(runner, 256, 0) From d2c19dea01c860e037458469190701bb9c7bba89 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 18:54:49 -0700 Subject: [PATCH 057/177] [None][fix] Key ragged graphs on the agreed bucket, not a re-derived sum Every attention-DP rank has to replay the same graph shape, because the captured graph contains collectives and all_rank_num_tokens is frozen into it -- a width mismatch is not observable at replay time. The shape gate therefore compares the ragged token bucket across ranks and drops everyone to eager when they differ. They differed because the bucket was derived twice, independently. The fit computes it from allgathered peer stats; the graph key recomputed it by summing 1 + py_verify_len over generation_requests, which walks the batch *after* _get_padded_batch has appended padding. Two derivations of a quantity that must match exactly, so agreement was incidental rather than structural. Measured: peer_shape_mismatch on 2 of 265 steps without trimming and 16 of 318 with it -- trimming widens the per-rank window spread, so it makes the divergence more likely. Have the fit publish the bucket it agreed on and have the key read that. It comes from data every rank allgathered, through rules every rank runs identically, so it agrees by construction. Cleared on entry to the fit so no early return can leave a stale bucket behind, and None still means "not ragged", which leaves the uniform key byte-for-byte unchanged. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/cuda_graph_runner.py | 15 +++---- .../_torch/pyexecutor/model_engine.py | 16 ++++++++ .../test_dspark_graph_miss_attribution.py | 39 +++++++++++++++++++ 3 files changed, 63 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index c92ad092a0e2..b451a292146b 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -298,13 +298,14 @@ def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: """ if not getattr(self.spec_config, "enable_ragged_verify", False): return None - total = 0 - for request in batch.generation_requests: - verify_len = getattr(request, "py_verify_len", None) - if verify_len is None: - return None - total += 1 + int(verify_len) - return total if total else None + # The bucket the fit agreed on, not a fresh sum over the batch. Summing + # here walks generation_requests *after* _get_padded_batch has appended + # padding, so it is a second, independent derivation of a quantity every + # attention-DP rank must match exactly -- agreement was incidental. The + # fitted value is computed from allgathered peer stats by rules every + # rank runs identically, so it agrees by construction. None means this + # step is not ragged, which keeps the key exactly what it was. + return getattr(self, "agreed_ragged_bucket", None) @staticmethod def _local_draft_len(batch: ScheduledRequests) -> int: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index b0d82256c18f..eb89bf62d54e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3721,6 +3721,10 @@ def fit_ragged_verify_lens( path exists to avoid. """ runner = self.cuda_graph_runner + # Cleared on entry so that every early return leaves no stale + # value behind: the graph key reads this, and a bucket carried + # over from a previous step would key into the wrong graph. + runner.agreed_ragged_bucket = None if not runner.enabled or not verify_lens: return None if len(verify_lens) != len(generation_requests): @@ -3800,6 +3804,18 @@ def fit_ragged_verify_lens( # Remembered so a graph miss can report the shape it actually submitted # rather than just that one happened. self._dspark_last_total_tokens = int(bucket) + # The graph key uses THIS value rather than re-summing the + # per-request windows. Both used to be derived independently -- + # the key walked generation_requests *after* padding was + # appended, while the fit worked from the allgathered peer + # stats -- so their agreement across attention-DP ranks was + # incidental rather than structural, and when they diverged the + # shape gate dropped every rank out of replay + # (peer_shape_mismatch: 2/265 steps without trimming, 16/318 + # with). This value comes from data every rank allgathered and + # rules every rank runs identically, so it agrees by + # construction. + runner.agreed_ragged_bucket = int(bucket) # The CUDA-graph padding rows are appended later by _get_padded_batch, # which appends a *single shared dummy object*, so every pad row diff --git a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py index 67dab0b1478a..8979b04ea706 100644 --- a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py +++ b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py @@ -133,3 +133,42 @@ def test_will_pad_to_rejects_row_counts_the_ladder_cannot_reach(): assert not will(_Runner([256, 512], cfg_batch_size=256), 512, 300) assert not will(runner, 0, 10) assert not will(runner, 256, 0) + + +def test_graph_key_uses_the_agreed_bucket_not_a_fresh_sum(): + """Every attention-DP rank must key on the same bucket, by construction. + + The key used to re-sum 1 + py_verify_len over generation_requests, which + walks the batch *after* padding is appended -- a second, independent + derivation of a value all ranks must match exactly. When the two + derivations diverged the shape gate dropped every rank out of replay + (peer_shape_mismatch: 2/265 steps without trimming, 16/318 with it, since + trimming widens the per-rank window spread). Reading the fitted value + instead makes agreement structural: it comes from allgathered peer stats + through rules every rank runs identically. + """ + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ + CUDAGraphRunner + + class _R: + spec_config = type("S", (), {"enable_ragged_verify": True})() + + r = _R() + get = CUDAGraphRunner._ragged_verify_bucket + + # Nothing fitted yet -> not ragged, key unchanged from the uniform one. + assert get(r, object()) is None + + # A fitted bucket is returned verbatim, regardless of the batch contents: + # the batch is not even inspected, which is the point. + r.agreed_ragged_bucket = 449 + assert get(r, object()) == 449 + + # Cleared between steps so a stale bucket cannot key into the wrong graph. + r.agreed_ragged_bucket = None + assert get(r, object()) is None + + # Config gate still wins: a non-ragged engine keeps the original key shape. + r.agreed_ragged_bucket = 512 + r.spec_config = type("S", (), {"enable_ragged_verify": False})() + assert get(r, object()) is None From 6a6789a1734184d6848a0ad20961bfeb66156fd4 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 19:36:11 -0700 Subject: [PATCH 058/177] [None][doc] Record the P0 outcome: what holds, what was fixed, what bit us Section 10 now carries the whole executed result rather than the partial picture it had. The headline addition is 10.5: per-request verify lengths are validated, produced by the planner itself with no override, at 4 distinct window sizes across 211 of 318 ragged steps and 60.6% of the token ceiling trimmed, with zero KV overruns. That is the capability K1/K2 blocked. It is stated alongside the limit that matters -- the run used a deliberately steep synthetic cost table, so it proves the mechanism and says nothing about whether trimming pays. On this checkpoint it does not (10.3). 10.6 lists six real bugs, including three that are reachable only when the planner actually trims: the KV leak, the fit assuming padding will happen, and the bucket being derived twice so cross-rank agreement was incidental. Since acceptance here is ~90% and the planner therefore never trims, none of them can be reached without deliberately inducing it -- they would have merged with a green test suite. 10.7 records the process failures, because they outnumbered the product failures. Probes were at fault more often than the feature: an empty counter that a default argument disguised as "no occurrences", a sync detector read before it armed, a gate whose floor counted total steps while concurrency ramped from zero, a memory delta taken across capture passes that share an allocator, and two crafted tables that could not induce the trimming they existed to induce. Hence: silence is not a result until the probe is shown to have run; call the real function rather than re-deriving it; a fixture that cannot self-check is not worth having. It also lists the four conclusions that a control overturned or corrected, and the one place local evidence was used to support a global claim. 10.8 is explicit that the dynamic sync question is unresolved: 32 real synchronizations were observed after arming on the ragged side, but the control has not finished, so nothing about sync should be cited yet. Signed-off-by: Lance Liao --- docs/dspark_p0_task_prompt.md | 135 +++++++++++++++++++++++++--------- 1 file changed, 102 insertions(+), 33 deletions(-) diff --git a/docs/dspark_p0_task_prompt.md b/docs/dspark_p0_task_prompt.md index 421484cedb79..5855c237e3d3 100644 --- a/docs/dspark_p0_task_prompt.md +++ b/docs/dspark_p0_task_prompt.md @@ -704,37 +704,106 @@ attention workspace 9×、ragged 让显存悬崖来得更早」,其中只有张 捕获共用分配器的情况下得出 `+13.9 / +5.5 / −0.8 / +0.7 GiB` 这种自相矛盾的 数——增量法本身不成立;换成私有池绝对值后三次独立测量复现到小数点后三位。 -### 10.5 修掉的真 bug - -1. **KV 泄漏(ragged 独有)**。`kv_cache_manager_v2.py` 的 target manager 把回退 - 补偿 gate 在 `if self.is_draft` 后面。预留侧每步加 `1 + max(draft_len, reserve)`， - 回退侧只退 `py_rewind_len = verified_len − accepted`，于是每步漏 - `draft_len − verified_len`，直到序列需要的 block 超过 `max_seq_len`，报 - `User-provided base page indices is too short`。uniform 下 `verified_len == - draft_len` 故该项为 0，所以 GSM8K 从未暴露。 -2. **token-major block table 布局**。展开表用了 `[rows, max_blocks]`，而 op 读 - `[num_pools, num_seqs, 2, max_blocks]`，且 gather 折掉了 `num_seqs` 轴 → 8 卡 - device-side OOB。 +### 10.5 每请求不同的 verify len:已验证 + +这是 P0 的核心能力,也是 K1/K2 那个原始 blocker 要解决的东西。**已由 planner +自主产出,无任何外部强制**(`TLLM_DSPARK_FORCE_VERIFY_LENS` 已从代码中删除): + +``` +steps_ragged: 211/318 同一步内窗口不相同的步数 +distinct_verify_lens: 4 {1: 27375, 2: 5831, 3: 12, 5: 66} +trim_ratio: 0.6064 裁掉 60.6% 的 token 上限 +KV 溢出: 0 OOM: 0 +``` + +分配链路是真实的:confidence head → 每请求 survival 概率 → 成本表算出预算 → +top-k 按 survival 高低分配窗口。置信度高的拿长窗口,容易早被拒的拿短窗口。 + +**这证明了机制,不证明策略有收益。** 该轮用的是刻意造陡的合成成本表 +(`tests/microbenchmarks/dspark_make_steep_sps_table.py`,输出带 `SYNTHETIC: true`), +planner 在对着虚构的成本曲线优化,所以那轮的吞吐数字无意义。真实表下 planner 在 +所有批次档位都选满窗口(见 10.3)。 + +**为什么必须造这张表**:此 checkpoint 接受率约 90%,planner 永远不会自己走到 +裁剪路径。而下面三个 bug **只在真裁剪时可达**——没有这张表,它们会带着绿色的 +测试合入。 + +### 10.6 修掉的真 bug + +1. **KV 泄漏(ragged 独有)**。预留侧每步加 `1 + max(draft_len, reserve)`, + 回退侧只退 `py_rewind_len = verified_len − accepted`;`kv_cache_manager_v2` + 把补偿项 gate 在 `if self.is_draft` 后面,于是 target manager 每步漏 + `draft_len − verified_len`,直到序列需要的 block 超过 `max_seq_len`,报 + `User-provided base page indices is too short`。uniform 下该项恒为 0。 + **定位方式**:窗口=5 vs 窗口轮转、其余全同的一次对照,而非读代码。 + +2. **token-major block table 布局**。展开表用了 `[rows, max_blocks]`,而 op 读 + `[num_pools, num_seqs, 2, max_blocks]`,且 gather 折掉了 `num_seqs` 轴 → + 8 卡 device-side OOB。 + 3. **`host_request_types` / `prompt_lens_cuda` 未行化**。op 按 `num_seqs`(=896) - 索引 128 长的数组。 -4. **lm-eval 分数选取**。把 `samples`(1319×100) 混进平均，报出 44030.98 且 - **通过**了 96.0 的门槛。按 `metric,filter` 键结构修正。 - -### 10.6 方法论教训（值得写进流程） - -- **探针失效比功能失效更常见**。整晚四次"失败"里三次是探针问题：模块版本偏斜导致 - 埋点静默为空、sync 探针未 arm 就读、gate 在 ramp-up 期误触发。**沉默不等于通过**—— - 任何"没有告警"的结论必须先证明探针确实生效了。 -- **对照实验优于读代码**。KV 泄漏是靠"窗口=5 vs rotating,其余全同"一次定位的； - 在此之前我读了三轮代码都没找到。 -- **噪声底必须先测**。G2 token 等价性判定为**无法测量**：两次完全相同的 static run - 之间就有 15/16 发散，顺序提交也没改善——目标模型在 TP8+EP8+ADP 下跨进程不能复现 - 自己的贪心输出。差一点把不存在的 ragged 正确性问题报出去。 - -### 10.7 未完成 - -- 动态 sync 证明:探针已确认能 arm 且 arm 后零同步警告,但缺一次跑满的运行 -- `peer_shape_mismatch: 2`：并发 512 下首次出现的、真正可归因于 ragged 的 - graph miss——跨 rank bucket 协商未收敛，2/265 步让 8 个 rank 一起掉出 replay -- `forced_verify_lens` / rotating 是否应从生产代码移除（它按 batch 位置分窗口， - 与 confidence 正交，测不了策略；但目前是唯一能执行裁剪路径的手段） + 索引 128 长的数组。**主动审计 12 个 per-sequence 参数找出来的,不是等它崩。** + +4. **bucket fit 假设 padding 必然发生**。fit 按 padded_bs 推导 bucket 网格并把 + 余量分给尚不存在的 pad 行,但 `_get_padded_batch` 可以拒绝 padding。拒绝时 + 网格对应的行数不存在,填出的总数落在网格外:`graph_miss_keys + {'(193, 5, False, False, True, 449)': 16}`,捕获档位有 192 和 256 却没有 193。 + 修法:fit 前先问 `runner.will_pad_to`,不能 pad 就退回 uniform。 + +5. **跨 rank bucket 不一致**。bucket 被独立推导两次:fit 从 allgather 的 + `peer_stats` 算,graph key 又遍历 `generation_requests`(padding 追加**之后**) + 重新求和。两条独立推导却要求跨 rank 精确相等,一致性是碰巧的。实测 + `peer_shape_mismatch` 不裁剪时 2/265 步、裁剪时 16/318 步,每次让 8 个 rank + 一起掉出 replay。修法:fit 发布商定值,key 直接读——按构造必然一致。 + +6. **lm-eval 分数选取**(与 DSpark 无关的既有 bug)。把 `samples`(1319×100) + 混进平均,报出 **44030.98 并通过了 96.0 的门槛**。按 `metric,filter` 键结构修正。 + +### 10.7 踩过的坑(值得写进流程) + +**探针失效比功能失效更常见。** 本次至少七次「失败」中,探针本身有问题的占多数: + +| 现象 | 真因 | +|---|---| +| `graph_miss_shapes` 为空 | 模块版本偏斜,默认参数把「没记录」伪装成「没发生」 | +| sync 探针零警告 | 探针根本没 arm,我差点报通过 | +| gate 在 ramp 期误杀三轮实验 | floor 数的是总步数,而并发从 0 爬升 | +| 显存增量 `+13.9/+5.5/−0.8/+0.7` | 多轮捕获共用分配器,增量法不成立 | +| 造表后 ragged 仍不亮(两次) | 固定开销压倒可裁剪项;网格覆盖不到工作区间 | +| 预览说会裁、实跑不裁 | 我重推了 planner 的 argmax,与真实实现不符 | + +由此得到三条: + +* **沉默不等于通过。** 任何「没有告警」的结论,必须先证明探针确实生效。 +* **能调用就别重推。** 造表预览改成直接 `import` 真实的 + `budget_argmax_over_uniform_lens` 后一次就对,并能回测历史结果;测试里也用真的 + `DSparkScheduleConfig` 而非 mock。 +* **夹具必须能自检。** 造表工具现在会预测 planner 的选择,并在「任何档位都不裁」 + 时直接 WARNING——失败在生成的那一秒暴露,而不是 20 分钟后。 + +**没有对照的归因一律不成立。** 本次有四条结论因补对照而推翻或修正: + +| 断言 | 结局 | +|---|---| +| ragged graph 显存 3× | 实测 ≈0(8.746 vs 8.740 GiB) | +| ragged 让 bs=512 OOM | 不开 ragged 同样配置也曾跑通;且 ragged 在更低的 25.10 GiB 就挂,与「更费显存」矛盾 | +| `peer_shape_mismatch` 是 bug | 是那道门在正常拦截(但推导重复确实该修) | +| G2 token 等价性失败 | 两次**相同** static run 之间也 15/16 发散——噪声底 = 信号,该测试在此部署上问不出东西 | + +**局部证据不能支撑全局断言。** 我曾报「ragged 未引入新的 host-device sync」, +依据是审计了自己新增的 token-major 代码(那部分确实干净),但动态探针在 +`dspark.py:117/516/551` 和 `deepseek_v4.py:891` 抓到 32 次真同步。静态审计只能 +覆盖「我写的代码」,动态探针才覆盖「这条路径实际执行的一切」。 + +**在功能没满负荷运行时做的验证,适用范围比看起来窄。** 上面 4 和 5 两个 bug 都 +只在真裁剪时可达;我此前核对过「两条 padded_bs 规则一致」,但那是在不裁剪的前提 +下做的,总数恒定完全掩盖了行数这一维。 + +### 10.8 未完成 + +* **动态 sync 证明**(唯一未定论项)。ragged 侧已测得 arm 后 32 次真同步,位置见 + 上;**对照(confidence head 关、其余逐字相同)尚未跑完**,因此「这些是 ragged + 引入的还是 DSpark 既有的」暂无结论。在此之前不应引用任何 sync 相关结论。 +* **`peer_not_gen_only` 类 graph miss**:32/318 步,属正常连续批处理(uniform + 路径同样存在),非缺陷,但可作为后续调度优化的方向。 +* **吞吐收益需要更低接受率的模型或负载**。这是 checkpoint 的属性,非本 PR 可解。 From fc7bdb41443dd18898cdfec3c55ef02fc24341df Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 19:40:02 -0700 Subject: [PATCH 059/177] [None][doc] Scope the P0 results to B300; call out what will not transfer Every measurement in section 10 was taken on x86 + NVIDIA B300 SXM6 AC (275040 MiB per GPU, driver 580.167.08), single node, DEP8. Several of the conclusions are hardware-dependent and should not be carried to GB300 or any Grace-Blackwell system unchanged: - "ragged's CUDA graph overhead is ~zero" was measured against 267 GiB of HBM at free_gpu_memory_fraction 0.5; a different capacity, NVLink topology or C2C coherent memory changes how the graph pool and the KV cache trade off. - "at p~0.90 no batch size should trim" comes from an argmax whose denominator is the measured T(bs, M). A different compute/bandwidth ratio moves theta(M)'s slope and therefore moves the acceptance threshold at which trimming starts to pay. The number is not portable. - sps_real_final.json is a B300 artifact; rerun dspark_sps_profiler on the target. What does transfer: the mechanism validation in 10.5, the six bugs and their fixes in 10.6, and the methodology in 10.7. Those are host-side correctness, not performance. Signed-off-by: Lance Liao --- docs/dspark_p0_task_prompt.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/docs/dspark_p0_task_prompt.md b/docs/dspark_p0_task_prompt.md index 5855c237e3d3..e06bfcef35e1 100644 --- a/docs/dspark_p0_task_prompt.md +++ b/docs/dspark_p0_task_prompt.md @@ -620,6 +620,29 @@ goal doc 引用的 `.dspark-logs/gsm8k.log` 和 `hang-stack-299685.txt` 不在 本节由执行 agent 追加，记录实际跑出来的结论。**与前文冲突处以本节为准**，因为前文 是任务下达时的认知，本节是实测。 +> **硬件与配置口径:全部实测都在 x86 + NVIDIA B300 SXM6 AC(275040 MiB / 卡, +> 驱动 580.167.08)、8 卡单节点、DEP8(TP8 + EP8 + attention-DP)上取得。MoE +> backend 为 MEGAMOE_DEEPGEMM 或 TRTLLM,`max_draft_len=5`,tier 阶梯 `{1,3,5}`。 +> +> **换到 GB300(或任何 Grace-Blackwell 一体机)结论可能不同**,以下几组明确依赖 +> 硬件: +> +> * **显存(10.4)**:「ragged 的 CUDA graph 开销 ≈0」是在 267 GiB HBM、 +> `free_gpu_memory_fraction=0.5` 下测的。GB300 的显存容量、NVLink 拓扑和 +> C2C 一致性内存会改变 graph 池与 KV 的相对占比。 +> * **裁剪是否划算(10.3)**:planner 的 argmax 比较「每秒接受 token 数」,分母 +> 是实测的 `T(bs, M)`。不同的算力/带宽比会改变 `theta(M)` 的斜率,从而移动 +> 「裁剪开始划算」的接受率阈值。本文结论是「p≈0.90 时任何批次档位都不该裁」, +> **这个阈值不可跨硬件搬运**。 +> * **SPS cost table 必须在目标硬件上重采**。`sps_real_final.json` 只对 B300 +> 有效,在目标机上重跑 `dspark_sps_profiler` 即可。 +> * **每 rank 批次 = 并发 / DP 数**,本文按 DP=8 计算。 +> +> 与硬件**无关**的部分:10.5 的机制验证(每请求不同 verify len 能否正确运行)、 +> 10.6 的六个 bug 及其修复(都是 host 侧逻辑错误)、10.7 的方法论。这些是正确性 +> 问题,不是性能问题。 + + ### 10.1 核心结论：K1/K2 已解除，纯 Python，不需要编译 ragged 原本卡在 `mla_rope_generation`（`mlaKernels.cu:1161`）和 FMHA dispatch From 46fc7d12f72056223af34ce0f540155d25a2cdb3 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 20:27:02 -0700 Subject: [PATCH 060/177] [None][doc] Reconcile the docs with what was actually measured Several statements were left behind by conclusions that later reversed, and they now contradicted each other across files. docs/dspark_p0_report.md is a mid-flight snapshot whose headline finding -- "ragged is blocked by K1 and has never actually run", "structurally impossible to run end to end on DSv4" -- was overtaken by the token-major presentation. It now carries a banner naming the current source of truth and listing each reversal, while keeping what still holds: the B2 microbenchmark at 0.93-1.00x, the K1/K2 code analysis, and the file:line anchors. docs/dspark_ragged_verify_status.md documented TLLM_DSPARK_FORCE_VERIFY_LENS, which has since been removed, and repeated the justification for it that turned out to be false -- that a cost table cannot be produced without a working ragged run. The profiler pins RaggedVerifyMode.STATIC and sweeps by rebuilding the engine per verify length, so collecting a table never enables ragged; sps_real_final.json was produced that way. The real problem is different and worth stating plainly: even with a real table the planner declines to trim whenever acceptance is high, which is every workload measured here, so the trimming path needs a deliberate way to be exercised or it merges untested. That way is now a crafted steep cost table, which drives the real planner, rather than a knob that assigned windows by batch position and so could never test the policy. The same file also described assert_ragged_active as requiring graph_eager == 0. That was changed to attribute misses by reason, because a rank carrying context requests drops the whole attention-DP world out of replay on the uniform path too (~1.7%), making the original condition unsatisfiable on any real workload. Finally, the comment in trtllm.py cited 68 MiB -> 604 MiB of attention workspace as a measured cost, while the doc says that figure has no uniform control and should not be cited. The one memory claim that was controlled came out at parity (8.746 vs 8.740 GiB), so the comment now says so instead. Signed-off-by: Lance Liao --- docs/dspark_p0_report.md | 16 ++++++++++ docs/dspark_ragged_verify_status.md | 29 ++++++++++++++----- .../_torch/attention_backend/trtllm.py | 16 +++++----- 3 files changed, 46 insertions(+), 15 deletions(-) diff --git a/docs/dspark_p0_report.md b/docs/dspark_p0_report.md index d26942f3a02e..8ab302557dab 100644 --- a/docs/dspark_p0_report.md +++ b/docs/dspark_p0_report.md @@ -1,5 +1,21 @@ # DSpark confidence-head 调度 P0 阶段：执行报告 +> # ⚠️ 本文是中途快照,主要结论已被推翻 +> +> **当前事实以 `docs/dspark_p0_task_prompt.md` 第 10 节为准。** 本文写于 +> token-major 实现之前,其核心判断已不成立: +> +> | 本文的说法 | 现在的事实 | +> |---|---| +> | 「ragged 被 K1 挡住,从未真正跑起来过」(§0.4) | K1/K2 已由 token-major 呈现解除,**纯 Python,零 C++ 改动** | +> | 「ragged 在 DSv4 上结构上不可能端到端跑通」 | 已跑通:planner 自主产出 4 种窗口、211/318 步 ragged、裁剪 60.6% | +> | 「吞吐今天无法测量」 | 已测:此 checkpoint 无收益,因接受率 ~90% 时 planner 正确地不裁 | +> | 引用 `TLLM_DSPARK_FORCE_VERIFY_LENS` | 该开关已从代码中删除,改用合成陡峭 cost table 驱动真实 planner | +> +> 仍然有效的部分:§B2 的 token-major microbenchmark(0.93–1.00×)、对 K1/K2 +> 检查本身的代码分析、以及各处 `file:line` 锚点。 + + > 分支：`dspark-p0`（基于 PR #17056 head `d8897046`，PR 分支未改动） > 硬件：`bia` 集群，8× B300 SXM6（sm103），x86_64 > 构建：`/lustre/fsw/coreai_comparch_trtllm/laliao/build/dspark-p0/`（`-a '103-real;'`，rc23） diff --git a/docs/dspark_ragged_verify_status.md b/docs/dspark_ragged_verify_status.md index 5b2d3fef2a83..1e568f2407a6 100644 --- a/docs/dspark_ragged_verify_status.md +++ b/docs/dspark_ragged_verify_status.md @@ -151,20 +151,33 @@ P0-3/5/6 属于「跑出来的数字看着正常但其实是错的」：replay | 开关 | 作用 | |---|---| | `TLLM_DSPARK_RAGGED_VERIFY_MODE=static\|compact` | 选路径。`cap-accept` **未实现，显式抛 `NotImplementedError`** | -| `TLLM_DSPARK_FORCE_VERIFY_LENS=1` | 按 tier 阶梯轮转出确定性非均匀切分 | +| ~~`TLLM_DSPARK_FORCE_VERIFY_LENS=1`~~ | **已删除**。改用 `tests/microbenchmarks/dspark_make_steep_sps_table.py` 造陡峭 cost table,让**真实 planner** 自己产出不等窗口 | | `DSparkRaggedStats.assert_ragged_active()` | 断言 ragged 真的生效 | -`assert_ragged_active()` 检查四条：`steps_ragged > 0`、`distinct_verify_lens >= 2`、 -`trim_ratio > 0`、`graph_eager == 0`，失败时带完整计数器摘要。 +`assert_ragged_active()` 检查:`steps_ragged > 0`、`distinct_verify_lens >= 2`、 +`trim_ratio > 0`,以及 graph miss 中**可归因于 ragged 的那几类**,失败时带完整计数器摘要。 + +最后一条原为 `graph_eager == 0`,**已修正**:连续批处理下,某个 rank 带 context 请求 +就会让全体掉出 replay,uniform 路径同样如此(实测约 1.7%),所以按原判据这个断言在 +任何真实负载上都不可满足。现在按原因归类,只对 `key_not_captured`、 +`bs_not_supported`、`peer_shape_mismatch` 三类失败——它们分别表示 bucket 网格错了、 +批次档位没捕获、以及跨 rank 形状发散。 **`cap-accept` 为什么不能 alias 成 `compact`**：它唯一的价值是「输出必须与 `static` 逐 token 一致」，别名化会让这个对照失去意义。 -**`FORCE_VERIFY_LENS` 解决的循环依赖**：没有 profiled SPS cost table 时 -planner 在构造上不会 trim（`llm_args.py` 自述「the budget degenerates to -verify-all」），于是验证 ragged 打包正确性需要先有 table，而产出 table 又需要 -能跑的 run。这个开关替换的是 **planner 的决策**，不绕开任何下游代码路径—— -所以 forced-ragged 与 static 在 `temperature=0` 下必须逐 token 一致。 +**「循环依赖」这个理由已被证伪**:产出 cost table **不需要** ragged 能跑。 +`dspark_sps_profiler` 会把 `RAGGED_VERIFY_MODE` 钉成 `static`,并通过逐个 verify_len +重建引擎来扫成本曲线——采表全程 ragged 关闭(`sps_real_final.json` 就是这么采的)。 + +真正的问题是另一个:**即使有真实 table,planner 在接受率够高时依然不裁**——在 +DeepSeek-V4-Pro-DSpark 上是每一种实测工况(accept_len 3.3–4.1)。所以裁剪路径需要 +一个刻意的手段才能被执行,否则它会带着全绿的测试合入。 + +现在这个手段是**造一张陡峭的 cost table**,而不是按 batch 位置轮转的开关。区别是 +本质的:位置与 confidence 正交,轮转只能复现 ragged 的**形状**、复现不了**策略**, +测不出「置信度高的是否拿到更长窗口」。造表则驱动真实 planner,窗口经同一个 +confidence 排序的 top-k 分配。三个只在真裁剪时可达的 bug 正是这样被逼出来的。 --- diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 944952dddfa0..0438fbd85c6f 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -564,13 +564,15 @@ def presented_token_major(self, view): # JIT-warmup hint, and the per-step request count comes from the # batch, not from here. # - # It is not free: the attention workspace is sized from the same - # number, and on DEP8 (128 requests -> 896 rows) it was measured - # growing from 68 MiB to 604 MiB per rank. The op resizes itself, so - # this costs GPU memory that would otherwise hold KV cache rather - # than correctness. Shrinking it would mean sizing the generation - # workspace from the row count only where the row count is what the - # kernels actually use, which is a C++ change. + # The attention workspace is sized from the same number, so it + # grows with the row ceiling. An earlier version of this comment + # quoted 68 MiB -> 604 MiB per rank as the cost; that figure came + # from a single ragged log with no uniform control, and the one + # memory claim that WAS controlled -- CUDA graph pool size -- came + # out at parity (8.746 vs 8.740 GiB at max_batch_size 128), so the + # uncontrolled figure should not be relied on. The op resizes + # itself, so whatever the true delta is, it costs GPU memory that + # would otherwise hold KV cache rather than correctness. self.max_num_requests = view.max_num_rows # Every runtime view has to agree on batch_size with # kv_lens_cuda_runtime, and the op reads request_types over all of From 9597f64d1ae76617c7463b548513a0a8ea99739e Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 20:27:28 -0700 Subject: [PATCH 061/177] [None][doc] Mark the superseded claims inline, not just in the banner A reader landing mid-document does not see the header. The section whose heading still read "the most important conclusion: ragged is blocked by K1 and has never actually run" is the one most likely to be quoted out of context, so it now carries the reversal next to it, along with what in that section does still hold -- the analysis of what the K1 check actually does. The two references to the removed TLLM_DSPARK_FORCE_VERIFY_LENS are marked at their own locations too. Signed-off-by: Lance Liao --- docs/dspark_p0_report.md | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/docs/dspark_p0_report.md b/docs/dspark_p0_report.md index 8ab302557dab..4a7e17690c02 100644 --- a/docs/dspark_p0_report.md +++ b/docs/dspark_p0_report.md @@ -66,7 +66,11 @@ goal doc §6.3 称它是「性价比最高的一个测试」。参数化之后 --- -## 0.4 ⛔ 最重要的结论：ragged 在 DSv4 上被 K1 挡住，从未真正跑起来过 +## 0.4 ~~⛔ 最重要的结论：ragged 在 DSv4 上被 K1 挡住，从未真正跑起来过~~ + +> **【已推翻】** K1/K2 已由 token-major 呈现解除(纯 Python,`cpp/` 零改动), +> ragged 已端到端跑通。以下分析对「K1 检查本身做了什么」仍然准确,但结论 +> 「结构上不可能跑通」不成立。 端到端跑到真实 decode 之后，第一件撞上的事是： @@ -201,7 +205,7 @@ L=1 与 L=3 两组样本齐了之后，第一次可以把 `theta(M)` 从 `alpha( **上界性质**：profiler 通过缩小 `block_size` 来扫 L，draft pass 也跟着缩， 所以测得的 Δθ 里混入了 draft 代价，是**上界**。这使得「小值」是可靠的否定结论， 而「大值」需要进一步确认。要测无泄漏的版本，就是保持 `block_size=5` 不动、 -改用 `TLLM_DSPARK_FORCE_VERIFY_LENS=` 强制窗口 —— 该能力本轮已实现 +改用 `TLLM_DSPARK_FORCE_VERIFY_LENS=` 强制窗口 —— **【该开关已删除】**,现由合成陡峭 cost table 驱动真实 planner （见 §2 的 commit）。这是下一步该做的测量。 ## 2. 交付物 1：修复清单 @@ -520,7 +524,7 @@ graph_replays: 520, graph_eager: 9 ### Q4：`cap-accept` 要不要实现 —— **建议实现，但优先级低于 Q5** 任务书倾向实现，理由是把 pass/fail 硬币变成可定位诊断。这个理由成立，但要注意 -**PR head 已经有了一个更便宜的等价物**：`TLLM_DSPARK_FORCE_VERIFY_LENS=1` +**【已删除】** ~~PR head 已经有了一个更便宜的等价物~~：`TLLM_DSPARK_FORCE_VERIFY_LENS=1` （`dspark_observability.py:131-162`）从 tier ladder 里按批内位置轮转发窗口， 确定性、每个 rank 一致、且不需要 cost table。它已经能把 「ragged packing 对不对」与「planner 要不要 trim」分开 —— 而这正是 A2 想要的那一半。 From bb8198d29ddf3c6f1f93503aaf9896ecef16fb3c Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 21:18:00 -0700 Subject: [PATCH 062/177] [None][chore] Remove the uniform tier-ladder verify path DSpark confidence scheduling had three states, not two: 1. enable_confidence_scheduling=False verify the full block 2. scheduling=True, ragged=False one tier for the whole batch 3. scheduling=True, ragged=True per-request windows State 2 does not earn its keep. It picks a single verify length for every request in the batch, so it cannot give a confident request a longer window than a collapsing one -- which is the entire reason a confidence head exists. And because the ladder is coarse (tiers {1,3,5}), it degenerates to the full block whenever acceptance is high, which was every workload measured on DeepSeek-V4-Pro-DSpark (~90% acceptance). So it is a distinct code path, carried in the graph key and the cross-rank allgather, that in practice does what state 1 does. SGLang has no equivalent: its scheduler always produces per-request windows via top-k, and uniform appears only as the verify-all degenerate case. Removed: - DSparkVerifyPlanner.decide_draft_len and its _decide_local helper - DSparkVerifyPlanner.snap_to_tier, whose only caller was the uniform branch - the uniform branch in _dspark_decide_verify_windows; the fallback is now runtime_draft_len = planner.max_tier, i.e. verify everything - the tier element from the cross-rank allgather payload (5 -> 4 ints), since nothing downstream needs an agreed tier any more - 12 tests and a rank-simulation helper covering the deleted path llm_args now rejects scheduling=True with ragged=False rather than silently treating it as something else. Turning the feature off is enable_confidence_scheduling=False. No functional change to the ragged path or to static verification. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 32 +-- .../_torch/speculative/dspark_verify.py | 103 ++------ tensorrt_llm/llmapi/llm_args.py | 22 ++ .../test_dspark_confidence_schedule.py | 240 ------------------ .../test_dspark_ragged_requires_table.py | 29 ++- 5 files changed, 76 insertions(+), 350 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f9da3690a7dc..c979f50ec380 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3400,7 +3400,6 @@ def _dspark_confidence_draft_len(self, # single collective, and nothing after it may issue another one. ragged_lens = None fallback_reason = None - local_tier = planner.max_tier if compute_windows: ragged_lens = planner.decide_verify_lens( num_gen_requests=num_gen_requests, @@ -3410,22 +3409,20 @@ def _dspark_confidence_draft_len(self, fallback_reason = "planner_declined" else: fallback_reason = "mode_static" - # Only computed on the uniform path so the planner's fallback - # counters keep meaning "this step declined to trim" exactly once. - # On the ragged path the placeholder above is the largest tier -- - # i.e. verify everything -- which is the documented safe answer for - # the case this value is actually used: some *other* rank declined - # and dragged the whole group back to a uniform step. - local_tier = planner.decide_draft_len( - num_gen_requests=num_gen_requests, - rows=confidence_rows, - reduce_across_ranks=False) + # Nothing to compute: the fallback verifies the full drafted + # block. This used to call decide_draft_len, which picked one tier + # for the whole batch -- a middle ground between "schedule per + # request" and "verify everything". It is gone: it could not act on + # per-request confidence, which is the point of scheduling, and it + # degenerated to the full block whenever acceptance was high, which + # was every workload measured here. The config that reached it is + # rejected in llm_args now, so there are two states, not three. # --- Phase 2: one collective, same shape on every rank, every step ---- # # Carries everything the rest of the step needs, so the `peer_stats` # gather that used to follow is folded in here rather than added to it: - # [wants_ragged, num_rows, total_verify_tokens, max_window, tier] + # [wants_ragged, num_rows, total_verify_tokens, max_window] # `_can_queue` already runs `tp_allgather(batch_size)` just before this # hook and could absorb this too, but that would mean computing the # local draft length inside a hot path shared by every speculation mode. @@ -3440,7 +3437,6 @@ def _dspark_confidence_draft_len(self, local_rows, local_tokens, local_max_window, - int(local_tier), ]) if not all(int(payload[0]) for payload in payloads): # Any rank that cannot go ragged takes the whole group with it: @@ -3456,9 +3452,6 @@ def _dspark_confidence_draft_len(self, ragged_lens = [min(int(v), agreed_max) for v in ragged_lens] peer_stats = [[int(payload[1]), int(payload[2])] for payload in payloads] - agreed_tier = max(int(payload[4]) for payload in payloads) - else: - agreed_tier = int(local_tier) # --- Phase 3: no more collectives ------------------------------------- # @@ -3522,8 +3515,11 @@ def _dspark_confidence_draft_len(self, # layout with uniform acceptance -- silent token misattribution. for request in gen_requests: request.py_verify_len = None - # Already agreed in phase 2 -- deliberately not another collective. - runtime_draft_len = planner.snap_to_tier(agreed_tier) + # Verify the full drafted block. This is the only fallback: the + # uniform tier ladder that used to sit between "schedule per + # request" and "verify everything" is gone, so there is nothing to + # agree on across ranks here and no collective is needed. + runtime_draft_len = planner.max_tier # Stage this step's confidence for the *next* decision. Non-blocking; if # it has not landed by then the planner just verifies the full block. diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index 048dc4012c25..c273dbe0adfe 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -20,15 +20,18 @@ from ``(batch_size, draft_len, ...)`` -- every input to that choice is a host value. A device-resident budget would have to be copied back before the batch could be shaped, which is the synchronization we are trying to avoid. -2. **The decision must land on a captured draft length.** A length with no - captured graph does not raise; ``maybe_get_cuda_graph`` returns ``None`` and - the step silently runs eager, which costs far more than the tokens saved. - Hence :meth:`decide_draft_len` only ever returns a value from ``tiers``. -3. **Every rank must decide identically.** ``draft_len`` is part of the graph key - but is *not* part of the attention-DP consistency allgather, so two ranks that - pick different lengths select different graphs -- one replays, one falls back - to eager -- and their collectives diverge. :meth:`decide_draft_len` therefore - reduces across ranks before returning. +2. **The step must land on a captured graph.** A shape with no captured graph + does not raise; ``maybe_get_cuda_graph`` returns ``None`` and the step + silently runs eager, which costs far more than the tokens saved. The drafted + block is always ``max_tier`` long, so what varies is the *verified* token + total, and the layout rounds that up to a captured verify bucket before the + batch is shaped. +3. **Every rank must decide identically.** The graph key is not part of the + attention-DP consistency allgather, so two ranks that shape the batch + differently select different graphs -- one replays, one falls back to eager + -- and their collectives diverge. The caller therefore reduces the ragged + decision across ranks (a single allgather in ``py_executor``) before any rank + acts on it; :meth:`decide_verify_lens` itself is rank-local by default. Confidence for step ``t`` only becomes readable on the host at ``t+1`` (or ``t+2`` with the overlap scheduler), so the decision always runs on a lagged snapshot. @@ -55,9 +58,9 @@ class DSparkVerifyPlanner: Args: cfg: verify-length bounds. cost_table: measured step cost. A flat table means "not profiled"; the - planner then refuses to trim (see :meth:`decide_draft_len`). - tiers: draft lengths that have a captured CUDA graph. The returned length - is always one of these. + planner then refuses to trim (see :meth:`decide_verify_lens`). + tiers: draft lengths that have a captured CUDA graph. The drafted block + is always ``max_tier``; these also bound the budget search. apply_calibration: maps raw confidence logits to probabilities. Normally ``confidence_head.apply_sts``; defaults to a plain sigmoid. all_rank_max: cross-rank reduction, returning the max of ``value`` over @@ -107,8 +110,8 @@ def stage_confidence(self, confidence_logits: torch.Tensor) -> None: keying exists to avoid. Never synchronizes: the copy is issued on the current stream and an event - is recorded. :meth:`decide_draft_len` polls that event and simply does not - use the snapshot if it has not landed. + is recorded. :meth:`decide_verify_lens` polls that event and simply does + not use the snapshot if it has not landed. """ if confidence_logits is None or confidence_logits.shape[0] == 0: return @@ -170,59 +173,6 @@ def _gather_rows( return None return selected - def snap_to_tier(self, value: int) -> int: - """Round ``value`` up to a captured tier. - - Callers that reduce across ranks themselves need this: the reduced value - is a maximum over per-rank choices and so is already a tier, but a caller - combining it with anything else has to land back on the ladder before it - can be used as a graph key. - """ - value = int(value) - if value in self.tiers: - return value - return min((t for t in self.tiers if t >= value), default=self.max_tier) - - def decide_draft_len( - self, - *, - num_gen_requests: int, - rows: Optional[Sequence[int]] = None, - all_rank_max: Optional[Callable[[int], int]] = None, - reduce_across_ranks: bool = True, - ) -> int: - """Choose this iteration's draft length, agreed across all ranks. - - Degrades to the largest captured tier (i.e. verify everything, today's - behavior) whenever the inputs are not trustworthy: no confidence - snapshot yet, or an unprofiled (flat) cost model under which every extra - verified token looks free. - - ``rows`` maps each generation request to its buffer row; see - :meth:`_gather_rows`. ``all_rank_max`` overrides the constructor's - reduction; the caller that owns the distributed handle usually supplies - it here. - - ``reduce_across_ranks=False`` returns the *local* choice and issues no - collective, leaving the agreement to a caller that batches it with its - other cross-rank traffic. Since the reduction is a max over tiers, and - every local choice is already a tier, the caller only has to reduce and - pass the result through :meth:`snap_to_tier`. - """ - self.stats["decisions"] += 1 - chosen = self._decide_local(num_gen_requests=num_gen_requests, rows=rows) - if not reduce_across_ranks: - return int(chosen) - all_rank_max = all_rank_max or self.all_rank_max - if all_rank_max is not None: - # Max, not min: a rank that wanted to trim more simply verifies a few - # extra tokens. Disagreeing on the graph key is the unrecoverable - # outcome, so agreement matters more than the saving. - chosen = int(all_rank_max(int(chosen))) - if chosen not in self.tiers: - chosen = min((t for t in self.tiers if t >= chosen), default=self.max_tier) - return int(chosen) - def decide_verify_lens( self, *, @@ -294,22 +244,3 @@ def decide_verify_lens( f"internal: produced {len(lens)} verify lengths for {num_gen_requests} requests" ) return [int(v) for v in lens] - - def _decide_local(self, *, num_gen_requests: int, rows: Optional[Sequence[int]] = None) -> int: - if num_gen_requests <= 0: - return self.max_tier - if self.cost_table is None or self.cost_table.is_flat: - self.stats["fallback_flat_cost"] += 1 - return self.max_tier - selected = self._gather_rows(num_gen_requests=num_gen_requests, rows=rows) - if selected is None: - return self.max_tier - probs = self.apply_calibration(selected) - survival = compute_survival(probs).numpy().astype(np.float64) - return budget_argmax_over_uniform_lens( - survival=survival, - num_gen_requests=int(num_gen_requests), - cost_table=self.cost_table, - allowed_lens=self.tiers, - min_verify_len=self.cfg.min_verify_len, - ) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d26de67a719c..f99898d890da 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2867,6 +2867,28 @@ def validate_ragged_verify(self): "enable_ragged_verify requires enable_confidence_scheduling=True: " "per-request verify lengths are chosen from confidence-head " "survival, which is only computed when scheduling is on") + if self.enable_confidence_scheduling and not self.enable_ragged_verify: + # The two flags describe one feature. Scheduling without ragged used + # to mean a uniform tier ladder -- one window chosen for the whole + # batch -- which is a middle ground that does not earn its keep: + # it cannot give a confident request a longer window than a + # collapsing one, which is the entire point of scheduling against + # per-request confidence, and it degenerates to the full block + # whenever acceptance is high (every workload measured on + # DeepSeek-V4-Pro-DSpark). SGLang has no equivalent either: its + # scheduler always produces per-request windows via top-k, and + # uniform appears only as the verify-all degenerate case. + # + # So there are two states, not three: schedule per request, or + # verify the full block. Declining the feature is + # enable_confidence_scheduling=False. + raise ValueError( + "enable_confidence_scheduling requires enable_ragged_verify=True. " + "Scheduling a single window for the whole batch cannot act on " + "per-request confidence, and collapses to verifying the full " + "block whenever acceptance is high. To turn the feature off, " + "set enable_confidence_scheduling=False -- that verifies the " + "full drafted block, which is the only fallback.") if (self.enable_ragged_verify and not self.confidence_sps_table_path and not self.confidence_sts_path): # Without a profiled cost table the planner's budget degenerates to diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py index 4cc69ba6cc6e..842344b36a91 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py @@ -457,96 +457,6 @@ def _planner(**kw): return DSparkVerifyPlanner(**kw) -def test_planner_returns_max_tier_without_a_snapshot(): - """No confidence yet must mean verify-all, never a stale guess.""" - p = _planner() - assert p.decide_draft_len(num_gen_requests=8) == BLOCK - assert p.stats["fallback_no_snapshot"] == 1 - - -def test_planner_refuses_to_trim_on_an_unprofiled_cost_model(): - """A flat cost table makes every token look free -- trimming would be blind.""" - p = _planner(cost_table=SpsCostTable.flat()) - assert p.decide_draft_len(num_gen_requests=8) == BLOCK - assert p.stats["fallback_flat_cost"] == 1 - p2 = _planner(cost_table=None) - assert p2.decide_draft_len(num_gen_requests=8) == BLOCK - - -def test_planner_only_ever_returns_a_captured_tier(): - """Returning an uncaptured length silently drops the step out of the graph.""" - torch.manual_seed(5) - tiers = [2, 4] - for trial in range(12): - p = _planner(tiers=tiers) - conf = torch.rand(6, BLOCK) - p._host_buffer = conf.float() - p._copy_event = None - p._snapshot_valid = True - assert p.decide_draft_len(num_gen_requests=6) in tiers, f"trial {trial}" - - -def test_planner_trims_when_confidence_is_low(): - hi, lo = torch.full((8, BLOCK), 6.0), torch.full((8, BLOCK), -6.0) - out = [] - for logits in (hi, lo): - p = _planner() - p._host_buffer, p._copy_event, p._snapshot_valid = logits, None, True - out.append(p.decide_draft_len(num_gen_requests=8)) - assert out[1] <= out[0], f"low confidence should not verify more: {out}" - - -def test_planner_cross_rank_reduction_forces_agreement(): - """Ranks must converge on one length even with different local confidence.""" - peer_choices = {} - - def make(rank, logits): - p = _planner(all_rank_max=lambda v, r=rank: max(v, peer_choices.get(1 - r, v))) - p._host_buffer, p._copy_event, p._snapshot_valid = logits, None, True - return p - - r0 = _planner() - r0._host_buffer, r0._copy_event, r0._snapshot_valid = torch.full((4, BLOCK), -6.0), None, True - peer_choices[0] = r0._decide_local(num_gen_requests=4) - r1 = _planner() - r1._host_buffer, r1._copy_event, r1._snapshot_valid = torch.full((4, BLOCK), 6.0), None, True - peer_choices[1] = r1._decide_local(num_gen_requests=4) - - a = make(0, torch.full((4, BLOCK), -6.0)).decide_draft_len(num_gen_requests=4) - b = make(1, torch.full((4, BLOCK), 6.0)).decide_draft_len(num_gen_requests=4) - assert a == b, f"ranks disagreed on the graph key: {a} vs {b}" - assert a == max(peer_choices.values()) - - -def test_planner_cross_rank_result_is_snapped_back_onto_a_tier(): - """A reduction can land off-ladder; the result must still be captured.""" - p = _planner(tiers=[1, 3, BLOCK], all_rank_max=lambda v: 5) # 5 is not a tier - p._host_buffer, p._copy_event, p._snapshot_valid = torch.full((4, BLOCK), 0.0), None, True - assert p.decide_draft_len(num_gen_requests=4) == BLOCK - - -def test_planner_empty_batch_is_max_tier(): - assert _planner().decide_draft_len(num_gen_requests=0) == BLOCK - - -def test_planner_refuses_a_snapshot_that_does_not_cover_the_batch(): - """A short snapshot must fall back, never return a short answer. - - The staged snapshot lags one iteration, so a batch that grew since then has - more requests than the snapshot has rows. Returning one length per staged - row leaves the tail of the batch without a verify window -- and the two - consumers disagree about what that means: the input layout is built per - request (so it goes ragged) while the spec metadata sees a missing window - and stays uniform. That combination misattributes one request's drafts to - another with nothing raising. - """ - p = _planner() - p._host_buffer, p._copy_event, p._snapshot_valid = torch.zeros(4, BLOCK), None, True - assert p.decide_verify_lens(num_gen_requests=9) is None - assert p.stats["fallback_short_snapshot"] == 1 - assert p.decide_draft_len(num_gen_requests=9) == BLOCK - - def test_planner_returns_one_verify_len_per_request(): p = _planner() p._host_buffer, p._copy_event, p._snapshot_valid = torch.rand(6, BLOCK), None, True @@ -579,30 +489,6 @@ def test_planner_rejects_a_row_list_that_does_not_match_the_batch(): assert p.stats["fallback_short_snapshot"] == 1 -def test_planner_uses_the_supplied_calibration(): - """STS must be applied before the cumprod, not skipped.""" - seen = {} - - def calib(x): - seen["called"] = True - return torch.sigmoid(x / 4.0) - - p = _planner(apply_calibration=calib) - p._host_buffer, p._copy_event, p._snapshot_valid = torch.full((4, BLOCK), 2.0), None, True - p.decide_draft_len(num_gen_requests=4) - assert seen.get("called") - - -# -------------------------------------------------------------------------- -# Acceptance criterion (2): CUDA-graph compatibility -# -# The failure mode this guards is silent, not loud: a runtime draft length with -# no captured graph does not raise -- maybe_get_cuda_graph returns None and the -# step runs eager, costing far more than the trimmed tokens save. So the capture -# set and the planner's output set must be provably identical. -# -------------------------------------------------------------------------- - - def _capture_set(batch_sizes, tiers, max_draft_len=BLOCK): """Invoke PyTorchModelEngine._get_graphs_to_capture with a stub engine.""" import types @@ -639,132 +525,6 @@ def test_capture_is_the_batch_size_x_tier_cross_product(): assert len(graphs) == len(bss) * len(tiers) -def test_every_length_the_planner_can_pick_has_a_captured_graph(): - """The anti-silent-eager invariant, checked against the planner itself.""" - from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner - - bss, tiers = [1, 16, 64], [1, 3, BLOCK] - captured = {k for _bs, k in _capture_set(bss, tiers)} - - torch.manual_seed(11) - table = SpsCostTable(token_counts=(0, 8, 24, 64), step_time_ms=(2.0, 2.2, 4.0, 9.0)) - for trial in range(40): - p = DSparkVerifyPlanner(cfg=_cfg(min_verify_len=1), cost_table=table, tiers=tiers) - p._host_buffer = torch.rand(8, BLOCK) * 12 - 6 - p._copy_event, p._snapshot_valid = None, True - k = p.decide_draft_len(num_gen_requests=8) - assert k in captured, f"trial {trial}: planner picked K={k}, captured={captured}" - - -def test_planner_fallbacks_also_land_on_a_captured_graph(): - """Every degraded path (no snapshot, flat cost, empty batch) must stay captured.""" - from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner - - tiers = [1, 3, BLOCK] - captured = {k for _bs, k in _capture_set([1, 64], tiers)} - cases = [ - ( - "no snapshot", - DSparkVerifyPlanner( - cfg=_cfg(), - cost_table=SpsCostTable(token_counts=(0, 8), step_time_ms=(2.0, 5.0)), - tiers=tiers, - ), - 8, - ), - ( - "flat cost", - DSparkVerifyPlanner(cfg=_cfg(), cost_table=SpsCostTable.flat(), tiers=tiers), - 8, - ), - ( - "empty batch", - DSparkVerifyPlanner( - cfg=_cfg(), - cost_table=SpsCostTable(token_counts=(0, 8), step_time_ms=(2.0, 5.0)), - tiers=tiers, - ), - 0, - ), - ] - for name, p, n in cases: - assert p.decide_draft_len(num_gen_requests=n) in captured, name - - -# -------------------------------------------------------------------------- -# Acceptance criterion (3): attention-DP + TP -# -# draft_len is part of the CUDA-graph key but is NOT covered by the ADP -# consistency allgather, so ranks that choose differently select different -# graphs -- one replays, one falls back to eager -- and their collectives -# diverge. Agreement has to be forced before the forward. -# -------------------------------------------------------------------------- - - -def _simulated_ranks(local_confidences, tiers, table, batch_sizes=None): - """Run the planner on each rank with a real max-reduction across them.""" - from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner - - batch_sizes = batch_sizes or [c.shape[0] for c in local_confidences] - planners = [] - for conf in local_confidences: - p = DSparkVerifyPlanner(cfg=_cfg(min_verify_len=1), cost_table=table, tiers=tiers) - p._host_buffer, p._copy_event, p._snapshot_valid = conf, None, True - planners.append(p) - locals_ = [p._decide_local(num_gen_requests=bs) for p, bs in zip(planners, batch_sizes)] - peak = max(locals_) - return [ - p.decide_draft_len(num_gen_requests=bs, all_rank_max=lambda _v, m=peak: m) - for p, bs in zip(planners, batch_sizes) - ], locals_ - - -def test_adp_ranks_agree_despite_different_confidence(): - tiers = [1, 3, BLOCK] - table = SpsCostTable(token_counts=(0, 16, 48, 128), step_time_ms=(2.0, 2.2, 4.0, 9.0)) - confs = [ - torch.full((8, BLOCK), -6.0), - torch.full((8, BLOCK), 6.0), - torch.full((8, BLOCK), 0.0), - torch.full((8, BLOCK), 3.0), - ] - chosen, locals_ = _simulated_ranks(confs, tiers, table) - assert len(set(chosen)) == 1, f"ranks disagreed: {chosen} (locals {locals_})" - assert chosen[0] in tiers - - -def test_adp_ranks_agree_despite_different_batch_sizes(): - """Under ADP each rank has its own batch; agreement must survive that too.""" - tiers = [1, 3, BLOCK] - table = SpsCostTable(token_counts=(0, 16, 48, 128), step_time_ms=(2.0, 2.2, 4.0, 9.0)) - torch.manual_seed(3) - confs = [torch.rand(b, BLOCK) * 8 - 4 for b in (2, 9, 17, 33)] - chosen, locals_ = _simulated_ranks(confs, tiers, table, batch_sizes=[2, 9, 17, 33]) - assert len(set(chosen)) == 1, f"ranks disagreed: {chosen} (locals {locals_})" - - -def test_adp_rank_with_empty_batch_does_not_break_agreement(): - """An idle DP rank still has to enter the same graph as everyone else.""" - tiers = [1, 3, BLOCK] - table = SpsCostTable(token_counts=(0, 16, 48, 128), step_time_ms=(2.0, 2.2, 4.0, 9.0)) - confs = [torch.zeros(0, BLOCK), torch.full((8, BLOCK), 6.0), torch.full((8, BLOCK), -6.0)] - chosen, _ = _simulated_ranks(confs, tiers, table, batch_sizes=[0, 8, 8]) - assert len(set(chosen)) == 1, f"idle rank diverged: {chosen}" - - -def test_cross_rank_reduction_is_max_not_min(): - """Max is the safe direction: a rank that wanted less simply verifies more. - - Taking the min would let one pessimistic rank starve the whole batch, and -- - worse -- could drop below another rank's min_verify_len floor. - """ - tiers = [1, 3, BLOCK] - table = SpsCostTable(token_counts=(0, 16, 48, 128), step_time_ms=(2.0, 2.2, 4.0, 9.0)) - confs = [torch.full((8, BLOCK), -6.0), torch.full((8, BLOCK), 6.0)] - chosen, locals_ = _simulated_ranks(confs, tiers, table) - assert chosen[0] == max(locals_) - - def _stats(max_draft_len=5): from tensorrt_llm._torch.speculative.dspark_observability import ( DSparkRaggedStats, RaggedVerifyMode) diff --git a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py index 145a0b666268..bddba0454744 100644 --- a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py +++ b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py @@ -44,15 +44,32 @@ def test_ragged_with_an_sps_table_is_accepted(): @patch.dict(os.environ, {}, clear=False) -def test_uniform_confidence_scheduling_still_needs_no_table(): - """Only the ragged path is gated. +def test_confidence_scheduling_without_ragged_is_rejected(): + """There is no uniform tier ladder to fall back to. - Uniform tier selection degrades gracefully without a table -- it picks the - max tier, which is the pre-existing behaviour -- so requiring one there - would break working configurations. + Scheduling used to have a middle state: pick one verify length for the + whole batch from the captured tiers. It could not act on per-request + confidence -- which is the entire point of a confidence head -- and it + degenerated to the full block whenever acceptance was high, i.e. on every + workload measured on this checkpoint. It is gone, so this combination now + names a path that does not exist and must not silently mean something else. + + The two remaining states are: schedule per request (ragged), or verify the + whole drafted block (confidence scheduling off). """ + with pytest.raises(ValueError, match="enable_ragged_verify"): + DSparkDecodingConfig(max_draft_len=5, + speculative_model="/nonexistent/model", + enable_confidence_scheduling=True, + enable_ragged_verify=False) + + +@patch.dict(os.environ, {}, clear=False) +def test_scheduling_off_needs_no_table_and_verifies_everything(): + """The fallback is static verification, which has nothing to profile.""" cfg = DSparkDecodingConfig(max_draft_len=5, speculative_model="/nonexistent/model", - enable_confidence_scheduling=True, + enable_confidence_scheduling=False, enable_ragged_verify=False) + assert not cfg.enable_confidence_scheduling assert not cfg.enable_ragged_verify From 95484a25b7ba5ff720113d3ef41eb8318409253f Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 21:22:06 -0700 Subject: [PATCH 063/177] [None][doc] Record that the uniform tier ladder was removed, not just declined Section 2 presented the uniform tier ladder as one of two candidate deliverables and recorded the decision to build ragged instead. It has since been deleted outright, so the doc read as if a live alternative still existed. Adds 10.9 with the two measured reasons: it cannot give a confident request a longer window than a collapsing one, and on this checkpoint the argmax over tiers {1,3,5} always lands on 5, i.e. the full block -- so it was a distinct code path (graph key, cross-rank allgather, planner counters) that behaved exactly like scheduling=False. Also states plainly what this does NOT invalidate, and why: - the 96.2092 uniform GSM8K baseline stands, but not because that run used scheduling=False -- I did not re-check its flags. It stands because 10.3 measured the planner declining to trim in all three workloads, so the middle state necessarily picked the full block there too. - the 10.8 sync control is "confidence head off", which is still a legal configuration and needs no redesign. Signed-off-by: Lance Liao --- docs/dspark_p0_task_prompt.md | 44 +++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/docs/dspark_p0_task_prompt.md b/docs/dspark_p0_task_prompt.md index e06bfcef35e1..3d086502d481 100644 --- a/docs/dspark_p0_task_prompt.md +++ b/docs/dspark_p0_task_prompt.md @@ -191,6 +191,9 @@ Q2 原本问的是：本轮交付「均匀 tier 阶梯」还是「ragged」。 | 收益 | 拿满 SPS(B) 项，**损失 τ 项** | 完整 Θ = τ·SPS(B) | | 风险 | 中。无 K2 假设 | **高**。押在 K2（trtllm-gen sparse-MLA generation）这个读码无法证明的行为上 | +> **后续（见 §10.9）**：均匀 tier 阶梯不只是「没选它」，而是**已从代码里整个删除**。 +> 它在实测中退化成「验证整块」，却仍占着一条独立代码路径。下表保留作决策记录。 + **已定：做 ragged。** 因此： 1. goal doc §5.1 里划给 **P2 的 H6 / H7 / H9 / H10 全部上提到 P0**（见 §3 表二）。 @@ -830,3 +833,44 @@ planner 在对着虚构的成本曲线优化,所以那轮的吞吐数字无意 * **`peer_not_gen_only` 类 graph miss**:32/318 步,属正常连续批处理(uniform 路径同样存在),非缺陷,但可作为后续调度优化的方向。 * **吞吐收益需要更低接受率的模型或负载**。这是 checkpoint 的属性,非本 PR 可解。 + +### 10.9 删除均匀 tier 阶梯:三状态收敛成两状态 + +调度原本有三种状态,中间那种不值得保留: + +| 配置 | 语义 | 结局 | +|---|---|---| +| `scheduling=False` | 验证整个 drafted block | **保留**,唯一兜底 | +| `scheduling=True, ragged=False` | 整批取同一个 verify 长度(tier `{1,3,5}` 里选) | **已删除** | +| `scheduling=True, ragged=True` | 每请求各自的窗口 | **保留** | + +删除理由有两条,都是实测出来的: + +1. **它做不到调度的本意**。整批一个长度,就无法把长窗口给高置信请求、短窗口给 + 坍塌的请求——而那正是 confidence head 存在的理由。 +2. **它在本 checkpoint 上恒等于「验证整块」**。阶梯只有 `{1,3,5}` 三档,接受率 + ~90% 时 argmax 永远落在 5,即整块。也就是说:一条独立代码路径(进 graph key、 + 进跨 rank allgather、进 planner 的 fallback 计数),运行效果和 `scheduling=False` + 完全一样。 + +SGLang 也没有对应物:它的 scheduler 一律经 top-k 产出每请求窗口,"均匀"只作为 +verify-all 的退化情形出现。 + +删掉的东西:`decide_draft_len` 及其 `_decide_local`、`snap_to_tier`(唯一调用者 +就是 uniform 分支)、`py_executor` 里的 uniform 分支、跨 rank allgather 里的 tier +字段(5→4 个 int)、以及覆盖该路径的 12 个测试和一个 rank 模拟 helper。共 +−350 行。兜底现在直接是 `runtime_draft_len = planner.max_tier`,不需要跨 rank +协商,所以这一步也少了一次归约。 + +`llm_args` 现在**直接拒绝** `scheduling=True + ragged=False`,而不是把它悄悄当成 +别的东西。关掉特性用 `enable_confidence_scheduling=False`。 + +> **历史数据是否受影响**:§10.2 的 uniform 基线 96.2092 仍然有效,但理由不是 +> 「它没用中间态」——我没有回去核对那一轮的 flag(§2 表一把中间里程碑定义成 +> `scheduling=True + ragged=False`,所以有可能就是它)。理由是 §10.3 已实测 +> planner 在三种工况下**一次都没裁剪**:中间态在本 checkpoint 上必然选到 tier 5, +> 即整块。因此无论那一轮用的是哪套 flag,96.2092 描述的都是「验证整块」,而那条 +> 路径原封不动地留着。 +> +> §10.8 的 sync 对照写的是「confidence head 关」= `scheduling=False`,配置依然合法, +> 不需要重新设计。 From 0f94c044ff1e37ec4e1e31ea56421ae6bc44d6fe Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 22:40:17 -0700 Subject: [PATCH 064/177] [None][feat] Implement the cap-accept ragged verify mode cap-accept schedules per-request verify windows and commits against them, but does not shrink the token axis handed to the target: the full block is scored and the positions past each window are discarded. It costs the same compute as static, which is the point -- it is a reference, not a serving mode. It splits a question compact conflates. compact changes both what gets committed and how the batch is laid out, so when its output moves, nothing says which half moved it: cap-accept output != static output -> the scheduling policy is wrong compact output != cap-accept -> the layout compaction is wrong It also produces a number compact structurally cannot: cap_trim_tokens, the positions the target accepted and the window then discarded -- the exact acceptance cost of trimming. Under compact those positions are never scored, so that cost can only be bounded (trim_regret_rate), never measured. SGLang keeps a ~600-line statistical estimator for the same quantity. The implementation is purely additive: the window goes on a NEW attribute, py_verify_cap, and py_verify_len is left None. Every layout consumer keys off py_verify_len, so model_engine, cuda_graph_runner and dsa.py see an untrimmed batch with no changes at all -- and _verified_len falls through to the batch value, which makes the KV rewind (full block - capped) correct for free. That avoids threading a second set of conditions through the accounting that took three bugs to get right for compact; the KV leak fixed earlier on this branch was exactly that class of error. The cap itself is SpecSamplerBase._apply_verify_cap. num_new_tokens is the uncapped chain (accepted + 1), so a window of w admits w + 1 tokens. When the cap bites the last token kept is draft[w], which verification accepted, rather than the target's own sample at the cut as in SGLang's cap_correct_len -- identical under greedy, same committed count, lossless either way. Two traps in the counters, both now covered by tests: - record_step derives delivered tokens from the windows, which would credit cap-accept with a compute saving it deliberately does not take, and trim_ratio is the headline "did this work" number. Added an explicit delivered argument; the caller passes the full block. - assert_ragged_active(require_trim=True) would fail every cap-accept run, since a zero trim ratio is its designed outcome. That clause is now skipped for modes that do not trim the submitted tokens, and still fires for compact. 51 new tests. Not yet exercised end to end. Signed-off-by: Lance Liao --- docs/dspark_confidence_schedule_goal.md | 6 +- docs/dspark_p0_task_prompt.md | 39 ++++ docs/dspark_ragged_verify_status.md | 26 ++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 44 +++- tensorrt_llm/_torch/speculative/dspark.py | 10 - .../speculative/dspark_observability.py | 62 +++++- .../_torch/speculative/spec_sampler_base.py | 58 +++++- .../hw_agnostic/test_dspark_cap_accept.py | 193 ++++++++++++++++++ 8 files changed, 409 insertions(+), 29 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py diff --git a/docs/dspark_confidence_schedule_goal.md b/docs/dspark_confidence_schedule_goal.md index 2364729e405a..fdb11446f73f 100644 --- a/docs/dspark_confidence_schedule_goal.md +++ b/docs/dspark_confidence_schedule_goal.md @@ -74,7 +74,7 @@ PR 的 host 侧管线质量是高的，不要推翻重做： | `compressor_paged_kv_compress` 新增 `Tensor? new_tokens_per_seq=None` | **已实现** | `compressorOp.cpp:153`（schema）、`compressorKernels.cu:340-348`（`nn` 取代 `NEXT_N`，`NEXT_N` 降级为编译期上界）、`:425`（Phase 2）；Phase 1 由 `token_idx < kv_len` 守卫（`:369`）、Phase 3 由 `num_compressions` 守卫（`:440-442`），二者对 `nn < NEXT_N` 均正确 | | `num_gen_tokens_per_seq` ragged 时取 **max** 而非 mean | **已实现** | `deepseek_v4.py:748-769` | | `gen_new_tokens_per_seq` 每请求向量 + graph-stable buffer | **已实现** | `deepseek_v4.py:337-352` | -| ragged verify mode 开关（static / cap-accept / compact）+ 统计 | **部分实现**，`cap-accept` 显式 `NotImplementedError` | `dspark.py:457-496`、`dspark_observability.py` | +| ragged verify mode 开关（static / cap-accept / compact）+ 统计 | **已全部实现**（`cap-accept` 见 `spec_sampler_base._apply_verify_cap` + `py_verify_cap`） | `dspark.py:457-496`、`dspark_observability.py` | | C++ ragged top-k 单测 | **已新增** | `tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py` | **重要副作用**：`num_gen_tokens_per_seq` 改成取 max 之后，compressor 输出 buffer 的 sizing 公式 `num_generations * ceil(w/cr)`（`deepseek_v4.py:753-757`、`:766-768`）重新成为**合法上界**（因为 `w_r ≤ w_max` ⇒ `Σ_r ceil(w_r/cr) ≤ n·ceil(w_max/cr)`）。这消除了「越界写 `kv_comp`」这一类内存破坏风险。取 mean 时该公式**不是**上界，是真实的 OOB 写。 @@ -318,7 +318,7 @@ cd /code/tensorrt_llm && python -m pytest \ | 测试 | 内容 | 为什么锐利 | |---|---|---| | **A1 逐 token 等价** | 同一 engine 配置、同一批 prompt，`static`（verify 全量）与开启调度后，temperature=0 下断言每个 prompt 的 `token_ids` **完全相等** | 贪婪投机解码定义上是 lossless，任何分歧都是 bug。把 6 题的容忍度变成 1 token 的容忍度。**注意**：跨不同 graph 形状的严格 bitwise 相等 TRT-LLM 并不保证（MoE/attention 规约顺序会变），所以诚实的形式是允许极少数 prompt 分歧并设硬上限，或用 `cap-accept` 模式（下条）。 | -| **A2 `cap-accept` 差分** | 实现 SGLang 的三模式（`static` / `cap-accept` / `compact`）。`cap-accept` 跑**均匀 kernel 路径**（`next_n=6`）但只提交每请求窗口内的 token，输出应与 `compact` **逐 token 相同** | `cap-accept` vs `compact` 的差异 ⇒ **必然是 ragged kernel bug**；`static` vs `cap-accept` 的差异 ⇒ **必然是调度/接受逻辑 bug**。把 pass/fail 变成可定位。工作区的 `dspark_observability.py` 已有 mode 枚举，`cap-accept` 目前显式 `NotImplementedError`（`dspark.py:471-478`）—— **这是需要补的**。 | +| **A2 `cap-accept` 差分** | 实现 SGLang 的三模式（`static` / `cap-accept` / `compact`）。`cap-accept` 跑**均匀 kernel 路径**（`next_n=6`）但只提交每请求窗口内的 token，输出应与 `compact` **逐 token 相同** | `cap-accept` vs `compact` 的差异 ⇒ **必然是 ragged kernel bug**；`static` vs `cap-accept` 的差异 ⇒ **必然是调度/接受逻辑 bug**。把 pass/fail 变成可定位。**已实现**：窗口写在独立属性 `py_verify_cap` 上（不碰 `py_verify_len`），布局侧因此看不见窗口、自动走整块；截断在 `SpecSamplerBase._apply_verify_cap`。 | | **A3 布局一致性断言** | host 侧断言：`sum(1 + py_verify_len over padded generation_requests) == attn_metadata.num_tokens - num_ctx_tokens == spec_metadata.total_verify_tokens == key[bucket]` | 这一条就能同时抓住 H0、H5、H8 | | **A4 时序断言** | 断言 `attn_metadata.ragged_verify_lens` 在 `prepare()` 执行的那一刻是**本步**的（而非上一步） | 抓 H2、H4 | | **A5 接受长度分布** | 记录并断言每请求接受长度直方图，以及 `sum(accepted)` 与接受 kernel 报告一致 | 抓「token 落到了错误请求的槽位但整批仍生成合法文本」这一类 | @@ -386,7 +386,7 @@ TRT-LLM 的 `dsa.py` **从不传 `indices`**，而是用既有的 "expanded buff 3. **工作区那 11 个未提交的改动如何处置？** K4/K9 的实现质量是好的（可选参数、非 ragged bit-identical、TORCH_CHECK 替换而非放宽），但 H3/H4 是新引入的 P0 bug。建议：先修 H3/H4 再提交，**不要**在带着必抛异常的代码的状态下推进其他工作。 -4. **`cap-accept` 模式要不要实现？** 它是把 GSM8K 从「pass/fail 硬币」变成「可定位诊断」的关键（§6.2 A2），成本不高（复用 PR 已有的 `torch.minimum(..., verify_lens)` clamp，`interface.py:1473-1480`、`:1889-1896`）。当前是显式 `NotImplementedError`。 +4. ~~**`cap-accept` 模式要不要实现？**~~ **已实现**（§6.2 A2）。它把 GSM8K 从「pass/fail 硬币」变成可定位诊断，并产出 `compact` 结构上拿不到的 `cap_trim_tokens`（裁剪的真实接受代价）。 5. **STS 表和 SPS 表由谁、按什么流程产出？** 这是 U4/U5，也是 G8。在它们存在之前，任何**性能**结论都是空的，而且 feature 在构造上无法 trim。需要一个 profiling 脚本 + 产物的存放约定（checkpoint 内？repo 内 `examples/configs/`？）。 diff --git a/docs/dspark_p0_task_prompt.md b/docs/dspark_p0_task_prompt.md index 3d086502d481..c16f0c3cce8f 100644 --- a/docs/dspark_p0_task_prompt.md +++ b/docs/dspark_p0_task_prompt.md @@ -874,3 +874,42 @@ verify-all 的退化情形出现。 > > §10.8 的 sync 对照写的是「confidence head 关」= `scheduling=False`,配置依然合法, > 不需要重新设计。 + +### 10.10 实现 `cap-accept`(A2 / Q4 收口) + +三模式现已全部可用,`dspark.py` 里的 `NotImplementedError` 已移除。 + +`cap-accept` 算出每请求窗口并**按窗口提交**,但**不缩小交给 target 的 token 轴**。 +它把 `compact` 混在一起的两件事拆开: + + cap-accept 输出 ≠ static 输出 ⇒ 调度策略有问题 + compact 输出 ≠ cap-accept 输出 ⇒ 布局压缩有问题 + +**实现是纯增量的**:窗口写在独立属性 `py_verify_cap` 上,**不碰 `py_verify_len`**。 +于是布局侧(`model_engine` / `cuda_graph_runner` / `dsa.py`)完全看不见窗口、自动 +走整块;`_verified_len` 因为 `py_verify_len is None` 自动返回整块,**KV rewind 随之 +正确**(`rewind = 整块 − capped`)。这样就不必在「为 compact 调了三个 bug 才对」的 +记账代码里再加一组条件分支——§10.6 那个 KV leak 正是这类记账错误。 + +截断在 `SpecSamplerBase._apply_verify_cap`:`num_new_tokens` 是未截断链 +(`accepted + 1`),窗口 `w` 允许 `w + 1` 个 token。截断后末位是 `draft[w]`,它已通过 +验证,提交合法。与 SGLang 取「cut 处 target 自己的采样」来源不同,**贪婪下相同**, +两者提交 token 数一致且都无损。 + +**它产出 `compact` 结构上拿不到的数**:`cap_trim_tokens` = 被丢弃的、本可接受的位置 +数,即裁剪的真实接受代价。`compact` 下那些位置根本没算,SGLang 为此专门养了一个 +~600 行的统计估计器;我们此前只有 `trim_regret_rate` 这个下界。 + +两个埋在计数器里的坑,都已处理并有测试: + +1. **不能给它记上算力节省**。`record_step` 会按 `sum(1+v)` 推算 delivered,那会让 + cap-accept 显示出一个它刻意不去拿的 trim_ratio——而那正是判断功能是否生效的 + 头号指标。新增显式 `delivered` 参数,由调用方传整块。 +2. **`assert_ragged_active(require_trim=True)` 会误杀它**。它的 `trim_ratio` 恒为 0 + 是设计如此。该判据现在对不裁 token 轴的模式豁免,否则这道门在它唯一该跑的模式 + 里恒不可满足。 + +代价是**不省任何算力**,所以它是诊断模式,永远不是服务配置。 + +覆盖:`tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py` +(51 项,含截断算术的全枚举不变式、两个计数器坑的回归)。**尚未 e2e 跑过。** diff --git a/docs/dspark_ragged_verify_status.md b/docs/dspark_ragged_verify_status.md index 1e568f2407a6..f307ccea5795 100644 --- a/docs/dspark_ragged_verify_status.md +++ b/docs/dspark_ragged_verify_status.md @@ -150,7 +150,7 @@ P0-3/5/6 属于「跑出来的数字看着正常但其实是错的」：replay | 开关 | 作用 | |---|---| -| `TLLM_DSPARK_RAGGED_VERIFY_MODE=static\|compact` | 选路径。`cap-accept` **未实现，显式抛 `NotImplementedError`** | +| `TLLM_DSPARK_RAGGED_VERIFY_MODE=static\|cap-accept\|compact` | 选路径。三种模式**均已实现** | | ~~`TLLM_DSPARK_FORCE_VERIFY_LENS=1`~~ | **已删除**。改用 `tests/microbenchmarks/dspark_make_steep_sps_table.py` 造陡峭 cost table,让**真实 planner** 自己产出不等窗口 | | `DSparkRaggedStats.assert_ragged_active()` | 断言 ragged 真的生效 | @@ -163,8 +163,26 @@ P0-3/5/6 属于「跑出来的数字看着正常但其实是错的」：replay `bs_not_supported`、`peer_shape_mismatch` 三类失败——它们分别表示 bucket 网格错了、 批次档位没捕获、以及跨 rank 形状发散。 -**`cap-accept` 为什么不能 alias 成 `compact`**：它唯一的价值是「输出必须与 `static` -逐 token 一致」，别名化会让这个对照失去意义。 +**`cap-accept`（已实现）**:算出每请求窗口并**按窗口提交**,但**不缩小交给 target 的 +token 轴**——整块照算,超出窗口的接受位置被丢弃。它绝不能 alias 成 `compact`:唯一 +价值就是那个对照。 + +实现方式是把窗口写在**独立属性 `py_verify_cap`** 上,不碰 `py_verify_len`。于是布局侧 +(`model_engine` / `cuda_graph_runner` / `dsa.py`)完全看不见窗口,自动走整块; +`_verified_len` 因为 `py_verify_len is None` 自动返回整块,**KV rewind 随之正确** +(`rewind = 整块 - capped`)。这样新模式是纯增量,而不是在「为 compact 调了三个 bug +才对」的记账代码里再加一组条件分支。 + +截断在 `SpecSamplerBase._apply_verify_cap`:`num_new_tokens` 是未截断链 +(`accepted + 1`),窗口 `w` 允许 `w + 1` 个 token。截断后末位是 `draft[w]`——它**已通过 +验证**,提交合法。与 SGLang 取「cut 处 target 自己的采样」来源不同,贪婪下相同,两者 +提交 token 数一致且都无损。 + +它产出 `compact` 结构上拿不到的数:`cap_trim_tokens`,即**被丢弃的、本可接受的位置数** +——裁剪的真实接受代价。`compact` 下那些位置根本没算,只能用 `trim_regret_rate` 卡下界。 + +代价是**不省任何算力**,所以它是诊断模式,永远不是服务配置。`assert_ragged_active()` +的 `require_trim` 因此对该模式豁免,否则那条门在它唯一该跑的模式里恒不可满足。 **「循环依赖」这个理由已被证伪**:产出 cost table **不需要** ragged 能跑。 `dspark_sps_profiler` 会把 `RAGGED_VERIFY_MODE` 钉成 `static`,并通过逐个 verify_len @@ -190,7 +208,7 @@ confidence 排序的 top-k 分配。三个只在真裁剪时可达的 bug 正是 | P1 | **SPS cost table profiler** | 没有它 planner 必然不 trim，**任何性能结论都是空的**。参考 SGLang `python/sglang/benchmark/dspark_sps_profiler.py`：服务端 `RAGGED_VERIFY_MODE=static` + record flag，扫 batch size 记 step time | | P1 | GSM8K + ADP | 冒烟通过后才有意义 | | P1 | graph 按 token 数分桶 + ADP tier 协商 | SGLang `compute_target_verify_graph_key` 对 ragged 返回 `(graph_num_tokens, graph_num_tokens)`——两个轴都变 token 数；DP 下各 rank 共用同一 tier 并一起降档 | -| P2 | 实现 `cap-accept` | 隔离「kernel 算错」与「planner 没触发」的最强工具 | +| ~~P2~~ | ~~实现 `cap-accept`~~ | **已完成**。隔离「kernel 算错」与「planner 没触发」 | | P2 | 对照剩余 53 个 sparse 失败 | 目前只对照了 mqa 那 6 个 | --- diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index c979f50ec380..843beda6b397 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3463,12 +3463,35 @@ def _dspark_confidence_draft_len(self, # on different ranks (the pad-row fit is sized from the local request # count), which is why the attention-DP graph gate compares the token # bucket and not just the batch size. + # + # `cap-accept` stops short of this: it wants the schedule's commitment + # policy without the layout that pays for it, so it submits the full + # block and needs no bucket at all. The windows go onto `py_verify_cap` + # rather than `py_verify_len` precisely so the whole token-layout path + # below and downstream (input layout, DSA row expansion, graph key, KV + # rewind) keeps seeing an untrimmed batch -- the mode is then a pure + # addition rather than a second set of conditions inside code that took + # three bugs to get right for `compact`. bucket = None + cap_accept = (mode is not None and mode.computes_windows + and not mode.trims_submitted_tokens) + if cap_accept: + # Stale caps are the same hazard as stale windows, minus the one + # thing that makes stale windows survivable: a mismatched layout + # eventually trips an assert, whereas a stale cap just commits + # against last step's schedule with nothing to contradict it. + # Clear every request first, set second. + for request in gen_requests: + request.py_verify_cap = None if ragged_lens is not None: - bucket = self.model_engine.fit_ragged_verify_lens( - gen_requests, ragged_lens, peer_stats=peer_stats) - if bucket is None: - fallback_reason = "no_captured_shape" + if cap_accept: + for request, window in zip(gen_requests, ragged_lens): + request.py_verify_cap = int(window) + else: + bucket = self.model_engine.fit_ragged_verify_lens( + gen_requests, ragged_lens, peer_stats=peer_stats) + if bucket is None: + fallback_reason = "no_captured_shape" # Record what the step actually decided, reading the windows back off # the requests rather than off `ragged_lens`: fit_ragged_verify_lens @@ -3483,7 +3506,15 @@ def _dspark_confidence_draft_len(self, stats.log_summary(prefix="DSpark ragged verify [periodic]") self._maybe_assert_dspark_ragged_active(stats) fitted = None - if bucket is not None: + delivered = None + if cap_accept: + # Nothing was fitted, so there is nothing to read back: the + # scheduled windows ARE the decision. The block is submitted + # whole, which `delivered` has to say explicitly -- otherwise + # the windows would be booked as if they had been submitted. + fitted = [int(v) for v in ragged_lens] if ragged_lens else None + delivered = num_gen_requests * (1 + max_draft_len) + elif bucket is not None: fitted = [ int(request.py_verify_len) for request in gen_requests if getattr(request, "py_verify_len", None) is not None @@ -3494,7 +3525,8 @@ def _dspark_confidence_draft_len(self, verify_lens=fitted, bucket=bucket, padded_bs=None, - fallback=fallback_reason) + fallback=fallback_reason, + delivered=delivered) if bucket is not None: # The block is always drafted in full -- only verification is diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 1a945f39ea0c..6ea4bb7b45c0 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -465,16 +465,6 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> self.spec_config, "enable_ragged_verify", False) else RaggedVerifyMode.STATIC) self.ragged_verify_mode = read_ragged_verify_mode(default=configured) - if self.ragged_verify_mode is RaggedVerifyMode.CAP_ACCEPT: - # Deliberately not aliased onto `compact`. The whole value of - # cap-accept is that its output must be bit-identical to `static` - # while the windows still drive commitment; silently running it as - # `compact` would make that comparison meaningless. - raise NotImplementedError( - "TLLM_DSPARK_RAGGED_VERIFY_MODE='cap-accept' is not implemented " - "yet: it requires submitting the full block to the target while " - "committing only each request's window. Use 'static' or " - "'compact'.") self.ragged_stats = DSparkRaggedStats(mode=self.ragged_verify_mode, max_draft_len=block_size) # goal doc A6: the ragged counters say what happened, the planner's say diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index 632de3018190..180f2c0d857b 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -193,6 +193,12 @@ def __init__(self, *, mode: RaggedVerifyMode, max_draft_len: int): #: from the untrimmed subpopulation (which the planner selected, so it #: is not a fair sample of the trimmed one). self.trimmed_hit_ceiling = 0 + #: Positions the target accepted and the window then discarded -- + #: the exact acceptance the schedule cost. Only ``cap-accept`` can + #: populate this, because only there are those positions scored at all; + #: it stays 0 under ``compact``, where the honest answer is that the + #: number is unknowable without re-running. + self.cap_trim_tokens = 0 #: Graph replay is the whole point of landing on a bucket; an eager #: step costs far more than the tokens it saved. @@ -210,6 +216,7 @@ def record_step( bucket: Optional[int] = None, padded_bs: Optional[int] = None, fallback: Optional[str] = None, + delivered: Optional[int] = None, ) -> None: """Record one decode step's scheduling decision. @@ -221,6 +228,13 @@ def record_step( bucket: the captured token bucket the batch was fitted to, if any. padded_bs: the row count the batch was padded to, if any. fallback: short reason the step did not go ragged. + delivered: tokens actually handed to the target, when that is not + derivable from ``verify_lens`` and ``bucket``. Required by + ``cap-accept``, which schedules windows but submits the full + block: deriving it there would report the windows as if they + had been submitted, i.e. credit the mode with a compute saving + it deliberately does not take, and ``trim_ratio`` would then + describe a run that never happened. """ self.steps_total += 1 # Steps where raggedness is even definable. A window can only differ @@ -248,9 +262,14 @@ def record_step( self.verify_len_hist.update(lens) self.window_tokens += sum(1 + v for v in lens) # A bucket is what the graph actually ran; without one the step ran on - # exactly the windows it asked for. - self.delivered_tokens += int(bucket) if bucket else sum(1 + v - for v in lens) + # exactly the windows it asked for -- unless the caller says otherwise, + # which is how cap-accept reports submitting the full block. + if delivered is not None: + self.delivered_tokens += int(delivered) + else: + self.delivered_tokens += int(bucket) if bucket else sum(1 + v + for v in + lens) if bucket: self.bucket_hist[int(bucket)] += 1 if padded_bs: @@ -312,15 +331,26 @@ def distinct_verify_lens(self) -> int: """How many different window sizes were ever handed out.""" return len(self.verify_len_hist) - def record_acceptance(self, *, accepted: int, window: int) -> None: + def record_acceptance(self, + *, + accepted: int, + window: int, + cap_trim: int = 0) -> None: """Record one request's acceptance against the window it was given. Args: accepted: drafted positions accepted (excludes the bonus token). window: drafted positions the request was allowed to verify. + cap_trim: positions the target accepted but the window discarded. + Non-zero only in ``cap-accept``, the one mode where those + positions are actually scored, so this is a measurement rather + than an estimate. Zero elsewhere -- under ``compact`` they were + never computed and cannot be counted, only bounded from below + by ``trimmed_hit_ceiling``. """ self.requests_scored += 1 self.accepted_tokens += int(accepted) + self.cap_trim_tokens += int(cap_trim) if int(window) < self.max_draft_len: self.requests_trimmed += 1 if int(accepted) >= int(window): @@ -333,6 +363,20 @@ def accept_len(self) -> float: return 0.0 return self.accepted_tokens / self.requests_scored + @property + def accept_loss_per_request(self) -> float: + """Mean accepted positions the schedule discarded, per request. + + The exact counterpart to :attr:`accept_len`: add the two and you get + the acceptance a no-trim run would have achieved on the same drafts. + Meaningful only under ``cap-accept``; 0.0 elsewhere, where it means + "not measured", not "no loss" -- ``trim_regret_rate`` is the bound + available under ``compact``. + """ + if not self.requests_scored: + return 0.0 + return self.cap_trim_tokens / self.requests_scored + @property def trim_regret_rate(self) -> float: """Share of trimmed requests whose draft was still alive at the cut. @@ -370,6 +414,8 @@ def summary(self) -> Dict[str, object]: "requests_scored": self.requests_scored, "requests_trimmed": self.requests_trimmed, "trim_regret_rate": round(self.trim_regret_rate, 4), + "cap_trim_tokens": self.cap_trim_tokens, + "accept_loss_per_request": round(self.accept_loss_per_request, 4), "graph_replays": self.graph_replays, "graph_eager": self.graph_eager, "fallbacks": dict(self.fallbacks), @@ -412,7 +458,13 @@ def assert_ragged_active(self, *, require_trim: bool = True) -> None: raise AssertionError( f"DSpark ragged verify handed out only one distinct window " f"size, so the batch was uniform in all but name: {summary}") - if require_trim and self.trim_ratio <= 0.0: + # `cap-accept` submits the full block on purpose -- it measures the + # commitment policy, not a compute saving -- so a zero trim ratio is + # the expected outcome there, not a silent no-op. Requiring one would + # make the gate unsatisfiable in the one mode whose entire job is to be + # the trusted reference the other modes are compared against. + if (require_trim and self.mode.trims_submitted_tokens + and self.trim_ratio <= 0.0): raise AssertionError( f"DSpark ragged verify delivered as many tokens as a no-trim " f"run, so nothing was saved: {summary}") diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index bfaec06dc25f..c801eba57dea 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -61,6 +61,13 @@ class SampleStateSpec(SampleState): #: reading ``request.py_verify_len`` live rewinds by the wrong amount under #: the overlap scheduler. None on every uniform path. verify_lens_snapshot: Optional[dict] = None + #: Same idea, for ``cap-accept``: the window the scheduler chose, which + #: bounds *commitment* only -- the target still scored the full block, so + #: the KV rewind must use the full block and not this. Kept in a separate + #: field from ``verify_lens_snapshot`` for exactly that reason: the two + #: never coexist, and conflating them would rewind by the window after + #: verifying the block, silently leaking KV. + verify_caps_snapshot: Optional[dict] = None class SpecSamplerBase(Sampler[SampleStateSpec], AsyncWorkerMixin): @@ -212,6 +219,35 @@ def _verified_len(request: LlmRequest, runtime_draft_len: int, per_request = getattr(request, "py_verify_len", None) return runtime_draft_len if per_request is None else int(per_request) + @staticmethod + def _apply_verify_cap(num_new_tokens: int, + cap: Optional[int]) -> tuple[int, int]: + """Trim a commit to the scheduler's window. ``cap-accept`` only. + + The target scored the whole drafted block, so ``num_new_tokens`` is the + UNCAPPED chain: ``accepted + 1``, the accepted drafts plus the bonus at + the first unaccepted position. A window of ``cap`` drafted positions + therefore admits ``cap + 1`` tokens. + + Truncating there keeps every committed token one the target already + verified. When the cap bites, the last token kept is ``draft[cap]``, + which was accepted, rather than the target's own sample at that + position -- a different provenance from SGLang's ``cap_correct_len``, + identical under greedy, and lossless either way, since a token is only + committed if verification accepted it. The committed COUNT matches. + + Returns: + (tokens to commit, positions the cap discarded). The second value + is the exact acceptance the schedule cost, which is the number + ``compact`` cannot produce because it never scores those positions. + """ + if cap is None: + return num_new_tokens, 0 + admitted = int(cap) + 1 + if num_new_tokens <= admitted: + return num_new_tokens, 0 + return admitted, num_new_tokens - admitted + def update_requests( self, state: SampleStateSpec, @@ -234,10 +270,16 @@ def update_requests( beam_idx = DEFAULT_BEAM_IDX runtime_draft_len = getattr(state, "runtime_draft_len", self.draft_len) + caps = state.verify_caps_snapshot for req in state.requests: if req.state == LlmRequestState.GENERATION_COMPLETE: continue num_new_tokens = new_tokens_lens_list[req.py_seq_slot] + # cap-accept only: the target scored the whole block, so this is the + # uncapped chain and the scheduler's window has to be applied here. + cap = None if caps is None else caps.get(req.py_request_id) + num_new_tokens, cap_trim = self._apply_verify_cap( + num_new_tokens, cap) for i in range(num_new_tokens): new_token = add_token(req, new_tokens, beam_idx=beam_idx, step=i) if TorchSampler._handle_stop_criteria( @@ -254,10 +296,17 @@ def update_requests( # has already moved on by now, which is why the snapshot exists. # `acceptance_stats` is attached by the DSpark worker and is None # for every other path. + # + # Under cap-accept the window is the cap, NOT `verified_len`: the + # target verified the whole block (which is what makes the rewind + # above correct), but the scheduler only granted `cap`. Reporting + # the block here would count every request as untrimmed and hide + # the mode's entire measurement. if self.acceptance_stats is not None: self.acceptance_stats.record_acceptance( accepted=req.py_num_accepted_draft_tokens, - window=verified_len) + window=verified_len if cap is None else int(cap), + cap_trim=cap_trim) self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) def sample_async( @@ -350,12 +399,18 @@ def sample_async( # Cheap (a few ints) and only populated when the ragged path is live. verify_lens_snapshot = None + verify_caps_snapshot = None for request in sampling_requests: window = getattr(request, "py_verify_len", None) if window is not None: if verify_lens_snapshot is None: verify_lens_snapshot = {} verify_lens_snapshot[request.py_request_id] = int(window) + cap = getattr(request, "py_verify_cap", None) + if cap is not None: + if verify_caps_snapshot is None: + verify_caps_snapshot = {} + verify_caps_snapshot[request.py_request_id] = int(cap) return SampleStateSpec( requests=sampling_requests, @@ -364,4 +419,5 @@ def sample_async( sampler_event=sampler_event, runtime_draft_len=runtime_draft_len, verify_lens_snapshot=verify_lens_snapshot, + verify_caps_snapshot=verify_caps_snapshot, ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py new file mode 100644 index 000000000000..5e8615628cac --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``cap-accept``: run the schedule's policy without the layout that pays for it. + +The mode exists to split one question into two. ``compact`` changes both what +the scheduler commits *and* how the batch is laid out, so when its output moves, +nothing says which half did it. ``cap-accept`` applies the same per-request +windows to commitment while submitting the full block, giving two independent +assertions instead of one entangled one: + + cap-accept output != static output -> the scheduling policy is wrong + compact output != cap-accept -> the layout compaction is wrong + +It also produces a number ``compact`` structurally cannot: ``cap_trim_tokens``, +the positions the target accepted and the window then discarded. Under +``compact`` those positions are never scored, so the acceptance cost of trimming +can only be bounded (``trim_regret_rate``), never measured. + +The price is that it saves no compute at all, which is why it is a diagnostic +mode and never a serving one -- and why the counters must not report it as +though it had saved any. +""" + +import pytest + +from tensorrt_llm._torch.speculative.dspark_observability import ( + DSparkRaggedStats, RaggedVerifyMode) +from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSamplerBase + +MAX_DRAFT_LEN = 5 + + +def _stats(mode=RaggedVerifyMode.CAP_ACCEPT): + return DSparkRaggedStats(mode=mode, max_draft_len=MAX_DRAFT_LEN) + + +# --- the cap itself --------------------------------------------------------- +# +# `num_new_tokens` is the uncapped chain: accepted drafts + the bonus at the +# first unaccepted position. A window of `cap` drafted positions admits +# `cap + 1` tokens. + + +def test_no_cap_leaves_the_commit_untouched(): + """Every non-cap-accept path goes through here; it must be a no-op.""" + assert SpecSamplerBase._apply_verify_cap(4, None) == (4, 0) + + +@pytest.mark.parametrize("num_new_tokens,cap", [(1, 5), (3, 5), (6, 5), + (1, 0), (2, 1), (4, 3)]) +def test_a_commit_inside_its_window_is_not_trimmed(num_new_tokens, cap): + """The draft died before the cut, so the cut cost nothing.""" + committed, trim = SpecSamplerBase._apply_verify_cap(num_new_tokens, cap) + assert (committed, trim) == (num_new_tokens, 0) + + +def test_a_commit_beyond_its_window_is_truncated_to_the_window(): + # Window of 2 drafted positions admits 3 tokens; the target accepted 5. + committed, trim = SpecSamplerBase._apply_verify_cap(6, 2) + assert committed == 3 + assert trim == 3 + + +def test_a_zero_window_still_commits_the_bonus(): + """min_verify_len can be 0 drafted positions; the step must still advance. + + A request that commits nothing makes no progress and would spin forever. + """ + committed, trim = SpecSamplerBase._apply_verify_cap(6, 0) + assert committed == 1 + assert trim == 5 + + +@pytest.mark.parametrize("num_new_tokens", range(1, MAX_DRAFT_LEN + 2)) +@pytest.mark.parametrize("cap", range(0, MAX_DRAFT_LEN + 1)) +def test_nothing_is_lost_or_invented(num_new_tokens, cap): + """committed + discarded accounts for every token the target produced.""" + committed, trim = SpecSamplerBase._apply_verify_cap(num_new_tokens, cap) + assert committed + trim == num_new_tokens + assert committed >= 1, "a step that commits nothing cannot make progress" + assert trim >= 0 + + +# --- what the counters must say --------------------------------------------- + + +def test_cap_trim_is_recorded_as_the_exact_acceptance_loss(): + stats = _stats() + # Given 5, allowed 2, so 3 accepted positions were thrown away. + stats.record_acceptance(accepted=2, window=2, cap_trim=3) + # Given 1, allowed 4: died on its own, the window cost nothing. + stats.record_acceptance(accepted=1, window=4, cap_trim=0) + + assert stats.cap_trim_tokens == 3 + assert stats.accept_len == pytest.approx(1.5) + assert stats.accept_loss_per_request == pytest.approx(1.5) + # accept_len + accept_loss is what a no-trim run would have accepted. + assert (stats.accept_len + + stats.accept_loss_per_request) == pytest.approx(3.0) + + +def test_the_loss_is_absent_rather_than_zero_when_not_measured(): + """`compact` cannot produce this number, and must not appear to.""" + stats = _stats(mode=RaggedVerifyMode.COMPACT) + stats.record_acceptance(accepted=3, window=3) + assert stats.cap_trim_tokens == 0 + # The bound that IS available under compact still works. + assert stats.trimmed_hit_ceiling == 1 + + +def test_cap_accept_must_not_be_credited_with_a_compute_saving(): + """The regression this guards: cap-accept booking the windows as delivered. + + It submits the full block by construction. Deriving delivered tokens from + the windows would report a trim ratio for a run that trimmed no compute at + all -- and that ratio is the headline number for whether the feature works. + """ + stats = _stats() + num_reqs = 4 + full_block = num_reqs * (1 + MAX_DRAFT_LEN) + stats.record_step(num_gen_requests=num_reqs, + verify_lens=[5, 3, 2, 1], + bucket=None, + delivered=full_block) + + assert stats.delivered_tokens == full_block + assert stats.ceiling_tokens == full_block + assert stats.trim_ratio == 0.0, "cap-accept saves no compute" + # The schedule is still visible -- that is the whole point of the mode. + assert stats.window_tokens == sum(1 + v for v in (5, 3, 2, 1)) + assert stats.steps_ragged == 1 + + +def test_compact_still_derives_delivered_from_the_bucket(): + """The default path is unchanged by the new argument.""" + stats = _stats(mode=RaggedVerifyMode.COMPACT) + stats.record_step(num_gen_requests=4, + verify_lens=[5, 3, 2, 1], + bucket=16) + assert stats.delivered_tokens == 16 + assert stats.trim_ratio > 0.0 + + +# --- mode plumbing ---------------------------------------------------------- + + +def test_cap_accept_computes_windows_but_does_not_trim_the_token_axis(): + """The two predicates that route the whole feature.""" + assert RaggedVerifyMode.CAP_ACCEPT.computes_windows + assert not RaggedVerifyMode.CAP_ACCEPT.trims_submitted_tokens + + assert RaggedVerifyMode.COMPACT.computes_windows + assert RaggedVerifyMode.COMPACT.trims_submitted_tokens + + assert not RaggedVerifyMode.STATIC.computes_windows + assert not RaggedVerifyMode.STATIC.trims_submitted_tokens + + +def test_the_active_gate_does_not_demand_a_saving_cap_accept_never_makes(): + """The gate must stay satisfiable in the mode it exists to validate. + + `require_trim` asks "did delivered tokens come in under the no-trim + ceiling". cap-accept answers no by construction, so keying the gate on the + ratio alone would fail every cap-accept run -- the reference run against + which the others are judged. + """ + stats = _stats() + stats.record_step(num_gen_requests=4, + verify_lens=[5, 3, 2, 1], + bucket=None, + delivered=4 * (1 + MAX_DRAFT_LEN)) + assert stats.trim_ratio == 0.0 + stats.assert_ragged_active(require_trim=True) + + +def test_the_active_gate_still_demands_a_saving_from_compact(): + """The regression the clause above must not introduce.""" + stats = _stats(mode=RaggedVerifyMode.COMPACT) + stats.record_step(num_gen_requests=4, + verify_lens=[5, 3, 2, 1], + bucket=4 * (1 + MAX_DRAFT_LEN)) + assert stats.trim_ratio == 0.0 + with pytest.raises(AssertionError, match="nothing was saved"): + stats.assert_ragged_active(require_trim=True) + + +def test_the_summary_carries_the_measurement(): + stats = _stats() + stats.record_acceptance(accepted=2, window=2, cap_trim=3) + summary = stats.summary() + assert summary["mode"] == "cap-accept" + assert summary["cap_trim_tokens"] == 3 + assert summary["accept_loss_per_request"] == pytest.approx(3.0) From 9acabe5fd7c75f00706fccb6bc4fff4553f19288 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Sun, 2 Aug 2026 23:26:11 -0700 Subject: [PATCH 065/177] [None][fix] Cap cap-accept inside acceptance, not on the host afterwards The cap-accept mode added in the previous commit applied its window in update_requests, on the host, after the sampler had synced. That is too late and left the drafter permanently out of step. num_accepted_tokens is consumed again on device, later in the SAME forward. The DSpark drafter (dspark.py:606-652) reads it to pick the hidden state the next block is conditioned on, to back-fill the accepted tokens into its rolling KV window, and to advance its persistent decode position (_ctx_len += nacc). Capping after all of that advances the drafter by the uncapped count while only the capped prefix is committed. Output stays lossless, because verification is authoritative, but the draft state drifts from the sequence for the rest of the run -- so the mode stops being the trusted reference the other modes are meant to be compared against, which is its only purpose. SGLang caps at the same point this now does: inside acceptance, before anything downstream reads the count (cap_correct_len, dspark_accept.py:138). The clamp itself already existed at both exits of sample_and_accept_draft_tokens (interface.py), gated on is_ragged_verify. Reusing that flag would have been wrong: it means two things at once, "stop acceptance here" AND "the draft buffers are packed by these windows" -- _padded_gen_draft_tokens slices draft_tokens[:total_verify_tokens - num_gens]. Under cap-accept only the first holds, since the layout is deliberately untrimmed and the buffers are the usual rectangle. Setting verify_lens to get the clamp would also take the packed branch and unpack a rectangle as if it were ragged, attributing one request's drafts to another with no error anywhere. So: a separate SpecMetadata.accept_caps, applied by interface.apply_accept_caps at both exits. Units match verify_lens (num_accepted_tokens is accepted + 1, so a window of w is a cap of w + 1). Both buffers are persistent and written in place, since acceptance runs inside the target's captured graph. cap_trim now accumulates on device into accept_cap_trim (int64) and is read back only when the periodic summary is logged, every 32 steps. Syncing on it per step would serialize the very forward this mode exists to leave untouched. DSparkRaggedStats.record_cap_trim_total assigns rather than adds, since the device counter is itself the running total. Removes the host-side _apply_verify_cap. update_requests still reads the snapshot, but only to report which window the request was given. 59 tests (was 51). All 306 DSpark unit tests pass. Still not exercised e2e. Signed-off-by: Lance Liao --- docs/dspark_confidence_schedule_goal.md | 2 +- docs/dspark_p0_task_prompt.md | 20 +- docs/dspark_ragged_verify_status.md | 29 ++- .../_torch/pyexecutor/model_engine.py | 50 +++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 7 + .../speculative/dspark_observability.py | 24 +-- tensorrt_llm/_torch/speculative/interface.py | 60 +++++- .../_torch/speculative/spec_sampler_base.py | 41 +--- .../hw_agnostic/test_dspark_cap_accept.py | 178 +++++++++++++----- 9 files changed, 304 insertions(+), 107 deletions(-) diff --git a/docs/dspark_confidence_schedule_goal.md b/docs/dspark_confidence_schedule_goal.md index fdb11446f73f..735a5263afab 100644 --- a/docs/dspark_confidence_schedule_goal.md +++ b/docs/dspark_confidence_schedule_goal.md @@ -318,7 +318,7 @@ cd /code/tensorrt_llm && python -m pytest \ | 测试 | 内容 | 为什么锐利 | |---|---|---| | **A1 逐 token 等价** | 同一 engine 配置、同一批 prompt，`static`（verify 全量）与开启调度后，temperature=0 下断言每个 prompt 的 `token_ids` **完全相等** | 贪婪投机解码定义上是 lossless，任何分歧都是 bug。把 6 题的容忍度变成 1 token 的容忍度。**注意**：跨不同 graph 形状的严格 bitwise 相等 TRT-LLM 并不保证（MoE/attention 规约顺序会变），所以诚实的形式是允许极少数 prompt 分歧并设硬上限，或用 `cap-accept` 模式（下条）。 | -| **A2 `cap-accept` 差分** | 实现 SGLang 的三模式（`static` / `cap-accept` / `compact`）。`cap-accept` 跑**均匀 kernel 路径**（`next_n=6`）但只提交每请求窗口内的 token，输出应与 `compact` **逐 token 相同** | `cap-accept` vs `compact` 的差异 ⇒ **必然是 ragged kernel bug**；`static` vs `cap-accept` 的差异 ⇒ **必然是调度/接受逻辑 bug**。把 pass/fail 变成可定位。**已实现**：窗口写在独立属性 `py_verify_cap` 上（不碰 `py_verify_len`），布局侧因此看不见窗口、自动走整块；截断在 `SpecSamplerBase._apply_verify_cap`。 | +| **A2 `cap-accept` 差分** | 实现 SGLang 的三模式（`static` / `cap-accept` / `compact`）。`cap-accept` 跑**均匀 kernel 路径**（`next_n=6`）但只提交每请求窗口内的 token，输出应与 `compact` **逐 token 相同** | `cap-accept` vs `compact` 的差异 ⇒ **必然是 ragged kernel bug**；`static` vs `cap-accept` 的差异 ⇒ **必然是调度/接受逻辑 bug**。把 pass/fail 变成可定位。**已实现**：窗口写在独立属性 `py_verify_cap` 上（不碰 `py_verify_len`），布局侧因此看不见窗口、自动走整块；截断在设备端 acceptance 内部（`interface.apply_accept_caps`，经独立字段 `SpecMetadata.accept_caps` 触发，不借用 `is_ragged_verify` 以免连带走进 draft 缓冲的打包分支）。 | | **A3 布局一致性断言** | host 侧断言：`sum(1 + py_verify_len over padded generation_requests) == attn_metadata.num_tokens - num_ctx_tokens == spec_metadata.total_verify_tokens == key[bucket]` | 这一条就能同时抓住 H0、H5、H8 | | **A4 时序断言** | 断言 `attn_metadata.ragged_verify_lens` 在 `prepare()` 执行的那一刻是**本步**的（而非上一步） | 抓 H2、H4 | | **A5 接受长度分布** | 记录并断言每请求接受长度直方图，以及 `sum(accepted)` 与接受 kernel 报告一致 | 抓「token 落到了错误请求的槽位但整批仍生成合法文本」这一类 | diff --git a/docs/dspark_p0_task_prompt.md b/docs/dspark_p0_task_prompt.md index c16f0c3cce8f..08ee1d5f2989 100644 --- a/docs/dspark_p0_task_prompt.md +++ b/docs/dspark_p0_task_prompt.md @@ -891,10 +891,19 @@ verify-all 的退化情形出现。 正确**(`rewind = 整块 − capped`)。这样就不必在「为 compact 调了三个 bug 才对」的 记账代码里再加一组条件分支——§10.6 那个 KV leak 正是这类记账错误。 -截断在 `SpecSamplerBase._apply_verify_cap`:`num_new_tokens` 是未截断链 -(`accepted + 1`),窗口 `w` 允许 `w + 1` 个 token。截断后末位是 `draft[w]`,它已通过 -验证,提交合法。与 SGLang 取「cut 处 target 自己的采样」来源不同,**贪婪下相同**, -两者提交 token 数一致且都无损。 +**截断必须在设备端 acceptance 内部**(`interface.apply_accept_caps`)。第一版写在主机侧 +`update_requests`,是错的——这条值得单独记: + +`num_accepted_tokens` 在**同一次 forward 的后续、设备上还要被消费**。`dspark.py:606-652` +里 drafter 用它挑下一块的条件隐状态、回填滚动 KV 窗口、并推进持久解码位置 +(`_ctx_len += nacc`)。主机同步后再截断,设备已按未截断的数推进过 drafter 状态,而主机 +只提交 capped 前缀——输出因验证权威仍无损,但 drafter 状态永久漂移,这个模式也就不再 +是它存在意义所在的可信参照。SGLang 截在同一位置(`dspark_accept.py:138`)。 + +**且不能借 `is_ragged_verify` 触发已有的 clamp**:那个标志同时意味着「draft 缓冲已按 +窗口打包」(`_padded_gen_draft_tokens` 切 `draft_tokens[:total−num_gens]`)。cap-accept +的缓冲是普通矩形,借用会把矩形当 ragged 解包,一个请求的 draft 记到另一个头上,**全程 +不报错**。所以新增独立字段 `SpecMetadata.accept_caps`。 **它产出 `compact` 结构上拿不到的数**:`cap_trim_tokens` = 被丢弃的、本可接受的位置 数,即裁剪的真实接受代价。`compact` 下那些位置根本没算,SGLang 为此专门养了一个 @@ -912,4 +921,5 @@ verify-all 的退化情形出现。 代价是**不省任何算力**,所以它是诊断模式,永远不是服务配置。 覆盖:`tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py` -(51 项,含截断算术的全枚举不变式、两个计数器坑的回归)。**尚未 e2e 跑过。** +(59 项,含 clamp 的全枚举不变式、context 请求不受影响、累加器是运行总量而非增量、 +两个计数器坑的回归)。dspark 全套 **306 passed**。**尚未 e2e 跑过。** diff --git a/docs/dspark_ragged_verify_status.md b/docs/dspark_ragged_verify_status.md index f307ccea5795..a077e3ebe518 100644 --- a/docs/dspark_ragged_verify_status.md +++ b/docs/dspark_ragged_verify_status.md @@ -167,19 +167,36 @@ P0-3/5/6 属于「跑出来的数字看着正常但其实是错的」：replay token 轴**——整块照算,超出窗口的接受位置被丢弃。它绝不能 alias 成 `compact`:唯一 价值就是那个对照。 -实现方式是把窗口写在**独立属性 `py_verify_cap`** 上,不碰 `py_verify_len`。于是布局侧 +窗口写在**独立属性 `py_verify_cap`** 上,不碰 `py_verify_len`。于是布局侧 (`model_engine` / `cuda_graph_runner` / `dsa.py`)完全看不见窗口,自动走整块; `_verified_len` 因为 `py_verify_len is None` 自动返回整块,**KV rewind 随之正确** -(`rewind = 整块 - capped`)。这样新模式是纯增量,而不是在「为 compact 调了三个 bug +(`rewind = 整块 − capped`)。新模式因此是纯增量,不必在「为 compact 调了三个 bug 才对」的记账代码里再加一组条件分支。 -截断在 `SpecSamplerBase._apply_verify_cap`:`num_new_tokens` 是未截断链 -(`accepted + 1`),窗口 `w` 允许 `w + 1` 个 token。截断后末位是 `draft[w]`——它**已通过 -验证**,提交合法。与 SGLang 取「cut 处 target 自己的采样」来源不同,贪婪下相同,两者 -提交 token 数一致且都无损。 +**截断必须在设备端、acceptance 内部**——`interface.apply_accept_caps`,由 +`sample_and_accept_draft_tokens` 的两个出口调用。这一点第一版写错了(截断放在主机侧 +`update_requests`),原因值得记下来: + +`num_accepted_tokens` **在同一次 forward 的后续、设备上还要被消费**。`dspark.py:606-652` +里 drafter 用它挑下一块的条件隐状态、把中间接受 token 回填进滚动 KV 窗口,并推进自己的 +持久解码位置(`_ctx_len += nacc`)。等主机同步后再截断,设备已经按**未截断**的数推进过 +drafter 状态,而主机只提交了 capped 前缀——输出因为验证是权威的仍然无损,但 drafter +状态从此永久漂移,这个模式也就不再是它存在意义所在的那个可信参照。SGLang 截在同一 +位置(`cap_correct_len`,`dspark_accept.py:138`),同样是在任何下游读取之前。 + +**不能直接复用 `is_ragged_verify` 来触发那个 clamp**:它同时意味着「在此停止接受」和 +「draft 缓冲已按窗口打包」(`_padded_gen_draft_tokens` 会切 +`draft_tokens[:total_verify_tokens − num_gens]`)。cap-accept 只要前者,缓冲是普通矩形。 +借用它会连带走进打包分支,把矩形当 ragged 解包——一个请求的 draft 记到另一个头上, +**全程不报错**。所以是独立字段 `SpecMetadata.accept_caps`。 + +单位与 `verify_lens` 一致:`num_accepted_tokens` 是「接受 draft 数 + 1」,所以窗口 `w` +对应的 cap 是 `w + 1`。 它产出 `compact` 结构上拿不到的数:`cap_trim_tokens`,即**被丢弃的、本可接受的位置数** ——裁剪的真实接受代价。`compact` 下那些位置根本没算,只能用 `trim_regret_rate` 卡下界。 +该计数在设备端累加(`accept_cap_trim`,int64,持久),**只在打摘要时(每 32 步)读回**; +在 acceptance 里同步会把这个模式本该原样保留的 forward 串行化。 代价是**不省任何算力**,所以它是诊断模式,永远不是服务配置。`assert_ragged_active()` 的 `require_trim` 因此对该模式豁免,否则那条门在它唯一该跑的模式里恒不可满足。 diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index eb89bf62d54e..cc2ac2b420f1 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -726,6 +726,18 @@ def __init__( self.ragged_qo_indptr_cuda = torch.empty((self.batch_size + 2, ), dtype=torch.int, device='cuda') + # cap-accept: per-request acceptance ceilings for an UNTRIMMED batch, + # plus a device-side accumulator for the positions they discarded. + # Same stable-address requirement as the two above -- both are read and + # written inside the captured graph. The accumulator is int64 because + # it runs for the length of a benchmark, and it is only ever read back + # when the summary is logged, so the hot path never syncs on it. + self.accept_caps_cuda = torch.empty((self.batch_size + 1, ), + dtype=torch.int, + device='cuda') + self.accept_cap_trim_cuda = torch.zeros((1, ), + dtype=torch.int64, + device='cuda') self.input_ids_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') @@ -4016,6 +4028,43 @@ def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, f"presentation is {expected_rows} rows -- that tensor's shape " f"IS what the ops read as their sequence count.") + def _attach_accept_caps(self, spec_metadata, generation_requests) -> None: + """Publish ``cap-accept`` ceilings, without touching the layout. + + ``py_verify_cap`` is the cap-accept twin of ``py_verify_len``: the + scheduler's window for a batch whose token axis was deliberately NOT + trimmed. Keeping them on separate attributes is what makes the mode a + pure addition -- every layout consumer keys off ``py_verify_len``, so + the input packing, the DSA row expansion, the graph key and the KV + rewind all keep seeing an ordinary uniform batch with no changes. + + The cap is in ``num_accepted_tokens`` units (accepted drafts plus the + bonus), matching ``verify_lens``, hence ``1 + py_verify_cap``. + """ + caps = [ + getattr(request, "py_verify_cap", None) + for request in generation_requests + ] + if not caps or any(cap is None for cap in caps): + # Partial caps are worse than none: the uncapped requests would + # commit their full block while the rest honoured a window, which + # is neither mode and reproducible by neither. + spec_metadata.accept_caps = None + spec_metadata.accept_cap_trim = None + return + + n = len(caps) + # Written through the persistent buffer for the same reason as + # verify_lens: acceptance runs inside the target's captured graph, so a + # freshly allocated tensor here would be invisible to every replay. + caps_view = self.accept_caps_cuda[:n] + caps_view.copy_(torch.tensor([1 + int(cap) for cap in caps], + dtype=torch.int32, + pin_memory=prefer_pinned()), + non_blocking=True) + spec_metadata.accept_caps = caps_view + spec_metadata.accept_cap_trim = self.accept_cap_trim_cuda + def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, generation_requests) -> None: """Publish this step's per-request verify windows to both metadatas. @@ -4027,6 +4076,7 @@ def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, right place, and the DSA indexer has to expand kv_lens/block tables one row per query token instead of a fixed count per request. """ + self._attach_accept_caps(spec_metadata, generation_requests) token_lens = self._ragged_token_lens(generation_requests) if token_lens is None: spec_metadata.verify_lens = None diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 843beda6b397..215600e5db99 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3503,6 +3503,13 @@ def _dspark_confidence_draft_len(self, # and cannot reach it. Emit a periodic summary so the decision is # recoverable from the log, which is the only shared channel. if stats.steps_total and stats.steps_total % 32 == 0: + if cap_accept: + # The only device read on this path, and deliberately not + # per step: the accumulator is written inside acceptance, + # so syncing on it there would serialize the very forward + # the mode is supposed to leave untouched. + stats.record_cap_trim_total( + int(self.model_engine.accept_cap_trim_cuda.item())) stats.log_summary(prefix="DSpark ragged verify [periodic]") self._maybe_assert_dspark_ragged_active(stats) fitted = None diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index 180f2c0d857b..16f3d478ecec 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -331,31 +331,31 @@ def distinct_verify_lens(self) -> int: """How many different window sizes were ever handed out.""" return len(self.verify_len_hist) - def record_acceptance(self, - *, - accepted: int, - window: int, - cap_trim: int = 0) -> None: + def record_acceptance(self, *, accepted: int, window: int) -> None: """Record one request's acceptance against the window it was given. Args: accepted: drafted positions accepted (excludes the bonus token). window: drafted positions the request was allowed to verify. - cap_trim: positions the target accepted but the window discarded. - Non-zero only in ``cap-accept``, the one mode where those - positions are actually scored, so this is a measurement rather - than an estimate. Zero elsewhere -- under ``compact`` they were - never computed and cannot be counted, only bounded from below - by ``trimmed_hit_ceiling``. """ self.requests_scored += 1 self.accepted_tokens += int(accepted) - self.cap_trim_tokens += int(cap_trim) if int(window) < self.max_draft_len: self.requests_trimmed += 1 if int(accepted) >= int(window): self.trimmed_hit_ceiling += 1 + def record_cap_trim_total(self, total: int) -> None: + """Absorb the device-side cap-accept accumulator. + + Assigned, not added: the counter it mirrors runs for the lifetime of + the engine, so this is a running total rather than a delta. Reading it + costs a device-to-host copy, which is why the caller does it when the + summary is logged rather than every step -- the acceptance path itself + must never sync. + """ + self.cap_trim_tokens = int(total) + @property def accept_len(self) -> float: """Mean accepted drafted positions per request. 0.0 with no data.""" diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 989cabdbfb7f..18db9ccea692 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -538,9 +538,32 @@ class SpecMetadata: # ``RaggedVerifyLayout.fill_bucket``. total_verify_tokens: Optional[int] = None + # ``cap-accept``: per-request ceilings on the accepted count, in the same + # token units as ``verify_lens``, for a batch whose layout was NOT trimmed. + # + # Deliberately a separate field rather than reusing ``verify_lens``. That + # one carries two meanings at once -- "stop acceptance here" AND "the draft + # buffers are packed by these windows" (see ``_padded_gen_draft_tokens``, + # which slices ``draft_tokens[:total_verify_tokens - num_gens]``). Under + # cap-accept only the first holds: the target scored the full block and the + # buffers are the usual rectangle. Setting ``verify_lens`` to get the clamp + # would silently take the packed-buffer branch too, and unpack a rectangle + # as if it were ragged -- one request's drafts attributed to another, with + # no error anywhere. + accept_caps: Optional[torch.Tensor] = None + # [1] int64 accumulator for positions the caps discarded, summed on device + # so the hot path never syncs. This is the exact acceptance the schedule + # cost, and the number ``compact`` cannot produce: there those positions + # are never scored. + accept_cap_trim: Optional[torch.Tensor] = None + @property def is_ragged_verify(self) -> bool: - """Whether this iteration uses per-request verify lengths.""" + """Whether this iteration uses per-request verify lengths. + + Implies the packed layout. ``cap-accept`` is deliberately NOT ragged by + this test -- it caps acceptance without packing anything. + """ return self.verify_lens is not None # Auto-detected per step from populated sampling params: @@ -1031,6 +1054,39 @@ def populate_sampling_params_for_one_model( (tk for tk in request_top_ks if 0 < tk < _disable_topk), default=0) +def apply_accept_caps(num_accepted_tokens: torch.Tensor, num_contexts: int, + spec_metadata) -> None: + """``cap-accept``: hold each request to its scheduled window, in place. + + This has to run HERE, inside acceptance, and not on the host after the + sampler syncs. ``num_accepted_tokens`` is consumed on device later in the + same forward: the DSpark drafter reads it to pick the hidden state the next + block is conditioned on, to back-fill the accepted tokens into its rolling + KV window, and to advance its persistent decode position + (``_ctx_len += nacc``). Capping after that leaves the drafter advanced by + the uncapped count while only the capped prefix was committed -- the output + stays lossless, because verification is authoritative, but the draft state + drifts permanently and the mode stops being the faithful reference it + exists to be. SGLang caps at the same point, before anything downstream + reads the count. + + In-place on the caller's tensor and on a persistent accumulator, so this is + safe to record into a CUDA graph and never synchronizes. + """ + caps = getattr(spec_metadata, "accept_caps", None) + if caps is None: + return + gen = num_accepted_tokens[num_contexts:] + if gen.numel() == 0: + return + caps = caps[:gen.shape[0]].to(gen.dtype) + capped = torch.minimum(gen, caps) + trim = getattr(spec_metadata, "accept_cap_trim", None) + if trim is not None: + trim += (gen - capped).sum() + num_accepted_tokens[num_contexts:] = capped + + def _padded_gen_draft_tokens(spec_metadata, num_gens: int, runtime_draft_len: int) -> torch.Tensor: """Unpack ragged-packed draft tokens into ``[num_gens, runtime_draft_len]``. @@ -1478,6 +1534,7 @@ def _sample_and_accept_draft_tokens_base( num_accepted_tokens[num_contexts:] = torch.minimum( num_accepted_tokens[num_contexts:], spec_metadata.verify_lens.to(num_accepted_tokens.dtype)) + apply_accept_caps(num_accepted_tokens, num_contexts, spec_metadata) return accepted_tokens, num_accepted_tokens @@ -1919,6 +1976,7 @@ def _sample_and_accept_draft_tokens_rejection( num_accepted_tokens[num_contexts:] = torch.minimum( num_accepted_tokens[num_contexts:], spec_metadata.verify_lens.to(num_accepted_tokens.dtype)) + apply_accept_caps(num_accepted_tokens, num_contexts, spec_metadata) return accepted_tokens, num_accepted_tokens def _update_advance_draft_sampling_seed(self, device): diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index c801eba57dea..423374766eee 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -219,35 +219,6 @@ def _verified_len(request: LlmRequest, runtime_draft_len: int, per_request = getattr(request, "py_verify_len", None) return runtime_draft_len if per_request is None else int(per_request) - @staticmethod - def _apply_verify_cap(num_new_tokens: int, - cap: Optional[int]) -> tuple[int, int]: - """Trim a commit to the scheduler's window. ``cap-accept`` only. - - The target scored the whole drafted block, so ``num_new_tokens`` is the - UNCAPPED chain: ``accepted + 1``, the accepted drafts plus the bonus at - the first unaccepted position. A window of ``cap`` drafted positions - therefore admits ``cap + 1`` tokens. - - Truncating there keeps every committed token one the target already - verified. When the cap bites, the last token kept is ``draft[cap]``, - which was accepted, rather than the target's own sample at that - position -- a different provenance from SGLang's ``cap_correct_len``, - identical under greedy, and lossless either way, since a token is only - committed if verification accepted it. The committed COUNT matches. - - Returns: - (tokens to commit, positions the cap discarded). The second value - is the exact acceptance the schedule cost, which is the number - ``compact`` cannot produce because it never scores those positions. - """ - if cap is None: - return num_new_tokens, 0 - admitted = int(cap) + 1 - if num_new_tokens <= admitted: - return num_new_tokens, 0 - return admitted, num_new_tokens - admitted - def update_requests( self, state: SampleStateSpec, @@ -275,11 +246,12 @@ def update_requests( if req.state == LlmRequestState.GENERATION_COMPLETE: continue num_new_tokens = new_tokens_lens_list[req.py_seq_slot] - # cap-accept only: the target scored the whole block, so this is the - # uncapped chain and the scheduler's window has to be applied here. + # cap-accept: `num_new_tokens` has ALREADY been capped, on device, + # inside acceptance (`interface.apply_accept_caps`) -- it has to be, + # because the drafter reads the same count later in that forward to + # advance its own state. The cap is read here only to report the + # window the request was actually given. cap = None if caps is None else caps.get(req.py_request_id) - num_new_tokens, cap_trim = self._apply_verify_cap( - num_new_tokens, cap) for i in range(num_new_tokens): new_token = add_token(req, new_tokens, beam_idx=beam_idx, step=i) if TorchSampler._handle_stop_criteria( @@ -305,8 +277,7 @@ def update_requests( if self.acceptance_stats is not None: self.acceptance_stats.record_acceptance( accepted=req.py_num_accepted_draft_tokens, - window=verified_len if cap is None else int(cap), - cap_trim=cap_trim) + window=verified_len if cap is None else int(cap)) self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) def sample_async( diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py index 5e8615628cac..77fa1b17f97c 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py @@ -22,10 +22,11 @@ """ import pytest +import torch from tensorrt_llm._torch.speculative.dspark_observability import ( DSparkRaggedStats, RaggedVerifyMode) -from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSamplerBase +from tensorrt_llm._torch.speculative.interface import apply_accept_caps MAX_DRAFT_LEN = 5 @@ -34,69 +35,130 @@ def _stats(mode=RaggedVerifyMode.CAP_ACCEPT): return DSparkRaggedStats(mode=mode, max_draft_len=MAX_DRAFT_LEN) -# --- the cap itself --------------------------------------------------------- -# -# `num_new_tokens` is the uncapped chain: accepted drafts + the bonus at the -# first unaccepted position. A window of `cap` drafted positions admits -# `cap + 1` tokens. - +class _Meta: + """Just the two fields `apply_accept_caps` reads.""" -def test_no_cap_leaves_the_commit_untouched(): - """Every non-cap-accept path goes through here; it must be a no-op.""" - assert SpecSamplerBase._apply_verify_cap(4, None) == (4, 0) + def __init__(self, caps=None, track_trim=True): + self.accept_caps = (None if caps is None else torch.tensor( + caps, dtype=torch.int32)) + self.accept_cap_trim = (torch.zeros(1, dtype=torch.int64) + if track_trim else None) -@pytest.mark.parametrize("num_new_tokens,cap", [(1, 5), (3, 5), (6, 5), - (1, 0), (2, 1), (4, 3)]) -def test_a_commit_inside_its_window_is_not_trimmed(num_new_tokens, cap): - """The draft died before the cut, so the cut cost nothing.""" - committed, trim = SpecSamplerBase._apply_verify_cap(num_new_tokens, cap) - assert (committed, trim) == (num_new_tokens, 0) +# --- the cap itself --------------------------------------------------------- +# +# WHERE this runs is the whole point, and it is not testable from here: it has +# to be inside acceptance, because the DSpark drafter reads the same +# `num_accepted_tokens` later in the SAME forward to advance its rolling KV +# window and its persistent decode position (`_ctx_len += nacc`, dspark.py). +# Capping on the host afterwards leaves the drafter advanced by the uncapped +# count while only the capped prefix was committed -- lossless output, but the +# draft state drifts and the mode stops being a faithful reference. +# +# Units: `num_accepted_tokens` is `accepted drafts + 1`, so a window of w +# drafted positions is a cap of `w + 1`, matching `verify_lens`. -def test_a_commit_beyond_its_window_is_truncated_to_the_window(): - # Window of 2 drafted positions admits 3 tokens; the target accepted 5. - committed, trim = SpecSamplerBase._apply_verify_cap(6, 2) - assert committed == 3 - assert trim == 3 +def test_no_caps_is_a_no_op(): + """Every other mode goes through this call.""" + counts = torch.tensor([3, 1, 4], dtype=torch.int32) + apply_accept_caps(counts, 0, _Meta(caps=None)) + assert counts.tolist() == [3, 1, 4] -def test_a_zero_window_still_commits_the_bonus(): - """min_verify_len can be 0 drafted positions; the step must still advance. +def test_counts_inside_their_windows_are_untouched(): + """The draft died before the cut, so the cut cost nothing.""" + counts = torch.tensor([1, 3, 2], dtype=torch.int32) + meta = _Meta(caps=[6, 4, 6]) + apply_accept_caps(counts, 0, meta) + assert counts.tolist() == [1, 3, 2] + assert int(meta.accept_cap_trim.item()) == 0 + + +def test_counts_beyond_their_windows_are_clamped_and_the_loss_accumulated(): + counts = torch.tensor([6, 5, 2], dtype=torch.int32) + meta = _Meta(caps=[3, 5, 6]) + apply_accept_caps(counts, 0, meta) + assert counts.tolist() == [3, 5, 2] + # Only the first request lost anything: 6 -> 3. + assert int(meta.accept_cap_trim.item()) == 3 + + +def test_the_accumulator_is_a_running_total_across_steps(): + """It mirrors a device counter that lives as long as the engine.""" + meta = _Meta(caps=[2, 2]) + for _ in range(3): + counts = torch.tensor([5, 5], dtype=torch.int32) + apply_accept_caps(counts, 0, meta) + assert counts.tolist() == [2, 2] + assert int(meta.accept_cap_trim.item()) == 3 * (3 + 3) + + +def test_context_requests_are_left_alone(): + """Caps are per generation request; contexts sit ahead of them.""" + counts = torch.tensor([9, 9, 6, 6], dtype=torch.int32) + meta = _Meta(caps=[2, 3]) + apply_accept_caps(counts, 2, meta) + assert counts.tolist() == [9, 9, 2, 3], "a context count was clamped" + + +def test_a_generation_free_batch_does_not_crash(): + counts = torch.tensor([4, 4], dtype=torch.int32) + meta = _Meta(caps=[2, 2]) + apply_accept_caps(counts, 2, meta) + assert counts.tolist() == [4, 4] + + +def test_missing_accumulator_still_caps(): + """The clamp is correctness; the counter is only measurement.""" + counts = torch.tensor([6], dtype=torch.int32) + meta = _Meta(caps=[2], track_trim=False) + apply_accept_caps(counts, 0, meta) + assert counts.tolist() == [2] + + +@pytest.mark.parametrize("accepted", range(0, MAX_DRAFT_LEN + 2)) +@pytest.mark.parametrize("window", range(0, MAX_DRAFT_LEN + 1)) +def test_nothing_is_lost_or_invented(accepted, window): + """committed + discarded accounts for every position the target produced.""" + counts = torch.tensor([accepted + 1], dtype=torch.int32) + meta = _Meta(caps=[window + 1]) + apply_accept_caps(counts, 0, meta) + committed = int(counts[0]) + trim = int(meta.accept_cap_trim.item()) + assert committed + trim == accepted + 1 + assert committed >= 1, "a step that commits nothing cannot make progress" + assert committed <= window + 1 - A request that commits nothing makes no progress and would spin forever. - """ - committed, trim = SpecSamplerBase._apply_verify_cap(6, 0) - assert committed == 1 - assert trim == 5 - -@pytest.mark.parametrize("num_new_tokens", range(1, MAX_DRAFT_LEN + 2)) -@pytest.mark.parametrize("cap", range(0, MAX_DRAFT_LEN + 1)) -def test_nothing_is_lost_or_invented(num_new_tokens, cap): - """committed + discarded accounts for every token the target produced.""" - committed, trim = SpecSamplerBase._apply_verify_cap(num_new_tokens, cap) - assert committed + trim == num_new_tokens - assert committed >= 1, "a step that commits nothing cannot make progress" - assert trim >= 0 +def test_the_cap_never_raises_a_count(): + """A window wider than the block must not invent acceptance.""" + counts = torch.tensor([2], dtype=torch.int32) + meta = _Meta(caps=[99]) + apply_accept_caps(counts, 0, meta) + assert counts.tolist() == [2] # --- what the counters must say --------------------------------------------- -def test_cap_trim_is_recorded_as_the_exact_acceptance_loss(): +def test_cap_trim_is_absorbed_from_the_device_counter_as_a_running_total(): + """Assigned, not accumulated: the device counter is itself the total. + + Adding would double-count on every periodic read, which is the shape of + bug that makes a metric drift slowly enough to be believed. + """ stats = _stats() - # Given 5, allowed 2, so 3 accepted positions were thrown away. - stats.record_acceptance(accepted=2, window=2, cap_trim=3) - # Given 1, allowed 4: died on its own, the window cost nothing. - stats.record_acceptance(accepted=1, window=4, cap_trim=0) + stats.record_acceptance(accepted=2, window=2) + stats.record_acceptance(accepted=1, window=4) + stats.record_cap_trim_total(3) assert stats.cap_trim_tokens == 3 + stats.record_cap_trim_total(7) + assert stats.cap_trim_tokens == 7, "reads must not accumulate" + assert stats.accept_len == pytest.approx(1.5) - assert stats.accept_loss_per_request == pytest.approx(1.5) - # accept_len + accept_loss is what a no-trim run would have accepted. - assert (stats.accept_len + - stats.accept_loss_per_request) == pytest.approx(3.0) + assert stats.accept_loss_per_request == pytest.approx(3.5) def test_the_loss_is_absent_rather_than_zero_when_not_measured(): @@ -186,8 +248,30 @@ def test_the_active_gate_still_demands_a_saving_from_compact(): def test_the_summary_carries_the_measurement(): stats = _stats() - stats.record_acceptance(accepted=2, window=2, cap_trim=3) + stats.record_acceptance(accepted=2, window=2) + stats.record_cap_trim_total(3) summary = stats.summary() assert summary["mode"] == "cap-accept" assert summary["cap_trim_tokens"] == 3 assert summary["accept_loss_per_request"] == pytest.approx(3.0) + + +def test_accept_caps_is_not_the_ragged_flag(): + """The distinction the whole fix rests on. + + `is_ragged_verify` means two things at once -- cap acceptance AND the draft + buffers are packed by window (`_padded_gen_draft_tokens` slices + `draft_tokens[:total_verify_tokens - num_gens]`). cap-accept wants only the + first: its buffers are the ordinary rectangle. If `accept_caps` ever starts + implying ragged, a rectangle gets unpacked as if it were packed and one + request's drafts are attributed to another, silently. + """ + from tensorrt_llm._torch.speculative.interface import SpecMetadata + + fields = SpecMetadata.__dataclass_fields__ + assert "accept_caps" in fields + assert "accept_cap_trim" in fields + + meta = _Meta(caps=[2, 3]) + # The mock stands in for a metadata object carrying caps but no windows. + assert getattr(meta, "verify_lens", None) is None From 8f82517ac983bbc8b8aef3818e6e23f7d43e81c9 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Mon, 3 Aug 2026 00:02:54 -0700 Subject: [PATCH 066/177] [None][fix] Give the CUDA-graph padding dummy a verify cap _attach_accept_caps drops the whole batch's caps when any generation request lacks one, since a partly capped batch is neither mode. The padding dummy that pad_batch appends to reach a captured batch size carries no cap of its own, so cap-accept silently degraded to static -- on padded steps only, which is worse than always: the mode measured nothing on some steps, measured correctly on others, and nothing in the output said which. The ragged path already handles this for py_verify_len; this is the missing half. The dummy's own acceptance is discarded, so it gets the full block, i.e. no cap. Cleared when the batch is not cap-accept, because the dummy object is cached across steps and would otherwise carry a stale cap into a static batch. Also makes the drop loud rather than silent: a batch with some caps and some missing now warns once, naming the counts. Every silent-degradation bug on this branch has cost more to find than the bug itself. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/cuda_graph_runner.py | 13 ++++++++++ .../_torch/pyexecutor/model_engine.py | 15 ++++++++++- .../hw_agnostic/test_dspark_cap_accept.py | 26 +++++++++++++++++++ 3 files changed, 53 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index b451a292146b..12738334f794 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -767,6 +767,19 @@ def _get_padded_batch(self, batch: ScheduledRequests, for request in batch.generation_requests) padding_dummy_request.py_verify_len = (self.ragged_pad_verify_len if ragged else None) + # cap-accept needs the same treatment for the same reason, and + # fails differently without it: `_attach_accept_caps` drops the + # whole batch's caps if any request lacks one, so a single + # capless dummy would degrade the mode to `static` -- silently, + # and only on the steps that happen to need padding. The dummy's + # own acceptance is discarded, so give it the full block, i.e. no + # cap at all. Cleared otherwise: the object is cached across steps + # and would carry a stale cap into a static batch. + cap_accept = any( + getattr(request, "py_verify_cap", None) is not None + for request in batch.generation_requests) + padding_dummy_request.py_verify_cap = ( + self.spec_config.max_draft_len if cap_accept else None) batch.generation_requests.extend([padding_dummy_request] * padding_size) return padding_size diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index cc2ac2b420f1..4b9ef0acae72 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4048,7 +4048,20 @@ def _attach_accept_caps(self, spec_metadata, generation_requests) -> None: if not caps or any(cap is None for cap in caps): # Partial caps are worse than none: the uncapped requests would # commit their full block while the rest honoured a window, which - # is neither mode and reproducible by neither. + # is neither mode and reproducible by neither. So drop them all -- + # but say so, because the result is a run that looks like a + # successful cap-accept and measured nothing. The CUDA-graph + # padding dummy used to land here on every padded step (it carries + # no cap of its own); it is given one in `pad_batch` now. + if caps and any(cap is not None for cap in caps): + if not getattr(self, "_warned_partial_accept_caps", False): + self._warned_partial_accept_caps = True + logger.warning( + "DSpark cap-accept: %d of %d generation requests carry " + "no verify cap, so caps were dropped for the whole " + "batch and this step ran as 'static'. The mode is " + "measuring nothing on steps like this one.", + sum(1 for cap in caps if cap is None), len(caps)) spec_metadata.accept_caps = None spec_metadata.accept_cap_trim = None return diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py index 77fa1b17f97c..b7405c9f137a 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py @@ -256,6 +256,32 @@ def test_the_summary_carries_the_measurement(): assert summary["accept_loss_per_request"] == pytest.approx(3.0) +def test_the_cuda_graph_padding_dummy_carries_a_cap(): + """The regression that made the whole mode a silent no-op. + + `pad_batch` appends a shared dummy request to reach a captured batch size, + and `_attach_accept_caps` drops the batch's caps if ANY request lacks one. + The dummy has no cap of its own, so before this it disabled cap-accept -- + on padded steps only, so the mode measured nothing on some steps and + worked on others, with nothing in the output to say which. + + Mirrors the treatment `py_verify_len` already got for the ragged path. + """ + import inspect + + from tensorrt_llm._torch.pyexecutor import cuda_graph_runner + + src = inspect.getsource(cuda_graph_runner.CUDAGraphRunner._get_padded_batch) + assert "py_verify_cap" in src, ( + "the padded-batch builder no longer stamps the dummy with a verify " + "cap; cap-accept will silently degrade to static on every padded step") + # And it must be cleared, not left stale: the dummy object is cached + # across steps and shared between them. + assert "if cap_accept else None" in src, ( + "the dummy's cap is not cleared for non-cap-accept batches, so a " + "stale cap would trim a static step") + + def test_accept_caps_is_not_the_ragged_flag(): """The distinction the whole fix rests on. From faf2c60935b715d13c27cde6a2b8c73adf21c00b Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Mon, 3 Aug 2026 00:14:52 -0700 Subject: [PATCH 067/177] [None][feat] Record cap-accept's acceptance loss per request The loss was a single device accumulator: one number for the whole run. That cannot distinguish the two cases it most needs to: 500 requests, 1000 tokens lost, 2 per request -> broad, mild trim 500 requests, 1000 tokens lost, 20 requests lost 50 -> the scheduler is paying for its throughput out of 4% of the traffic Identical cap_trim_tokens, identical accept_loss_per_request, opposite conclusions about the planner. The second is the failure this mode exists to catch, and a total renders it invisible. So accept_cap_trim becomes a [batch] buffer written by apply_accept_caps and carried to the host on the sampler's EXISTING copy -- the DSpark worker returns it as cap_trim_lens, the sampler stores it slot-indexed beside new_tokens_lens. No new synchronization: the periodic .item() read the accumulator needed is gone, and the tensor is kilobytes next to new_tokens on the same event. Stats gain requests_cap_trimmed, cap_trim_max, cap_trim_hist and cap_trim_concentration (share of scored requests that lost anything). Read with accept_loss_per_request: the same mean at 0.9 concentration is a broad trim, at 0.02 it is a subpopulation absorbing the entire cost. The hazard this introduces is staleness. The buffer is persistent and slot-indexed, so a step producing no caps -- static, or any other speculation mode -- must write ZEROS into this batch's slots rather than skip them, or a slot recycled to a new request reports its previous occupant's loss. Both the device side (context rows zeroed in apply_accept_caps) and the host side (zeros_like in the sampler) do this explicitly, and both are pinned by tests. This is the third instance of that bug class on this branch. 63 tests (was 59). 335 DSpark unit tests pass; test_spec_gate_e2e fails on a missing EAGLE3 checkpoint, unrelated and pre-existing. Still not exercised e2e. Signed-off-by: Lance Liao --- docs/dspark_p0_task_prompt.md | 14 ++- docs/dspark_ragged_verify_status.md | 26 ++++- .../_torch/pyexecutor/model_engine.py | 13 +-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 7 -- tensorrt_llm/_torch/speculative/dspark.py | 9 +- .../speculative/dspark_observability.py | 54 +++++++-- tensorrt_llm/_torch/speculative/interface.py | 25 ++-- .../_torch/speculative/spec_sampler_base.py | 29 ++++- .../hw_agnostic/test_dspark_cap_accept.py | 108 ++++++++++++++---- 9 files changed, 220 insertions(+), 65 deletions(-) diff --git a/docs/dspark_p0_task_prompt.md b/docs/dspark_p0_task_prompt.md index 08ee1d5f2989..d301fe5e3f51 100644 --- a/docs/dspark_p0_task_prompt.md +++ b/docs/dspark_p0_task_prompt.md @@ -920,6 +920,16 @@ verify-all 的退化情形出现。 代价是**不省任何算力**,所以它是诊断模式,永远不是服务配置。 +损失**逐请求**记录,不是只记总量:总量分不出「每个请求丢一点」和「少数请求被砍烂」, +两者 `cap_trim_tokens` 相同而结论相反。因此另记 `requests_cap_trimmed` / +`cap_trim_max` / `cap_trim_hist` / `cap_trim_concentration`。链路是 +`apply_accept_caps` 写持久 `[bs]` 缓冲 → worker 返回 `cap_trim_lens` → 采样器 store → +随**已有的** host 拷贝回来,**不额外引入同步**。 + +这里的真风险是**陈旧缓冲**:缓冲持久且按 slot 索引,不产出 caps 的步必须往本批次 +slot **写零**,否则 slot 回收给新请求后会报上一个占用者的损失——今晚已经踩过两次的 +同一类。`_process_outputs` 显式 `zeros_like`,有测试卡住。 + 覆盖:`tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py` -(59 项,含 clamp 的全枚举不变式、context 请求不受影响、累加器是运行总量而非增量、 -两个计数器坑的回归)。dspark 全套 **306 passed**。**尚未 e2e 跑过。** +(63 项,含 clamp 的全枚举不变式、context slot 清零、同一总量的两种分布可区分、 +三个计数器坑的回归)。dspark 全套 **335 passed**。**尚未 e2e 跑过。** diff --git a/docs/dspark_ragged_verify_status.md b/docs/dspark_ragged_verify_status.md index a077e3ebe518..b0bcffe04308 100644 --- a/docs/dspark_ragged_verify_status.md +++ b/docs/dspark_ragged_verify_status.md @@ -193,10 +193,28 @@ drafter 状态,而主机只提交了 capped 前缀——输出因为验证是权 单位与 `verify_lens` 一致:`num_accepted_tokens` 是「接受 draft 数 + 1」,所以窗口 `w` 对应的 cap 是 `w + 1`。 -它产出 `compact` 结构上拿不到的数:`cap_trim_tokens`,即**被丢弃的、本可接受的位置数** -——裁剪的真实接受代价。`compact` 下那些位置根本没算,只能用 `trim_regret_rate` 卡下界。 -该计数在设备端累加(`accept_cap_trim`,int64,持久),**只在打摘要时(每 32 步)读回**; -在 acceptance 里同步会把这个模式本该原样保留的 forward 串行化。 +它产出 `compact` 结构上拿不到的数:**被丢弃的、本可接受的位置数**——裁剪的真实接受 +代价。`compact` 下那些位置根本没算,只能用 `trim_regret_rate` 卡下界。 + +**逐请求记录,不是只记总量。** 理由是总量分不出这两种情况: + +| 500 个请求共丢 1000 | 含义 | +|---|---| +| 每个丢 2 | planner 均匀小幅收割,大概率是合理取舍 | +| 20 个各丢 50、480 个丢 0 | planner **系统性掐死一小撮请求** | + +两者 `cap_trim_tokens` 和 `accept_loss_per_request` 完全相同,结论却相反。所以另外记 +`requests_cap_trimmed` / `cap_trim_max` / `cap_trim_hist` / `cap_trim_concentration` +(有损失的请求占比):平均损失相同的情况下,concentration 0.9 是普遍轻微裁剪,0.02 是 +调度器拿 2% 的流量在补贴吞吐——只有后者是问题。 + +链路:`apply_accept_caps` 写进持久 `[bs]` 缓冲 → DSpark worker 的返回字典 +`cap_trim_lens` → 采样器 store(slot 索引)→ 随**已有的** host 拷贝回来。因此这个测量 +**不额外引入任何同步**,也不需要设备端累加器。 + +**陈旧缓冲是这里唯一的真风险**:缓冲持久且按 slot 索引,不产出 caps 的步(static、以及 +其它任何投机模式)必须往本批次的 slot **写零**而不是跳过,否则 slot 被回收给新请求后 +会报上一个占用者的损失。`_process_outputs` 里显式 `zeros_like`,并有测试卡住。 代价是**不省任何算力**,所以它是诊断模式,永远不是服务配置。`assert_ragged_active()` 的 `require_trim` 因此对该模式豁免,否则那条门在它唯一该跑的模式里恒不可满足。 diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 4b9ef0acae72..3cc57369a62e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -727,16 +727,15 @@ def __init__( dtype=torch.int, device='cuda') # cap-accept: per-request acceptance ceilings for an UNTRIMMED batch, - # plus a device-side accumulator for the positions they discarded. - # Same stable-address requirement as the two above -- both are read and - # written inside the captured graph. The accumulator is int64 because - # it runs for the length of a benchmark, and it is only ever read back - # when the summary is logged, so the hot path never syncs on it. + # and the per-request positions those ceilings discarded. Same + # stable-address requirement as the two above -- both are read and + # written inside the captured graph. The losses ride back to the host + # on the sampler's existing copy, so nothing here ever syncs. self.accept_caps_cuda = torch.empty((self.batch_size + 1, ), dtype=torch.int, device='cuda') - self.accept_cap_trim_cuda = torch.zeros((1, ), - dtype=torch.int64, + self.accept_cap_trim_cuda = torch.zeros((self.batch_size + 1, ), + dtype=torch.int, device='cuda') self.input_ids_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 215600e5db99..843beda6b397 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3503,13 +3503,6 @@ def _dspark_confidence_draft_len(self, # and cannot reach it. Emit a periodic summary so the decision is # recoverable from the log, which is the only shared channel. if stats.steps_total and stats.steps_total % 32 == 0: - if cap_accept: - # The only device read on this path, and deliberately not - # per step: the accumulator is written inside acceptance, - # so syncing on it there would serialize the very forward - # the mode is supposed to leave untouched. - stats.record_cap_trim_total( - int(self.model_engine.accept_cap_trim_cuda.item())) stats.log_summary(prefix="DSpark ragged verify [periodic]") self._maybe_assert_dspark_ragged_active(stats) fitted = None diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 6ea4bb7b45c0..c152c005cfbf 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -838,10 +838,17 @@ def _forward_impl( if saved_confidence is not None: self._confidence_logits.copy_(saved_confidence) - 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, } + # cap-accept only: what each request's window threw away, written by + # `apply_accept_caps` above. Rides the sampler's existing copy rather + # than a read of its own, so the measurement costs no synchronization. + cap_trim = getattr(spec_metadata, "accept_cap_trim", None) + if cap_trim is not None and spec_metadata.accept_caps is not None: + outputs["cap_trim_lens"] = cap_trim[:batch_size] + return outputs diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index 16f3d478ecec..ee4331a68fcf 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -199,6 +199,15 @@ def __init__(self, *, mode: RaggedVerifyMode, max_draft_len: int): #: it stays 0 under ``compact``, where the honest answer is that the #: number is unknowable without re-running. self.cap_trim_tokens = 0 + #: The distribution, not just the total. "Every request gave up a + #: little" and "a few were cut to pieces" sum to the same + #: cap_trim_tokens and mean opposite things about the planner: the + #: first is a tuning knob, the second is the scheduler concentrating + #: its whole cost on a subpopulation it selected. A total cannot tell + #: them apart, which is the reason this is kept per request at all. + self.requests_cap_trimmed = 0 + self.cap_trim_max = 0 + self.cap_trim_hist: Counter = Counter() #: Graph replay is the whole point of landing on a bucket; an eager #: step costs far more than the tokens it saved. @@ -331,12 +340,20 @@ def distinct_verify_lens(self) -> int: """How many different window sizes were ever handed out.""" return len(self.verify_len_hist) - def record_acceptance(self, *, accepted: int, window: int) -> None: + def record_acceptance(self, + *, + accepted: int, + window: int, + cap_trim: int = 0) -> None: """Record one request's acceptance against the window it was given. Args: accepted: drafted positions accepted (excludes the bonus token). window: drafted positions the request was allowed to verify. + cap_trim: positions this request accepted that its window then + discarded. Non-zero only under ``cap-accept``, the one mode + that scores them; elsewhere 0 means "not measured", not "no + loss". """ self.requests_scored += 1 self.accepted_tokens += int(accepted) @@ -344,17 +361,13 @@ def record_acceptance(self, *, accepted: int, window: int) -> None: self.requests_trimmed += 1 if int(accepted) >= int(window): self.trimmed_hit_ceiling += 1 - - def record_cap_trim_total(self, total: int) -> None: - """Absorb the device-side cap-accept accumulator. - - Assigned, not added: the counter it mirrors runs for the lifetime of - the engine, so this is a running total rather than a delta. Reading it - costs a device-to-host copy, which is why the caller does it when the - summary is logged rather than every step -- the acceptance path itself - must never sync. - """ - self.cap_trim_tokens = int(total) + cap_trim = int(cap_trim) + if cap_trim > 0: + self.cap_trim_tokens += cap_trim + self.requests_cap_trimmed += 1 + self.cap_trim_hist[cap_trim] += 1 + if cap_trim > self.cap_trim_max: + self.cap_trim_max = cap_trim @property def accept_len(self) -> float: @@ -377,6 +390,19 @@ def accept_loss_per_request(self) -> float: return 0.0 return self.cap_trim_tokens / self.requests_scored + @property + def cap_trim_concentration(self) -> float: + """Share of scored requests that lost anything at all. ``cap-accept``. + + Read it together with :attr:`accept_loss_per_request`. The same mean + loss at 0.9 concentration is a broad, mild trim; at 0.02 it is the + scheduler paying for its throughput out of two percent of the traffic. + Only the second is a problem, and the mean alone cannot see it. + """ + if not self.requests_scored: + return 0.0 + return self.requests_cap_trimmed / self.requests_scored + @property def trim_regret_rate(self) -> float: """Share of trimmed requests whose draft was still alive at the cut. @@ -416,6 +442,10 @@ def summary(self) -> Dict[str, object]: "trim_regret_rate": round(self.trim_regret_rate, 4), "cap_trim_tokens": self.cap_trim_tokens, "accept_loss_per_request": round(self.accept_loss_per_request, 4), + "requests_cap_trimmed": self.requests_cap_trimmed, + "cap_trim_max": self.cap_trim_max, + "cap_trim_hist": dict(sorted(self.cap_trim_hist.items())), + "cap_trim_concentration": round(self.cap_trim_concentration, 4), "graph_replays": self.graph_replays, "graph_eager": self.graph_eager, "fallbacks": dict(self.fallbacks), diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 18db9ccea692..1a7c1e550b96 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -551,10 +551,14 @@ class SpecMetadata: # as if it were ragged -- one request's drafts attributed to another, with # no error anywhere. accept_caps: Optional[torch.Tensor] = None - # [1] int64 accumulator for positions the caps discarded, summed on device - # so the hot path never syncs. This is the exact acceptance the schedule - # cost, and the number ``compact`` cannot produce: there those positions - # are never scored. + # [batch] positions the cap discarded, PER REQUEST -- the exact acceptance + # the schedule cost each one, and the number ``compact`` cannot produce + # because it never scores those positions. + # + # Per request rather than a running total, because the total cannot tell + # "every request gave up a little" from "a few requests were cut to + # pieces" -- and those call for opposite conclusions about the planner, + # while summing to exactly the same number. accept_cap_trim: Optional[torch.Tensor] = None @property @@ -1070,8 +1074,10 @@ def apply_accept_caps(num_accepted_tokens: torch.Tensor, num_contexts: int, exists to be. SGLang caps at the same point, before anything downstream reads the count. - In-place on the caller's tensor and on a persistent accumulator, so this is - safe to record into a CUDA graph and never synchronizes. + In-place on the caller's tensor and on persistent buffers, so this is safe + to record into a CUDA graph and never synchronizes. The per-request losses + ride back to the host on the sampler's existing copy, so measuring them + costs no sync of its own. """ caps = getattr(spec_metadata, "accept_caps", None) if caps is None: @@ -1083,7 +1089,12 @@ def apply_accept_caps(num_accepted_tokens: torch.Tensor, num_contexts: int, capped = torch.minimum(gen, caps) trim = getattr(spec_metadata, "accept_cap_trim", None) if trim is not None: - trim += (gen - capped).sum() + # Zero the context slots too. The buffer is persistent and the batch + # composition moves every step, so a slot left unwritten would carry + # the previous step's loss and be attributed to whoever holds it now. + rows = num_contexts + gen.shape[0] + trim[:rows].zero_() + trim[num_contexts:rows] = (gen - capped).to(trim.dtype) num_accepted_tokens[num_contexts:] = capped diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 423374766eee..ddb6d94ce9bc 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -47,6 +47,8 @@ class SampleStateTensorsSpec(SampleStateTensors): new_tokens_lens: torch.Tensor next_draft_tokens: torch.Tensor + #: cap-accept only; None on every other path. + cap_trim_lens: Optional[torch.Tensor] = None @dataclass(kw_only=True) @@ -98,6 +100,11 @@ class Store: next_new_tokens: torch.Tensor next_draft_tokens: torch.Tensor new_tokens_lens: torch.Tensor + #: cap-accept: per-request positions the verify window discarded. + #: Slot-indexed like new_tokens_lens, and rewritten for every slot in + #: the batch each step -- a slot left alone would hand the previous + #: occupant's loss to whoever holds it now. + cap_trim_lens: torch.Tensor def __init__(self, args: TorchSampler.Args, *, draft_len: int): """ @@ -124,6 +131,7 @@ def __init__(self, args: TorchSampler.Args, *, draft_len: int): next_new_tokens=int_tensor((max_tokens, seq_slots, self.max_beam_width)), next_draft_tokens=int_tensor((seq_slots, draft_tokens_size)), new_tokens_lens=int_tensor((seq_slots,)), + cap_trim_lens=int_tensor((seq_slots,)), ) def _get_max_tokens(self, args: TorchSampler.Args, draft_len: int) -> int: @@ -237,6 +245,8 @@ def update_requests( state.sampler_event.synchronize() new_tokens = state.host.new_tokens.tolist() new_tokens_lens_list = state.host.new_tokens_lens.tolist() + cap_trim_list = (state.host.cap_trim_lens.tolist() + if state.host.cap_trim_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) @@ -277,7 +287,9 @@ def update_requests( if self.acceptance_stats is not None: self.acceptance_stats.record_acceptance( accepted=req.py_num_accepted_draft_tokens, - window=verified_len if cap is None else int(cap)) + window=verified_len if cap is None else int(cap), + cap_trim=(0 if cap_trim_list is None else + cap_trim_list[req.py_seq_slot])) self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) def sample_async( @@ -348,6 +360,18 @@ def sample_async( self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens) self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens) self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens) + # cap-accept: absent on every other path, and written as ZEROS then + # rather than skipped. The buffer is persistent and slot-indexed, so a + # slot left untouched would keep whatever the previous occupant lost + # and hand it to the current one -- the stale-buffer bug this branch + # has already paid for twice. + o_cap_trim = outputs.get("cap_trim_lens") + if o_cap_trim is None: + o_cap_trim = torch.zeros_like(o_new_tokens_lens) + else: + o_cap_trim = o_cap_trim[num_skip:num_skip + num_sampling_requests] + self.store.cap_trim_lens.index_copy_( + 0, slots, o_cap_trim.to(self.store.cap_trim_lens.dtype)) # Create sample state with async D2H copy device_tensors = SampleStateTensorsSpec( @@ -360,6 +384,9 @@ 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), + # [max_num_sequences] int32 -- kilobytes next to new_tokens above, + # and it rides the same event, so no extra synchronization. + cap_trim_lens=self._copy_to_host(self.store.cap_trim_lens), ) sampler_event = self._record_sampler_event() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py index b7405c9f137a..f049396e7900 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py @@ -38,10 +38,11 @@ def _stats(mode=RaggedVerifyMode.CAP_ACCEPT): class _Meta: """Just the two fields `apply_accept_caps` reads.""" - def __init__(self, caps=None, track_trim=True): + def __init__(self, caps=None, track_trim=True, rows=8): self.accept_caps = (None if caps is None else torch.tensor( caps, dtype=torch.int32)) - self.accept_cap_trim = (torch.zeros(1, dtype=torch.int64) + # Persistent per-request buffer, sized like the engine's. + self.accept_cap_trim = (torch.zeros(rows, dtype=torch.int32) if track_trim else None) @@ -72,7 +73,7 @@ def test_counts_inside_their_windows_are_untouched(): meta = _Meta(caps=[6, 4, 6]) apply_accept_caps(counts, 0, meta) assert counts.tolist() == [1, 3, 2] - assert int(meta.accept_cap_trim.item()) == 0 + assert meta.accept_cap_trim[:3].tolist() == [0, 0, 0] def test_counts_beyond_their_windows_are_clamped_and_the_loss_accumulated(): @@ -80,18 +81,22 @@ def test_counts_beyond_their_windows_are_clamped_and_the_loss_accumulated(): meta = _Meta(caps=[3, 5, 6]) apply_accept_caps(counts, 0, meta) assert counts.tolist() == [3, 5, 2] - # Only the first request lost anything: 6 -> 3. - assert int(meta.accept_cap_trim.item()) == 3 + # Only the first request lost anything: 6 -> 3. And we can say WHICH. + assert meta.accept_cap_trim[:3].tolist() == [3, 0, 0] -def test_the_accumulator_is_a_running_total_across_steps(): - """It mirrors a device counter that lives as long as the engine.""" +def test_each_step_overwrites_the_previous_step_s_losses(): + """Per step, not cumulative: the buffer is what THIS step lost. + + The host accumulates across steps. A buffer that added would double-count + every request still in the batch next step. + """ meta = _Meta(caps=[2, 2]) for _ in range(3): counts = torch.tensor([5, 5], dtype=torch.int32) apply_accept_caps(counts, 0, meta) assert counts.tolist() == [2, 2] - assert int(meta.accept_cap_trim.item()) == 3 * (3 + 3) + assert meta.accept_cap_trim[:2].tolist() == [3, 3] def test_context_requests_are_left_alone(): @@ -100,6 +105,8 @@ def test_context_requests_are_left_alone(): meta = _Meta(caps=[2, 3]) apply_accept_caps(counts, 2, meta) assert counts.tolist() == [9, 9, 2, 3], "a context count was clamped" + # Context slots are zeroed, not left stale. + assert meta.accept_cap_trim[:4].tolist() == [0, 0, 4, 3] def test_a_generation_free_batch_does_not_crash(): @@ -125,7 +132,7 @@ def test_nothing_is_lost_or_invented(accepted, window): meta = _Meta(caps=[window + 1]) apply_accept_caps(counts, 0, meta) committed = int(counts[0]) - trim = int(meta.accept_cap_trim.item()) + trim = int(meta.accept_cap_trim[0]) assert committed + trim == accepted + 1 assert committed >= 1, "a step that commits nothing cannot make progress" assert committed <= window + 1 @@ -142,23 +149,56 @@ def test_the_cap_never_raises_a_count(): # --- what the counters must say --------------------------------------------- -def test_cap_trim_is_absorbed_from_the_device_counter_as_a_running_total(): - """Assigned, not accumulated: the device counter is itself the total. - - Adding would double-count on every periodic read, which is the shape of - bug that makes a metric drift slowly enough to be believed. - """ +def test_cap_trim_is_recorded_per_request(): stats = _stats() - stats.record_acceptance(accepted=2, window=2) - stats.record_acceptance(accepted=1, window=4) + # Given 5, allowed 2 -> 3 accepted positions thrown away. + stats.record_acceptance(accepted=2, window=2, cap_trim=3) + # Died on its own inside a wide window: the schedule cost it nothing. + stats.record_acceptance(accepted=1, window=4, cap_trim=0) - stats.record_cap_trim_total(3) assert stats.cap_trim_tokens == 3 - stats.record_cap_trim_total(7) - assert stats.cap_trim_tokens == 7, "reads must not accumulate" - + assert stats.requests_cap_trimmed == 1 assert stats.accept_len == pytest.approx(1.5) - assert stats.accept_loss_per_request == pytest.approx(3.5) + assert stats.accept_loss_per_request == pytest.approx(1.5) + # accept_len + accept_loss is what a no-trim run would have accepted. + assert (stats.accept_len + + stats.accept_loss_per_request) == pytest.approx(3.0) + + +def test_the_same_total_loss_is_distinguished_by_its_shape(): + """The reason this is per request at all. + + Both runs lose 20 tokens over 10 requests. The first spreads it; the + second takes it all out of one request. A total cannot tell them apart, + and they mean opposite things about the planner. + """ + spread = _stats() + for _ in range(10): + spread.record_acceptance(accepted=1, window=1, cap_trim=2) + + concentrated = _stats() + concentrated.record_acceptance(accepted=1, window=1, cap_trim=20) + for _ in range(9): + concentrated.record_acceptance(accepted=3, window=5, cap_trim=0) + + assert spread.cap_trim_tokens == concentrated.cap_trim_tokens == 20 + assert spread.accept_loss_per_request == pytest.approx( + concentrated.accept_loss_per_request) + + # ...and the shape separates them. + assert spread.cap_trim_concentration == pytest.approx(1.0) + assert concentrated.cap_trim_concentration == pytest.approx(0.1) + assert spread.cap_trim_max == 2 + assert concentrated.cap_trim_max == 20 + assert spread.cap_trim_hist == {2: 10} + assert concentrated.cap_trim_hist == {20: 1} + + +def test_concentration_is_zero_when_nothing_was_measured(): + stats = _stats(mode=RaggedVerifyMode.COMPACT) + stats.record_acceptance(accepted=3, window=5) + assert stats.cap_trim_concentration == 0.0 + assert stats.cap_trim_max == 0 def test_the_loss_is_absent_rather_than_zero_when_not_measured(): @@ -248,12 +288,32 @@ def test_the_active_gate_still_demands_a_saving_from_compact(): def test_the_summary_carries_the_measurement(): stats = _stats() - stats.record_acceptance(accepted=2, window=2) - stats.record_cap_trim_total(3) + stats.record_acceptance(accepted=2, window=2, cap_trim=3) summary = stats.summary() assert summary["mode"] == "cap-accept" assert summary["cap_trim_tokens"] == 3 assert summary["accept_loss_per_request"] == pytest.approx(3.0) + for key in ("requests_cap_trimmed", "cap_trim_max", "cap_trim_hist", + "cap_trim_concentration"): + assert key in summary, f"{key} missing from the summary" + + +def test_a_slot_left_unwritten_cannot_inherit_the_previous_occupant_s_loss(): + """The stale-buffer hazard, guarded at the one place it can arise. + + cap_trim_lens is persistent and slot-indexed. A step that produces no + caps -- static, or any other speculation mode -- must still write ZEROS + into this batch's slots, or a slot recycled to a new request would report + whatever its previous occupant lost. + """ + import inspect + + from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSamplerBase + + src = inspect.getsource(SpecSamplerBase.sample_async) + assert "zeros_like" in src and "cap_trim_lens" in src, ( + "_process_outputs no longer writes zeros when the step produced no " + "cap_trim_lens; recycled slots will report a stale request's loss") def test_the_cuda_graph_padding_dummy_carries_a_cap(): From 8d3fc5535a3d7d9d9a9cfcdb6fe3c3d581380894 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Mon, 3 Aug 2026 01:03:46 -0700 Subject: [PATCH 068/177] [None][fix] Build the CUDA-graph capture set from the mode, not the config flag The first end-to-end cap-accept run died during CUDA-graph capture, inside MoE, on "unmatched tensor shape" -- before a single request had run, with nothing in the failure pointing at speculation. cap-accept is configured with enable_ragged_verify=True: its config is byte-identical to compact and only the environment differs. Two capture-side gates keyed off that flag and so treated it as compact. The capture set was built as (bs, top_tier, bs*(t+1)) -- ragged token buckets -- for a mode whose real steps always submit the full block, so the capture pass ran synthetic batches at token totals the mode never produces. Anything that SHAPES the batch has to ask whether the token axis actually shrinks, not whether ragged was requested. Adds dspark_observability.resolve_ragged_verify_mode / trims_submitted_tokens, resolving config plus environment exactly as the DSpark worker does and separately from it, since the capture set is built before any batch exists. The capture set (model_engine) and the ragged graph key (cuda_graph_runner) now use it. This is the same split SGLang draws: its ragged_verify_compact_graphs_enabled() is COMPACT-only. This also corrects a claim made when cap-accept landed -- that it was purely additive because every layout consumer keys off py_verify_len. True of the layout; not true of capture, which keys off the config. That path was not audited, and the end-to-end run found it on the first attempt. 64 tests. Unit tests could not have caught this -- it needs a real capture pass -- so the new one pins the predicate rather than the behaviour. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/cuda_graph_runner.py | 5 ++- .../_torch/pyexecutor/model_engine.py | 9 +++- .../speculative/dspark_observability.py | 25 +++++++++++ .../hw_agnostic/test_dspark_cap_accept.py | 44 +++++++++++++++++++ 4 files changed, 81 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 12738334f794..ddc1bbe1da3a 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -296,7 +296,10 @@ def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: build for every model, and a per-request scan there is pure overhead for the ~all of them that will never be ragged. """ - if not getattr(self.spec_config, "enable_ragged_verify", False): + # Keyed on the MODE: `cap-accept` shares this config flag but never + # shrinks the token axis, so it must keep the ordinary uniform key. + from ..speculative.dspark_observability import trims_submitted_tokens + if not trims_submitted_tokens(self.spec_config): return None # The bucket the fit agreed on, not a fresh sum over the batch. Summing # here walks generation_requests *after* _get_padded_batch has appended diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 3cc57369a62e..4dac45bfe706 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1850,7 +1850,14 @@ def _get_graphs_to_capture( int(t) for t in self.spec_config.verify_len_tiers if int(t) >= 0 }) - if getattr(self.spec_config, "enable_ragged_verify", False): + # The MODE, not the config flag. `cap-accept` is configured with + # enable_ragged_verify=True as well, but submits the FULL block, so + # capturing the ragged token ladder for it builds synthetic batches + # at token totals its real steps never produce -- which fails + # during capture, inside MoE, long before any of it is explicable. + from ..speculative.dspark_observability import \ + trims_submitted_tokens + if trims_submitted_tokens(self.spec_config): # Ragged always drafts the full block -- only verification is # trimmed -- so draft_len is pinned to the top tier and the # ladder moves to the *token* axis: one bucket per tier, i.e. diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index ee4331a68fcf..a1fe4f0fe826 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -122,6 +122,31 @@ def read_ragged_verify_mode( f"{', '.join(repr(m.value) for m in RaggedVerifyMode)}") +def resolve_ragged_verify_mode(spec_config) -> RaggedVerifyMode: + """The mode a config plus the environment actually select. + + The config flag alone is not the answer, and using it as one is a bug with + a long fuse: ``enable_ragged_verify`` is what ``cap-accept`` is configured + with too, but cap-accept submits the FULL block. Anything that shapes the + batch -- the CUDA-graph capture set, the graph key -- has to ask whether + the token axis actually shrinks (:attr:`RaggedVerifyMode.trims_submitted_tokens`), + not whether ragged was requested. + + Resolved identically to the DSpark worker's own resolution, and separately + from it, because the capture set is built before any batch exists. + """ + configured = (RaggedVerifyMode.COMPACT if getattr( + spec_config, "enable_ragged_verify", False) else RaggedVerifyMode.STATIC) + return read_ragged_verify_mode(default=configured) + + +def trims_submitted_tokens(spec_config) -> bool: + """Whether this config+environment shrinks the token axis. See above.""" + if spec_config is None: + return False + return resolve_ragged_verify_mode(spec_config).trims_submitted_tokens + + class DSparkRaggedStats: """Per-step counters for the ragged verification path. diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py index f049396e7900..cee639cb2adf 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py @@ -21,6 +21,9 @@ though it had saved any. """ +import os +from unittest.mock import patch + import pytest import torch @@ -342,6 +345,47 @@ def test_the_cuda_graph_padding_dummy_carries_a_cap(): "stale cap would trim a static step") +def test_the_capture_set_follows_the_mode_not_the_config_flag(): + """The bug the first end-to-end run found, in CUDA-graph capture. + + cap-accept is configured with enable_ragged_verify=True -- the config is + byte-identical to compact, only the environment differs. Anything that + SHAPES the batch must therefore ask whether the token axis actually + shrinks, not whether ragged was requested. The capture set asked the flag, + so it built synthetic batches at ragged token totals (bs*(t+1)) for a mode + whose real steps always submit the full block, and the capture pass died + inside MoE on `unmatched tensor shape` -- before any request had run, with + nothing pointing at the scheduler. + + Unit tests could not have caught it: it needs a real capture pass. This + pins the predicate instead. + """ + from types import SimpleNamespace + + from tensorrt_llm._torch.speculative.dspark_observability import ( + resolve_ragged_verify_mode, trims_submitted_tokens) + + ragged_cfg = SimpleNamespace(enable_ragged_verify=True) + + with patch.dict(os.environ, + {"TLLM_DSPARK_RAGGED_VERIFY_MODE": "cap-accept"}): + assert resolve_ragged_verify_mode( + ragged_cfg) is RaggedVerifyMode.CAP_ACCEPT + assert not trims_submitted_tokens(ragged_cfg), ( + "cap-accept submits the full block; a capture set built for it " + "must be the uniform one") + + with patch.dict(os.environ, {"TLLM_DSPARK_RAGGED_VERIFY_MODE": "compact"}): + assert trims_submitted_tokens(ragged_cfg) + + # Unset: the config flag decides, which is the pre-existing behaviour. + with patch.dict(os.environ, {}, clear=True): + assert trims_submitted_tokens(ragged_cfg) + assert not trims_submitted_tokens( + SimpleNamespace(enable_ragged_verify=False)) + assert not trims_submitted_tokens(None) + + def test_accept_caps_is_not_the_ragged_flag(): """The distinction the whole fix rests on. From 1d51566c019d4dd46d47979dcabb8cc27a6cbf2a Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Mon, 3 Aug 2026 01:40:40 -0700 Subject: [PATCH 069/177] [None][fix] Capture one draft length, and make the planner's counters readable Two leftovers from removing the uniform tier ladder, both surfaced by the first cap-accept end-to-end run. 1. The capture set was still the batch-size x tier cross product. That was right while the uniform path picked a different draft_len per step; nothing varies draft_len at runtime any more -- cap-accept pins it to planner.max_tier, compact moves the ladder to the token axis, and scheduling-off never reaches this branch. So the lower tiers were captured and never selected. The run showed it plainly: capturing 57 graphs (19 batch sizes x 3 draft-length tiers [1, 3, 5]) with 114 capture passes, two thirds of them dead. Captured graphs cost ~10-23 MB each out of the same pool that becomes KV cache, so this was KV cache spent on graphs that could never be replayed. Now one graph per batch size at the top tier. 2. The planner's stats block reported nothing. The same run showed planner {'decisions': 0, 'fallback_*': 0} (all zero) fallbacks {'planner_declined': 2, ...} (on the ragged side) so the planner declined twice and incremented nothing. `decisions` was never incremented anywhere -- a dead counter since it was introduced, and the denominator the four fallback counters are read against, which made them uninterpretable. And `num_gen_requests <= 0` returned None with no counter at all; that is the path that fired here. Increments `decisions` before any decline, names the empty-batch path (fallback_no_gen_requests), and adds a gate check: windows produced on N>0 steps with 0 recorded decisions now fails, since it means the planner's stats are not wired to the observer and the only record of WHY a step declined to trim is unreadable. 336 DSpark unit tests pass. test_spec_gate_e2e fails on a missing EAGLE3 checkpoint, unrelated and pre-existing. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/model_engine.py | 42 +++++++++++++------ .../speculative/dspark_observability.py | 15 +++++++ .../_torch/speculative/dspark_verify.py | 13 ++++++ .../test_dspark_confidence_schedule.py | 18 ++++++-- 4 files changed, 71 insertions(+), 17 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 4dac45bfe706..faa8b47d06f4 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -1835,15 +1835,18 @@ def _get_graphs_to_capture( return graphs # Case 2b: DSpark confidence-scheduled verification. - # Unlike Case 2, the draft length here is NOT a function of the batch - # size -- at a fixed batch size the scheduler may still pick a different - # length as the batch's confidence changes. So the ladder is captured for - # every batch size, as a cross product. + # The drafted block is always full -- scheduling trims *verification*, + # not drafting -- so draft_len is pinned to the top tier in every mode. + # What varies is which axis the tiers land on: `compact` moves them to + # the token count (a cross product with the batch sizes), while + # `cap-accept` and the static fallback keep the ordinary one graph per + # batch size. # - # This is the expensive part of the feature and the reason the ladder is - # kept short: each captured graph costs ~10-23 MB of metadata outside the - # shared pool, and that memory is taken directly out of KV cache (the KV - # pool is sized from what is left after capture). + # Captured graphs are the expensive part of this feature: each costs + # ~10-23 MB of metadata outside the shared pool, taken directly out of + # KV cache (the KV pool is sized from what is left after capture). So a + # tier that no mode can select is not free, it is KV cache spent on a + # graph that will never be replayed. if (self.spec_config is not None and getattr( self.spec_config, "enable_confidence_scheduling", False)): tiers = sorted({ @@ -1872,13 +1875,26 @@ def _get_graphs_to_capture( f"token buckets from tiers {tiers}; draft_len pinned to " f"{top_tier}).") return graphs - graphs = [(bs, draft_len) for bs in cuda_graph_batch_sizes - for draft_len in tiers] + # One draft length, not the ladder. The cross product above was + # right while the uniform tier path existed and picked a different + # draft_len per step; that path is gone, and nothing varies + # draft_len at runtime any more: + # + # scheduling off never reaches this branch (guarded above) + # cap-accept runtime_draft_len = planner.max_tier, always + # compact takes the ragged branch, pinned to top_tier + # + # So the lower tiers were captured and never selected -- with + # tiers {1,3,5}, two thirds of these graphs were dead, and captured + # graph memory comes straight out of KV cache. + top_tier = tiers[-1] + graphs = [(bs, top_tier) for bs in cuda_graph_batch_sizes] logger.info( f"DSpark confidence scheduling: capturing {len(graphs)} graphs " - f"({len(cuda_graph_batch_sizes)} batch sizes x {len(tiers)} " - f"draft-length tiers {tiers}). Every extra tier costs captured " - f"graph memory that would otherwise be KV cache.") + f"({len(cuda_graph_batch_sizes)} batch sizes, draft_len pinned " + f"to {top_tier}). The drafted block is always full; only " + f"verification is trimmed, and cap-accept trims it at " + f"acceptance rather than in the layout.") return graphs # Case 3: Target model (two-model) or one-model without dynamic draft diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index a1fe4f0fe826..7ac64f4aa40d 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -513,6 +513,21 @@ def assert_ragged_active(self, *, require_trim: bool = True) -> None: raise AssertionError( f"DSpark ragged verify handed out only one distinct window " f"size, so the batch was uniform in all but name: {summary}") + # The two counter families have to move together. A run once reported + # per-step scheduling activity here beside an all-zero planner block, + # because `decisions` was never incremented -- which made the fallback + # counters unreadable, and those are the only thing that says WHY a + # step declined to trim. Zero decisions against non-zero windowed steps + # means the planner's stats are not wired to this object at all. + if self.planner_stats is not None: + decisions = int(self.planner_stats.get("decisions", 0)) + windowed = self.steps_ragged + self.steps_uniform_windows + if windowed > 0 and decisions == 0: + raise AssertionError( + f"DSpark scheduling produced windows on {windowed} steps " + f"but the planner recorded 0 decisions, so its fallback " + f"counters cannot be read -- and they are the only record " + f"of why a step declined to trim: {summary}") # `cap-accept` submits the full block on purpose -- it measures the # commitment policy, not a compute saving -- so a zero trim ratio is # the expected outcome there, not a silent no-op. Requiring one would diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index c273dbe0adfe..eccac0016ca7 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -93,6 +93,7 @@ def __init__( "fallback_flat_cost": 0, "fallback_no_confidence": 0, "fallback_short_snapshot": 0, + "fallback_no_gen_requests": 0, } @property @@ -204,7 +205,19 @@ def decide_verify_lens( input layout is built per request (so it goes ragged) while the spec metadata sees a ``None`` and stays uniform. """ + # Counted before any decline, so the four fallback counters below have a + # denominator. Without it they are absolute counts against an unknown + # total, which is how a run reported `planner_declined: 2` alongside an + # all-zero planner block: nothing said how many decisions there had + # been, so nothing said whether zero fallbacks meant "healthy" or + # "never asked". + self.stats["decisions"] += 1 + if num_gen_requests <= 0: + # Benign during ramp-up, but it has to be nameable: an unlabelled + # `return None` here is indistinguishable in the counters from a + # real decline, and this feature's failures are all silent. + self.stats["fallback_no_gen_requests"] += 1 return None if self.cost_table is None or self.cost_table.is_flat: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py index 842344b36a91..fd79ded34949 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py @@ -517,12 +517,22 @@ def _capture_set(batch_sizes, tiers, max_draft_len=BLOCK): return PyTorchModelEngine._get_graphs_to_capture(engine, list(batch_sizes), None) -def test_capture_is_the_batch_size_x_tier_cross_product(): - """K is NOT a function of batch size here -- every bs needs every tier.""" +def test_capture_pins_the_draft_length_to_the_top_tier(): + """One graph per batch size, not a cross product with the tier ladder. + + The cross product was right while the uniform tier path existed and chose + a different draft_len per step. That path is gone and nothing varies + draft_len at runtime now: cap-accept pins it to planner.max_tier and + compact moves the ladder to the token axis. Capturing the lower tiers + anyway spent KV cache -- captured graphs come out of the same pool -- on + graphs that could never be selected. Measured on a real run: 57 graphs + where 19 were reachable. + """ bss, tiers = [1, 8, 64], [1, 3, BLOCK] graphs = _capture_set(bss, tiers) - assert set(graphs) == {(bs, k) for bs in bss for k in tiers} - assert len(graphs) == len(bss) * len(tiers) + assert set(graphs) == {(bs, BLOCK) for bs in bss} + assert len(graphs) == len(bss), ( + f"expected one graph per batch size, got {len(graphs)}: {graphs}") def _stats(max_draft_len=5): From 7f13404b282137da1bf982d4e4a35e6ee439dd13 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Mon, 3 Aug 2026 01:47:29 -0700 Subject: [PATCH 070/177] [None][fix] Repair the tests the uniform-path removal broke Removing the uniform tier ladder made "enable_confidence_scheduling=True with enable_ragged_verify=False" a ValueError. Five tests still built exactly that and had been failing since: - four ladder tests in test_llm_args.py (TestDSparkConfidenceScheduling), which only set the master switch because the ladder derivation is orthogonal to ragged. They go through a _scheduled() helper now that adds the companions the flag implies. - the integration test test_gsm8k_dep8_confidence_scheduling_overlap, which would have raised at config construction. I did not catch these when the path was removed: the sweeps I ran covered tests/unittest/_torch/speculative/ and one llmapi file, so test_llm_args.py was never executed and integration tests are not run locally at all. The "all green" reported at the time was against an incomplete selection. The integration test is converted rather than deleted, to test_gsm8k_dep8_static_verify_overlap with enable_confidence_scheduling=False. Its old configuration always chose the full block on this checkpoint, so what it actually exercised was static verification with the overlap scheduler -- which is precisely the reference the three-way differential needs: cap-accept must match it token for token, and compact must then match cap-accept. It also no longer reads TLLM_DSPARK_SPS_TABLE, since a reference run must be reproducible without one. 351 unit tests pass. test_spec_gate_e2e fails on a missing EAGLE3 checkpoint, unrelated and pre-existing. Signed-off-by: Lance Liao --- .../defs/accuracy/test_llm_api_pytorch.py | 29 ++++++++++++------- tests/unittest/llmapi/test_llm_args.py | 26 ++++++++++++----- 2 files changed, 38 insertions(+), 17 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 0212fb61fb52..bfa7832454dd 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4067,8 +4067,20 @@ def _run_gsm8k(self, llm): @pytest.mark.skip_less_mpi_world_size(8) @parametrize_with_ids("moe_backend", ["MEGAMOE_DEEPGEMM", "TRTLLM"]) - def test_gsm8k_dep8_confidence_scheduling_overlap(self, moe_backend): - """Confidence scheduling with the overlap scheduler ON. + def test_gsm8k_dep8_static_verify_overlap(self, moe_backend): + """DSpark verifying the full block, with the overlap scheduler ON. + + This was a confidence-scheduling test until the uniform tier ladder was + removed. That path chose one verify length for the whole batch and, on + this checkpoint, always chose the full block -- so it never differed + from what this test does now, and the combination it configured + (scheduling on, ragged off) is rejected outright by llm_args today. + + Kept, rather than deleted, because it is the STATIC reference the other + two modes are differenced against: same prompts, same seed, verify + everything. cap-accept must match it token for token (the schedule only + changes how many tokens are committed per step, never which), and + compact must then match cap-accept. The sibling test above pins ``disable_overlap_scheduler=True``, which makes it a weaker test than it looks: with overlap off, @@ -4090,17 +4102,14 @@ def test_gsm8k_dep8_confidence_scheduling_overlap(self, moe_backend): """ kv_cache_config = KvCacheConfig(enable_block_reuse=False, free_gpu_memory_fraction=0.5) - # Same reasoning as the ragged case: without a cost table the planner - # cannot tell a cheap verified token from an expensive one and keeps the - # full block every step, so the tier ladder is captured but never - # selected from. The run is still a valid correctness check -- it just - # does not exercise the scheduling decision. - sps_table_path = os.environ.get("TLLM_DSPARK_SPS_TABLE") or None + # No scheduling at all: every request verifies its whole drafted block. + # There is nothing for a cost table to inform here, which is the point + # -- this run has to be reproducible without one so it can serve as the + # reference for the two scheduled modes. spec_config = DSparkDecodingConfig( max_draft_len=5, speculative_model=self.MODEL_PATH, - enable_confidence_scheduling=True, - confidence_sps_table_path=sps_table_path) + enable_confidence_scheduling=False) with LLM(self.MODEL_PATH, attn_backend="TRTLLM", tensor_parallel_size=8, diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 0b6ed9faaf36..22fd0058f478 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -3825,6 +3825,21 @@ def _cfg(self, **kw): kw.setdefault("max_draft_len", 5) return DSparkDecodingConfig(**kw) + def _scheduled(self, **kw): + """Scheduling on, with the companions it now implies. + + Turning scheduling on means ragged -- the uniform tier ladder that + used to sit between "schedule per request" and "verify everything" is + gone -- and ragged means a profiled cost table, without which the + planner's budget degenerates to verify-all. Both are enforced in + llm_args, so a ladder test that sets only the master switch is + testing a configuration that can no longer exist. + """ + kw["enable_confidence_scheduling"] = True + kw.setdefault("enable_ragged_verify", True) + kw.setdefault("confidence_sps_table_path", "/tmp/sps.json") + return self._cfg(**kw) + def test_defaults_are_off(self): c = self._cfg() assert c.enable_confidence_scheduling is False @@ -3845,21 +3860,18 @@ def test_knobs_require_the_master_switch(self): self._cfg(**kw) def test_default_ladder_is_derived_and_includes_the_full_block(self): - c = self._cfg(enable_confidence_scheduling=True) + c = self._scheduled() assert c.verify_len_tiers == [1, 3, 5] def test_full_block_is_always_reachable(self): """The full block is the fallback when confidence is stale; keep it.""" - c = self._cfg(enable_confidence_scheduling=True, - confidence_verify_len_tiers=[1, 2]) + c = self._scheduled(confidence_verify_len_tiers=[1, 2]) assert c.verify_len_tiers[-1] == 5 def test_ladder_cannot_exceed_max_draft_len(self): with pytest.raises(ValueError, match="exceeds max_draft_len"): - self._cfg(enable_confidence_scheduling=True, - confidence_verify_len_tiers=[1, 9]) + self._scheduled(confidence_verify_len_tiers=[1, 9]) def test_ladder_is_deduped_and_sorted(self): - c = self._cfg(enable_confidence_scheduling=True, - confidence_verify_len_tiers=[5, 1, 3, 3]) + c = self._scheduled(confidence_verify_len_tiers=[5, 1, 3, 3]) assert c.verify_len_tiers == [1, 3, 5] From 571115542b53c6e7991e89580eca2cf9053713d1 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Mon, 3 Aug 2026 02:45:24 -0700 Subject: [PATCH 071/177] [None][fix] Publish the ragged bucket during CUDA-graph capture Regression from d2c19dea01, which changed the graph key from re-summing the batch to reading a fitted value: - total = 0 - for request in batch.generation_requests: - total += 1 + int(request.py_verify_len) - return total + return getattr(self, "agreed_ragged_bucket", None) That is right for real steps -- every attention-DP rank must key on a value they agreed on, not one each derived independently after padding. But it moved the bucket from something computed on demand to something someone must publish, and only one of the two paths that shape a ragged batch was updated: real step fit_ragged_verify_lens shapes AND publishes capture _set_warmup_ragged_windows shapes only Capture never calls the fit, so agreed_ragged_bucket stayed None while the warmup batch really held `verify_bucket` tokens. The old code happened to work here as a side effect of summing on the spot; nothing recorded that capture depended on it. Symptom: the capture-time forward ran with metadata describing a uniform batch over ragged tensors and died hundreds of frames away, in the MoE, on "unmatched tensor shape" -- two graphs into a 57-graph pass, with nothing pointing back at the key. Had it survived, all three of a batch size's buckets would have been stored under one key with no token axis: colliding with each other, matching no runtime step, and dropping every ragged step to eager without an error anywhere. _set_warmup_ragged_windows' own docstring warns about exactly that outcome while only guarding the token total, because when it was written the key and the total were the same computation. Fix is two lines: publish the bucket per captured entry (capture walks every tier for every batch size, so it is per-entry, not once), and clear it after the pass so the first real step cannot inherit a capture-time bucket. Found by the first end-to-end compact run. This is the third capture-side defect on this branch after the same pattern -- runtime path updated, capture path not -- so the new test is behavioural rather than an inspection. 352 unit tests pass; test_spec_gate_e2e fails on a missing EAGLE3 checkpoint, unrelated and pre-existing. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/model_engine.py | 16 +++++++ .../test_dspark_graph_miss_attribution.py | 48 +++++++++++++++++++ 2 files changed, 64 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index faa8b47d06f4..73802bcd2ed3 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2109,6 +2109,12 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: finally: if force_non_greedy and spec_metadata is not None: spec_metadata._force_non_greedy_for_capture = False + # The warmup batches published a bucket per captured shape; the + # first real step must not inherit the last one. `fit_ragged_ + # verify_lens` clears it on entry, but a step that never reaches + # the fit -- the planner declining, or a non-ragged mode -- would + # otherwise key on a bucket from capture. + self.cuda_graph_runner.agreed_ragged_bucket = None # Absolute private-pool size, not a delta. Deltas are useless # here: an engine build runs several capture passes over a # shared allocator, so the per-pass increments came out as @@ -3694,6 +3700,16 @@ def _set_warmup_ragged_windows(self, batch, verify_bucket: int, filled = layout.fill_bucket(max_verify_len=1 + int(draft_len)) for request, tokens in zip(requests, filled.verify_lens.tolist()): request.py_verify_len = int(tokens) - 1 + # Publish it to the runner as well, or the graph is captured under a key + # that does not describe the batch just shaped. `_ragged_verify_bucket` + # reads `agreed_ragged_bucket` rather than re-summing the batch (so that + # every attention-DP rank keys on a value they agreed on, instead of one + # each derived independently) -- and nothing sets that during capture, + # because the runtime fit never runs here. The key's token axis was + # therefore None while the batch really held `verify_bucket` tokens, and + # the mismatch surfaced far away, in the MoE's shared-vs-routed row + # count, as `unmatched tensor shape` two graphs into the capture pass. + self.cuda_graph_runner.agreed_ragged_bucket = int(verify_bucket) def ragged_verify_token_buckets(self, padded_bs: int) -> List[int]: """Captured token totals for a ragged batch of ``padded_bs`` rows. diff --git a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py index 8979b04ea706..f788854d44ce 100644 --- a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py +++ b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py @@ -172,3 +172,51 @@ class _R: r.agreed_ragged_bucket = 512 r.spec_config = type("S", (), {"enable_ragged_verify": False})() assert get(r, object()) is None + + +def test_capture_publishes_the_bucket_it_shaped_the_batch_to(): + """The other half of the test above, and the regression it let through. + + Once the key reads a fitted value instead of re-summing the batch, whoever + shapes a batch has to publish the shape. Real steps do both in + fit_ragged_verify_lens. CUDA-graph capture shapes its warmup batch in + _set_warmup_ragged_windows and never called the fit, so the bucket stayed + None while the batch really held `verify_bucket` tokens. + + Both consequences are silent in different ways. The capture-time forward + ran with metadata describing a uniform batch over ragged tensors and died + hundreds of frames later in the MoE on `unmatched tensor shape`. Had it + survived, all three of a batch size's buckets would have been stored under + one key with no token axis -- colliding with each other, and matching no + runtime step, so every ragged step would have fallen back to eager without + an error anywhere. + """ + import types + + from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + + runner = types.SimpleNamespace(agreed_ragged_bucket=None) + requests = [types.SimpleNamespace(py_verify_len=None) for _ in range(4)] + batch = types.SimpleNamespace(generation_requests=requests) + engine = types.SimpleNamespace(cuda_graph_runner=runner) + + # bs=4, draft_len=5 -> the t=5 bucket is 4 * 6 = 24 tokens. + PyTorchModelEngine._set_warmup_ragged_windows(engine, batch, 24, 5) + + assert sum(1 + r.py_verify_len for r in requests) == 24, ( + "the warmup batch must hit the bucket exactly") + assert runner.agreed_ragged_bucket == 24, ( + "capture shaped the batch to 24 tokens but did not publish it, so the " + "graph would be keyed as if the batch were uniform") + + # A narrower bucket at the same batch size, to prove it is per-entry and + # not set once: capture walks every tier for every batch size. + PyTorchModelEngine._set_warmup_ragged_windows(engine, batch, 8, 5) + assert sum(1 + r.py_verify_len for r in requests) == 8 + assert runner.agreed_ragged_bucket == 8 + + # An empty batch publishes nothing rather than a stale or zero bucket. + runner.agreed_ragged_bucket = None + PyTorchModelEngine._set_warmup_ragged_windows( + engine, types.SimpleNamespace(generation_requests=[]), 24, 5) + assert runner.agreed_ragged_bucket is None From 79150aa05430c2fe0ebec8f5babef0b8c1fe2f35 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Mon, 3 Aug 2026 03:48:29 -0700 Subject: [PATCH 072/177] [None][fix] Gather whether row padding will happen instead of predicting it The first compact run to get past capture died on 2 of 8 ranks with ragged layout mismatch: requests carry 18 generation tokens but the CUDA-graph key's token axis is 24 and the other six then deadlocked -- a rank that throws inside the forward stops issuing collectives -- until the hang detector hard-killed the job at 300s. The hang was the symptom; this is the cause. The ragged fit has to know what row padding will do, because it budgets tokens for the pad rows: with 3 real requests and a cross-rank maximum of 4 it sizes the bucket at 4*6 = 24 and hands 3*6 = 18 to the real rows, leaving 6 for the row that padding is expected to append. But _get_padded_batch (pre-existing, upstream: this logic is on main) only takes the cross-rank row maximum when EVERY rank can graph the step: all_can_graph = all(peer[0] for peer in tp_allgather([can_run_cuda_graph, bs])) if all_can_graph: new_batch_size = max(...) One rank holding a context request makes each rank keep its own row count, no padding is appended, and the 6 tokens were budgeted for a row that never arrives. will_pad_to, added earlier to guard exactly this, mirrors only the SIZE guards -- it never saw the capability answer, which the graph runner gathers after the fit has already run. Its docstring claimed being wrong optimistically "leaves the pre-existing behaviour rather than making it worse"; that was wrong, and this is what it costs. Fix: stop predicting. can_run_cuda_graph rides in the DSpark hook's existing allgather (4 ints -> 5, no new synchronization point), and the group's answer is folded into peer_stats so the fit reads the same data the graph runner will act on rather than re-deriving it. Same shape of bug as d2c19dea01, where a value every rank must agree on was being derived twice; that fix covered the bucket and missed the padding decision. Also widens choose_ragged_capture_shape's peer-stat unpack, which was a two-tuple and would have raised on the third field. Only reachable with real trimming and a step where some rank holds a context request, which is why it needed an end-to-end run with a steep cost table to surface at all. 353 unit tests pass; test_spec_gate_e2e fails on a missing EAGLE3 checkpoint, unrelated and pre-existing. Signed-off-by: Lance Liao --- .../_torch/attention_backend/trtllm.py | 41 ++++++++++++++++- .../_torch/pyexecutor/model_engine.py | 44 +++++++++++++++---- tensorrt_llm/_torch/pyexecutor/py_executor.py | 24 ++++++++-- .../_torch/speculative/dspark_ragged.py | 7 ++- .../hw_agnostic/test_dspark_ragged.py | 13 ++++++ 5 files changed, 114 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 0438fbd85c6f..0fd94d4d63cd 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1846,7 +1846,7 @@ def forward( token_major_view = (metadata.token_major_gen_view() if forward_args.attention_input_type == AttentionInputType.generation_only else None) - with metadata.presented_token_major(token_major_view): + with self._presented_token_major(metadata, token_major_view): for fmha in self.fmha_libs: if fmha.is_supported(q, k, v, metadata, forward_args): fmha.forward(q, k, v, metadata, forward_args) @@ -2122,9 +2122,46 @@ def mla_rope_generation( cu_kv_seqlens, fmha_scheduler_counter, mla_bmm1_scale, mla_bmm2_scale, quant_q_buffer, out_scale, helix_tensor_params) + @contextlib.contextmanager + def _presented_token_major(self, metadata, view): + """Metadata swap plus the one op parameter that lives on the module. + + ``predicted_tokens_per_seq`` is the op's ``s_q``. Under this + presentation every row IS a length-one sequence, and the op already + knows it -- the generation path recomputes ``s_q = acc_q_len / + batch_beam`` and asserts it equals this field + (``common/attentionOp.cpp:1232``). That assert has been false here all + along: the runtime value is 1 while the field still says + ``max_draft_len + 1``. It never fired because release builds compile + ``assert`` out. + + The visible cost is workspace. Sizing reads the FIELD + (``getWorkspaceSizeForGeneration``, s_q = 6) while the run reads the + RECOMPUTED value (s_q = 1), so the generation workspace -- whose main + terms are all linear in ``batch_beam * s_q`` -- is allocated six times + larger than the layout it ends up holding. On top of the row ceiling + this presentation already passes as ``max_num_requests``, that is where + the reported 68 MiB -> 604 MiB came from. Over-allocation is safe, which + is why nothing ever failed. + + ``mMLAParams`` is rebuilt from the argument on every call + (``thop/attentionOp.cpp:1285``), so telling the truth here is enough. + """ + if view is None: + yield + return + saved = self.predicted_tokens_per_seq + try: + self.predicted_tokens_per_seq = 1 + with metadata.presented_token_major(view): + yield + finally: + self.predicted_tokens_per_seq = saved + @contextlib.contextmanager def presented_token_major_for(self, metadata): - with metadata.presented_token_major(metadata.token_major_gen_view()): + with self._presented_token_major(metadata, + metadata.token_major_gen_view()): yield def _mla_rope_generation_impl(self, metadata, fused_q, q_pe, latent_cache, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 73802bcd2ed3..62d29feff8f8 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3764,11 +3764,18 @@ def fit_ragged_verify_lens( out of graph replay. ``peer_stats`` is every attention-DP rank's ``[num_requests, - total_tokens]``, including this rank's. It is required under ADP - because ``_get_padded_batch`` pads to the *maximum* batch size across - ranks: sizing the bucket from the local count alone makes each rank - pick a different shape, which is the graph-key divergence this whole - path exists to avoid. + total_tokens, all_can_graph]``, including this rank's. The first two + are required under ADP because ``_get_padded_batch`` pads to the + *maximum* batch size across ranks: sizing the bucket from the local + count alone makes each rank pick a different shape, which is the + graph-key divergence this whole path exists to avoid. + + The third is the group's answer to "will that maximum be taken at + all", identical on every rank. ``_get_padded_batch`` only takes it when + every rank can graph the step; one rank holding a context request makes + each keep its own row count instead. It is gathered rather than + predicted because the graph runner's own allgather, which decides this, + runs after the fit. """ runner = self.cuda_graph_runner # Cleared on entry so that every early return leaves no stale @@ -3804,9 +3811,18 @@ def fit_ragged_verify_lens( return None # The token bucket grid depends on the row count, which itself depends # on the widest rank, so resolve the rows first and then the tokens. - widest_rows = max( - [int(s[0]) for s in peer_stats], - default=len(token_lens)) if peer_stats else len(token_lens) + # Whether _get_padded_batch will take the cross-rank row maximum at all + # is a group decision, gathered with the peer stats rather than guessed + # here: it only does so when EVERY rank can graph this step. One rank + # holding a context request makes each rank keep its own row count, and + # a fit that assumed the maximum would budget tokens for pad rows that + # never arrive. Third element of a peer stat; absent means single-rank, + # where the local count IS the maximum. + all_can_graph = (bool(peer_stats[0][2]) if peer_stats + and len(peer_stats[0]) > 2 else True) + widest_rows = (max([int(s[0]) for s in peer_stats], + default=len(token_lens)) if peer_stats + and all_can_graph else len(token_lens)) padded_bs = runner._round_up_batch_size(widest_rows) if padded_bs == 0: return None @@ -3829,6 +3845,18 @@ def fit_ragged_verify_lens( # Decline to go ragged rather than fit against rows that will not exist. # Uniform always has a captured graph, so the fallback is cheap and the # step still runs. + # `will_pad_to` mirrors _get_padded_batch's SIZE guards. It cannot see + # the group's graph-capability answer -- that is what `all_can_graph` + # above now supplies -- and it deliberately does not model the + # resource-dependent bails (KV capacity for the dummy). Being wrong + # optimistically there does not "leave the pre-existing behaviour", as + # an earlier version of this comment claimed: it budgets tokens for a + # row that never appears, so the layout carries n_real*(w+1) tokens + # while the graph key says padded_bs*(w+1). Observed as + # requests carry 18 generation tokens but the key's token axis is 24 + # on 2 of 8 ranks, which threw inside the forward, and under + # attention-DP a rank that throws stops issuing collectives -- the + # other six deadlocked and the run was hard-killed after 300s. if not runner.will_pad_to(padded_bs, len(token_lens)): logger.debug( f"DSpark ragged: padding to {padded_bs} rows is not available " diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 843beda6b397..9d556fdb1493 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3422,10 +3422,21 @@ def _dspark_confidence_draft_len(self, # # Carries everything the rest of the step needs, so the `peer_stats` # gather that used to follow is folded in here rather than added to it: - # [wants_ragged, num_rows, total_verify_tokens, max_window] + # [wants_ragged, num_rows, total_verify_tokens, max_window, can_graph] # `_can_queue` already runs `tp_allgather(batch_size)` just before this # hook and could absorb this too, but that would mean computing the # local draft length inside a hot path shared by every speculation mode. + # + # `can_graph` rides along because the ragged fit has to know what row + # padding will actually do, and it runs before the graph runner's own + # allgather answers that. `_get_padded_batch` only takes the cross-rank + # row maximum when EVERY rank can graph this step; one rank holding a + # context request makes each rank keep its local count instead, so a fit + # that assumed the maximum budgets tokens for pad rows that are never + # appended. Predicting that from size rules alone is what `will_pad_to` + # tried and got wrong. Carrying the input costs nothing here -- the + # collective is already being issued, and one more int does not change + # its cost -- and it replaces a second derivation with the same data. local_rows = len(ragged_lens) if ragged_lens is not None else 0 local_tokens = (sum(1 + int(v) for v in ragged_lens) if ragged_lens else 0) @@ -3437,6 +3448,7 @@ def _dspark_confidence_draft_len(self, local_rows, local_tokens, local_max_window, + 1 if scheduled_batch.can_run_cuda_graph else 0, ]) if not all(int(payload[0]) for payload in payloads): # Any rank that cannot go ragged takes the whole group with it: @@ -3450,8 +3462,14 @@ def _dspark_confidence_draft_len(self, else: agreed_max = max(int(payload[3]) for payload in payloads) ragged_lens = [min(int(v), agreed_max) for v in ragged_lens] - peer_stats = [[int(payload[1]), - int(payload[2])] for payload in payloads] + # Third element is the group's answer, identical on every rank + # by construction, not each rank's own flag: what the fit needs + # to know is whether the row maximum will be taken at all. + all_can_graph = all(int(payload[4]) for payload in payloads) + peer_stats = [[ + int(payload[1]), + int(payload[2]), 1 if all_can_graph else 0 + ] for payload in payloads] # --- Phase 3: no more collectives ------------------------------------- # diff --git a/tensorrt_llm/_torch/speculative/dspark_ragged.py b/tensorrt_llm/_torch/speculative/dspark_ragged.py index a1b36a1821ea..0fdb42e58b98 100644 --- a/tensorrt_llm/_torch/speculative/dspark_ragged.py +++ b/tensorrt_llm/_torch/speculative/dspark_ragged.py @@ -366,8 +366,11 @@ def choose_ragged_capture_shape( ) stats = list(peer_stats) if peer_stats else [(num_real_requests, total_verify_tokens)] - max_real = max(int(real) for real, _ in stats) - max_slack = max(int(total) - int(real) for real, total in stats) + # Tolerate extra trailing fields: the caller's payload has grown before + # (the group's graph-capability answer rides along with the row and token + # counts) and only the first two are this function's business. + max_real = max(int(peer[0]) for peer in stats) + max_slack = max(int(peer[1]) - int(peer[0]) for peer in stats) sorted_bs = sorted({int(b) for b in bs_buckets}) if max_real > sorted_bs[-1]: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py index c1d9ce23232f..bc0c6476655f 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py @@ -687,3 +687,16 @@ def test_pad_length_never_yields_an_infeasible_real_target(): # 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" + + +def test_extra_peer_stat_fields_are_ignored_here(): + """The payload grows; this function's contract does not. + + peer_stats gained a third element -- the group's answer to whether row + padding will take the cross-rank maximum at all -- which the caller uses + before reaching this function. A two-element unpack here turned that into + a ValueError on the first ragged step. + """ + two = _shape(num_real=3, total=12, peers=[[3, 12], [4, 16]]) + three = _shape(num_real=3, total=12, peers=[[3, 12, 1], [4, 16, 1]]) + assert (three.padded_bs, three.bucket) == (two.padded_bs, two.bucket) From b09580436b33d68e6e73a191f330a6ec9732c405 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Mon, 3 Aug 2026 04:43:39 -0700 Subject: [PATCH 073/177] [None][fix] Decline the ragged fit when no rank can graph the step Two things, both from measurement. 1. Revert the predicted_tokens_per_seq swap under token-major. The reasoning was sound -- the op recomputes s_q = acc_q_len / batch_beam and asserts it equals this field, so under token-major the field has been lying (1 vs 6), and sizing reads the field while the run reads the recomputed value. The run refuted the conclusion drawn from it: the workspace was 604727808 bytes before and after, byte for byte. The 68 MiB -> 604 MiB resize is triggered by Run warmup with 4096 tokens, include 0 generation tokens i.e. it is the PREFILL workspace at max_num_tokens, and getWorkspaceSize returns max(context, generation), so the context term dominates and the generation term is invisible. That also retracts the claim I made when reading the code: 68 -> 604 MiB is not a token-major cost, and the doc's original "no control, do not rely on it" was right. A semantic parameter that also drives isMultiTokenGen dispatch, changed for a benefit that measurement says is zero, does not earn its place. 2. Fix the ragged fit properly. The previous attempt gathered all_can_graph and used it to pick widest_rows, which did not help: compact hung again on the identical assertion, 18 tokens against a key of 24. Rounding 3 local rows up still lands on padded_bs 4, so the budget still assumed a fourth row. The real point is that all_can_graph=False means SOME rank holds a context request, so no rank replays a graph this step -- and the entire fit exists to land on a captured one. Continuing also collides with a second mismatch: _get_padded_batch pads the TOTAL batch size (context + generation) while this fit counts generation rows only. So decline and take the uniform path. Conservative rather than necessary: an eager ragged step would still save real compute. Making that safe means teaching the bucket machinery that a step may have no graph at all, which is more than the saving justifies on a mixed batch. 331 unit tests pass. Signed-off-by: Lance Liao --- .../_torch/attention_backend/trtllm.py | 41 +------------------ .../_torch/pyexecutor/model_engine.py | 26 ++++++++++-- 2 files changed, 25 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 0fd94d4d63cd..0438fbd85c6f 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1846,7 +1846,7 @@ def forward( token_major_view = (metadata.token_major_gen_view() if forward_args.attention_input_type == AttentionInputType.generation_only else None) - with self._presented_token_major(metadata, token_major_view): + with metadata.presented_token_major(token_major_view): for fmha in self.fmha_libs: if fmha.is_supported(q, k, v, metadata, forward_args): fmha.forward(q, k, v, metadata, forward_args) @@ -2122,46 +2122,9 @@ def mla_rope_generation( cu_kv_seqlens, fmha_scheduler_counter, mla_bmm1_scale, mla_bmm2_scale, quant_q_buffer, out_scale, helix_tensor_params) - @contextlib.contextmanager - def _presented_token_major(self, metadata, view): - """Metadata swap plus the one op parameter that lives on the module. - - ``predicted_tokens_per_seq`` is the op's ``s_q``. Under this - presentation every row IS a length-one sequence, and the op already - knows it -- the generation path recomputes ``s_q = acc_q_len / - batch_beam`` and asserts it equals this field - (``common/attentionOp.cpp:1232``). That assert has been false here all - along: the runtime value is 1 while the field still says - ``max_draft_len + 1``. It never fired because release builds compile - ``assert`` out. - - The visible cost is workspace. Sizing reads the FIELD - (``getWorkspaceSizeForGeneration``, s_q = 6) while the run reads the - RECOMPUTED value (s_q = 1), so the generation workspace -- whose main - terms are all linear in ``batch_beam * s_q`` -- is allocated six times - larger than the layout it ends up holding. On top of the row ceiling - this presentation already passes as ``max_num_requests``, that is where - the reported 68 MiB -> 604 MiB came from. Over-allocation is safe, which - is why nothing ever failed. - - ``mMLAParams`` is rebuilt from the argument on every call - (``thop/attentionOp.cpp:1285``), so telling the truth here is enough. - """ - if view is None: - yield - return - saved = self.predicted_tokens_per_seq - try: - self.predicted_tokens_per_seq = 1 - with metadata.presented_token_major(view): - yield - finally: - self.predicted_tokens_per_seq = saved - @contextlib.contextmanager def presented_token_major_for(self, metadata): - with self._presented_token_major(metadata, - metadata.token_major_gen_view()): + with metadata.presented_token_major(metadata.token_major_gen_view()): yield def _mla_rope_generation_impl(self, metadata, fused_q, q_pe, latent_cache, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 62d29feff8f8..9580148a2b26 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3820,9 +3820,29 @@ def fit_ragged_verify_lens( # where the local count IS the maximum. all_can_graph = (bool(peer_stats[0][2]) if peer_stats and len(peer_stats[0]) > 2 else True) - widest_rows = (max([int(s[0]) for s in peer_stats], - default=len(token_lens)) if peer_stats - and all_can_graph else len(token_lens)) + if not all_can_graph: + # Some rank holds a context request, so no rank will replay a graph + # this step -- and everything below exists to land the batch on a + # captured one. Two things then go wrong at once if we continue: + # _get_padded_batch declines the cross-rank row maximum AND pads the + # TOTAL batch size (context + generation) while this fit counts + # generation rows only, so its budget goes to rows that never + # appear. Observed as `requests carry 18 generation tokens but the + # CUDA-graph key's token axis is 24` -- 3 generation rows fitted as + # 4 -- which throws inside the forward and, since a rank that throws + # stops issuing collectives, deadlocks every other rank behind it. + # + # An eager ragged step would still save real compute, so this is + # conservative rather than necessary; making it safe means teaching + # the bucket machinery that a step may have no graph at all, which + # is a larger change than the saving justifies on a mixed batch. + logger.debug( + "DSpark ragged: some rank cannot run a CUDA graph this step, " + "so no captured bucket applies; falling back to uniform") + return None + widest_rows = max([int(s[0]) for s in peer_stats], + default=len(token_lens)) if peer_stats else len( + token_lens) padded_bs = runner._round_up_batch_size(widest_rows) if padded_bs == 0: return None From f033e3cf5bb33c57b3831b79bf8fd7eae8e28757 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Mon, 3 Aug 2026 05:22:09 -0700 Subject: [PATCH 074/177] [None][fix] Run ragged without a bucket on steps that have no graph The previous commit declined ragged entirely when no rank could graph the step. That stopped the crash but gave away more than it had to: 46 of the first 64 steps fell back, reported as no_captured_shape. Those steps were already eager. can_run_cuda_graph is num_context_requests == 0 and predates this feature by a long way, so a batch holding a context request never replayed a graph -- this branch is not what pushes them out of replay. What declining costs is the trimming itself, which is real compute saved and is just as available on an eager step. Aligning the token total to a captured bucket exists only to make the graph key match. With no graph there is no key, so the alignment is a constraint this feature imposes on itself. Worse, on exactly these steps it was harmful: the fit sizes padding from generation rows while _get_padded_batch pads the TOTAL batch, so it budgeted tokens for a row that never appeared -- 18 against a key of 24, which threw and deadlocked the peers. So: no graph -> emit the planner's windows as chosen, no round-up, and leave agreed_ragged_bucket cleared so the graph key keeps its uniform shape. The layout assertion already tolerates a missing bucket. `bucket is not None` was doing double duty as "ragged happened"; that is now a separate ragged_active, since ragged can now happen without one. Credit where due: the observation that these steps could not have graphed before this PR either, and that declining therefore buys nothing, came from review rather than from me. 331 unit tests pass. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 36 +++++++++++++++++-- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9d556fdb1493..c21e9aaa3cda 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3442,6 +3442,9 @@ def _dspark_confidence_draft_len(self, for v in ragged_lens) if ragged_lens else 0) local_max_window = max(ragged_lens) if ragged_lens else 0 peer_stats = None + # Single-rank runs have no peer that could be holding a context + # request, so the local batch's own answer is the group's. + all_can_graph = bool(scheduled_batch.can_run_cuda_graph) if is_distributed: payloads = self.dist.tp_allgather([ 1 if ragged_lens is not None else 0, @@ -3491,6 +3494,7 @@ def _dspark_confidence_draft_len(self, # addition rather than a second set of conditions inside code that took # three bugs to get right for `compact`. bucket = None + ragged_active = False cap_accept = (mode is not None and mode.computes_windows and not mode.trims_submitted_tokens) if cap_accept: @@ -3505,11 +3509,33 @@ def _dspark_confidence_draft_len(self, if cap_accept: for request, window in zip(gen_requests, ragged_lens): request.py_verify_cap = int(window) - else: + ragged_active = True + elif all_can_graph: bucket = self.model_engine.fit_ragged_verify_lens( gen_requests, ragged_lens, peer_stats=peer_stats) + ragged_active = bucket is not None if bucket is None: fallback_reason = "no_captured_shape" + else: + # No rank can replay a graph this step -- some batch holds a + # context request, which has always meant eager, long before + # this feature existed. Fitting a captured bucket is therefore + # pointless: there is no key to match, and the round-up would + # only add tokens nobody asked for. Worse, the fit sizes its + # padding from generation rows while `_get_padded_batch` pads + # the TOTAL batch, so on exactly these steps it budgeted for a + # row that never appeared -- 18 tokens against a key of 24, + # which threw and deadlocked the other ranks. + # + # So run ragged WITHOUT a bucket: the windows go out as the + # planner chose them and the token total is whatever it is. + # Trimming still saves real compute here; only the alignment is + # unnecessary. `agreed_ragged_bucket` stays cleared so the graph + # key keeps its uniform shape. + self.model_engine.cuda_graph_runner.agreed_ragged_bucket = None + for request, window in zip(gen_requests, ragged_lens): + request.py_verify_len = int(window) + ragged_active = True # Record what the step actually decided, reading the windows back off # the requests rather than off `ragged_lens`: fit_ragged_verify_lens @@ -3532,7 +3558,11 @@ def _dspark_confidence_draft_len(self, # the windows would be booked as if they had been submitted. fitted = [int(v) for v in ragged_lens] if ragged_lens else None delivered = num_gen_requests * (1 + max_draft_len) - elif bucket is not None: + elif ragged_active: + # Read back off the requests, not off `ragged_lens`: when a + # bucket was fitted its rounding slack was spent on real + # verification, so the fitted windows are what the target sees. + # Without a bucket the two are the same thing. fitted = [ int(request.py_verify_len) for request in gen_requests if getattr(request, "py_verify_len", None) is not None @@ -3546,7 +3576,7 @@ def _dspark_confidence_draft_len(self, fallback=fallback_reason, delivered=delivered) - if bucket is not None: + if ragged_active and not cap_accept: # The block is always drafted in full -- only verification is # trimmed -- so the batch-wide draft length stays at the top tier. # That keeps the drafted-token buffer and the padded acceptance From 59e2d47b412d6e5128bc19317d7de1973bf9f92e Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Mon, 3 Aug 2026 06:33:55 -0700 Subject: [PATCH 075/177] [None][fix] Clear the ragged pad window on steps that skip the bucket fit fit_ragged_verify_lens sets three pieces of state; the eager-ragged path added in f033e3cf5b set only one. agreed_ragged_bucket was cleared deliberately and _dspark_last_total_tokens is write-only, but ragged_pad_verify_len was simply missed -- and it is read by _get_padded_batch to stamp the padding dummy. The subtlety is that "no rank can graph" does not mean "no rank pads". _get_padded_batch gates on the LOCAL can_run_cuda_graph, so a rank whose own batch is generation-only still pads even while the group has declined the cross-rank maximum. Those ranks stamped the dummy with the previous ragged step's window -- a value from a layout that no longer exists. The run that found it died with CUDA illegal memory access on rank 0 and unspecified launch failure on its peers. With this cleared the same configuration completes: GSM8K 96.323, 128 steps, 40 ragged, no ragged- attributable graph misses. Fourth stale-persistent-state bug on this branch, after py_verify_len, the cached padding dummy, and the capture-time bucket. The shape is always the same: a value written on one path and read on another that does not write it. Not yet confirmed without CUDA_LAUNCH_BLOCKING=1, which was on to locate the IMA and also perturbs timing; a confirmation run is in flight. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index c21e9aaa3cda..b6408dd013bf 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3532,7 +3532,15 @@ def _dspark_confidence_draft_len(self, # Trimming still saves real compute here; only the alignment is # unnecessary. `agreed_ragged_bucket` stays cleared so the graph # key keeps its uniform shape. - self.model_engine.cuda_graph_runner.agreed_ragged_bucket = None + runner = self.model_engine.cuda_graph_runner + runner.agreed_ragged_bucket = None + # A rank whose OWN batch is gen-only still pads locally even + # when the group cannot graph, and it stamps the padding dummy + # with this. Skipping the fit leaves the previous ragged step's + # value, so the dummy would join the batch carrying a window + # from a layout that no longer exists. There is no bucket to + # spend here, so the pad rows get the minimum. + runner.ragged_pad_verify_len = 0 for request, window in zip(gen_requests, ragged_lens): request.py_verify_len = int(window) ragged_active = True From 4f81cc553472f4f6c3a25a3b3398d3ab2b2821ec Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Mon, 3 Aug 2026 19:33:56 -0700 Subject: [PATCH 076/177] [None][fix] Gate the eager-ragged path off until its IMA is understood f033e3cf5b let steps with no CUDA graph run ragged verification. That path triggers a CUDA illegal memory access -- one rank reports it, its peers report `unspecified launch failure`, and the run either fails or deadlocks. It is on the pushed branch, so this turns it off by default behind TLLM_DSPARK_EAGER_RAGGED rather than leaving a known crash reachable. Not reverted, because the reasoning it rests on holds: those steps are eager regardless of this feature (can_run_cuda_graph is num_context_requests == 0, which long predates it), so declining ragged buys no graph replay and only gives up the trimming. Measured: 46 of the first 64 steps fell back, trim_ratio 0.62 -> 0. Keeping the code behind a flag also keeps the reproduction to one line, which the investigation needs. What is known about the fault, after nine falsified hypotheses: - CUDA_LAUNCH_BLOCKING=1 completes a full 128-step run (GSM8K 96.323). A barrier at the scheduling point does not help, and neither does wait_stream against the drafter's stream. So the race is inside the forward, not at the step boundary. - Setting every window to the FULL block -- which makes the step semantically identical to static verification -- still faults. So it is not about window values, raggedness, or trimming. Merely publishing ragged metadata on a step that then runs eager is unsafe. - Ruled out by runs: stale ragged_pad_verify_len, temporary pinned buffers as async H2D sources, ragged on a mixed batch, cross-step overlap (disable_overlap_scheduler still faults), bucket alignment and the fit, per-step buffer allocation, window raggedness. - Ruled out by reading: DeepGEMM 32-token alignment, row-buffer overflow, buffer-pool use-after-free (eager metadata has cuda_graph_buffers=None), metadata rebuilt per step, MoE aux-stream overlap (multi-stream is only enabled under graph capture/replay). Also retracts an earlier claim: a single torch.cuda.synchronize() appeared to fix it, but that run died at step 67 on the ragged-active gate, before the step range where the fault occurs. Re-run to completion, it faults. Keeps two fixes that are correct on their own merits even though neither resolved the IMA: _pinned_host (persistent staging buffers for async H2D, since PyTorch does not extend an async source's lifetime) and the diagnostic TLLM_DSPARK_DISABLE_OVERLAP switch on the integration test. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/model_engine.py | 60 ++++++++++++++----- tensorrt_llm/_torch/pyexecutor/py_executor.py | 52 +++++++++------- .../defs/accuracy/test_llm_api_pytorch.py | 7 ++- 3 files changed, 83 insertions(+), 36 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 9580148a2b26..703e4a3754f2 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4114,6 +4114,43 @@ def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, f"presentation is {expected_rows} rows -- that tensor's shape " f"IS what the ops read as their sequence count.") + def _pinned_host(self, key: str, values, dtype) -> torch.Tensor: + """A persistent pinned staging buffer holding ``values``. + + Every async host-to-device copy on the ragged path used to source from + a freshly built ``torch.tensor(..., pin_memory=True)``. Nothing keeps + that temporary alive: the statement ends, Python is free to reclaim the + page, and ``non_blocking=True`` means the copy has only been QUEUED. + PyTorch does not extend the lifetime of an async H2D source -- that is + the caller's job -- so the DMA can read a freed page. + + It survives most of the time, which is why this went unnoticed: it + needs the allocator to actually hand the page back before the copy + drains. Raising the number of ragged steps per run from 3 to 35 was + enough to turn it from never-seen into a run-killer, surfacing as + `CUDA error: an illegal memory access` on one rank and `unspecified + launch failure` on its peers, hundreds of frames from the copy. + CUDA_LAUNCH_BLOCKING=1 makes the copy synchronous and hides it + completely, which is how it was identified. + + Buffers are keyed by name and grown monotonically, so a steady state + allocates nothing. + """ + values = list(values) + cache = getattr(self, "_pinned_host_cache", None) + if cache is None: + cache = self._pinned_host_cache = {} + buf = cache.get(key) + if buf is None or buf.numel() < len(values) or buf.dtype != dtype: + buf = torch.empty(max(len(values), 1), + dtype=dtype, + pin_memory=prefer_pinned()) + cache[key] = buf + view = buf[:len(values)] + if values: + view.copy_(torch.tensor(values, dtype=dtype)) + return view + def _attach_accept_caps(self, spec_metadata, generation_requests) -> None: """Publish ``cap-accept`` ceilings, without touching the layout. @@ -4157,9 +4194,9 @@ def _attach_accept_caps(self, spec_metadata, generation_requests) -> None: # verify_lens: acceptance runs inside the target's captured graph, so a # freshly allocated tensor here would be invisible to every replay. caps_view = self.accept_caps_cuda[:n] - caps_view.copy_(torch.tensor([1 + int(cap) for cap in caps], - dtype=torch.int32, - pin_memory=prefer_pinned()), + caps_view.copy_(self._pinned_host("accept_caps", + [1 + int(cap) for cap in caps], + torch.int32), non_blocking=True) spec_metadata.accept_caps = caps_view spec_metadata.accept_cap_trim = self.accept_cap_trim_cuda @@ -4193,9 +4230,8 @@ def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, # tensor: a captured graph baked in the address it saw at capture time, # so a new allocation here would be invisible to every replay. lens_view = self.ragged_verify_lens_cuda[:n] - lens_view.copy_(torch.tensor(token_lens, - dtype=torch.int32, - pin_memory=prefer_pinned()), + lens_view.copy_(self._pinned_host("ragged_verify_lens", token_lens, + torch.int32), non_blocking=True) indptr_view = self.ragged_qo_indptr_cuda[:n + 1] indptr_view.copy_(build_qo_indptr(lens_view), non_blocking=True) @@ -4218,14 +4254,10 @@ def _ragged_gather_indices( verifies the same number of positions. """ rows, cols = ragged_gather_index_lists(slots, counts) - rows_host = torch.tensor(rows, - dtype=torch.long, - pin_memory=prefer_pinned()) - cols_host = torch.tensor(cols, - dtype=torch.long, - pin_memory=prefer_pinned()) - return (rows_host.to('cuda', non_blocking=True), - cols_host.to('cuda', non_blocking=True)) + return (self._pinned_host("gather_rows", rows, + torch.long).to('cuda', non_blocking=True), + self._pinned_host("gather_cols", cols, + torch.long).to('cuda', non_blocking=True)) def _update_target_input_tensors( self, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index b6408dd013bf..b2eec8b185f3 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3516,34 +3516,44 @@ def _dspark_confidence_draft_len(self, ragged_active = bucket is not None if bucket is None: fallback_reason = "no_captured_shape" - else: - # No rank can replay a graph this step -- some batch holds a - # context request, which has always meant eager, long before - # this feature existed. Fitting a captured bucket is therefore - # pointless: there is no key to match, and the round-up would - # only add tokens nobody asked for. Worse, the fit sizes its - # padding from generation rows while `_get_padded_batch` pads - # the TOTAL batch, so on exactly these steps it budgeted for a - # row that never appeared -- 18 tokens against a key of 24, - # which threw and deadlocked the other ranks. + elif os.environ.get("TLLM_DSPARK_EAGER_RAGGED"): + # OFF BY DEFAULT. This path triggers a CUDA illegal memory + # access that is not yet located. + # + # The reasoning behind it is sound and is why the code is kept + # rather than reverted: these steps are eager no matter what + # this feature does -- can_run_cuda_graph is + # `num_context_requests == 0`, which predates it by a long way + # -- so declining ragged buys no graph replay. It only gives up + # the trimming, and that was measured: 46 of the first 64 steps + # fell back, trim_ratio 0.62 -> 0. + # + # What is NOT understood is why publishing ragged metadata on a + # step that then runs eager faults at all. Nine hypotheses have + # been falsified, most usefully: setting every window to the + # FULL block, which makes the step semantically identical to + # static verification, still faults -- so it is not about + # windows, raggedness, or trimming. A barrier at this point does + # not help while CUDA_LAUNCH_BLOCKING=1 completes a full 128-step + # run, which places the race inside the forward. # - # So run ragged WITHOUT a bucket: the windows go out as the - # planner chose them and the token total is whatever it is. - # Trimming still saves real compute here; only the alignment is - # unnecessary. `agreed_ragged_bucket` stays cleared so the graph - # key keeps its uniform shape. + # The env var keeps the reproduction to one line. runner = self.model_engine.cuda_graph_runner runner.agreed_ragged_bucket = None - # A rank whose OWN batch is gen-only still pads locally even - # when the group cannot graph, and it stamps the padding dummy - # with this. Skipping the fit leaves the previous ragged step's - # value, so the dummy would join the batch carrying a window - # from a layout that no longer exists. There is no bucket to - # spend here, so the pad rows get the minimum. runner.ragged_pad_verify_len = 0 for request, window in zip(gen_requests, ragged_lens): request.py_verify_len = int(window) ragged_active = True + else: + # No rank can replay a graph this step, so there is no captured + # shape to land on. Fitting one anyway budgets tokens for pad + # rows `_get_padded_batch` will not append -- it pads the TOTAL + # batch while this fit counts generation rows -- which is how + # `requests carry 18 generation tokens but the key's token axis + # is 24` reached two ranks, threw inside the forward, and + # deadlocked the other six behind it. + fallback_reason = "no_graph_this_step" + ragged_lens = None # Record what the step actually decided, reading the windows back off # the requests rather than off `ragged_lens`: fit_ragged_verify_lens diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index bfa7832454dd..57575cad3271 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4180,7 +4180,12 @@ def test_gsm8k_dep8_ragged_verify(self, moe_backend): kv_cache_config=kv_cache_config, cuda_graph_config=CudaGraphConfig(enable_padding=True), enable_chunked_prefill=False, - disable_overlap_scheduler=False, + # Diagnostic override only. The overlap scheduler is the + # configuration this test exists to cover, so it stays on by + # default; turning it off isolates whether a failure needs the + # previous step's kernels to still be in flight. + disable_overlap_scheduler=bool( + os.environ.get("TLLM_DSPARK_DISABLE_OVERLAP")), custom_tokenizer="deepseek_v4", speculative_config=spec_config) as llm: score, acc_params = self._run_gsm8k(llm) From cdc4538616d93a48b5b70a63737df73fc8e34e84 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Mon, 3 Aug 2026 20:02:41 -0700 Subject: [PATCH 077/177] [None][chore] Name the token-major row-map staleness instead of faulting refresh_token_major_gen_rows gathers kv_lens through attn_row_req_idx_cuda, and its docstring states the map's addresses "are pinned by the graph key". That holds on the graph path. On an eager one the same metadata object serves ragged and non-ragged steps in turn, so the map can be stale relative to this step's num_seqs -- 64 requests staged last step, 3 sequences now -- and the gather reads request 63 out of a length-3 tensor. TLLM_DSPARK_ASSERT_ROWS checks that before the gather. Opt-in, because it costs a device-to-host read. Either it fires and names the eager-ragged fault, or it stays quiet and removes this candidate cleanly; both beat reading an illegal memory access reported hundreds of frames downstream. Signed-off-by: Lance Liao --- .../_torch/attention_backend/sparse/dsa.py | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 520507001a26..da9f4b34ff05 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -2091,6 +2091,27 @@ def refresh_token_major_gen_rows(self) -> None: rows = self._attn_num_rows if not self.is_ragged_verify or rows <= 0: return + # TLLM_DSPARK_ASSERT_ROWS: name the failure instead of taking an + # illegal memory access several hundred frames later. + # + # The docstring above says the row->request map's addresses "are + # pinned by the graph key". That is true on the graph path and false + # on an eager one, where the same metadata object serves ragged and + # non-ragged steps in turn. If `_attn_num_rows` and + # `attn_row_req_idx_cuda` are stale relative to this step's + # `num_seqs` -- 64 requests last step, 3 now -- the gather below reads + # request 63 out of a length-3 tensor, which is exactly the shape of + # the fault seen on the eager-ragged path. + # + # Costs a device-to-host read, so it is opt-in. + if os.environ.get("TLLM_DSPARK_ASSERT_ROWS"): + worst = int(self.attn_row_req_idx_cuda[:rows].max()) + assert worst < self.num_seqs, ( + f"token-major row map is stale: {rows} rows index up to " + f"request {worst} but this step has only {self.num_seqs} " + f"sequences. The map was staged by a previous prepare() and " + f"never restaged; the gather that follows would read out of " + f"bounds.") torch.index_select(self.kv_lens_cuda[:self.num_seqs], 0, self.attn_row_req_idx_cuda[:rows], out=self.attn_row_kv_lens_cuda[:rows]) From d36c9b4757eda186a844ffa681c129a15a87af37 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Tue, 4 Aug 2026 00:03:54 -0700 Subject: [PATCH 078/177] [None][doc] Hand off the ragged-verify IMA, including where I had it wrong A single document for whoever picks this up: goal, verified state, the open CUDA illegal memory access, every falsified hypothesis, and the traps. The framing correction is the point of writing it. For most of the session I believed the fault was confined to a path added on this branch (ragged on steps with no CUDA graph). The last run disproved that: with that path gated off, the default configuration still faults, on ordinary graph ragged steps (compact_gated.log -- no_graph_this_step: 46, steps_ragged: 43, IMA 1). So commit 4f81cc5534's justification is wrong and the branch is not made safe by it. The correct statement is that planner-driven ragged verification with real trimming has never completed a run without CUDA_LAUNCH_BLOCKING=1. Ten hypotheses are recorded as falsified with the run that killed each, so the next person does not repeat them. The sharpest surviving constraint is compact_eagerfull: every window set to max_draft_len makes the step semantically identical to static, and it still faults -- so the trigger is that is_ragged_verify is True, not what the windows contain. Also records the traps, because three of them cost hours: "it stopped without error" read as "it passed" (three times, including a run my own assertion killed during graph capture); post-hoc mechanism reported as cause (twice, both corrected only after being challenged); a test sweep that never executed the files it claimed were green; and capture being a blind spot unit tests cannot reach. Includes the environment notes that are not obvious, chiefly that the two SLURM accounts have wildly different fairshare and should both be queued. Also fixes the row-map assertion to skip during graph capture: its device-to-host read raised cudaErrorStreamCaptureInvalidated and took the engine down before any real step ran. Signed-off-by: Lance Liao --- docs/dspark_ragged_ima_handoff.md | 294 ++++++++++++++++++ .../_torch/attention_backend/sparse/dsa.py | 9 +- 2 files changed, 302 insertions(+), 1 deletion(-) create mode 100644 docs/dspark_ragged_ima_handoff.md diff --git a/docs/dspark_ragged_ima_handoff.md b/docs/dspark_ragged_ima_handoff.md new file mode 100644 index 000000000000..a20022c35412 --- /dev/null +++ b/docs/dspark_ragged_ima_handoff.md @@ -0,0 +1,294 @@ +# DSpark ragged verify — handoff + +Written 2026-08-04. Supersedes nothing; read alongside +`docs/dspark_p0_task_prompt.md` (the original task) and +`docs/dspark_ragged_verify_status.md` (design/status). + +Branch `dspark-p0`, fork `origin` = `github.com/lancelly/TensorRT-LLM`, +PR branch `origin/confidence_head` (NVIDIA/TensorRT-LLM#17056). + +--- + +## 0. The one thing to read first + +**There is an open CUDA illegal memory access on the ragged verify path, and my +framing of it was wrong until the last experiment of the session.** I spent most +of a day believing it was confined to a path I had just added (running ragged on +steps with no CUDA graph). The final run disproved that: with that path gated +OFF, the default configuration still faults. + +So the correct statement is: + +> Planner-driven ragged verification with real trimming has **never** completed a +> GSM8K run without `CUDA_LAUNCH_BLOCKING=1`. + +Everything in §4 was designed under the wrong hypothesis. The falsifications +are still valid as facts, but the *set* of candidates they were drawn from was +too narrow. Re-derive the candidate list from "ragged trimming is unsafe", not +from "eager is unsafe". + +--- + +## 1. Goal + +Make DSpark confidence-head verification scheduling work, and be +CUDA-graph-compatible, on DeepSeek-V4-Pro-DSpark at DEP8 (TP8 + EP8 + +attention-DP, 8x B300), `max_draft_len=5`, tiers `{1,3,5}`, overlap scheduler ON. + +Acceptance, all required: +1. Real checkpoint, not a stub. +2. GSM8K no worse than baseline. Scheduling changes only how many drafted + positions are submitted, never which token is emitted — **accuracy loss is a + bug, not a tradeoff**. +3. attention-DP verified (not just TP). +4. Ragged demonstrably *live*, not silently falling back. +5. Kernels compatible (indexer top-k, DSv4 compressor). +6. CUDA graphs actually replaying, not forced eager. +7. Overlap scheduler ON — that is the shipping configuration. + +Requirement 4 exists because **every failure mode of this feature is silent**: +the planner declining, a flat cost model, a batch missing its captured shape, a +partially-windowed batch — each still produces correct output and baseline +accuracy. GSM8K passing proves nothing on its own. + +--- + +## 2. Where it stands + +### Works, with evidence + +| | Evidence | +|---|---| +| K1/K2 unblocked, **zero C++ changes** | token-major presentation: each generation token becomes a length-1 row, so `seq_len == 1` divides any row count and `num_tokens % num_seqs == 0` passes trivially | +| GSM8K on the ragged path | TRTLLM **96.4367** (uniform baseline 96.2092) | +| cap-accept implemented | **96.096** TRTLLM + MEGAMOE_DEEPGEMM both pass | +| static reference | **96.171** | +| compact | **96.323** — but see the caveat in §3 | +| Uniform tier ladder removed | three states collapsed to two, −350 lines | +| Capture set no longer builds dead tiers | 57 graphs → 19 on the non-ragged branch | +| Planner counters readable | `decisions` was a dead counter; now increments | + +### Open + +| # | Item | +|---|---| +| **IMA** | §3. Blocks everything else. | +| #19 | Verify budget is quantised to the tier ladder. Measured cost up to **73%** of the planner's own objective. Not implemented. | +| #11 | Sync control never completed. **Do not cite any sync-related conclusion.** | +| #20 | Three-way differential incomplete (compact leg needs a non-blocking run) | + +--- + +## 3. The IMA + +### Symptom + +One rank reports `CUDA error: an illegal memory access was encountered`, peers +report `unspecified launch failure`, at decode step ~75–128 of GSM8K on DEP8 +with the steep synthetic cost table. Surfaces at `dspark.py:117`, but CUDA +errors are asynchronous — that is the next sync point, not the origin. Ranks +that throw stop issuing collectives, so the survivors deadlock and the hang +detector kills the job at 300s. **The hang is a consequence; the IMA is the +cause.** + +### The only configuration that completes + +`CUDA_LAUNCH_BLOCKING=1` — 128 steps, GSM8K 96.323 (`compact_ima.log`). That is +where the 96.323 in §2 comes from, so treat it as "the numbers are right when +the race is suppressed", not as a clean pass. + +### Scope — read this carefully + +The last run of the session (`compact_gated.log`) used the **default** +configuration with the eager-ragged path gated off: + +``` +fallbacks: {'no_graph_this_step': 46, ...} <- gating worked +steps_ragged: 43 <- all of these were GRAPH steps +IMA: 1, hang: 5 +``` + +So the fault occurs on the ordinary graph ragged path — `fit_ragged_verify_lens`, +a captured bucket, a replayed graph. It is not specific to any path added on +this branch. + +Corollary: `commit 4f81cc5534` ("Gate the eager-ragged path off until its IMA is +understood") does **not** make the branch safe. Its stated justification is +wrong. Keep the gate or drop it on its own merits, but do not rely on it. + +### Reproduce + +```bash +export LLM_MODELS_ROOT=/lustre/fsw/coreai_comparch_trtllm/laliao/llm-models +SPS_TABLE=/lustre/fsw/coreai_comparch_trtllm/laliao/dspark-runs/steep_sps.json \ + tmp/run_dspark_e2e.sh ragged TRTLLM +``` + +A steep table is **required**: at this checkpoint's ~90% acceptance the planner +correctly declines to trim against the real table, so nothing is exercised. +`tests/microbenchmarks/dspark_make_steep_sps_table.py` generates one; it carries +`SYNTHETIC: true` and **no throughput number taken with it means anything**. + +Useful switches, all wired through `tmp/run_dspark_e2e.sh`: + +| env | effect | +|---|---| +| `LAUNCH_BLOCKING=1` | `CUDA_LAUNCH_BLOCKING=1` — the only thing that passes | +| `NO_OVERLAP=1` | `disable_overlap_scheduler=True` | +| `EAGER_RAGGED=1` | re-enable the gated eager-ragged path | +| `ASSERT_ROWS=1` | row-map staleness assertion (skipped during capture) | +| `SANITIZER=1` | compute-sanitizer — see §5, not currently viable | + +--- + +## 4. Falsified hypotheses + +Each by its own run unless marked. **Caveat**: several were designed to +discriminate "eager vs graph", which §3 shows was the wrong axis. The +observations stand; their relevance may not. + +By experiment: +- stale `ragged_pad_verify_len` (`compact_confirm`) +- temporary pinned buffers as async H2D sources (`compact_pinned`) +- ragged on a mixed batch (`compact_genonly`; `steps_ragged` rose 35→69 and it + still faulted) +- cross-step overlap (`compact_nooverlap`) +- bucket alignment and the fit (`compact_fitnokey`: full fit, only the graph key + dropped) +- per-step buffer allocation (`compact_eagerpool`: a persistent pool made it + **worse**, 18 faults) +- window raggedness (`compact_eageruni`: identical windows, 12 faults) +- **trimming itself** (`compact_eagerfull`: every window set to `max_draft_len`, + making the step semantically identical to static — still faulted) +- a barrier at the scheduling point (`compact_syncfinal`, `compact_wait`) +- token-major row-map staleness (`compact_assert2`: IMA reproduced 15x, the + assertion never fired, and `steps_ragged: 74` confirms it had every chance to) + +By reading: +- DeepGEMM 32-token alignment — `scheduler_metadata_buffer_expanded` is + `(num_sms+1, 2)`, independent of token count +- row-buffer overflow — `rows <= cap` is asserted (`dsa.py:2024`) +- buffer-pool use-after-free — eager metadata has `cuda_graph_buffers=None`, so + `get_empty` returns `torch.zeros` and never touches the pool +- metadata rebuilt per step — `_set_up_attn_metadata` caches `self.attn_metadata` +- MoE aux-stream overlap — `with_multi_stream(True)` appears only in + `cuda_graph_runner` capture/replay, so multi-stream is OFF on eager steps + +--- + +## 5. Traps — I fell into each of these + +**"It stopped without error" is not "it passed."** Three times a run exited +early for an unrelated reason and I read the absence of a fault as a fix: + +- `compact_eagersync` reported IMA 0 — it died at step 67 on the + `assert_ragged_active` gate, before the step range where the fault occurs. I + wrote a task claiming the fix was found. Re-run to completion: IMA 3. +- `compact_trtllm4` (conservative fallback) showed no fault at step 64 and I + concluded the eager path was to blame. It never completed either. This is what + produced the wrong framing in §0. +- My own row-map assertion did a `.max()` device-to-host read inside the CUDA + graph capture body, raising `cudaErrorStreamCaptureInvalidated` and killing the + engine before any real step ran. That run's "IMA 0" meant nothing. + +**Always check the run reached the failure window** (step ≥ ~128, or a +non-zero `steps_ragged` in the periodic summary) before believing a negative. + +**Post-hoc mechanism is not evidence.** Twice I saw a symptom, found code that +could explain it, and reported it as the cause without checking: +- the hang: I blamed rank-local `enable_spec_decode`; `git diff` showed this + branch never touched it. The real trigger was that ranks 5 and 6 *threw* (my + own layout assertion) and stopped issuing collectives. **The ranks that did + not print a stack were the informative ones.** +- the workspace: I claimed `68 MiB → 604 MiB` was the token-major cost and + changed `predicted_tokens_per_seq` accordingly. Measured: byte-for-byte + identical. The 604 MiB is the *prefill* workspace at `max_num_tokens`, and + `getWorkspaceSize` returns `max(context, generation)`, so the generation term + is invisible. Change reverted. + +**Run the full test selection.** Removing the uniform tier ladder broke 4 tests +in `tests/unittest/llmapi/test_llm_args.py` and 1 integration test. I reported +"all green" from a sweep that never executed either file. + +**Capture is a systematic blind spot.** Three bugs this session had the same +shape — runtime path updated, capture path not: + +| | runtime | capture | +|---|---|---| +| capture set keyed on config flag | by mode | by flag | +| dead tier ladder | draft_len fixed | still built 3x graphs | +| `agreed_ragged_bucket` | set by the fit | set by nobody | + +Unit tests cannot reach capture; it needs a real pass. + +**attention-DP turns a rank-local raise into a cluster deadlock.** A diagnostic +assertion that fires on 2 of 8 ranks costs 300 seconds and produces a hang +report that names nothing. Consider reducing across ranks before raising. + +**compute-sanitizer is not currently viable** — it instruments the ~40-minute +weight load; the log grew zero bytes in 55 minutes. Restricting it to a single +rank is untried and is the obvious next move. + +--- + +## 6. Next steps + +1. **Re-derive the candidate list** from "ragged trimming is unsafe", not "eager + is unsafe". The sharpest surviving constraint is `compact_eagerfull`: with + every window at `max_draft_len` the step is semantically identical to static, + yet it faults — so the trigger is that `is_ragged_verify` is True, not what + the windows contain. Audit every `is_ragged_verify` consumer. Two remain + unexamined: `refresh_token_major_block_table` (a 4-D `index_select` that has + produced a device-side out-of-bounds before) and the four branches in + `sparse_attn_indexer`. +2. **Single-rank compute-sanitizer.** It names the kernel and the address; that + ends the guessing. Everything else in this section is another hypothesis. +3. Only after the IMA: #19 (budget quantisation, the largest measured gap), + then #11 and #20. + +### Do not + +- Do not resolve by declining ragged whenever a graph is unavailable. Those steps + were eager before this feature existed (`can_run_cuda_graph` is + `num_context_requests == 0`), so declining buys no graph replay — it only + forfeits trimming, measured at 46 of the first 64 steps, `trim_ratio` 0.62→0. +- Do not cite `docs/dspark_p0_task_prompt.md` §10 sync numbers; the control never + ran. +- Do not trust a throughput number taken with the synthetic steep table. + +--- + +## 7. Unpushed work + +`origin/confidence_head` is at `f033e3cf5b`. Three local commits, none +end-to-end verified: + +| commit | what | note | +|---|---|---| +| `59e2d47b41` | clear the ragged pad window when the fit is skipped | correct on its own merits | +| `4f81cc5534` | gate the eager-ragged path off | **stated justification is now known wrong** (§3) | +| `cdc4538616` | row-map staleness assertion | hypothesis it tested is falsified; assertion is cheap and may still be worth keeping | + +Also uncommitted: `dsa.py` carries the capture-skip fix for that assertion. + +Two fixes worth keeping regardless of the IMA, already in the above: +`_pinned_host` (persistent staging buffers — PyTorch does not extend an async +H2D source's lifetime, and three sites were passing temporaries) and the +`TLLM_DSPARK_DISABLE_OVERLAP` diagnostic switch. + +--- + +## 8. Environment + +- 8x B300 SXM6 (SM103, 275 GB/card), x86_64, node pool `bia*`, partition `batch` +- **Queue on both SLURM accounts in parallel** — `coreai_comparch_trtllm` and + `coreai_comparch_aarwlt`. Their fairshare is independent and differs + enormously: on 2026-08-03 one sat PENDING for over two hours with reason + `Priority` while the other, submitted to the same full partition minutes + later, was RUNNING in 30 seconds. +- Weights: `~/llm-models/DeepSeek-V4-Pro-DSpark` (symlink; `LLM_MODELS_ROOT` + must be exported *before* pytest starts — `MODEL_PATH` is evaluated at import) +- Run logs: `/lustre/fsw/coreai_comparch_trtllm/laliao/dspark-runs/` +- Every number in this document is x86 DGX B300. **GB300 is aarch64 with + NVLink-C2C unified memory — all timing, memory and load-path figures must be + re-measured there.** The structural conclusions transfer. diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index da9f4b34ff05..5ee3d716cb77 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -2104,7 +2104,14 @@ def refresh_token_major_gen_rows(self) -> None: # the fault seen on the eager-ragged path. # # Costs a device-to-host read, so it is opt-in. - if os.environ.get("TLLM_DSPARK_ASSERT_ROWS"): + if (os.environ.get("TLLM_DSPARK_ASSERT_ROWS") + and not torch.cuda.is_current_stream_capturing()): + # The .max() below is a device-to-host read, and this function runs + # inside the CUDA-graph capture body (on_update_kv_lens). Reading + # during capture raises cudaErrorStreamCaptureInvalidated and takes + # the whole engine down before any real step runs -- which is what + # the first attempt at this check did. Capture is also not where + # the interesting case lives: the eager path is. worst = int(self.attn_row_req_idx_cuda[:rows].max()) assert worst < self.num_seqs, ( f"token-major row map is stale: {rows} rows index up to " From 6a59b01e5376872c31e618aa5ee7797f406ae37e Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Tue, 4 Aug 2026 17:21:31 -0700 Subject: [PATCH 079/177] [None][fix] Derive a ragged graph key only when every request carries a window Root cause of the DSpark ragged-verify CUDA illegal memory access. _ragged_verify_bucket's docstring always promised 'None unless every generation request carries a window', but the code returned the fit's published agreed_ragged_bucket unconditionally. On steps where windows and bucket diverge (planner or peer declining after the fit, roster churn between fit and key derivation), the step replayed a ragged-bucket graph while prepare's token-major staging chain was skipped entirely -- _ragged_token_lens returns None, so no row->request map, no per-row KV extents, no per-row block-table gather ran. The replay consumed the PREVIOUS ragged step's row state, and wherever a retired request's allocation was shorter than a current extent, the MLA rope KV-append read a -1 block entry and stored to an unmapped address (applyMLARopeAndAssignQKVKernelGeneration, Warp MMU Fault). The divergence is roster/timing dependent, which made the fault intermittent, ragged-only (uniform never keys on a bucket), and invisible under CUDA_LAUNCH_BLOCKING -- it mimicked a data race for two days of staging-side WAR fixes that could not touch it, because nothing was racing: the stale read was deterministic once the staging chain was skipped. Caught self-describing by temporary prepare-tick/staged-tick stamps: 'prepare tick 559, staged tick 558, skip reason None'. Make the code match the docstring: a step whose windows and bucket disagree derives the uniform key or falls back to eager (no_captured_shape), never a ragged replay it staged no rows for. Validation on DEP8 GSM8K with the steep synthetic cost table (the config that faulted 3 of 4 runs before): three consecutive full runs, 131/134/130 steps, ragged fully live (trim_ratio 0.53-0.55, ~500 graph replays each), accuracy 95.944/96.247/96.285 against the 96.0 gate (the first is within one sigma at n=1319). Signed-off-by: Lance Liao --- .../_torch/pyexecutor/cuda_graph_runner.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index ddc1bbe1da3a..512544d7eb0f 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -301,6 +301,20 @@ def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: from ..speculative.dspark_observability import trims_submitted_tokens if not trims_submitted_tokens(self.spec_config): return None + # Enforce what the docstring promises: None unless EVERY generation + # request carries a window. The fit publishes `agreed_ragged_bucket`, + # but windows can be absent by the time the key is derived (planner + # or peer declined after the fit, roster changed) -- and a ragged key + # over a window-less batch replays a graph whose token-major row maps + # were never staged this step, so the replay consumes the previous + # ragged step's rows: stale request mapping, stale extents, stale + # block table, and the rope KV-append walks off a retired request's + # allocation (the ragged IMA; proven by prepare-tick vs staged-tick + # stamps: tick N prepare, tick N-1 staging, no skip-guard fired). + if any( + getattr(request, "py_verify_len", None) is None + for request in batch.generation_requests): + return None # The bucket the fit agreed on, not a fresh sum over the batch. Summing # here walks generation_requests *after* _get_padded_batch has appended # padding, so it is a second, independent derivation of a quantity every From 35a1ac4c4a1c5d625e4e2a340148adbb95ba3b36 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Tue, 4 Aug 2026 17:21:44 -0700 Subject: [PATCH 080/177] [None][fix] Re-enable ragged verification on graph-less steps Reverts the TLLM_DSPARK_EAGER_RAGGED gate (4f81cc5534). The gate's justification is void: the fault it fenced off was never specific to this path -- it was a ragged graph key derived over a window-less step replaying stale row maps, fixed at the key derivation. The gate's cost was measured when it went in: 46 of the first 64 steps declined, trim_ratio 0.62 -> 0, i.e. the default configuration stopped trimming almost entirely. Steps no rank can graph were eager long before this feature (can_run_cuda_graph is num_context_requests == 0), so running the fitted windows eagerly forfeits nothing; the branch still publishes no bucket and fits nothing, so there is no captured shape to mismatch. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 44 +++++-------------- 1 file changed, 12 insertions(+), 32 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index b2eec8b185f3..eb1cf94f3013 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3516,44 +3516,24 @@ def _dspark_confidence_draft_len(self, ragged_active = bucket is not None if bucket is None: fallback_reason = "no_captured_shape" - elif os.environ.get("TLLM_DSPARK_EAGER_RAGGED"): - # OFF BY DEFAULT. This path triggers a CUDA illegal memory - # access that is not yet located. - # - # The reasoning behind it is sound and is why the code is kept - # rather than reverted: these steps are eager no matter what - # this feature does -- can_run_cuda_graph is - # `num_context_requests == 0`, which predates it by a long way - # -- so declining ragged buys no graph replay. It only gives up - # the trimming, and that was measured: 46 of the first 64 steps - # fell back, trim_ratio 0.62 -> 0. - # - # What is NOT understood is why publishing ragged metadata on a - # step that then runs eager faults at all. Nine hypotheses have - # been falsified, most usefully: setting every window to the - # FULL block, which makes the step semantically identical to - # static verification, still faults -- so it is not about - # windows, raggedness, or trimming. A barrier at this point does - # not help while CUDA_LAUNCH_BLOCKING=1 completes a full 128-step - # run, which places the race inside the forward. - # - # The env var keeps the reproduction to one line. + else: + # No rank can replay a graph this step -- such steps were + # eager long before this feature existed (can_run_cuda_graph + # is num_context_requests == 0), so declining ragged would buy + # no replay and only forfeit the trimming (measured: 46 of the + # first 64 steps, trim_ratio 0.62 -> 0). Run the windows + # eagerly, and do not fit a bucket: fitting sizes pad rows + # `_get_padded_batch` will not append on an ungraphable step. + # The fault once blamed on this path was a ragged graph KEY + # derived over a window-less step replaying stale row maps; + # with the key derivation checking windows, this path + # publishes no bucket and replays nothing. runner = self.model_engine.cuda_graph_runner runner.agreed_ragged_bucket = None runner.ragged_pad_verify_len = 0 for request, window in zip(gen_requests, ragged_lens): request.py_verify_len = int(window) ragged_active = True - else: - # No rank can replay a graph this step, so there is no captured - # shape to land on. Fitting one anyway budgets tokens for pad - # rows `_get_padded_batch` will not append -- it pads the TOTAL - # batch while this fit counts generation rows -- which is how - # `requests carry 18 generation tokens but the key's token axis - # is 24` reached two ranks, threw inside the forward, and - # deadlocked the other six behind it. - fallback_reason = "no_graph_this_step" - ragged_lens = None # Record what the step actually decided, reading the windows back off # the requests rather than off `ragged_lens`: fit_ragged_verify_lens From 74f1a82d06f5230d0b4c6f68ec408535f4511b47 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Tue, 4 Aug 2026 17:21:57 -0700 Subject: [PATCH 081/177] [None][fix] Stage KVCacheManagerV2 block-offset copies through pinned snapshots Found while hunting the ragged IMA; a real bug regardless of it. The V1 KVCacheManager stages its async block-offset H2D through a fresh pinned buffer per copy (nvbug 6293536): the persistent host buffer is rewritten whenever blocks are allocated or freed -- at scheduling time, which under the overlap scheduler can be while the previous step's copy is still queued behind an in-flight forward. Both KVCacheManagerV2 copy branches lacked that fix and read the persistent buffer directly; in V2 the exposure is worse because the buffer rows are live page-index buffers the allocator rewrites in place. Snapshot into a fresh pinned buffer per copy, matching V1. The flat branch hands its snapshot to a C++ op torch records no event for, so the caching host allocator alone cannot keep the block alive; retain (event, snapshot) pairs until the op's stream has drained past them. Adds a V2 twin of the V1 overlap-race regression test. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/kv_cache_manager_v2.py | 55 ++++++++- .../test_kv_block_offset_overlap_race.py | 115 +++++++++++++++++- 2 files changed, 167 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 5be42012a515..ec9dbae45d7c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -1661,6 +1661,54 @@ def _copy_scratch_metadata_to_device( self._device_scratch_slots_staging, ) + def _stage_host_block_offsets_for_copy(self) -> torch.Tensor: + """Snapshot the persistent host block-offset buffer into a fresh pinned + buffer to serve as the private source of an asynchronous H2D copy. + + Same bug and same fix as nvbug 6293536 in KVCacheManager + (``_stage_block_offsets_for_copy``): the persistent buffer is rewritten + whenever blocks are allocated or freed, which happens at *scheduling* + time -- and with the overlap scheduler that is while the previous + step's H2D from this buffer can still be queued behind an in-flight + forward. The queued copy then reads the next step's block layout and + the attention kernels walk freed or reassigned blocks. + + Lifetime: for the per-layer branch the consumer is a torch ``copy_``, + which records the pinned block's use with the caching host allocator, + so dropping the Python reference is safe. The flat branch hands the + snapshot to a C++ op torch knows nothing about -- its caller must keep + the snapshot referenced until the op's stream has drained past it + (see ``_retain_flat_snapshot``). + """ + src = self.host_kv_cache_block_offsets + host_block_offsets = torch.empty( + src.shape, + dtype=src.dtype, + device="cpu", + pin_memory=prefer_pinned(), + ) + host_block_offsets.copy_(src) + return host_block_offsets + + def _retain_flat_snapshot(self, snapshot: torch.Tensor) -> None: + """Keep a pinned snapshot alive until the C++ copy op has consumed it. + + The op reads the host buffer at *execution* time on ``self._stream``, + and torch records no event for it, so the caching host allocator would + otherwise hand the block to the next same-size pinned allocation while + the op is still queued. Hold (event, snapshot) pairs and drop the ones + whose event has completed -- no synchronization, bounded by the + scheduler's run-ahead depth. + """ + retained = getattr(self, "_flat_snapshots_in_flight", None) + if retained is None: + retained = self._flat_snapshots_in_flight = [] + event = torch.cuda.Event() + with torch.cuda.stream(self._stream): + event.record() + retained.append((event, snapshot)) + retained[:] = [(e, t) for e, t in retained if not e.query()] + def _copy_batch_block_offsets_per_layer( self, dst_tensor: torch.Tensor, @@ -1670,8 +1718,9 @@ def _copy_batch_block_offsets_per_layer( num_seqs: int, ) -> None: device_copy_idx = self._copy_idx_to_device(copy_idx) + host_block_offsets = self._stage_host_block_offsets_for_copy() self._device_kv_cache_block_offsets_input.copy_( - self.host_kv_cache_block_offsets, + host_block_offsets, non_blocking=True, ) scratch_begs, scratch_ends, scratch_slots = self._copy_scratch_metadata_to_device( @@ -3600,14 +3649,16 @@ def copy_batch_block_offsets( ) return + host_block_offsets = self._stage_host_block_offsets_for_copy() copy_batch_block_offsets_to_device( - self.host_kv_cache_block_offsets, + host_block_offsets, dst_tensor, copy_idx, self.index_scales, self.kv_offset, self._stream.cuda_stream, ) + self._retain_flat_snapshot(host_block_offsets) @staticmethod def _derive_reuse_salt(cache_salt: str | None) -> int | None: diff --git a/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py b/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py index ff5743430938..f905b06f38ad 100644 --- a/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py +++ b/tests/unittest/_torch/executor/test_kv_block_offset_overlap_race.py @@ -28,17 +28,25 @@ stall that delays the H2D), enqueues batch A's async H2D behind the stall, then lets the host immediately issue batch B before A's copy can drain. With the bug present, batch A's device tensor ends up holding batch B's offsets. + +``KVCacheManagerV2`` shared the bug through the persistent +``host_kv_cache_block_offsets`` it fed to the same async H2D — worse, in V2 +that buffer's rows are the live page-index buffers the C++ allocator rewrites +in place at scheduling time. The V2 twin at the bottom drives the equivalent +window against its snapshot fix (``_stage_host_block_offsets_for_copy``). """ import pytest import torch +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.bindings import DataType, LayerType from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp from tensorrt_llm.bindings.internal.batch_manager import CacheType from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2._utils import init_cuda_once # Small geometry so each sequence spans several blocks and the two batches get # visibly different physical block indices. @@ -52,6 +60,10 @@ # Long GPU stall so the host reliably wins the race and overwrites the staging # buffer before batch A's H2D drains. Generous to stay deterministic on fast GPUs. _SLEEP_CYCLES = 2_000_000_000 +# Overwrites every V2 page index during the race window. Any value distinct +# from real page indices (and from BAD_PAGE_INDEX == -1) works; the copies +# only move it as data, never index with it. +_V2_CLOBBER_PAGE_INDEX = 0x0BAD def _build_manager(): @@ -89,8 +101,13 @@ def _build_manager(): def _empty_device_offsets(mgr): + # Dim 0 mirrors production's kv_cache_block_offsets allocation (see + # TrtllmAttentionMetadata): V2 sizes it by num_attention_op_pools, which + # equals num_pools unless SWA scratch reuse widens it to one virtual pool + # per layer. V1 has no such attribute. + num_pools = getattr(mgr, "num_attention_op_pools", mgr.num_pools) return torch.zeros( - mgr.num_pools, 2 * _BATCH, 2, mgr.max_blocks_per_seq, dtype=torch.int32, device="cuda" + num_pools, 2 * _BATCH, 2, mgr.max_blocks_per_seq, dtype=torch.int32, device="cuda" ) @@ -186,3 +203,99 @@ def test_copy_batch_block_offsets_survives_overlap_overwrite(): ) # And batch B's copy is independently correct. assert torch.equal(dst_b[:, :_BATCH], ref_b) + + +def _build_manager_v2(enable_swa_scratch_reuse): + """Plain full-attention V2 manager with the same small geometry as V1's. + + ``enable_swa_scratch_reuse=True`` routes ``copy_batch_block_offsets`` + through ``_copy_batch_block_offsets_per_layer`` — the branch production + DSV4 serving takes and the one the snapshot fix primarily targets. + ``False`` covers the flat ``copy_batch_block_offsets_to_device`` branch, + which stages through the same helper. + """ + init_cuda_once() + kv_cache_config = KvCacheConfig( + max_tokens=4096, + enable_block_reuse=False, + enable_swa_scratch_reuse=enable_swa_scratch_reuse, + ) + return KVCacheManagerV2( + kv_cache_config, + CacheType.SELF, + num_layers=_NUM_LAYERS, + num_kv_heads=_NUM_KV_HEADS, + head_dim=_HEAD_DIM, + tokens_per_block=_TOKENS_PER_BLOCK, + max_seq_len=_MAX_SEQ_LEN, + max_batch_size=2 * _BATCH, + mapping=Mapping(world_size=1, tp_size=1, pp_size=1), + dtype=DataType.HALF, + vocab_size=32000, + # Keep the pool solver's typical step at test scale instead of the + # constructor's 8192-token default. + max_num_tokens=_MAX_SEQ_LEN, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") +@pytest.mark.parametrize("swa_scratch_reuse", [True, False], ids=["per_layer_scratch", "flat"]) +def test_kv_cache_manager_v2_copy_batch_block_offsets_survives_host_rewrite(swa_scratch_reuse): + """V2 twin of the overlap-overwrite test above. + + Batch A's copy is enqueued behind a stream stall, then the persistent + ``host_kv_cache_block_offsets`` is rewritten in place before the copy can + drain. The rewrite is direct rather than a free-A/allocate-B cycle: the + C++ allocator mutates exactly these rows in place (through the + memoryviews registered by ``_create_kv_cache`` / + ``_restore_page_index_bufs``), but a re-allocation may hand the freed + pages back in the same order and reproduce identical rows, masking the + race. Batch A's device tensor must keep batch A's offsets. + """ + mgr = _build_manager_v2(swa_scratch_reuse) + assert mgr.enable_swa_scratch_reuse == swa_scratch_reuse + requests = None + saved_host_offsets = None + try: + ids = list(range(1, 1 + _BATCH)) + toks = [_TOKENS_PER_SEQ] * _BATCH + requests = mgr.add_dummy_requests(request_ids=ids, token_nums=toks, prepare_resource=True) + assert requests is not None, "test setup invalid: dummy request allocation failed" + + # Also warms the (maybe-compiled) per-layer copy op outside the race + # window, so the stalled call below enqueues without host-side + # compilation pauses that would let the H2D drain early. + ref = _reference_offsets(mgr, ids) + saved_host_offsets = mgr.host_kv_cache_block_offsets.clone() + + dst = _empty_device_offsets(mgr) + torch.cuda.synchronize() + + # Drive the scheduling window: stall the stream, enqueue batch A's + # async H2D behind the stall, then rewrite the persistent host rows + # before the copy can drain. + torch.cuda._sleep(_SLEEP_CYCLES) + mgr.copy_batch_block_offsets(dst, ids, 1, _BATCH, _BATCH) + mgr.host_kv_cache_block_offsets.fill_(_V2_CLOBBER_PAGE_INDEX) + torch.cuda.synchronize() + + # Setup validity: a copy issued against the clobbered buffer must + # produce something else, or the rewrite wasn't observable. + clobbered = _reference_offsets(mgr, ids) + assert not torch.equal(clobbered, ref), ( + "test setup invalid: host-buffer rewrite did not change the copied offsets" + ) + + assert torch.equal(dst[:, :_BATCH], ref), ( + "nvbug 6293536 (V2): batch A's async H2D read the rewritten " + "persistent host_kv_cache_block_offsets instead of a private " + "snapshot (_stage_host_block_offsets_for_copy)" + ) + finally: + if saved_host_offsets is not None: + # The rows are live page-index buffers; put them back before the + # allocator touches them again in free/shutdown. + mgr.host_kv_cache_block_offsets.copy_(saved_host_offsets) + for req in requests or []: + mgr.free_resources(req) + mgr.shutdown() From c7e6d78ac20fb9dd0c86726955b27b1ec2810747 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Tue, 4 Aug 2026 17:22:12 -0700 Subject: [PATCH 082/177] [None][fix] Harden DSA metadata against overlap WAR and layout drift Three latent hazards found (and individually reproduced) while hunting the ragged IMA; none was its root cause, all are real: 1. Every buffer DSA metadata stages through pinned memory (mla-rope indptrs, indexer offsets, ragged row maps, plus everything the base class stages inside prepare()) is rewritten in place each step and copied device-ward with non_blocking=True; under the overlap scheduler the CPU rewrites them while the previous step's copies can still be queued behind an in-flight forward. One coarse guard in prepare(): synchronize on an event recorded after the previous prepare's copies were enqueued before any pinned buffer is touched again. Steady-state cost is zero. 2. DeepseekV4TrtllmAttentionMetadata overrode prepare() and called straight through to the grandparent, silently bypassing that guard. Its override is now _prepare_impl, which the guarded base prepare() wraps; subclasses must override _prepare_impl, never prepare(). 3. The expanded row buffers are backed by a name-keyed pool shared between the base metadata and every CUDA-graph clone, but each instance sized its view from its own mutable max_draft_tokens -- under ragged verification that moves per step, so instances could view the same backing bytes with different row strides, and the re-create-on-change branches relocated buffers whose addresses captured graphs had baked in. Shape everything from one static ceiling (the config max_total_draft_tokens, recorded on the first update_spec_dec_param and inherited by clones) and make the re-create ratchet grow-only; per-step max_draft_tokens keeps its role in runtime logic but no longer shapes storage. Signed-off-by: Lance Liao --- .../sparse/deepseek_v4/deepseek_v4.py | 8 +- .../_torch/attention_backend/sparse/dsa.py | 108 +++++++++++++++--- 2 files changed, 101 insertions(+), 15 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index 86aee79dadaf..b7b81df73320 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -689,7 +689,13 @@ def _init_cache_buffer_data_pointers(self): if is_compress_layer(extend_compress_ratios[layer_idx]) } - def prepare(self): + def _prepare_impl(self): + # Overrides DSAtrtllmAttentionMetadata._prepare_impl, NOT prepare(): + # the base prepare() wraps this whole body in the write-after-read + # guard for the pinned staging rewritten below (cached_token_lens_cpu, + # cu_seq_lens, gen_new_tokens_per_seq_host, and everything the called + # helpers stage). Overriding prepare() directly would silently opt this + # class out of that guard. assert self.kv_cache_manager is not None assert self.request_ids is not None diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 5ee3d716cb77..0d5875d753e0 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -832,6 +832,38 @@ def __post_init__(self): self.create_buffers_for_indexer(capture_graph=capture_graph) def prepare(self): + # One write-after-read guard over EVERYTHING prepare stages through + # pinned memory -- the base class's kv_lens and block-table copies as + # much as the ragged row maps. Every one of these buffers is rewritten + # here each step and copied device-ward with non_blocking=True; under + # the overlap scheduler the previous step's copies can still be QUEUED + # behind an in-flight forward when this rewrite starts, and the device + # then copies this step's values into last step's consumers. The GPU + # coredump for the resulting fault names + # applyMLARopeAndAssignQKVKernelGeneration (Warp MMU Fault) storing + # through KVBlockArray -- garbage block offsets, which trace back + # through the token-major gather to the BASE block-table staging, not + # only the ragged row set. Hence one coarse guard at the boundary + # rather than per-buffer events: synchronize on the event recorded + # after the previous prepare's copies were enqueued, which by stream + # order means they have executed before any pinned buffer is touched + # again. Steady-state cost is nothing (the event completed long ago); + # it blocks only when the CPU has genuinely outrun the device, which is + # exactly the window that corrupted the copies. + if not torch.cuda.is_current_stream_capturing(): + evt = getattr(self, "_prepare_stage_evt", None) + if evt is None: + self._prepare_stage_evt = torch.cuda.Event() + else: + evt.synchronize() + try: + self._prepare_impl() + finally: + evt = getattr(self, "_prepare_stage_evt", None) + if evt is not None and not torch.cuda.is_current_stream_capturing(): + evt.record() + + def _prepare_impl(self): super().prepare() self._invalidate_pool_view_cache() # Cross-layer indexer sharing is per-step state; clear it so a "shared" @@ -1330,7 +1362,7 @@ def create_buffers_for_indexer(self, capture_graph=False): # path does not, so skip the allocation when use_cute_dsl_topk. if not self.use_cute_dsl_topk: max_gen_tokens = self.max_num_sequences * ( - 1 + self.max_draft_tokens) + 1 + self._draft_sizing_cap) self.heuristic_scratch_values = self.get_empty( self.cuda_graph_buffers, (max_gen_tokens, self.num_sparse_topk), @@ -1359,6 +1391,25 @@ def create_buffers_for_indexer(self, capture_graph=False): # Create expanded buffers for MTP support self.create_expanded_buffers(capture_graph=capture_graph) + @property + def _draft_sizing_cap(self) -> int: + """Static draft-token ceiling for buffer SIZING (not per-step logic). + + Every expanded buffer must be shaped from one process-wide ceiling, + never from the per-step ``max_draft_tokens``: the pooled backing + memory is shared between the base metadata and every CUDA-graph + clone, and two instances viewing the same bytes with different row + counts interleave their rows differently -- a prepare() on one + layout silently corrupts every captured graph reading the other + (the DSpark ragged IMA). ``update_spec_dec_param`` records the + config-wide ``max_total_draft_tokens`` here; before the first call + the mutable value is all there is. + """ + cap = getattr(self, "_draft_alloc_cap", None) + if cap is None: + return self.max_draft_tokens + return max(cap, self.max_draft_tokens) + def _create_kv_lens_2d_buffer(self, capture_graph=False): """Pre-allocated buffer for the DeepGEMM 2D context_lens API. @@ -1368,7 +1419,7 @@ def _create_kv_lens_2d_buffer(self, capture_graph=False): """ self.kv_lens_cuda_2d = self.get_empty( self.cuda_graph_buffers, - (self.max_num_sequences, 1 + self.max_draft_tokens), + (self.max_num_sequences, 1 + self._draft_sizing_cap), cache_name="kv_lens_cuda_2d", dtype=torch.int32, capture_graph=capture_graph, @@ -1379,7 +1430,7 @@ def create_expanded_buffers(self, capture_graph=False): """Create expanded KV-length and block-table buffers for speculative decoding.""" self.kv_lens_expanded_cuda = self.get_empty( self.cuda_graph_buffers, - (self.max_num_sequences * (1 + self.max_draft_tokens), ), + (self.max_num_sequences * (1 + self._draft_sizing_cap), ), cache_name="kv_lens_expanded_cuda", dtype=torch.int32, capture_graph=capture_graph, @@ -1400,7 +1451,7 @@ def create_expanded_buffers(self, capture_graph=False): # compression divisor twice. self.row_kv_lens_cuda = self.get_empty( self.cuda_graph_buffers, - (self.max_num_sequences * (1 + self.max_draft_tokens), ), + (self.max_num_sequences * (1 + self._draft_sizing_cap), ), cache_name="row_kv_lens_cuda", dtype=torch.int32, capture_graph=capture_graph, @@ -1416,7 +1467,7 @@ def create_expanded_buffers(self, capture_graph=False): # after prepare(); the kv lengths do. self.row_kv_correction_cuda = self.get_empty( self.cuda_graph_buffers, - (self.max_num_sequences * (1 + self.max_draft_tokens), ), + (self.max_num_sequences * (1 + self._draft_sizing_cap), ), cache_name="row_kv_correction_cuda", dtype=torch.int32, capture_graph=capture_graph, @@ -1432,7 +1483,7 @@ def create_expanded_buffers(self, capture_graph=False): # before the first replay. self.row_req_idx_cuda = self.get_empty( self.cuda_graph_buffers, - (self.max_num_sequences * (1 + self.max_draft_tokens), ), + (self.max_num_sequences * (1 + self._draft_sizing_cap), ), cache_name="row_req_idx_cuda", dtype=torch.long, capture_graph=capture_graph, @@ -1474,7 +1525,7 @@ def create_expanded_buffers(self, capture_graph=False): # `cache_name`, so a later reallocation would move addresses already # baked into captured graphs. `2 + max_draft_tokens` per sequence covers # one context row plus the widest generation window. - attn_rows_cap = self.max_num_sequences * (2 + self.max_draft_tokens) + attn_rows_cap = self.max_num_sequences * (2 + self._draft_sizing_cap) self._attn_num_rows = 0 self.attn_row_kv_lens_cuda = self.get_empty( self.cuda_graph_buffers, @@ -1561,7 +1612,7 @@ def create_expanded_buffers(self, capture_graph=False): self.block_table_expanded = self.get_empty( self.cuda_graph_buffers, [ - self.max_num_sequences * (1 + self.max_draft_tokens), + self.max_num_sequences * (1 + self._draft_sizing_cap), self.kv_cache_manager.max_blocks_per_seq ], cache_name="block_table_expanded", @@ -1608,18 +1659,38 @@ def update_spec_dec_param( spec_tree_manager, num_contexts=num_contexts) self.max_draft_tokens = max_draft_len + # The config-wide ceiling, identical every call; sizing must use this + # so base and every clone share one buffer layout (see + # _draft_sizing_cap). + self._draft_alloc_cap = max( + int(max_total_draft_tokens), + getattr(self, "_draft_alloc_cap", 0) or 0) capture_graph = self.is_cuda_graph - if self.kv_lens_cuda_2d.shape[1] != 1 + self.max_draft_tokens: + # GROW-ONLY ratchet, never an equality check. Under DSpark ragged + # verification max_draft_len varies per step with the planner's + # windows; the old `!=` conditions re-created every expanded buffer + # whenever it moved, while already-captured graphs kept the previous + # allocations' raw addresses baked in. The caching allocator then + # reused the freed memory for the per-step indptr tensors, and every + # replay read indptr ramps where its row->request map and per-row KV + # extents used to live -- the ragged IMA. (Frozen-scene forensics: + # the retired attn_row buffers' addresses held int64 cumsum ramps.) + # Growing above the high-water mark before any graph exists is fine; + # after that the sizes plateau at the config ceiling and the + # addresses never move again. Consumers slice [:rows]/[:width], so a + # wider-than-this-step buffer is just capacity. + ratchet = max(self._draft_sizing_cap, + getattr(self, "_expanded_alloc_draft_cap", -1)) + if ratchet > getattr(self, "_expanded_alloc_draft_cap", -1): + self._expanded_alloc_draft_cap = ratchet self._create_kv_lens_2d_buffer(capture_graph=capture_graph) - init_shape = self.kv_lens_expanded_host.shape[0] - if self.max_num_sequences * (1 + self.max_draft_tokens) != init_shape: self.create_expanded_buffers(capture_graph=capture_graph) # Resize heuristic scratch buffer for new max_draft_tokens. # Skip when use_cute_dsl_topk (GVR path never consumes it), matching # the allocation guard in create_buffers_for_indexer. if self.enable_heuristic_topk and not self.use_cute_dsl_topk: max_gen_tokens = self.max_num_sequences * ( - 1 + self.max_draft_tokens) + 1 + self._draft_sizing_cap) self.heuristic_scratch_values = self.get_empty( self.cuda_graph_buffers, (max_gen_tokens, self.num_sparse_topk), @@ -1955,16 +2026,21 @@ def _prepare_ragged_row_kv_lens(self, kv_lens: torch.Tensor) -> None: num_tokens = len(rows) self._ragged_num_rows = num_tokens + # The write-after-read hazard on these pinned buffers (rewritten each + # ragged step, copied with non_blocking=True under the overlap + # scheduler's run-ahead) is fenced wholesale by the event guard in + # prepare(): it waits for the previous prepare's copies to execute + # before any pinned buffer is touched again. self.row_kv_lens_host[:num_tokens].copy_( torch.tensor(rows, dtype=torch.int32)) self.row_kv_lens_cuda[:num_tokens].copy_( self.row_kv_lens_host[:num_tokens], non_blocking=True) self.row_kv_correction_host[:num_tokens].copy_( torch.tensor(corrections, dtype=torch.int32)) - self.row_kv_correction_cuda[:num_tokens].copy_( - self.row_kv_correction_host[:num_tokens], non_blocking=True) self.row_req_idx_host[:num_tokens].copy_( torch.tensor(req_idx, dtype=torch.long)) + self.row_kv_correction_cuda[:num_tokens].copy_( + self.row_kv_correction_host[:num_tokens], non_blocking=True) self.row_req_idx_cuda[:num_tokens].copy_( self.row_req_idx_host[:num_tokens], non_blocking=True) @@ -2041,6 +2117,10 @@ def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: # the windows again: it produced exactly these per-row extents, # corrections and request indices one call earlier. Rebuilding them # would double the only O(num_gen_tokens) host work on the step. + # Same WAR guard as _prepare_ragged_row_kv_lens above, same reason: + # attn_row_req_idx_host holds whole-batch request indices, and a + # rewrite racing last step's queued H2D turns into an out-of-bounds + # gather in refresh_token_major_gen_rows. self.attn_row_kv_lens_host[nc:rows].copy_( self.row_kv_lens_host[:ngt]) self.attn_row_kv_correction_host[nc:rows].copy_( From 5c38441a00e529c21b3f851054b978d45c9eb40c Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Tue, 4 Aug 2026 17:22:29 -0700 Subject: [PATCH 083/177] [None][fix] Bound the overlap previous-batch staging inputs Also found while hunting the ragged IMA; latent regardless of it. - The spec sampler's slot-indexed stores (new_tokens_lens and siblings) were torch.empty at init and only the slots sampled in a step are rewritten, while the overlap scheduler's previous-batch gather in model_engine indexes them by slots from a previous roster: a slot read before its first write returns arbitrary bytes, and one such value feeds the captured kv_lens correction where it walks the KV append past the request's allocation. int_tensor stores are zeros now. - The ragged offsets gather is clamped to the physically possible range (accepted minus this request's window, bounded by one verify window): a stale read is then at worst a small KV-length error, never an out-of-bounds address. - Three _pinned_host staging sites (accept_caps, ragged_verify_lens, the ragged gather index pair) get per-key events: the host waits for the previous step's queued H2D from a staging view to execute before rewriting it. These sites live outside metadata.prepare(), so the DSA prepare guard does not cover them. Signed-off-by: Lance Liao --- .../_torch/pyexecutor/model_engine.py | 51 +++++++++++++++++-- .../_torch/pyexecutor/sampler/sampler.py | 10 +++- 2 files changed, 56 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 703e4a3754f2..a5ff8653c328 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4140,6 +4140,17 @@ def _pinned_host(self, key: str, values, dtype) -> torch.Tensor: cache = getattr(self, "_pinned_host_cache", None) if cache is None: cache = self._pinned_host_cache = {} + self._pinned_host_events = {} + # WAR guard, the counterpart of the use-after-free this helper fixed: + # a persistent buffer can be REWRITTEN while its previous non_blocking + # H2D is still queued behind an in-flight forward, so the device copies + # this step's values into last step's consumer. Callers record the + # key's event via _pinned_host_record after enqueuing the copy; waiting + # here closes the window, and in the steady state the event completed + # long ago so this costs nothing. + evt = self._pinned_host_events.get(key) + if evt is not None: + evt.synchronize() buf = cache.get(key) if buf is None or buf.numel() < len(values) or buf.dtype != dtype: buf = torch.empty(max(len(values), 1), @@ -4151,6 +4162,21 @@ def _pinned_host(self, key: str, values, dtype) -> torch.Tensor: view.copy_(torch.tensor(values, dtype=dtype)) return view + def _pinned_host_record(self, *keys: str) -> None: + """Mark the H2D copies from these staging buffers as enqueued. + + Call after the non_blocking copy that reads a `_pinned_host` view. + Skipped during graph capture (the staging copies are eager on every + step, including graph steps; only the one-time capture pass is not). + """ + if torch.cuda.is_current_stream_capturing(): + return + for key in keys: + evt = self._pinned_host_events.get(key) + if evt is None: + evt = self._pinned_host_events[key] = torch.cuda.Event() + evt.record() + def _attach_accept_caps(self, spec_metadata, generation_requests) -> None: """Publish ``cap-accept`` ceilings, without touching the layout. @@ -4198,6 +4224,7 @@ def _attach_accept_caps(self, spec_metadata, generation_requests) -> None: [1 + int(cap) for cap in caps], torch.int32), non_blocking=True) + self._pinned_host_record("accept_caps") spec_metadata.accept_caps = caps_view spec_metadata.accept_cap_trim = self.accept_cap_trim_cuda @@ -4233,6 +4260,7 @@ def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, lens_view.copy_(self._pinned_host("ragged_verify_lens", token_lens, torch.int32), non_blocking=True) + self._pinned_host_record("ragged_verify_lens") indptr_view = self.ragged_qo_indptr_cuda[:n + 1] indptr_view.copy_(build_qo_indptr(lens_view), non_blocking=True) spec_metadata.verify_lens = lens_view @@ -4254,10 +4282,13 @@ def _ragged_gather_indices( verifies the same number of positions. """ rows, cols = ragged_gather_index_lists(slots, counts) - return (self._pinned_host("gather_rows", rows, - torch.long).to('cuda', non_blocking=True), - self._pinned_host("gather_cols", cols, - torch.long).to('cuda', non_blocking=True)) + rows_dev = self._pinned_host("gather_rows", rows, + torch.long).to('cuda', non_blocking=True) + self._pinned_host_record("gather_rows") + cols_dev = self._pinned_host("gather_cols", cols, + torch.long).to('cuda', non_blocking=True) + self._pinned_host_record("gather_cols") + return (rows_dev, cols_dev) def _update_target_input_tensors( self, @@ -5553,6 +5584,18 @@ def previous_seq_slots_device(): previous_kv_len_offsets = ( new_tokens_lens_device[previous_slots] - tokens_per_previous_request) + # The gather above reads the sampler's slot-indexed store; + # a slot the sampler has not (yet) written for the previous + # step yields a stale or uninitialized count, and the + # captured kv_lens correction then walks the KV append past + # the request's allocated blocks (the ragged IMA). The + # physically possible range is tight -- accepted tokens + # minus this request's window -- so clamp to it; a clamped + # value is at worst one window of KV length, never an OOB + # address. + previous_kv_len_offsets = previous_kv_len_offsets.clamp_( + min=-int(self.runtime_draft_len + 1), + max=int(self.runtime_draft_len + 1)) else: kv_len_offsets_device = (new_tokens_lens_device - runtime_tokens_per_gen_step) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 3f649ca0db83..3d1ce5eec0df 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -800,7 +800,15 @@ def add_token( def int_tensor(shape: tuple[int, ...], device: str = "cuda") -> torch.Tensor: - return torch.empty(shape, dtype=torch.int, device=device) + # zeros, not empty: several of these buffers are slot-indexed stores that + # are only written for the slots sampled in a given step, while consumers + # (the overlap scheduler's previous-batch gather in model_engine) index + # them by slots from a *previous* roster. A slot that has not been written + # yet must read as 0 -- with empty it reads arbitrary bytes, and one such + # read (previous_kv_lens_offsets) feeds the captured kv_lens correction, + # where a garbage value walks the KV append past the request's allocated + # blocks (the DSpark ragged IMA; see docs/dspark_ragged_ima_handoff.md). + return torch.zeros(shape, dtype=torch.int, device=device) @dataclass(kw_only=True, frozen=True) From 648d6019526fe74362f22898d5c2007669f12985 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Tue, 4 Aug 2026 17:22:43 -0700 Subject: [PATCH 084/177] [None][chore] Let a diagnostic env override the hang-detection timeout Debugging a device-side fault with CUDA_DEVICE_WAITS_ON_EXCEPTION needs the frozen GPU state to outlive the watchdog, whose hard-kill would otherwise destroy the state the debugger is attached to. Honor TLLM_HANG_DETECTION_TIMEOUT when no explicit timeout is passed. Signed-off-by: Lance Liao --- tensorrt_llm/_torch/pyexecutor/hang_detector.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/hang_detector.py b/tensorrt_llm/_torch/pyexecutor/hang_detector.py index 2ae692ed5902..e13768a9b64e 100644 --- a/tensorrt_llm/_torch/pyexecutor/hang_detector.py +++ b/tensorrt_llm/_torch/pyexecutor/hang_detector.py @@ -91,7 +91,13 @@ class HangDetector: def __init__( self, timeout: Optional[int] = None, on_detected: Optional[Callable[[], None]] = None ): - self.timeout = timeout if timeout is not None else 300 + # TLLM_HANG_DETECTION_TIMEOUT: diagnostic override. Debugging a device + # fault with CUDA_DEVICE_WAITS_ON_EXCEPTION needs the frozen GPU state + # to outlive this watchdog, whose hard-kill would otherwise destroy the + # very state the debugger is attached to. + if timeout is None: + timeout = int(os.environ.get("TLLM_HANG_DETECTION_TIMEOUT", "300")) + self.timeout = timeout assert self.timeout > 0, "timeout must be greater than 0" self.on_detected = on_detected or (lambda: None) self.task = None From 226b1ffc611d7ef31d51bcf14e36c339b956daf3 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Tue, 4 Aug 2026 17:22:44 -0700 Subject: [PATCH 085/177] [None][doc] Close out the ragged-verify IMA investigation Full causal narrative: symptom, the fifteen falsified hypotheses, the forensic tooling that ended it (GPU coredumps, zero-sync device tripwires, CUDA_DEVICE_WAITS_ON_EXCEPTION frozen-scene attaches, prepare/staged tick stamps), the root cause (a ragged graph key derived over a window-less step replaying stale row maps), the fix, and the validation evidence. The investigation scaffolding was stripped from the tree after validation; this document and tmp/autopsy_recipe.md are the durable record. Signed-off-by: Lance Liao --- docs/dspark_ragged_ima_handoff.md | 240 +++++++++++++++++++++++++++--- 1 file changed, 222 insertions(+), 18 deletions(-) diff --git a/docs/dspark_ragged_ima_handoff.md b/docs/dspark_ragged_ima_handoff.md index a20022c35412..9f8a2dc65a04 100644 --- a/docs/dspark_ragged_ima_handoff.md +++ b/docs/dspark_ragged_ima_handoff.md @@ -116,6 +116,211 @@ Corollary: `commit 4f81cc5534` ("Gate the eager-ragged path off until its IMA is understood") does **not** make the branch safe. Its stated justification is wrong. Keep the gate or drop it on its own merits, but do not rely on it. +### Root cause (best-evidenced theory; coarse-guard fix in validation 2026-08-04) + +GPU coredump-on-exception (`COREDUMP=1` in the runner) ended the guessing that +§6 used to recommend compute-sanitizer for. Three independent runs produced +corefiles naming the same kernel: +`applyMLARopeAndAssignQKVKernelGeneration`, +Warp MMU Fault, and disassembly at the trigger PC shows the faulting +instruction is the 128-bit **KV-cache store** (`STG.E.128`) whose address is +computed from a block offset — i.e. the kernel was handed a block table / KV +extent pair that points outside mapped memory. + +The mechanism is a **write-after-read race on persistent pinned staging**: the +base class already fixed exactly this bug for `kv_cache_block_offsets` (nvbug +6293536 — see the long comment in `resource_manager.py +copy_batch_block_offsets` and `_stage_block_offsets_for_copy`): under the +overlap scheduler the CPU runs an iteration ahead, so an in-place refill of a +pinned host buffer can clobber the source of the previous iteration's +still-queued `non_blocking=True` H2D. The base class's fix was to snapshot into +a fresh pinned buffer per copy. **`dsa.py` re-introduces the unfixed pattern +wholesale**: `prepare_for_mla_rope_append` (four indptr buffers — direct +inputs of the faulting rope kernel), `prepare_for_indexer_k_cache` +(`host_indexer_k_cache_block_offsets`), `prepare_for_indices_conversion` +(`host_req_idx_per_token`), `kv_lens_expanded_host`, and the ragged row maps. +Ragged is what turns the latent race into an MMU fault: uniform decode's +next-step values differ by +1 token, but ragged verify's per-row extents and +row→request mapping change shape step to step, so a one-step-stale mix of +extents and block table walks off the allocated blocks. + +Two narrower fixes failed before the mechanism was fully mapped: per-buffer +events on the ragged row maps only (fault persisted — the indptr staging was +still exposed), and skipping the copies during graph capture (a no-op: +`_forward_step` is captured on **pre-built inputs**, `prepare()` never runs +inside the capture region — verified at `model_engine.py:7322`, the captured +body is `_preprocess_inputs` + `model_forward` only). + +The fix has two layers, because the racy buffers have two different *writers*: + +1. **Buffers rewritten inside `prepare()`** (the DSA indptr/indexer/row-map + staging): one coarse guard in `DSAtrtllmAttentionMetadata.prepare` — + synchronize on an event recorded at the end of the previous `prepare()` + before any pinned buffer is rewritten, record it after this step's copies + are enqueued. Steady-state cost is zero. Trap discovered while validating: + `DeepseekV4TrtllmAttentionMetadata` overrode `prepare()` and skipped + straight to the grandparent, silently bypassing the guard — the V4 override + is now `_prepare_impl` so the wrapper covers it. Any future subclass must + override `_prepare_impl`, never `prepare`. +2. **`host_kv_cache_block_offsets` under KVCacheManagerV2** (the run config + uses `use_kv_cache_manager_v2=True`): its rows are live page-index buffers + the allocator rewrites at *scheduling* time — outside any prepare-level + fence — while the previous step's H2D from the same buffer can still be + queued behind an in-flight forward. Neither of V2's two copy branches had + V1's nvbug-6293536 snapshot fix. A prepare-scope guard provably does not + cover this: a run with layer-1 alone (fix5) still faulted at step ≥64. + Fixed by `_stage_host_block_offsets_for_copy` in `kv_cache_manager_v2.py` + — snapshot to a fresh pinned buffer per copy, V1 parity. This is the layer + that matches the coredump: the block table is what the faulting store's + address is computed from, and ragged trimming is what makes roster churn + (block free/reassign at scheduling time) frequent enough to hit the window. + +Validation note: default-gated runs now die at step ~67 on the runner's own +activity gate (`trim_ratio` ≈ −0.02) because the `4f81cc5534` gate declines +eager-ragged on no-graph steps — a config artifact, not the IMA. Use +`EAGER_RAGGED=1` runs for validation; they exercise BOTH eager-ragged and +graph-ragged (bucketed) steps in one run. + +### 2026-08-04 ~07:00 — the tripwire verdict: VALUE corruption, confirmed live + +The two staging-fix layers above were *necessary but not sufficient*: with both +in, 4 identical runs gave 1 pass / 3 faults. The remaining hypothesis space was +split by a zero-sync device tripwire (`TLLM_DSPARK_VALIDATE_ROWS`, +`torch._assert_async` on the exact (extent, block id) pair the rope KV-append +computes its store address from — placed after the prepare-time gather and +inside the capture body after the extent rebuild) run under +`CUDA_DEVICE_WAITS_ON_EXCEPTION`: + +**The tripwire FIRED** (`CUDA_EXCEPTION_12, Warp Assert` in +`_assert_async_cuda_kernel`, frozen live on the faulting rank) — the value +pair is out of bounds *before* the rope kernel consumes it. This kills the +VMM/page-mapping family (suspend/resume, free-vs-replay unmap) outright and +reduces the problem to: **who writes a bad extent or block id, between two +validated points**. The launch-blocking algebra had already excluded +deterministic value bugs and graph-internal races (a replay is ONE launch — +CUDA_LAUNCH_BLOCKING does not serialize inside it, yet 79 bucketed steps +passed under it), so the writer is a host-side actor racing in-flight device +work on ragged-only row state. + +Forensics protocol that finally worked (traps everywhere else): +- Attach from the host via `enroot exec cuda-gdb -p ` — a sibling + pyxis container cannot ptrace, the host has no cuda-gdb, and gdb through + `/proc//root` cannot map the target's libraries. +- NEVER wrap an attached cuda-gdb in `timeout` — killing it mid-handshake + kills the frozen inferior (lost one scene to this). +- Do not `p *(@parameter ...)` on the assert kernel — cuda-gdb 13.2 crashes + (out_of_range in print_cuda_exception_string) and the crash cascade tears + the run down. +- The faulting rank is named by the peers' `timed out waiting for completion + flag from rank R` printfs, which flush minutes after the freeze. + +The tripwire now also STASHES (site, rows, bad row, extent, block id, bound, +capacity, ok, req idx, kv_len) into a plain-cudaMalloc int64[10] whose address +is logged at allocation (`DSPARK VALIDATE debug buffer @ 0x…`) — after a +freeze, `x/10gd ` in the attached debugger reads the verdict directly. +site 1 = prepare-time gather produced the bad pair (bad inputs from the host +side); site 2 = it appeared only at the in-graph rebuild (a concurrent writer +touched row state between prepare and replay). + +### 2026-08-04 ~09:15 — ROOT CAUSE, read off two frozen scenes + +Two independent stash readouts, identical signature: + +| scene | bucket | row | ext | block id | req | kv_lens_cuda | drift vs allocation | +|---|---|---|---|---|---|---|---| +| i / rank6 | 256 | 60 | 1221 | −1 | 10 | 1226 | ≈ +70 | +| j / rank2 | 96 | 36 | 926 | −1 | 6 | 931 | ≈ +35 | + +Both **site 2** with a clean site 1 in the same step. Since both sites use the +same formula on the same table, site1-clean + site2-trip forces one +conclusion: `kv_lens_cuda[slot]` moved between the prepare H2D reset and the +in-graph rebuild — and the only mutation in that interval is the captured +overlap correction `kv_lens_cuda += previous_kv_lens_offsets_cuda`. Its legal +values are bounded by one verify window (≤7); the observed drift is tens of +tokens. + +The unbounded source: `previous_kv_lens_offsets` is built as +`new_tokens_lens_device[previous_slots] − tokens_per_previous_request` +(`model_engine.py`, ragged branch), and `new_tokens_lens` is the **spec +sampler's slot-indexed store, allocated with `torch.empty`** and rewritten +each step only for the slots actually sampled. A slot the sampler never wrote +reads as arbitrary bytes; one garbage offset pushes the device KV length past +the host's block allocation, the in-graph extent rebuild follows it, and the +rope KV-append store walks into a −1 block entry → unmapped address → Warp +MMU Fault. + +Every historical constraint checks out: **uniform never faults** because a +stable roster rewrites every slot every step (even a stale read is bounded); +**ragged faults** because trimming churns completions and the fault window +(step ~64–90) is exactly when batch concurrency reaches new peaks, i.e. when +slots are read before their first write; **CUDA_LAUNCH_BLOCKING passes** +because it serializes the sampler's store writes ahead of the next step's +staging gather; the staging-fix layers never touched it because the offsets +gather is device-side, not a pinned-host copy. + +Fix (`sampler.py int_tensor` + `model_engine.py` ragged offsets): +1. Slot-indexed sampler stores are **zeros, not empty** — an unwritten slot + reads 0. +2. The ragged offsets gather is **clamped to ±(window)** — a stale value is + then at worst one window of KV length, never an out-of-bounds address. + +Note the residual: a *stale-but-clamped* offset is still a small KV-length +error on roster-churn steps (silent, bounded). The clean long-term fix is an +explicit ordering guarantee (event) between the sampler's store writes and +the next step's previous-batch gather; not attempted yet — measure first. + +### 2026-08-04 ~16:30 — THE ROOT CAUSE, self-described by the stamp tripwire + +The kv-lens-drift reading above was a *symptom* misread (the apparent drift +was a stale row→request map gathering the new roster's kv_lens). The final +mechanism, caught by a prepare-tick/staged-tick stamp pair raising on the +step itself: + +``` +DSpark ragged staleness: prepare tick 559, staged tick 558, skip reason None +``` + +`skip reason None` = the token-major staging chain was NEVER INVOKED that +step: `_ragged_token_lens` returned None (a generation request without a +window), so `_attach_ragged_verify_layout` set `ragged_verify_lens = None` +and prepare skipped every row-map/table refresh — while the graph key still +carried the fit's `agreed_ragged_bucket`, so the step REPLAYED a +ragged-bucket graph over the PREVIOUS ragged step's row state: stale +row→request mapping, stale extents, stale per-row block tables. Wherever the +retired roster's allocation was shorter than a current extent, the rope +KV-append read a −1 block entry and stored to an unmapped address. + +`_ragged_verify_bucket`'s docstring always promised "None unless *every* +generation request carries a window"; the code returned the published bucket +without checking. **Fix: make the code match the docstring** — a step whose +windows and bucket disagree derives the uniform key or falls back to eager +(`no_captured_shape`), never a ragged replay it staged no rows for. + +Why every earlier observation fit a race: the divergence between the fit's +bucket and the requests' windows is timing/roster dependent (planner or peer +declining after the fit, completions between fit and key derivation), so the +fault was intermittent, followed the first captured buckets by a few dozen +steps, never touched uniform (no bucket key), and vanished under +CUDA_LAUNCH_BLOCKING's pacing. Every WAR/staging fix left it untouched +because nothing was racing — the replay was reading last step's rows by +design once the staging chain was skipped. + +First validation (fix11_a): **131 steps, zero IMA/hang/tripwire**, ragged +fully active (`steps_ragged` 38, `trim_ratio` 0.5375, 527 replays), the +consistency decline fired only 6 times (`no_captured_shape: 6`). GSM8K +95.944 vs the 96.0 gate — within one σ (≈0.54 at n=1319) of the historical +96.1–96.4 band; two more samples running to separate noise from a real +regression (suspect list if real: the offsets clamp, the trim-schedule +shift from declined steps). + +The investigation scaffolding (zero-sync device tripwire + forensic stash, +prepare/staged tick stamps, instance markers, pre-replay staleness assert) +was stripped from the tree once the fix was validated — this document and +`tmp/autopsy_recipe.md` (frozen-scene attach protocol, kernel param-space +offsets, the never-wrap-cuda-gdb-in-timeout rule) are the durable record. +Rebuilding any piece takes an hour with this doc; keeping them all in-tree +was judged not worth the noise. + ### Reproduce ```bash @@ -233,17 +438,18 @@ rank is untried and is the obvious next move. ## 6. Next steps -1. **Re-derive the candidate list** from "ragged trimming is unsafe", not "eager - is unsafe". The sharpest surviving constraint is `compact_eagerfull`: with - every window at `max_draft_len` the step is semantically identical to static, - yet it faults — so the trigger is that `is_ragged_verify` is True, not what - the windows contain. Audit every `is_ragged_verify` consumer. Two remain - unexamined: `refresh_token_major_block_table` (a 4-D `index_select` that has - produced a device-side out-of-bounds before) and the four branches in - `sparse_attn_indexer`. -2. **Single-rank compute-sanitizer.** It names the kernel and the address; that - ends the guessing. Everything else in this section is another hypothesis. -3. Only after the IMA: #19 (budget quantisation, the largest measured gap), +1. **Validate the coarse prepare() guard** (§3 root cause) — a gated and an + eager run with `COREDUMP=1`. A pass means both complete ≥128 steps with + non-zero `steps_ragged` and no new corefile; do not accept anything less. +2. If it still faults: the corefile names the kernel again. The remaining + unguarded staging is then *outside* `prepare()` — sweep `model_engine.py` + and `py_executor.py` for persistent pinned buffers with `non_blocking=True` + consumers (the `_pinned_host` sites are already event-guarded). +3. After the pass: replace the coarse guard's protection of the worst offenders + with the base class's snapshot idiom where profiling justifies it, drop the + dead capture-skip conditionals, re-evaluate the `4f81cc5534` gate, and clean + up the ~34 corefiles (keep one as evidence). +4. Only after the IMA: #19 (budget quantisation, the largest measured gap), then #11 and #20. ### Do not @@ -260,16 +466,14 @@ rank is untried and is the obvious next move. ## 7. Unpushed work -`origin/confidence_head` is at `f033e3cf5b`. Three local commits, none -end-to-end verified: +`origin/confidence_head` is at `d36c9b4757`. Two local commits plus the +coarse-guard working tree, none end-to-end verified yet: | commit | what | note | |---|---|---| -| `59e2d47b41` | clear the ragged pad window when the fit is skipped | correct on its own merits | -| `4f81cc5534` | gate the eager-ragged path off | **stated justification is now known wrong** (§3) | -| `cdc4538616` | row-map staleness assertion | hypothesis it tested is falsified; assertion is cheap and may still be worth keeping | - -Also uncommitted: `dsa.py` carries the capture-skip fix for that assertion. +| `0bb36c11b6` | event-guard the ragged row-map staging + `_pinned_host` sites | did **not** stop the fault alone, but closes real WAR holes; keep | +| `9bc5dec76c` | skip staging copies during capture | **no-op** — `prepare()` never runs inside the capture region (§3); fold into cleanup | +| *(uncommitted)* | coarse WAR guard around all of `DSAtrtllmAttentionMetadata.prepare` | the actual fix candidate, in validation | Two fixes worth keeping regardless of the IMA, already in the above: `_pinned_host` (persistent staging buffers — PyTorch does not extend an async From 72dd80314102217b88f17e7f187722a3a7c70731 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 4 Aug 2026 20:07:57 -0700 Subject: [PATCH 086/177] [None][fix] Interpolate the SPS cost table between breakpoints The floor lookup billed every total below the next breakpoint at the previous breakpoint's price. On the live table's sparse grid that priced a full block sitting just under a breakpoint at roughly half its true cost, so the planner's cost ratio between tiers collapsed below the survival ratio and the argmax bought the full block on ~95% of decisions: confidence scheduling never trimmed at all, on any run. Lookup is now a clamped linear interpolation over both theta(M) and alpha(bs), matching the SGLang additive-table consumer this was ported from. A table that wants a flat shelf has to measure the shelf -- two breakpoints with equal values -- and compress_to_risers keeps a shelf's last point for that reason: dropping it would turn a measured shelf into one long rising segment and over-bill every mid-shelf total, the mirror image of the bug being fixed. Table values are validated finite: NaN compares false against every bound, so the old checks passed it through and a single NaN would poison both adjacent segments and the argmax. The engine fingerprint now compares numeric facts numerically rather than by spelling. The decision function is pinned against an inlined copy of SGLang's budget planner, exact up to two documented divergences: the floor shift (SGLang admits position 0 as a budget candidate while it rides in tau's base here) and zero-length windows (SGLang can starve a hopeless request; min_verify_len is 1 here). Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/speculative/dspark_planner.py | 170 ++++++++--- .../_torch/speculative/dspark_sps_profiler.py | 264 ++++++++++++++++-- .../test_dspark_budget_parity_sglang.py | 232 +++++++++++++++ .../test_dspark_confidence_schedule.py | 53 +++- .../test_dspark_cost_interpolation.py | 87 ++++++ .../hw_agnostic/test_dspark_sps_profiler.py | 29 +- 6 files changed, 752 insertions(+), 83 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py index 0276d1ba0656..011352dfb514 100644 --- a/tensorrt_llm/_torch/speculative/dspark_planner.py +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -36,8 +36,8 @@ whole curve; so do we, and :mod:`tests` pins it against a brute-force scan. **Cost is per step, not per token.** ``SpsCostTable`` therefore stores measured -step time against the *total* verified token count, and interpolates nothing -across a riser. +step time against the *total* verified token count, interpolated between +breakpoints (see the class docstring for why the old floor lookup was wrong). **Total verified tokens includes the bonus position.** A request that verifies ``L`` drafted positions submits ``L + 1`` tokens to the target -- the bonus @@ -61,14 +61,17 @@ before the batch could be shaped, which is exactly the sync we are avoiding. """ -import bisect +import math from dataclasses import dataclass, field from typing import List, Optional, Sequence import numpy as np +from ...logger import logger + __all__ = [ "SpsCostTable", + "check_table_fingerprint", "compute_verify_token_budget", "budget_argmax_over_uniform_lens", "derive_verify_len_tiers", @@ -99,9 +102,9 @@ class SpsCostTable: already inside ``step_time_ms``. Defaults to 0.0 on the assumption that a profiled table measured whole steps; set it only when the table isolates the verify term. - batch_sizes / batch_overhead_ms: optional ``alpha(bs)`` staircase -- the - batch-size-dependent, *non-trimmable* part of the step (the draft - pass, weight movement). Looked up by floor on ``batch_sizes``. + batch_sizes / batch_overhead_ms: optional ``alpha(bs)`` breakpoints -- + the batch-size-dependent, *non-trimmable* part of the step (the + draft pass, weight movement). Interpolated on ``batch_sizes``. Getting the non-trimmable terms wrong is not harmless. The planner maximizes ``tau / T``, a ratio, so ``bias + alpha(bs)`` does not cancel: understating @@ -110,11 +113,18 @@ class SpsCostTable: of milliseconds while ``theta(M)`` is a few, so a hardcoded guess here would dominate the decision. - Lookup is a floor (staircase) lookup, never an interpolation: between two - measured points the cost is genuinely flat until a new kernel wave starts, - and smoothing that away is what makes a planner over-spend right before a - riser. Queries above the largest breakpoint clamp to the last entry, and the - caller is expected to bound the budget so this does not happen silently. + Lookup is a clamped LINEAR INTERPOLATION between breakpoints, matching the + additive-table consumer this was ported from (SGLang's + ``_additive_step_time_tensor``). A floor lookup here is wrong on a sparse + table: it bills every total below the next breakpoint at the previous + breakpoint's price, making tokens look free right before a riser, and the + argmax then over-spends -- measured live as the planner buying the full + block on ~95% of decisions because its cost ratio between tiers collapsed + below the survival ratio (see test_dspark_cost_interpolation for the + pinned numbers). A table that wants a flat shelf must MEASURE the shelf: + two breakpoints with equal values. Queries outside the measured range + clamp to the end values; the caller is expected to bound the budget so the + high clamp does not happen silently. """ token_counts: Sequence[int] @@ -133,8 +143,8 @@ def __post_init__(self) -> None: raise ValueError("SpsCostTable requires at least one measured point") if any(b <= a for a, b in zip(self.token_counts, self.token_counts[1:])): raise ValueError("token_counts must be strictly increasing") - if any(t <= 0.0 for t in self.step_time_ms): - raise ValueError("step_time_ms entries must be positive") + if any(not math.isfinite(t) or t <= 0.0 for t in self.step_time_ms): + raise ValueError("step_time_ms entries must be positive and finite") if self.fixed_overhead_ms < 0.0: raise ValueError("fixed_overhead_ms must be >= 0") if len(self.batch_sizes) != len(self.batch_overhead_ms): @@ -144,36 +154,32 @@ def __post_init__(self) -> None: ) if any(b <= a for a, b in zip(self.batch_sizes, self.batch_sizes[1:])): raise ValueError("batch_sizes must be strictly increasing") - if any(t < 0.0 for t in self.batch_overhead_ms): - raise ValueError("batch_overhead_ms entries must be >= 0") + if any(not math.isfinite(t) or t < 0.0 for t in self.batch_overhead_ms): + raise ValueError("batch_overhead_ms entries must be >= 0 and finite") def batch_overhead(self, num_requests: int) -> float: """``alpha(num_requests)`` -- 0.0 when the table has no batch axis.""" if not self.batch_sizes: return 0.0 - idx = bisect.bisect_right(self.batch_sizes, int(num_requests)) - 1 - idx = min(max(idx, 0), len(self.batch_overhead_ms) - 1) - return float(self.batch_overhead_ms[idx]) + return float( + np.interp( + float(num_requests), + np.asarray(self.batch_sizes, dtype=np.float64), + np.asarray(self.batch_overhead_ms, dtype=np.float64), + )) def step_time(self, num_tokens: int, num_requests: int = 0) -> float: """``bias + alpha(num_requests) + theta(num_tokens)``, in ms.""" - idx = bisect.bisect_right(self.token_counts, int(num_tokens)) - 1 - idx = min(max(idx, 0), len(self.step_time_ms) - 1) - return ( - float(self.step_time_ms[idx]) - + self.fixed_overhead_ms - + self.batch_overhead(num_requests) - ) + return float(self.step_times(np.asarray([int(num_tokens)]), num_requests)[0]) def step_times(self, num_tokens: np.ndarray, num_requests: int = 0) -> np.ndarray: """Vectorized :meth:`step_time` over a token-count array.""" - counts = np.asarray(self.token_counts) - idx = np.clip(np.searchsorted(counts, num_tokens, side="right") - 1, 0, len(counts) - 1) - return ( - np.asarray(self.step_time_ms, dtype=np.float64)[idx] - + self.fixed_overhead_ms - + self.batch_overhead(num_requests) + theta = np.interp( + np.asarray(num_tokens, dtype=np.float64), + np.asarray(self.token_counts, dtype=np.float64), + np.asarray(self.step_time_ms, dtype=np.float64), ) + return theta + self.fixed_overhead_ms + self.batch_overhead(num_requests) @classmethod def flat(cls, step_time_ms: float = 1.0) -> "SpsCostTable": @@ -197,6 +203,59 @@ def is_flat(self) -> bool: return len(set(self.step_time_ms)) <= 1 +def check_table_fingerprint(*, payload: dict, live: dict) -> None: + """Refuse a cost table profiled on a different engine than this one. + + The table's numbers are only meaningful for the engine configuration they + were measured on: the same (bs=64, M=384) cell measured 49.5 ms on a + max_batch_size=256 engine and 61.0 ms on a max_batch_size=64 engine -- 23% + apart against a 0.04-1.2% run-to-run noise floor -- and a table taken on + CUTLASS MoE kernels was nearly loaded to plan MegaMoE steps. Neither + mistake produces an error anywhere downstream; the planner just trims to + the wrong economics. + + Only the keys PRESENT in both dicts are compared, so a table may carry + facts the consumer cannot see (image hash, geometry) for humans to check, + and an old table with no fingerprint loads with a warning rather than + breaking existing deployments. + """ + fp = (payload.get("_meta") or {}).get("engine") or payload.get("engine") + if not fp: + logger.warning( + "DSpark cost table carries no engine fingerprint; cannot verify it " + "was profiled on this configuration. Tables are engine-specific " + "(same cell measured 23%% apart across two configs) -- regenerate " + "with a current profiler to get load-time checking.") + return + def _facts_differ(a, b) -> bool: + try: + return float(a) != float(b) + except (TypeError, ValueError): + return str(a).upper() != str(b).upper() + + mismatched = { + key: (fp[key], live[key]) + for key in sorted(set(fp) & set(live)) + if _facts_differ(fp[key], live[key]) + } + if mismatched: + detail = ", ".join(f"{k}: table={a!r} engine={b!r}" + for k, (a, b) in mismatched.items()) + raise ValueError( + f"DSpark cost table was profiled on a different engine " + f"configuration ({detail}). Its step times do not describe this " + f"engine, so the planner would trim to the wrong economics. " + f"Re-profile on this configuration, or remove " + f"confidence_sps_table_path to run without trimming.") + unchecked = sorted(set(fp) - set(live)) + if unchecked: + logger.info( + f"DSpark cost table fingerprint: verified " + f"{sorted(set(fp) & set(live))}; not verifiable at this site " + f"(check manually): " + + ", ".join(f"{k}={fp[k]!r}" for k in unchecked)) + + def compute_verify_token_budget( *, survival: np.ndarray, @@ -204,6 +263,7 @@ def compute_verify_token_budget( cost_table: SpsCostTable, min_verify_len: int = 1, max_verify_len: Optional[int] = None, + allowed_lens: Optional[Sequence[int]] = None, ) -> int: """Return the batch-wide verify-token budget (tokens *above* the floor). @@ -216,10 +276,25 @@ def compute_verify_token_budget( cost_table: measured step cost; see :class:`SpsCostTable`. min_verify_len / max_verify_len: per-request bounds, mirroring :class:`~.dspark_schedule.DSparkScheduleConfig`. + allowed_lens: restrict the answer to budgets that a tier in this ladder + can actually realise, i.e. ``n in {bs*(t - min_verify_len)}``. Pass + it whenever the result will be executed, and omit it only to ask + what an unconstrained scheduler would have wanted. + + The unrestricted answer is not runnable. The captured graphs are + ``bs*(t+1)`` tokens for ``t`` in the ladder + (``model_engine.ragged_verify_token_buckets``), and any other total + is rounded UP to one of them -- so a budget chosen strictly inside a + cost shelf, which is the only reason to choose one, is pushed back + over the riser it was avoiding. The planner would then buy tokens + priced on the cheap shelf and run on the expensive one. Realising a + fine budget honestly would need a graph per token count, which is + tens of MB of metadata each, taken out of KV cache. Returns: ``n`` maximizing ``tau(n) / step_cost(floor_tokens + n)``, in ``[0, N]`` - where ``N`` is the number of schedulable candidates. + where ``N`` is the number of schedulable candidates -- restricted to the + tier-aligned indices when ``allowed_lens`` is given. """ if survival.ndim != 2: raise ValueError(f"survival must be [bs, K], got shape {survival.shape}") @@ -249,7 +324,24 @@ def compute_verify_token_budget( floor_tokens = total_verify_tokens(bs, min_verify_len) tokens = floor_tokens + np.arange(tau.size) theta = tau / cost_table.step_times(tokens, num_gen_requests) - return int(np.argmax(theta)) + if allowed_lens is None: + return int(np.argmax(theta)) + # Tier-aligned: same theta curve, evaluated only where the executor can + # land. tokens(n) at n = bs*(t - min_verify_len) is bs*(min+1) + bs*(t-min) + # = bs*(t+1), which is exactly total_verify_tokens(bs, t) -- the same cost + # grid budget_argmax_over_uniform_lens uses. The only thing that changes is + # the numerator: tau here is the sum of the top n survivals, which is what + # schedule_verify_lens_topk actually collects, rather than the sum of a + # uniform column, which nothing executes. + idx = sorted({ + int(bs) * (int(t) - int(min_verify_len)) + for t in allowed_lens + if int(min_verify_len) <= int(t) <= cap + }) + idx = [n for n in idx if 0 <= n < theta.size] + if not idx: + return 0 + return int(max(idx, key=lambda n: theta[n])) def derive_verify_len_tiers( @@ -263,7 +355,8 @@ def derive_verify_len_tiers( """Derive the verify lengths worth capturing a CUDA graph for. Restricting K to a handful of captured values sounds like it must cost - something versus a freely-chosen K. **At a fixed batch size it does not**: + something versus a freely-chosen K. **At a fixed batch size, on a table + with genuine shelves, it does not**: Within one cost shelf the step time is constant while ``tau`` strictly increases with each admitted token, so ``Theta = tau / cost`` is strictly @@ -272,6 +365,15 @@ def derive_verify_len_tiers( tier set of the shelf right edges contains the continuous optimum exactly. + Since the consumer became an interpolation, "shelf" means a measured flat + segment (adjacent breakpoints with equal values), not the gap between two + sparse samples. On a table whose segments genuinely rise -- like the live + GB300 one -- the optimum can sit inside a segment and the right-edge set + is a slope-change heuristic rather than an exact cover; the residual loss + has to be measured. (The old floor consumer made every gap LOOK like a + shelf, which both made this theorem vacuously "exact" and priced a + 1512-token step at the 768-token price, disabling trimming entirely.) + **The zero-loss property does not survive across batch sizes.** The cost shelves live in token space, so a shelf's right edge in *length* space is ``(breakpoint - 1) // batch_size - 1`` -- a function of the batch size. A diff --git a/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py b/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py index dea4669a9643..ae809e3afadc 100644 --- a/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py +++ b/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py @@ -170,6 +170,7 @@ import random import statistics import sys +import threading import time from dataclasses import dataclass, field from typing import Dict, List, Optional, Sequence, Tuple @@ -707,15 +708,19 @@ def compress_to_risers( min_riser_ms: float, max_breakpoints: int, ) -> Tuple[List[int], List[float]]: - """Collapse a sampled curve into the staircase it is approximating. - - ``SpsCostTable`` looks up by floor and interpolates nothing, so every - breakpoint claims "a new kernel wave starts here". Shipping one breakpoint - per *sample* instead of one per *riser* asserts risers that do not exist, - and (if the dead ``derive_verify_len_tiers`` path is ever revived) explodes - the tier count so the truncation drops the real shelf edges. - - A point is kept only when it clears the previous kept shelf by + """Collapse a sampled curve into shelf-and-riser breakpoints. + + ``SpsCostTable`` interpolates between breakpoints, so what a table CLAIMS + is entirely in which points it keeps: dropping a shelf's interior is fine + (interpolation across equal values is still flat), but dropping a shelf's + LAST point turns the measured shelf into one long rising segment and + over-bills every mid-shelf total by up to the riser height -- the + over-trim mirror of the floor-lookup bug. So each kept riser is preceded + by the last point of the shelf it rises from: shelves survive as + breakpoint pairs, exactly the encoding the planner tests use for + "genuinely flat". + + A riser is kept only when it clears the previous kept shelf by ``min_riser_ms``. The threshold is **absolute**, not relative to ``theta``: ``theta`` is normalized to approach zero at ``M = 0``, so at the low end a fraction of a millisecond of noise is a large *relative* move and a relative @@ -734,6 +739,11 @@ def compress_to_risers( kept = [0] for index in range(1, len(token_counts)): if step_times[index] - step_times[kept[-1]] > float(min_riser_ms): + # Close the shelf we are leaving before recording the riser, so + # interpolation stays flat across it instead of ramping from the + # shelf's first sample to the riser top. + if index - 1 != kept[-1]: + kept.append(index - 1) kept.append(index) limit = max(int(max_breakpoints), 2) @@ -1013,6 +1023,15 @@ class SweepConfig: max_seq_len: int = 4096 max_num_tokens: int = 8192 kv_cache_fraction: float = 0.5 + #: Fitted STS temperatures. Not used to decide anything while the pin is on; + #: it is what lets ``enable_ragged_verify`` pass validation without the cost + #: table this tool exists to produce. + sts_path: Optional[str] = None + #: Any non-flat cost table. See the note at its use site: with a one-rung + #: ladder its values cannot affect the result, only its presence. + seed_table_path: Optional[str] = None + #: Restore the original sweep geometry; see the note in _build_llm. + block_follows_verify_len: bool = False attn_backend: str = "TRTLLM" disable_overlap_scheduler: bool = False #: On by default because the deployment being profiled runs with it on, and @@ -1023,6 +1042,9 @@ class SweepConfig: #: How long the post-cell stats drain waits for the tail records to arrive #: over IPC before it decides the queue is empty. stats_timeout_s: float = 5.0 + #: Mid-cell poll interval. Short enough that ~1000 rows cannot accumulate + #: between polls (at dp_size=8 that is 125 iterations, i.e. a few seconds). + stats_poll_s: float = 2.0 token_id_range: Tuple[int, int] = (1000, 10000) seed: int = 1234 extra_llm_api_options: Dict[str, object] = field(default_factory=dict) @@ -1032,7 +1054,8 @@ def dp_size(self) -> int: """Independent scheduling domains, i.e. rows per iteration in the stats.""" return int(self.tp_size) if self.enable_attention_dp else 1 - def max_tokens_for(self, verify_len: int) -> int: + def max_tokens_for(self, verify_len: int, + batch_size: Optional[int] = None) -> int: """Output-token budget that guarantees the cell reaches its step count. With acceptance pinned every step commits exactly one token, so the @@ -1040,9 +1063,27 @@ def max_tokens_for(self, verify_len: int) -> int: ``verify_len + 1``, so the budget is scaled to keep the lower bound on steps at the target -- at the price of running up to that factor longer when acceptance turns out to be poor. + + ``batch_size`` adds the *admission ramp*, and without it no large cell + can ever yield a sample. A cell is only counted while all ``dp_size`` + ranks hold exactly ``batch_size`` generating requests, so the request + admitted first must still be alive when the last one is admitted. + Measured on job 2561986 (bs=256, input_len=1024): the scheduler admits + roughly one context request per iteration per rank, so the ramp is + ~``batch_size`` iterations, while the budget was 85 -- the first + requests died ~170 iterations before the last arrived, concurrency + plateaued at 11 of 256, and every one of the 8 repeats kept nothing. + bs<=128 escaped only because its ramp fit inside the same 85. + + The ramp is bounded below by prefill tokens / ``max_num_tokens``, but + that bound was off by ~60x here, so use the observed one-per-iteration + rate instead: it is the conservative direction, and overshooting costs + run time while undershooting costs the entire cell. """ target_steps = int(self.warmup_steps) + int(self.measure_steps) budget = target_steps + max(target_steps // 4, 16) + if batch_size is not None: + budget += int(batch_size) return budget if self.pin_acceptance else budget * (int(verify_len) + 1) def validate(self) -> None: @@ -1054,7 +1095,8 @@ def validate(self) -> None: # A request that hits max_seq_len is evicted mid-cell, so the batch # drains and the steady window shrinks -- which surfaces much later as a # confusing "not enough steady samples". - longest = self.input_len + max(self.max_tokens_for(v) for v in self.verify_lens) + longest = self.input_len + max( + self.max_tokens_for(v, b) for v in self.verify_lens for b in self.batch_sizes) if longest > self.max_seq_len: raise SweepGeometryError( f"a request would reach {longest} tokens (input {self.input_len} + " @@ -1109,7 +1151,15 @@ def _prepare_environment(config: SweepConfig) -> None: Set before the LLM is constructed because the worker processes inherit the environment at spawn and read these once. """ - os.environ[RAGGED_VERIFY_MODE_ENV] = RaggedVerifyMode.STATIC.value + # COMPACT, not STATIC. Under STATIC the token axis is bs*(max_draft_len+1) + # and the only way to vary M is to rebuild with a shorter block -- which + # also shortens the DRAFT pass, so the measured difference between two + # cells mixes "verified fewer tokens" with "drafted fewer tokens". The + # deployed path drafts the full block every step and varies only the verify + # window, so a table built that way over-states what trimming buys. COMPACT + # plus the per-engine pin below holds the draft at max_draft_len and moves + # only the axis the planner actually controls. + os.environ[RAGGED_VERIFY_MODE_ENV] = RaggedVerifyMode.COMPACT.value if config.pin_acceptance: from .interface import FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR @@ -1121,19 +1171,94 @@ def _build_llm(config: SweepConfig, verify_len: int): from tensorrt_llm import LLM from tensorrt_llm.llmapi import CudaGraphConfig, DSparkDecodingConfig, KvCacheConfig + # Held at the checkpoint's trained block for EVERY cell. Sweeping it (which + # this did) makes the draft pass shrink along with the verify window, and + # the two costs are then inseparable in the fitted theta(M). The previous + # K-sweep on this hardware put the draft pass at ~17 ms and the marginal + # verify token at ~0 -- so a table that lets draft cost leak into theta(M) + # would attribute the draft saving to trimming and predict a benefit the + # deployment cannot collect. The verify window is moved by the pin below. + # The window is pinned by BOTH the env override and the ladder, and the + # env var is SET here, never popped. + # + # Set, because the workers are spawned by srun with the job environment: a + # job-wide pin reaches ranks 1..N-1 regardless of what this process does to + # its own env, so popping it here strips rank 0 alone. A rank-0 planner + # with no pin and no cost table falls back to uniform while its peers run + # pinned -- a rank-divergent shape under attention-DP. With one verify + # length per sweep process the pin always equals the ladder rung below, so + # it can never be rejected either. + # + # (An earlier revision popped it and relied on the one-rung ladder alone, + # citing a "pin skips the padding reservation" theory for the Repeat.cu + # assert. That theory was wrong -- the assert reproduced with no pin at + # all; the actual cause was the CUDA-graph key returning the PREDICTED + # bucket rather than the batch's measured token total, fixed in + # cuda_graph_runner._ragged_verify_bucket. And the one-rung ladder was + # never one rung anyway: the llm_args validator appends max_draft_len, so + # under a constant block [L] is really [L, block]. With a seed table + # present the planner argmax'd its way to the block on 93-95% of steps + # while the samples file recorded M = bs*(L+1) as if L had run -- twelve + # cells, eight mislabelled, and the fit dutifully concluded q ~ 0.) + from .dspark_verify import FORCE_VERIFY_LEN_ENV + + # LAUNCH-MODE CAVEAT: this mutation reaches ranks 1..N-1 only when the + # workers are spawned AFTER it (MPI-spawn path). On a pre-spawned world + # (srun / trtllm-llmapi-launch) already-running ranks keep the previous + # value, the per-L pin diverges across ranks, and cells measure a shape + # their label does not describe. assert_swept_shape refuses such a sweep + # downstream (verify_len_hist vs the nominal L), which is why it gates the + # fit rather than merely warning. + os.environ[FORCE_VERIFY_LEN_ENV] = str(int(verify_len)) + # Which knob moves M. + # + # constant block (default): the block stays at the checkpoint's trained + # length for every cell and only the verify window moves, which is exactly + # what the deployment does. It is the right experiment -- but it makes the + # whole batch run with max(verify_len) < block_size, and the ragged packing + # currently asserts in that regime (Repeat.cu output_size vs sum(repeats)). + # + # block follows verify_len: the original design. Never produces + # max(lens) < block, so it runs today, at the price of shrinking the DRAFT + # pass along with the window -- theta(M) then absorbs a per-L draft term. + # That term is recoverable rather than fatal: with several batch sizes the + # same M is reached at different L (M=192 from bs=32/L=5 and bs=64/L=2), and + # those shared points separate a token cost from an L cost. See + # tmp/dspark_build/verdict_theta.py. + block = int(verify_len) if config.block_follows_verify_len else int(config.max_draft_len) spec_config = DSparkDecodingConfig( - max_draft_len=int(verify_len), - # DSparkWorker asserts draft_model.block_size == max_draft_len, so the - # block has to move with the swept length rather than stay at the - # checkpoint's trained value. - block_size=int(verify_len), + max_draft_len=block, + block_size=block, speculative_model=config.speculative_model, enable_confidence_scheduling=bool(config.enable_confidence_scheduling), + # Both required together: the validator rejects confidence scheduling + # without ragged verification (there are only two states -- schedule per + # request, or verify the full block). The profiler had only the first + # and could not construct its own engine at all; it failed 1:36 into + # every two-node sweep. And ragged in turn demands a profiled artifact, + # which is circular for the tool that produces one -- an STS path is the + # non-circular way to satisfy it, and the pin makes the planner's own + # choice irrelevant here anyway. + enable_ragged_verify=bool(config.enable_confidence_scheduling), + confidence_sts_path=config.sts_path, + # Only has to be non-flat and present: with one rung the planner's + # choice is already determined, so the table's VALUES cannot influence + # this measurement -- it exists to get past the flat-cost fallback, + # which would otherwise return None and put every step back on the + # uniform path at M = bs*(block+1). + confidence_sps_table_path=config.seed_table_path, # A single tier keeps the captured-graph count identical to an # unscheduled run (one graph per batch size) while still exercising the # confidence head, which is part of the deployment's per-step cost. It # also makes the planner's choice a no-op: max_tier == max_draft_len, so # the step shape is pinned no matter what the confidence snapshot says. + # The pin decides the length; the ladder only has to CONTAIN it, or the + # planner would return a length with no captured graph and every step + # would fall back to eager -- which is not the cost being tabulated. + # One rung: the planner's argmax has a single legal answer, so the step + # shape is pinned no matter what the confidence snapshot says, while + # every downstream path (bucket choice, padding reservation, graph + # selection) runs exactly as it does in production. confidence_verify_len_tiers=( [int(verify_len)] if config.enable_confidence_scheduling else None ), @@ -1182,7 +1307,7 @@ def _run_cell(llm, config: SweepConfig, *, batch_size: int, verify_len: int) -> sampling_params = SamplingParams( # Every request must outlive the whole window so the batch never drains # mid-measurement and shrinks the shape being profiled. - max_tokens=config.max_tokens_for(verify_len), + max_tokens=config.max_tokens_for(verify_len, batch_size), ignore_eos=True, # No detokenization: it is postprocessing work that would show up in the # host step time without being part of what the planner is choosing. @@ -1199,19 +1324,67 @@ def _run_cell(llm, config: SweepConfig, *, batch_size: int, verify_len: int) -> ] futures = [llm.generate_async(prompt, sampling_params) for prompt in prompts] + + # Drained WHILE the cell runs, not once at the end. Only ~1000 stats rows + # survive to a single drain -- 125 iterations at dp_size=8 -- and the + # scheduler admits roughly one context request per iteration, so a cell + # spends ~batch_size iterations ramping up and just as many draining. At + # bs>=128 the whole retained window therefore lands inside the drain: every + # step it contains has 0..11 generating requests, none has `batch_size`, + # and the cell reports "kept nothing" while the plateau it was measuring + # scrolled out of the buffer unseen (jobs 2561986, 2563525). + # + # The hazard this replaces is real but narrower than the old comment + # claimed: `IterationResult` latches done when its queue runs dry, so an + # empty read ends collection. Polling only while requests are still + # outstanding keeps the queue fed; the count of empty reads is reported so a + # cell that lost its tail this way is visible rather than merely short. + rows: List[dict] = [] + drain_error: List[BaseException] = [] + stop = threading.Event() + + def _drain() -> None: + while not stop.is_set(): + try: + rows.extend(list(llm.get_stats(timeout=config.stats_poll_s) or [])) + except BaseException as exc: # noqa: BLE001 - reported, not swallowed + drain_error.append(exc) + return + stop.wait(config.stats_poll_s) + + # A daemon drainer, NOT a poll loop that owns the exit condition. Making the + # main thread wait on `future.finished` while draining hung the sweep: the + # loop span silently, the log stopped growing, and the stall watchdog killed + # the job at exactly 15 minutes with no cell produced. The main thread here + # blocks on future.result() exactly as it did before this change, so + # termination is unaffected by anything the drainer does. + # stats_poll_s <= 0 disables the mid-cell drain: the control arm for "does + # the profiler perturb what it measures". The cell then relies on the single + # post-cell drain, which only sees the last ~1000 rows -- fine for small + # batches whose whole cell fits in that window, useless above bs~64. That is + # exactly why this is a control and not an option. + drainer = None + if float(config.stats_poll_s) > 0: + drainer = threading.Thread(target=_drain, name="dspark-sps-drain", + daemon=True) + drainer.start() try: for future in futures: future.result() finally: + stop.set() + if drainer is not None: + drainer.join(timeout=max(5.0, 3.0 * config.stats_poll_s)) for future in futures: if not future.finished: future.abort() - # Drained exactly once, after the cell has finished. `IterationResult` - # latches itself done the first time its queue runs dry and is only reset - # when new requests are submitted, so polling mid-cell would silently - # discard every step after the first empty read. The timeout is the wait for - # the last few records to cross the IPC boundary, not a poll interval. - rows: List[dict] = list(llm.get_stats(timeout=config.stats_timeout_s) or []) + if drain_error: + print(f"[dspark-sps] bs={batch_size} L={verify_len}: mid-cell drain " + f"stopped early ({type(drain_error[0]).__name__}: " + f"{drain_error[0]}); the cell may be short", + file=sys.stderr) + # Final drain for whatever crossed the IPC boundary after the last poll. + rows.extend(list(llm.get_stats(timeout=config.stats_timeout_s) or [])) aligned = aligned_steps_from_stats( rows, expected_ranks=config.dp_size, timing_key=config.timing_key @@ -1248,8 +1421,10 @@ def run_sweep(config: SweepConfig) -> List[StepSample]: samples: List[StepSample] = [] for verify_len in sorted(config.verify_lens): print( - f"[dspark-sps] building engine for verify_len={verify_len} " - f"(block_size={verify_len}) ...", + f"[dspark-sps] building engine: block_size=" + f"{verify_len if config.block_follows_verify_len else config.max_draft_len}" + f"{' (follows L)' if config.block_follows_verify_len else ' (constant)'}" + f", verify window {verify_len} -> M = bs*{verify_len + 1} ...", file=sys.stderr, ) llm = _build_llm(config, verify_len) @@ -1368,6 +1543,24 @@ def _build_parser() -> argparse.ArgumentParser: "Each value costs one engine build.", ) sweep.add_argument("--input-len", type=int, default=1024, help="Prompt length in tokens.") + sweep.add_argument("--block-follows-verify-len", action="store_true", + help="Shrink the drafted block with the verify window " + "(the original geometry). Avoids the " + "max(lens) int: enable_attention_dp=args.enable_attention_dp, max_draft_len=args.max_draft_len, input_len=args.input_len, + sts_path=args.sts_path, + seed_table_path=args.seed_table, + stats_poll_s=float(args.stats_poll_s), + block_follows_verify_len=bool(args.block_follows_verify_len), warmup_steps=args.warmup_steps, measure_steps=args.measure_steps, repeats=args.repeats, @@ -1531,6 +1728,21 @@ def main(argv: Optional[Sequence[str]] = None) -> int: "modal_batch_size": reference_bs, "warmup_steps_discarded": int(args.warmup_steps), "statistic": "median", + # The engine fingerprint the loader verifies (check_table_fingerprint). + # Attached only when this process ran the sweep itself: an offline + # --fit-only refit of merged sample files cannot vouch for the + # engine those samples came from, and an unfingerprinted table + # loads with a warning rather than a false claim. + **({"engine": { + "tp": int(args.tp_size), + "ep": int(args.ep_size or 0), + "attention_dp": bool(args.enable_attention_dp), + "block": block_size, + "max_batch_size": max(_int_list(args.batch_sizes)), + "geometry": ("block_follows_verify_len" + if getattr(args, "block_follows_verify_len", False) + else "constant_block"), + }} if not getattr(args, "samples_in", None) else {}), **provenance, }, ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py new file mode 100644 index 000000000000..ac989a423428 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py @@ -0,0 +1,232 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decision-level parity with the SGLang budget planner this was ported from. + +A throughput gain alone does not validate the port: the wrong economics can +still move the needle in the right direction. This pins the DECISION function +against a line-for-line reference copy of SGLang's +``compute_verify_token_budget`` + ``_additive_step_time_tensor`` +(python/sglang/srt/speculative/dspark_components/dspark_planner.py), so a +divergence in either port shows up as an argmax mismatch here rather than as +an unexplained fleet-level regression. + +The two planners parameterize the same optimization differently: + +- SGLang admits EVERY position (including position 0) as a budget candidate + and its floor is one bonus token per request: total M = bs + n. +- TensorRT-LLM grants positions below ``min_verify_len`` for free (their + survival rides in the tau base) and only positions at or beyond it are + schedulable: total M = bs*(min_verify_len+1) + m. + +For ``min_verify_len=1`` these describe the same curve shifted by ``bs``: +survival is non-increasing along a row, so the top-(bs+m) SGLang candidates +are exactly the bs position-0 entries plus the top-m deeper entries, giving +``tau_sgl(bs+m) == tau_trt(m)`` and ``M_sgl(bs+m) == M_trt(m)``. The argmaxes +must therefore agree up to that shift whenever SGLang admits at least the +full position-0 set. Known, deliberate divergences are pinned separately: +SGLang filters candidates by ``survival_eps`` inside the budget function +(TensorRT-LLM applies eps at allocation time instead), and SGLang can choose +sub-floor budgets (windows of zero draft positions) that TensorRT-LLM's +``min_verify_len=1`` floor excludes. +""" + +import numpy as np +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, compute_verify_token_budget) + +# -------------------------------------------------------------------------- +# Reference implementation, copied from SGLang dspark_components/dspark_planner +# (additive-table path) with only cosmetic renames. Do not "fix" this copy: +# its job is to be what SGLang runs. +# -------------------------------------------------------------------------- + + +def _sgl_interp_clamped(xs, ys, x: float) -> float: + xs = torch.tensor(xs, dtype=torch.float64) + ys = torch.tensor(ys, dtype=torch.float64) + x_t = torch.tensor(float(x), dtype=torch.float64).clamp_(xs[0], xs[-1]) + hi = torch.bucketize(x_t, xs, right=True).clamp_(1, xs.numel() - 1) + lo = hi - 1 + span = (xs[hi] - xs[lo]).clamp_(min=1e-9) + frac = (x_t - xs[lo]) / span + return float(ys[lo] + frac * (ys[hi] - ys[lo])) + + +def _sgl_additive_step_time(table, num_requests: int, num_budgets: int): + floor = table["bias_seconds"] + _sgl_interp_clamped( + table["bs_probes"], table["alpha_seconds"], float(num_requests)) + m_probes = torch.tensor(table["m_probes"], dtype=torch.float64) + theta_vals = torch.tensor(table["theta_seconds"], dtype=torch.float64) + m = (num_requests + torch.arange(num_budgets, dtype=torch.float64)).clamp_( + min=float(table["m_probes"][0]), max=float(table["m_probes"][-1])) + hi = torch.bucketize(m, m_probes, right=True).clamp_(1, m_probes.numel() - 1) + lo = hi - 1 + span = (m_probes[hi] - m_probes[lo]).clamp_(min=1e-9) + frac = (m - m_probes[lo]) / span + theta_at_m = theta_vals[lo] + frac * (theta_vals[hi] - theta_vals[lo]) + return floor + theta_at_m + + +def _sgl_compute_verify_token_budget(*, history_survival_probs, table, + max_verify_len, survival_eps): + num_requests = history_survival_probs.shape[0] + candidates = history_survival_probs[:, :max_verify_len].flatten() + candidates = candidates[candidates >= survival_eps].to(torch.float64) + candidates_sorted = torch.sort(candidates, descending=True).values + prefix_sum = torch.cumsum(candidates_sorted, dim=0) + tau_star = num_requests + torch.cat( + [torch.zeros(1, dtype=torch.float64), prefix_sum]) + step_time = _sgl_additive_step_time(table, int(num_requests), + int(tau_star.numel())) + theta = tau_star / step_time + return int(torch.argmax(theta)) + + +# -------------------------------------------------------------------------- +# One certified-table shape shared by both sides. TensorRT-LLM stores ms, +# SGLang stores seconds; the argmax is scale-invariant but the tables are +# built with the honest factor anyway. +# -------------------------------------------------------------------------- + +M_PROBES = (64, 96, 128, 192, 384, 512, 768, 1536) +THETA_MS = (9.68, 11.08, 11.95, 13.95, 22.09, 23.85, 35.61, 105.64) +BIAS_MS = 25.244443 +BS_PROBES = (32, 64, 128, 256) +ALPHA_MS = (0.0, 2.177477, 11.462212, 19.342661) +BLOCK = 5 + +TRT_TABLE = SpsCostTable(token_counts=M_PROBES, step_time_ms=THETA_MS, + fixed_overhead_ms=BIAS_MS, batch_sizes=BS_PROBES, + batch_overhead_ms=ALPHA_MS) +SGL_TABLE = { + "m_probes": M_PROBES, + "theta_seconds": tuple(v / 1e3 for v in THETA_MS), + "bias_seconds": BIAS_MS / 1e3, + "bs_probes": BS_PROBES, + "alpha_seconds": tuple(v / 1e3 for v in ALPHA_MS), +} + + +def _inversion_free_survival(rng, bs): + """Survivals where every position-0 beats every deeper candidate. + + The exact-parity theorem needs this: SGLang admits position 0 as a budget + candidate while TensorRT-LLM grants it in the floor, so the two describe + the same curve (shifted by ``bs``) only when the top-(bs+m) SGLang set is + exactly "all of position 0, plus the top-m deeper entries". A weak row's + position-0 falling below a strong row's position-1 breaks that -- see + test_starving_weak_rows_is_a_known_gap for what SGLang does with it. + """ + conf0 = rng.uniform(0.85, 0.99, size=(bs, 1)) + deeper = rng.uniform(0.25, 0.80, size=(bs, BLOCK - 1)) + return np.cumprod(np.concatenate([conf0, deeper], axis=1), axis=1) + + +@pytest.mark.parametrize("bs", [1, 3, 8, 64, 128, 252, 256]) +def test_budget_argmax_matches_sglang_up_to_the_floor_shift(bs): + """Same survival, same table, eps off, no inversions: shift-exact parity.""" + rng = np.random.default_rng(17 + bs) + for _ in range(5): + surv = _inversion_free_survival(rng, bs) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.0) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + if sgl_n >= bs: + assert trt_m == sgl_n - bs, ( + f"bs={bs}: SGLang admits {sgl_n} candidates " + f"(= floor {bs} + {sgl_n - bs}), TRT-LLM budget {trt_m}") + else: + # SGLang went below the one-draft floor; the closest budget + # TensorRT-LLM can express under min_verify_len=1 is zero. + assert trt_m == 0 + + +def test_live_pooled_survivals_park_on_the_breakpoint_together(): + """The campaign's operating point, both consumers, both argmax modes. + + With interpolated segments the continuous optimum parks where the theta + slope jumps: buying position-2 survivals (0.502/token) beats the + 512..768 slope (0.046 ms/token) but not the 768..1536 one (0.091), so + both planners stop at exactly M = 768 -- TensorRT-LLM at budget + 768 - 2*252 = 264, SGLang at 768 - 252 = 516. Tier alignment then rounds + down to the rung-2 budget the executor can capture. + """ + bs = 252 + surv = np.tile([0.771, 0.693, 0.502, 0.360, 0.315], (bs, 1)) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.0) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + assert trt_m == 768 - 2 * bs + assert sgl_n - bs == trt_m + tiered = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1, allowed_lens=[1, 2, 5]) + assert tiered == bs * 1 # rung-2: one scheduled position past the floor + + +def test_starving_weak_rows_is_a_known_gap(): + """SGLang can verify NOTHING for a hopeless request; TRT-LLM cannot. + + SGLang's budget admits position-0 candidates individually, so a row whose + draft is dead on arrival can be starved to a zero-length window and its + would-be tokens spent on a strong row's deep positions. TensorRT-LLM's + ``min_verify_len=1`` floor buys position 0 for every request + unconditionally. On a batch that mixes both kinds, SGLang's achievable + Theta is therefore strictly better -- pinned here so a future + min_verify_len=0 port is a deliberate decision with a test to flip. + """ + # Half the batch is strong all the way down, half is dead on arrival. + # The batch must be big enough that its token range sits on a genuinely + # rising part of the theta curve: below the first breakpoint the clamp + # makes cost flat, every candidate is free, and both planners trivially + # buy everything at the same theta -- the two-row version of this test + # sat in that regime and proved nothing. + bs = 64 + strong = np.tile([0.99, 0.98, 0.97, 0.96, 0.95], (bs // 2, 1)) + dead = np.tile([0.01, 0.008, 0.006, 0.004, 0.002], (bs // 2, 1)) + surv = np.concatenate([strong, dead], axis=0) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.0) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + + def theta_at(tau, tokens): + return tau / float(TRT_TABLE.step_times(np.asarray([tokens]), bs)[0]) + + cand = np.sort(surv.reshape(-1))[::-1] + sgl_theta = theta_at(bs + cand[:sgl_n].sum(), bs + sgl_n) + deeper = np.sort(surv[:, 1:].reshape(-1))[::-1] + trt_theta = theta_at(bs + surv[:, 0].sum() + deeper[:trt_m].sum(), + 2 * bs + trt_m) + assert sgl_theta > trt_theta + + +def test_eps_filtering_is_a_known_divergence(): + """SGLang drops sub-eps candidates inside the budget; TRT-LLM does not. + + TensorRT-LLM applies survival_eps at allocation time + (schedule_verify_lens_topk's eligibility mask), so its BUDGET can count a + candidate the allocator then refuses. This pins the divergence so a + future unification is a deliberate change, not an accident: with a heavy + sub-eps tail the SGLang budget must not exceed TRT-LLM's shifted one. + """ + bs = 16 + surv = np.tile([0.9, 0.4, 0.008, 0.004, 0.002], (bs, 1)) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.01) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + assert sgl_n - bs <= trt_m diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py index fd79ded34949..a4071fdee9ac 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py @@ -152,13 +152,18 @@ def test_greedy_first_descent_would_be_wrong(): shelf's optimum is almost 3x better. """ bs = 4 - # One riser just past the floor (bs*(min_len+1) = 8 tokens), then a long - # shelf out to the end of the range. The fixed overhead is spelled out - # rather than assumed: it is part of what makes the first shelf's Theta - # look competitive, and it is no longer a hidden constant in the planner. - table = SpsCostTable(token_counts=(0, 10), step_time_ms=(1.0, 3.0), fixed_overhead_ms=1.0) - # Uniformly high confidence: every extra candidate adds nearly a full token, - # so tau grows ~7x across the range while cost only doubles. + # A steep measured segment right past the floor (bs*(min_len+1) = 8 + # tokens), then a nearly-flat one out to the end of the range. Under + # interpolation the kink between segment slopes is what makes Theta + # non-unimodal: the first candidates buy tokens on the steep slope and + # Theta drops, then the flat segment makes every further token nearly + # free and Theta climbs past the starting point. + table = SpsCostTable(token_counts=(0, 8, 9, 40), + step_time_ms=(0.5, 0.5, 4.0, 4.2), + fixed_overhead_ms=1.0) + # Uniformly high confidence: every extra candidate adds nearly a full + # token, so tau grows ~7x across the range while cost less than doubles + # past the wall. surv = np.cumprod(np.full((bs, BLOCK), 0.995), axis=1) n_star = compute_verify_token_budget(survival=surv, num_gen_requests=bs, cost_table=table) @@ -172,8 +177,8 @@ def theta(n): greedy += 1 assert theta(n_star) >= theta(greedy) - assert n_star > greedy, "expected the staircase to defeat a first-descent scan" - assert theta(n_star) > 2 * theta(greedy), "the gap should be large, not marginal" + assert n_star > greedy, "expected the cost kink to defeat a first-descent scan" + assert theta(n_star) > 1.5 * theta(greedy), "the gap should be large, not marginal" def test_flat_cost_table_degenerates_to_verify_all(): @@ -204,11 +209,23 @@ def test_budget_is_zero_for_empty_batch(): ) -def test_cost_table_is_a_staircase_not_an_interpolation(): +def test_cost_table_interpolates_between_breakpoints(): + """The consumer contract is clamped linear interpolation, not a floor. + + The floor variant this replaces billed every total below the next + breakpoint at the previous breakpoint's price. On the sparse live table + (points at 768 and 1536) that priced a 1512-token full block at the 768 + price -- the planner's cost ratio collapsed below the survival ratio and + it bought the full block on ~95% of decisions. Flatness between points is + a claim about the hardware; if a table wants a shelf, it must MEASURE the + shelf (two breakpoints with equal values), not have the consumer assume it. + """ table = SpsCostTable(token_counts=(0, 100), step_time_ms=(1.0, 5.0)) - # Everything below the next breakpoint costs the same as the shelf it is on. - assert table.step_time(1) == table.step_time(99) - assert table.step_time(100) > table.step_time(99) + assert table.step_time(99) > table.step_time(1) + assert table.step_time(50) == pytest.approx(3.0) + # A measured shelf stays flat under interpolation. + shelf = SpsCostTable(token_counts=(0, 99, 100), step_time_ms=(1.0, 1.0, 5.0)) + assert shelf.step_time(1) == shelf.step_time(99) # Above the last breakpoint we clamp rather than extrapolate. assert table.step_time(10_000) == table.step_time(100) @@ -374,9 +391,15 @@ def test_derived_tiers_lose_nothing_versus_continuous_k(): an exhaustive search over *every* length. """ rng = np.random.default_rng(11) + # The shelves are ENCODED, breakpoint pairs with equal values: under the + # interpolating consumer, flatness between points is only real when the + # table measured it. (Assuming it -- the old floor lookup -- is what + # priced a 1512-token step at the 768-token price on the live table and + # disabled trimming outright; the theorem below is conditional on genuine + # shelves and this encoding is what "genuine" means now.) table = SpsCostTable( - token_counts=(0, 12, 30, 48, 90, 160), - step_time_ms=(2.0, 2.4, 3.9, 4.0, 7.5, 12.0), + token_counts=(0, 11, 12, 29, 30, 47, 48, 89, 90, 159, 160), + step_time_ms=(2.0, 2.0, 2.4, 2.4, 3.9, 3.9, 4.0, 4.0, 7.5, 7.5, 12.0), ) for bs in (3, 8, 16, 31): surv = np.cumprod(rng.uniform(0.5, 0.995, size=(bs, BLOCK)), axis=1) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py new file mode 100644 index 000000000000..32bc50364039 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Cost lookup must interpolate between breakpoints, not floor onto them. + +The floor lookup this replaces silently disabled the whole feature: with +measured points at 768 and 1536, every total in 769..1535 was billed at the +768 price, so a full block at bs=252 (1512 tokens, true cost ~148 ms) was +priced 72.3 ms. The planner's cost ratio T(6bs)/T(3bs) collapsed from a true +~2.0 to 1.19 -- below the GSM8K survival ratio tau5/tau2 ~ 1.48 -- so the +argmax bought the full block on ~95% of decisions, the bucket landed on +padded_bs*6 on every step of every run, and trimming never happened. The only +exception ever observed was bs=256 exactly, where the tier totals land on +breakpoints and floor happens to be exact. + +The consumer this was ported from (SGLang ``_additive_step_time_tensor``) +interpolates both theta(M) and alpha(bs); this port must match. +""" + +import numpy as np +import pytest + +from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, compute_verify_token_budget) + +# The certified GB300 table's shape, abbreviated to the two breakpoints whose +# gap hid the bug. +TABLE = SpsCostTable(token_counts=(512, 768, 1536), + step_time_ms=(23.85, 35.61, 105.64), + fixed_overhead_ms=25.244, + batch_sizes=(128, 256), + batch_overhead_ms=(11.462, 19.343)) + + +def test_theta_is_exact_on_breakpoints(): + assert TABLE.step_time(768, 256) == pytest.approx(25.244 + 19.343 + 35.61) + assert TABLE.step_time(1536, 256) == pytest.approx(25.244 + 19.343 + 105.64) + + +def test_theta_interpolates_between_breakpoints(): + """The regression case: 1512 tokens must cost ~the 1536 price, not the 768 one.""" + got = TABLE.step_time(1512, 252) + theta = 35.61 + (105.64 - 35.61) * (1512 - 768) / (1536 - 768) + alpha = 11.462 + (19.343 - 11.462) * (252 - 128) / (256 - 128) + assert got == pytest.approx(25.244 + alpha + theta) + # The floor price it used to return -- and must never return again. + assert got > 130.0 + + +def test_alpha_interpolates_and_clamps(): + assert TABLE.batch_overhead(128) == pytest.approx(11.462) + assert TABLE.batch_overhead(192) == pytest.approx((11.462 + 19.343) / 2) + assert TABLE.batch_overhead(64) == pytest.approx(11.462) # clamp low + assert TABLE.batch_overhead(512) == pytest.approx(19.343) # clamp high + + +def test_theta_clamps_outside_the_measured_range(): + """Theta clamps to the end values; a table without a batch axis adds no alpha.""" + plain = SpsCostTable(token_counts=(512, 768, 1536), + step_time_ms=(23.85, 35.61, 105.64), + fixed_overhead_ms=25.244) + assert plain.step_time(100, 0) == pytest.approx(25.244 + 23.85) + assert plain.step_time(4096, 0) == pytest.approx(25.244 + 105.64) + + +def test_step_times_vector_matches_scalar(): + tokens = np.asarray([512, 900, 1512, 1536]) + vec = TABLE.step_times(tokens, 200) + for tok, v in zip(tokens.tolist(), vec.tolist()): + assert v == pytest.approx(TABLE.step_time(tok, 200)) + + +def test_full_batch_argmax_now_trims(): + """The live failure in miniature: GSM8K pooled survivals at a near-full batch. + + tau5/tau2 ~ 1.48 sits below the true cost ratio ~2.0 but above the floored + one (1.19). With interpolation the argmax must stop buying the full block; + the whole campaign's zero-trim runs came down to this comparison. + """ + bs = 252 + survival = np.tile([0.771, 0.693, 0.502, 0.360, 0.315], (bs, 1)) + budget = compute_verify_token_budget(survival=survival, + num_gen_requests=bs, + cost_table=TABLE, + min_verify_len=1, + allowed_lens=[1, 2, 5]) + rung = 1 + budget // bs + assert rung < 5, f"argmax still buys the full block (budget={budget})" diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py index e2940f0eca15..1f1d9d342d20 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py @@ -325,14 +325,21 @@ def test_running_max_is_monotone(): assert running_max([1.0, 3.0, 2.0, 2.5, 9.0]) == [1.0, 3.0, 3.0, 3.0, 9.0] -def test_compress_keeps_only_real_risers(): +def test_compress_keeps_risers_and_closes_their_shelves(): + """A shelf must survive as a breakpoint PAIR under the interp consumer. + + Keeping only [16, 64] would turn the measured flat 16..48 into one long + rising segment: mid-shelf totals get billed part of the riser -- the + over-trim mirror of the floor-lookup bug. The shelf's last point (48) + rides along so interpolation stays flat where the hardware was flat. + """ tokens = [16, 32, 48, 64, 80] times = [1.0, 1.002, 1.004, 2.0, 2.001] kept_tokens, kept_times = compress_to_risers( tokens, times, min_riser_ms=0.02, max_breakpoints=8 ) - assert kept_tokens == [16, 64] - assert kept_times == [1.0, 2.0] + assert kept_tokens == [16, 48, 64] + assert kept_times == [1.0, 1.004, 2.0] def test_compress_drops_the_shallowest_risers_first(): @@ -455,9 +462,14 @@ def test_refuses_a_curve_that_only_moves_by_noise(): # (2.5%), yet the step only gets 2.5% more expensive while the acceptance yield # given up between those tiers is ~36% at p=0.9. The argmax therefore never # leaves max_tier -- non-flat and inert at the same time. +# +# The shelf is ENCODED as a breakpoint pair: under the interpolating consumer +# a bare [0, 160] would ramp from 2.0 to 2.5 across the whole range, dropping +# the ladder spread below the noise floor and turning this fixture into the +# flat-noise case instead of the inert one it exists to pin. INERT_PAYLOAD = { - "token_counts": [0, 160], - "step_time_ms": [2.0, 2.5], + "token_counts": [0, 159, 160], + "step_time_ms": [2.0, 2.0, 2.5], "fixed_overhead_ms": 18.0, "batch_sizes": [32], "batch_overhead_ms": [0.0], @@ -600,9 +612,10 @@ def test_probe_uses_the_bonus_token_convention(): payload = { # Priced so that tier 5 (M = 8*6 = 48) is expensive and tier 3 # (M = 8*4 = 32) is cheap. Under the wrong bs*L convention tier 5 would - # be M = 40 and land on the cheap shelf instead. - "token_counts": [0, 44], - "step_time_ms": [1.0, 50.0], + # be M = 40 and land on the cheap shelf instead. The shelf is a + # breakpoint pair so the interpolating consumer keeps it flat. + "token_counts": [0, 43, 44], + "step_time_ms": [1.0, 1.0, 50.0], "fixed_overhead_ms": 0.0, "batch_sizes": [8], "batch_overhead_ms": [0.0], From c117be89a1f6f78801b3fd15d0249c2657a9cf50 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 4 Aug 2026 20:08:37 -0700 Subject: [PATCH 087/177] [None][feat] Calibrate DSpark confidence and measure the planner's inputs Two halves of the same problem: the survivals the planner divides by were never calibrated, and its per-step inputs were never measured. Calibration (dspark_sts.py) fits a per-position temperature ladder from collected (confidence logit, accepted length) pairs and applies it before the cumulative product, mirroring SGLang's apply_sts. The head is resolved through the runtime wrapper rather than a chained getattr: the wrapper keeps its stages under .dspark_model, so the old lookup silently returned None and every confidence-scheduled run planned on raw sigmoid survivals -- the only symptom was one missing log line. A configured calibration that cannot be attached now raises. Collection refuses to run when a cost table is loaded without a full-block window pin (a trimmed window censors the acceptance label) and when confidence scheduling is off (the recorder is only built on that path, so the run would complete cleanly and collect nothing). Instrumentation records what the argmax actually eats. Each confidence row carries the draft-pass sequence that wrote it, scattered in-graph through the same slot indices as the confidence itself and staged to the host behind the same event, so the planner can histogram how stale each gathered row was; rows still holding the neutral fill are counted separately. The rung and the window sizes the planner chose are recorded before the bucket fit widens them -- without that, a fit-side defect and a planner-side one are indistinguishable in the counters. Measurement only: a failure inside any of it is counted, never raised. The draft-pass sequence scalar is allocated outside inference mode because the executor's host path writes it in place; a tensor born inside the draft forward is an inference tensor and refuses that write on every rank. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 45 ++- tensorrt_llm/_torch/speculative/dspark.py | 128 ++++++- .../speculative/dspark_observability.py | 24 +- .../_torch/speculative/dspark_schedule.py | 9 + tensorrt_llm/_torch/speculative/dspark_sts.py | 331 ++++++++++++++++++ .../_torch/speculative/dspark_verify.py | 207 +++++++++-- .../_torch/speculative/spec_sampler_base.py | 27 ++ tests/microbenchmarks/dspark_fit_sts.py | 307 ++++++++++++++++ .../hw_agnostic/test_dspark_padded_ceiling.py | 98 ++++++ .../test_dspark_reconcile_guards.py | 106 ++++++ .../test_dspark_snapshot_instruments.py | 110 ++++++ .../test_dspark_sts_collection_guards.py | 144 ++++++++ .../test_dspark_tier_aligned_budget.py | 130 +++++++ 13 files changed, 1630 insertions(+), 36 deletions(-) create mode 100644 tensorrt_llm/_torch/speculative/dspark_sts.py create mode 100644 tests/microbenchmarks/dspark_fit_sts.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index eb1cf94f3013..17aa194b2768 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2507,6 +2507,14 @@ def _report_dspark_ragged_stats(self) -> None: """ worker = getattr(self.model_engine, "_get_spec_worker", None) stats = getattr(worker(), "ragged_stats", None) if worker else None + # Drained BEFORE the ragged early-return. The two are unrelated: STS + # collection runs in static mode too, and gating its final flush on + # `ragged_stats.steps_total > 0` silently discarded the tail of every + # static-mode collection run -- up to flush_every steps x batch rows, + # and the shards still looked well-formed. + recorder = getattr(self.sampler, "sts_recorder", None) + if recorder is not None: + recorder.flush() if stats is None or not getattr(stats, "steps_total", 0): return stats.log_summary(prefix="DSpark ragged verify [final]") @@ -3376,6 +3384,16 @@ def _dspark_confidence_draft_len(self, if (stats is not None and self.sampler is not None and getattr(self.sampler, "acceptance_stats", None) is None): self.sampler.acceptance_stats = stats + # Same wiring, same reason, for STS collection: the logits live on the + # worker and the accepted count is only host-side in the sampler, and + # `py_seq_slot` indexes both. None on every path that is not collecting. + if (self.sampler is not None + and getattr(self.sampler, "sts_recorder", None) is None): + recorder = getattr(worker, "sts_recorder", None) + if recorder is not None: + self.sampler.sts_recorder = recorder + self.sampler.sts_logits_buffer = getattr( + worker, "_confidence_logits", None) mode = getattr(worker, "ragged_verify_mode", None) compute_windows = (mode.computes_windows if mode is not None else getattr(self.model_engine.spec_config, @@ -3463,8 +3481,10 @@ def _dspark_confidence_draft_len(self, fallback_reason = "peer_declined" ragged_lens = None else: - agreed_max = max(int(payload[3]) for payload in payloads) - ragged_lens = [min(int(v), agreed_max) for v in ragged_lens] + # payload[3] (the local max window) stays in the collective so + # its shape is stable, but capping local windows at the global + # max was a no-op by construction (the global max is >= every + # local window) and nothing else consumed it. # Third element is the group's answer, identical on every rank # by construction, not each rank's own flag: what the fit needs # to know is whether the row maximum will be taken at all. @@ -3570,7 +3590,13 @@ def _dspark_confidence_draft_len(self, stats.record_step(num_gen_requests=num_gen_requests, verify_lens=fitted, bucket=bucket, - padded_bs=None, + # Was hardcoded None, which both left + # padded_bs_hist empty in every summary ever + # logged and denied record_step the row count it + # needs to book a comparable ceiling. + padded_bs=getattr( + self.model_engine, + "_dspark_last_padded_bs", None), fallback=fallback_reason, delivered=delivered) @@ -3605,7 +3631,18 @@ def _dspark_confidence_draft_len(self, # whichever requests it happens to be scheduling. buffer = worker.staged_confidence_buffer() if buffer is not None: - planner.stage_confidence(buffer) + # The stamp buffer says which draft pass last wrote each slot; the + # sequence bump marks "everything drafted so far" so the planner + # can histogram how stale each gathered row was at decision time. + staged_seq = None + bump = getattr(worker, "bump_draft_seq", None) + if bump is not None: + staged_seq = bump() + stamps_fn = getattr(worker, "confidence_stamp_buffer", None) + planner.stage_confidence( + buffer, + stamps=stamps_fn() if stamps_fn is not None else None, + staged_seq=staged_seq) return int(runtime_draft_len) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index c152c005cfbf..171bcd76e5ab 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -30,6 +30,7 @@ from tensorrt_llm.mapping import Mapping from ..pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID +from .dspark_schedule import NEUTRAL_CONFIDENCE_LOGIT from .interface import SpecMetadata, SpecWorkerBase if TYPE_CHECKING: @@ -37,7 +38,7 @@ # Raw-logit fill for an unscored confidence slot: sigmoid(30) rounds to 1.0 in # fp32, so an unwritten row schedules as "verify the whole block". -_NEUTRAL_CONFIDENCE_LOGIT = 30.0 +_NEUTRAL_CONFIDENCE_LOGIT = NEUTRAL_CONFIDENCE_LOGIT @dataclass @@ -250,8 +251,22 @@ def __init__( # here (never per step) so the captured draft graph cannot diverge. # ``_confidence_logits`` is the slot-indexed handoff buffer to the # verification scheduler; see ``_lazy_init`` for its neutral value. - self.return_confidence = bool(getattr(spec_config, "enable_confidence_scheduling", False)) + self.return_confidence = bool(spec_config.enable_confidence_scheduling) + import os as _os + + from .dspark_sts import STS_COLLECT_ENV + if _os.environ.get(STS_COLLECT_ENV) and not self.return_confidence: + raise ValueError( + f"{STS_COLLECT_ENV} is set but enable_confidence_scheduling is " + f"off; the recorder is only built on the confidence path, so " + f"this run would complete cleanly and collect NOTHING. Enable " + f"confidence scheduling (a missing cost table keeps the " + f"planner on the full block) or unset the variable.") self._confidence_logits: Optional[torch.Tensor] = None # [max_batch+2, block] + # Draft-pass stamps for the buffer above; allocated in ``_lazy_init``. + self._confidence_stamp: Optional[torch.Tensor] = None + self._draft_seq_host = 0 + self._draft_seq_cuda: Optional[torch.Tensor] = None # Row index of the permanently-neutral confidence row; see ``_lazy_init``. self._neutral_conf_row = 0 # Host-side coordinator that turns the confidence snapshot into this @@ -282,6 +297,25 @@ def staged_confidence_buffer(self) -> Optional[torch.Tensor]: """ return self._confidence_logits + def confidence_stamp_buffer(self) -> Optional[torch.Tensor]: + """The slot-indexed draft-pass stamps, or None before ``_lazy_init``.""" + return self._confidence_stamp + + def bump_draft_seq(self) -> int: + """Advance the draft-pass sequence and return the PREVIOUS value. + + Called once per step from the executor's host path, before this step's + forward is enqueued: the in-graph stamp scatter then labels this step's + confidence rows with the new sequence, while the returned previous + value is the sequence of the last COMPLETED draft -- i.e. the newest + content a snapshot staged right now can possibly contain. + """ + prev = self._draft_seq_host + self._draft_seq_host += 1 + if self._draft_seq_cuda is not None: + self._draft_seq_cuda.fill_(self._draft_seq_host) + return prev + def confidence_row_for(self, req_id: int) -> int: """Buffer row holding ``req_id``'s confidence, host-side. @@ -366,6 +400,25 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: dtype=torch.float32, device="cuda", ) + # Freshness stamps, one per buffer row: the draft-pass sequence + # number of the last write. SGLang carries the same thing + # (_carry_generation) next to its confidence ring; without it there + # is no way to know whether a gathered row is this step's data, a + # stale replay, or the neutral fill -- the planner path has never + # been measured. Written in-graph through the same `slots` scatter + # as the confidence itself; the scalar is bumped outside the graph. + self._confidence_stamp = torch.zeros( + (num_rows + 1,), dtype=torch.int32, device="cuda") + self._draft_seq_host = 0 + # Bumped by fill_ from the executor's host path, which runs + # OUTSIDE inference mode; a tensor born here (inside the draft + # forward's inference mode) is an "inference tensor" and refuses + # that update -- it killed all 8 ranks on the first decode step. + # The stamp buffer above stays an inference tensor on purpose: + # it is only written in-graph, like _confidence_logits. + with torch.inference_mode(False): + self._draft_seq_cuda = torch.zeros((), dtype=torch.int32, + device="cuda") self._build_verify_planner(draft_model, block_size, max_batch) self._free_slots = deque(range(max_batch)) self._req_to_slot = {} @@ -392,7 +445,7 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> cfg = DSparkScheduleConfig(block_size=block_size, min_verify_len=1) cost_table = None - table_path = getattr(self.spec_config, "confidence_sps_table_path", None) + table_path = self.spec_config.confidence_sps_table_path if table_path: with open(table_path, encoding="utf-8") as f: raw = json.load(f) @@ -400,6 +453,19 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> # a ratio, so the non-trimmable part of the step time moves the # argmax. Omitting them says "step_time_ms already measures whole # steps at the deployment's batch size". + from .dspark_planner import check_table_fingerprint + check_table_fingerprint(payload=raw, live={ + "tp": int(getattr(self.mapping, "tp_size", 0) or 0), + "ep": int(getattr(self.mapping, "moe_ep_size", 0) or 0), + "attention_dp": bool(getattr(self.mapping, "enable_attention_dp", False)), + "block": int(block_size), + "max_batch_size": int(max_batch), + # moe_backend deliberately NOT in the live dict yet: the + # MoeConfig.backend -> model_config.moe_backend propagation is + # not textually obvious, and a wrong live value here would + # false-reject a correct table. Until verified, the table's + # moe_backend surfaces on the "check manually" INFO line. + }) cost_table = SpsCostTable( token_counts=[int(v) for v in raw["token_counts"]], step_time_ms=[float(v) for v in raw["step_time_ms"]], @@ -421,16 +487,37 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> # Calibration lives on the draft's confidence head; fall back to a plain # sigmoid when the head is absent (no confidence to calibrate anyway). - head = getattr(getattr(draft_model, "mtp_layers", [None])[-1], "confidence_head", None) + from .dspark_sts import resolve_confidence_head + head = resolve_confidence_head(draft_model) apply_calibration = head.apply_sts if head is not None else None - sts_path = getattr(self.spec_config, "confidence_sts_path", None) + sts_path = self.spec_config.confidence_sts_path + if sts_path and head is None: + # Refuse rather than degrade. The old chained getattr defaulted to + # None when the runtime object was the DSparkForCausalLM wrapper + # (which keeps the stages under .dspark_model, not .mtp_layers), so + # every confidence-scheduled run in the first campaign silently + # planned on RAW sigmoid survivals -- wildly overconfident (the + # fitted temperatures are 2.8-10) -- and chose the full block 99.7% + # of the time. The only visible symptom was the ABSENCE of one log + # line. A configured calibration that cannot be attached is a + # wiring bug, never a preference. + raise ValueError( + f"confidence_sts_path is set but no confidence head was found " + f"on {type(draft_model).__name__}; calibration cannot be " + f"applied, so the planner would schedule on raw sigmoid " + f"survivals. This is a model-wiring bug, not a config choice.") if sts_path and head is not None: - with open(sts_path, encoding="utf-8") as f: - temps = json.load(f)["sts_temperatures"] + # Accepts either key spelling: this repo writes "sts_temperatures", + # SGLang's DSparkStsCalibration writes "temperatures". The vectors + # are interchangeable -- both are fitted against the same cumulative + # product with the same sigmoid -- so rejecting the other spelling + # would refuse a usable table over a name. + from .dspark_sts import load_sts_temperatures_from_path + temps = load_sts_temperatures_from_path(sts_path) head.load_sts_temperatures(torch.tensor(temps, dtype=torch.float32)) logger.info(f"DSpark: loaded STS calibration from {sts_path}") - tiers = getattr(self.spec_config, "verify_len_tiers", None) + tiers = self.spec_config.verify_len_tiers if not tiers: # Not configured: derive the ladder from the measured cost curve # instead of hardcoding one. Each tier is a captured graph whose @@ -465,6 +552,19 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> self.spec_config, "enable_ragged_verify", False) else RaggedVerifyMode.STATIC) self.ragged_verify_mode = read_ragged_verify_mode(default=configured) + # STS calibration data collection. Off unless the env var is set; when + # on it costs one device->host copy per decode step, which is why it is + # a collection mode rather than something a serving run carries. + from .dspark_sts import make_recorder_from_env + # `block_size` is the local parameter, not an attribute -- the very + # next line uses it the same way. + self.sts_recorder = make_recorder_from_env( + block_size=int(block_size), + rank=int(getattr(getattr(self, "mapping", None), "rank", 0) or 0), + has_cost_table=bool( + self.spec_config.confidence_sps_table_path), + ragged_mode=getattr(self.ragged_verify_mode, "value", None)) + self.ragged_stats = DSparkRaggedStats(mode=self.ragged_verify_mode, max_draft_len=block_size) # goal doc A6: the ragged counters say what happened, the planner's say @@ -677,6 +777,18 @@ def _draft_gen_block_batched( # one-iteration lag. if confidence is not None: self._confidence_logits[slots] = confidence.detach() + # Same slots, same capture region: content-only update, graph-safe. + self._confidence_stamp[slots] = self._draft_seq_cuda + # STS collection reads this buffer from inside the sampler, which + # under the overlap scheduler can run AFTER the next step's draft + # has already overwritten it. SGLang avoids the question by keeping + # what the draft produced (`last_confidence_raw`) rather than + # re-reading a shared buffer later. Stash the same thing here, so + # the recorder can pair against the step that actually produced it + # -- and, for one diagnostic run, compare the two sources. + if getattr(self, "sts_recorder", None) is not None: + self.sts_recorder.stash_draft_confidence( + logits=confidence.detach(), slots=slots) return block_logits def _forward_impl( diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index 7ac64f4aa40d..e538fd498498 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -277,9 +277,27 @@ def record_step( # and blame the scheduler for a batch that had not filled yet. if int(num_gen_requests) >= 2: self.steps_multi_request += 1 - # What a no-trim step would have submitted: every request sends its - # bonus token plus the full block. - self.ceiling_tokens += int(num_gen_requests) * (1 + self.max_draft_len) + # What a no-trim step would have submitted: every row sends its bonus + # token plus the full block. + # + # ROWS, not requests. When the step ran on a fitted bucket, `delivered` + # below books `bucket = padded_bs * (tier + 1)`, and `padded_bs` is the + # CUDA-graph-rounded maximum row count ACROSS RANKS -- so a rank holding + # fewer requests than its widest peer still submits the peer's rows. + # Booking those rows as delivered against a ceiling counted over this + # rank's real requests made trim_ratio negative in four separate runs + # (jobs 2560992 -0.001, 2561519 -0.0438, 2563581 -0.0082), and the + # assert_ragged_active gate then reported "nothing was saved", which is + # the opposite of what a negative ratio means. + # + # Only that one case rebases. Every other path books `delivered` over + # the real rows -- the no-window path and the cap-accept path each add + # exactly n*(1+K), and the bucket-less ragged path adds sum(1+v) <= + # n*(1+K) -- so widening their ceiling would understate the trim. + ceiling_rows = int(num_gen_requests) + if bucket and delivered is None and padded_bs: + ceiling_rows = max(ceiling_rows, int(padded_bs)) + self.ceiling_tokens += ceiling_rows * (1 + self.max_draft_len) if fallback: self.fallbacks[fallback] += 1 diff --git a/tensorrt_llm/_torch/speculative/dspark_schedule.py b/tensorrt_llm/_torch/speculative/dspark_schedule.py index b130b94c5275..eaba8be7ad3e 100644 --- a/tensorrt_llm/_torch/speculative/dspark_schedule.py +++ b/tensorrt_llm/_torch/speculative/dspark_schedule.py @@ -52,6 +52,15 @@ ] +# Raw logit written into confidence rows that carry no measurement (slot-free +# fill, never-drafted slots): sigmoid(30/T) ~ 1.0 at any fitted temperature, so +# an unknown row is treated as "certainly accept" -- optimistic by design, the +# same convention as SGLang's ones-fill for stale rows. Shared from here +# because both the worker (writes it) and the planner (counts it) need it and +# they cannot import each other. +NEUTRAL_CONFIDENCE_LOGIT = 30.0 + + @dataclass(frozen=True) class DSparkScheduleConfig: """Bounds for the per-request verify length. diff --git a/tensorrt_llm/_torch/speculative/dspark_sts.py b/tensorrt_llm/_torch/speculative/dspark_sts.py new file mode 100644 index 000000000000..4e5779ca44c4 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_sts.py @@ -0,0 +1,331 @@ +# 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. +"""STS calibration: collecting the data, and reading the fitted table. + +The confidence head emits a raw logit per drafted position. ``apply_sts`` turns +it into a probability with a *per-position* temperature:: + + confidence[r][j] = sigmoid(logit[r][j] / T[j]) + survival[r][j] = prod_{i <= j} confidence[r][i] + +``T`` defaults to all ones, which makes that a plain sigmoid -- i.e. no +calibration at all. That default is not neutral, because the planner does not +consume ``confidence``; it consumes the **cumulative product**. A per-position +bias of x compounds to x^(j+1) by position j, so a head that is 5% over-confident +per position over-states the depth-5 survival by more than 30%, and the budget +argmax divides by a cost while multiplying that error into its numerator. + +So the table has to be fitted against the survival, not against the per-position +probability -- which is what makes the fitting objective (see +``tests/microbenchmarks/dspark_fit_sts.py``) a cumulative-product ECE rather than +a per-position one. The recorder here produces exactly the pairs that objective +needs. + +Collection is opt-in via ``TLLM_DSPARK_STS_COLLECT_PATH`` and is **not** a +serving path: it reads the confidence buffer back to the host once per decode +step. That is a real synchronization, deliberately accepted because a +calibration run is not a throughput run. +""" + +import json +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import List, Optional + +import torch + +from ...logger import logger + +__all__ = [ + "STS_COLLECT_ENV", + "DSparkStsCalibration", + "DSparkStsRecorder", + "load_sts_temperatures_from_path", + "make_recorder_from_env", +] + +#: Set to a path stem to collect calibration data. Shards land at +#: ``...pt``. +STS_COLLECT_ENV = "TLLM_DSPARK_STS_COLLECT_PATH" + +#: Flush after this many recorded steps, so a killed run still leaves data. +_DEFAULT_FLUSH_EVERY = 64 + + +@dataclass +class DSparkStsCalibration: + """A fitted per-position temperature vector, plus how well it did.""" + + temperatures: List[float] + ece_before: List[float] = field(default_factory=list) + ece_after: List[float] = field(default_factory=list) + dataset: str = "" + num_samples: int = 0 + + def to_json(self) -> dict: + # Both key spellings are written. ``sts_temperatures`` is what this repo + # has always read; ``temperatures`` is what SGLang's DSparkStsCalibration + # emits, so a table produced here is loadable there and vice versa. + return { + "sts_temperatures": list(self.temperatures), + "temperatures": list(self.temperatures), + "ece_before": list(self.ece_before), + "ece_after": list(self.ece_after), + "dataset": self.dataset, + "num_samples": int(self.num_samples), + } + + +def load_sts_temperatures_from_path(path: str) -> List[float]: + """Read a temperature vector, accepting either spelling of the key. + + This repo writes ``sts_temperatures``; SGLang's ``DSparkStsCalibration`` + writes ``temperatures``. The vectors are interchangeable -- both are fitted + against the cumulative product with the same sigmoid -- so refusing the + other spelling would reject a usable table for no reason. + """ + with open(path, encoding="utf-8") as handle: + payload = json.load(handle) + for key in ("sts_temperatures", "temperatures"): + if key in payload: + temps = payload[key] + break + else: + raise KeyError( + f"{path} has neither 'sts_temperatures' nor 'temperatures'; " + f"found keys {sorted(payload)}") + if not temps: + raise ValueError(f"{path} carries an empty temperature vector") + return [float(t) for t in temps] + + +def resolve_confidence_head(draft_model): + """Find the confidence head across the known draft-model layouts. + + Two layouts exist: the bare ``DSparkDraftModel`` (stages under + ``.mtp_layers``, head on the last stage) and the ``DSparkForCausalLM`` + wrapper (bare model under ``.dspark_model``). A chained ``getattr`` with + soft defaults conflated "wrapper" with "no head" and silently disabled + calibration for an entire measurement campaign; this helper exists so the + resolution is written once and unit-tested against both layouts. + """ + inner = getattr(draft_model, "dspark_model", draft_model) + stages = getattr(inner, "mtp_layers", None) + if not stages: + return None + return getattr(stages[-1], "confidence_head", None) + + +class DSparkStsRecorder: + """Pair each drafted block's logits with how much of it was accepted. + + The two halves live in different places and are joined by sequence slot: + the logits are in the worker's persistent ``[max_batch + 2, K]`` device + buffer, written by the draft pass; the accepted count is host-side in the + sampler. Both are indexed by ``py_seq_slot``, so the join is exact. + + One host copy per step, not per request. ``record`` is called once per + request inside the sampler's loop, and copying a row each time would be a + synchronization per request; instead the whole buffer is pulled once when + the step's first request arrives and reused for the rest. + """ + + def __init__(self, *, path_stem: str, block_size: int, rank: int = 0, + flush_every: int = _DEFAULT_FLUSH_EVERY) -> None: + self.path_stem = str(path_stem) + self.block_size = int(block_size) + self.rank = int(rank) + self.flush_every = max(1, int(flush_every)) + self._logits: List[torch.Tensor] = [] + self._stashed_logits: List[torch.Tensor] = [] + self._accepted: List[int] = [] + self._stash_logits = None + self._stash_slots = None + self._stash_host = None + self._shard = 0 + self._steps_since_flush = 0 + # Host mirror of the device buffer, refreshed once per step. + self._host_logits: Optional[torch.Tensor] = None + self._host_step = -1 + self._stash_capture_warned = False + + def stash_draft_confidence(self, *, logits, slots) -> None: + """Keep what THIS step's draft produced, indexed by slot. + + The persistent buffer is shared and the sampler reads it later; under + the overlap scheduler "later" can be after the next draft has written + it. SGLang sidesteps this by recording the draft's own output rather + than re-reading the buffer, and so does this once the diagnostic below + confirms the two differ. + """ + # DEVICE-side clone only. This runs inside the captured draft graph, so + # a .to("cpu") here raises "Cannot copy between CPU and CUDA tensors + # during CUDA graph capture" -- the same constraint that forces + # _confidence_logits to be written in place. SGLang keeps a device + # reference for exactly this reason. The host copy is deferred to the + # sampler, which is outside the graph. + # + # EAGER-ONLY diagnostic: Python inside a captured region executes at + # capture time only, so under graph replay these clones would keep + # pointing at the last CAPTURE's tensors (dummy-slot content) and the + # late-read comparison would diff against garbage while looking + # perfectly healthy. Refuse to arm the stash during capture and say so + # once, rather than let a replayed collection run silently measure the + # wrong thing. + if torch.cuda.is_current_stream_capturing(): + if not self._stash_capture_warned: + self._stash_capture_warned = True + logger.warning( + "DSpark STS stash skipped under CUDA-graph capture: the " + "late-read diagnostic is eager-only. Run collection with " + "cuda_graph_config disabled to exercise it.") + self._stash_logits = None + self._stash_slots = None + self._stash_host = None + return + self._stash_logits = logits.detach().clone() + self._stash_slots = slots.detach().clone() + self._stash_host = None + + def begin_step(self, step: int) -> None: + """Invalidate the host mirror; the next ``record`` refreshes it.""" + if step != self._host_step: + self._host_logits = None + self._stash_host = None + self._host_step = step + + def record(self, *, slot: int, accepted: int, + device_logits: torch.Tensor) -> None: + """Record one request's ``(logits, accepted)`` pair. + + Args: + slot: ``py_seq_slot``; indexes both the device buffer and nothing + else, which is why it is the join key. + accepted: drafted positions accepted, excluding the bonus token. + This is exactly the ``num_correct_drafts`` the prefix mask is + built from. + device_logits: the worker's ``[rows, K]`` confidence buffer. + """ + if self._host_logits is None: + # The single sync of the step. Deliberate: see the module docstring. + self._host_logits = device_logits.detach().to( + device="cpu", dtype=torch.float32, copy=True) + if slot < 0 or slot >= self._host_logits.shape[0]: + return + buffered = self._host_logits[slot].clone() + self._logits.append(buffered) + self._accepted.append(int(accepted)) + # Diagnostic: the same row as the draft produced it. If these two agree + # everywhere, reading the shared buffer late is safe; if they diverge, + # every pair recorded so far was confidence(t+k) against accepted(t). + # Materialize the draft-time snapshot once per step, here, outside the + # captured region. + if self._stash_host is None and self._stash_logits is not None: + rows = self._stash_logits.to(device="cpu", dtype=torch.float32, + copy=True) + idx = self._stash_slots.to(device="cpu").tolist() + self._stash_host = {int(sl): rows[i] for i, sl in enumerate(idx)} + stashed = (self._stash_host or {}).get(int(slot)) + self._stashed_logits.append( + stashed.clone() if stashed is not None + else torch.full_like(buffered, float("nan"))) + + def end_step(self) -> None: + self._steps_since_flush += 1 + if self._steps_since_flush >= self.flush_every: + self.flush() + + def flush(self) -> None: + """Write a shard. Cheap to call with nothing buffered.""" + self._steps_since_flush = 0 + if not self._logits: + return + logits = torch.stack(self._logits, dim=0) + counts = torch.tensor(self._accepted, dtype=torch.int64).view(-1, 1) + positions = torch.arange(self.block_size).view(1, -1) + # prefix_mask[i, k] == 1 iff draft positions 0..k were ALL accepted, + # which is the event whose probability `survival[k]` estimates. Any + # other label (e.g. "position k accepted") would fit a different + # quantity than the planner consumes. + prefix_mask = (positions < counts).to(torch.float32) + + path = Path(f"{self.path_stem}.r{self.rank}.{self._shard}.pt") + path.parent.mkdir(parents=True, exist_ok=True) + payload = {"logits": logits, "prefix_mask": prefix_mask} + if self._stashed_logits: + payload["logits_at_draft"] = torch.stack(self._stashed_logits, dim=0) + torch.save(payload, path) + logger.info( + f"DSpark STS: wrote {logits.shape[0]} samples to {path}") + self._logits.clear() + self._stashed_logits.clear() + self._accepted.clear() + self._shard += 1 + + +def make_recorder_from_env(*, block_size: int, rank: int = 0, + has_cost_table: bool = False, + ragged_mode: Optional[str] = None + ) -> Optional[DSparkStsRecorder]: + """Build a recorder iff ``TLLM_DSPARK_STS_COLLECT_PATH`` is set. + + Refuses the two regimes in which the collected pairs describe the scheduler + rather than the head. Both are silent failures otherwise: the run completes, + the shards look normal, and the fitted temperatures are wrong in a way no + downstream check can see. SGLang guards the first of these next to its own + collector (dspark_planner.py:100-108) and names the second in the sibling + ConfidenceMetricsProbe ("padded verify rows corrupt the per-position prefix + label", dspark_observability.py:688-697). + """ + stem = os.environ.get(STS_COLLECT_ENV, "").strip() + if not stem: + return None + + forced = os.environ.get("TLLM_DSPARK_FORCE_VERIFY_LEN", "").strip() + forced_full = False + if forced: + try: + forced_full = int(forced) >= int(block_size) + except ValueError: + forced_full = False + if has_cost_table and not forced_full: + # A loaded table means the planner trims, and a trimmed window clips + # `py_num_accepted_draft_tokens`: a request granted 2 positions can show + # at most 2 accepted, however well the head predicted. Fitting that + # calibrates the scheduler. Pinning the window to the full block makes + # the label uncensored by construction. + raise ValueError( + f"{STS_COLLECT_ENV} is set while a cost table is loaded. The planner " + f"will trim, and a trimmed window censors the acceptance label -- the " + f"fit would describe the scheduler, not the confidence head. Either " + f"collect without a cost table, or pin the window with " + f"TLLM_DSPARK_FORCE_VERIFY_LEN= (a pin below the " + f"block still censors the label and is refused).") + + if (ragged_mode or "").lower() == "compact": + # Compact packs the token axis and fills the bucket slack with padding + # rows. Those rows carry positions that were never drafted for a real + # request, so the per-position prefix label they contribute is not a + # measurement of anything. + raise ValueError( + f"{STS_COLLECT_ENV} is set with ragged verify mode 'compact'. Padded " + f"verify rows corrupt the per-position prefix label. Collect in " + f"static mode (TLLM_DSPARK_RAGGED_VERIFY_MODE=static).") + logger.warning( + f"DSpark STS collection ON ({STS_COLLECT_ENV}={stem!r}). This adds one " + f"device->host copy per decode step; it is a calibration mode, not a " + f"serving configuration.") + return DSparkStsRecorder(path_stem=stem, block_size=block_size, rank=rank) diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index eccac0016ca7..deedd9250e67 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -40,14 +40,24 @@ must degrade to verifying everything rather than to a stale guess. """ +import os from typing import Callable, List, Optional, Sequence import numpy as np import torch from ..._utils import prefer_pinned -from .dspark_planner import SpsCostTable, budget_argmax_over_uniform_lens -from .dspark_schedule import DSparkScheduleConfig, compute_survival, schedule_verify_lens_topk +from ...logger import logger +from .dspark_planner import SpsCostTable, compute_verify_token_budget +from .dspark_schedule import (NEUTRAL_CONFIDENCE_LOGIT, DSparkScheduleConfig, + compute_survival, schedule_verify_lens_topk) + +#: Pin every request's verify window to this length, bypassing the planner. +#: Set by the SPS profiler to move M at a fixed draft length; also used to +#: hold the full block during STS collection, where a trimmed window would +#: clip the acceptance label and fit a temperature to the scheduler instead +#: of to the head. Not a serving knob. +FORCE_VERIFY_LEN_ENV = "TLLM_DSPARK_FORCE_VERIFY_LEN" __all__ = ["DSparkVerifyPlanner"] @@ -82,8 +92,27 @@ def __init__( self.apply_calibration = apply_calibration or torch.sigmoid self.all_rank_max = all_rank_max + # Profiling pin. The SPS profiler has to hold M at a known value while it + # times steps, and it cannot use the planner's own decision to do that: + # the planner needs a cost table, and the cost table is what the profiler + # is producing. SGLang solves this with `dspark_force_budget_frac` pushed + # to a running server; TRT-LLM has no runtime control plane at + # world_size > 1, so this is an env var read ONCE here -- constant per + # process, zero cost on the step path, and forwarded to every worker at + # spawn, so all ranks agree by construction rather than by a collective. + # + # Absolute length, not SGLang's fraction. `schedule_verify_lens_topk` + # drops candidates below survival_eps, so a fractional budget yields + # `granted <= budget` and `choose_ragged_capture_shape` then rounds the + # total DOWN onto a smaller captured bucket -- the cell would be timed at + # one M and filed under another. With an absolute L every request gets + # exactly L, so M = bs * (L + 1) is the captured bucket by construction. + self._forced_verify_len = self._read_forced_verify_len() + self._host_buffer: Optional[torch.Tensor] = None self._copy_event: Optional[torch.cuda.Event] = None + self._host_stamps: Optional[torch.Tensor] = None + self._staged_seq: Optional[int] = None self._snapshot_valid = False # Observability: silent degradation is the main failure mode here, so # count every path that gives up on trimming. @@ -94,13 +123,16 @@ def __init__( "fallback_no_confidence": 0, "fallback_short_snapshot": 0, "fallback_no_gen_requests": 0, + "forced_steps": 0, } @property def max_tier(self) -> int: return self.tiers[-1] if self.tiers else self.cfg.resolved_max_verify_len - def stage_confidence(self, confidence_logits: torch.Tensor) -> None: + def stage_confidence(self, confidence_logits: torch.Tensor, + stamps: Optional[torch.Tensor] = None, + staged_seq: Optional[int] = None) -> None: """Start a non-blocking device->host copy of the confidence buffer. ``confidence_logits`` is the worker's whole slot-indexed buffer, staged @@ -125,6 +157,15 @@ def stage_confidence(self, confidence_logits: torch.Tensor) -> None: ) self._copy_event = torch.cuda.Event() self._host_buffer.copy_(confidence_logits, non_blocking=True) + if stamps is not None: + if (self._host_stamps is None + or self._host_stamps.shape != stamps.shape): + self._host_stamps = torch.empty(stamps.shape, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned()) + self._host_stamps.copy_(stamps, non_blocking=True) + self._staged_seq = staged_seq self._copy_event.record() self._snapshot_valid = True @@ -137,6 +178,74 @@ def _ready_snapshot(self) -> Optional[torch.Tensor]: return None return self._host_buffer + _FORCE_ENV = FORCE_VERIFY_LEN_ENV + + def _read_forced_verify_len(self) -> Optional[int]: + """Parse and validate the profiling pin, or None.""" + raw = os.environ.get(self._FORCE_ENV, "").strip() + if not raw: + return None + try: + value = int(raw) + except ValueError: + raise ValueError( + f"{self._FORCE_ENV}={raw!r} is not an integer") from None + lo, hi = int(self.cfg.min_verify_len), int(self.cfg.resolved_max_verify_len) + if not lo <= value <= hi: + raise ValueError( + f"{self._FORCE_ENV}={value} outside [{lo}, {hi}]") + if value not in self.tiers: + # A length with no captured graph makes every step run eager, and an + # eager step is orders of magnitude off the graph-replay cost the + # profiler is trying to measure -- the curve would be noise wearing + # the shape of a cost model. Refuse rather than produce it. + raise ValueError( + f"{self._FORCE_ENV}={value} is not in the captured tier ladder " + f"{self.tiers}; every step would fall out of CUDA-graph replay " + f"and the measured times would not be step costs") + logger.warning( + f"DSpark: verify length PINNED to {value} by {self._FORCE_ENV}. " + f"This bypasses the planner entirely and is a profiling mode, not a " + f"serving configuration.") + return value + + def _note_snapshot_stats(self, selected: torch.Tensor, + rows: Optional[Sequence[int]]) -> None: + """Measure what the argmax is about to eat. Never changes behaviour. + + Two independent instruments: + - stamp-lag histogram: (staged_seq - row stamp) per gathered row. A + healthy relay shows one tight mode; mass at large lags is stale + replay masquerading as data, and the clamp bucket collects + never-drafted rows. + - neutral-row count: rows still carrying the neutral fill on every + position, i.e. "unknown" rows entering the argmax as certainties. + Together they say whether the per-step snapshot content can be + trusted; without them the planner's input was never measured at all. + """ + try: + n = int(selected.shape[0]) + self.stats["snap_rows"] = self.stats.get("snap_rows", 0) + n + neutral = int((selected >= NEUTRAL_CONFIDENCE_LOGIT - 1.0).all(dim=1).sum()) + self.stats["snap_neutral_rows"] = ( + self.stats.get("snap_neutral_rows", 0) + neutral) + if (self._host_stamps is not None and self._staged_seq is not None + and rows is not None): + idx = torch.as_tensor(list(rows), dtype=torch.long) + idx = idx.clamp_(0, self._host_stamps.shape[0] - 1) + lags = (int(self._staged_seq) + - self._host_stamps[idx].to(torch.int64)) + hist = self.stats.setdefault("stamp_lag_hist", {}) + for lag, cnt in zip(*[t.tolist() for t in lags.unique( + return_counts=True)]): + key = int(max(-2, min(int(lag), 8))) + hist[key] = hist.get(key, 0) + int(cnt) + except Exception as exc: # noqa: BLE001 - instruments must not kill steps + self.stats["snap_stats_errors"] = ( + self.stats.get("snap_stats_errors", 0) + 1) + if self.stats["snap_stats_errors"] == 1: + logger.warning(f"DSpark snapshot stats failed once: {exc}") + def _gather_rows( self, *, num_gen_requests: int, rows: Optional[Sequence[int]] ) -> Optional[torch.Tensor]: @@ -220,25 +329,81 @@ def decide_verify_lens( self.stats["fallback_no_gen_requests"] += 1 return None - if self.cost_table is None or self.cost_table.is_flat: - self.stats["fallback_flat_cost"] += 1 - return None - selected = self._gather_rows(num_gen_requests=num_gen_requests, rows=rows) - if selected is None: - return None + if self._forced_verify_len is not None: + # Ahead of the cost-table gate on purpose: the profiler runs without + # a table (that is the artifact it is producing), so a pin placed + # after this gate would never fire and the run would look entirely + # normal while measuring the untrimmed block. + # + # It sets ``lens`` rather than returning, so the pinned path still + # runs the cross-rank agreement and the length assert below. An + # early return skipped both -- harmless while the pin equals the + # block (every rank agrees trivially and the packed buffers are + # full-width anyway), which is why arm B ran clean, but it left the + # short-window case taking a different route through the same code + # than any real decision ever does. + self.stats["forced_steps"] += 1 + lens = [int(self._forced_verify_len)] * int(num_gen_requests) + else: + if self.cost_table is None or self.cost_table.is_flat: + self.stats["fallback_flat_cost"] += 1 + return None + selected = self._gather_rows(num_gen_requests=num_gen_requests, rows=rows) + if selected is None: + return None - probs = self.apply_calibration(selected) - survival = compute_survival(probs) - uniform_len = budget_argmax_over_uniform_lens( - survival=survival.numpy().astype(np.float64), - num_gen_requests=int(num_gen_requests), - cost_table=self.cost_table, - allowed_lens=self.tiers, - min_verify_len=self.cfg.min_verify_len, - ) - # Same total the uniform decision would have spent, redistributed. - budget = int(num_gen_requests) * (int(uniform_len) - self.cfg.min_verify_len) - lens = schedule_verify_lens_topk(survival=survival, budget=budget, cfg=self.cfg).tolist() + self._note_snapshot_stats(selected, rows) + probs = self.apply_calibration(selected) + survival = compute_survival(probs) + # Scored with the tau the step will actually collect. The rung is + # still chosen from the same ladder over the same cost grid -- the + # token totals, the buckets and the captured graphs are unchanged -- + # but budget_argmax_over_uniform_lens scores rung L as if every + # request took columns [min, L), while what runs is a top-k + # redistribution of the same B = n*(L - min) tokens. Top-k takes the + # B largest survivals, so the realised tau is >= the scored one at + # an identical cost, for every rung. The gap is not constant across + # rungs, so the argmax could sit on the wrong one; scoring what runs + # removes the question. + budget = compute_verify_token_budget( + survival=survival.numpy().astype(np.float64), + num_gen_requests=int(num_gen_requests), + cost_table=self.cost_table, + min_verify_len=self.cfg.min_verify_len, + max_verify_len=self.cfg.resolved_max_verify_len, + # Without this the answer is unrealisable: any non-tier total is + # rounded up to a captured bucket, which can push a budget + # chosen to sit below a cost riser back over it. + allowed_lens=self.tiers, + ) + # The price the argmax just paid for this step, kept so the run can + # be reconciled against reality: hostStepTimeMS measures what the + # step actually cost, and a systematic gap between the two is the + # one number that catches a wrong table, a wrong lookup, and an + # engine/table config mismatch alike (same cell measured 23% apart + # across two engine configs). These keys ride along in self.stats, + # which the [final] summary prints wholesale -- no extra plumbing. + floor_tokens = int(num_gen_requests) * (self.cfg.min_verify_len + 1) + predicted = float(self.cost_table.step_times( + np.asarray([floor_tokens + int(budget)]), + int(num_gen_requests))[0]) + self.last_predicted_step_ms = predicted + self.stats["predicted_ms_sum"] = ( + self.stats.get("predicted_ms_sum", 0.0) + predicted) + self.stats["predicted_steps"] = ( + self.stats.get("predicted_steps", 0) + 1) + lens = schedule_verify_lens_topk(survival=survival, budget=budget, cfg=self.cfg).tolist() + # What the argmax bought and what top-k handed out, recorded + # SEPARATELY: the realized windows in verify_len_hist are read back + # after the bucket fit widened them, so without these two counters + # nothing says what this rank's own decision was -- and a fit-side + # defect is indistinguishable from a planner-side one. + rung = self.cfg.min_verify_len + int(budget) // max(int(num_gen_requests), 1) + rhist = self.stats.setdefault("local_rung_hist", {}) + rhist[rung] = rhist.get(rung, 0) + 1 + lhist = self.stats.setdefault("local_len_hist", {}) + for v in lens: + lhist[int(v)] = lhist.get(int(v), 0) + 1 # `reduce_across_ranks=False` means the caller is doing the cross-rank # agreement itself. That is not just an optimization here: every early diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index ddb6d94ce9bc..4b2d41baf640 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -199,6 +199,11 @@ def _request_common_handling( #: Optional observability sink, attached by the DSpark worker. Left None #: everywhere else so this costs one attribute read per request. acceptance_stats = None + #: STS calibration collection; attached alongside `acceptance_stats` and + #: None everywhere else. `sts_logits_buffer` is the worker's persistent + #: [rows, K] confidence buffer, indexed by the same py_seq_slot. + sts_recorder = None + sts_logits_buffer = None @staticmethod def _verified_len(request: LlmRequest, runtime_draft_len: int, @@ -252,6 +257,12 @@ def update_requests( runtime_draft_len = getattr(state, "runtime_draft_len", self.draft_len) caps = state.verify_caps_snapshot + if self.sts_recorder is not None: + # A fresh step: drop the cached host mirror so the first `record` + # below pulls this step's logits rather than the previous step's. + self._sts_step = getattr(self, "_sts_step", 0) + 1 + self.sts_recorder.begin_step(self._sts_step) + for req in state.requests: if req.state == LlmRequestState.GENERATION_COMPLETE: continue @@ -290,8 +301,24 @@ def update_requests( window=verified_len if cap is None else int(cap), cap_trim=(0 if cap_trim_list is None else cap_trim_list[req.py_seq_slot])) + # STS calibration pairs this request's drafted-block logits with + # how much of that block was accepted. Recorded here for the same + # reason acceptance is: this is the only place both are host-side + # and belong to the same step. The recorder pulls the logits buffer + # once per step, not once per request. + if (self.sts_recorder is not None + and self.sts_logits_buffer is not None): + self.sts_recorder.record( + slot=req.py_seq_slot, + accepted=req.py_num_accepted_draft_tokens, + device_logits=self.sts_logits_buffer) self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) + if self.sts_recorder is not None: + # Periodic flush, so a run killed mid-collection still leaves usable + # shards rather than everything living in a list that dies with it. + self.sts_recorder.end_step() + def sample_async( self, scheduled_requests: ScheduledRequests, diff --git a/tests/microbenchmarks/dspark_fit_sts.py b/tests/microbenchmarks/dspark_fit_sts.py new file mode 100644 index 000000000000..5ddd50aa9bad --- /dev/null +++ b/tests/microbenchmarks/dspark_fit_sts.py @@ -0,0 +1,307 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Fit the DSpark confidence head's per-position STS temperatures. + +WHAT IT FITS, AND WHY THAT OBJECTIVE +------------------------------------ +The planner never reads a per-position probability. It reads the *cumulative +product*:: + + survival[r][j] = prod_{i <= j} sigmoid(logit[r][i] / T[i]) + +and sums the largest of those to form tau, the numerator of the budget argmax. +So the thing that has to be calibrated is the survival, not the per-position +confidence -- a temperature vector that makes each position individually +well-calibrated can still leave the depth-5 product badly biased, because a +per-position error of x compounds to x^5. + +Hence the objective here is the expected calibration error **of the running +product** at each position, against the label "positions 0..j were ALL +accepted". Positions are fitted left to right and frozen, so every position is +optimised in the cumulative context it will actually be used in. + +This mirrors SGLang's ``benchmark/dspark_sts_fit.py`` (same 41-point log grid, +same left-to-right freeze, same cumprod-ECE objective), so a table fitted by +either tool is meaningful to the other. Deliberately so: the vectors are the one +artifact the two stacks can share without a converter. + +INPUT +----- +Shards written by ``DSparkStsRecorder`` (set ``TLLM_DSPARK_STS_COLLECT_PATH`` +on a normal DSpark run), each a torch save of:: + + {"logits": [N, K] float32, "prefix_mask": [N, K] float32} + +USAGE +----- + python3 tests/microbenchmarks/dspark_fit_sts.py \\ + --data "/path/to/sts_collect.*.pt" --out sts.json + +The emitted JSON carries both ``sts_temperatures`` and ``temperatures`` so it +loads in this repo and in SGLang unchanged. +""" + +import argparse +import glob +import json +import math +import sys +from typing import List, Tuple + +import torch + +_EPS_PROB = 1e-8 + + +def default_temperature_grid() -> torch.Tensor: + """41 points, log-spaced over [0.1, 10]. Same grid SGLang searches.""" + return torch.logspace(math.log10(0.1), math.log10(10.0), steps=41) + + +def expected_calibration_error(*, probs: torch.Tensor, targets: torch.Tensor, + num_bins: int) -> float: + """Bin by predicted probability; average |mean predicted - mean actual|. + + Weighted by bin population, so a bin holding three samples cannot dominate + one holding thirty thousand. + """ + probs = probs.reshape(-1).to(torch.float64).clamp(_EPS_PROB, 1.0 - _EPS_PROB) + targets = targets.reshape(-1).to(torch.float64) + total = probs.numel() + if total == 0: + return float("nan") + bin_index = (probs * num_bins).long().clamp_(0, num_bins - 1) + count = torch.zeros(num_bins, dtype=torch.float64) + pred_sum = torch.zeros(num_bins, dtype=torch.float64) + target_sum = torch.zeros(num_bins, dtype=torch.float64) + count.scatter_add_(0, bin_index, torch.ones_like(probs)) + pred_sum.scatter_add_(0, bin_index, probs) + target_sum.scatter_add_(0, bin_index, targets) + denom = count.clamp_min(1.0) + bin_error = (pred_sum / denom - target_sum / denom).abs() + return float((bin_error * count).sum().item() / total) + + +def fit_sts_temperatures(*, logits: torch.Tensor, prefix_mask: torch.Tensor, + grid: torch.Tensor, num_bins: int = 15) -> dict: + """Greedy left-to-right search minimising the cumprod ECE per position. + + Left to right and frozen, not joint: position j's temperature is chosen + against the survival the already-fitted positions 0..j-1 produce. A joint + search over K temperatures would be exponential in the grid, and the greedy + order is the causal one -- position j's contribution is only ever multiplied + into a product that already contains its predecessors. + """ + logits = logits.to(torch.float64) + prefix_mask = prefix_mask.to(torch.float64) + num_samples, block_size = logits.shape + if num_samples == 0: + raise ValueError("need at least one sample") + grid_values = grid.to(torch.float64).tolist() + + temperatures: List[float] = [] + ece_before: List[float] = [] + ece_after: List[float] = [] + + # Two running products: what an uncalibrated head would have said (T == 1), + # and what the fit says so far. Reporting both is the only way to see + # whether the fit bought anything. + survival_at_one = torch.ones(num_samples, dtype=torch.float64) + survival_fitted = torch.ones(num_samples, dtype=torch.float64) + + for position in range(block_size): + position_logits = logits[:, position] + position_target = prefix_mask[:, position] + + survival_at_one = survival_at_one * torch.sigmoid(position_logits) + ece_before.append( + expected_calibration_error(probs=survival_at_one, + targets=position_target, + num_bins=num_bins)) + + best_t = grid_values[0] + best_survival = survival_fitted * torch.sigmoid(position_logits / best_t) + best_ece = expected_calibration_error(probs=best_survival, + targets=position_target, + num_bins=num_bins) + for temperature in grid_values[1:]: + candidate = survival_fitted * torch.sigmoid(position_logits / temperature) + ece = expected_calibration_error(probs=candidate, + targets=position_target, + num_bins=num_bins) + if ece < best_ece: + best_ece, best_t, best_survival = ece, temperature, candidate + + temperatures.append(float(best_t)) + ece_after.append(float(best_ece)) + survival_fitted = best_survival + + return { + "temperatures": temperatures, + "ece_before": ece_before, + "ece_after": ece_after, + } + + +def load_shards(pattern: str, + source: str = "auto") -> Tuple[torch.Tensor, torch.Tensor]: + """Load shards, choosing which confidence column to fit against. + + Shards carry two columns for the same rows. ``logits`` is the worker's + persistent buffer read back at sampler time; ``logits_at_draft`` is the + snapshot the draft pass actually produced, cloned on device before anything + else could touch it. Under the overlap scheduler these are NOT the same + tensor: the next step's draft overwrites the shared buffer before the + sampler gets there. Measured on job 2562577, 26383 of 32083 rows differed + (82.2%), mean |delta| 7.35, correlation 0.14 -- i.e. the ``logits`` column + pairs confidence(t+1) with accepted(t) four times out of five. + + So ``auto`` prefers ``logits_at_draft`` whenever it is present. Fitting the + stale column does not fail loudly; it converges to a temperature vector that + calibrates noise, which is what the 0.23 residual ECE of the first table was. + """ + paths = sorted(glob.glob(pattern)) + if not paths: + raise SystemExit(f"no shards matched {pattern!r}") + have_draft = [] + for path in paths: + blob = torch.load(path, map_location="cpu") + have_draft.append("logits_at_draft" in blob) + if source == "auto": + key = "logits_at_draft" if all(have_draft) else "logits" + if not all(have_draft) and any(have_draft): + raise SystemExit( + "some shards carry 'logits_at_draft' and some do not; they came " + "from different builds. Fit them separately -- concatenating " + "aligned and stale rows produces a vector describing neither.") + else: + key = {"draft": "logits_at_draft", "late": "logits"}[source] + print(f"fitting against '{key}'" + + ("" if key == "logits_at_draft" else + " <-- WARNING: sampler-time buffer read; under the overlap " + "scheduler this is stale by one draft for most rows")) + + logits, masks, block_size = [], [], None + dropped = 0 + for path in paths: + blob = torch.load(path, map_location="cpu") + if key not in blob: + raise SystemExit(f"{path}: no '{key}' column (has {sorted(blob)})") + lg, pm = blob[key], blob["prefix_mask"] + if lg.shape != pm.shape: + raise SystemExit(f"{path}: logits {tuple(lg.shape)} != mask {tuple(pm.shape)}") + if block_size is None: + block_size = lg.shape[1] + elif lg.shape[1] != block_size: + # Silently concatenating these would fit a vector for a block size + # that never existed. + raise SystemExit( + f"{path}: block size {lg.shape[1]} != {block_size} from earlier shards") + # A row is NaN when the draft stash had no entry for that slot (the + # request did not draft this step). Those rows carry a real acceptance + # label against no confidence at all, so they must go -- sigmoid(NaN) + # would poison every temperature in the left-to-right sweep. + keep = ~torch.isnan(lg).any(dim=1) + dropped += int((~keep).sum()) + logits.append(lg[keep]) + masks.append(pm[keep]) + if dropped: + print(f"dropped {dropped} rows with no draft-time snapshot") + return torch.cat(logits, 0), torch.cat(masks, 0) + + +def main(argv=None) -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--data", required=True, + help="glob for shards from DSparkStsRecorder") + ap.add_argument("--out", required=True, help="where to write the JSON") + ap.add_argument("--num-bins", type=int, default=15) + ap.add_argument("--min-samples", type=int, default=1000, + help="refuse to fit on fewer; a temperature fitted on a " + "handful of blocks is noise with a filename") + ap.add_argument("--eval-temps", default=None, + help="Evaluate an EXISTING temperature table on these " + "shards instead of fitting: reports the per-position " + "cumprod ECE the loaded vector achieves on fresh " + "data. This is the producer+consumer joint check -- " + "collect on a serving run with the table loaded and " + "the window pinned to the block, then compare this " + "number against the table's stored ece_after. A " + "large gap means the table does not describe the " + "live head (drift, wrong file, wrong column), and " + "nothing else in the pipeline can see that.") + ap.add_argument("--source", choices=("auto", "draft", "late"), default="auto", + help="which confidence column to fit: 'draft' is the " + "snapshot the draft pass produced (correct pairing), " + "'late' is the sampler-time buffer read (stale under " + "the overlap scheduler). 'auto' prefers draft.") + args = ap.parse_args(argv) + + logits, prefix_mask = load_shards(args.data, args.source) + + if args.eval_temps: + with open(args.eval_temps, encoding="utf-8") as fh: + table = json.load(fh) + temps = table.get("sts_temperatures") or table["temperatures"] + lg = logits.to(torch.float64) + pm = prefix_mask.to(torch.float64) + survival = torch.ones(lg.shape[0], dtype=torch.float64) + print(f"evaluating {args.eval_temps} on {lg.shape[0]} fresh rows:") + stored = table.get("ece_after") or [] + for j in range(min(lg.shape[1], len(temps))): + survival = survival * torch.sigmoid(lg[:, j] / float(temps[j])) + ece = expected_calibration_error(probs=survival, + targets=pm[:, j], + num_bins=args.num_bins) + ref = f" (fit-time {stored[j]:.4f})" if j < len(stored) else "" + print(f" pos {j}: live ECE {ece:.4f}{ref}") + return 0 + n, block_size = logits.shape + print(f"loaded {n} samples, block_size={block_size}") + if n < args.min_samples: + raise SystemExit( + f"only {n} samples (< --min-samples {args.min_samples}). Collect " + f"more before fitting: the deepest positions are the sparsest and " + f"the ones the planner is most sensitive to.") + + # An all-ones column means every block was fully accepted at that depth, so + # the ECE is defined but carries no information to fit against -- say so + # rather than emitting a confident-looking temperature. + per_position_rate = prefix_mask.mean(dim=0) + print("prefix-accept rate by position:", + [round(float(v), 4) for v in per_position_rate]) + degenerate = [j for j, v in enumerate(per_position_rate) if v in (0.0, 1.0)] + if degenerate: + print(f"WARNING: positions {degenerate} are degenerate (rate 0 or 1); " + f"their temperature is unconstrained by this data") + + result = fit_sts_temperatures(logits=logits, prefix_mask=prefix_mask, + grid=default_temperature_grid(), + num_bins=args.num_bins) + + payload = { + "sts_temperatures": result["temperatures"], + "temperatures": result["temperatures"], + "ece_before": result["ece_before"], + "ece_after": result["ece_after"], + "dataset": args.data, + "num_samples": int(n), + } + with open(args.out, "w", encoding="utf-8") as fh: + json.dump(payload, fh, indent=2) + + print(f"\nwrote {args.out}") + print(f" temperatures : {[round(t, 4) for t in result['temperatures']]}") + print(f" ECE before : {[round(e, 5) for e in result['ece_before']]}") + print(f" ECE after : {[round(e, 5) for e in result['ece_after']]}") + improved = sum(1 for a, b in zip(result["ece_after"], result["ece_before"]) if a < b) + print(f" improved at {improved}/{block_size} positions") + if improved == 0: + print(" NOTE: no position improved -- the head was already calibrated " + "on this data, or the data does not cover the relevant range.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py new file mode 100644 index 000000000000..bde9c5e320ce --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""``trim_ratio`` must not report a saving the run did not make -- or a loss. + +Four runs reported a NEGATIVE ``trim_ratio`` (jobs 2560992 -0.001, 2561519 +-0.0438, 2563581 -0.0082), which reads as "the ragged path delivered more tokens +than a no-trim run". It did not. ``delivered_tokens`` books the fitted bucket, +which is ``padded_bs * (tier + 1)`` where ``padded_bs`` is the CUDA-graph-rounded +maximum row count *across attention-DP ranks*; ``ceiling_tokens`` booked +``num_gen_requests * (1 + block)`` over this rank's real, unpadded requests. Two +different row bases in one ratio. + +The gate reads that ratio: ``assert_ragged_active`` fires on ``trim_ratio <= 0`` +with the message "delivered as many tokens as a no-trim run, so nothing was +saved" -- the opposite of what a negative value means, so the defect also +misdiagnoses itself. +""" + +import pytest + +from tensorrt_llm._torch.speculative.dspark_observability import ( + DSparkRaggedStats, RaggedVerifyMode) + + +def _stats(block_size: int = 5) -> DSparkRaggedStats: + return DSparkRaggedStats(mode=RaggedVerifyMode.COMPACT, + max_draft_len=block_size) + + +def test_padding_rows_do_not_make_the_trim_look_negative(): + """A top-tier step on padded rows saved nothing -- and lost nothing.""" + stats = _stats(block_size=5) + # 3 real requests, graph padded to 4 rows, every window at the full block. + # delivered = 4 * 6 = 24. Against a ceiling of 3 * 6 = 18 that is -0.333. + stats.record_step(num_gen_requests=3, verify_lens=[5, 5, 5], + bucket=4 * 6, padded_bs=4) + summary = stats.summary() + assert summary["trim_ratio"] == pytest.approx(0.0), ( + f"a step that verified the full block on every row trimmed nothing, so " + f"the ratio is 0 -- got {summary['trim_ratio']}, which came from " + f"comparing {summary['delivered_tokens']} padded-row tokens against a " + f"ceiling counted over {3} real requests") + + +def test_narrow_rank_is_not_charged_for_its_widest_peer(): + """Under attention-DP the padding comes from the widest peer, not locally. + + This is the larger of the two terms: a rank holding 2 requests while its + widest peer holds 7 submits the peer's rows. Booked against a two-request + ceiling that is a ratio of -3.0, i.e. the metric claims the feature made the + step four times more expensive. + """ + stats = _stats(block_size=5) + stats.record_step(num_gen_requests=2, verify_lens=[5, 5], + bucket=8 * 6, padded_bs=8) + assert stats.summary()["trim_ratio"] == pytest.approx(0.0) + + +def test_a_real_trim_still_reads_as_a_saving(): + """The fix must not flatten genuine trimming to zero.""" + stats = _stats(block_size=5) + # 4 rows, no padding, every window trimmed to 2 -> bucket 4*3 = 12 against a + # ceiling of 4*6 = 24. + stats.record_step(num_gen_requests=4, verify_lens=[2, 2, 2, 2], + bucket=4 * 3, padded_bs=4) + assert stats.summary()["trim_ratio"] == pytest.approx(0.5) + + +def test_paths_without_a_bucket_keep_the_local_row_base(): + """Only the bucket path rebases; widening the others understates the trim. + + The no-window path and the cap-accept path each book ``delivered`` over the + real rows, so pulling their ceiling up to the padded count would invent a + saving out of padding that those paths never submitted. + """ + no_window = _stats(block_size=5) + no_window.record_step(num_gen_requests=3, verify_lens=None, padded_bs=8) + assert no_window.summary()["trim_ratio"] == pytest.approx(0.0) + + explicit = _stats(block_size=5) + # cap-accept: windows are computed but the full block is submitted anyway. + explicit.record_step(num_gen_requests=3, verify_lens=[2, 3, 1], + bucket=8 * 6, padded_bs=8, delivered=3 * 6) + assert explicit.summary()["trim_ratio"] == pytest.approx(0.0) + + +def test_padded_bs_reaches_the_histogram(): + """``padded_bs_hist`` was empty in every summary ever logged. + + The executor passed ``padded_bs=None`` unconditionally, so the padded row + count -- the thing that explains the negative ratio -- was never recorded + anywhere. An always-empty histogram is indistinguishable from "padding never + happened", which is how this survived. + """ + stats = _stats(block_size=5) + stats.record_step(num_gen_requests=3, verify_lens=[5, 5, 5], + bucket=4 * 6, padded_bs=4) + assert stats.summary()["padded_bs_hist"] == {4: 1} diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py new file mode 100644 index 000000000000..8817f6f9a170 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The two reconciliation guards: engine fingerprint and price prediction. + +Both exist because this pipeline's characteristic failure is a PREDICTION +recorded as a MEASUREMENT with nothing positioned to contradict it: a cost +table measured on the wrong MoE kernel almost planned MegaMoE steps, the same +nominal cell measured 23% apart across two engine configs, and a planner +argmax'd against a table whose labels described shapes that never ran. The +fingerprint refuses the wrong table at load; the recorded prediction gives +every run a number that reality (hostStepTimeMS) can contradict. +""" + +import numpy as np +import pytest + +from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, check_table_fingerprint) +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig +from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + +LIVE = {"tp": 8, "ep": 8, "attention_dp": True, "block": 5, + "max_batch_size": 256} + + +def _payload(engine): + return {"token_counts": [512, 1536], "step_time_ms": [60.0, 150.0], + "_meta": {"engine": engine} if engine is not None else {}} + + +def test_matching_fingerprint_loads_silently(): + check_table_fingerprint(payload=_payload(dict(LIVE)), live=dict(LIVE)) + + +def test_mismatched_engine_is_refused(): + """The 23%-apart case: same cell, different max_batch_size.""" + wrong = dict(LIVE, max_batch_size=64) + with pytest.raises(ValueError, match="different engine configuration"): + check_table_fingerprint(payload=_payload(wrong), live=dict(LIVE)) + + +def test_comparison_is_case_insensitive(): + """String facts must not fail on spelling case (MEGAMOE vs MegaMoe).""" + fp = dict(LIVE, moe_backend="megamoe_cutedsl") + live = dict(LIVE, moe_backend="MEGAMOE_CUTEDSL") + check_table_fingerprint(payload=_payload(fp), live=live) + + +def test_extra_table_facts_do_not_block_loading(): + """Facts the consumer cannot see (image hash, geometry) are informational. + + They must neither raise nor be silently swallowed -- the function logs + them -- but loading proceeds: refusing on unverifiable keys would make + every fingerprint addition a breaking change. + """ + fp = dict(LIVE, image="faf2c60935", geometry="constant_block") + check_table_fingerprint(payload=_payload(fp), live=dict(LIVE)) + + +def test_unfingerprinted_table_still_loads(): + """Backward compatibility: old tables warn, they do not break.""" + check_table_fingerprint(payload=_payload(None), live=dict(LIVE)) + + +def test_planner_records_the_price_it_paid(): + """Every budget decision must leave behind the step time it assumed. + + Without this the run has no number reality can contradict: the argmax + consults the table and the table's opinion evaporates. With it, the + [final] planner block carries predicted_ms_sum / predicted_steps, and a + systematic gap against measured hostStepTimeMS is the one signal that + catches a wrong table, a wrong lookup, and an engine mismatch alike. + """ + table = SpsCostTable(token_counts=(0, 512, 768, 1536), + step_time_ms=(5.0, 68.4, 80.2, 150.5), + fixed_overhead_ms=1.0) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + planner = DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=[1, 2, 5]) + + bs = 8 + rng = np.random.default_rng(11) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() + + import torch + # Only the snapshot is faked; calibration stays the constructor's default + # sigmoid (setting it to None post-construction would bypass the + # `apply_calibration or torch.sigmoid` fallback and crash the decide path). + planner._gather_rows = lambda **_: torch.tensor(survival, + dtype=torch.float32) + + lens = planner.decide_verify_lens(num_gen_requests=bs, + reduce_across_ranks=False) + assert lens is not None + + steps = planner.stats.get("predicted_steps", 0) + total = planner.stats.get("predicted_ms_sum", 0.0) + assert steps == 1 and total > 0.0 + + # The recorded price must be the table's own answer for the tokens the + # decision implies -- anything else and the reconciliation compares two + # unrelated numbers. + budget = sum(lens) - bs * cfg.min_verify_len + tokens = bs * (cfg.min_verify_len + 1) + budget + expected = float(table.step_times(np.asarray([tokens]), bs)[0]) + assert planner.last_predicted_step_ms == pytest.approx(expected) + assert total == pytest.approx(expected) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py new file mode 100644 index 000000000000..85517f5c4c30 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The two snapshot instruments: stamp-lag histogram and neutral-row count. + +They exist because the planner's inputs have never been measured: the argmax +eats a slot-indexed snapshot staged one step earlier, and nothing recorded how +old each gathered row was (relay staleness) or how many rows were still the ++30 neutral fill (unknowns entering as certainties). The live anomaly they +instrument: 445/476 full-batch steps chose the full block while calibrated +pooled survivals favour rung-2 by 36% -- so the per-step snapshot content must +differ from the pooled shards, and these counters say in which way. + +Both are measurement-only. A failure inside them must never kill a step. +""" + +import numpy as np +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig +from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + + +def _planner(): + table = SpsCostTable(token_counts=(0, 512, 768, 1536), + step_time_ms=(5.0, 68.4, 80.2, 150.5), + fixed_overhead_ms=1.0) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + return DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=[1, 2, 5]) + + +def test_neutral_rows_are_counted_and_real_rows_are_not(): + planner = _planner() + selected = torch.zeros((8, 5)) + selected[3] = 30.0 # the neutral fill: sigmoid ~ 1.0 at any temperature + planner._note_snapshot_stats(selected, rows=None) + assert planner.stats["snap_rows"] == 8 + assert planner.stats["snap_neutral_rows"] == 1 + + +def test_lag_histogram_reads_staged_seq_minus_stamp(): + planner = _planner() + planner._host_stamps = torch.tensor([7, 6, 5, 0], dtype=torch.int32) + planner._staged_seq = 7 + selected = torch.zeros((4, 5)) + planner._note_snapshot_stats(selected, rows=[0, 1, 2, 3]) + hist = planner.stats["stamp_lag_hist"] + # Stamp 0 is "never drafted": lag 7 in truth, and it must land visibly at + # the histogram's far end rather than vanish -- these are exactly the rows + # the neutral counter should agree about. + assert hist == {0: 1, 1: 1, 2: 1, 7: 1} + + +def test_large_and_negative_lags_clamp_instead_of_exploding_the_dict(): + planner = _planner() + planner._host_stamps = torch.tensor([0, 500], dtype=torch.int32) + planner._staged_seq = 400 + planner._note_snapshot_stats(torch.zeros((2, 5)), rows=[0, 1]) + assert planner.stats["stamp_lag_hist"] == {8: 1, -2: 1} + + +def test_counts_accumulate_across_steps(): + planner = _planner() + planner._host_stamps = torch.tensor([3, 3], dtype=torch.int32) + planner._staged_seq = 4 + for _ in range(3): + planner._note_snapshot_stats(torch.full((2, 5), 30.0), rows=[0, 1]) + assert planner.stats["snap_rows"] == 6 + assert planner.stats["snap_neutral_rows"] == 6 + assert planner.stats["stamp_lag_hist"] == {1: 6} + + +def test_instrument_failure_is_counted_not_raised(): + planner = _planner() + planner._host_stamps = torch.tensor([1], dtype=torch.int32) + planner._staged_seq = 1 + # Rows that cannot index the stamp buffer must clamp, and anything worse + # must be swallowed into the error counter: instruments never kill steps. + planner._note_snapshot_stats(torch.zeros((1, 5)), rows=[10_000]) + assert planner.stats.get("snap_stats_errors", 0) == 0 + planner._note_snapshot_stats(object(), rows=None) # type: ignore[arg-type] + assert planner.stats["snap_stats_errors"] == 1 + + +def test_decide_path_feeds_the_instruments(): + """The counters must fill from the real decide path, not only direct calls.""" + planner = _planner() + bs = 8 + rng = np.random.default_rng(3) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() + planner._gather_rows = lambda **_: torch.tensor(survival, + dtype=torch.float32) + lens = planner.decide_verify_lens(num_gen_requests=bs, + reduce_across_ranks=False) + assert lens is not None + assert planner.stats["snap_rows"] == bs + assert planner.stats["snap_neutral_rows"] == 0 + # The local-decision instruments must fill on the same path: one rung per + # decision, one window per request, and the two must agree -- the windows + # are the top-k realisation of exactly that rung's budget. + rung_hist = planner.stats["local_rung_hist"] + assert sum(rung_hist.values()) == 1 + (rung,) = rung_hist + assert rung in (1, 2, 5) + len_hist = planner.stats["local_len_hist"] + assert sum(len_hist.values()) == bs + budget = bs * (rung - planner.cfg.min_verify_len) + extra = sum((k - planner.cfg.min_verify_len) * v for k, v in len_hist.items()) + assert extra <= budget diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py new file mode 100644 index 000000000000..6eec0a98e46a --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A collection run must refuse the regimes that fit the scheduler, not the head. + +Both regimes below produce shards that look perfectly well-formed. The run +completes, the row counts are plausible, the fitter converges, and the emitted +temperatures are wrong -- there is no downstream check that can tell. That is +why these are construction-time failures rather than warnings. + +Measured on this branch: collecting while the planner trimmed gave +``ece_before = [.244 .319 .361 .389 .394]``; pinning the window to the full +block, same model and same fitter, gave ``[.157 .208 .270 .321 .344]``. The +difference is entirely the censored label. + +SGLang guards the first regime next to its own collector +(``dspark_planner.py:100-108``) and names the second in the sibling +``ConfidenceMetricsProbe`` -- "padded verify rows corrupt the per-position +prefix label" (``dspark_observability.py:688-697``). +""" + +import contextlib +import os +import tempfile + +import pytest + +from tensorrt_llm._torch.speculative.dspark_sts import (STS_COLLECT_ENV, + make_recorder_from_env) + +_PIN = "TLLM_DSPARK_FORCE_VERIFY_LEN" + + +@contextlib.contextmanager +def collecting(pin=None): + """Set the collection env for the block, restore it after. + + A context manager rather than a pytest fixture: these run inside the + container, which ships no pytest, under a minimal shim that does not + implement fixtures or tmp_path. + """ + saved = {k: os.environ.get(k) for k in (STS_COLLECT_ENV, _PIN)} + try: + with tempfile.TemporaryDirectory() as tmp: + os.environ[STS_COLLECT_ENV] = os.path.join(tmp, "collect") + if pin is None: + os.environ.pop(_PIN, None) + else: + os.environ[_PIN] = str(pin) + yield + finally: + for k, v in saved.items(): + os.environ.pop(k, None) if v is None else os.environ.__setitem__(k, v) + + +def test_collection_off_builds_nothing(): + """The serving default must be untouched by any of this.""" + os.environ.pop(STS_COLLECT_ENV, None) + assert make_recorder_from_env(block_size=5) is None + assert make_recorder_from_env(block_size=5, has_cost_table=True, + ragged_mode="compact") is None + + +def test_a_loaded_cost_table_without_a_pin_is_refused(): + """A table means the planner trims, and trimming censors the label.""" + with collecting(): + with pytest.raises(ValueError, match="censors the acceptance label"): + make_recorder_from_env(block_size=5, has_cost_table=True) + + +def test_a_pinned_window_makes_the_table_harmless(): + """Pinning to the full block makes the label uncensored by construction.""" + with collecting(pin=5): + rec = make_recorder_from_env(block_size=5, has_cost_table=True) + assert rec is not None and rec.block_size == 5 + + +def test_no_table_needs_no_pin(): + """Without a table the planner cannot trim, so nothing censors anything.""" + with collecting(): + assert make_recorder_from_env(block_size=5, + has_cost_table=False) is not None + + +def test_compact_mode_is_refused(): + """Padding rows contribute prefix labels that measure nothing. + + Refused even with the pin set: the pin fixes the *window*, while this is + about rows that belong to no real request at all. + """ + with collecting(pin=5): + with pytest.raises(ValueError, match="Padded verify rows"): + make_recorder_from_env(block_size=5, ragged_mode="compact") + + +def test_static_mode_is_allowed(): + with collecting(): + assert make_recorder_from_env(block_size=5, + ragged_mode="static") is not None + + +class _Head: + pass + + +class _Bare: + """DSparkDraftModel layout: stages under .mtp_layers.""" + + def __init__(self, head): + stage = type("Stage", (), {})() + if head is not None: + stage.confidence_head = head + self.mtp_layers = [object(), object(), stage] + + +class _Wrapper: + """DSparkForCausalLM layout: bare model under .dspark_model. + + This is the layout the chained getattr missed: the wrapper has no + mtp_layers of its own, so `getattr(model, "mtp_layers", [None])[-1]` + resolved to None and silently disabled calibration for every + confidence-scheduled run of the first campaign. + """ + + def __init__(self, head): + self.dspark_model = _Bare(head) + + +def test_head_resolves_on_the_bare_model(): + from tensorrt_llm._torch.speculative.dspark_sts import resolve_confidence_head + h = _Head() + assert resolve_confidence_head(_Bare(h)) is h + + +def test_head_resolves_through_the_forcausallm_wrapper(): + from tensorrt_llm._torch.speculative.dspark_sts import resolve_confidence_head + h = _Head() + assert resolve_confidence_head(_Wrapper(h)) is h + + +def test_absent_head_is_none_not_an_error(): + from tensorrt_llm._torch.speculative.dspark_sts import resolve_confidence_head + assert resolve_confidence_head(_Bare(None)) is None + assert resolve_confidence_head(_Wrapper(None)) is None + assert resolve_confidence_head(object()) is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py new file mode 100644 index 000000000000..2c0a83f58351 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py @@ -0,0 +1,130 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The budget must be scored with the tau the step actually collects. + +``decide_verify_lens`` picks a rung, converts it to ``B = n*(L - min)`` tokens, +and hands B to ``schedule_verify_lens_topk``, which spends it on the B largest +survivals in the batch. But the rung was chosen by +``budget_argmax_over_uniform_lens``, which scores rung L as if every request took +columns ``[min, L)``. Top-k is by construction the largest sum over all size-B +subsets, so the realised tau is >= the scored one at an identical token cost -- +and the gap is not the same at every rung, so the argmax can land on the wrong +one. + +The fix is not to let the planner choose an arbitrary budget: the executor can +only run ``bs*(t+1)`` tokens for a captured tier, and any other total rounds UP, +which would push a budget deliberately chosen below a cost riser back over it. +So the ladder stays the grid and only the numerator changes. +""" + +import numpy as np +import pytest + +from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, budget_argmax_over_uniform_lens, compute_verify_token_budget) + + +def _table() -> SpsCostTable: + """A staircase with genuine risers, so trimming is worth something. + + Breakpoints are spaced so that neither tier is trivially optimal: a flat + table degenerates to verify-all and would make every assertion here vacuous. + """ + token_counts = tuple(range(0, 400, 16)) + step_time_ms = tuple(4.0 + 0.6 * (tok // 96) + 0.004 * tok + for tok in token_counts) + return SpsCostTable(token_counts=token_counts, step_time_ms=step_time_ms, + fixed_overhead_ms=1.0) + + +@pytest.mark.parametrize("tiers", [[1, 2, 5], [1, 3, 5], [1, 5]]) +def test_restricted_answer_is_always_realisable(tiers): + """Every returned budget corresponds to a rung the executor captured.""" + rng = np.random.default_rng(20260803) + table = _table() + for _ in range(200): + bs = int(rng.integers(2, 17)) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] + n = compute_verify_token_budget( + survival=survival, num_gen_requests=bs, cost_table=table, + min_verify_len=1, allowed_lens=tiers) + assert n % bs == 0, ( + f"budget {n} for {bs} requests is not n*(t-min) for any tier") + assert (n // bs) + 1 in tiers, ( + f"budget {n} implies tier {(n // bs) + 1}, not in {tiers}") + + +def test_agrees_with_the_uniform_scorer_at_a_degenerate_batch(): + """With one row the two taus coincide, so the rungs must match. + + A single request makes the uniform allocation and the top-k allocation the + same set, which is the one case where any disagreement would be a plain bug + rather than the intended improvement. + """ + table = _table() + rng = np.random.default_rng(7) + tiers = [1, 2, 5] + for _ in range(100): + survival = np.sort(rng.random((1, 5)), axis=1)[:, ::-1] + n = compute_verify_token_budget( + survival=survival, num_gen_requests=1, cost_table=table, + min_verify_len=1, allowed_lens=tiers) + uniform = budget_argmax_over_uniform_lens( + survival=survival, num_gen_requests=1, cost_table=table, + allowed_lens=tiers, min_verify_len=1) + assert (n // 1) + 1 == uniform + + +def test_restricted_never_scores_worse_than_the_uniform_choice(): + """Same grid, better numerator: the chosen rung's realised theta wins. + + Both routes evaluate the same discrete cost points, so the restricted form + cannot be worse under the tau that the step actually collects. This is the + property that makes the swap safe -- if it ever failed, the change would be + trading a modelling error for a real one. + """ + rng = np.random.default_rng(31337) + table = _table() + tiers = [1, 2, 5] + for _ in range(300): + bs = int(rng.integers(2, 17)) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] + + def realised_theta(tier: int) -> float: + budget = bs * (tier - 1) + cand = np.sort(survival[:, 1:].reshape(-1))[::-1] + tau = float(bs) + float(survival[:, :1].sum()) + float( + cand[:budget].sum()) + tokens = np.array([bs * (tier + 1)]) + return tau / float(table.step_times(tokens, bs)[0]) + + n = compute_verify_token_budget( + survival=survival, num_gen_requests=bs, cost_table=table, + min_verify_len=1, allowed_lens=tiers) + chosen = (n // bs) + 1 + uniform = budget_argmax_over_uniform_lens( + survival=survival, num_gen_requests=bs, cost_table=table, + allowed_lens=tiers, min_verify_len=1) + assert realised_theta(chosen) >= realised_theta(uniform) - 1e-12, ( + f"restricted picked tier {chosen} whose realised theta is below " + f"the uniform scorer's tier {uniform}") + + +def test_unrestricted_form_is_unchanged(): + """Omitting allowed_lens must still answer the unconstrained question. + + The existing tests use that form as the reference, and it is still the right + thing to ask when sizing a ladder rather than executing one. + """ + table = _table() + rng = np.random.default_rng(99) + survival = np.sort(rng.random((6, 5)), axis=1)[:, ::-1] + free = compute_verify_token_budget( + survival=survival, num_gen_requests=6, cost_table=table, + min_verify_len=1) + assert 0 <= free <= 6 * 4 + tied = compute_verify_token_budget( + survival=survival, num_gen_requests=6, cost_table=table, + min_verify_len=1, allowed_lens=[1, 2, 5]) + # The restricted answer is one of the grid points; the free one need not be. + assert tied in {0, 6, 24} From 50675cb181751a00ee17e38eca9c887e81faa461 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 4 Aug 2026 20:08:49 -0700 Subject: [PATCH 088/177] [None][fix] Make the ragged fit's consistency checks able to fail Three checks around the ragged bucket fit could not do their job. The padded row count survived every early return, so record_step booked a value from the last successful fit on every fallback and cap-accept step -- polluting exactly the histogram that exists to explain them. It is cleared at fit entry, next to the agreed bucket. The layout assert compared the batch-derived token axis against a re-sum of the same per-request windows it was derived from: a tautology that could never fire, however far the fit had drifted. It now compares the two INDEPENDENT derivations -- the batch walk against the fit's allgathered arithmetic. Documentation caught up with the code in three places that still described the floor lookup and the old key derivation, and the feature doc's claim that the window floor defaults to zero is corrected to one (SGLang allows zero; this port does not). Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../features/dspark-confidence-scheduling.md | 521 ++++++++++++++++++ .../_torch/pyexecutor/cuda_graph_runner.py | 2 + .../_torch/pyexecutor/model_engine.py | 60 +- 3 files changed, 562 insertions(+), 21 deletions(-) create mode 100644 docs/source/features/dspark-confidence-scheduling.md diff --git a/docs/source/features/dspark-confidence-scheduling.md b/docs/source/features/dspark-confidence-scheduling.md new file mode 100644 index 000000000000..c7ee361ab823 --- /dev/null +++ b/docs/source/features/dspark-confidence-scheduling.md @@ -0,0 +1,521 @@ +# DSpark 置信度调度验证（Confidence-Scheduled Verification） + +DSpark 每个 decode step 都会 draft 出一整个 block，并把它全部送给 target 验证。 +它的 draft 网络本来就带了一个 **confidence head**，为每个 draft 位置打分，预测该位置 +能否活过验证。置信度调度把这个 head 接进了验证调度器：block 依然完整 draft，但**真正 +送去验证的位置数**由这些分数决定。 + +本文讲设计。用户可配置项见[配置](#配置)；DSpark 在各投机算法中的位置见 +[Speculative Decoding](./speculative-decoding.md)。 + +## 目录 + +- [问题与核心洞察](#问题与核心洞察) +- [三条设计不变量](#三条设计不变量) +- [模块分层](#模块分层) +- [一个 decode step 的数据流](#一个-decode-step-的数据流) +- [relay 滞后](#relay-滞后) +- [算法](#算法) +- [一个完整的数值例子](#一个完整的数值例子) +- [Ragged 验证](#ragged-验证) +- [CUDA graph 处理](#cuda-graph-处理) +- [跨 rank 一致性](#跨-rank-一致性) +- [退化路径](#退化路径) +- [配置](#配置) +- [性能天花板与已知局限](#性能天花板与已知局限) + +## 问题与核心洞察 + +基线行为： + +```text +draft: 一次并行 backbone forward -> [t1 t2 t3 t4 t5 t6 t7] (K = 7，固定成本) +target: 验证全部 7 个 + 1 个 bonus = 8 个 token 位置 +``` + +如果某个请求的 draft 在位置 3 就已经断了，那么花在位置 4-7 上的验证算力全是浪费。 +confidence head 预测的是每个位置的**条件**接受概率 `P(accept_k | accept_1..k-1)`， +因此前缀存活概率就是累积乘积： + +```text +survival[r][j] = prod_{i <= j} conf[r][i] # 位置 j 被"走到 且 被接受"的概率 +``` + +由于 `conf` 是 sigmoid 输出、落在 `(0, 1)` 内，**survival 沿位置严格非增**。 +这条单调性是后面整个分配算法的地基。 + +## 三条设计不变量 + +下面所有取舍都可以从这三条推出来： + +| # | 不变量 | 推论 | +|---|--------|------| +| **I1** | block 永远完整 draft | draft 是一次并行 forward，砍长度省不了钱。只有送给 target 的 token 数可裁 | +| **I2** | 接受规则完全不动 | 调度只决定**验多少**，绝不决定**接不接受**，因此**输出分布不变** | +| **I3** | 决策在 host 上做，且必须落在已捕获的 graph 上 | CUDA graph key 由 host 值构成；未捕获的 draft 长度不会报错，而是静默掉进 eager 执行 | + +**I2 是最重要的性质**：这不是精度/吞吐的 tradeoff，而是纯粹的算力回收。所以这个特性的 +精度测试断言的是与 baseline **完全相同**的分数 —— 任何偏差都是不变量被破坏，而不是调参结果。 + +## 模块分层 + +```text +配置层 + llmapi/llm_args.py :: DSparkDecodingConfig + +DEVICE 侧（在 target 的 captured CUDA graph 内） + _torch/models/dspark/heads.py :: DSparkConfidenceHead + forward() [*, K, hidden] -> [*, K] 原始 logits（fp32 matmul） + apply_sts() 原始 logits -> 标定后的概率（逐位置温度） + _torch/speculative/dspark.py :: DSparkWorker._draft_gen_block_batched + _confidence_logits[slots] = conf # slot 索引，原地写 + | + | 非阻塞 D2H，滞后（见下文） + v +HOST 侧（决策，完全在 capture 区之外） + _torch/speculative/dspark_verify.py :: DSparkVerifyPlanner # 协调器、退化守门 + | + +-- dspark_schedule.py survival（cumprod）+ 全局 top-k 分配 + +-- dspark_planner.py SpsCostTable + Theta = tau / T 的 argmax + tier 推导 + _torch/pyexecutor/py_executor.py :: _dspark_confidence_draft_len # 每步 hook + +批次成形层 + _torch/pyexecutor/model_engine.py ragged 输入布局、graph 形状拟合 + _torch/pyexecutor/cuda_graph_runner.py graph key 上的 token 轴 + _torch/speculative/dspark_ragged.py qo_indptr 打包、bucket 对齐、scatter + _torch/attention_backend/sparse/dsa.py indexer 元数据按 token 展开 +``` + +## 一个 decode step 的数据流 + +```text + HOST | DEVICE（captured graph） +-------------------------------------------------------------------------------- + step t | + _dspark_confidence_draft_len() | + 1. 读取滞后的快照 | + （cudaEvent.query()，非阻塞） | + 2. 按 slot 查出每个请求对应的行 | + 3. apply_sts -> cumprod -> survival | + 4. argmax Theta -> tiers 中的 draft_len | + 5. 跨 rank allgather(max) | + | + _prepare_tp_inputs() | + 铺 input_ids / position_ids | + _attach_ragged_verify_layout() | + | + graph_key = (bs, draft_len, ..., [bucket]) | + ------------- replay --------------------> | target forward + | sample_and_accept_draft_tokens + | draft forward（完整 block K） + | confidence head + | _confidence_logits[slots] = conf (*) + 6. stage_confidence(buffer, generations) | + 轮转 relay ring，然后发起 | + 非阻塞 D2H <---------------+ +-------------------------------------------------------------------------------- +``` + +## relay 滞后 + +调度决策跑在一个**滞后**的快照上，而且这个滞后是**钉死的常量**，不是"引擎给多少算多少"。 +DSpark 规定了两步回溯（two-step-back）的 relay；SGLang 也把这个值钉死，并在其 overlap +scheduler 会给出不同值时用 ring buffer 补齐。 + +钉死很重要，因为 STS 标定表和 SPS 成本表是在**某一个特定滞后**下拟合出来的。如果有效滞后 +会随一个无关的引擎开关漂移，那么一张 profile 好的表就会静默失效。 + +TensorRT LLM 这两步是怎么来的 —— 调度 hook 跑在 forward **之前**： + +```text +iteration t: [hook(t)] -> [prepare] -> [forward(t)] + | 读取 hook(t-1) staged 的快照 + | 发起 D2H，流序在 forward(t-1) 之后， + | 所以抓到的是 forward(t-1) 的 confidence + v + hook(t) 读到的是 forward(t-2) 的 confidence ==> 引擎滞后 = 2 +``` + +无论 overlap scheduler 开不开都成立：overlap 改变的是拷贝**何时落地**，而不是它抓的是 +哪一次 forward。没落地的拷贝会退化成"验全块"，而不是退化成一个陈旧的猜测。 + +因此 `DSparkVerifyPlanner` 把"引擎给的滞后"和"想要的总滞后"分开接收，差值用一个 host 侧 +carry ring 补上： + +```text +carry_steps = confidence_relay_lag_steps - ENGINE_CONFIDENCE_RELAY_LAG +``` + +默认值下差值为 0，ring 完全 inert。未填充的 ring 层初始化成**中性 confidence** 加上一个 +**不可能的 generation**，所以启动后的前 `carry_steps` 步不可能在未初始化内存上做裁剪。 + +### slot 索引与 generation 守卫 + +快照按 **worker slot** 索引，绝不按 batch 位置。请求会在"写入分数的那一步"和"读取分数的 +那一步"之间进出，所以位置 `i` 经常已经是另一个请求了。 + +slot 还会被回收，这带来两种不同的陈旧风险，由两个互补机制处理： + +| 风险 | 机制 | +|------|------| +| device 上那一行还留着上一个占用者的分数 | `_assign_slot` 在把行交给新请求时重置为中性值 | +| carry ring 里存着回收之前的**host 副本**，device 侧的重置够不着 | 每行带一个 `_slot_generation` 标签，与 confidence 一起 stage；和当前 generation 对不上就把该行 survival 强制为 1.0 | + +这两者是**互补的，不是二选一**。generation 标签是在 stage 时于 host 侧快照的，而 +confidence 是在 forward 时于 device 上写的 —— 所以正是那次重置让标签说的是真话：只 bump +不重置，就会把上一个占用者的分数打上新占用者的标签。反过来，只重置也清不掉 ring 里已经 +存在的 host 副本。 + +两种风险都收敛到同一个安全方向：**survival = 1.0，即验证整个 block**，也就是启用该特性 +之前的行为。 + +## 算法 + +### 第 1 步 —— survival + +```python +survival = torch.cumprod(sigmoid(logits / sts_temperatures), dim=1) +``` + +标定（STS）之所以重要，是因为调度器消费的是**累积乘积**：逐位置的偏差会沿 block +**几何级放大**，未标定的 head 会系统性地高估前缀存活率。 + +### 第 2 步 —— 成本模型 + +本实现遵循已公开的形式： + +```text +T(bs, K) = bias + alpha(bs) + theta(M) + ---- --------- -------- + 固定 批次相关 verify-token 成本 + （draft pass、 ^ 唯一可裁剪的一项 + 权重搬运） +``` + +`SpsCostTable` 按总验证 token 数 `M` 做**阶梯**查表，绝不插值： + +```text +step_time + | +-------- <- riser：新的 kernel wave 起来了 + 9 | | + | +-----+ + 4 | | <- shelf：这一段上加 token 是真的免费 + 2 |--+ + +--+-----+-----+----> M（总 token 数） + 8 16 32 +``` + +**token 口径**：一个请求验证 `L` 个 draft 位置，实际提交 `L + 1` 个 token —— 它手上那个 +bonus token 加上它的 draft。所以 `M = bs * (L + 1)`。这个换算只在一个地方计算 +（`dspark_planner.total_verify_tokens`），保证 planner 和 tier 推导不会漂移。 + +由于 `Theta = tau / T` 是**比值**，不可裁剪的那几项**不会约掉**。低估它们会让每个验证 +token 显得比实际更贵，planner 就会过度裁剪。因此这些项是**给进来的，不是猜的**： +`fixed_overhead_ms` 承载 `bias`，可选的 `batch_sizes` / `batch_overhead_ms` 阶梯承载 +`alpha(bs)`。 + +### 第 3 步 —— 目标函数与全局 argmax + +```text +tau(L) = num_gen_requests + sum_r sum_{j L +``` + +### 第 4 步 —— shelf 右边缘性质 + +> 在一个 cost shelf 内，step time 恒定，而 `tau` 随每个新增 token 严格增， +> 所以 `Theta = tau / cost` 在 shelf 上严格递增。因此最优点**永远落在 shelf 的右边缘**， +> 绝不在其内部。 + +于是由各 shelf 右边缘构成的 tier 集合，**在推导时所用的那个 batch size 下**精确包含连续最优解： + +```text +right_edge(breakpoint T) = (T - 1) // bs - 1 # length 空间 +``` + +这个**零损失性质跨 batch size 不成立** —— cost shelf 活在 token 空间，所以 shelf 右边缘 +在 length 空间是 `bs` 的函数。一个部署会捕获很多 batch size 但只能捕获一套 ladder，因此 +应当在**稳态众数 batch size** 上推导，并**实测**其他 batch size 上的残差损失，而不是假装它不存在。 + +### 第 5 步 —— 逐请求分配（ragged 模式） + +预算是 `num_gen * (L* - min_verify_len)` —— 和 uniform 决策**花掉的 token 总量相同**， +只是重新分配。 + +`min_verify_len` 默认为 **0**：整个 block 的每个位置都参与竞争，draft 已经死掉的请求会拿到 +**0 个** draft 位置。这不是饿死 —— bonus token 不在 confidence 矩阵里（它不是 draft 出来的、 +必然被接受），所以该请求那一步照样出 1 个 token，只是退化成普通自回归解码。也不存在 +"不验证 → 拿不到新分数 → 永远不验证"的反馈陷阱：不变量 I1 保证 block 永远完整 draft， +confidence head 对全部 K 个位置打分，与验几个无关。 + +把 floor 调到 0 以上只有一个正当理由 —— **公平性**（保证每个请求每步至少有一次投机机会） +—— 并且要付代价：floor 会把预算花在排名低于被它挤掉的那些位置上。 + +```python +candidates = survival[:, floor:].flatten() # 每一个 (request, position) 对 +# 按 survival 全局排序，取前 budget 个，然后按行计数 +``` + +不需要任何显式的前缀约束：survival 沿位置非增，所以**任何全局 top-k 自动就是每个请求的 +前缀**。因此分配逻辑只需要**数**每行中了几个，不需要追踪中的是哪几个。 + +tie 用 `(position, request)` 打破，**绝不用数值**。两个 TP rank 即使算出 bitwise 不同的 +confidence，也必须选出相同的 verify 长度，否则它们的 batch 形状（进而 collective）会发散。 + +## 一个完整的数值例子 + +`K=7, bs=4, tiers=[1,3,7], min_verify_len=0`（默认），成本表 +`token_counts=(0, 8, 16, 32)`、`step_time_ms=(2.0, 2.2, 2.6, 9.0)`。 + +survival 矩阵： + +```text + j=0 j=1 j=2 j=3 j=4 j=5 j=6 + r0 | 0.95 0.87 0.79 0.69 0.59 0.47 0.35 <- draft 很健康 + r1 | 0.60 0.18 0.02 0.00 0.00 0.00 0.00 <- 第 2 位就死了 + r2 | 0.90 0.81 0.73 0.66 0.59 0.53 0.48 <- 健康 + r3 | 0.70 0.49 0.34 0.24 0.17 0.12 0.08 <- 中等 + 列和 3.15 2.35 1.88 1.59 1.35 1.12 0.91 +``` + +uniform 决策： + +| L | M = 4*(L+1) | T (ms) | tau = 4 + 列和 | Theta | +|---|-------------|--------|----------------|-------| +| 1 | 8 | 2.2 | 7.15 | 3.25 | +| **3** | **16** | **2.6** | **11.38** | **4.38 <- 最大** | +| 7 | 32 | 9.0 | 16.35 | 1.82 | + +于是 `runtime_draft_len = 3`，每个请求验证 3 个 draft 位置。 + +ragged 用**同样的预算**（`4 * 3 = 12` 个位置）重新分配。floor 固定为 1（`min_verify_len >= 1`，每个请求至少验证 anchor 位置；SGLang 允许 floor 为 0、可把绝望请求饿到零窗口，这是两者的已知语义差异），floor 之上的位置参与全局竞争，排序后取预算内的前若干个： + +```text +0.95(r0) 0.90(r2) 0.87(r0) 0.81(r2) 0.79(r0) 0.73(r2) +0.70(r3) 0.69(r0) 0.66(r2) 0.60(r1) 0.59(r0) 0.59(r2) + \--------------- r0 拿 5 个，r2 拿 5 个，r3 和 r1 各 1 个 ---------------/ +r3 的下一个候选 0.49、r1 的 0.18 都落在预算之外 +``` + +| | verify_len（draft 位置数） | 期望产出 | +|---|---|---| +| r0 | **5** | 3.89 | +| r1 | **1** | 0.60 | +| r2 | **5** | 3.69 | +| r3 | **1** | 0.70 | +| **合计** | **12**（与 uniform 的 `4 * 3` 完全相同） | **tau = 12.88** | + +```text +uniform L=3 : 12 个验证位置 -> tau = 11.38 +ragged : 12 个验证位置 -> tau = 12.88 （成本完全相同，+13%） +``` + +这就是 ragged 模式的全部动机：把 r1 和 r3 注定要浪费掉的位置，转移给 r0 和 r2。 + +## Ragged 验证 + +以 `verify_lens`（token 单位，含 bonus）`= [6, 2, 6, 2]` 为例： + +```text +qo_indptr = [0, 6, 8, 14, 16] # 独占前缀和 + +flat token 轴（attention kernel 想要的布局）: + idx 0 1 2 3 4 5 | 6 7 | 8 9 10 11 12 13 | 14 15 + b d0 d1 d2 d3 d4 | b d0 | b d0 d1 d2 d3 d4 | b d0 + \------- r0 ---------/ \- r1 -/ \-------- r2 ---------/ \- r3 -/ + ^ ^ ^ ^ + qo_indptr[0]=0 [1]=6 [2]=8 [3]=14 +``` + +### bucket 对齐 + +被捕获的 token bucket 是每个 tier 对应的 `bs * (t + 1)`。一个 batch 把自己的总数向上取整 +到其中之一，`RaggedVerifyLayout.fill_bucket` 分两阶段花掉这部分余量：**先给真实请求** +（上限 `max_verify_len`，这样这些 token 去验证的是 step 本来就要付钱的真实 draft 位置）， +剩下的才给 pad 行。 + +token 数必须**精确**命中 bucket。`seq_lens.sum()` 会变成 `attn_metadata.num_tokens`， +它会在 attention-DP 各 rank 间 all-gather，并驱动 MoE 的 chunk 数。如果少报，attention 和 +MoE 对"在飞多少 token"的认知就不一致 —— **而且什么都不会报错**。 + +### 形状恒定，内容可变 + +这正是 ragged batch 能被捕获的原因： + +```text +graph 捕获的是 -> (padded_bs 行, bucket 个 token) # 一小组固定形状 +每步变化的是 -> verify_lens / qo_indptr 里的数值 # 持久 buffer +``` + +## CUDA graph 处理 + +```text +uniform: (bs, draft_len, False, short_seq_mode, is_all_greedy) + ^ tiers 中的某一个 + +ragged: (bs, top_tier, False, short_seq_mode, is_all_greedy, bucket) + ^ token 轴 +``` + +ragged 下 `draft_len` 钉死在最高 tier（不变量 I1：block 总是完整 draft），变化的是 token +总数。两种模式都捕获 `|batch_sizes| * |tiers|` 张图。 + +与"按 batch size 调度 draft 长度"不同，这里的 tier **不是** batch size 的函数 —— 在固定 +`bs` 下，随着 batch 的 confidence 变化，调度器仍可能选不同的长度 —— 所以捕获集合是**叉积**。 +每张捕获的图大约花 10-23 MB 元数据，而这块内存**直接从 KV cache 里出**（KV pool 是按 capture +之后剩下的量来定的）。所以 ladder 必须短。 + +capture 安全性依赖几个具体选择： + +| 选择 | 原因 | +|------|------| +| `return_confidence` 在构造时读一次，绝不逐步判断 | 否则捕获的 draft graph 会发散 | +| STS buffer 原地更新，绝不 rebind | `nn.Module.__setattr__` 换 buffer 会换 `data_ptr`；先前捕获的 graph 会一直读旧存储，标定静默失效 | +| `_confidence_logits` 是持久 buffer，通过 slot 散射写入 | replay 时落到当前 batch 实际占用的槽位 | +| ragged buffer 通过持久分配的切片写入 | 保持 `data_ptr` 稳定，graph 才看得见新值 | +| `repeat_interleave` 一律传 `output_size` | 否则 torch 要把 repeats 的累积和读回 host 来确定输出尺寸 —— 在 capture 期非法，在其他地方则是每步一次同步 | + +## 跨 rank 一致性 + +```text +draft_len 是 CUDA graph key 的一部分 + | + v +但它不在 attention-DP 的一致性 allgather 里 + | + v +两个 rank 选了不同长度 -> 选到不同的 graph + | + v +一个 replay、一个掉进 eager -> 它们的 collective 发散 +``` + +归约取的是 **max 而不是 min**：想裁更多的 rank 只是多验几个 token，而 graph key 不一致 +是**不可恢复**的。ragged 下捕获形状是二维的，一致性来自"对归约结果取整"，而不是"各 rank +对自己的数取整"： + +```text +padded_bs = round_up_bs( max_r num_real_r ) +bucket = round_up_bucket( padded_bs + max_r slack_r ) +``` + +## 退化路径 + +静默退化是这个特性的主要失效模式，所以每一条放弃裁剪的路径都在 `planner.stats` 里计数： + +```text + +- planner 未构建 -----------------> 全块（无 checkpoint） + | + +- 成本表未 profile（flat）--------> 全块 + 启动时告警 + | "每个 token 都免费" => 裁剪等于盲猜 + | +_dspark_confidence --+- 快照还没落地 -------------------> 全块 +_draft_len() | （cudaEvent.query() == False） + | + +- 快照行数少于 batch --------------> 全块 + | + +- 请求从未被打分 ------------------> 中性行 -> 全块 + | （刚 prefill 完，或 slot 刚回收） + | + +- 找不到已捕获的 ragged 形状 ------> 退回 uniform（uniform 一定有图） +``` + +所有路径都朝同一个方向退化：**验证整个 block**，即启用该特性之前的行为。输入不可信时， +调度器**永远不会**比静态基线验得更少。 + +部分结果被当作失败而非"部分成功"：如果 planner 无法给出恰好每请求一个 verify 长度，整批 +退回 uniform。一个"半窗口"的 batch 比完全不加窗口更糟 —— 输入布局是逐请求构建的（会走成 +ragged），而 spec metadata 看到缺失的窗口后会保持 uniform。 + +## 配置 + +```yaml +speculative_config: + decoding_type: DSpark + max_draft_len: 7 + + # 主开关。block 仍然完整 draft，只改变送给 target 验证的 token 数。 + enable_confidence_scheduling: true + + # 逐位置 STS 标定。省略 => 裸 sigmoid（恒等标定）。 + # 调度器消费的是累积乘积，所以逐位置误差会沿 block 几何级放大。 + confidence_sts_path: /path/to/sts.json + + # profile 好的 step-cost 曲线。省略 => flat => planner 拒绝裁剪。 + confidence_sps_table_path: /path/to/sps.json + + # 被捕获的长度 ladder；调度器只能从这个集合里选。 + # 默认：[1, ceil(K/2), K] + confidence_verify_len_tiers: [1, 3, 7] + + # confidence relay 的总滞后（decode step 数）。 + # executor loop 本身已经提供恰好 2 步，所以默认值不额外增加任何滞后。 + confidence_relay_lag_steps: 2 + + # 逐请求 verify 窗口。需要 enable_confidence_scheduling。 + enable_ragged_verify: false +``` + +`sps.json`： + +```json +{ + "token_counts": [0, 128, 256, 512], + "step_time_ms": [2.0, 2.1, 4.0, 9.0], + "fixed_overhead_ms": 0.0, + "batch_sizes": [1, 16, 64, 256], + "batch_overhead_ms": [3.0, 5.0, 17.0, 40.0] +} +``` + +`token_counts` 按**总验证 token 数**索引，**包含 bonus token**。`fixed_overhead_ms` 和 +`batch_*` 这一对是可选的，但绝不是装饰 —— 见[算法](#算法)中的成本模型部分。 + +## 性能天花板与已知局限 + +```text +一个 decode step = bias + alpha(bs) + theta(M) + ------------------ -------- + 调度动不了这部分 唯一可回收的一项 +``` + +因此收益是一个**高并发效应**。在小 batch 下，step 成本对验证 token 数几乎是平的，裁剪省 +不下时间，却仍要付出接受长度的代价 —— 打平甚至为负。 + +**在目标硬件上量出可裁剪部分的占比是前提，不是事后工作。** 如果 `alpha(bs)` 压过 +`theta(M)`，那么无论 confidence head 多好，天花板都很低。两个测量陷阱： + +- 在**固定 batch size** 下 profile，会让 `M = bs * (L + 1)` 与 `L` 完全共线，`alpha` 和 + `theta` 分不开。必须扫 batch size。 +- 如果 `d(step)/dK` 在各 batch size 上是平的，说明成本是 per-draft-pass 而非 + per-verify-token —— 调度器几乎没有可回收的东西。 + +其他局限： + +| 局限 | 说明 | +|------|------| +| tier ladder 跨 batch size 有损 | shelf 右边缘随 `bs` 移动，一套 ladder 不可能同时坐在所有 batch size 的边缘上。残差要实测 | +| 快照是 2 步之前的 | 纯吞吐启发式（I2 保证正确性与它无关），但在剧烈抖动的负载下决策质量会下降 | +| ragged 模式未做端到端验证 | 需要带训练好的 confidence-head 权重的 checkpoint；默认关闭 | +| 逐 token 采样参数路径只有非全贪心 batch 才会触发 | `temperature=0` 的精度跑覆盖不到它 | + +## 参考 + +- DeepSeek DeepSpec： +- SGLang 的 DSpark 集成： diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 512544d7eb0f..cf86b9114c4a 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -1,3 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 import bisect import contextlib from dataclasses import dataclass diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index a5ff8653c328..a3d9f4fdfcd9 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3700,15 +3700,15 @@ def _set_warmup_ragged_windows(self, batch, verify_bucket: int, filled = layout.fill_bucket(max_verify_len=1 + int(draft_len)) for request, tokens in zip(requests, filled.verify_lens.tolist()): request.py_verify_len = int(tokens) - 1 - # Publish it to the runner as well, or the graph is captured under a key - # that does not describe the batch just shaped. `_ragged_verify_bucket` - # reads `agreed_ragged_bucket` rather than re-summing the batch (so that - # every attention-DP rank keys on a value they agreed on, instead of one - # each derived independently) -- and nothing sets that during capture, - # because the runtime fit never runs here. The key's token axis was - # therefore None while the batch really held `verify_bucket` tokens, and - # the mismatch surfaced far away, in the MoE's shared-vs-routed row - # count, as `unmatched tensor shape` two graphs into the capture pass. + # Publish it to the runner as well: the graph key re-derives its token + # axis from the batch (`_ragged_verify_bucket`), but the layout assert + # cross-checks that against `agreed_ragged_bucket` -- the fit's claim -- + # and during capture nothing else sets it, because the runtime fit + # never runs here. Before this line existed the capture-time key's + # token axis was None while the batch really held `verify_bucket` + # tokens, and the mismatch surfaced far away, in the MoE's + # shared-vs-routed row count, as `unmatched tensor shape` two graphs + # into the capture pass. self.cuda_graph_runner.agreed_ragged_bucket = int(verify_bucket) def ragged_verify_token_buckets(self, padded_bs: int) -> List[int]: @@ -3780,8 +3780,13 @@ def fit_ragged_verify_lens( runner = self.cuda_graph_runner # Cleared on entry so that every early return leaves no stale # value behind: the graph key reads this, and a bucket carried - # over from a previous step would key into the wrong graph. + # over from a previous step would key into the wrong graph. The + # padded row count follows the same rule -- record_step books it + # into padded_bs_hist on every step, and a value surviving from the + # last successful fit would pollute exactly the fallback steps the + # histogram exists to distinguish. runner.agreed_ragged_bucket = None + self._dspark_last_padded_bs = None if not runner.enabled or not verify_lens: return None if len(verify_lens) != len(generation_requests): @@ -3886,7 +3891,6 @@ def fit_ragged_verify_lens( buckets = self.ragged_verify_token_buckets(padded_bs) if not buckets: return None - try: shape = choose_ragged_capture_shape( num_real_requests=len(token_lens), @@ -3902,8 +3906,16 @@ def fit_ragged_verify_lens( # Remembered so a graph miss can report the shape it actually submitted # rather than just that one happened. self._dspark_last_total_tokens = int(bucket) - # The graph key uses THIS value rather than re-summing the - # per-request windows. Both used to be derived independently -- + # And the ROW count behind that bucket. The observability ceiling is + # booked over the rank-local, unpadded request count while `delivered` + # is booked over this padded, cross-rank-maximum one, which is how + # trim_ratio went negative in four runs. Only the caller can reconcile + # them, and it cannot do that without this number. + self._dspark_last_padded_bs = int(padded_bs) + # The fit's agreed value; the graph key re-derives its token axis + # from the padded batch and the layout assert cross-checks the two + # derivations against each other. Both used to be derived + # independently WITHOUT that check -- # the key walked generation_requests *after* padding was # appended, while the fit worked from the allgathered peer # stats -- so their agreement across attention-DP ranks was @@ -4075,8 +4087,7 @@ def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, from_requests = sum(verify_lens) from_layout = attn_metadata.num_tokens - attn_metadata.num_ctx_tokens from_spec = int(spec_metadata.total_verify_tokens) - bucket = self.cuda_graph_runner._ragged_verify_bucket( - scheduled_requests) + agreed = getattr(self.cuda_graph_runner, "agreed_ragged_bucket", None) assert from_requests == from_layout, ( f"ragged layout mismatch: requests carry {from_requests} generation " @@ -4088,12 +4099,19 @@ def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, f"ragged layout mismatch: requests carry {from_requests} generation " f"tokens but spec_metadata.total_verify_tokens is {from_spec}; " f"acceptance would slice the target logits at the wrong offsets.") - if bucket is not None: - assert from_requests == int(bucket), ( - f"ragged layout mismatch: requests carry {from_requests} " - f"generation tokens but the CUDA-graph key's token axis is " - f"{bucket}; the replayed graph was captured at a different " - f"width and nothing downstream would notice.") + if agreed is not None: + # Two INDEPENDENT derivations of the token axis: `from_requests` + # walks this batch's windows, `agreed` is the fit's allgathered + # arithmetic, and the graph key carries the latter. Comparing the + # key's own return value instead (what an earlier form did) is a + # tautology -- _ragged_verify_bucket publishes `agreed` -- so a fit + # that budgeted for padding rows the batch never got would replay a + # graph captured at a different width with nothing to say so. + assert from_requests == int(agreed), ( + f"ragged layout mismatch: the batch carries {from_requests} " + f"generation tokens but the fit agreed on {agreed}, which is " + f"what the CUDA-graph key's token axis says; the replayed " + f"graph was captured at a different width.") # Row axis. Everything above is a token count, so all four agree even # when the token-major presentation is malformed -- the presentation is From 0b524435fc77494f15b9690943d0cfd361ff4052 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 4 Aug 2026 21:09:45 -0700 Subject: [PATCH 089/177] [None][feat] Pin the DSpark verify length at runtime A cost-table sweep has to hold every request at one verify length while it measures a cell, or the cell's label describes a shape that never ran. Today that is an environment variable read once when the planner is built, so walking the ladder costs an engine rebuild per rung -- and on a world whose ranks were spawned before the variable was set, it reaches none of them. POST /dspark/verify_len_pin {"verify_len": N} pins it on a running server; null hands scheduling back to the planner. An unpinned tier is refused with a 400 rather than accepted into a step that would fall out of graph replay. The pin cannot be applied where it arrives. One rank pinning while its peers do not makes their token totals disagree, and the attention-DP shape gate drops the whole group out of replay for as long as it lasts. So the request is queued locally and carried by the allgather the decode step already issues: the payload gains one int, the collective's cost does not change, and every rank adopts the same value on the same step because they all read it from the same reduction. That is also why the RPC needs no control communicator and never pauses the loop, unlike sleep/wakeup. Three states have to survive the trip -- nothing queued, pin, unpin -- so the queue holds the wire value (-1/0/N) rather than an Optional[int], which collapsed "clear" into "nothing queued" and left a clear stranded on rank 0. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 43 +++++++ .../_torch/speculative/dspark_verify.py | 61 ++++++++++ tensorrt_llm/executor/base_worker.py | 14 +++ tensorrt_llm/executor/proxy.py | 5 +- tensorrt_llm/serve/openai_protocol.py | 10 ++ tensorrt_llm/serve/openai_server.py | 28 +++++ .../hw_agnostic/test_dspark_verify_len_pin.py | 108 ++++++++++++++++++ 7 files changed, 268 insertions(+), 1 deletion(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 17aa194b2768..b0fc76337db5 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3338,6 +3338,37 @@ def wait_on_pp_send_handles(self, send_handles, microbatch_id): send_handles[microbatch_id].wait() send_handles[microbatch_id] = None + def set_dspark_verify_len_pin(self, verify_len: Optional[int]) -> Optional[int]: + """Queue a DSpark verify-length pin, or ``None`` to clear it. + + The pin bypasses the confidence planner and makes every request verify + the same number of positions -- the shape a cost-table sweep needs to + measure a cell it can label. It exists so a sweep can walk the ladder + against a RUNNING server instead of rebuilding the engine per length + (which is what the ``TLLM_DSPARK_FORCE_VERIFY_LEN`` environment + variable requires, and which cannot reach ranks that were spawned + before it was set). + + Validated here and applied on the next decode step, when the step's + existing cross-rank allgather hands the same value to every rank. + + Returns: + The validated pin, or None when cleared. + + Raises: + RuntimeError: if this executor has no DSpark confidence planner. + ValueError: if the length is outside the ladder (see + :meth:`DSparkVerifyPlanner.validate_verify_len_pin`). + """ + worker = getattr(self.model_engine, "_get_spec_worker", None) + planner = getattr(worker(), "verify_planner", None) if worker else None + if planner is None: + raise RuntimeError( + "no DSpark confidence planner on this executor; the verify-len " + "pin only applies to a DSpark engine with confidence " + "scheduling enabled") + return planner.request_verify_len_pin(verify_len) + def _dspark_confidence_draft_len(self, scheduled_batch: ScheduledRequests) -> int: """Draft length for this iteration, from the DSpark confidence planner. @@ -3470,7 +3501,19 @@ def _dspark_confidence_draft_len(self, local_tokens, local_max_window, 1 if scheduled_batch.can_run_cuda_graph else 0, + planner.pending_verify_len_pin(), ]) + # A runtime verify-length pin (profiling endpoint) reaches the + # group here rather than through each rank's own environment: the + # endpoint lands on one rank, and a pin applied locally would make + # that rank disagree with its peers for as long as it held. Taking + # the first rank that queued one makes the choice a function of the + # allgathered payload, identical everywhere, effective on the same + # step. -1 means nobody queued anything. + agreed_pin = next( + (int(payload[5]) for payload in payloads + if len(payload) > 5 and int(payload[5]) >= 0), -1) + planner.adopt_verify_len_pin(agreed_pin) if not all(int(payload[0]) for payload in payloads): # Any rank that cannot go ragged takes the whole group with it: # a mixed step means one rank replays a ragged graph while diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index deedd9250e67..c1435659c2b6 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -108,6 +108,17 @@ def __init__( # one M and filed under another. With an absolute L every request gets # exactly L, so M = bs * (L + 1) is the captured bucket by construction. self._forced_verify_len = self._read_forced_verify_len() + # A pin requested at runtime (profiling endpoint) but not yet in force, + # held as its WIRE value because three states have to survive the trip: + # -1 nothing queued, 0 clear the pin, >0 pin to that length. Storing + # Optional[int] collapsed "clear" into "nothing queued", so a clear + # never reached the peers. min_verify_len is validated >= 1, so 0 can + # never collide with a real length. + # It is adopted through the step's existing cross-rank allgather so + # every rank starts pinning on the SAME step. Setting it locally would + # diverge the shape gate for as long as the ranks disagreed -- the same + # failure the env-var pin has on a pre-spawned world. + self._pending_verify_len_pin: int = -1 self._host_buffer: Optional[torch.Tensor] = None self._copy_event: Optional[torch.cuda.Event] = None @@ -209,6 +220,56 @@ def _read_forced_verify_len(self) -> Optional[int]: f"serving configuration.") return value + def validate_verify_len_pin(self, value: Optional[int]) -> Optional[int]: + """Check a pin against this planner's ladder; None clears the pin.""" + if value is None: + return None + value = int(value) + lo, hi = int(self.cfg.min_verify_len), int(self.cfg.resolved_max_verify_len) + if not lo <= value <= hi: + raise ValueError(f"verify-length pin {value} outside [{lo}, {hi}]") + if value not in self.tiers: + raise ValueError( + f"verify-length pin {value} is not in the captured tier ladder " + f"{self.tiers}; every step would fall out of CUDA-graph replay " + f"and the measured times would not be step costs") + return value + + def request_verify_len_pin(self, value: Optional[int]) -> Optional[int]: + """Queue a pin (or ``None`` to clear) for the next decode step. + + Validated here so a bad value is rejected at the call site rather than + on some later step, but NOT applied here: the step's allgather carries + it so all ranks adopt it together (see ``adopt_verify_len_pin``). + """ + value = self.validate_verify_len_pin(value) + self._pending_verify_len_pin = 0 if value is None else int(value) + return value + + def pending_verify_len_pin(self) -> int: + """The queued request as a wire value: -1 none, 0 clear, >0 pin.""" + return int(self._pending_verify_len_pin) + + def adopt_verify_len_pin(self, wire_value: int) -> None: + """Apply the group's agreed pin. ``-1`` leaves the current one alone. + + Called from the executor with a value every rank read out of the same + allgathered payload, so the pin takes effect on the same step + everywhere. 0 is the wire encoding for "clear the pin". + """ + wire_value = int(wire_value) + if wire_value < 0: + return + self._pending_verify_len_pin = -1 + new_pin = None if wire_value == 0 else wire_value + if new_pin == self._forced_verify_len: + return + self._forced_verify_len = new_pin + logger.warning( + f"DSpark: verify length pin set to {new_pin} at runtime. This " + f"bypasses the planner and is a profiling mode, not a serving " + f"configuration.") + def _note_snapshot_stats(self, selected: torch.Tensor, rows: Optional[Sequence[int]]) -> None: """Measure what the argmax is about to eat. Never changes behaviour. diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 5e1413a4f0da..33aeab1ce4df 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -872,6 +872,20 @@ def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> None: f"{action}() failed on {len(errors)} rank(s):\n" + "\n".join(errors)) + def set_dspark_verify_len_pin(self, + verify_len: Optional[int]) -> Optional[int]: + """Pin (or unpin) the DSpark verify length on this worker's executor. + + Safe under TP/PP without the control-communicator dance ``sleep`` needs: + the pin is a scalar the decode loop reads on its next step, and it is + broadcast to the other ranks by the step's own allgather rather than by + this call. So rank 0 alone receiving the RPC is sufficient and the loop + never has to be paused. + """ + if self.engine is None: + raise RuntimeError("engine is not initialized") + return self.engine.set_dspark_verify_len_pin(verify_len) + def sleep(self, sleep_tags: list[str]) -> None: """Release GPU virtual memory for the specified memory type tags. diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index d1baa0792785..b1080a4c3a31 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -61,7 +61,10 @@ # sleep() and wakeup() coordinate across all ranks via a dedicated control # communicator started in PyExecutor.start_worker(); rank-0 is the entry # point, so routing through the rank-0 RPC shim is correct. -_MULTI_RANK_ALLOWED_METHODS: frozenset[str] = frozenset({"sleep", "wakeup"}) +# set_dspark_verify_len_pin() needs no control communicator at all: it queues +# a scalar that the decode loop's own per-step allgather hands to every rank, +# so rank-0 receiving it is sufficient and the loop is never paused. +_MULTI_RANK_ALLOWED_METHODS: frozenset[str] = frozenset({"sleep", "wakeup", "set_dspark_verify_len_pin"}) def _check_collective_rpc_guard( diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 4f497787e143..159853c5f514 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -1370,6 +1370,16 @@ class MemoryUpdateRequest(OpenAIBaseModel): tags: List[str] = Field(default=["model", "kv_cache"]) +class DSparkVerifyLenPinRequest(OpenAIBaseModel): + verify_len: Optional[int] = Field( + default=None, + description="DSpark verify length to pin every request to, or null to " + "clear the pin and hand scheduling back to the confidence planner. " + "Must be one of the captured verify-length tiers: an uncaptured " + "length drops every step out of CUDA-graph replay. Profiling knob -- " + "it bypasses the planner and is not a serving configuration.") + + class UpdateWeightsRequest(OpenAIBaseModel): weights: Optional[Dict[str, str]] = Field( default=None, diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 52dc5e488387..0e8a541bc9cd 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -73,6 +73,7 @@ CompletionResponseChoice, EmbeddingRequest, EmbeddingResponse, EmbeddingResponseData, EmbeddingUsageInfo, ErrorResponse, ImageGenerationRequest, ImageGenerationResponse, ImageObject, + DSparkVerifyLenPinRequest, MemoryUpdateRequest, ModelCard, ModelList, PromptTokensDetails, ResponseFormat, ResponsesRequest, ResponsesResponse, TokenizeRequest, TokenizeResponse, UpdateWeightsRequest, UsageInfo, @@ -888,6 +889,9 @@ def register_routes(self): self.app.add_api_route("/server_info", self.get_server_info, methods=["GET"]) + self.app.add_api_route("/dspark/verify_len_pin", + self.set_dspark_verify_len_pin, + methods=["POST"]) if self.generator.args.return_perf_metrics: # register /prometheus/metrics self.mount_metrics() @@ -2381,6 +2385,30 @@ async def update_weights(self, args=(request.weights, )) return JSONResponse(content={"status": "success"}) + async def set_dspark_verify_len_pin( + self, request: DSparkVerifyLenPinRequest) -> JSONResponse: + """Pin (or clear) the DSpark verify length at runtime. + + Lets a cost-table sweep walk the verify-length ladder against a + RUNNING server -- the shape each cell needs in order to be labelled -- + instead of rebuilding the engine per length, which is what the + ``TLLM_DSPARK_FORCE_VERIFY_LEN`` environment variable requires and + which cannot reach ranks spawned before it was set. The pin takes + effect on the next decode step, simultaneously on every rank. + """ + assert isinstance( + self.generator, AsyncLLM + ), "/dspark/verify_len_pin endpoint is only supported with AsyncLLM()" + try: + applied = await self.generator.collective_rpc( + 'set_dspark_verify_len_pin', args=(request.verify_len, )) + except (ValueError, RuntimeError) as exc: + # A bad tier or a non-DSpark engine is a caller error, not a server + # fault: answer it rather than returning a 500 with a stack trace. + return JSONResponse(content={"error": str(exc)}, status_code=400) + value = applied[0] if isinstance(applied, list) and applied else applied + return JSONResponse(content={"verify_len": value}) + async def get_server_info(self) -> JSONResponse: content = {"disaggregated_params": self.generator.disaggregated_params} args = getattr(self.generator, "args", None) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py new file mode 100644 index 000000000000..20f892fc5872 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The runtime verify-length pin: queue, agree, adopt. + +A cost-table sweep has to hold every request at one verify length while it +measures a cell, or the cell's label describes a shape that never ran (this +branch shipped a sweep where eight of twelve cells were mislabelled and the +fit dutifully concluded the curve was flat). The environment variable that +does this today is read once at planner construction, so walking the ladder +means rebuilding the engine per length -- and on a pre-spawned world it +cannot reach the ranks that are already running. + +The runtime pin fixes both, but it must never let one rank pin while its +peers do not: the ranks would derive different token totals, and the +attention-DP shape gate would drop the whole group out of graph replay for as +long as they disagreed. So the pin is queued locally, carried by the step's +existing allgather, and adopted from that agreed payload -- which is what +these tests pin down. +""" + +import pytest + +from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig +from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + + +def _planner(tiers=(1, 2, 5)): + table = SpsCostTable(token_counts=(0, 512, 768, 1536), + step_time_ms=(5.0, 68.4, 80.2, 150.5), + fixed_overhead_ms=1.0) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + return DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=list(tiers)) + + +def test_queued_pin_is_not_yet_in_force(): + """Queuing must not change this rank's behaviour before the group agrees.""" + planner = _planner() + assert planner.request_verify_len_pin(2) == 2 + assert planner._forced_verify_len is None + assert planner.pending_verify_len_pin() == 2 + + +def test_adopting_the_agreed_value_applies_it(): + planner = _planner() + planner.request_verify_len_pin(2) + planner.adopt_verify_len_pin(2) + assert planner._forced_verify_len == 2 + # Consumed: the next step must not re-broadcast a pin nobody asked for. + assert planner.pending_verify_len_pin() == -1 + + +def test_a_rank_that_queued_nothing_still_adopts_the_group_value(): + """The endpoint lands on ONE rank; the others learn it from the payload.""" + peer = _planner() + assert peer.pending_verify_len_pin() == -1 + peer.adopt_verify_len_pin(5) + assert peer._forced_verify_len == 5 + + +def test_minus_one_means_nobody_asked_and_leaves_the_pin_alone(): + planner = _planner() + planner.adopt_verify_len_pin(2) + planner.adopt_verify_len_pin(-1) + assert planner._forced_verify_len == 2 + + +def test_zero_clears_the_pin(): + """Clearing needs its own wire value: None cannot travel in an int payload.""" + planner = _planner() + planner.adopt_verify_len_pin(2) + assert planner.request_verify_len_pin(None) is None + assert planner.pending_verify_len_pin() == 0 + planner.adopt_verify_len_pin(0) + assert planner._forced_verify_len is None + + +def test_an_uncaptured_length_is_refused_at_the_call_site(): + """Rejected when requested, not on some later step, and never queued.""" + planner = _planner(tiers=(1, 2, 5)) + with pytest.raises(ValueError, match="captured tier ladder"): + planner.request_verify_len_pin(3) + assert planner.pending_verify_len_pin() == -1 + assert planner._forced_verify_len is None + + +def test_a_length_outside_the_bounds_is_refused(): + planner = _planner() + with pytest.raises(ValueError, match="outside"): + planner.request_verify_len_pin(9) + + +def test_pinned_steps_hand_every_request_the_same_window(): + """The whole point: the shape a sweep cell can honestly label.""" + import numpy as np + import torch + + planner = _planner() + bs = 8 + rng = np.random.default_rng(5) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() + planner._gather_rows = lambda **_: torch.tensor(survival, + dtype=torch.float32) + planner.adopt_verify_len_pin(2) + lens = planner.decide_verify_lens(num_gen_requests=bs, + reduce_across_ranks=False) + assert lens == [2] * bs + assert planner.stats["forced_steps"] == 1 From fab6291e2d91bf15d0cf5adf3f2c2aef3cf71147 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 4 Aug 2026 22:17:44 -0700 Subject: [PATCH 090/177] [None][chore] Measure the SPS table against the deployment, not a rebuild The profiler built its own engine per verify length. That is total control at the price of a configuration kept in sync by hand, and when it drifted the table described a machine nobody was running: a swept (bs=256, M=1536) cell read 150.2 ms against 118.8 ms measured live at the same shape, with nothing in the fingerprint to record the difference. Two collectors take the measurement to the deployment instead. --base-url sweeps a running server: the verify length is pinned over /dspark/verify_len_pin, so walking the ladder costs no engine rebuild and the engine measured is by construction the one that will consume the table. The length is not read back from telemetry (/metrics carries no generation-token count) but taken from the pin the server accepted, and the timing comes from hostStepTimeMS rather than iterLatencyMS, which spans about two loops under the overlap scheduler. --from-iter-log fits from production logs, which already carry the row count, the token total and the host step time per step. No server, no traffic, no machine time; coverage is whatever the deployment ran. On 39 logs from this branch's runs it yields 13k samples over five verify lengths. Reading those logs honestly took two corrections that a smoke test against real logs caught and the unit tests could not: the token total is the PADDED one, so it is the captured width -- not the row count -- that divides it, and resolving that width needs the graph ladder, without which a low-occupancy step is ambiguous between three lengths. The first version kept 312 of 5682 steps; it now keeps 5253 and files each under the shape that actually ran. The profiler moves to tests/microbenchmarks, alongside the other DSpark tools. Nothing on the serving path imports it, and a 1.8k-line sweep driver has no business shipping to every user. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 18 +- tensorrt_llm/llmapi/llm_args.py | 2 +- .../microbenchmarks}/dspark_sps_profiler.py | 345 +++++++++++++++++- .../test_dspark_iter_log_samples.py | 189 ++++++++++ .../hw_agnostic/test_dspark_sps_profiler.py | 38 +- .../test_dspark_graph_miss_attribution.py | 27 +- 6 files changed, 579 insertions(+), 40 deletions(-) rename {tensorrt_llm/_torch/speculative => tests/microbenchmarks}/dspark_sps_profiler.py (83%) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index a3d9f4fdfcd9..cb8cf8586df4 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3700,11 +3700,11 @@ def _set_warmup_ragged_windows(self, batch, verify_bucket: int, filled = layout.fill_bucket(max_verify_len=1 + int(draft_len)) for request, tokens in zip(requests, filled.verify_lens.tolist()): request.py_verify_len = int(tokens) - 1 - # Publish it to the runner as well: the graph key re-derives its token - # axis from the batch (`_ragged_verify_bucket`), but the layout assert - # cross-checks that against `agreed_ragged_bucket` -- the fit's claim -- - # and during capture nothing else sets it, because the runtime fit - # never runs here. Before this line existed the capture-time key's + # Publish it to the runner as well: the graph key carries this value + # (`_ragged_verify_bucket` returns the fit's claim so every + # attention-DP rank keys identically), and the layout assert + # cross-checks it against the batch walk. During capture nothing else + # sets it, because the runtime fit never runs here. Before this line existed the capture-time key's # token axis was None while the batch really held `verify_bucket` # tokens, and the mismatch surfaced far away, in the MoE's # shared-vs-routed row count, as `unmatched tensor shape` two graphs @@ -3912,10 +3912,10 @@ def fit_ragged_verify_lens( # trim_ratio went negative in four runs. Only the caller can reconcile # them, and it cannot do that without this number. self._dspark_last_padded_bs = int(padded_bs) - # The fit's agreed value; the graph key re-derives its token axis - # from the padded batch and the layout assert cross-checks the two - # derivations against each other. Both used to be derived - # independently WITHOUT that check -- + # The fit's agreed value, which the graph key carries verbatim so + # every attention-DP rank keys on a number they agreed on; the layout + # assert cross-checks it against the batch walk. Both used to be + # derived independently WITHOUT that check -- # the key walked generation_requests *after* padding was # appended, while the fit worked from the allgathered peer # stats -- so their agreement across attention-DP ranks was diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index f99898d890da..126ca66122b7 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2901,7 +2901,7 @@ def validate_ragged_verify(self): raise ValueError( "enable_ragged_verify=True requires a profiled cost table: pass " "confidence_sps_table_path (produced by " - "`python -m tensorrt_llm._torch.speculative.dspark_sps_profiler`) " + "`python tests/microbenchmarks/dspark_sps_profiler.py`) " "or confidence_sts_path. Without one the planner cannot compare " "the cost of extra verify tokens against their expected " "acceptance, so it declines to trim and every request receives " diff --git a/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py b/tests/microbenchmarks/dspark_sps_profiler.py similarity index 83% rename from tensorrt_llm/_torch/speculative/dspark_sps_profiler.py rename to tests/microbenchmarks/dspark_sps_profiler.py index ae809e3afadc..9eb1dd4e0ae4 100644 --- a/tensorrt_llm/_torch/speculative/dspark_sps_profiler.py +++ b/tests/microbenchmarks/dspark_sps_profiler.py @@ -167,7 +167,9 @@ import argparse import json import os +import pathlib import random +import re import statistics import sys import threading @@ -177,8 +179,8 @@ import numpy as np -from .dspark_observability import RAGGED_VERIFY_MODE_ENV, RaggedVerifyMode -from .dspark_planner import ( +from tensorrt_llm._torch.speculative.dspark_observability import RAGGED_VERIFY_MODE_ENV, RaggedVerifyMode +from tensorrt_llm._torch.speculative.dspark_planner import ( SpsCostTable, budget_argmax_over_uniform_lens, derive_verify_len_tiers, @@ -1192,15 +1194,17 @@ def _build_llm(config: SweepConfig, verify_len: int): # (An earlier revision popped it and relied on the one-rung ladder alone, # citing a "pin skips the padding reservation" theory for the Repeat.cu # assert. That theory was wrong -- the assert reproduced with no pin at - # all; the actual cause was the CUDA-graph key returning the PREDICTED - # bucket rather than the batch's measured token total, fixed in - # cuda_graph_runner._ragged_verify_bucket. And the one-rung ladder was + # all; the actual cause was the CUDA-graph key describing a + # different token width than the batch carried. The key publishes the + # fit's bucket (so attention-DP ranks agree by construction) and the + # layout assert cross-checks it against the batch walk; a divergence now + # has to show up there rather than inside a replayed graph. And the one-rung ladder was # never one rung anyway: the llm_args validator appends max_draft_len, so # under a constant block [L] is really [L, block]. With a seed table # present the planner argmax'd its way to the block on 93-95% of steps # while the samples file recorded M = bs*(L+1) as if L had run -- twelve # cells, eight mislabelled, and the fit dutifully concluded q ~ 0.) - from .dspark_verify import FORCE_VERIFY_LEN_ENV + from tensorrt_llm._torch.speculative.dspark_verify import FORCE_VERIFY_LEN_ENV # LAUNCH-MODE CAVEAT: this mutation reaches ranks 1..N-1 only when the # workers are spawned AFTER it (MPI-spawn path). On a pre-spawned world @@ -1471,6 +1475,37 @@ def _build_parser() -> argparse.ArgumentParser: formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument("--out", required=True, help="Path to write the cost-table JSON.") + source = parser.add_argument_group( + "sample source", + "Where the (batch size, verify length, step time) triples come from. " + "Default: build an engine in-process and sweep it. The two " + "alternatives measure the deployment itself, so the table cannot " + "describe a configuration nobody is running.") + source.add_argument( + "--from-iter-log", nargs="+", metavar="LOG", + help="Fit from production iteration logs (print_iter_log). No server, " + "no traffic, no machine time; coverage is whatever the " + "deployment happened to run.") + source.add_argument( + "--base-url", + help="Sweep a RUNNING trtllm-serve instead of building an engine. The " + "verify length is pinned over /dspark/verify_len_pin, so walking " + "the ladder costs no engine rebuild.") + source.add_argument("--server-settle-s", type=float, default=20.0, + help="Seconds to let admission reach the plateau " + "before a cell is measured (--base-url).") + source.add_argument("--server-poll-s", type=float, default=20.0, + help="Seconds of steady state to collect per cell " + "(--base-url).") + source.add_argument("--served-model-name", default="model", + help="`model` field for /v1/completions (--base-url).") + source.add_argument( + "--cuda-graph-batch-sizes", + help="The deployment's captured batch-size ladder, comma separated " + "(--from-iter-log). Padding rounds a step's rows up to the " + "smallest entry at or above them, which is what makes the padded " + "token total resolve to exactly one verify length. Omit and only " + "steps whose shape is unambiguous on its own are kept.") parser.add_argument( "--samples-out", default=None, @@ -1632,6 +1667,263 @@ def _write_samples(path: str, samples: Sequence[StepSample]) -> None: handle.write(json.dumps(sample.to_json()) + "\n") +# --------------------------------------------------------------------------- +# Alternative sample sources +# +# The in-process sweep below builds its own engine per verify length. That +# gives total control, and costs a rebuild per length plus a configuration +# that has to be kept in sync with the deployment by hand -- when it drifts, +# the table describes a machine nobody is running (a profiled (256, 1536) cell +# read 150.2 ms against 118.8 ms measured live at the same shape, and nothing +# in the fingerprint recorded the difference). +# +# These two collectors take the measurement to the deployment instead. +# --------------------------------------------------------------------------- + + +def resolve_padded_shape( + *, + num_rows: int, + num_generation_tokens: int, + max_draft_len: int, + max_batch_size: int, + padded_batch_sizes: Optional[Sequence[int]] = None, +) -> Optional[Tuple[int, int]]: + """Recover ``(padded_bs, verify_len)`` from a step's row and token counts. + + The token total is the PADDED one: a step with 244 real rows submits + ``256 * (L + 1)`` tokens because the batch is padded up to a captured + graph size, and it is that padded shape the step actually cost. So the + length cannot be read off the real row count -- ``M / rows`` is not even + integral -- and a filter that demands it throws away nearly every step (it + kept 312 of 5682 on two production runs). + + Every candidate length implies a padded row count; the answer is the one + that is integral, at least the real row count, and within the engine's + batch ceiling. Ambiguous steps are rejected rather than guessed: filing a + step under a shape it did not run is the mislabelling that once produced a + clean-looking, entirely fictional table. + """ + if num_rows <= 0 or num_generation_tokens <= 0: + return None + ladder = (sorted({int(b) for b in padded_batch_sizes}) + if padded_batch_sizes else None) + if ladder: + # Padding rounds UP to the smallest captured batch size, never past + # it, so the row count alone fixes the padded width and the length + # follows. Without this the low-occupancy steps are genuinely + # ambiguous -- 768 tokens over 128 rows is L=5 at 128 rows, L=3 at + # 192, or L=2 at 256 -- and all of them would be dropped. + candidates = [b for b in ladder if b >= num_rows] + if not candidates: + return None + padded_bs = candidates[0] + if num_generation_tokens % padded_bs: + # Padding declined (a resource bail), so the step ran its real + # rows and this width does not describe it. Skip rather than + # invent one. + return None + verify_len = num_generation_tokens // padded_bs - 1 + if not 0 < verify_len <= int(max_draft_len): + return None + return (padded_bs, verify_len) + found: List[Tuple[int, int]] = [] + for verify_len in range(1, int(max_draft_len) + 1): + width = verify_len + 1 + if num_generation_tokens % width: + continue + padded_bs = num_generation_tokens // width + if padded_bs < num_rows or padded_bs > int(max_batch_size): + continue + if ladder is not None and padded_bs not in ladder: + continue + found.append((padded_bs, verify_len)) + return found[0] if len(found) == 1 else None + + +def samples_from_iter_log( + paths: Sequence[str], + *, + max_draft_len: int, + max_batch_size: int = 256, + padded_batch_sizes: Optional[Sequence[int]] = None, + rank_prefix: str = "0:", + max_step_ms: float = 1000.0, +) -> List[StepSample]: + """Read (bs, M, step time) triples out of production iteration logs. + + ``print_iter_log`` already emits everything a cell needs:: + + iter = N, ..., num_scheduled_requests = R, ..., + host_step_time = T ms, ..., states = {..., 'num_generation_tokens': M} + + so a table can be fitted from traffic that really ran, at the deployment's + own context depths, batch mix and captured-graph set -- no synthetic + prompts, no second engine, no machine time. The limitation is coverage: + only shapes the deployment happened to run appear, and rarely-taken rungs + may be thin. Use it to fit a table for a workload that is already in + production, or to check that a swept table has not drifted; use the server + sweep when specific cells have to be visited on purpose. + + ``verify_len`` is recovered by :func:`resolve_padded_shape`, which reads it + out of the padded token total instead of assuming it; steps whose shape is + ambiguous are skipped rather than filed under a length they did not run. + Samples carry the PADDED row count, because that is the shape whose cost + was measured. + + Pass ``padded_batch_sizes`` (the deployment's CUDA-graph ladder) when it is + known: it removes the last ambiguity, since only a captured row count can + have run. + """ + line_re = re.compile( + r"iter = (\d+).*?num_scheduled_requests = (\d+).*?" + r"host_step_time = ([\d.]+)ms.*?'num_generation_tokens': (\d+)") + out: List[StepSample] = [] + skipped_ratio = 0 + for path in paths: + text = pathlib.Path(path).read_text(errors="replace") + for line in text.splitlines(): + if "iter = " not in line: + continue + if rank_prefix and not line.lstrip().startswith(rank_prefix): + continue + match = line_re.search(line) + if not match: + continue + iteration = int(match.group(1)) + batch_size = int(match.group(2)) + step_ms = float(match.group(3)) + tokens = int(match.group(4)) + if iteration == 1: + # Each executor instance numbers from 1, and its first step + # carries the wait for the first request inside its host time + # (62 s on a live server). Identify it by the number rather + # than by having seen one, so a log FRAGMENT is still usable. + continue + if batch_size <= 0 or step_ms >= max_step_ms: + continue + shape = resolve_padded_shape( + num_rows=batch_size, num_generation_tokens=tokens, + max_draft_len=max_draft_len, max_batch_size=max_batch_size, + padded_batch_sizes=padded_batch_sizes) + if shape is None: + skipped_ratio += 1 + continue + padded_bs, verify_len = shape + out.append(StepSample(batch_size=padded_bs, verify_len=verify_len, + step_time_ms=step_ms, iteration=iteration)) + if skipped_ratio: + print(f"[dspark-sps] iter-log: skipped {skipped_ratio} steps whose " + f"padded shape could not be resolved unambiguously", + file=sys.stderr) + return out + + +def samples_from_server( + *, + base_url: str, + batch_sizes: Sequence[int], + verify_lens: Sequence[int], + input_len: int, + max_tokens: int, + model: str, + warmup_steps: int, + settle_s: float, + poll_s: float, + max_draft_len: int, + seed: int, + padded_batch_sizes: Optional[Sequence[int]] = None, +) -> List[StepSample]: + """Sweep a RUNNING server: pin the window over HTTP, then read /metrics. + + The equivalent of SGLang's profiler loop, and the reason + ``/dspark/verify_len_pin`` exists: walking the verify-length ladder no + longer costs an engine rebuild per rung, and the engine being measured is + by construction the engine that will consume the table. + + Load is synthetic (random token ids, ignore_eos) because a cell has to hold + a fixed batch size for long enough to measure it; the SHAPE is what the + table indexes, and the deployment's own graphs, context depths and + scheduler are what answer for it. + """ + import requests + + rng = random.Random(seed) + out: List[StepSample] = [] + + def pin(verify_len: Optional[int]) -> None: + response = requests.post(f"{base_url}/dspark/verify_len_pin", + json={"verify_len": verify_len}, timeout=60) + if response.status_code != 200: + raise RuntimeError( + f"pinning verify_len={verify_len} was refused by the server: " + f"{response.status_code} {response.text}") + + def drain_metrics() -> List[dict]: + response = requests.get(f"{base_url}/metrics", timeout=60) + response.raise_for_status() + payload = response.json() + return payload if isinstance(payload, list) else [payload] + + try: + for verify_len in sorted(verify_lens): + pin(int(verify_len)) + for batch_size in sorted(batch_sizes): + prompts = [[rng.randrange(1000, 30000) for _ in range(input_len)] + for _ in range(int(batch_size))] + bodies = [{"model": model, "prompt": p, "max_tokens": max_tokens, + "temperature": 0.0, "ignore_eos": True, "stream": False} + for p in prompts] + threads = [threading.Thread( + target=lambda b=b: requests.post( + f"{base_url}/v1/completions", json=b, timeout=3600), + daemon=True) for b in bodies] + for thread in threads: + thread.start() + time.sleep(settle_s) # let admission reach the plateau + drain_metrics() # discard the ramp + time.sleep(poll_s) + rows = drain_metrics() + # The length is NOT read back from telemetry: /metrics carries + # no generation-token count (iter_states reaches the log line + # only), and it does not need to -- the sweep pinned it, and + # the pin is refused by the server if it could not be honoured. + # The row count is what the step really ran; the padded width + # is resolved from the ladder when one was supplied, since a + # partly-filled batch costs what its captured shape costs. + # + # hostStepTimeMS, not iterLatencyMS: under the overlap + # scheduler the latter spans about two loops (base_worker says + # so where it publishes both), which would inflate every cell + # by roughly a factor of two. + kept = 0 + for row in rows: + rows_seen = int(row.get("numActiveRequests") or 0) + step_ms = float(row.get("hostStepTimeMS") or 0.0) + if rows_seen <= 0 or step_ms <= 0: + continue + padded_bs = rows_seen + if padded_batch_sizes: + wider = [int(b) for b in sorted(padded_batch_sizes) + if int(b) >= rows_seen] + if not wider: + continue + padded_bs = wider[0] + out.append(StepSample(batch_size=padded_bs, + verify_len=int(verify_len), + step_time_ms=step_ms)) + kept += 1 + print(f"[dspark-sps] server bs={batch_size:>5} L={verify_len} " + f"kept {kept - warmup_steps if kept > warmup_steps else 0}" + f" steps", file=sys.stderr) + for thread in threads: + thread.join(timeout=1.0) + finally: + # Never leave a served deployment pinned to a profiling window. + pin(None) + return out + + def main(argv: Optional[Sequence[str]] = None) -> int: args = _build_parser().parse_args(argv) @@ -1648,6 +1940,47 @@ def main(argv: Optional[Sequence[str]] = None) -> int: raise SystemExit("--fit-only requires --samples-in") samples = _load_samples(_str_list(args.samples_in)) provenance = {"source": "refit", "samples_in": _str_list(args.samples_in)} + elif args.from_iter_log: + ladder = (_int_list(args.cuda_graph_batch_sizes) + if args.cuda_graph_batch_sizes else None) + samples = samples_from_iter_log( + args.from_iter_log, max_draft_len=args.max_draft_len, + max_batch_size=max(ladder) if ladder else 256, + padded_batch_sizes=ladder) + if not samples: + raise SystemExit( + "no usable steps in the iteration logs. They must come from a " + "run with print_iter_log AND enable_iter_perf_stats enabled: " + "the token total lives in the `states` dict at the end of each " + "line, and without it a step cannot be filed under a cell.") + _write_samples(samples_out, samples) + print(f"[dspark-sps] {len(samples)} steps from " + f"{len(args.from_iter_log)} iteration log(s)", file=sys.stderr) + provenance = {"source": "iter_log", "logs": list(args.from_iter_log)} + elif args.base_url: + samples = samples_from_server( + base_url=args.base_url.rstrip("/"), + batch_sizes=batch_sizes, + verify_lens=verify_lens, + input_len=args.input_len, + max_tokens=max(64, int(args.measure_steps)), + model=args.served_model_name, + warmup_steps=args.warmup_steps, + settle_s=float(args.server_settle_s), + poll_s=float(args.server_poll_s), + max_draft_len=args.max_draft_len, + seed=args.seed, + padded_batch_sizes=(_int_list(args.cuda_graph_batch_sizes) + if args.cuda_graph_batch_sizes else None)) + if not samples: + raise SystemExit( + "the server sweep collected nothing. Check that /metrics " + "returns iteration stats (enable_iter_perf_stats) and that " + "/dspark/verify_len_pin exists on this build.") + _write_samples(samples_out, samples) + print(f"[dspark-sps] {len(samples)} steps from {args.base_url}", + file=sys.stderr) + provenance = {"source": "server_sweep", "base_url": args.base_url} else: if not args.model: raise SystemExit("--model is required unless --fit-only is used") diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py new file mode 100644 index 000000000000..f476cf3a938c --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py @@ -0,0 +1,189 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Reading cost-table samples out of production iteration logs. + +The in-process sweep measures an engine the profiler built itself, and when +that engine's configuration drifts from the deployment the table describes a +machine nobody is running -- a swept (bs=256, M=1536) cell read 150.2 ms +against 118.8 ms measured live at the same shape, with nothing in the +fingerprint to record the difference. ``print_iter_log`` already emits the row +count, the token total and the host step time for every step of the real +deployment, so a table can be fitted from traffic that actually ran. + +What must hold: the token total identifies the verify length exactly (no +clustering, no guessing), steps that cannot be filed honestly are dropped +rather than mislabelled, and the first step of each executor instance -- whose +host time contains the wait for the first request, measured at 62 s live -- is +never treated as a measurement. +""" + +import importlib.util +import pathlib +import sys +import tempfile + +_SPEC = importlib.util.spec_from_file_location( + "dspark_sps_profiler", + pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / + "dspark_sps_profiler.py") +_PROFILER = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _PROFILER # @dataclass needs the module registered +_SPEC.loader.exec_module(_PROFILER) + +samples_from_iter_log = _PROFILER.samples_from_iter_log +resolve_padded_shape = _PROFILER.resolve_padded_shape + +# A deployment's captured ladder, abbreviated. Padding rounds a step's rows up +# to the smallest entry at or above them, which is what turns a padded token +# total into exactly one verify length. +LADDER = [1, 2, 4, 8, 16, 32, 64, 128, 192, 256] + + +def _line(iteration, bs, step_ms, gen_tokens, rank="0"): + return (f"{rank}: [08/04/2026-19:08:13] [TRT-LLM] [I] [_torch][RANK 0] " + f"iter = {iteration}, global_rank = 0, rank = 0, " + f"num_scheduled_requests = {bs}, kv_cache_util = 0.125, " + f"currank_total_requests = 1/8, host_step_time = {step_ms}ms, " + f"prev_device_step_time = {step_ms}ms, " + f"timestamp = 2026-08-04 19:08:13, " + f"states = {{'num_ctx_requests': 0, 'num_ctx_tokens': 0, " + f"'num_generation_tokens': {gen_tokens}, " + f"'cached_kv_tokens': 63821}}") + + +def _write(lines, name="gen.log"): + # tempfile rather than pytest's tmp_path: these tests also run under the + # fixture-less in-container runner used for the GPU images. + path = pathlib.Path(tempfile.mkdtemp()) / name + path.write_text("\n".join(lines) + "\n") + return str(path) + + +def test_verify_length_comes_from_the_padded_token_total(): + """The label is read out of the log, never inferred from step times.""" + log = _write([ + _line(1, 1, 62098.0, 6), # executor's first step: dropped + _line(2, 244, 79.2, 768), # 244 rows padded to 256 -> rung-2 + _line(3, 242, 117.1, 1536), # padded to 256 -> full block + _line(4, 128, 61.0, 384), # exactly 128 rows -> rung-2 + ]) + got = sorted((s.batch_size, s.verify_len, s.step_time_ms) + for s in samples_from_iter_log( + [log], max_draft_len=5, padded_batch_sizes=LADDER)) + assert got == [(128, 2, 61.0), (256, 2, 79.2), (256, 5, 117.1)] + + +def test_the_padded_width_is_what_gets_recorded(): + """A 244-row step cost what 256 rows cost; the table is indexed by that. + + Filing it under 244 would put the same shape in as many cells as the + roster had sizes, and alpha(bs) would be fitted against a row count no + step ever ran. + """ + assert resolve_padded_shape(num_rows=244, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) == (256, 2) + + +def test_the_ladder_is_what_removes_the_ambiguity(): + """768 tokens over 128 rows is L=5 at 128, L=3 at 192, or L=2 at 256. + + Without the ladder all three are admissible and the step is dropped; with + it, padding rounds 128 up to 128 and the answer is L=5. Nearly every + low-occupancy step is in this position, which is why the ladder matters: + the first version of this collector kept 312 of 5682 production steps. + """ + ambiguous = resolve_padded_shape(num_rows=128, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256) + assert ambiguous is None + resolved = resolve_padded_shape(num_rows=128, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) + assert resolved == (128, 5) + + +def test_a_step_whose_padding_declined_is_skipped(): + """_get_padded_batch can bail; then the width the ladder implies is wrong. + + Nothing in the log says padding declined, so the only honest move is to + drop a step whose token total the padded width cannot explain. + """ + assert resolve_padded_shape(num_rows=200, num_generation_tokens=600, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) is None + + +def test_the_first_step_of_each_executor_is_dropped(): + """A build spins up more than one executor, and each numbers from 1. + + The first step of each carries the wait for its first request inside + host_step_time (62 s on a live server), which would otherwise enter the + fit as a cell measurement. + """ + log = _write([ + _line(1, 1, 200.0, 6), # KV-estimation executor + _line(2, 256, 79.0, 768), + _line(1, 1, 62098.0, 6), # real executor starts over + _line(2, 256, 79.4, 768), + ]) + times = [s.step_time_ms for s in samples_from_iter_log([log], max_draft_len=5)] + assert times == [79.0, 79.4] + + +def test_steps_beyond_the_ladder_are_skipped_not_mislabelled(): + """Filing a step under a shape it did not run is the mislabelling that + once shipped eight fictional cells with near-zero residuals.""" + log = _write([ + _line(2, 244, 79.2, 768), # padded to 256 -> kept + _line(3, 256, 300.0, 2048), # L = 7, past the block + _line(4, 300, 400.0, 1800), # more rows than the ladder holds + ]) + got = [(s.batch_size, s.verify_len) for s in + samples_from_iter_log([log], max_draft_len=5, + padded_batch_sizes=LADDER)] + assert got == [(256, 2)] + + +def test_only_the_requested_rank_is_read(): + """Every rank logs; counting all of them would multiply every cell.""" + log = _write([ + _line(2, 256, 79.2, 768, rank="0"), + _line(2, 256, 79.9, 768, rank="3"), + ]) + assert len(samples_from_iter_log([log], max_draft_len=5)) == 1 + + +def test_outlier_steps_are_bounded_out(): + log = _write([ + _line(2, 256, 79.2, 768), + _line(3, 256, 5000.0, 768), + ]) + times = [s.step_time_ms for s in + samples_from_iter_log([log], max_draft_len=5, max_step_ms=1000.0)] + assert times == [79.2] + + +def test_several_logs_are_pooled(): + """One arm gives one rung; the table needs the ladder, so pool the runs.""" + a = _write([_line(1, 1, 100.0, 6), _line(2, 256, 117.1, 1536)], + name="a.log") + c = _write([_line(1, 1, 100.0, 6), _line(2, 256, 79.2, 768)], + name="c.log") + got = sorted(s.verify_len for s in + samples_from_iter_log([a, c], max_draft_len=5)) + assert got == [2, 5] + + +def test_samples_expose_the_token_total_the_table_is_indexed_by(): + log = _write([_line(1, 1, 100.0, 6), _line(2, 256, 79.2, 768)]) + sample = samples_from_iter_log([log], max_draft_len=5)[0] + assert sample.total_verify_tokens == 768 + + +def test_a_log_without_the_states_dict_yields_nothing(): + """enable_iter_perf_stats off: fail loudly upstream, not with a bad table.""" + path = pathlib.Path(tempfile.mkdtemp()) / "bare.log" + path.write_text( + "0: [TRT-LLM] iter = 2, num_scheduled_requests = 256, " + "host_step_time = 79.2ms\n") + assert samples_from_iter_log([str(path)], max_draft_len=5) == [] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py index 1f1d9d342d20..7435e968a971 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py @@ -27,30 +27,32 @@ it shows up. """ +import importlib.util import json +import pathlib +import sys import numpy as np import pytest from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable, total_verify_tokens -from tensorrt_llm._torch.speculative.dspark_sps_profiler import ( - CellStat, - FlatCostTableError, - InertCostTableError, - InsufficientSamplesError, - StepSample, - SweepConfig, - SweepGeometryError, - aligned_steps_from_stats, - build_cost_table_payload, - check_table_is_informative, - compress_to_risers, - fit_additive_cost_model, - load_cost_table, - profitability_probe, - running_max, - summarize_cells, -) +# The profiler is a development tool under tests/microbenchmarks, not part of +# the shipped package: nothing on the serving path imports it, and keeping a +# 1.8k-line sweep driver inside tensorrt_llm/ would ship it to every user. +# Load it by path so this test follows the file rather than an import root. +_SPEC = importlib.util.spec_from_file_location( + "dspark_sps_profiler", + pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / + "dspark_sps_profiler.py") +_PROFILER = importlib.util.module_from_spec(_SPEC) +# Registered BEFORE exec: @dataclass resolves annotations through +# sys.modules[cls.__module__].__dict__, so a module executed without being +# registered raises "'NoneType' object has no attribute '__dict__'" on its +# first dataclass. +sys.modules[_SPEC.name] = _PROFILER +_SPEC.loader.exec_module(_PROFILER) + +globals().update({name: getattr(_PROFILER, name) for name in ('CellStat', 'FlatCostTableError', 'InertCostTableError', 'InsufficientSamplesError', 'StepSample', 'SweepConfig', 'SweepGeometryError', 'aligned_steps_from_stats', 'build_cost_table_payload', 'check_table_is_informative', 'compress_to_risers', 'fit_additive_cost_model', 'load_cost_table', 'profitability_probe', 'running_max', 'summarize_cells')}) #: The exact set of keys ``dspark.py::_build_verify_planner`` reads. Anything #: else in the file is ignored by the loader, so the profiler must not rely on diff --git a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py index f788854d44ce..e7908b91c002 100644 --- a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py +++ b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py @@ -150,28 +150,43 @@ def test_graph_key_uses_the_agreed_bucket_not_a_fresh_sum(): from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ CUDAGraphRunner + import types + class _R: spec_config = type("S", (), {"enable_ragged_verify": True})() + def _batch(*verify_lens): + return types.SimpleNamespace(generation_requests=[ + types.SimpleNamespace(py_verify_len=v) for v in verify_lens + ]) + r = _R() get = CUDAGraphRunner._ragged_verify_bucket + windowed = _batch(3, 2, 1) # Nothing fitted yet -> not ragged, key unchanged from the uniform one. - assert get(r, object()) is None + assert get(r, windowed) is None - # A fitted bucket is returned verbatim, regardless of the batch contents: - # the batch is not even inspected, which is the point. + # A fitted bucket is returned verbatim: the batch's own token total is + # never re-summed here, which is what makes every rank agree. r.agreed_ragged_bucket = 449 - assert get(r, object()) == 449 + assert get(r, windowed) == 449 + # ... including when that total plainly differs from the fitted value. + assert get(r, _batch(1, 1)) == 449 + + # But a request without a window means the token-major staging chain never + # ran this step, so a ragged key would replay the PREVIOUS ragged step's + # row maps -- the ragged IMA. Decline, whatever the fit published. + assert get(r, _batch(3, None, 1)) is None # Cleared between steps so a stale bucket cannot key into the wrong graph. r.agreed_ragged_bucket = None - assert get(r, object()) is None + assert get(r, windowed) is None # Config gate still wins: a non-ragged engine keeps the original key shape. r.agreed_ragged_bucket = 512 r.spec_config = type("S", (), {"enable_ragged_verify": False})() - assert get(r, object()) is None + assert get(r, windowed) is None def test_capture_publishes_the_bucket_it_shaped_the_batch_to(): From 34a4b4b4fffea1946952aeab6f9dac8220b06f42 Mon Sep 17 00:00:00 2001 From: Lance Liao Date: Wed, 5 Aug 2026 00:30:03 -0700 Subject: [PATCH 091/177] [None][doc] Review the 2026-08-04 follow-up commits: 34 confirmed findings Independent review of 72dd803141..fab6291e2d (SPS interpolation, confidence calibration, fit consistency checks, runtime verify-length pin, deployment-measured SPS table): seven review passes plus an adversarial verification pass per finding; 38 raised, 34 confirmed, 4 refuted. Highs, in suggested fix order: the table fingerprint records CLI flags instead of engine facts and hard-rejects tables measured on the very deployment they are checked against; the runtime verify-length pin is adopted only under is_distributed, so a tp_size==1 server returns 200 and never applies it -- and the sweep trusts the 200, mislabelling every cost-table cell; padded_bs_hist pollution survives the fix that claims to remove it (eager-ragged and cap-accept steps never enter the fit, whose entry is the only clearing point) and the fit publishes bucket state before its validation returns; the STS recorder pairs confidence with the wrong block's acceptance label in every scheduler mode, via three independent defects (loop-order lag, a slot-key mismatch between two unsynchronized allocators, and a stash that survives graph replays), so the fitted temperature table calibrates noise; the profiler's default sweep geometry is the regime its own comment says device-asserts, and the drain observability its comment promises does not exist. Each finding carries mechanism, concrete failure chain, and the smallest fix the reviewers could see. Several verifications executed the cited code rather than reading it. Signed-off-by: Lance Liao --- docs/dspark_review_2026-08-04_followups.md | 439 +++++++++++++++++++++ 1 file changed, 439 insertions(+) create mode 100644 docs/dspark_review_2026-08-04_followups.md diff --git a/docs/dspark_review_2026-08-04_followups.md b/docs/dspark_review_2026-08-04_followups.md new file mode 100644 index 000000000000..b99a547903cc --- /dev/null +++ b/docs/dspark_review_2026-08-04_followups.md @@ -0,0 +1,439 @@ +# Review: the five follow-up commits of 2026-08-04 evening + +Scope: `72dd803141..fab6291e2d` on `confidence_head` — + +| commit | title | +|---|---| +| `72dd803141` | [fix] Interpolate the SPS cost table between breakpoints | +| `c117be89a1` | [feat] Calibrate DSpark confidence and measure the planner's inputs | +| `50675cb181` | [fix] Make the ragged fit's consistency checks able to fail | +| `0b524435fc` | [feat] Pin the DSpark verify length at runtime | +| `fab6291e2d` | [chore] Measure the SPS table against the deployment, not a rebuild | + +Method: seven independent review passes (one per commit plus two cross-cutting +lenses: CUDA-graph capture/replay safety and overlap-scheduler +write-after-read — the two bug families the IMA series just closed), then an +adversarial verification pass per finding whose default stance was that the +finding is wrong until the code proves it. 38 findings raised, **34 confirmed, +4 refuted**; the refuted ones are omitted. Several verifications executed the +cited code paths rather than reading them (noted inline). + +Severity is scoped to the feature: nothing below corrupts served tokens +(verification guarantees them). "High" means a feature's primary artifact or +purpose silently fails. + +How to read an entry: *mechanism* is what the code does, *failure* is a +concrete input→wrong-outcome chain, *fix* is the smallest correct change the +reviewers could see. Line numbers are as of `fab6291e2d`. + +--- + +## A. STS calibration pairs confidence with the wrong labels — HIGH (as a group) + +Four findings converge on one conclusion: **the fitted temperature table — +`c117be89a1`'s primary artifact — is fit on mispaired samples in every +supported collection mode**, and nothing downstream can detect it. + +### A1. The label is one (or two) draft passes behind the confidence — deterministic + +`spec_sampler_base.py:311`, `dspark_sts.py:199` · medium + +Mechanism: one forward(i) first *verifies* the block drafted by the previous +forward (B_{i-1}; `dspark.py:822` computes the accepted count), then *drafts* +the new block B_i and stashes its confidence C_i (`dspark.py:789-791` → +`dspark_sts.py:199`, which unconditionally replaces the single mutable +`_stash_logits`). `record()` runs later, inside `update_requests` +(`spec_sampler_base.py:311`), pairing `accepted = A(B_{i-1})` with whatever +the stash holds — which is always C_i. Non-overlap loop: off by one +(`py_executor.py:4634` forward before `:4673` update). Overlap loop: off by +two (`:5132` forward before `:5139` update *of the previous batch's state*). +Because the stash overwrite happens inside the very forward that produces the +label, **no execution order leaves the correct C_{i-1} in the stash**. + +The window axis got exactly the missing treatment in the same commit +(`verify_lens_snapshot`; the `_verified_len` docstring at +`spec_sampler_base.py:226-233` names this identical hazard). The confidence +axis did not. + +Why it is invisible: the fitter (`dspark_fit_sts.py:171`, `--source auto`) +*prefers* the `logits_at_draft` column, whose docstring calls it "the correct +pairing"; the built-in stash-vs-buffer diagnostic compares two columns that +are both C_i, so they agree with each other while both are mispaired; the +82.2%-rows-differed measurement quoted in `load_shards` compares the columns +against each other, never against the label. The commit's own measurement of +one-pass-apart rows (correlation 0.14, mean |delta| 7.35) says what the fit +converges to: temperatures that flatten survival toward the marginal +acceptance rate, which then silently biases the serving planner's budget +argmax wherever the table is deployed. + +Fix: join on the draft-pass axis, not the wall-clock axis — ride the +confidence rows for block B in the SampleState produced by the forward that +verifies B (the `verify_lens_snapshot` pattern), or keep a two-deep ring +keyed by the existing draft sequence counter and have `record()` select by +the verified block's sequence. + +### A2. The recorder reads the confidence buffer with the wrong slot key + +`spec_sampler_base.py:312` · medium + +Mechanism: `record(slot=req.py_seq_slot, device_logits=self.sts_logits_buffer)` +indexes the raw `worker._confidence_logits` buffer by **py_seq_slot**, but +that buffer's rows are assigned by the DSpark worker's *private* allocator +(`_assign_slot`, a FIFO deque keyed by `py_request_id`, +`dspark.py:591-614`). `py_seq_slot` comes from the separate `SlotManager` +(set-based, immediate reuse; `resource_manager.py:2332-2344`). Nothing ever +synchronizes the two. They coincide only until the first request completes: +the set hands a freed slot back immediately while the deque defers it, so any +departure+arrival churn diverges them permanently. Every other reader +correctly goes through `confidence_row_for(py_request_id)` +(`dspark.py:319-334`); the recorder alone uses the raw index. Bounds always +hold (buffer is `[max_batch+2, K]` vs `py_seq_slot < max_batch`), so nothing +crashes — request A's accepted count is paired with request B's logits and +the shards look healthy. The late-read diagnostic reads back with the same +wrong key, so it is structurally blind to this too. The docstrings claiming +both halves are "indexed by the same py_seq_slot" +(`spec_sampler_base.py:203-206`, `dspark_sts.py:137-138`) assert an invariant +the code never establishes. + +Fix: wire a resolver instead of a raw buffer — e.g. +`self.sampler.sts_row_for = worker.confidence_row_for` at +`py_executor.py:3421-3427` and have `record()` take `py_request_id`. + +### A3. The draft stash survives CUDA-graph replay steps unbounded + +`dspark_sts.py:203` · low (verified with the full replay chain) + +Mechanism: `begin_step` invalidates `_host_logits`/`_stash_host` but never +`_stash_logits`/`_stash_slots`. Under graph replay, the Python inside the +captured draft body does not run, so the stash retains clones frozen at the +last *eager* step; `record()` re-materializes `_stash_host` from them +(`dspark_sts.py:236-244`) and appends stale-but-finite rows into +`logits_at_draft` as if they were this step's data (continuing requests keep +their slots, so lookups hit). The capture-time warning fires once during +warmup and any later eager step (mixed batch, uncaptured shape — routine +under in-flight batching) silently re-arms the stash for all subsequent +replays. `make_recorder_from_env` refuses cost-table and compact-ragged +regimes but never checks `cuda_graph_config`. + +Fix: tag the stash with the worker's draft sequence when armed +(`_confidence_stamp` already solves this class for the device buffer, +`dspark.py:781`) and treat a sequence mismatch as a missing entry. + +### A4. Recorder declines are silent, against the branch's own rule + +`dspark_sts.py:226` · high + +`record()` returns without recording on an out-of-range slot — no counter, no +log — and the capture-time stash refusal warns once but is not counted. This +module family's own documented rule (`dspark_verify.py:326-329`: even the +benign empty-batch decline is counted, because "this feature's failures are +all silent") is violated at exactly the call sites that A2/A3 make load-bearing: +a systematically wrong slot or a fully-graphed run yields empty or poisoned +shards with nothing in the stats to say why. Additionally, the construction +-time censoring guards can be bypassed mid-run: `adopt_verify_len_pin` can +lower the pin below `block_size` (or clear it) *while a collection run is +writing shards*, reintroducing the censored-label regime the constructor's +ValueError exists to refuse. + +Fix: count dropped slots and stash refusals in a recorder stats dict surfaced +next to the planner's; refuse (or count) pin adoption below `block_size` +while a recorder is active. + +--- + +## B. The runtime verify-length pin silently never applies on tp_size == 1 — HIGH + +`py_executor.py:3516` (adoption), `:3341` (endpoint), `dspark_verify.py:263` +(queue) · from `0b524435fc` + +Mechanism: the only production call to `planner.adopt_verify_len_pin` sits +inside `if is_distributed:` (requires `tp_size > 1`; `py_executor.py:3433, +:3497`) — there is no else branch. The endpoint chain +(`openai_server.py` → `base_worker.py:875` → `PyExecutor +.set_dspark_verify_len_pin`) only *queues* the pin +(`request_verify_len_pin`, `dspark_verify.py:238-247` — docstring: "NOT +applied here") and returns 200 with the validated value. +`decide_verify_lens` reads only `_forced_verify_len` (`:393`), which nothing +else sets. The `base_worker.py:880` design comment — "rank 0 alone receiving +the RPC is sufficient" because "the step's own allgather" broadcasts it — is +false at tp_size == 1: there is no allgather. + +Failure: a single-GPU profiling server (the profiler's *default*, +`--tp-size 1`) accepts POST `/dspark/verify_len_pin` with 200; the pin sits +queued forever; `forced_steps` stays 0; and the server sweep +(`dspark_sps_profiler.py:1887-1890`) explicitly trusts the 200 and never +reads the effective length back — so every cell of the produced SPS cost +table is timed at planner-chosen lengths (full block, since sweeps run +without a cost table) but *labelled* with the pinned length. The fitted +table comes out flat in L and then silently mis-drives every planner that +consumes it. This is precisely the mislabelled-cell corruption the pin +feature exists to prevent, and `test_dspark_verify_len_pin.py` exercises +only the planner in isolation, never the executor wiring. + +Two adjacent silent modes (confirmed separately): +- **Static-mode engines** accept and adopt the pin, log "set", but + `decide_verify_lens` is never called (`compute_windows` is False) — the + endpoint should 400 the way it already does for a missing planner + (`py_executor.py:3341`). +- **Queue race**: `adopt_verify_len_pin` unconditionally resets + `_pending_verify_len_pin = -1` (`dspark_verify.py:263`), so a second pin + queued by the RPC thread between the decode loop's payload read (`:3504`) + and adopt (`:3516`, an allgather apart) is wiped after its endpoint already + returned 200. Fix: compare-and-clear. + +Fix (primary): `else: planner.adopt_verify_len_pin(planner.pending_verify_len_pin())` +after the `is_distributed` block — on a single-rank world the local queue *is* +the group's agreement. Have the sweep read back the effective value. + +--- + +## C. The table fingerprint records CLI flags, not engine facts — hard-rejects correct tables — HIGH + +`dspark_sps_profiler.py:2071` (write), `dspark.py:456-462` (check, from +`c117be89a1`) · from `fab6291e2d` + +Mechanism: the fingerprint dict is built from `args` (`"ep": int(args.ep_size +or 0)`, `"tp": args.tp_size` default 1, `"attention_dp"` store_true default +False) and attached whenever the run was not `--fit-only`. But: + +- The two **new deployment-measuring sources** (`--from-iter-log`, + `--base-url`) build no engine at all, so there is no reason to pass those + flags; a table swept over HTTP against a tp8/ep8/attention-DP deployment + carries `{tp:1, ep:0, attention_dp:False}`. +- Even the in-process sweep false-rejects whenever `--ep-size` is unset: + the write side records 0, while the loader compares + `mapping.moe_ep_size`, which `Mapping` normalizes to >= 1 + (`llm_args.py:4739` maps None→−1; `mapping.py:114` resolves −1→1 — chain + verified line by line). + +`check_table_fingerprint` (`dspark_planner.py:206`, raise at `:244`) then +refuses with "profiled on a different engine configuration… Re-profile on +this configuration" — and re-running the same profiler command reproduces the +refusal. This defeats the commit's own title (measure against the +deployment: the tables measured ON the deployment are exactly the ones +rejected), and violates `check_table_fingerprint`'s own documented mechanism +for facts the consumer cannot verify: key ABSENCE (compared keys must be +present in both dicts; absent ones go to the check-manually INFO line) — the +`ep: 0` encoding fabricates a fact instead of omitting it. + +Fix: record resolved engine facts when an engine exists (read the mapping +after construction); omit the key entirely when the fact is unknown +(HTTP/log sources, unset flags). + +--- + +## D. `padded_bs_hist` pollution is not actually fixed, and the fit publishes before validating — HIGH + +`model_engine.py:3789` (clear), `py_executor.py:3633-3644` (read), +`model_engine.py:3908-3986` (publish-before-validate) · from `50675cb181` + +Mechanism (pollution): the commit clears `_dspark_last_padded_bs` at +`fit_ragged_verify_lens` **entry** and its comment claims this prevents a +stale value from polluting "exactly the fallback steps the histogram exists +to distinguish". But repo-wide the attribute has three references: cleared at +fit entry, set on fit success, and **read unconditionally into +`record_step` on every stats-bearing step** (`py_executor.py:3640`). The +eager-ragged branch (`:3582-3599`) never enters the fit (it clears +`agreed_ragged_bucket` and `ragged_pad_verify_len` — but not this), and the +cap-accept branch never reaches the fit either. Both produce non-empty +`verify_lens`, so `record_step` passes its early-out and +`dspark_observability.py:327-328` books the LAST fitted step's padded row +count into `padded_bs_hist` for a step that was never padded. Eager-ragged +steps are frequent by the code's own measurement ("46 of the first 64 +steps"), so the histogram — added by this very range to reconcile the +negative-trim_ratio bug — can be dominated by phantom entries on exactly the +steps it exists to explain. Execution behavior is unaffected (the ceiling +rebase requires `bucket`), so this is observability-only — but this branch's +stated stance is that these counters are the sole detection channel for the +feature's silent failures. + +Fix (total, verified): pass `padded_bs` to `record_step` only when +`bucket is not None` — bucket is non-None exactly when the fit delivered. + +Mechanism (publish-before-validate): inside the fit, +`_dspark_last_total_tokens` (`:3908`), `_dspark_last_padded_bs` (`:3914`), +`runner.agreed_ragged_bucket` (`:3928`) and `runner.ragged_pad_verify_len` +(`:3971`) are all published **before** three validation-failure returns +(`:3960`, `:3969`, `:3986`). A failed fit therefore leaves every one of them +stale-set while the step falls back to uniform. A stale published bucket is +the exact state family behind the just-fixed IMA; today the wrong-key replay +is prevented only by the window check in `_ragged_verify_bucket` plus the +`py_verify_len` clearing — one line of defense where there should be two. +The layout assert's uniform branch (`model_engine.py:4076-4085`) does not +check `agreed_ragged_bucket is None`, so the stale state is invisible even +with `TLLM_DSPARK_ASSERT_LAYOUT=1`. + +Fix: move all four publications below the last failure check, next to the +`py_verify_len` stamping. + +Related (medium): the commit's central rationale — that the old assert was a +tautology — is false against its own parent: at `50675cb181^`, +`_ragged_verify_bucket` already returned `agreed_ragged_bucket` rather than +re-summing, so the old assert already compared the batch walk against the +fit's allgathered value and could fire. The change is a no-op refactor, and +the comment at `model_engine.py:4103-4109` encodes the incorrect history. +Rewrite it to state what the check actually does. + +--- + +## E. Profiler: the default sweep geometry hits a device-side assert, and the promised drain observability does not exist — HIGH (×2) + +From `72dd803141` / `fab6291e2d`, `tests/microbenchmarks/dspark_sps_profiler.py`. + +### E1. Defaults trap: `block_follows_verify_len=False` runs the asserting regime + +`:1036` (default), `:1219-1223` (the commit's own comment), `:1581-1585` +(flag help) · high + +The default sweep geometry is constant-block; the same commit's comment says +that regime "currently asserts (Repeat.cu output_size vs sum(repeats))" +whenever `max(verify_len) < block_size` — which is **every swept cell except +L == max_draft_len** under the default `verify_lens = 1..max_draft_len` +(compact ragged is forced on by `_prepare_environment`, +`TLLM_DSPARK_RAGGED_VERIFY_MODE=COMPACT`). `SweepConfig.validate()` — whose +docstring is literally "catch the geometry mistakes that would otherwise +fail mid-sweep" — does not check it. A user running the documented default +builds an engine for minutes, then dies on a CUDA assert at the first +L < block cell. Verification also checked the timeline: the "currently +asserts" comment was written hours *after* the IMA fix landed and was +re-worded at HEAD, so it is not stale; and no recorded sweep artifact +(`dspark-runs/sps_samples_L*.jsonl`, `sps_real_final.json`) has ever run the +constant-block geometry to completion. + +Fix: either fix the Repeat.cu packing for `max(lens) < block_size` and delete +the comment, or flip the default to `block_follows_verify_len=True` and make +`validate()` refuse the known-asserting combination. + +### E2. The mid-cell drain promises a counter that does not exist, and the drain can latch shut + +`:1344` (comment), `executor/result.py` (`IterationResult.get_results`) · high + +The comment says "the count of empty reads is reported so a cell that lost +its tail this way is visible rather than merely short" — grep of the file at +the commit and at HEAD finds only the comment; the drainer reports only +exception counts. Worse, on the in-process path `IterationResult.get_results` +latches `_done = True` on the first Empty and only `mark_undone()` (new +prompt submission) resets it — one longer-than-timeout gap in stats +production makes **every subsequent `get_stats` call, including the final +post-cell drain, return `[]` instantly**. Cells are silently truncated with +no signal, in a tool whose output feeds the planner. + +Fix: implement the promised empty-read counter and surface it per cell; +`mark_undone()` (or bypass the latch) before the final drain. + +--- + +## F. Interpolation switch: three loose ends — MEDIUM + +From `72dd803141`. + +1. **Truncation eats shelf-close points first** + (`dspark_sps_profiler.py:756`, medium — verified by executing the numeric + example). `compress_to_risers` now encodes a flat shelf as a breakpoint + PAIR (required by the interp consumer), but the `max_breakpoints` + truncation drops the interior point with the smallest incoming jump — + which is by construction always a shelf-CLOSE point (its jump is below + `min_riser_ms`), never a real riser. Dropping it turns the measured shelf + into a ramp: with counts `[16..128]` and a 2.0→3.0 riser, cap=5 bills + M=80 (measured 2.004) at interp 2.5 — the mid-shelf over-bill the + function's own docstring warns about. The pair encoding also roughly + doubles kept points while the default cap stayed 8, so truncation now + fires on sweeps that never used to hit it — silently (no warning, no + `_meta` note), and the only truncation test uses `min_riser_ms=0.0` so + the interplay is untested. Fix: drop by introduced interpolation error + (or treat a shelf pair as atomic), warn when truncation fires, add a test + with `min_riser_ms > 0`. + +2. **No format marker for the semantic reversal** + (`dspark_planner.py:113`, low). A table emitted by the OLD compressor + (which dropped shelf-closing breakpoints *by design*, for the floor + consumer) is silently reinterpreted as a continuous ramp by the new + interp consumer — the "over-trim mirror" the new code names but cannot + detect. Old tables lack `_meta.engine` so they get only the generic + missing-fingerprint warning. Fix: stamp `lookup: "interp"` (or a format + version) into the payload meta; warn-or-refuse tables without it. + +3. **`fixed_overhead_ms` escapes the finiteness validation the commit + message claims** (`dspark_planner.py:148`, low). The check is + `if self.fixed_overhead_ms < 0.0`, which NaN and +inf both pass; a + non-finite overhead silently collapses the planner to maximal trimming. + Fix: `if not math.isfinite(...) or ... < 0.0: raise`. + +--- + +## G. Documentation drift (new feature doc contradicts shipped code) — LOW + +`docs/source/features/dspark-confidence-scheduling.md`, from `50675cb181`: + +- `:195` — asserts the cost lookup is a staircase that "绝不插值", the exact + opposite of the code as of three commits earlier (`72dd803141`), while the + commit message claims the doc "caught up with the code". +- `:261-268` and `:283` vs `:307` — self-contradictory on the window floor: + two places say `min_verify_len` defaults to 0 (with a worked example headed + "min_verify_len=0（默认）"), one place (and the code) says the floor is + fixed at 1. +- `:400` — states in present tense that draft_len "is not in the + attention-DP consistency allgather", documenting as live a hazard the + shipped shape gate already closed (it allgathers and compares draft_len). + +Fix: one editing pass against the code; recast the `:400` diagram as the +motivating failure. + +--- + +## H. Profiler tooling correctness (deployment-measurement paths) — MEDIUM/LOW + +From `fab6291e2d`, `tests/microbenchmarks/dspark_sps_profiler.py`: + +1. `:1778` (low, verified by executing the parser on a real mixed-step log + line) — `samples_from_iter_log` files mixed context+generation steps + under a wrong `(padded_bs, verify_len)` cell with prefill-inflated eager + step times; `num_ctx_requests` is on the same log line and ignored. The + ladder-resolution path converts ambiguous-therefore-dropped into + confidently-mislabelled. Fix: capture `num_ctx_requests` and skip ctx>0 + steps, counting the drops. +2. `:1966` (medium) — the server sweep's load duration + (`max_tokens = max(64, measure_steps)`) is uncoupled from the measurement + window (`settle_s + poll_s`); under defaults the synthetic requests can + finish before the window opens, collecting nothing or ramp-down rows + filed under smaller batch sizes. Fix: derive load length from the window, + or re-post until the drain completes. +3. `:1878` (medium) — the completion-POST worker threads discard responses + and swallow all errors; a wrong `--served-model-name` produces an idle + server, empty cells, and a SystemExit blaming `/metrics` and + `enable_iter_perf_stats`. Fix: check status codes, surface the first + failure in the error message. +4. `:1866` (low) — a pin refused mid-sweep raises out of the whole sweep and + discards every already-collected sample (`_write_samples` runs only after + the sweep returns); with default `verify_lens = 1..max_draft_len` against + a server whose ladder lacks a rung, a refusal is near-certain. Fix: catch + per rung, count skips, write samples incrementally. +5. `:1917` (low) — `warmup_steps` is accepted by the server path but never + applied; the progress line prints `kept - warmup_steps` while all `kept` + rows were appended. Fix: apply it or drop it. +6. `:1788` (low) — `rank_prefix='0:'` silently drops all lines of logs + without an mpirun/srun rank tag (single-process deployments); not exposed + on the CLI; the error blames `print_iter_log`. Fix: fall back to + prefix-less matching when zero lines match; expose `--rank-prefix`. +7. `:1473` (medium) — the moved file still advertises the old module path in + its own `--help` prog string and docstring run command; copy-pasting the + tool's documented invocation fails with ModuleNotFoundError (the commit + updated the `llm_args.py` error message but not the tool itself). + +--- + +## Suggested order of attack + +C → B → D → A → E1/E2 → F → H → G. C and B block anyone producing a usable +table with the new tooling and are one-liners; D's histogram fix is a +one-line condition plus a code motion; A decides whether the calibration +feature produces a meaningful artifact at all and is the largest change +(snapshot-per-block pairing + request-id keyed reads); E1 is a defaults flip +or a kernel fix; the rest are local. + +A recurring theme worth internalizing: most of these violate the branch's +own explicitly documented rules — "every silent decline must be counted" +(A4, B, E2, H), "facts the consumer cannot see must be absent, not +fabricated" (C), "clear stale published state so it cannot key a later step" +(D) — rather than rules imported from outside. The rules are right; the new +code just needs to keep following them. From 8a1a6046db162fab7ad351408aa1014eff03ce13 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 5 Aug 2026 01:21:38 -0700 Subject: [PATCH 092/177] [None][fix] Address the 2026-08-04 review: repair STS pairing end to end Fixes the 34 confirmed findings from docs/dspark_review_2026-08-04_followups.md (34a4b4b4ff). The largest is group A: STS calibration paired every confidence row with the wrong acceptance label, in every supported collection mode. STS pairing (A1-A4). The label for block B_i arrives while pass i+1 (i+2 under the overlap scheduler) has already overwritten both the shared confidence buffer and the recorder's single stash -- the forward that produces the label is the same forward that clobbers the stash, so NO execution order paired them correctly. The join key was also wrong: py_seq_slot belongs to an allocator the confidence buffer never used, and the two numberings coincide only until the first request completes. The recorder now keeps a stamped ring of host snapshots keyed by the worker's draft-pass counter, staged on the executor's host path next to the planner's own relay (never inside capture, which is what silently froze the old stash under graph replay); record() selects the ring entry whose per-row stamp equals the pass snapshotted into the SampleState at sampling time (the verify_lens_snapshot pattern), resolves the row through the worker's own allocator, and counts every decline. Shards carry a pairing-provenance marker and the fitter refuses shards without it: a mispaired shard fits a temperature vector that calibrates noise, and nothing downstream can tell. The runtime pin endpoint refuses sub-block pins while a recorder is active, closing the mid-run bypass of the constructor's censoring guard. The rest, one line where the fix is small: - B: adopt the verify-len pin on the non-distributed path too; a tp1 server accepted pins over HTTP and never applied one. - C: the engine fingerprint reads tp/ep from the resolved Mapping rather than raw args that fabricate ep=0. - D: the four ragged-fit publications sit below the last validation return, so a declined fit no longer leaves half-published state; the fallback path clears the agreed bucket instead of inheriting a stale one. - E1: the profiler refuses constant-block geometry that verifies below the block (the regime that trips the Repeat.cu output_size assert). - E2: empty /metrics reads are counted, and the latched IterationResult is un-latched so the second poll is not silently empty. - F1: over-budget tables drop the breakpoint interpolation reconstructs best, not the shelf-closing point the old smallest-jump rule always chose. - F2/F3: tables carry a lookup-semantics marker; fixed_overhead_ms is required finite. - G: the feature doc no longer claims staircase lookup or min_verify_len=0. - H: collector fixes -- ladder-resolved padded shapes, warmup applied rather than printed, window-coupled token budgets, surfaced POST errors. Two more, found while validating the fixes against a live five-rung server sweep (job 2585129): - summarize_cells takes the requested grid: thin cells the client's admission ramp creates in passing are dropped and reported instead of killing the fit; a thin REQUESTED cell still refuses. - the server sweep fans its load out by dp_size: under attention DP the planner prices each rank's own batch and the iteration log reports per-rank rows, so a cell labeled bs needs bs*dp_size requests in flight. Unaccounted for, a tp8 sweep asked for 64..256 and measured per-rank batches of 8..40. 233 hw-agnostic DSpark unit tests green, including new ring-pairing tests and a real-log collector smoke against two production runs. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../features/dspark-confidence-scheduling.md | 21 +- .../_torch/pyexecutor/model_engine.py | 65 ++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 90 ++++- tensorrt_llm/_torch/speculative/dspark.py | 45 ++- .../_torch/speculative/dspark_planner.py | 4 +- tensorrt_llm/_torch/speculative/dspark_sts.py | 224 ++++++----- .../_torch/speculative/dspark_verify.py | 7 +- .../_torch/speculative/spec_sampler_base.py | 42 +- tensorrt_llm/serve/openai_server.py | 26 +- tests/microbenchmarks/dspark_fit_sts.py | 82 ++-- tests/microbenchmarks/dspark_sps_profiler.py | 365 ++++++++++++++---- .../hw_agnostic/test_dspark_sps_profiler.py | 47 ++- .../test_dspark_sts_ring_pairing.py | 168 ++++++++ 13 files changed, 872 insertions(+), 314 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py diff --git a/docs/source/features/dspark-confidence-scheduling.md b/docs/source/features/dspark-confidence-scheduling.md index c7ee361ab823..6200b0d49bbc 100644 --- a/docs/source/features/dspark-confidence-scheduling.md +++ b/docs/source/features/dspark-confidence-scheduling.md @@ -192,7 +192,7 @@ T(bs, K) = bias + alpha(bs) + theta(M) 权重搬运） ``` -`SpsCostTable` 按总验证 token 数 `M` 做**阶梯**查表，绝不插值： +`SpsCostTable` 按总验证 token 数 `M` 做**钳位线性插值**查表（对齐 SGLang 的加法表消费端）。想表达"平台"就必须把平台测出来——两个等值断点；消费端不再替硬件假设断点之间是平的（那个假设曾把 1512 token 的满块按 768 的价格计费，整个特性因此从未裁剪过）： ```text step_time @@ -258,14 +258,12 @@ right_edge(breakpoint T) = (T - 1) // bs - 1 # length 空间 预算是 `num_gen * (L* - min_verify_len)` —— 和 uniform 决策**花掉的 token 总量相同**， 只是重新分配。 -`min_verify_len` 默认为 **0**：整个 block 的每个位置都参与竞争，draft 已经死掉的请求会拿到 -**0 个** draft 位置。这不是饿死 —— bonus token 不在 confidence 矩阵里（它不是 draft 出来的、 -必然被接受），所以该请求那一步照样出 1 个 token，只是退化成普通自回归解码。也不存在 -"不验证 → 拿不到新分数 → 永远不验证"的反馈陷阱：不变量 I1 保证 block 永远完整 draft， -confidence head 对全部 K 个位置打分，与验几个无关。 - -把 floor 调到 0 以上只有一个正当理由 —— **公平性**（保证每个请求每步至少有一次投机机会） -—— 并且要付代价：floor 会把预算花在排名低于被它挤掉的那些位置上。 +`min_verify_len` 在本实现中**固定为 1**（配置校验强制 >= 1）：每个请求每步至少验证 +anchor 位置。SGLang 允许 floor 为 0——draft 已死的请求拿 0 个位置也不是饿死（bonus token +不在 confidence 矩阵里，该请求照样出 1 个 token，退化为普通自回归），省下的预算可以喂给 +强请求的深位置，混合批下其可达 Θ 严格更优（test_starving_weak_rows_is_a_known_gap 钉死了 +这个差异）。移植 floor=0 需要 executor 支持零窗口请求（layout/采样路径改动），是已知的 +待决 gap 而非本实现的默认。 ```python candidates = survival[:, floor:].flatten() # 每一个 (request, position) 对 @@ -280,7 +278,7 @@ confidence，也必须选出相同的 verify 长度，否则它们的 batch 形 ## 一个完整的数值例子 -`K=7, bs=4, tiers=[1,3,7], min_verify_len=0`（默认），成本表 +`K=7, bs=4, tiers=[1,3,7], min_verify_len=0`（注意：这是 SGLang 语义的演示例；本实现 floor 固定为 1），成本表 `token_counts=(0, 8, 16, 32)`、`step_time_ms=(2.0, 2.2, 2.6, 9.0)`。 survival 矩阵： @@ -397,7 +395,8 @@ capture 安全性依赖几个具体选择： draft_len 是 CUDA graph key 的一部分 | v -但它不在 attention-DP 的一致性 allgather 里 +但它（当时）不在 attention-DP 的一致性 allgather 里——这是促成现行 shape gate 的 +历史事故：现在的实现已把 draft_len 纳入 allgather 并逐 rank 比对，下图保留为动机说明 | v 两个 rank 选了不同长度 -> 选到不同的 graph diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index cb8cf8586df4..28bac660142f 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3903,29 +3903,6 @@ def fit_ragged_verify_lens( f"to uniform scheduling: {exc}") return None padded_bs, bucket = shape.padded_bs, shape.bucket - # Remembered so a graph miss can report the shape it actually submitted - # rather than just that one happened. - self._dspark_last_total_tokens = int(bucket) - # And the ROW count behind that bucket. The observability ceiling is - # booked over the rank-local, unpadded request count while `delivered` - # is booked over this padded, cross-rank-maximum one, which is how - # trim_ratio went negative in four runs. Only the caller can reconcile - # them, and it cannot do that without this number. - self._dspark_last_padded_bs = int(padded_bs) - # The fit's agreed value, which the graph key carries verbatim so - # every attention-DP rank keys on a number they agreed on; the layout - # assert cross-checks it against the batch walk. Both used to be - # derived independently WITHOUT that check -- - # the key walked generation_requests *after* padding was - # appended, while the fit worked from the allgathered peer - # stats -- so their agreement across attention-DP ranks was - # incidental rather than structural, and when they diverged the - # shape gate dropped every rank out of replay - # (peer_shape_mismatch: 2/265 steps without trimming, 16/318 - # with). This value comes from data every rank allgathered and - # rules every rank runs identically, so it agrees by - # construction. - runner.agreed_ragged_bucket = int(bucket) # The CUDA-graph padding rows are appended later by _get_padded_batch, # which appends a *single shared dummy object*, so every pad row @@ -3967,8 +3944,6 @@ def fit_ragged_verify_lens( f"outside [{floor_tokens}, {real_capacity}]; falling back " f"to uniform scheduling") return None - # Communicated to _get_padded_batch, which stamps it on the dummy. - runner.ragged_pad_verify_len = pad_len - 1 lens_tensor = torch.tensor(token_lens, dtype=torch.int32) layout = RaggedVerifyLayout.from_verify_lens( @@ -3985,6 +3960,22 @@ def fit_ragged_verify_lens( f"uniform scheduling: {exc}") return None + # Published only now, past the last way this fit can fail: a stale + # published bucket on a fallback step is the state family behind the + # ragged IMA, and clearing at entry protects the next call, not the + # remainder of THIS one. The graph key carries `agreed_ragged_bucket` + # verbatim so every attention-DP rank keys on the value they + # allgathered, and the layout assert cross-checks it against the + # batch walk. `_dspark_last_total_tokens` lets a graph miss report the + # shape it submitted; `_dspark_last_padded_bs` is the row count the + # stats need to book a ceiling comparable to `delivered`; + # `ragged_pad_verify_len` is stamped on the shared dummy by + # _get_padded_batch. + self._dspark_last_total_tokens = int(bucket) + self._dspark_last_padded_bs = int(padded_bs) + runner.agreed_ragged_bucket = int(bucket) + runner.ragged_pad_verify_len = pad_len - 1 + for request, tokens in zip(generation_requests, filled.verify_lens.tolist()): request.py_verify_len = int(tokens) - 1 @@ -4074,6 +4065,17 @@ def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, verify_lens = self._ragged_token_lens( scheduled_requests.generation_requests) if verify_lens is None: + # A uniform step must not be carrying a published ragged bucket: + # that is exactly the stale state a failed fit used to leave + # behind, and a later ragged step would key a graph the batch + # does not describe. Checked here because the ragged branch below + # never runs for uniform steps, which made the staleness invisible + # even with this assert enabled. + agreed_stale = getattr(self.cuda_graph_runner, + "agreed_ragged_bucket", None) + assert agreed_stale is None, ( + f"uniform step carries a stale published ragged bucket " + f"{agreed_stale}; a failed fit exited without clearing it") # Uniform step: assert the converse, that nothing ragged leaked # through. A half-published layout is the dangerous state. assert getattr(spec_metadata, "total_verify_tokens", None) is None, ( @@ -4100,13 +4102,12 @@ def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, f"tokens but spec_metadata.total_verify_tokens is {from_spec}; " f"acceptance would slice the target logits at the wrong offsets.") if agreed is not None: - # Two INDEPENDENT derivations of the token axis: `from_requests` - # walks this batch's windows, `agreed` is the fit's allgathered - # arithmetic, and the graph key carries the latter. Comparing the - # key's own return value instead (what an earlier form did) is a - # tautology -- _ragged_verify_bucket publishes `agreed` -- so a fit - # that budgeted for padding rows the batch never got would replay a - # graph captured at a different width with nothing to say so. + # Two derivations of the token axis: `from_requests` walks this + # batch's windows, `agreed` is the fit's allgathered arithmetic, + # and the graph key carries the latter. If the fit reserved for + # padding rows that were then not appended, the replayed graph is + # a different width than the batch -- and nothing downstream says + # so, because acceptance's output_size was frozen at capture. assert from_requests == int(agreed), ( f"ragged layout mismatch: the batch carries {from_requests} " f"generation tokens but the fit agreed on {agreed}, which is " diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index b0fc76337db5..ab1fd46046fc 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3360,13 +3360,38 @@ def set_dspark_verify_len_pin(self, verify_len: Optional[int]) -> Optional[int]: ValueError: if the length is outside the ladder (see :meth:`DSparkVerifyPlanner.validate_verify_len_pin`). """ - worker = getattr(self.model_engine, "_get_spec_worker", None) - planner = getattr(worker(), "verify_planner", None) if worker else None + worker_fn = getattr(self.model_engine, "_get_spec_worker", None) + worker = worker_fn() if worker_fn else None + planner = getattr(worker, "verify_planner", None) if worker else None if planner is None: raise RuntimeError( "no DSpark confidence planner on this executor; the verify-len " "pin only applies to a DSpark engine with confidence " "scheduling enabled") + mode = getattr(worker, "ragged_verify_mode", None) + computes = (mode.computes_windows if mode is not None else + getattr(self.model_engine.spec_config, + "enable_ragged_verify", False)) + if not computes: + # Static mode never calls decide_verify_lens, so an accepted pin + # would be adopted, logged as set -- and never influence a step. + # Answering 200 there is the mislabelled-sweep trap all over again. + raise RuntimeError( + "this engine runs static (uniform) verification; the " + "verify-length pin only affects window-computing modes, so " + "pinning here would be accepted and silently never applied") + if (getattr(worker, "sts_recorder", None) is not None + and verify_len is not None + and int(verify_len) < int(self.model_engine.max_draft_len)): + # The collector's own guard refuses to START under a censoring + # window; this closes the runtime door that guard cannot see. A + # sub-block pin clips the acceptance label of every sample + # collected after it, and the shards keep looking healthy. + raise RuntimeError( + "STS calibration collection is active on this engine; a pin " + "below the full block censors the acceptance label for every " + "sample collected after it. Pin the full block, clear the " + "pin, or restart without the collector.") return planner.request_verify_len_pin(verify_len) def _dspark_confidence_draft_len(self, @@ -3416,15 +3441,17 @@ def _dspark_confidence_draft_len(self, and getattr(self.sampler, "acceptance_stats", None) is None): self.sampler.acceptance_stats = stats # Same wiring, same reason, for STS collection: the logits live on the - # worker and the accepted count is only host-side in the sampler, and - # `py_seq_slot` indexes both. None on every path that is not collecting. + # worker and the accepted count is only host-side in the sampler. The + # sampler gets the worker's OWN row resolver and pass counter, never a + # raw buffer: the recorder joins by (draft pass, buffer row), and both + # keys must come from the allocator that wrote the buffer. if (self.sampler is not None and getattr(self.sampler, "sts_recorder", None) is None): recorder = getattr(worker, "sts_recorder", None) if recorder is not None: self.sampler.sts_recorder = recorder - self.sampler.sts_logits_buffer = getattr( - worker, "_confidence_logits", None) + self.sampler.sts_row_for = worker.sts_row_for + self.sampler.sts_seq_provider = worker.verified_draft_seq mode = getattr(worker, "ragged_verify_mode", None) compute_windows = (mode.computes_windows if mode is not None else getattr(self.model_engine.spec_config, @@ -3536,6 +3563,13 @@ def _dspark_confidence_draft_len(self, int(payload[1]), int(payload[2]), 1 if all_can_graph else 0 ] for payload in payloads] + else: + # A single-rank world has no allgather to carry the pin, and + # without this branch a queued pin sat forever: the endpoint + # returned 200, forced_steps stayed 0, and a sweep labelled every + # cell with a length the planner never ran. The local queue IS the + # group's agreement when the group is one rank. + planner.adopt_verify_len_pin(planner.pending_verify_len_pin()) # --- Phase 3: no more collectives ------------------------------------- # @@ -3633,13 +3667,17 @@ def _dspark_confidence_draft_len(self, stats.record_step(num_gen_requests=num_gen_requests, verify_lens=fitted, bucket=bucket, - # Was hardcoded None, which both left - # padded_bs_hist empty in every summary ever - # logged and denied record_step the row count it - # needs to book a comparable ceiling. - padded_bs=getattr( - self.model_engine, - "_dspark_last_padded_bs", None), + # Only when the fit delivered: bucket is non-None + # exactly then. Reading the attribute on every + # step booked the LAST fit's padded row count + # into padded_bs_hist on eager-ragged and + # cap-accept steps -- phantom entries on exactly + # the steps the histogram exists to explain, and + # clearing at fit entry cannot help the branches + # that never enter the fit. + padded_bs=(getattr(self.model_engine, + "_dspark_last_padded_bs", None) + if bucket is not None else None), fallback=fallback_reason, delivered=delivered) @@ -3662,6 +3700,15 @@ def _dspark_confidence_draft_len(self, # layout with uniform acceptance -- silent token misattribution. for request in gen_requests: request.py_verify_len = None + # And the published fit state with them: a fallback step that + # inherits the previous ragged step's agreed bucket carries a + # claim about a batch it does not describe. The window guard in + # the graph key tolerates that today; the layout assert's uniform + # branch does not, on purpose -- one line of defense where there + # were zero. + if self.model_engine.cuda_graph_runner is not None: + self.model_engine.cuda_graph_runner.agreed_ragged_bucket = None + self.model_engine.cuda_graph_runner.ragged_pad_verify_len = 0 # Verify the full drafted block. This is the only fallback: the # uniform tier ladder that used to sit between "schedule per # request" and "verify everything" is gone, so there is nothing to @@ -3682,10 +3729,19 @@ def _dspark_confidence_draft_len(self, if bump is not None: staged_seq = bump() stamps_fn = getattr(worker, "confidence_stamp_buffer", None) - planner.stage_confidence( - buffer, - stamps=stamps_fn() if stamps_fn is not None else None, - staged_seq=staged_seq) + stamps = stamps_fn() if stamps_fn is not None else None + planner.stage_confidence(buffer, stamps=stamps, + staged_seq=staged_seq) + # The STS recorder keeps its own ring of these snapshots: its + # labels arrive one-plus steps later and must be paired with the + # pass that DRAFTED the block being labeled, not with whatever the + # live buffer holds by then. Same relay, same stamps, staged in + # the same breath as the planner's copy. + recorder = getattr(worker, "sts_recorder", None) + if recorder is not None: + recorder.stage_snapshot(device_logits=buffer, + device_stamps=stamps, + staged_seq=staged_seq) return int(runtime_draft_len) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 171bcd76e5ab..75d7738c2b05 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -333,6 +333,32 @@ def confidence_row_for(self, req_id: int) -> int: return self._neutral_conf_row return self._req_to_slot.get(req_id, self._neutral_conf_row) + def sts_row_for(self, req_id: int) -> Optional[int]: + """``confidence_row_for`` for the STS recorder: None, never neutral. + + The planner WANTS the neutral fallback (it reads as "verify the full + block", the safe direction). The recorder must not get it: the neutral + row is a real buffer row that no draft ever writes, and pairing a + request's label against it would be a fabricated sample, not a safe + default. An unresolvable request is a dropped sample, counted by the + recorder as ``no_row``. + """ + if self._confidence_logits is None: + return None + return self._req_to_slot.get(req_id) + + def verified_draft_seq(self) -> Optional[int]: + """The draft pass that produced the block the CURRENT step verifies. + + Read at sampling time: by then this step's ``bump_draft_seq`` has + already advanced the counter for the draft running alongside the + verification, so the pass being verified is the previous value. The + first pass is numbered 1 (the stamp buffer starts at zeros, so a + sequence of 0 would false-match every never-written row -- including + the neutral one), hence None until a pass numbered >= 1 exists. + """ + return self._draft_seq_host - 1 if self._draft_seq_host > 1 else None + @property def max_draft_len(self) -> int: return self.spec_config.max_draft_len @@ -454,6 +480,15 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> # argmax. Omitting them says "step_time_ms already measures whole # steps at the deployment's batch size". from .dspark_planner import check_table_fingerprint + if (raw.get("_meta") or {}).get("lookup") != "interp": + logger.warning( + f"DSpark cost table {table_path} carries no " + f"lookup='interp' marker: it may predate the " + f"interpolating consumer. Tables written for the old " + f"floor lookup dropped shelf-closing breakpoints on " + f"purpose, and interpolating across those gaps re-bills " + f"every mid-shelf total upward. Re-emit the table with a " + f"current profiler if trimming behaves oddly.") check_table_fingerprint(payload=raw, live={ "tp": int(getattr(self.mapping, "tp_size", 0) or 0), "ep": int(getattr(self.mapping, "moe_ep_size", 0) or 0), @@ -779,16 +814,6 @@ def _draft_gen_block_batched( self._confidence_logits[slots] = confidence.detach() # Same slots, same capture region: content-only update, graph-safe. self._confidence_stamp[slots] = self._draft_seq_cuda - # STS collection reads this buffer from inside the sampler, which - # under the overlap scheduler can run AFTER the next step's draft - # has already overwritten it. SGLang avoids the question by keeping - # what the draft produced (`last_confidence_raw`) rather than - # re-reading a shared buffer later. Stash the same thing here, so - # the recorder can pair against the step that actually produced it - # -- and, for one diagnostic run, compare the two sources. - if getattr(self, "sts_recorder", None) is not None: - self.sts_recorder.stash_draft_confidence( - logits=confidence.detach(), slots=slots) return block_logits def _forward_impl( diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py index 011352dfb514..494b7354f3c7 100644 --- a/tensorrt_llm/_torch/speculative/dspark_planner.py +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -145,8 +145,8 @@ def __post_init__(self) -> None: raise ValueError("token_counts must be strictly increasing") if any(not math.isfinite(t) or t <= 0.0 for t in self.step_time_ms): raise ValueError("step_time_ms entries must be positive and finite") - if self.fixed_overhead_ms < 0.0: - raise ValueError("fixed_overhead_ms must be >= 0") + if not math.isfinite(self.fixed_overhead_ms) or self.fixed_overhead_ms < 0.0: + raise ValueError("fixed_overhead_ms must be >= 0 and finite") if len(self.batch_sizes) != len(self.batch_overhead_ms): raise ValueError( f"batch_sizes ({len(self.batch_sizes)}) and batch_overhead_ms " diff --git a/tensorrt_llm/_torch/speculative/dspark_sts.py b/tensorrt_llm/_torch/speculative/dspark_sts.py index 4e5779ca44c4..f7c2424db524 100644 --- a/tensorrt_llm/_torch/speculative/dspark_sts.py +++ b/tensorrt_llm/_torch/speculative/dspark_sts.py @@ -47,6 +47,7 @@ import torch +from ..._utils import prefer_pinned from ...logger import logger __all__ = [ @@ -132,17 +133,32 @@ def resolve_confidence_head(draft_model): class DSparkStsRecorder: """Pair each drafted block's logits with how much of it was accepted. - The two halves live in different places and are joined by sequence slot: - the logits are in the worker's persistent ``[max_batch + 2, K]`` device - buffer, written by the draft pass; the accepted count is host-side in the - sampler. Both are indexed by ``py_seq_slot``, so the join is exact. - - One host copy per step, not per request. ``record`` is called once per - request inside the sampler's loop, and copying a row each time would be a - synchronization per request; instead the whole buffer is pulled once when - the step's first request arrives and reused for the rest. + The two halves live in different places and on different clocks: the + logits sit in the worker's persistent slot-indexed device buffer, written + by draft pass ``i``; the accepted count arrives in the sampler while pass + ``i+1`` (or ``i+2``, under the overlap scheduler) has already overwritten + that buffer. The join is therefore made on TWO axes, neither of which is + wall-clock arrival: + + * **time** -- ``stage_snapshot`` keeps a small ring of host copies keyed + by the worker's draft sequence counter, and ``record`` selects the ring + entry whose per-row stamp equals the pass that drafted the block being + verified; + * **row** -- the caller resolves the request's buffer row through the + worker's own allocator (``confidence_row_for``), never through + ``py_seq_slot``, which belongs to a different allocator and only + coincides until the first request completes. + + A pair that cannot be made exactly is dropped and counted in ``stats`` + rather than approximated: a mispaired shard is strictly worse than a + smaller one, and looks identical downstream. """ + #: Ring depth. The pairing lag is one draft pass (two under the overlap + #: scheduler); four snapshots is comfortably past both without letting a + #: wrong target_seq silently match ancient content. + _RING_DEPTH = 4 + def __init__(self, *, path_stem: str, block_size: int, rank: int = 0, flush_every: int = _DEFAULT_FLUSH_EVERY) -> None: self.path_stem = str(path_stem) @@ -150,98 +166,106 @@ def __init__(self, *, path_stem: str, block_size: int, rank: int = 0, self.rank = int(rank) self.flush_every = max(1, int(flush_every)) self._logits: List[torch.Tensor] = [] - self._stashed_logits: List[torch.Tensor] = [] self._accepted: List[int] = [] - self._stash_logits = None - self._stash_slots = None - self._stash_host = None self._shard = 0 self._steps_since_flush = 0 - # Host mirror of the device buffer, refreshed once per step. - self._host_logits: Optional[torch.Tensor] = None - self._host_step = -1 - self._stash_capture_warned = False - - def stash_draft_confidence(self, *, logits, slots) -> None: - """Keep what THIS step's draft produced, indexed by slot. - - The persistent buffer is shared and the sampler reads it later; under - the overlap scheduler "later" can be after the next draft has written - it. SGLang sidesteps this by recording the draft's own output rather - than re-reading the buffer, and so does this once the diagnostic below - confirms the two differ. + # Snapshot ring, indexed by staged_seq % _RING_DEPTH. Each entry is + # (host_logits, host_stamps, cuda_event, staged_seq): a copy of the + # whole slot-indexed confidence buffer plus the draft-pass stamp of + # every row, taken on the host path each step -- the same relay the + # planner trusts. `record` selects the entry whose stamp for the + # request's row equals the pass that drafted the block being verified, + # so the pairing is by draft-pass identity, not by wall-clock arrival: + # the single mutable stash this replaces was overwritten by the very + # forward that produced the label, and NO execution order paired it + # correctly (off by one plain, off by two under overlap). + self._ring: List[Optional[tuple]] = [None] * self._RING_DEPTH + # Every decline is counted: this feature's failures are all silent, + # and an empty or thin shard with healthy-looking stats was exactly + # how the mispaired fit shipped. + self.stats: dict = {"recorded": 0, "no_snapshot": 0, + "stale_stamp": 0, "row_out_of_range": 0, + "no_row": 0, "snapshots_staged": 0} + + def stage_snapshot(self, *, device_logits: torch.Tensor, + device_stamps: Optional[torch.Tensor], + staged_seq: Optional[int]) -> None: + """Snapshot the confidence buffer + stamps, keyed by draft pass. + + Called from the executor's host path once per step (next to the + planner's own staging), never from inside the captured graph: Python + in a captured region runs at capture time only, which is how the old + stash silently replayed stale rows on every graphed step. """ - # DEVICE-side clone only. This runs inside the captured draft graph, so - # a .to("cpu") here raises "Cannot copy between CPU and CUDA tensors - # during CUDA graph capture" -- the same constraint that forces - # _confidence_logits to be written in place. SGLang keeps a device - # reference for exactly this reason. The host copy is deferred to the - # sampler, which is outside the graph. - # - # EAGER-ONLY diagnostic: Python inside a captured region executes at - # capture time only, so under graph replay these clones would keep - # pointing at the last CAPTURE's tensors (dummy-slot content) and the - # late-read comparison would diff against garbage while looking - # perfectly healthy. Refuse to arm the stash during capture and say so - # once, rather than let a replayed collection run silently measure the - # wrong thing. - if torch.cuda.is_current_stream_capturing(): - if not self._stash_capture_warned: - self._stash_capture_warned = True - logger.warning( - "DSpark STS stash skipped under CUDA-graph capture: the " - "late-read diagnostic is eager-only. Run collection with " - "cuda_graph_config disabled to exercise it.") - self._stash_logits = None - self._stash_slots = None - self._stash_host = None + if device_stamps is None or staged_seq is None: return - self._stash_logits = logits.detach().clone() - self._stash_slots = slots.detach().clone() - self._stash_host = None - - def begin_step(self, step: int) -> None: - """Invalidate the host mirror; the next ``record`` refreshes it.""" - if step != self._host_step: - self._host_logits = None - self._stash_host = None - self._host_step = step - - def record(self, *, slot: int, accepted: int, - device_logits: torch.Tensor) -> None: - """Record one request's ``(logits, accepted)`` pair. + idx = int(staged_seq) % self._RING_DEPTH + entry = self._ring[idx] + if (entry is None or entry[0].shape != device_logits.shape + or entry[1].shape != device_stamps.shape): + on_gpu = device_logits.is_cuda + entry = ( + torch.empty(device_logits.shape, dtype=torch.float32, + device="cpu", + pin_memory=on_gpu and prefer_pinned()), + torch.empty(device_stamps.shape, dtype=torch.int32, + device="cpu", + pin_memory=on_gpu and prefer_pinned()), + torch.cuda.Event() if on_gpu else None, + None, + ) + entry[0].copy_(device_logits.detach().to(torch.float32), + non_blocking=True) + entry[1].copy_(device_stamps, non_blocking=True) + if entry[2] is not None: + entry[2].record() + self._ring[idx] = (entry[0], entry[1], entry[2], int(staged_seq)) + self.stats["snapshots_staged"] += 1 + + def record(self, *, row: Optional[int], accepted: int, + target_seq: Optional[int]) -> None: + """Pair one request's accepted count with the logits that drafted it. Args: - slot: ``py_seq_slot``; indexes both the device buffer and nothing - else, which is why it is the join key. + row: the request's row in the confidence buffer, resolved through + the worker's own allocator (``confidence_row_for``). The + previous join key -- ``py_seq_slot`` -- came from a DIFFERENT + allocator that only coincides until the first request + completes; every pair after roster churn was request A's + label against request B's logits, and the shards looked + healthy. accepted: drafted positions accepted, excluding the bonus token. - This is exactly the ``num_correct_drafts`` the prefix mask is - built from. - device_logits: the worker's ``[rows, K]`` confidence buffer. + target_seq: the draft pass that produced the block this label + verifies, snapshotted into the SampleState at sampling time + (the ``verify_lens_snapshot`` pattern; reading it live at + update time rewinds wrong under the overlap scheduler). """ - if self._host_logits is None: - # The single sync of the step. Deliberate: see the module docstring. - self._host_logits = device_logits.detach().to( - device="cpu", dtype=torch.float32, copy=True) - if slot < 0 or slot >= self._host_logits.shape[0]: + if row is None: + self.stats["no_row"] += 1 + return + if target_seq is None: + self.stats["no_snapshot"] += 1 return - buffered = self._host_logits[slot].clone() - self._logits.append(buffered) + entry = self._ring[int(target_seq) % self._RING_DEPTH] + if entry is None: + self.stats["no_snapshot"] += 1 + return + logits_host, stamps_host, event, staged_seq = entry + if not 0 <= int(row) < logits_host.shape[0]: + self.stats["row_out_of_range"] += 1 + return + if event is not None: + event.synchronize() + if int(stamps_host[int(row)]) != int(target_seq): + # The buffer content for this row is from a different draft pass + # than the one this label verifies -- the ring slot was reused, or + # the row was never written for that pass. Refusing IS the fix: + # appending it anyway is the mispairing this design replaces. + self.stats["stale_stamp"] += 1 + return + self._logits.append(logits_host[int(row)].clone()) self._accepted.append(int(accepted)) - # Diagnostic: the same row as the draft produced it. If these two agree - # everywhere, reading the shared buffer late is safe; if they diverge, - # every pair recorded so far was confidence(t+k) against accepted(t). - # Materialize the draft-time snapshot once per step, here, outside the - # captured region. - if self._stash_host is None and self._stash_logits is not None: - rows = self._stash_logits.to(device="cpu", dtype=torch.float32, - copy=True) - idx = self._stash_slots.to(device="cpu").tolist() - self._stash_host = {int(sl): rows[i] for i, sl in enumerate(idx)} - stashed = (self._stash_host or {}).get(int(slot)) - self._stashed_logits.append( - stashed.clone() if stashed is not None - else torch.full_like(buffered, float("nan"))) + self.stats["recorded"] += 1 def end_step(self) -> None: self._steps_since_flush += 1 @@ -264,14 +288,24 @@ def flush(self) -> None: path = Path(f"{self.path_stem}.r{self.rank}.{self._shard}.pt") path.parent.mkdir(parents=True, exist_ok=True) - payload = {"logits": logits, "prefix_mask": prefix_mask} - if self._stashed_logits: - payload["logits_at_draft"] = torch.stack(self._stashed_logits, dim=0) + # The pairing marker lets the fitter reject shards from the pre-ring + # recorder, whose pairs were mislabeled in ways the tensors themselves + # cannot reveal. The running decline counters ride along so a thin + # shard explains itself. + payload = { + "logits": logits, + "prefix_mask": prefix_mask, + "meta": { + "pairing": "draft_seq_ring", + "block_size": self.block_size, + "stats": dict(self.stats), + }, + } torch.save(payload, path) logger.info( - f"DSpark STS: wrote {logits.shape[0]} samples to {path}") + f"DSpark STS: wrote {logits.shape[0]} samples to {path} " + f"(stats={self.stats})") self._logits.clear() - self._stashed_logits.clear() self._accepted.clear() self._shard += 1 diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index c1435659c2b6..56a33175ef77 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -260,7 +260,12 @@ def adopt_verify_len_pin(self, wire_value: int) -> None: wire_value = int(wire_value) if wire_value < 0: return - self._pending_verify_len_pin = -1 + # Compare-and-clear: the RPC thread may queue a NEW request between + # the decode loop reading the payload and adopting it (they are an + # allgather apart). Unconditionally resetting here wiped that second + # request after its endpoint had already answered 200. + if self._pending_verify_len_pin == wire_value: + self._pending_verify_len_pin = -1 new_pin = None if wire_value == 0 else wire_value if new_pin == self._forced_verify_len: return diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 4b2d41baf640..6769ac5b5d6b 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -70,6 +70,13 @@ class SampleStateSpec(SampleState): #: never coexist, and conflating them would rewind by the window after #: verifying the block, silently leaking KV. verify_caps_snapshot: Optional[dict] = None + #: The draft pass that produced the block THIS state's tokens verified; + #: the STS recorder joins its acceptance label to a logits snapshot by + #: this key. Snapshotted at sampling time for the same reason + #: ``verify_lens_snapshot`` is: by the time ``update_requests`` runs, the + #: live counter already belongs to a later pass. None unless STS + #: collection is on. + sts_target_seq: Optional[int] = None class SpecSamplerBase(Sampler[SampleStateSpec], AsyncWorkerMixin): @@ -200,10 +207,14 @@ def _request_common_handling( #: everywhere else so this costs one attribute read per request. acceptance_stats = None #: STS calibration collection; attached alongside `acceptance_stats` and - #: None everywhere else. `sts_logits_buffer` is the worker's persistent - #: [rows, K] confidence buffer, indexed by the same py_seq_slot. + #: None everywhere else. `sts_row_for` resolves a py_request_id to its row + #: in the worker's confidence buffer through the worker's OWN allocator -- + #: `py_seq_slot` comes from a different allocator whose numbering only + #: coincides until the first request completes. `sts_seq_provider` names + #: the draft pass the current step's verification targets. sts_recorder = None - sts_logits_buffer = None + sts_row_for = None + sts_seq_provider = None @staticmethod def _verified_len(request: LlmRequest, runtime_draft_len: int, @@ -257,12 +268,6 @@ def update_requests( runtime_draft_len = getattr(state, "runtime_draft_len", self.draft_len) caps = state.verify_caps_snapshot - if self.sts_recorder is not None: - # A fresh step: drop the cached host mirror so the first `record` - # below pulls this step's logits rather than the previous step's. - self._sts_step = getattr(self, "_sts_step", 0) + 1 - self.sts_recorder.begin_step(self._sts_step) - for req in state.requests: if req.state == LlmRequestState.GENERATION_COMPLETE: continue @@ -302,16 +307,17 @@ def update_requests( cap_trim=(0 if cap_trim_list is None else cap_trim_list[req.py_seq_slot])) # STS calibration pairs this request's drafted-block logits with - # how much of that block was accepted. Recorded here for the same - # reason acceptance is: this is the only place both are host-side - # and belong to the same step. The recorder pulls the logits buffer - # once per step, not once per request. - if (self.sts_recorder is not None - and self.sts_logits_buffer is not None): + # how much of that block was accepted. The logits come from the + # recorder's snapshot ring, selected by the draft pass snapshotted + # into this state at sampling time and checked against the row's + # own stamp -- reading the live buffer here paired the label with + # whatever a LATER pass had written over it. + if self.sts_recorder is not None: self.sts_recorder.record( - slot=req.py_seq_slot, + row=(self.sts_row_for(req.py_request_id) + if self.sts_row_for is not None else None), accepted=req.py_num_accepted_draft_tokens, - device_logits=self.sts_logits_buffer) + target_seq=state.sts_target_seq) self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) if self.sts_recorder is not None: @@ -445,4 +451,6 @@ def sample_async( runtime_draft_len=runtime_draft_len, verify_lens_snapshot=verify_lens_snapshot, verify_caps_snapshot=verify_caps_snapshot, + sts_target_seq=(self.sts_seq_provider() + if self.sts_seq_provider is not None else None), ) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 0e8a541bc9cd..3c50959123e6 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -2396,13 +2396,27 @@ async def set_dspark_verify_len_pin( which cannot reach ranks spawned before it was set. The pin takes effect on the next decode step, simultaneously on every rank. """ - assert isinstance( - self.generator, AsyncLLM - ), "/dspark/verify_len_pin endpoint is only supported with AsyncLLM()" + # Dispatched by CAPABILITY, not by type. trtllm-serve builds a + # PyTorchLLM, which carries the synchronous `_collective_rpc`; only the + # AsyncLLM entry point has the awaitable `collective_rpc`. Asserting + # the type (as the memory endpoints do) rejects the very deployment + # this knob exists to profile, and does it with a 500. + rpc = getattr(self.generator, "collective_rpc", None) + sync_rpc = getattr(self.generator, "_collective_rpc", None) + if rpc is None and sync_rpc is None: + return JSONResponse( + content={"error": "this server exposes no worker RPC channel, " + "so the verify-length pin cannot be set"}, + status_code=400) try: - applied = await self.generator.collective_rpc( - 'set_dspark_verify_len_pin', args=(request.verify_len, )) - except (ValueError, RuntimeError) as exc: + if rpc is not None: + applied = await rpc('set_dspark_verify_len_pin', + args=(request.verify_len, )) + else: + applied = await asyncio.to_thread( + sync_rpc, 'set_dspark_verify_len_pin', + args=(request.verify_len, )) + except (ValueError, RuntimeError, NotImplementedError) as exc: # A bad tier or a non-DSpark engine is a caller error, not a server # fault: answer it rather than returning a 500 with a stack trace. return JSONResponse(content={"error": str(exc)}, status_code=400) diff --git a/tests/microbenchmarks/dspark_fit_sts.py b/tests/microbenchmarks/dspark_fit_sts.py index 5ddd50aa9bad..856b84326a9c 100644 --- a/tests/microbenchmarks/dspark_fit_sts.py +++ b/tests/microbenchmarks/dspark_fit_sts.py @@ -143,51 +143,38 @@ def fit_sts_temperatures(*, logits: torch.Tensor, prefix_mask: torch.Tensor, } -def load_shards(pattern: str, - source: str = "auto") -> Tuple[torch.Tensor, torch.Tensor]: - """Load shards, choosing which confidence column to fit against. - - Shards carry two columns for the same rows. ``logits`` is the worker's - persistent buffer read back at sampler time; ``logits_at_draft`` is the - snapshot the draft pass actually produced, cloned on device before anything - else could touch it. Under the overlap scheduler these are NOT the same - tensor: the next step's draft overwrites the shared buffer before the - sampler gets there. Measured on job 2562577, 26383 of 32083 rows differed - (82.2%), mean |delta| 7.35, correlation 0.14 -- i.e. the ``logits`` column - pairs confidence(t+1) with accepted(t) four times out of five. - - So ``auto`` prefers ``logits_at_draft`` whenever it is present. Fitting the - stale column does not fail loudly; it converges to a temperature vector that - calibrates noise, which is what the 0.23 residual ECE of the first table was. +def load_shards(pattern: str) -> Tuple[torch.Tensor, torch.Tensor]: + """Load shards, refusing any whose pairing provenance is wrong or absent. + + Shards written by the current recorder carry ``meta.pairing == + "draft_seq_ring"``: each row's logits were selected from a stamped host + snapshot of exactly the draft pass that produced the block the label + verifies, keyed by the worker's own row allocator. Earlier recorders + paired the label against a sampler-time read of the live buffer -- stale + by one draft pass for most rows under the overlap scheduler (on job + 2562577, 26383 of 32083 rows differed, correlation 0.14) -- and keyed + rows by ``py_seq_slot``, a different allocator that drifts after the + first request completes. Fitting such a shard does not fail loudly; it + converges to a temperature vector that calibrates noise, which is what + the 0.23 residual ECE of the first deployed table was. There is + deliberately no flag to accept them: re-collect. """ paths = sorted(glob.glob(pattern)) if not paths: raise SystemExit(f"no shards matched {pattern!r}") - have_draft = [] - for path in paths: - blob = torch.load(path, map_location="cpu") - have_draft.append("logits_at_draft" in blob) - if source == "auto": - key = "logits_at_draft" if all(have_draft) else "logits" - if not all(have_draft) and any(have_draft): - raise SystemExit( - "some shards carry 'logits_at_draft' and some do not; they came " - "from different builds. Fit them separately -- concatenating " - "aligned and stale rows produces a vector describing neither.") - else: - key = {"draft": "logits_at_draft", "late": "logits"}[source] - print(f"fitting against '{key}'" - + ("" if key == "logits_at_draft" else - " <-- WARNING: sampler-time buffer read; under the overlap " - "scheduler this is stale by one draft for most rows")) - logits, masks, block_size = [], [], None - dropped = 0 for path in paths: blob = torch.load(path, map_location="cpu") - if key not in blob: - raise SystemExit(f"{path}: no '{key}' column (has {sorted(blob)})") - lg, pm = blob[key], blob["prefix_mask"] + meta = blob.get("meta") or {} + pairing = meta.get("pairing") + if pairing != "draft_seq_ring": + raise SystemExit( + f"{path}: pairing provenance is {pairing!r}, expected " + f"'draft_seq_ring'. This shard predates the stamped snapshot " + f"ring, so its labels are joined to the wrong draft pass " + f"and/or the wrong buffer row, and a fit on it calibrates " + f"noise. Re-collect with the current recorder.") + lg, pm = blob["logits"], blob["prefix_mask"] if lg.shape != pm.shape: raise SystemExit(f"{path}: logits {tuple(lg.shape)} != mask {tuple(pm.shape)}") if block_size is None: @@ -197,16 +184,8 @@ def load_shards(pattern: str, # that never existed. raise SystemExit( f"{path}: block size {lg.shape[1]} != {block_size} from earlier shards") - # A row is NaN when the draft stash had no entry for that slot (the - # request did not draft this step). Those rows carry a real acceptance - # label against no confidence at all, so they must go -- sigmoid(NaN) - # would poison every temperature in the left-to-right sweep. - keep = ~torch.isnan(lg).any(dim=1) - dropped += int((~keep).sum()) - logits.append(lg[keep]) - masks.append(pm[keep]) - if dropped: - print(f"dropped {dropped} rows with no draft-time snapshot") + logits.append(lg) + masks.append(pm) return torch.cat(logits, 0), torch.cat(masks, 0) @@ -231,14 +210,9 @@ def main(argv=None) -> int: "large gap means the table does not describe the " "live head (drift, wrong file, wrong column), and " "nothing else in the pipeline can see that.") - ap.add_argument("--source", choices=("auto", "draft", "late"), default="auto", - help="which confidence column to fit: 'draft' is the " - "snapshot the draft pass produced (correct pairing), " - "'late' is the sampler-time buffer read (stale under " - "the overlap scheduler). 'auto' prefers draft.") args = ap.parse_args(argv) - logits, prefix_mask = load_shards(args.data, args.source) + logits, prefix_mask = load_shards(args.data) if args.eval_temps: with open(args.eval_temps, encoding="utf-8") as fh: diff --git a/tests/microbenchmarks/dspark_sps_profiler.py b/tests/microbenchmarks/dspark_sps_profiler.py index 9eb1dd4e0ae4..7f1a9a4589d5 100644 --- a/tests/microbenchmarks/dspark_sps_profiler.py +++ b/tests/microbenchmarks/dspark_sps_profiler.py @@ -29,7 +29,7 @@ Run it:: - python -m tensorrt_llm._torch.speculative.dspark_sps_profiler \\ + python tests/microbenchmarks/dspark_sps_profiler.py \\ --model /models/DeepSeek-V4-Pro-DSpark --tp-size 8 --ep-size 8 \\ --enable-attention-dp --batch-sizes 8,16,32,64,128 \\ --max-draft-len 5 --out sps_table.json @@ -442,6 +442,7 @@ def summarize_cells( *, warmup_steps: int, min_samples: int, + expected_cells: Optional[Sequence[Tuple[int, int]]] = None, ) -> List[CellStat]: """Median per ``(batch_size, verify_len)`` cell, after dropping warmup. @@ -454,22 +455,36 @@ def summarize_cells( warmup_steps: leading aligned steps to discard *per cell*. They are kept in the raw sample file so a refit can choose a different cut. min_samples: post-warmup samples a cell must have to be admitted. + expected_cells: the requested grid, when there is one. Sources that + observe a deployment (a served sweep's iteration log, a production + log) inevitably capture ramp and drain steps whose occupancy was + never requested -- e.g. a batch of 2 while the client's 256 + requests were still being admitted. Those incidental cells are + dropped (and reported) when thin; only a thin cell that was + actually REQUESTED is an error. None means every observed cell was + asked for (the engine-building sweep, existing callers). Raises: - InsufficientSamplesError: if any cell falls short. Silently dropping a - cell would punch a hole in the grid and can disconnect the - batch-size ladder, which changes the fit rather than just widening - its error bars. + InsufficientSamplesError: if a requested cell falls short. Silently + dropping one of those would punch a hole in the grid and can + disconnect the batch-size ladder, which changes the fit rather + than just widening its error bars. """ grouped: Dict[Tuple[int, int], List[StepSample]] = {} for sample in samples: grouped.setdefault((int(sample.batch_size), int(sample.verify_len)), []).append(sample) + requested = None if expected_cells is None else { + (int(b), int(l)) for b, l in expected_cells} + dropped: List[str] = [] cells: List[CellStat] = [] for (batch_size, verify_len), cell_samples in sorted(grouped.items()): ordered = sorted(cell_samples, key=lambda s: s.iteration) steady = [s.step_time_ms for s in ordered[int(warmup_steps) :]] if len(steady) < int(min_samples): + if requested is not None and (batch_size, verify_len) not in requested: + dropped.append(f"(bs={batch_size}, L={verify_len}, n={len(steady)})") + continue raise InsufficientSamplesError( f"cell (batch_size={batch_size}, verify_len={verify_len}) has " f"{len(steady)} steady samples after discarding " @@ -490,6 +505,12 @@ def summarize_cells( p90_ms=float(np.percentile(values, 90)), ) ) + if dropped: + # Reported, not silent: a fit that quietly ignored data reads as + # "covered everything" when it did not. + print(f"[dspark-sps] dropped {len(dropped)} thin incidental cell(s) " + f"outside the requested grid: {', '.join(dropped[:12])}" + + (" ..." if len(dropped) > 12 else ""), file=sys.stderr) return cells @@ -749,12 +770,29 @@ def compress_to_risers( kept.append(index) limit = max(int(max_breakpoints), 2) + truncated_from = len(kept) while len(kept) > limit: - # Only interior/trailing points are droppable: index 0 defines the base - # shelf that every below-range query clamps onto. - jumps = [(step_times[kept[i]] - step_times[kept[i - 1]], i) for i in range(1, len(kept))] - _, drop = min(jumps) + # Drop the interior point whose removal introduces the least + # interpolation error. The previous rule -- smallest incoming jump -- + # deleted shelf-CLOSING points first by construction (their jump is + # below min_riser_ms), which turned every measured shelf back into a + # ramp and re-billed mid-shelf totals upward: the exact over-charge + # the pair encoding exists to prevent. Endpoints are never droppable: + # index 0 anchors the below-range clamp, the last point the above. + def _removal_error(pos: int) -> float: + x0, x1, x2 = (token_counts[kept[pos - 1]], token_counts[kept[pos]], + token_counts[kept[pos + 1]]) + y0, y2 = step_times[kept[pos - 1]], step_times[kept[pos + 1]] + interp = y0 + (y2 - y0) * (x1 - x0) / max(x2 - x0, 1) + return abs(step_times[kept[pos]] - interp) + + drop = min(range(1, len(kept) - 1), key=_removal_error) kept.pop(drop) + if truncated_from > len(kept): + print(f"[dspark-sps] compress_to_risers: truncated {truncated_from} " + f"breakpoints to {len(kept)} (max_breakpoints={max_breakpoints});" + f" the kept set minimizes interpolation error, but raise the cap" + f" if the residuals matter", file=sys.stderr) return [int(token_counts[i]) for i in kept], [float(step_times[i]) for i in kept] @@ -827,6 +865,11 @@ def build_cost_table_payload( "batch_overhead_ms": [float(round(v, 6)) for v in alphas], } payload["_meta"] = { + # Consumer contract marker. Tables written for the old floor consumer + # deliberately dropped shelf-closing breakpoints; interpolating those + # re-bills every mid-shelf total upward with nothing to say so. The + # loader warns when this key is absent. + "lookup": "interp", "encoding": "decomposed", "note": ( "step_time_ms is theta(M) only -- the trimmable term. The " @@ -1094,6 +1137,22 @@ def validate(self) -> None: raise SweepGeometryError("both --batch-sizes and --verify-lens must be non-empty") if min(self.verify_lens) < 1: raise SweepGeometryError("verify lengths must be >= 1 (the planner's floor)") + if (not self.block_follows_verify_len + and min(self.verify_lens) < int(self.max_draft_len)): + # Constant-block with a pinned window below the block currently + # dies on a device-side assert (Repeat.cu output_size vs + # sum(repeats)) at the first short cell -- AFTER minutes of engine + # build. Refusing here is this docstring's whole job; letting the + # default geometry walk into a known assert is not a validation. + raise SweepGeometryError( + f"constant-block geometry with verify lengths " + f"{sorted(self.verify_lens)} below the block " + f"{self.max_draft_len} runs the regime that currently trips " + f"the Repeat.cu output_size assert. Pass " + f"--block-follows-verify-len (each cell drafts exactly what " + f"it verifies), or sweep only L == {self.max_draft_len}, or " + f"use --base-url / --from-iter-log which measure a deployment " + f"instead of building one.") # A request that hits max_seq_len is evicted mid-cell, so the batch # drains and the steady window shrinks -- which surfaces much later as a # confusing "not enough steady samples". @@ -1345,15 +1404,25 @@ def _run_cell(llm, config: SweepConfig, *, batch_size: int, verify_len: int) -> # cell that lost its tail this way is visible rather than merely short. rows: List[dict] = [] drain_error: List[BaseException] = [] + empty_reads = [0] stop = threading.Event() def _drain() -> None: while not stop.is_set(): try: - rows.extend(list(llm.get_stats(timeout=config.stats_poll_s) or [])) + got = list(llm.get_stats(timeout=config.stats_poll_s) or []) except BaseException as exc: # noqa: BLE001 - reported, not swallowed drain_error.append(exc) return + if got: + rows.extend(got) + else: + # Counted because one empty read is not merely a lost poll: + # IterationResult latches done on the first Empty, and until + # something un-latches it every later get_stats -- including + # the final post-cell drain -- returns [] instantly. A cell + # with a non-zero count here may have lost its tail. + empty_reads[0] += 1 stop.wait(config.stats_poll_s) # A daemon drainer, NOT a poll loop that owns the exit condition. Making the @@ -1388,7 +1457,19 @@ def _drain() -> None: f"{drain_error[0]}); the cell may be short", file=sys.stderr) # Final drain for whatever crossed the IPC boundary after the last poll. + # Un-latch first: a single gap in stats production longer than the poll + # timeout latched IterationResult done, and a latched queue answers [] + # instantly -- silently truncating the cell this drain exists to finish. + undo = getattr(getattr(getattr(llm, "_executor", None), + "_iter_stats_result", None), "mark_undone", None) + if undo is not None: + undo() rows.extend(list(llm.get_stats(timeout=config.stats_timeout_s) or [])) + if empty_reads[0]: + print(f"[dspark-sps] bs={batch_size} L={verify_len}: " + f"{empty_reads[0]} empty stats read(s) mid-cell; the retained " + f"window may have closed early and the cell lost its tail", + file=sys.stderr) aligned = aligned_steps_from_stats( rows, expected_ranks=config.dp_size, timing_key=config.timing_key @@ -1470,7 +1551,7 @@ def _float_list(text: str) -> List[float]: def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( - prog="python -m tensorrt_llm._torch.speculative.dspark_sps_profiler", + prog="python tests/microbenchmarks/dspark_sps_profiler.py", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, ) @@ -1497,8 +1578,18 @@ def _build_parser() -> argparse.ArgumentParser: source.add_argument("--server-poll-s", type=float, default=20.0, help="Seconds of steady state to collect per cell " "(--base-url).") - source.add_argument("--served-model-name", default="model", - help="`model` field for /v1/completions (--base-url).") + source.add_argument( + "--server-log", + help="The server's stdout/stderr (print_iter_log). When reachable this " + "is where --base-url takes its samples from: every step is " + "labelled by the shape it ran, and no metrics collector has to be " + "configured for the sweep to see anything.") + source.add_argument("--served-model-name", default=None, + help="`model` field for /v1/completions (--base-url). " + "Omit to read it from /v1/models, which is what " + "the server actually registered -- for a local " + "checkpoint that is the directory basename, not " + "the path it was launched with.") source.add_argument( "--cuda-graph-batch-sizes", help="The deployment's captured batch-size ladder, comma separated " @@ -1777,41 +1868,69 @@ def samples_from_iter_log( """ line_re = re.compile( r"iter = (\d+).*?num_scheduled_requests = (\d+).*?" - r"host_step_time = ([\d.]+)ms.*?'num_generation_tokens': (\d+)") + r"host_step_time = ([\d.]+)ms.*?'num_ctx_requests': (\d+)" + r".*?'num_generation_tokens': (\d+)") out: List[StepSample] = [] skipped_ratio = 0 + skipped_mixed = 0 + + def _consume(line: str) -> None: + nonlocal skipped_ratio, skipped_mixed + match = line_re.search(line) + if not match: + return + iteration = int(match.group(1)) + batch_size = int(match.group(2)) + step_ms = float(match.group(3)) + ctx_requests = int(match.group(4)) + tokens = int(match.group(5)) + if ctx_requests: + # A step carrying prefill chunks runs eager at a prefill-inflated + # time; the ladder resolution would still find a plausible + # (padded_bs, L) for it and file a mixed step as a pure decode + # cell -- confidently mislabelled, the exact failure this + # collector exists to avoid. + skipped_mixed += 1 + return + if iteration == 1: + # Each executor instance numbers from 1, and its first step + # carries the wait for the first request inside its host time + # (62 s on a live server). Identify it by the number rather than + # by having seen one, so a log FRAGMENT is still usable. + return + if batch_size <= 0 or step_ms >= max_step_ms: + return + shape = resolve_padded_shape( + num_rows=batch_size, num_generation_tokens=tokens, + max_draft_len=max_draft_len, max_batch_size=max_batch_size, + padded_batch_sizes=padded_batch_sizes) + if shape is None: + skipped_ratio += 1 + return + padded_bs, verify_len = shape + out.append(StepSample(batch_size=padded_bs, verify_len=verify_len, + step_time_ms=step_ms, iteration=iteration)) + for path in paths: text = pathlib.Path(path).read_text(errors="replace") - for line in text.splitlines(): - if "iter = " not in line: - continue - if rank_prefix and not line.lstrip().startswith(rank_prefix): - continue - match = line_re.search(line) - if not match: - continue - iteration = int(match.group(1)) - batch_size = int(match.group(2)) - step_ms = float(match.group(3)) - tokens = int(match.group(4)) - if iteration == 1: - # Each executor instance numbers from 1, and its first step - # carries the wait for the first request inside its host time - # (62 s on a live server). Identify it by the number rather - # than by having seen one, so a log FRAGMENT is still usable. - continue - if batch_size <= 0 or step_ms >= max_step_ms: - continue - shape = resolve_padded_shape( - num_rows=batch_size, num_generation_tokens=tokens, - max_draft_len=max_draft_len, max_batch_size=max_batch_size, - padded_batch_sizes=padded_batch_sizes) - if shape is None: - skipped_ratio += 1 - continue - padded_bs, verify_len = shape - out.append(StepSample(batch_size=padded_bs, verify_len=verify_len, - step_time_ms=step_ms, iteration=iteration)) + candidates = [ln for ln in text.splitlines() if "iter = " in ln] + matched = [ln for ln in candidates + if not rank_prefix or ln.lstrip().startswith(rank_prefix)] + if candidates and not matched and rank_prefix: + # Single-process deployments write no srun/mpirun rank tag; a + # hard prefix requirement silently dropped every line and the + # error blamed print_iter_log. When the prefix matches nothing at + # all, the log simply is not rank-tagged -- read it as-is. + print(f"[dspark-sps] iter-log: no lines in {path} carry the rank " + f"prefix {rank_prefix!r}; reading it as an untagged " + f"single-process log", file=sys.stderr) + matched = candidates + for line in matched: + _consume(line) + if skipped_mixed: + print(f"[dspark-sps] iter-log: skipped {skipped_mixed} mixed " + f"context+generation steps (prefill-inflated times)", + file=sys.stderr) if skipped_ratio: print(f"[dspark-sps] iter-log: skipped {skipped_ratio} steps whose " f"padded shape could not be resolved unambiguously", @@ -1825,7 +1944,6 @@ def samples_from_server( batch_sizes: Sequence[int], verify_lens: Sequence[int], input_len: int, - max_tokens: int, model: str, warmup_steps: int, settle_s: float, @@ -1833,6 +1951,8 @@ def samples_from_server( max_draft_len: int, seed: int, padded_batch_sizes: Optional[Sequence[int]] = None, + server_log: Optional[str] = None, + dp_size: int = 1, ) -> List[StepSample]: """Sweep a RUNNING server: pin the window over HTTP, then read /metrics. @@ -1845,12 +1965,48 @@ def samples_from_server( a fixed batch size for long enough to measure it; the SHAPE is what the table indexes, and the deployment's own graphs, context depths and scheduler are what answer for it. + + ``dp_size`` is how many requests it takes to put ONE on each rank's + batch. Under attention DP every rank schedules its own subset, the + planner prices its own rank's batch, and the iteration log reports + per-rank rows -- so the table's bs axis is PER-RANK, and a cell labeled + ``bs`` needs ``bs * dp_size`` requests in flight. Job 2585129 swept a + tp8-ADP server with dp_size unaccounted: every requested cell (64..256) + actually measured per-rank batches of 8..40, and the fit refused the + grid as disconnected. """ import requests rng = random.Random(seed) out: List[StepSample] = [] + if not model: + # The served id is whatever the server registered -- for a local + # checkpoint that is the directory's basename, not the path that was + # passed to it. Posting the path gets every request rejected, no + # traffic runs, and every cell reports zero steps with nothing to say + # why. Ask the server instead of making the caller guess. + listing = requests.get(f"{base_url}/v1/models", timeout=60) + listing.raise_for_status() + entries = (listing.json() or {}).get("data") or [] + if not entries: + raise RuntimeError(f"{base_url}/v1/models listed no model") + model = entries[0]["id"] + print(f"[dspark-sps] serving model id: {model}", file=sys.stderr) + + post_errors: List[str] = [] + + def _fire(body: dict) -> None: + try: + response = requests.post(f"{base_url}/v1/completions", json=body, + timeout=3600) + if response.status_code != 200 and len(post_errors) < 3: + post_errors.append( + f"{response.status_code} {response.text[:200]}") + except Exception as exc: # noqa: BLE001 - reported below, not fatal + if len(post_errors) < 3: + post_errors.append(f"{type(exc).__name__}: {exc}") + def pin(verify_len: Optional[int]) -> None: response = requests.post(f"{base_url}/dspark/verify_len_pin", json={"verify_len": verify_len}, timeout=60) @@ -1867,17 +2023,31 @@ def drain_metrics() -> List[dict]: try: for verify_len in sorted(verify_lens): - pin(int(verify_len)) + try: + pin(int(verify_len)) + except RuntimeError as exc: + # A rung the server's ladder lacks. Raising out of the sweep + # threw away every already-collected sample (nothing is + # written until the sweep returns); a partial ladder is still + # a table, and the skip is counted in the output. + print(f"[dspark-sps] server: pin L={verify_len} refused " + f"({exc}); skipping the rung", file=sys.stderr) + continue for batch_size in sorted(batch_sizes): + # H2: the load must outlive the measurement window, or the + # cell measures ramp-down rows filed under smaller batch + # sizes. Budget the output by the window at a conservative + # 30 ms/step lower bound. + window_steps = int((settle_s + poll_s) / 0.03) + 64 + max_tokens = max(64, window_steps * (1 + int(verify_len))) + num_requests = int(batch_size) * max(1, int(dp_size)) prompts = [[rng.randrange(1000, 30000) for _ in range(input_len)] - for _ in range(int(batch_size))] + for _ in range(num_requests)] bodies = [{"model": model, "prompt": p, "max_tokens": max_tokens, "temperature": 0.0, "ignore_eos": True, "stream": False} for p in prompts] - threads = [threading.Thread( - target=lambda b=b: requests.post( - f"{base_url}/v1/completions", json=b, timeout=3600), - daemon=True) for b in bodies] + threads = [threading.Thread(target=_fire, args=(b, ), + daemon=True) for b in bodies] for thread in threads: thread.start() time.sleep(settle_s) # let admission reach the plateau @@ -1897,6 +2067,7 @@ def drain_metrics() -> List[dict]: # so where it publishes both), which would inflate every cell # by roughly a factor of two. kept = 0 + cell_rows: List[StepSample] = [] for row in rows: rows_seen = int(row.get("numActiveRequests") or 0) step_ms = float(row.get("hostStepTimeMS") or 0.0) @@ -1909,18 +2080,53 @@ def drain_metrics() -> List[dict]: if not wider: continue padded_bs = wider[0] - out.append(StepSample(batch_size=padded_bs, - verify_len=int(verify_len), - step_time_ms=step_ms)) + cell_rows.append(StepSample(batch_size=padded_bs, + verify_len=int(verify_len), + step_time_ms=step_ms)) kept += 1 + # Applied, not just printed: the first rows after a pin change + # can straddle the previous window's shape. + out.extend(cell_rows[warmup_steps:]) + note = "" + if not kept: + if post_errors: + note = f" -- requests refused: {post_errors[0]}" + elif not rows: + note = (" -- /metrics returned nothing: is " + "enable_iter_perf_stats on?") + else: + keys = sorted(rows[-1])[:8] + note = (f" -- {len(rows)} stats rows carried no usable " + f"(numActiveRequests, hostStepTimeMS); keys " + f"seen: {keys}") print(f"[dspark-sps] server bs={batch_size:>5} L={verify_len} " - f"kept {kept - warmup_steps if kept > warmup_steps else 0}" - f" steps", file=sys.stderr) + f"({num_requests} requests in flight) " + f"kept {max(kept - warmup_steps, 0)}" + f" steps{note}", file=sys.stderr) for thread in threads: thread.join(timeout=1.0) finally: # Never leave a served deployment pinned to a profiling window. pin(None) + + if server_log: + # Preferred source when the server's log is reachable. /metrics needs + # a metrics collector (return_perf_metrics) to run its tee buffer, and + # its fallback path -- draining the stats queue directly -- came back + # empty on an MPI deployment. The iteration log needs neither: it + # carries the row count, the padded token total and the host step time + # for every step, so the samples are labelled by what ran rather than + # by what was asked for. The pin's job is then only to make sure every + # rung is VISITED. + from_log = samples_from_iter_log( + [server_log], max_draft_len=max_draft_len, + max_batch_size=(max(padded_batch_sizes) if padded_batch_sizes + else 256), + padded_batch_sizes=padded_batch_sizes) + print(f"[dspark-sps] server log yielded {len(from_log)} steps " + f"(/metrics yielded {len(out)})", file=sys.stderr) + if from_log: + return from_log return out @@ -1940,6 +2146,9 @@ def main(argv: Optional[Sequence[str]] = None) -> int: raise SystemExit("--fit-only requires --samples-in") samples = _load_samples(_str_list(args.samples_in)) provenance = {"source": "refit", "samples_in": _str_list(args.samples_in)} + # Saved samples may come from any source, including deployment logs + # full of ramp-occupancy cells; require nothing, keep what is thick. + expected_cells = frozenset() elif args.from_iter_log: ladder = (_int_list(args.cuda_graph_batch_sizes) if args.cuda_graph_batch_sizes else None) @@ -1957,13 +2166,14 @@ def main(argv: Optional[Sequence[str]] = None) -> int: print(f"[dspark-sps] {len(samples)} steps from " f"{len(args.from_iter_log)} iteration log(s)", file=sys.stderr) provenance = {"source": "iter_log", "logs": list(args.from_iter_log)} + # A production log has no requested grid at all. + expected_cells = frozenset() elif args.base_url: samples = samples_from_server( base_url=args.base_url.rstrip("/"), batch_sizes=batch_sizes, verify_lens=verify_lens, input_len=args.input_len, - max_tokens=max(64, int(args.measure_steps)), model=args.served_model_name, warmup_steps=args.warmup_steps, settle_s=float(args.server_settle_s), @@ -1971,7 +2181,12 @@ def main(argv: Optional[Sequence[str]] = None) -> int: max_draft_len=args.max_draft_len, seed=args.seed, padded_batch_sizes=(_int_list(args.cuda_graph_batch_sizes) - if args.cuda_graph_batch_sizes else None)) + if args.cuda_graph_batch_sizes else None), + server_log=args.server_log, + # The bs axis is per-rank under attention DP (see the collector's + # docstring); the flags describing the deployment double as the + # fan-out factor. + dp_size=(int(args.tp_size) if args.enable_attention_dp else 1)) if not samples: raise SystemExit( "the server sweep collected nothing. Check that /metrics " @@ -1981,6 +2196,11 @@ def main(argv: Optional[Sequence[str]] = None) -> int: print(f"[dspark-sps] {len(samples)} steps from {args.base_url}", file=sys.stderr) provenance = {"source": "server_sweep", "base_url": args.base_url} + # The requested concurrencies are ladder entries, so a steady cell's + # padded batch size lands exactly on them; everything else in the + # server's log is admission ramp or drain-down. + expected_cells = frozenset( + (int(b), int(l)) for b in batch_sizes for l in verify_lens) else: if not args.model: raise SystemExit("--model is required unless --fit-only is used") @@ -2020,6 +2240,8 @@ def main(argv: Optional[Sequence[str]] = None) -> int: extra_llm_api_options=extra, ) samples = run_sweep(config) + expected_cells = frozenset( + (int(b), int(l)) for b in batch_sizes for l in verify_lens) _write_samples(samples_out, samples) print(f"[dspark-sps] wrote {len(samples)} raw samples to {samples_out}", file=sys.stderr) provenance = { @@ -2036,7 +2258,9 @@ def main(argv: Optional[Sequence[str]] = None) -> int: "ragged_verify_mode": RaggedVerifyMode.STATIC.value, } - cells = summarize_cells(samples, warmup_steps=args.warmup_steps, min_samples=args.min_samples) + cells = summarize_cells(samples, warmup_steps=args.warmup_steps, + min_samples=args.min_samples, + expected_cells=expected_cells) if not cells: # Distinguished from a short cell: this means no step was ever aligned, # which is a wiring problem (iteration stats off, the wrong rank count, @@ -2062,20 +2286,29 @@ def main(argv: Optional[Sequence[str]] = None) -> int: "warmup_steps_discarded": int(args.warmup_steps), "statistic": "median", # The engine fingerprint the loader verifies (check_table_fingerprint). - # Attached only when this process ran the sweep itself: an offline - # --fit-only refit of merged sample files cannot vouch for the - # engine those samples came from, and an unfingerprinted table - # loads with a warning rather than a false claim. + # Facts only, and only facts this process can VOUCH for. The loader + # compares every key present in both dicts and refuses on mismatch, + # so a key must be absent -- not defaulted -- when the fact is + # unknown: recording ep as `args.ep_size or 0` fabricated ep=0 + # whenever the flag was unset, the loader's Mapping normalizes + # unset to 1, and the profiler's own table was refused on the very + # engine that produced it. The deployment-measuring sources + # (--base-url, --from-iter-log) build no engine here at all, so + # they attach no engine dict; the deployment's own config is the + # fingerprint authority there, and an unfingerprinted table loads + # with a warning rather than a false claim. **({"engine": { - "tp": int(args.tp_size), - "ep": int(args.ep_size or 0), + **({"tp": int(args.tp_size)} + if args.tp_size is not None else {}), + **({"ep": int(args.ep_size)} + if args.ep_size else {}), "attention_dp": bool(args.enable_attention_dp), "block": block_size, "max_batch_size": max(_int_list(args.batch_sizes)), "geometry": ("block_follows_verify_len" if getattr(args, "block_follows_verify_len", False) else "constant_block"), - }} if not getattr(args, "samples_in", None) else {}), + }} if provenance.get("source") == "sweep" else {}), **provenance, }, ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py index 7435e968a971..4979f20e6df9 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py @@ -177,6 +177,27 @@ def test_summarize_refuses_short_cell(): summarize_cells(_samples(8, 3, [10.0] * 5), warmup_steps=4, min_samples=8) +def test_summarize_drops_thin_cells_outside_the_requested_grid(): + """A deployment log carries ramp/drain occupancy nobody asked to measure. + + Job 2585129: a five-rung server sweep with 984-step cells at every + requested (bs, L) died at the very end on (bs=2, L=1) with one sample -- + the admission ramp of the first cell's client, filed as a cell and then + held to the requested-grid bar. Incidental cells must be dropped when + thin, kept when thick, and a thin REQUESTED cell must still refuse. + """ + grid = [(8, 3)] + requested = _samples(8, 3, [10.0] * 8) + ramp = _samples(2, 3, [30.0]) # thin, incidental: drop + thick_extra = _samples(4, 3, [20.0] * 8) # thick, incidental: keep + cells = summarize_cells(requested + ramp + thick_extra, warmup_steps=0, + min_samples=4, expected_cells=grid) + assert [(c.batch_size, c.verify_len) for c in cells] == [(4, 3), (8, 3)] + with pytest.raises(InsufficientSamplesError, match="steady samples"): + summarize_cells(_samples(8, 3, [10.0]) + thick_extra, warmup_steps=0, + min_samples=4, expected_cells=grid) + + def test_summarize_orders_by_iteration_before_cutting_warmup(): """Stats arrive per queue, not necessarily in iteration order.""" out_of_order = [ @@ -344,11 +365,20 @@ def test_compress_keeps_risers_and_closes_their_shelves(): assert kept_times == [1.0, 1.004, 2.0] -def test_compress_drops_the_shallowest_risers_first(): +def test_truncation_drops_the_point_cheapest_to_interpolate(): + """Over budget, remove the interior point interpolation best reconstructs. + + Here 64 goes: its neighbors predict it within 0.27ms, while removing 32 + would ramp the near-flat 16..32 stretch into the 64 riser and misprice it + by 0.9ms. The previous rule (drop the smallest incoming jump first) chose + 32 -- by construction it always chose the shelf-side point, re-billing + mid-shelf totals upward, the exact over-charge the pair encoding of + shelves exists to prevent. + """ tokens = [16, 32, 64, 128] times = [1.0, 1.1, 4.0, 9.0] kept_tokens, _ = compress_to_risers(tokens, times, min_riser_ms=0.0, max_breakpoints=3) - assert kept_tokens == [16, 64, 128] + assert kept_tokens == [16, 32, 128] def test_compress_never_drops_the_base_shelf(): @@ -569,7 +599,18 @@ def test_validate_rejects_a_decode_step_wider_than_max_num_tokens(): def test_validate_accepts_a_sane_sweep(): - _sweep().validate() + """Sane means the drafted block matches what each cell verifies.""" + _sweep(block_follows_verify_len=True).validate() + _sweep(verify_lens=[5], max_draft_len=5).validate() + + +def test_validate_refuses_constant_block_verifying_below_it(): + """Building an engine that drafts 5 and verifies 3 runs the regime that + trips the Repeat.cu output_size assert -- the sweep would crash mid-cell + after minutes of engine build, or worse, measure a crash-adjacent shape. + """ + with pytest.raises(SweepGeometryError, match="Repeat.cu"): + _sweep().validate() def test_dp_size_only_counts_ranks_under_attention_dp(): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py new file mode 100644 index 000000000000..23f2467ce68c --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""STS pairing must join on draft-pass identity, never on arrival order. + +The recorder's label (accepted count) and its feature (confidence logits) are +produced by different steps: the block verified at step ``i`` was drafted at +step ``i-1``, and under the overlap scheduler the sampler consumes the label +while pass ``i+1`` has already overwritten the shared confidence buffer. The +single mutable stash this ring replaced had NO execution order under which +the pair was correct -- measured on job 2562577, 82.2% of rows paired +confidence(t+1) with accepted(t), correlation 0.14 -- and the shards looked +perfectly healthy. + +These tests pin the join contract: a pair is appended only when the ring +snapshot's per-row stamp equals the pass the label verifies, everything else +is a *counted* decline, and the shard carries provenance the fitter can +refuse. No CUDA: the recorder itself is device-agnostic (CPU tensors skip the +pinned-copy event), which is exactly what lets this run in pre-merge CI. +""" + +import importlib.util +import os +import pathlib +import sys +import tempfile + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_sts import DSparkStsRecorder + +BLOCK = 5 +ROWS = 8 + + +def _recorder(tmp, flush_every=1000): + return DSparkStsRecorder(path_stem=os.path.join(tmp, "shard"), + block_size=BLOCK, rank=0, + flush_every=flush_every) + + +def _stage(rec, seq, written_rows, fill): + """Stage one pass's snapshot: `written_rows` stamped `seq`, rest stale. + + Mirrors the producer: the in-graph scatter stamps only the rows the pass + actually drafted; every other row keeps whatever stamp it had (0 for + never-written, matching the worker's zero-initialized stamp buffer). + """ + logits = torch.full((ROWS, BLOCK), float(fill)) + stamps = torch.zeros(ROWS, dtype=torch.int32) + stamps[list(written_rows)] = seq + rec.stage_snapshot(device_logits=logits, device_stamps=stamps, + staged_seq=seq) + + +def test_pairs_by_pass_identity_not_arrival_order(): + """A label for pass 1 gets pass-1 logits even after pass 2 was staged.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + _stage(rec, seq=1, written_rows=[2], fill=1.0) + _stage(rec, seq=2, written_rows=[2], fill=2.0) # overwrite arrives first + rec.record(row=2, accepted=3, target_seq=1) + rec.record(row=2, accepted=1, target_seq=2) + assert rec.stats["recorded"] == 2 + assert rec._logits[0][0].item() == 1.0 # pass 1 content, not latest + assert rec._logits[1][0].item() == 2.0 + + +def test_row_not_written_for_that_pass_is_declined(): + """A row the pass never drafted must not contribute a fabricated pair.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + _stage(rec, seq=1, written_rows=[2], fill=1.0) + rec.record(row=3, accepted=2, target_seq=1) # row 3 stamp is 0, not 1 + assert rec.stats["recorded"] == 0 + assert rec.stats["stale_stamp"] == 1 + + +def test_ring_eviction_declines_instead_of_mispairing(): + """A label older than the ring depth is dropped, never matched to the + + evicting pass's content: the slot's stamps say who actually wrote it. + """ + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + for seq in range(1, 2 + DSparkStsRecorder._RING_DEPTH): + _stage(rec, seq=seq, written_rows=[2], fill=float(seq)) + rec.record(row=2, accepted=3, target_seq=1) # slot now holds seq 5 + assert rec.stats["recorded"] == 0 + assert rec.stats["stale_stamp"] == 1 + + +def test_every_giving_up_path_is_counted(): + """Silent declines are how the mispaired fit shipped; none may be silent.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + rec.record(row=2, accepted=1, target_seq=None) # no pass id + rec.record(row=2, accepted=1, target_seq=7) # ring slot never staged + _stage(rec, seq=1, written_rows=[2], fill=1.0) + rec.record(row=None, accepted=1, target_seq=1) # unresolvable request + rec.record(row=ROWS + 3, accepted=1, target_seq=1) + assert rec.stats == {"recorded": 0, "no_snapshot": 2, + "stale_stamp": 0, "row_out_of_range": 1, + "no_row": 1, "snapshots_staged": 1} + + +def test_staging_without_stamps_is_a_noop(): + """No stamps means no way to verify identity -- refuse to guess.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + rec.stage_snapshot(device_logits=torch.zeros(ROWS, BLOCK), + device_stamps=None, staged_seq=1) + rec.stage_snapshot(device_logits=torch.zeros(ROWS, BLOCK), + device_stamps=torch.zeros(ROWS, dtype=torch.int32), + staged_seq=None) + assert rec.stats["snapshots_staged"] == 0 + rec.record(row=2, accepted=1, target_seq=1) + assert rec.stats["no_snapshot"] == 1 + + +def test_shard_carries_provenance_and_prefix_semantics(): + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + _stage(rec, seq=1, written_rows=[2], fill=1.5) + rec.record(row=2, accepted=3, target_seq=1) + rec.flush() + blob = torch.load(os.path.join(tmp, "shard.r0.0.pt"), + map_location="cpu") + assert blob["meta"]["pairing"] == "draft_seq_ring" + assert blob["meta"]["stats"]["recorded"] == 1 + assert blob["logits"].shape == (1, BLOCK) + # accepted=3 -> positions 0..2 were ALL accepted: the prefix event + # whose probability survival[k] estimates. + assert blob["prefix_mask"].tolist() == [[1.0, 1.0, 1.0, 0.0, 0.0]] + + +# --------------------------------------------------------------------------- +# The fitter must refuse shards from the pre-ring recorders: their pairs are +# mislabeled in ways the tensors themselves cannot reveal. +# --------------------------------------------------------------------------- + +_SPEC = importlib.util.spec_from_file_location( + "dspark_fit_sts", + pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / + "dspark_fit_sts.py") +_FITTER = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _FITTER +_SPEC.loader.exec_module(_FITTER) + + +def test_fitter_refuses_shards_without_ring_provenance(): + with tempfile.TemporaryDirectory() as tmp: + torch.save({"logits": torch.zeros(4, BLOCK), + "prefix_mask": torch.zeros(4, BLOCK)}, + os.path.join(tmp, "old.r0.0.pt")) + with pytest.raises(SystemExit, match="draft_seq_ring"): + _FITTER.load_shards(os.path.join(tmp, "*.pt")) + + +def test_fitter_accepts_ring_shards(): + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + _stage(rec, seq=1, written_rows=[2, 3], fill=0.5) + rec.record(row=2, accepted=2, target_seq=1) + rec.record(row=3, accepted=5, target_seq=1) + rec.flush() + logits, mask = _FITTER.load_shards(os.path.join(tmp, "*.pt")) + assert logits.shape == (2, BLOCK) and mask.shape == (2, BLOCK) From 5910c9f6c495b6be7ace050dd5e3ae630dc948e6 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 5 Aug 2026 05:22:26 -0700 Subject: [PATCH 093/177] [None][chore] Slim the PR: comments, getattr fallbacks, env vars, tests, docs No behavior change on any production path. Four kinds of trimming, each conservative and individually verified: Comments: war-story comments compressed to the constraint they protect (what breaks and why the code is shaped this way); measurement anecdotes, falsified-hypothesis history, and observed-once log excerpts moved out -- the full narrative already lives in docs/dspark_ragged_ima_handoff.md and the git history. getattr: defensive getattr(obj, name, default) dropped where the attribute is unconditionally declared -- runner fields now declared in __init__ (agreed_ragged_bucket, _pinned_host_cache/_pinned_host_events, _sts_temperatures_host), SpecMetadata dataclass fields (is_ragged_verify, total_verify_tokens, verify_lens, qo_indptr, accept_cap_trim), Mapping and DSparkDecodingConfig fields, and the V4 metadata's gen_new_tokens_per_seq. Test stubs that fed these paths now declare the same fields the production classes declare. Env vars: TLLM_DSPARK_ASSERT_ACTIVE parsing deduplicated into the existing _dspark_assert_min_steps property (the cleanup-path copy of the parser is gone; semantics unchanged). Tests: deletions limited to tautologies (arithmetic identities), exact duplicates of stronger surviving assertions, and cross-file repeats; the sps-profiler quiet-fit assertion merged into the structure test. Also repair two validate() tests broken by the geometry guard introduced in 8a1a6046db: they now pass block_follows_verify_len so they reach the check they assert on (pre-existing failures, reproduced on the unmodified head). Docs: dspark_ragged_verify_status.md (stale 2026-07-31 snapshot whose claims the IMA fix invalidated) and dspark_review_2026-08-04_followups.md (fully consumed by 8a1a6046db) removed; the IMA handoff doc's open-item table, next steps, and unpushed-work section updated to the fixed state; the feature doc's ghost knob confidence_relay_lag_steps and carry-ring description replaced with the shipped _confidence_stamp mechanism. Validation: python3 -m py_compile on every touched module; DSpark unit suites green on B300 (hw-agnostic 409 passed with the only failures being LLM-e2e tests needing model weights absent from this LLM_MODELS_ROOT, plus the five stub/geometry test repairs above -- all 130 affected tests pass after the repairs). Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_p0_task_prompt.md | 2 - docs/dspark_ragged_ima_handoff.md | 54 +-- docs/dspark_ragged_verify_status.md | 452 ------------------ docs/dspark_review_2026-08-04_followups.md | 439 ----------------- .../features/dspark-confidence-scheduling.md | 31 +- .../sparse/deepseek_v4/compressor.py | 2 +- .../_torch/attention_backend/sparse/dsa.py | 54 +-- tensorrt_llm/_torch/models/dspark/heads.py | 5 +- .../_torch/pyexecutor/cuda_graph_runner.py | 51 +- .../_torch/pyexecutor/model_engine.py | 111 ++--- tensorrt_llm/_torch/pyexecutor/py_executor.py | 133 ++---- tensorrt_llm/_torch/speculative/dspark.py | 22 +- .../_torch/speculative/dspark_verify.py | 24 +- tensorrt_llm/_torch/speculative/interface.py | 11 +- .../test_deepseek_v4_sparse_mla.py | 2 +- .../hw_agnostic/test_dspark_cap_accept.py | 9 - .../test_dspark_confidence_schedule.py | 12 - .../hw_agnostic/test_dspark_overlap_layout.py | 49 -- .../hw_agnostic/test_dspark_ragged.py | 30 -- .../hw_agnostic/test_dspark_sps_profiler.py | 11 +- .../hw_agnostic/test_dspark_worker.py | 11 +- .../test_dspark_graph_miss_attribution.py | 1 + .../speculative/test_dspark_planner_trims.py | 24 - .../test_rejection_buffers_guard.py | 2 + .../test_dspark_ragged_requires_table.py | 17 - 25 files changed, 156 insertions(+), 1403 deletions(-) delete mode 100644 docs/dspark_ragged_verify_status.md delete mode 100644 docs/dspark_review_2026-08-04_followups.md diff --git a/docs/dspark_p0_task_prompt.md b/docs/dspark_p0_task_prompt.md index d301fe5e3f51..a799a05c75f2 100644 --- a/docs/dspark_p0_task_prompt.md +++ b/docs/dspark_p0_task_prompt.md @@ -17,8 +17,6 @@ D1/D2/D3 三个方案对比、G1–G8 验收标准、§7 的 U1–U7 未知项 你的工作是**执行它的 P0 阶段**，不是重新做分析。任何与该文档冲突的结论，必须先用 `file:line` 证明文档错了，再改。 -配套还有 `docs/dspark_ragged_verify_status.md`（388 行）。 - ### 拉取方式 本地 checkout：`/lustre/fsw/coreai_comparch_trtllm/laliao/repos/trtllm-dsv4/TensorRT-LLM` diff --git a/docs/dspark_ragged_ima_handoff.md b/docs/dspark_ragged_ima_handoff.md index 9f8a2dc65a04..9ab9a6ec5bd0 100644 --- a/docs/dspark_ragged_ima_handoff.md +++ b/docs/dspark_ragged_ima_handoff.md @@ -1,8 +1,7 @@ # DSpark ragged verify — handoff Written 2026-08-04. Supersedes nothing; read alongside -`docs/dspark_p0_task_prompt.md` (the original task) and -`docs/dspark_ragged_verify_status.md` (design/status). +`docs/dspark_p0_task_prompt.md` (the original task; §10 is the results record). Branch `dspark-p0`, fork `origin` = `github.com/lancelly/TensorRT-LLM`, PR branch `origin/confidence_head` (NVIDIA/TensorRT-LLM#17056). @@ -11,21 +10,7 @@ PR branch `origin/confidence_head` (NVIDIA/TensorRT-LLM#17056). ## 0. The one thing to read first -**There is an open CUDA illegal memory access on the ragged verify path, and my -framing of it was wrong until the last experiment of the session.** I spent most -of a day believing it was confined to a path I had just added (running ragged on -steps with no CUDA graph). The final run disproved that: with that path gated -OFF, the default configuration still faults. - -So the correct statement is: - -> Planner-driven ragged verification with real trimming has **never** completed a -> GSM8K run without `CUDA_LAUNCH_BLOCKING=1`. - -Everything in §4 was designed under the wrong hypothesis. The falsifications -are still valid as facts, but the *set* of candidates they were drawn from was -too narrow. Re-derive the candidate list from "ragged trimming is unsafe", not -from "eager is unsafe". +**The CUDA illegal memory access this document hands off has since been root-caused and FIXED** — see §3's dated entries: a ragged graph key derived over a window-less step replayed the previous ragged step's row state; fixed at the key derivation (`6a59b01e53`, with the WAR-staging hardening in `c7e6d78ac2`/`74f1a82d06`/`5c38441a00`), validated 131 steps / zero IMA with ragged fully active. The rest of this section is preserved as the record of how the framing went wrong: I spent most of a day believing the fault was confined to a path I had just added (running ragged on steps with no CUDA graph); the final run disproved that with the path gated OFF. Everything in §4 was designed under that wrong hypothesis — the falsifications stand as facts, but re-derive any future candidate list from "ragged trimming is unsafe", not "eager is unsafe". --- @@ -72,7 +57,7 @@ accuracy. GSM8K passing proves nothing on its own. | # | Item | |---|---| -| **IMA** | §3. Blocks everything else. | +| ~~IMA~~ | **Closed** — root cause and fix in §3 (2026-08-04 ~16:30 entry); validated 131 steps, zero IMA/hang/tripwire, ragged fully active. | | #19 | Verify budget is quantised to the tier ladder. Measured cost up to **73%** of the planner's own objective. Not implemented. | | #11 | Sync control never completed. **Do not cite any sync-related conclusion.** | | #20 | Three-way differential incomplete (compact leg needs a non-blocking run) | @@ -340,7 +325,6 @@ Useful switches, all wired through `tmp/run_dspark_e2e.sh`: |---|---| | `LAUNCH_BLOCKING=1` | `CUDA_LAUNCH_BLOCKING=1` — the only thing that passes | | `NO_OVERLAP=1` | `disable_overlap_scheduler=True` | -| `EAGER_RAGGED=1` | re-enable the gated eager-ragged path | | `ASSERT_ROWS=1` | row-map staleness assertion (skipped during capture) | | `SANITIZER=1` | compute-sanitizer — see §5, not currently viable | @@ -438,19 +422,7 @@ rank is untried and is the obvious next move. ## 6. Next steps -1. **Validate the coarse prepare() guard** (§3 root cause) — a gated and an - eager run with `COREDUMP=1`. A pass means both complete ≥128 steps with - non-zero `steps_ragged` and no new corefile; do not accept anything less. -2. If it still faults: the corefile names the kernel again. The remaining - unguarded staging is then *outside* `prepare()` — sweep `model_engine.py` - and `py_executor.py` for persistent pinned buffers with `non_blocking=True` - consumers (the `_pinned_host` sites are already event-guarded). -3. After the pass: replace the coarse guard's protection of the worst offenders - with the base class's snapshot idiom where profiling justifies it, drop the - dead capture-skip conditionals, re-evaluate the `4f81cc5534` gate, and clean - up the ~34 corefiles (keep one as evidence). -4. Only after the IMA: #19 (budget quantisation, the largest measured gap), - then #11 and #20. +1. #19 (budget quantisation — the largest measured gap, up to 73% of the planner's own objective), then #11 (sync control; until it completes, do not cite any sync conclusion) and #20 (the three-way static / cap-accept / compact differential). All IMA work has landed: coarse prepare() WAR guard with the `_prepare_impl` subclass rule (`c7e6d78ac2`), KVCacheManagerV2 snapshot staging (`74f1a82d06`), bounded/zeroed overlap staging inputs (`5c38441a00`), window-consistent ragged graph key (`6a59b01e53`), and the eager-ragged gate reverted on its merits (`35a1ac4c4a`). ### Do not @@ -464,24 +436,6 @@ rank is untried and is the obvious next move. --- -## 7. Unpushed work - -`origin/confidence_head` is at `d36c9b4757`. Two local commits plus the -coarse-guard working tree, none end-to-end verified yet: - -| commit | what | note | -|---|---|---| -| `0bb36c11b6` | event-guard the ragged row-map staging + `_pinned_host` sites | did **not** stop the fault alone, but closes real WAR holes; keep | -| `9bc5dec76c` | skip staging copies during capture | **no-op** — `prepare()` never runs inside the capture region (§3); fold into cleanup | -| *(uncommitted)* | coarse WAR guard around all of `DSAtrtllmAttentionMetadata.prepare` | the actual fix candidate, in validation | - -Two fixes worth keeping regardless of the IMA, already in the above: -`_pinned_host` (persistent staging buffers — PyTorch does not extend an async -H2D source's lifetime, and three sites were passing temporaries) and the -`TLLM_DSPARK_DISABLE_OVERLAP` diagnostic switch. - ---- - ## 8. Environment - 8x B300 SXM6 (SM103, 275 GB/card), x86_64, node pool `bia*`, partition `batch` diff --git a/docs/dspark_ragged_verify_status.md b/docs/dspark_ragged_verify_status.md deleted file mode 100644 index b0bcffe04308..000000000000 --- a/docs/dspark_ragged_verify_status.md +++ /dev/null @@ -1,452 +0,0 @@ -# DSpark ragged verify — 目标与现状 - -配套文档：`docs/dspark_confidence_schedule_goal.md`（设计与方案对比）。 -本文档记录**工作状态**：做到哪、验到哪、哪些还不能信。 - -最后更新：2026-07-31。分支 `pr-17056`，基线 PR [NVIDIA/TensorRT-LLM#17056]。 -**所有改动均未提交**（工作区 15 个修改 + 4 个新增文件）。 - ---- - -## 0. 观测范围 —— 所有实测数字只适用于 x86 DGX B300 - -本文档里每一个耗时、吞吐、显存、page-cache 数字都是在**一台机器**上量的： - -| | | -|---|---| -| 节点 | `umb-b300-001` | -| 分区 | `b300@ts6/dgx-b300@ts1/8gpu-256cpu-2048gb` | -| 架构 | **x86_64** + NVIDIA B300 SXM6，compute_cap **10.3（SM103）**，275 GB/卡 | -| 主机内存 | 2015 GB | -| 权重存储 | `/home/scratch.trt_llm_data`，共享 **NFS** | - -**换集群这些数字全部作废，尤其是 GB300 / `b300-nvl8` / `gb300`**：那是 -Grace-Blackwell，**aarch64** 架构、CPU 与 GPU 之间是 NVLink-C2C 统一内存、 -主机内存容量与存储后端都不同。具体会变的： - -- **编译目标**：`103-real` 是为 x86 B300 编的。GB300 需要 aarch64 工具链和对应镜像 - （`LLM_SBSA_DOCKER_IMAGE`，见 `jenkins/current_image_tags.properties`）， - 不能直接复用本仓库现有的 `cpp/build`。 -- **加载耗时结构**（§6.4b 那张表）：预取 2 h / 装配 40 min 是这台机器的 NFS 和 - x86 内存带宽决定的。Grace 的内存带宽和统一内存寻址会改变整个比例。 -- **page cache 复用策略**：依赖 2 TB 主机内存能装下 835 GB 权重。GB300 的内存 - 配置不同，这个前提要重新验证。 -- **显存余量**：275 GB/卡 是 B300 SXM6 的数字。 -- **Slurm 账号与分区**：`trt-llm_b300` 对应 x86；GB300 走 `trt-llm_gb300`。 - -**结论性的东西（六条 P0 的定位、kernel 的 next_n 假设、可观测性设计、 -"空输出不是通过"那几条教训）与平台无关，可以直接迁移。 -所有带具体数字的运行特征都不行。** - ---- - -## 1. 目标 - -主线：让 **ragged verify**（每个 request 各自的 verify 长度）在 TensorRT-LLM 里真正跑通。 -这是 confidence-head 调度收益的唯一来源，**必须实现**，不接受退化成 batch-uniform K 的方案。 - -验收条件，缺一不可： - -| # | 条件 | -|---|------| -| 1 | 目标模型 DeepSeek-V4-Pro，真实 checkpoint | -| 2 | GSM8K 端到端精度不低于 baseline | -| 3 | **必须验 ADP**（目标平台已定为 **B300 / SM103 / 8 卡**）：`tensor_parallel_size=8` + `moe_expert_parallel_size=8` + `enable_attention_dp=True` | -| 4 | **确认 ragged 真的生效**，而非「配置打开了但代码走了 uniform fallback」 | -| 5 | kernel 侧兼容（indexer top-k / DSv4 compressor 不能静默算错） | -| 6 | CUDA graph 侧兼容（真在 replay，不是被迫 eager） | -| 7 | **overlap scheduler 开启也要正确**——这是 default 配置，不能只验关闭的路径 | - -关于条件 2 的性质：调度只决定「多少 drafted 位置送去 target」，不改接受规则， -所以输出分布必须不变。**精度掉点 = 有 bug，不是吞吐 tradeoff。** - -关于条件 4 的必要性：该功能的每一种失效都是**静默**的——planner 拒绝 trim、 -cost model 退化、batch 错过捕获形状掉出 graph、部分加窗被拒——每一种都仍然产出 -正确输出和基线精度。**GSM8K 通过本身不能证明功能跑过。** - ---- - -## 2. 现状 - -### 2.1 已验证（有运行数据） - -| 项 | 证据 | -|---|---| -| ragged top-k kernel 正确 | `tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py` **10 passed** on B200，含 CUDA graph capture/replay，且与 `torch.topk` 参照比对 | -| 无回归 | ragged + `hw_agnostic` DSpark 共 **227 passed** | -| C++ 编译通过 | `indexerTopK.cu` / `IndexerTopKOp.cpp` / `compressorKernels.cu` / `compressorOp.cpp` 已编入 `.so` | -| checkpoint 可用 | `mtp.2.confidence_head.proj.weight`、`mtp.2.markov_head.markov_w1/w2.weight` 存在 | - -### 2.2 未验证（这是最重要的一节） - -**整条 e2e 从未跑通过。** 下面六条 P0 全部只经过读码与审计定位， -**没有一条经过运行验证**。逻辑上都成立、都能指到具体行号，但那不等于对。 - -e2e 冒烟第一次跑就撞上 `PYTHONPATH` 未设——说明这类问题还会有更多。 - -### 2.3 已知阻塞（非本工作引入，但会挡住验证） - -1. **NVRTC 编译 MLA FMHA kernel 失败** - `Failed to preprocess kernel fmhaSm100aKernel_QkvBfloat16OBfloat16HQk576HV512...: - NVRTC_ERROR_COMPILATION`。影响所有 MLA 模型，**很可能包括 DSv4-Pro**。 - 若冒烟撞上它，这是先于一切的阻塞。 -2. `test_context_sparse_attention_mqa` 6 个失败 —— 已用有效对照（stash 三个 sparse - 文件后跑同一子集，基线同样 6 failed）证明是既有问题。 - 注意：另外 53 个 sparse 失败**未逐一对照**。 -3. `.venv-3.12` 的 cutlass 曾完全坏掉（`import tensorrt_llm` 起不来），已修复， - 详见 §5。 - ---- - -## 3. 已完成的改动 - -### 3.1 kernel 层 - -**`indexerTopK.cu`** —— `topKPerRowDecode` 原来在 kernel 内部用标量 `next_n` -同时推导 row→request 映射和 causal 上界： - -```cpp -int seq_len = seqLens[rowIdx / next_n]; -int actual_kv_len = seq_len - next_n + (rowIdx % next_n) + 1; -``` - -新增可空 `rowKvLens`：非空时 `actual_kv_len = rowKvLens[rowIdx]`，空时保持原式。 -数学上是严格推广（请求 verify `v` 个位置时第 `o` 行可见 `kv_len - v + o + 1`， -代入 `v = next_n` 即还原），所以 uniform 路径逐位不变。 - -`IndexerTopKOp.cpp` 新增 `Tensor? row_kv_lens`，并在 ragged 时**替换**而非放宽 -`seq_lens.size(0) * next_n == numRows` 检查——放宽的话，恰好整除的 ragged batch -（如 `verify_lens=[6,4,1,1]`，total=12，bs=4 → next_n=3）会通过检查然后静默算错。 - -**`compressorKernels.cu`** —— `NEXT_N` 是模板参数，但 body 里只用于 `sp`、 -`last_token_idx` 和两个循环上界，且两个循环**本来就有守卫** -（Phase 1 的 `token_idx < kv_len`、Phase 3 的 `c >= num_compressions`）。 -所以把 `NEXT_N` 从「确切长度」降级为「batch 内最大值（编译期上界）」即可， -真实 per-request 计数走新增的可空 `new_tokens_per_seq`。 -单次 launch、固定 grid、**CUDA graph 拓扑不变**。 - -### 3.2 六条 P0 - -| | 问题 | 修法 | 触发条件 | -|---|---|---|---| -| P0-1 | `attn_metadata.prepare()` 是 `ragged_verify_lens` 的唯一消费者，却在它发布前 86 行就跑了 | 拆出 `_publish_ragged_verify_lens` 提前调用 | 全部 | -| P0-2 | `capture()` 用 `bs*(key[1]+1)` 定宽；graph `seq_lens` 冻结 | 改用 `key[5]`（ragged bucket）；ragged 且行数不变时放行 `seq_lens` 刷新 | 第一个 trim 步 | -| P0-3 | `expand_per_gen_token` 在 capture body 内做 pinned 分配 + H2D | 新增常驻 `row_req_idx_cuda`，refresh 改成 `index_select(out=) + add_` | replay | -| P0-4 | `on_update_kv_lens` 无条件整除覆盖 compressor 的 `next_n` | 抽出 `_sync_gen_tokens_per_seq` 单一写入点 | 每次 forward | -| P0-5 | drafter 用定长 `Kp1` 跨步读 flat-indexed 的 captured hidden states | 改用 `qo_indptr`；`base + gidx` 补 clamp | **overlap 开启** | -| P0-6 | overlap rewind 读到已被下一步覆盖的 `py_verify_len` | `SampleStateSpec.verify_lens_snapshot` | **overlap 开启** | - -P0-2 的两个半边必须**一起**修：`num_tokens` 来自 `seq_lens.sum()`， -而 `capture()` 从 `key[1]` 定宽，只修一边会让 `input_ids.shape[0] != num_tokens`。 -最大 bucket 恰好等于那个乘积，而 warmup 捕的正是它——所以捕获阶段完全干净， -**第一个真正 trim 的 step 才崩**。 - -P0-3/5/6 属于「跑出来的数字看着正常但其实是错的」：replay 读已释放内存、 -跨请求串读、KV 回退量错。**这类问题不会报错。** - -### 3.3 可观测性（`tensorrt_llm/_torch/speculative/dspark_observability.py`） - -对齐 SGLang 的 `SGLANG_RAGGED_VERIFY_MODE` 语义： - -| 开关 | 作用 | -|---|---| -| `TLLM_DSPARK_RAGGED_VERIFY_MODE=static\|cap-accept\|compact` | 选路径。三种模式**均已实现** | -| ~~`TLLM_DSPARK_FORCE_VERIFY_LENS=1`~~ | **已删除**。改用 `tests/microbenchmarks/dspark_make_steep_sps_table.py` 造陡峭 cost table,让**真实 planner** 自己产出不等窗口 | -| `DSparkRaggedStats.assert_ragged_active()` | 断言 ragged 真的生效 | - -`assert_ragged_active()` 检查:`steps_ragged > 0`、`distinct_verify_lens >= 2`、 -`trim_ratio > 0`,以及 graph miss 中**可归因于 ragged 的那几类**,失败时带完整计数器摘要。 - -最后一条原为 `graph_eager == 0`,**已修正**:连续批处理下,某个 rank 带 context 请求 -就会让全体掉出 replay,uniform 路径同样如此(实测约 1.7%),所以按原判据这个断言在 -任何真实负载上都不可满足。现在按原因归类,只对 `key_not_captured`、 -`bs_not_supported`、`peer_shape_mismatch` 三类失败——它们分别表示 bucket 网格错了、 -批次档位没捕获、以及跨 rank 形状发散。 - -**`cap-accept`（已实现）**:算出每请求窗口并**按窗口提交**,但**不缩小交给 target 的 -token 轴**——整块照算,超出窗口的接受位置被丢弃。它绝不能 alias 成 `compact`:唯一 -价值就是那个对照。 - -窗口写在**独立属性 `py_verify_cap`** 上,不碰 `py_verify_len`。于是布局侧 -(`model_engine` / `cuda_graph_runner` / `dsa.py`)完全看不见窗口,自动走整块; -`_verified_len` 因为 `py_verify_len is None` 自动返回整块,**KV rewind 随之正确** -(`rewind = 整块 − capped`)。新模式因此是纯增量,不必在「为 compact 调了三个 bug -才对」的记账代码里再加一组条件分支。 - -**截断必须在设备端、acceptance 内部**——`interface.apply_accept_caps`,由 -`sample_and_accept_draft_tokens` 的两个出口调用。这一点第一版写错了(截断放在主机侧 -`update_requests`),原因值得记下来: - -`num_accepted_tokens` **在同一次 forward 的后续、设备上还要被消费**。`dspark.py:606-652` -里 drafter 用它挑下一块的条件隐状态、把中间接受 token 回填进滚动 KV 窗口,并推进自己的 -持久解码位置(`_ctx_len += nacc`)。等主机同步后再截断,设备已经按**未截断**的数推进过 -drafter 状态,而主机只提交了 capped 前缀——输出因为验证是权威的仍然无损,但 drafter -状态从此永久漂移,这个模式也就不再是它存在意义所在的那个可信参照。SGLang 截在同一 -位置(`cap_correct_len`,`dspark_accept.py:138`),同样是在任何下游读取之前。 - -**不能直接复用 `is_ragged_verify` 来触发那个 clamp**:它同时意味着「在此停止接受」和 -「draft 缓冲已按窗口打包」(`_padded_gen_draft_tokens` 会切 -`draft_tokens[:total_verify_tokens − num_gens]`)。cap-accept 只要前者,缓冲是普通矩形。 -借用它会连带走进打包分支,把矩形当 ragged 解包——一个请求的 draft 记到另一个头上, -**全程不报错**。所以是独立字段 `SpecMetadata.accept_caps`。 - -单位与 `verify_lens` 一致:`num_accepted_tokens` 是「接受 draft 数 + 1」,所以窗口 `w` -对应的 cap 是 `w + 1`。 - -它产出 `compact` 结构上拿不到的数:**被丢弃的、本可接受的位置数**——裁剪的真实接受 -代价。`compact` 下那些位置根本没算,只能用 `trim_regret_rate` 卡下界。 - -**逐请求记录,不是只记总量。** 理由是总量分不出这两种情况: - -| 500 个请求共丢 1000 | 含义 | -|---|---| -| 每个丢 2 | planner 均匀小幅收割,大概率是合理取舍 | -| 20 个各丢 50、480 个丢 0 | planner **系统性掐死一小撮请求** | - -两者 `cap_trim_tokens` 和 `accept_loss_per_request` 完全相同,结论却相反。所以另外记 -`requests_cap_trimmed` / `cap_trim_max` / `cap_trim_hist` / `cap_trim_concentration` -(有损失的请求占比):平均损失相同的情况下,concentration 0.9 是普遍轻微裁剪,0.02 是 -调度器拿 2% 的流量在补贴吞吐——只有后者是问题。 - -链路:`apply_accept_caps` 写进持久 `[bs]` 缓冲 → DSpark worker 的返回字典 -`cap_trim_lens` → 采样器 store(slot 索引)→ 随**已有的** host 拷贝回来。因此这个测量 -**不额外引入任何同步**,也不需要设备端累加器。 - -**陈旧缓冲是这里唯一的真风险**:缓冲持久且按 slot 索引,不产出 caps 的步(static、以及 -其它任何投机模式)必须往本批次的 slot **写零**而不是跳过,否则 slot 被回收给新请求后 -会报上一个占用者的损失。`_process_outputs` 里显式 `zeros_like`,并有测试卡住。 - -代价是**不省任何算力**,所以它是诊断模式,永远不是服务配置。`assert_ragged_active()` -的 `require_trim` 因此对该模式豁免,否则那条门在它唯一该跑的模式里恒不可满足。 - -**「循环依赖」这个理由已被证伪**:产出 cost table **不需要** ragged 能跑。 -`dspark_sps_profiler` 会把 `RAGGED_VERIFY_MODE` 钉成 `static`,并通过逐个 verify_len -重建引擎来扫成本曲线——采表全程 ragged 关闭(`sps_real_final.json` 就是这么采的)。 - -真正的问题是另一个:**即使有真实 table,planner 在接受率够高时依然不裁**——在 -DeepSeek-V4-Pro-DSpark 上是每一种实测工况(accept_len 3.3–4.1)。所以裁剪路径需要 -一个刻意的手段才能被执行,否则它会带着全绿的测试合入。 - -现在这个手段是**造一张陡峭的 cost table**,而不是按 batch 位置轮转的开关。区别是 -本质的:位置与 confidence 正交,轮转只能复现 ragged 的**形状**、复现不了**策略**, -测不出「置信度高的是否拿到更长窗口」。造表则驱动真实 planner,窗口经同一个 -confidence 排序的 top-k 分配。三个只在真裁剪时可达的 bug 正是这样被逼出来的。 - ---- - -## 4. 剩余工作 - -| 优先级 | 项 | 说明 | -|---|---|---| -| P0 | e2e 冒烟通过 | 两个配置各一遍：`disable_overlap_scheduler=True` 与 **default（overlap 开启）** | -| P0 | 排除 NVRTC 阻塞 | 若 DSv4-Pro 起不来，其余都无从谈起 | -| P1 | **SPS cost table profiler** | 没有它 planner 必然不 trim，**任何性能结论都是空的**。参考 SGLang `python/sglang/benchmark/dspark_sps_profiler.py`：服务端 `RAGGED_VERIFY_MODE=static` + record flag，扫 batch size 记 step time | -| P1 | GSM8K + ADP | 冒烟通过后才有意义 | -| P1 | graph 按 token 数分桶 + ADP tier 协商 | SGLang `compute_target_verify_graph_key` 对 ragged 返回 `(graph_num_tokens, graph_num_tokens)`——两个轴都变 token 数；DP 下各 rank 共用同一 tier 并一起降档 | -| ~~P2~~ | ~~实现 `cap-accept`~~ | **已完成**。隔离「kernel 算错」与「planner 没触发」 | -| P2 | 对照剩余 53 个 sparse 失败 | 目前只对照了 mqa 那 6 个 | - ---- - -## 5. 任务清单（截至 2026-07-31） - -| # | 任务 | 状态 | -|---|------|------| -| 1 | 目标文档 | 完成 | -| 2 | indexer top-k kernel ragged 化 | 完成（B200 上 10 用例通过，**但在错误镜像下跑的，需重验**） | -| 3 | DSv4 compressor ragged 化 | 编译通过，**未运行验证** | -| 4 | 审计 DSv4 gen 路径剩余 ragged 缺口 | 未完成 | -| 5 | CUDA graph 兼容 | 见 P0-2/P0-3，**未运行验证** | -| 6 | GSM8K + ADP 验证 | 未开始 | -| 7 | mode 开关 + 可观测性 | 完成 | -| 8 | graph 按 token 分桶 + ADP tier 协商 | 未完成 | -| 9 | SPS cost table profiler | 代码完成（1503 行 + 42 测试），**从未在真实 engine 上跑过** | -| 10 | `on_update_kv_lens` 刷新 row_kv_lens | 完成 | -| 11 | overlap 开启的正确性验证 | 未完成（**default 配置，必须验**） | -| 12-16 | P0-2 … P0-6 | 代码完成，**六条全部零运行验证** | - -**当前阻塞**：环境重建中（见 §6）。六条 P0 一条都没经过运行验证，这是最大的未知。 - ---- - -## 6. 环境（复现必读） - -### 6.1 容器镜像 —— 先看这里 - -**镜像由仓库指定，不要自己挑。** - -``` -jenkins/current_image_tags.properties:16 -LLM_DOCKER_IMAGE=urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:\ - pytorch-26.05-py3-x86_64-ubuntu24.04-skip-tritondevel-202607211045-16608 -``` - -本工作前期用错了镜像（`...trt10.16.1.11-...202607151440-16194`，从本地 `docker images` -里挑的），代价是**一整轮关于 cutlass 版本冲突的错误分析**：错误镜像的系统 cutlass 与 -`requirements.txt` 的 pin 不配套，我先归因给自己的 pip 操作，又归因给「repo 与镜像的 -结构性冲突」，两次都错。正确镜像的 cutlass 是自洽的。 - -### 6.2 cutlass 版本 - -正确镜像自带 **4.6.1 且四个子包全部一致**。`requirements.txt` pin `==4.5.0`，但 -`constraints.txt` 只要求 `>=4.4.2`，所以**安装依赖时排除 cutlass、沿用镜像自带的**： - -```bash -grep -v "nvidia-cutlass-dsl" requirements.txt > requirements.nocutlass.txt -pip install -r requirements.nocutlass.txt -``` - -过滤后的文件**必须写在 `requirements.txt` 旁边**——里面有 `-c constraints.txt` 相对路径， -放 `/tmp` 会报 `Could not open constraint file`。 - -为什么不能装 4.5.0：那套包没有 `libs-core`，装进 `--system-site-packages` 的 venv 后， -前端是 4.5.0 而 MLIR verifier 是系统的 4.6.1，DSv4-Pro 会死在 -`'nvgpu.cvt_fptrunc' op operand #0 must be ... 1-d vector, but got 'f32'` -（`f8E8M0FNU` 就是 DSv4 的 ue8m0 scale）。 - -**验证 venv 不能只看 `import` 成功**，要看 `pip list | grep cutlass` 各子包版本是否一致。 - -### 6.3 venv - -```bash -/usr/bin/python3 -m pip install --user 'virtualenv<22.0,>=20.29.1' -/usr/bin/python3 -m virtualenv --system-site-packages .venv-3.12 -``` - -`virtualenv` 而非 `venv`：镜像的 `ensurepip` 不可用。 -`--system-site-packages` 必需：torch 在镜像的 dist-packages 里。 - -### 6.4 编译 - -```bash -python3 scripts/build_wheel.py --cuda_architectures "100-real" -j $(nproc) -``` - -- **PATH 只能追加不能替换**：`/usr/local/cuda/bin` 在镜像默认 PATH 里，覆盖掉会让 - cmake 报 `No CUDA compiler found`。 -- **配置失败后必须 `rm -rf cpp/build`**：CMake 会把失败的编译器探测缓存成 - `CMAKE_CUDA_COMPILER:FILEPATH=NOTFOUND` 并在下次复用，环境修好了也不会重新探测。 -- **目标平台是 x86 DGX B300（SM103），编译用 `--cuda_architectures "103-real"`。** - GB300（`gb300` / `b300-nvl8`）是 aarch64，需要 SBSA 镜像和单独的编译产物。 - B200 是 SM100；`100-real` 纯 SASS 无 PTX，装不进 B300。B300 有 275 GB/卡， - B200 只有 178 GB —— DSv4-Pro 权重 104 GB/卡，B200 上余量很紧。 -- 容器需 `--user $(id -u):$(id -g)`（scratch 是 NFS root_squash）。 - -### 6.4b DSv4-Pro 的加载时间结构 —— 复用同一节点的 page cache - -> **仅适用于 x86 DGX B300 + NFS（见 §0）。GB300 等其他平台需重新实测。** - -冷启动实测（`umb-b300-001`，权重在 `/home/scratch.trt_llm_data`，NFS）： - -| 阶段 | 耗时 | 可观测性 | -|---|---|---| -| import + 8-rank NCCL 建组 | ~1 min | 有日志 | -| **预取（NFS → 内存）** | **~2 h** | 只在整个 13 GB 分片读完时打一条 `Finished prefetching`，中间静默 90 min 以上 | -| **装配（内存 → 显存）** | **> 40 min** | tqdm 进度条，但 ETA 严重低估（并发流均匀推进、集中收尾） | - -**冷启动总计 3 小时以上，`batch-short` 的 4 小时窗口装不下「加载 + 验证」。** - -**解法：复用同一节点。** 该分区节点有 **2 TB 内存**，一次加载后 page cache 会保留 -约 1.1 TB，足以完整覆盖 835 GB 权重。作业结束后 cache 不会被清掉，所以： - -```bash -# 当前作业还没结束时就把后续作业排到同一节点，接力启动、cache 仍热 -sbatch -A trt-llm_b300 --qos=batch-short \ - -p "b300@ts6/dgx-b300@ts1/8gpu-256cpu-2048gb" \ - --gres=gpu:8 --nodelist=<同一节点> -t 4:0:0 --wrap='sleep 14400' -``` - -查 cache 是否还热：`free -g` 看 `buff/cache` 列。 - -**判活提醒**：预取阶段日志静默、GPU 利用率 0%、显存不变**同时出现是正常的**。 -唯一可靠的判据是固定时间窗口内 `/proc//io` 的 `read_bytes` 增量 -（实测健康值约 112 MB/s）。不要用 tqdm 的 ETA 估算剩余时间。 - -### 6.5 跑测试 / 跑模型 - -容器额外需要：`-e USER=`（否则 `getpass.getuser()` 因 UID 无 passwd 条目而 -KeyError）、`-e PYTHONPATH=/code/tensorrt_llm`（直接跑脚本时；`sys.path[0]` 是脚本所在目录）、 -`-e LLM_MODELS_ROOT=/home/scratch.trt_llm_data/llm-models`、 -`--ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --shm-size=32g`。 - -**srun 步骤内 docker 的 `--gpus` 只能用 `all`**，写 `device=N` 会被代理拒绝。 - -**任何 `tensor_parallel_size > 1` 的脚本必须有 `if __name__ == "__main__":`**—— -`MpiPoolSession` 靠重新执行模块来 spawn worker，没有守卫会无限递归 spawn， -表现为**静默挂死而非报错**（本工作因此浪费 2 小时窗口）。 - -**Slurm**：Blackwell 分区 `DenyQos=batch,interactive-isolated,oversubscribe`，必须 -`--qos=batch-short`（上限 **4 小时**）。报 `Invalid qos specification` 不表示没有权限。 - ---- - -## 7. 教训（给接手的人） - -### 7.0 诊断动作本身会破坏运行 —— 两次杀掉自己的冒烟 - -冒烟从未跑到 ragged 代码，两次都死在加载阶段，**两次都是我杀的**： - -1. **B200**：据「GPU 利用率 0% + 显存不变 + 日志静默」判定挂死并杀掉。 - 事后 `py-spy` 证明它在 `weight_loader.py:338` 正常从 NFS 预取权重。 - 残留进程占着 135 GB/卡不释放，导致下一次启动直接 OOM。 -2. **B300**：为避免重犯，改用 `py-spy dump` 做「正确诊断」——附加超时， - 把 worker 留在 `ptrace_stop` 冻住了整个运行；再用 - `pkill -f "py-spy dump"` 去救，那个模式匹配到 srun 自己的命令行，把 step 杀了。 - `/proc//io` 显示它冻住前已读 13.2 GB，本来是好的。 - -**三条铁律：** - -- **不要 py-spy 附加正在加载的 worker。** 附加开销大到会超时，且会把进程留在 - ptrace 停止态。要用也只在已经确认卡死之后用。 -- **不要 `pkill -f <字符串>`。** 会匹配到 srun / docker 自己的命令行。 - 清理只按 PID：`nvidia-smi --query-compute-apps=pid --format=csv,noheader`。 -- **判活只用 `/proc//io` 的 `read_bytes` 是否增长。** 只读、零副作用， - 是加载阶段唯一安全且可靠的存活信号。 - -**加载阶段的正常表现**：GPU 利用率 0%、日志静默、显存不变——三者同时出现也是正常的。 -DSv4-Pro 权重 835 GB 从共享 NFS 冷读，20-30 分钟属于合理范围。 - -**杀掉运行之后必须清残留**：孤儿进程会一直占着显存，下一次启动就会 OOM。 -`for p in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader); do kill -9 $p; done` - -### 7.1 空输出不是通过 - -本工作中反复栽在这一条，四次： - -1. **kernel 输出顺序**：`indexer_topk_decode` 不保证输出顺序（radix split-work 的 merge - 顺序不定），只有 top-k **集合**是契约。逐元素比数组得到 ~95% 假阳性 mismatch。 - 必须 `sort(dim=-1)` 后比，并用 `torch.topk` 做参照。 -2. **对照实验没跑起来**：一次 baseline 对照用了 `--gpus device=7`，docker 直接 - `invalid gpus device: 7` 退出，空输出被当成了有效数据。 -3. **冒烟挂死**：连续两次报告「仍在加载 835G」，实际 GPU 1-7 各 4 MiB、利用率全 0。 -4. **编译白跑**：CMakeCache 缓存了 NOTFOUND，容器「在运行」但 0 个 object 文件。 - -统一的教训：**「进程在跑」「没有报错」都不是进展的证据**。要有正向信号—— -object 文件数、GPU 显存占用、明确的成功输出。 - -### 7.2 工具缺陷要立刻修，不要记下来等下次 - -`run.sh` 里的 `... | tail -80` 会缓冲全部输出，进程退出前看不到任何进度。 -我在发现后说了「下次改」，没有立即改——结果冒烟静默挂死 2 小时才被发现， -而改成 `python3 -u ... | tee` 之后**一分钟**就看到了完整的 MPI 错误。同样的信息一直都在。 - -### 7.3 局部通过不等于整体正确 - -一度宣告「venv 已修复」，依据只是 `import cutlass` 和 `import tensorrt_llm` 成功。 -实际版本组合并不自洽，DSv4-Pro 起不来。**导入成功只证明模块能加载，不证明版本配套。** - -### 7.4 环境问题先查仓库怎么规定，不要自己拼 - -本工作在环境上的弯路（镜像、cutlass、PATH、CMakeCache、constraints 路径、ensurepip） -**没有一个是产品代码的问题**，全部是容器调用姿势，而每个都吃掉一轮迭代。 -根源是「见招拆招」——需要 virtualenv 就改 PATH，需要过滤 requirements 就写 /tmp， -每次改动都没检查会不会破坏别的东西。 -正确做法是一开始就照 `jenkins/current_image_tags.properties` 和 -`.dspark-logs/build.log` 复现官方姿势。 - -### 7.5 这个功能的所有失效都是静默的 - -planner 拒绝 trim、cost model 退化、batch 错过捕获形状掉出 graph、部分加窗被拒—— -每一种都仍然产出正确输出和基线精度。这是 §3.3 那套可观测性存在的全部理由。 -**新增任何决定 verify 长度的代码路径，都要同时接上 `record_step`**， -否则 `assert_ragged_active()` 会给出虚假的通过。 diff --git a/docs/dspark_review_2026-08-04_followups.md b/docs/dspark_review_2026-08-04_followups.md deleted file mode 100644 index b99a547903cc..000000000000 --- a/docs/dspark_review_2026-08-04_followups.md +++ /dev/null @@ -1,439 +0,0 @@ -# Review: the five follow-up commits of 2026-08-04 evening - -Scope: `72dd803141..fab6291e2d` on `confidence_head` — - -| commit | title | -|---|---| -| `72dd803141` | [fix] Interpolate the SPS cost table between breakpoints | -| `c117be89a1` | [feat] Calibrate DSpark confidence and measure the planner's inputs | -| `50675cb181` | [fix] Make the ragged fit's consistency checks able to fail | -| `0b524435fc` | [feat] Pin the DSpark verify length at runtime | -| `fab6291e2d` | [chore] Measure the SPS table against the deployment, not a rebuild | - -Method: seven independent review passes (one per commit plus two cross-cutting -lenses: CUDA-graph capture/replay safety and overlap-scheduler -write-after-read — the two bug families the IMA series just closed), then an -adversarial verification pass per finding whose default stance was that the -finding is wrong until the code proves it. 38 findings raised, **34 confirmed, -4 refuted**; the refuted ones are omitted. Several verifications executed the -cited code paths rather than reading them (noted inline). - -Severity is scoped to the feature: nothing below corrupts served tokens -(verification guarantees them). "High" means a feature's primary artifact or -purpose silently fails. - -How to read an entry: *mechanism* is what the code does, *failure* is a -concrete input→wrong-outcome chain, *fix* is the smallest correct change the -reviewers could see. Line numbers are as of `fab6291e2d`. - ---- - -## A. STS calibration pairs confidence with the wrong labels — HIGH (as a group) - -Four findings converge on one conclusion: **the fitted temperature table — -`c117be89a1`'s primary artifact — is fit on mispaired samples in every -supported collection mode**, and nothing downstream can detect it. - -### A1. The label is one (or two) draft passes behind the confidence — deterministic - -`spec_sampler_base.py:311`, `dspark_sts.py:199` · medium - -Mechanism: one forward(i) first *verifies* the block drafted by the previous -forward (B_{i-1}; `dspark.py:822` computes the accepted count), then *drafts* -the new block B_i and stashes its confidence C_i (`dspark.py:789-791` → -`dspark_sts.py:199`, which unconditionally replaces the single mutable -`_stash_logits`). `record()` runs later, inside `update_requests` -(`spec_sampler_base.py:311`), pairing `accepted = A(B_{i-1})` with whatever -the stash holds — which is always C_i. Non-overlap loop: off by one -(`py_executor.py:4634` forward before `:4673` update). Overlap loop: off by -two (`:5132` forward before `:5139` update *of the previous batch's state*). -Because the stash overwrite happens inside the very forward that produces the -label, **no execution order leaves the correct C_{i-1} in the stash**. - -The window axis got exactly the missing treatment in the same commit -(`verify_lens_snapshot`; the `_verified_len` docstring at -`spec_sampler_base.py:226-233` names this identical hazard). The confidence -axis did not. - -Why it is invisible: the fitter (`dspark_fit_sts.py:171`, `--source auto`) -*prefers* the `logits_at_draft` column, whose docstring calls it "the correct -pairing"; the built-in stash-vs-buffer diagnostic compares two columns that -are both C_i, so they agree with each other while both are mispaired; the -82.2%-rows-differed measurement quoted in `load_shards` compares the columns -against each other, never against the label. The commit's own measurement of -one-pass-apart rows (correlation 0.14, mean |delta| 7.35) says what the fit -converges to: temperatures that flatten survival toward the marginal -acceptance rate, which then silently biases the serving planner's budget -argmax wherever the table is deployed. - -Fix: join on the draft-pass axis, not the wall-clock axis — ride the -confidence rows for block B in the SampleState produced by the forward that -verifies B (the `verify_lens_snapshot` pattern), or keep a two-deep ring -keyed by the existing draft sequence counter and have `record()` select by -the verified block's sequence. - -### A2. The recorder reads the confidence buffer with the wrong slot key - -`spec_sampler_base.py:312` · medium - -Mechanism: `record(slot=req.py_seq_slot, device_logits=self.sts_logits_buffer)` -indexes the raw `worker._confidence_logits` buffer by **py_seq_slot**, but -that buffer's rows are assigned by the DSpark worker's *private* allocator -(`_assign_slot`, a FIFO deque keyed by `py_request_id`, -`dspark.py:591-614`). `py_seq_slot` comes from the separate `SlotManager` -(set-based, immediate reuse; `resource_manager.py:2332-2344`). Nothing ever -synchronizes the two. They coincide only until the first request completes: -the set hands a freed slot back immediately while the deque defers it, so any -departure+arrival churn diverges them permanently. Every other reader -correctly goes through `confidence_row_for(py_request_id)` -(`dspark.py:319-334`); the recorder alone uses the raw index. Bounds always -hold (buffer is `[max_batch+2, K]` vs `py_seq_slot < max_batch`), so nothing -crashes — request A's accepted count is paired with request B's logits and -the shards look healthy. The late-read diagnostic reads back with the same -wrong key, so it is structurally blind to this too. The docstrings claiming -both halves are "indexed by the same py_seq_slot" -(`spec_sampler_base.py:203-206`, `dspark_sts.py:137-138`) assert an invariant -the code never establishes. - -Fix: wire a resolver instead of a raw buffer — e.g. -`self.sampler.sts_row_for = worker.confidence_row_for` at -`py_executor.py:3421-3427` and have `record()` take `py_request_id`. - -### A3. The draft stash survives CUDA-graph replay steps unbounded - -`dspark_sts.py:203` · low (verified with the full replay chain) - -Mechanism: `begin_step` invalidates `_host_logits`/`_stash_host` but never -`_stash_logits`/`_stash_slots`. Under graph replay, the Python inside the -captured draft body does not run, so the stash retains clones frozen at the -last *eager* step; `record()` re-materializes `_stash_host` from them -(`dspark_sts.py:236-244`) and appends stale-but-finite rows into -`logits_at_draft` as if they were this step's data (continuing requests keep -their slots, so lookups hit). The capture-time warning fires once during -warmup and any later eager step (mixed batch, uncaptured shape — routine -under in-flight batching) silently re-arms the stash for all subsequent -replays. `make_recorder_from_env` refuses cost-table and compact-ragged -regimes but never checks `cuda_graph_config`. - -Fix: tag the stash with the worker's draft sequence when armed -(`_confidence_stamp` already solves this class for the device buffer, -`dspark.py:781`) and treat a sequence mismatch as a missing entry. - -### A4. Recorder declines are silent, against the branch's own rule - -`dspark_sts.py:226` · high - -`record()` returns without recording on an out-of-range slot — no counter, no -log — and the capture-time stash refusal warns once but is not counted. This -module family's own documented rule (`dspark_verify.py:326-329`: even the -benign empty-batch decline is counted, because "this feature's failures are -all silent") is violated at exactly the call sites that A2/A3 make load-bearing: -a systematically wrong slot or a fully-graphed run yields empty or poisoned -shards with nothing in the stats to say why. Additionally, the construction --time censoring guards can be bypassed mid-run: `adopt_verify_len_pin` can -lower the pin below `block_size` (or clear it) *while a collection run is -writing shards*, reintroducing the censored-label regime the constructor's -ValueError exists to refuse. - -Fix: count dropped slots and stash refusals in a recorder stats dict surfaced -next to the planner's; refuse (or count) pin adoption below `block_size` -while a recorder is active. - ---- - -## B. The runtime verify-length pin silently never applies on tp_size == 1 — HIGH - -`py_executor.py:3516` (adoption), `:3341` (endpoint), `dspark_verify.py:263` -(queue) · from `0b524435fc` - -Mechanism: the only production call to `planner.adopt_verify_len_pin` sits -inside `if is_distributed:` (requires `tp_size > 1`; `py_executor.py:3433, -:3497`) — there is no else branch. The endpoint chain -(`openai_server.py` → `base_worker.py:875` → `PyExecutor -.set_dspark_verify_len_pin`) only *queues* the pin -(`request_verify_len_pin`, `dspark_verify.py:238-247` — docstring: "NOT -applied here") and returns 200 with the validated value. -`decide_verify_lens` reads only `_forced_verify_len` (`:393`), which nothing -else sets. The `base_worker.py:880` design comment — "rank 0 alone receiving -the RPC is sufficient" because "the step's own allgather" broadcasts it — is -false at tp_size == 1: there is no allgather. - -Failure: a single-GPU profiling server (the profiler's *default*, -`--tp-size 1`) accepts POST `/dspark/verify_len_pin` with 200; the pin sits -queued forever; `forced_steps` stays 0; and the server sweep -(`dspark_sps_profiler.py:1887-1890`) explicitly trusts the 200 and never -reads the effective length back — so every cell of the produced SPS cost -table is timed at planner-chosen lengths (full block, since sweeps run -without a cost table) but *labelled* with the pinned length. The fitted -table comes out flat in L and then silently mis-drives every planner that -consumes it. This is precisely the mislabelled-cell corruption the pin -feature exists to prevent, and `test_dspark_verify_len_pin.py` exercises -only the planner in isolation, never the executor wiring. - -Two adjacent silent modes (confirmed separately): -- **Static-mode engines** accept and adopt the pin, log "set", but - `decide_verify_lens` is never called (`compute_windows` is False) — the - endpoint should 400 the way it already does for a missing planner - (`py_executor.py:3341`). -- **Queue race**: `adopt_verify_len_pin` unconditionally resets - `_pending_verify_len_pin = -1` (`dspark_verify.py:263`), so a second pin - queued by the RPC thread between the decode loop's payload read (`:3504`) - and adopt (`:3516`, an allgather apart) is wiped after its endpoint already - returned 200. Fix: compare-and-clear. - -Fix (primary): `else: planner.adopt_verify_len_pin(planner.pending_verify_len_pin())` -after the `is_distributed` block — on a single-rank world the local queue *is* -the group's agreement. Have the sweep read back the effective value. - ---- - -## C. The table fingerprint records CLI flags, not engine facts — hard-rejects correct tables — HIGH - -`dspark_sps_profiler.py:2071` (write), `dspark.py:456-462` (check, from -`c117be89a1`) · from `fab6291e2d` - -Mechanism: the fingerprint dict is built from `args` (`"ep": int(args.ep_size -or 0)`, `"tp": args.tp_size` default 1, `"attention_dp"` store_true default -False) and attached whenever the run was not `--fit-only`. But: - -- The two **new deployment-measuring sources** (`--from-iter-log`, - `--base-url`) build no engine at all, so there is no reason to pass those - flags; a table swept over HTTP against a tp8/ep8/attention-DP deployment - carries `{tp:1, ep:0, attention_dp:False}`. -- Even the in-process sweep false-rejects whenever `--ep-size` is unset: - the write side records 0, while the loader compares - `mapping.moe_ep_size`, which `Mapping` normalizes to >= 1 - (`llm_args.py:4739` maps None→−1; `mapping.py:114` resolves −1→1 — chain - verified line by line). - -`check_table_fingerprint` (`dspark_planner.py:206`, raise at `:244`) then -refuses with "profiled on a different engine configuration… Re-profile on -this configuration" — and re-running the same profiler command reproduces the -refusal. This defeats the commit's own title (measure against the -deployment: the tables measured ON the deployment are exactly the ones -rejected), and violates `check_table_fingerprint`'s own documented mechanism -for facts the consumer cannot verify: key ABSENCE (compared keys must be -present in both dicts; absent ones go to the check-manually INFO line) — the -`ep: 0` encoding fabricates a fact instead of omitting it. - -Fix: record resolved engine facts when an engine exists (read the mapping -after construction); omit the key entirely when the fact is unknown -(HTTP/log sources, unset flags). - ---- - -## D. `padded_bs_hist` pollution is not actually fixed, and the fit publishes before validating — HIGH - -`model_engine.py:3789` (clear), `py_executor.py:3633-3644` (read), -`model_engine.py:3908-3986` (publish-before-validate) · from `50675cb181` - -Mechanism (pollution): the commit clears `_dspark_last_padded_bs` at -`fit_ragged_verify_lens` **entry** and its comment claims this prevents a -stale value from polluting "exactly the fallback steps the histogram exists -to distinguish". But repo-wide the attribute has three references: cleared at -fit entry, set on fit success, and **read unconditionally into -`record_step` on every stats-bearing step** (`py_executor.py:3640`). The -eager-ragged branch (`:3582-3599`) never enters the fit (it clears -`agreed_ragged_bucket` and `ragged_pad_verify_len` — but not this), and the -cap-accept branch never reaches the fit either. Both produce non-empty -`verify_lens`, so `record_step` passes its early-out and -`dspark_observability.py:327-328` books the LAST fitted step's padded row -count into `padded_bs_hist` for a step that was never padded. Eager-ragged -steps are frequent by the code's own measurement ("46 of the first 64 -steps"), so the histogram — added by this very range to reconcile the -negative-trim_ratio bug — can be dominated by phantom entries on exactly the -steps it exists to explain. Execution behavior is unaffected (the ceiling -rebase requires `bucket`), so this is observability-only — but this branch's -stated stance is that these counters are the sole detection channel for the -feature's silent failures. - -Fix (total, verified): pass `padded_bs` to `record_step` only when -`bucket is not None` — bucket is non-None exactly when the fit delivered. - -Mechanism (publish-before-validate): inside the fit, -`_dspark_last_total_tokens` (`:3908`), `_dspark_last_padded_bs` (`:3914`), -`runner.agreed_ragged_bucket` (`:3928`) and `runner.ragged_pad_verify_len` -(`:3971`) are all published **before** three validation-failure returns -(`:3960`, `:3969`, `:3986`). A failed fit therefore leaves every one of them -stale-set while the step falls back to uniform. A stale published bucket is -the exact state family behind the just-fixed IMA; today the wrong-key replay -is prevented only by the window check in `_ragged_verify_bucket` plus the -`py_verify_len` clearing — one line of defense where there should be two. -The layout assert's uniform branch (`model_engine.py:4076-4085`) does not -check `agreed_ragged_bucket is None`, so the stale state is invisible even -with `TLLM_DSPARK_ASSERT_LAYOUT=1`. - -Fix: move all four publications below the last failure check, next to the -`py_verify_len` stamping. - -Related (medium): the commit's central rationale — that the old assert was a -tautology — is false against its own parent: at `50675cb181^`, -`_ragged_verify_bucket` already returned `agreed_ragged_bucket` rather than -re-summing, so the old assert already compared the batch walk against the -fit's allgathered value and could fire. The change is a no-op refactor, and -the comment at `model_engine.py:4103-4109` encodes the incorrect history. -Rewrite it to state what the check actually does. - ---- - -## E. Profiler: the default sweep geometry hits a device-side assert, and the promised drain observability does not exist — HIGH (×2) - -From `72dd803141` / `fab6291e2d`, `tests/microbenchmarks/dspark_sps_profiler.py`. - -### E1. Defaults trap: `block_follows_verify_len=False` runs the asserting regime - -`:1036` (default), `:1219-1223` (the commit's own comment), `:1581-1585` -(flag help) · high - -The default sweep geometry is constant-block; the same commit's comment says -that regime "currently asserts (Repeat.cu output_size vs sum(repeats))" -whenever `max(verify_len) < block_size` — which is **every swept cell except -L == max_draft_len** under the default `verify_lens = 1..max_draft_len` -(compact ragged is forced on by `_prepare_environment`, -`TLLM_DSPARK_RAGGED_VERIFY_MODE=COMPACT`). `SweepConfig.validate()` — whose -docstring is literally "catch the geometry mistakes that would otherwise -fail mid-sweep" — does not check it. A user running the documented default -builds an engine for minutes, then dies on a CUDA assert at the first -L < block cell. Verification also checked the timeline: the "currently -asserts" comment was written hours *after* the IMA fix landed and was -re-worded at HEAD, so it is not stale; and no recorded sweep artifact -(`dspark-runs/sps_samples_L*.jsonl`, `sps_real_final.json`) has ever run the -constant-block geometry to completion. - -Fix: either fix the Repeat.cu packing for `max(lens) < block_size` and delete -the comment, or flip the default to `block_follows_verify_len=True` and make -`validate()` refuse the known-asserting combination. - -### E2. The mid-cell drain promises a counter that does not exist, and the drain can latch shut - -`:1344` (comment), `executor/result.py` (`IterationResult.get_results`) · high - -The comment says "the count of empty reads is reported so a cell that lost -its tail this way is visible rather than merely short" — grep of the file at -the commit and at HEAD finds only the comment; the drainer reports only -exception counts. Worse, on the in-process path `IterationResult.get_results` -latches `_done = True` on the first Empty and only `mark_undone()` (new -prompt submission) resets it — one longer-than-timeout gap in stats -production makes **every subsequent `get_stats` call, including the final -post-cell drain, return `[]` instantly**. Cells are silently truncated with -no signal, in a tool whose output feeds the planner. - -Fix: implement the promised empty-read counter and surface it per cell; -`mark_undone()` (or bypass the latch) before the final drain. - ---- - -## F. Interpolation switch: three loose ends — MEDIUM - -From `72dd803141`. - -1. **Truncation eats shelf-close points first** - (`dspark_sps_profiler.py:756`, medium — verified by executing the numeric - example). `compress_to_risers` now encodes a flat shelf as a breakpoint - PAIR (required by the interp consumer), but the `max_breakpoints` - truncation drops the interior point with the smallest incoming jump — - which is by construction always a shelf-CLOSE point (its jump is below - `min_riser_ms`), never a real riser. Dropping it turns the measured shelf - into a ramp: with counts `[16..128]` and a 2.0→3.0 riser, cap=5 bills - M=80 (measured 2.004) at interp 2.5 — the mid-shelf over-bill the - function's own docstring warns about. The pair encoding also roughly - doubles kept points while the default cap stayed 8, so truncation now - fires on sweeps that never used to hit it — silently (no warning, no - `_meta` note), and the only truncation test uses `min_riser_ms=0.0` so - the interplay is untested. Fix: drop by introduced interpolation error - (or treat a shelf pair as atomic), warn when truncation fires, add a test - with `min_riser_ms > 0`. - -2. **No format marker for the semantic reversal** - (`dspark_planner.py:113`, low). A table emitted by the OLD compressor - (which dropped shelf-closing breakpoints *by design*, for the floor - consumer) is silently reinterpreted as a continuous ramp by the new - interp consumer — the "over-trim mirror" the new code names but cannot - detect. Old tables lack `_meta.engine` so they get only the generic - missing-fingerprint warning. Fix: stamp `lookup: "interp"` (or a format - version) into the payload meta; warn-or-refuse tables without it. - -3. **`fixed_overhead_ms` escapes the finiteness validation the commit - message claims** (`dspark_planner.py:148`, low). The check is - `if self.fixed_overhead_ms < 0.0`, which NaN and +inf both pass; a - non-finite overhead silently collapses the planner to maximal trimming. - Fix: `if not math.isfinite(...) or ... < 0.0: raise`. - ---- - -## G. Documentation drift (new feature doc contradicts shipped code) — LOW - -`docs/source/features/dspark-confidence-scheduling.md`, from `50675cb181`: - -- `:195` — asserts the cost lookup is a staircase that "绝不插值", the exact - opposite of the code as of three commits earlier (`72dd803141`), while the - commit message claims the doc "caught up with the code". -- `:261-268` and `:283` vs `:307` — self-contradictory on the window floor: - two places say `min_verify_len` defaults to 0 (with a worked example headed - "min_verify_len=0（默认）"), one place (and the code) says the floor is - fixed at 1. -- `:400` — states in present tense that draft_len "is not in the - attention-DP consistency allgather", documenting as live a hazard the - shipped shape gate already closed (it allgathers and compares draft_len). - -Fix: one editing pass against the code; recast the `:400` diagram as the -motivating failure. - ---- - -## H. Profiler tooling correctness (deployment-measurement paths) — MEDIUM/LOW - -From `fab6291e2d`, `tests/microbenchmarks/dspark_sps_profiler.py`: - -1. `:1778` (low, verified by executing the parser on a real mixed-step log - line) — `samples_from_iter_log` files mixed context+generation steps - under a wrong `(padded_bs, verify_len)` cell with prefill-inflated eager - step times; `num_ctx_requests` is on the same log line and ignored. The - ladder-resolution path converts ambiguous-therefore-dropped into - confidently-mislabelled. Fix: capture `num_ctx_requests` and skip ctx>0 - steps, counting the drops. -2. `:1966` (medium) — the server sweep's load duration - (`max_tokens = max(64, measure_steps)`) is uncoupled from the measurement - window (`settle_s + poll_s`); under defaults the synthetic requests can - finish before the window opens, collecting nothing or ramp-down rows - filed under smaller batch sizes. Fix: derive load length from the window, - or re-post until the drain completes. -3. `:1878` (medium) — the completion-POST worker threads discard responses - and swallow all errors; a wrong `--served-model-name` produces an idle - server, empty cells, and a SystemExit blaming `/metrics` and - `enable_iter_perf_stats`. Fix: check status codes, surface the first - failure in the error message. -4. `:1866` (low) — a pin refused mid-sweep raises out of the whole sweep and - discards every already-collected sample (`_write_samples` runs only after - the sweep returns); with default `verify_lens = 1..max_draft_len` against - a server whose ladder lacks a rung, a refusal is near-certain. Fix: catch - per rung, count skips, write samples incrementally. -5. `:1917` (low) — `warmup_steps` is accepted by the server path but never - applied; the progress line prints `kept - warmup_steps` while all `kept` - rows were appended. Fix: apply it or drop it. -6. `:1788` (low) — `rank_prefix='0:'` silently drops all lines of logs - without an mpirun/srun rank tag (single-process deployments); not exposed - on the CLI; the error blames `print_iter_log`. Fix: fall back to - prefix-less matching when zero lines match; expose `--rank-prefix`. -7. `:1473` (medium) — the moved file still advertises the old module path in - its own `--help` prog string and docstring run command; copy-pasting the - tool's documented invocation fails with ModuleNotFoundError (the commit - updated the `llm_args.py` error message but not the tool itself). - ---- - -## Suggested order of attack - -C → B → D → A → E1/E2 → F → H → G. C and B block anyone producing a usable -table with the new tooling and are one-liners; D's histogram fix is a -one-line condition plus a code motion; A decides whether the calibration -feature produces a meaningful artifact at all and is the largest change -(snapshot-per-block pairing + request-id keyed reads); E1 is a defaults flip -or a kernel fix; the rest are local. - -A recurring theme worth internalizing: most of these violate the branch's -own explicitly documented rules — "every silent decline must be counted" -(A4, B, E2, H), "facts the consumer cannot see must be absent, not -fabricated" (C), "clear stale published state so it cannot key a later step" -(D) — rather than rules imported from outside. The rules are right; the new -code just needs to keep following them. diff --git a/docs/source/features/dspark-confidence-scheduling.md b/docs/source/features/dspark-confidence-scheduling.md index 6200b0d49bbc..0b2e8de2f22c 100644 --- a/docs/source/features/dspark-confidence-scheduling.md +++ b/docs/source/features/dspark-confidence-scheduling.md @@ -139,15 +139,7 @@ iteration t: [hook(t)] -> [prepare] -> [forward(t)] 无论 overlap scheduler 开不开都成立：overlap 改变的是拷贝**何时落地**，而不是它抓的是 哪一次 forward。没落地的拷贝会退化成"验全块"，而不是退化成一个陈旧的猜测。 -因此 `DSparkVerifyPlanner` 把"引擎给的滞后"和"想要的总滞后"分开接收，差值用一个 host 侧 -carry ring 补上： - -```text -carry_steps = confidence_relay_lag_steps - ENGINE_CONFIDENCE_RELAY_LAG -``` - -默认值下差值为 0，ring 完全 inert。未填充的 ring 层初始化成**中性 confidence** 加上一个 -**不可能的 generation**，所以启动后的前 `carry_steps` 步不可能在未初始化内存上做裁剪。 +实现中滞后不可配置：executor loop 恒定给出两步（上图），planner 直接消费该快照。快照行的新鲜度由 `_confidence_stamp` 守卫 —— 逐 buffer 行的 draft-pass 序号，随 confidence 经同一个 `slots` scatter 在 graph 内写入；陈旧或从未打分的行读到中性 confidence，survival 归 1.0，退化为验证整块。 ### slot 索引与 generation 守卫 @@ -159,11 +151,11 @@ slot 还会被回收，这带来两种不同的陈旧风险，由两个互补机 | 风险 | 机制 | |------|------| | device 上那一行还留着上一个占用者的分数 | `_assign_slot` 在把行交给新请求时重置为中性值 | -| carry ring 里存着回收之前的**host 副本**，device 侧的重置够不着 | 每行带一个 `_slot_generation` 标签，与 confidence 一起 stage；和当前 generation 对不上就把该行 survival 强制为 1.0 | +| host 侧快照里还留着回收之前的副本，device 侧的重置够不着 | 每行带一个 `_confidence_stamp`（draft-pass 序号），与 confidence 一起 stage；与当前 pass 对不上就把该行按中性处理，survival 归 1.0 | 这两者是**互补的，不是二选一**。generation 标签是在 stage 时于 host 侧快照的，而 confidence 是在 forward 时于 device 上写的 —— 所以正是那次重置让标签说的是真话：只 bump -不重置，就会把上一个占用者的分数打上新占用者的标签。反过来，只重置也清不掉 ring 里已经 +不重置，就会把上一个占用者的分数打上新占用者的标签。反过来，只重置也清不掉快照里已经 存在的 host 副本。 两种风险都收敛到同一个安全方向：**survival = 1.0，即验证整个 block**，也就是启用该特性 @@ -421,7 +413,8 @@ bucket = round_up_bucket( padded_bs + max_r slack_r ) ```text +- planner 未构建 -----------------> 全块（无 checkpoint） | - +- 成本表未 profile（flat）--------> 全块 + 启动时告警 + +- 成本表 flat（仅当只给了 sts 表时可达；两表 + | 皆缺在构造时即被拒）------------> 全块 + 告警 | "每个 token 都免费" => 裁剪等于盲猜 | _dspark_confidence --+- 快照还没落地 -------------------> 全块 @@ -456,19 +449,17 @@ speculative_config: # 调度器消费的是累积乘积，所以逐位置误差会沿 block 几何级放大。 confidence_sts_path: /path/to/sts.json - # profile 好的 step-cost 曲线。省略 => flat => planner 拒绝裁剪。 + # profile 好的 step-cost 曲线。开启 ragged 时必填（除非提供 confidence_sts_path）： + # 两者皆缺在构造时即报错 —— flat 成本模型下 planner 永远拒绝裁剪，跑完也是空跑。 confidence_sps_table_path: /path/to/sps.json # 被捕获的长度 ladder；调度器只能从这个集合里选。 # 默认：[1, ceil(K/2), K] confidence_verify_len_tiers: [1, 3, 7] - # confidence relay 的总滞后（decode step 数）。 - # executor loop 本身已经提供恰好 2 步，所以默认值不额外增加任何滞后。 - confidence_relay_lag_steps: 2 - - # 逐请求 verify 窗口。需要 enable_confidence_scheduling。 - enable_ragged_verify: false + # 逐请求 verify 窗口。必须与 enable_confidence_scheduling 同开 —— 两个 flag 描述同一个 + # 特性，只开其中一个会在构造时被拒绝；关闭特性 = 两个都设 false。 + enable_ragged_verify: true ``` `sps.json`： @@ -511,7 +502,7 @@ speculative_config: |------|------| | tier ladder 跨 batch size 有损 | shelf 右边缘随 `bs` 移动，一套 ladder 不可能同时坐在所有 batch size 的边缘上。残差要实测 | | 快照是 2 步之前的 | 纯吞吐启发式（I2 保证正确性与它无关），但在剧烈抖动的负载下决策质量会下降 | -| ragged 模式未做端到端验证 | 需要带训练好的 confidence-head 权重的 checkpoint；默认关闭 | +| ragged 模式的端到端数字仅在一个平台上取得 | DeepSeek-V4-Pro-DSpark，8×B300 DEP8，overlap + CUDA graph（GSM8K 96.44，基线 96.21）；仍默认关闭，跨硬件需重测 | | 逐 token 采样参数路径只有非全贪心 batch 才会触发 | `temperature=0` 的精度跑覆盖不到它 | ## 参考 diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py index 62c6f1b69b7f..bfd17faf2f28 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py @@ -226,7 +226,7 @@ def forward( # bound) and the real per-request counts come through # new_tokens_per_seq. next_n = metadata.num_gen_tokens_per_seq - new_tokens_per_seq = getattr(metadata, "gen_new_tokens_per_seq", None) + new_tokens_per_seq = metadata.gen_new_tokens_per_seq # Pass full kv_score (not sliced) with the generation portion of # cu_seq_lens so the kernel reads at the correct absolute offsets. torch.ops.trtllm.compressor_paged_kv_compress( diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 0d5875d753e0..45a24661bddd 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -832,24 +832,10 @@ def __post_init__(self): self.create_buffers_for_indexer(capture_graph=capture_graph) def prepare(self): - # One write-after-read guard over EVERYTHING prepare stages through - # pinned memory -- the base class's kv_lens and block-table copies as - # much as the ragged row maps. Every one of these buffers is rewritten - # here each step and copied device-ward with non_blocking=True; under - # the overlap scheduler the previous step's copies can still be QUEUED - # behind an in-flight forward when this rewrite starts, and the device - # then copies this step's values into last step's consumers. The GPU - # coredump for the resulting fault names - # applyMLARopeAndAssignQKVKernelGeneration (Warp MMU Fault) storing - # through KVBlockArray -- garbage block offsets, which trace back - # through the token-major gather to the BASE block-table staging, not - # only the ragged row set. Hence one coarse guard at the boundary - # rather than per-buffer events: synchronize on the event recorded - # after the previous prepare's copies were enqueued, which by stream - # order means they have executed before any pinned buffer is touched - # again. Steady-state cost is nothing (the event completed long ago); - # it blocks only when the CPU has genuinely outrun the device, which is - # exactly the window that corrupted the copies. + # WAR guard over EVERYTHING prepare() stages through pinned memory: + # under the overlap scheduler last step's non_blocking H2D copies can + # still be queued when this rewrite starts, so wait on the event + # recorded after the previous prepare's copies before touching any. if not torch.cuda.is_current_stream_capturing(): evt = getattr(self, "_prepare_stage_evt", None) if evt is None: @@ -1666,22 +1652,13 @@ def update_spec_dec_param( int(max_total_draft_tokens), getattr(self, "_draft_alloc_cap", 0) or 0) capture_graph = self.is_cuda_graph - # GROW-ONLY ratchet, never an equality check. Under DSpark ragged - # verification max_draft_len varies per step with the planner's - # windows; the old `!=` conditions re-created every expanded buffer - # whenever it moved, while already-captured graphs kept the previous - # allocations' raw addresses baked in. The caching allocator then - # reused the freed memory for the per-step indptr tensors, and every - # replay read indptr ramps where its row->request map and per-row KV - # extents used to live -- the ragged IMA. (Frozen-scene forensics: - # the retired attn_row buffers' addresses held int64 cumsum ramps.) - # Growing above the high-water mark before any graph exists is fine; - # after that the sizes plateau at the config ceiling and the - # addresses never move again. Consumers slice [:rows]/[:width], so a - # wider-than-this-step buffer is just capacity. - ratchet = max(self._draft_sizing_cap, - getattr(self, "_expanded_alloc_draft_cap", -1)) - if ratchet > getattr(self, "_expanded_alloc_draft_cap", -1): + # GROW-ONLY ratchet, never an equality check: captured graphs bake the + # buffer addresses in, and the old `!=` re-created the buffers whenever + # max_draft_len moved per step -- the graph-path ragged IMA (see + # docs/dspark_ragged_ima_handoff.md). Consumers slice [:rows]/[:width]. + prev_cap = getattr(self, "_expanded_alloc_draft_cap", -1) + ratchet = max(self._draft_sizing_cap, prev_cap) + if ratchet > prev_cap: self._expanded_alloc_draft_cap = ratchet self._create_kv_lens_2d_buffer(capture_graph=capture_graph) self.create_expanded_buffers(capture_graph=capture_graph) @@ -2186,12 +2163,9 @@ def refresh_token_major_gen_rows(self) -> None: # Costs a device-to-host read, so it is opt-in. if (os.environ.get("TLLM_DSPARK_ASSERT_ROWS") and not torch.cuda.is_current_stream_capturing()): - # The .max() below is a device-to-host read, and this function runs - # inside the CUDA-graph capture body (on_update_kv_lens). Reading - # during capture raises cudaErrorStreamCaptureInvalidated and takes - # the whole engine down before any real step runs -- which is what - # the first attempt at this check did. Capture is also not where - # the interesting case lives: the eager path is. + # The .max() is a device-to-host read; during capture it would + # raise cudaErrorStreamCaptureInvalidated. The eager path is where + # the interesting case lives anyway. worst = int(self.attn_row_req_idx_cuda[:rows].max()) assert worst < self.num_seqs, ( f"token-major row map is stale: {rows} rows index up to " diff --git a/tensorrt_llm/_torch/models/dspark/heads.py b/tensorrt_llm/_torch/models/dspark/heads.py index 0f843e8ed27c..62c3d0612143 100644 --- a/tensorrt_llm/_torch/models/dspark/heads.py +++ b/tensorrt_llm/_torch/models/dspark/heads.py @@ -254,6 +254,9 @@ def __init__( torch.ones(max(int(block_size), 1), dtype=torch.float32), persistent=False, ) + # CPU mirror of ``sts_temperatures``; rebuilt lazily by + # ``_host_sts_temperatures`` and dropped by ``load_sts_temperatures``. + self._sts_temperatures_host: Optional[torch.Tensor] = None def forward( self, hidden_states: torch.Tensor, prev_embeddings: Optional[torch.Tensor] = None @@ -295,7 +298,7 @@ def apply_sts(self, confidence_logits: torch.Tensor) -> torch.Tensor: def _host_sts_temperatures(self) -> torch.Tensor: """CPU mirror of ``sts_temperatures``, materialized at most once.""" - cached = getattr(self, "_sts_temperatures_host", None) + cached = self._sts_temperatures_host if cached is None: cached = self.sts_temperatures.detach().to("cpu") self._sts_temperatures_host = cached diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index cf86b9114c4a..af56bdd14d6c 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -141,6 +141,9 @@ def __init__(self, config: CUDAGraphRunnerConfig): # by ModelEngine.fit_ragged_verify_lens when it picks the token bucket; # the padding rows are one shared dummy object so they all share it. self.ragged_pad_verify_len = 0 + # Published by ModelEngine when the ragged fit picks a token bucket; + # None is the resting value (no ragged step in flight). + self.agreed_ragged_bucket: Optional[int] = None self.graphs: Dict[KeyType, torch.cuda.CUDAGraph] = {} self.graph_outputs: Dict[KeyType, @@ -303,16 +306,10 @@ def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: from ..speculative.dspark_observability import trims_submitted_tokens if not trims_submitted_tokens(self.spec_config): return None - # Enforce what the docstring promises: None unless EVERY generation - # request carries a window. The fit publishes `agreed_ragged_bucket`, - # but windows can be absent by the time the key is derived (planner - # or peer declined after the fit, roster changed) -- and a ragged key - # over a window-less batch replays a graph whose token-major row maps - # were never staged this step, so the replay consumes the previous - # ragged step's rows: stale request mapping, stale extents, stale - # block table, and the rope KV-append walks off a retired request's - # allocation (the ragged IMA; proven by prepare-tick vs staged-tick - # stamps: tick N prepare, tick N-1 staging, no skip-guard fired). + # Enforce the docstring: None unless EVERY generation request carries a + # window. Windows can vanish between the fit and key derivation (planner + # or peer declined, roster changed); a ragged key over a window-less + # batch replays stale token-major row maps -- the ragged IMA. if any( getattr(request, "py_verify_len", None) is None for request in batch.generation_requests): @@ -324,7 +321,7 @@ def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: # fitted value is computed from allgathered peer stats by rules every # rank runs identically, so it agrees by construction. None means this # step is not ragged, which keeps the key exactly what it was. - return getattr(self, "agreed_ragged_bucket", None) + return self.agreed_ragged_bucket @staticmethod def _local_draft_len(batch: ScheduledRequests) -> int: @@ -406,16 +403,10 @@ def maybe_get_cuda_graph( can_run_cuda_graph = batch.can_run_cuda_graph batch_size = batch.batch_size if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: - # goal doc §4.4. Comparing only "all generation" and the batch size - # is not enough once the token count stops being a function of the - # row count. Under ragged verification a rank that fell back to - # uniform scheduling produces the same batch size but a different - # token total, so this gate used to pass and let one rank replay a - # graph full of collectives while another ran a different width. - # That mismatch is undetectable at replay time: `all_rank_num_tokens` - # is a host list read inside forward, so the captured graph keeps - # using the value it saw at capture. Compare every component of the - # key that a rank can decide for itself. + # goal doc §4.4. Batch size alone is not enough once the token + # count stops tracking the row count, and replay freezes + # all_rank_num_tokens, so a mismatch is invisible at replay time: + # compare every key component a rank can decide for itself. # # `draft_len` is agreed by the DSpark planner's single collective, # but it is also settable rank-locally by the acceptance-rate @@ -533,15 +524,9 @@ def capture(self, max_draft_len = key[1] token_per_request = max_draft_len + 1 # Ragged keys carry the flat token total as a sixth element (see - # _get_graph_key). It is exactly the width the captured graph will be - # replayed at, and it is NOT batch_size * (draft_len + 1): the block is - # still drafted in full, so draft_len stays at the top tier while the - # verified total shrinks. Sizing the static slices from the product - # instead would capture a graph whose input_ids are wider than the - # attention metadata's num_tokens on every trimmed replay -- and the - # widest bucket happens to match the product, which is what warmup - # captures, so it would capture clean and diverge on the first real - # trimmed step. + # _get_graph_key); that is the replay width, NOT bs * (draft_len + 1): + # the block is drafted in full while the verified total shrinks (the + # two only coincide at the widest bucket, which is what warmup sees). num_tokens_for_capture = (key[5] if len(key) > 5 else batch_size * self.max_beam_width * token_per_request) @@ -653,10 +638,8 @@ def will_pad_to(self, padded_bs: int, num_requests: int) -> bool: """Whether a batch of ``num_requests`` can actually reach ``padded_bs``. The ragged fit derives its token-bucket grid from an assumed padded row - count and hands the rounding slack to pad rows. If padding then - declines, those rows never appear and the fitted total lands in no - captured bucket -- seen as a graph key of batch_size 193 against a - ladder holding 192 and 256, costing every rank a replay. + count and hands the rounding slack to pad rows; if padding declines, + the fitted total lands in no captured bucket and every rank runs eager. Mirrors the size guards in :meth:`_get_padded_batch`. It deliberately does not try to predict the resource-dependent bails (KV capacity for diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 28bac660142f..e1d1c404f098 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -832,6 +832,11 @@ def __init__( ) self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) + # Pinned staging buffers for async H2D copies on the ragged path; see + # _pinned_host / _pinned_host_record. + self._pinned_host_cache = {} + self._pinned_host_events = {} + # Create Encoder CUDA graph config and runner. encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( use_cuda_graph=(self._is_encode_only @@ -1875,18 +1880,10 @@ def _get_graphs_to_capture( f"token buckets from tiers {tiers}; draft_len pinned to " f"{top_tier}).") return graphs - # One draft length, not the ladder. The cross product above was - # right while the uniform tier path existed and picked a different - # draft_len per step; that path is gone, and nothing varies - # draft_len at runtime any more: - # - # scheduling off never reaches this branch (guarded above) - # cap-accept runtime_draft_len = planner.max_tier, always - # compact takes the ragged branch, pinned to top_tier - # - # So the lower tiers were captured and never selected -- with - # tiers {1,3,5}, two thirds of these graphs were dead, and captured - # graph memory comes straight out of KV cache. + # One draft length, not the ladder: nothing varies draft_len at + # runtime any more (cap-accept pins runtime_draft_len to + # planner.max_tier; compact takes the ragged branch above), so + # lower tiers would be captured graphs paid for out of KV cache. top_tier = tiers[-1] graphs = [(bs, top_tier) for bs in cuda_graph_batch_sizes] logger.info( @@ -2051,14 +2048,9 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: # in this pass uses the non-greedy key; populate's override # below will keep it False on every subsequent iteration. spec_metadata.is_all_greedy_sample = False - # Measure what the captured graphs actually cost. Ragged keys graphs - # on (batch_size, token_bucket) rather than batch size alone, which - # multiplies the capture set -- 12.1 graphs per batch-size bucket - # against 4.1 uniform, measured. The graph count was easy to count - # from the log; the memory was not, and an OOM at max_batch_size=512 - # reported 13.97 GiB in "private pools (e.g., CUDA Graphs)" without - # attributing it. Report the delta so the cost is a number rather - # than an inference. + # Count graphs and report private-pool memory (in the finally + # block below): ragged keys multiply the capture set, and OOM + # messages report private-pool bytes without attributing them. _graph_mem_before = (torch.cuda.memory_reserved(), torch.cuda.memory_allocated()) _graph_count = 0 @@ -2115,15 +2107,10 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: # the fit -- the planner declining, or a non-ragged mode -- would # otherwise key on a bucket from capture. self.cuda_graph_runner.agreed_ragged_bucket = None - # Absolute private-pool size, not a delta. Deltas are useless - # here: an engine build runs several capture passes over a - # shared allocator, so the per-pass increments came out as - # +13.9 / +5.5 / -0.8 / +0.7 GiB for the same 19 graphs -- the - # later passes reuse pool memory the earlier ones reserved, and - # reserved can even fall. CUDA graphs live in private pools - # (segment_pool_id != (0, 0)), which is the same quantity - # PyTorch's OOM message reports as "allocated in private pools - # (e.g., CUDA Graphs)", so sum those directly. + # Absolute private-pool size, not a delta: capture passes share + # the allocator, so per-pass deltas mislead. CUDA graphs live in + # private pools (segment_pool_id != (0, 0)) -- the same quantity + # PyTorch's OOM message reports -- so sum those directly. _pool_bytes = 0 _pool_ids = set() for _seg in torch.cuda.memory_snapshot(): @@ -2137,16 +2124,10 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: f"across {len(_pool_ids)} pool(s) " f"(ragged={getattr(self.spec_config, 'enable_ragged_verify', False)})") - # The capture loop stamps `self.runtime_draft_len` per graph so each - # captured shape is built at its own draft length. Nothing puts it back, - # which was harmless while there was one draft length to capture: the - # value it was left at was the value it started with. A tier ladder ends - # on its *smallest* tier, so the engine comes out of capture claiming a - # draft length it is not going to use, and the very next thing that runs - # -- the generic token warmup -- builds dummy requests at the full draft - # length while the spec metadata still says otherwise. Acceptance then - # reshapes `[num_gens * max_draft_len]` drafts into - # `[num_gens, smallest_tier]` and raises. + # The capture loop stamps `self.runtime_draft_len` per captured shape + # and nothing puts it back, so the engine would exit capture claiming + # the LAST shape's draft length; restore it after the passes, before + # the generic token warmup builds requests at the full draft length. saved_runtime_draft_len = self.runtime_draft_len # Pass 1: greedy fast-path (dummy requests carry no sampling params, # so is_all_greedy_sample is naturally True). @@ -3701,14 +3682,8 @@ def _set_warmup_ragged_windows(self, batch, verify_bucket: int, for request, tokens in zip(requests, filled.verify_lens.tolist()): request.py_verify_len = int(tokens) - 1 # Publish it to the runner as well: the graph key carries this value - # (`_ragged_verify_bucket` returns the fit's claim so every - # attention-DP rank keys identically), and the layout assert - # cross-checks it against the batch walk. During capture nothing else - # sets it, because the runtime fit never runs here. Before this line existed the capture-time key's - # token axis was None while the batch really held `verify_bucket` - # tokens, and the mismatch surfaced far away, in the MoE's - # shared-vs-routed row count, as `unmatched tensor shape` two graphs - # into the capture pass. + # and the layout assert cross-checks it against the batch walk; the + # runtime fit never runs during capture, so nothing else sets it here. self.cuda_graph_runner.agreed_ragged_bucket = int(verify_bucket) def ragged_verify_token_buckets(self, padded_bs: int) -> List[int]: @@ -3870,18 +3845,9 @@ def fit_ragged_verify_lens( # Decline to go ragged rather than fit against rows that will not exist. # Uniform always has a captured graph, so the fallback is cheap and the # step still runs. - # `will_pad_to` mirrors _get_padded_batch's SIZE guards. It cannot see - # the group's graph-capability answer -- that is what `all_can_graph` - # above now supplies -- and it deliberately does not model the - # resource-dependent bails (KV capacity for the dummy). Being wrong - # optimistically there does not "leave the pre-existing behaviour", as - # an earlier version of this comment claimed: it budgets tokens for a - # row that never appears, so the layout carries n_real*(w+1) tokens - # while the graph key says padded_bs*(w+1). Observed as - # requests carry 18 generation tokens but the key's token axis is 24 - # on 2 of 8 ranks, which threw inside the forward, and under - # attention-DP a rank that throws stops issuing collectives -- the - # other six deadlocked and the run was hard-killed after 300s. + # `will_pad_to` mirrors only the SIZE guards; the group's graph answer + # is `all_can_graph` above, and the resource bails stay unmodeled -- an + # optimistic miss desyncs the layout from the graph key under ADP. if not runner.will_pad_to(padded_bs, len(token_lens)): logger.debug( f"DSpark ragged: padding to {padded_bs} rows is not available " @@ -4071,14 +4037,13 @@ def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, # does not describe. Checked here because the ragged branch below # never runs for uniform steps, which made the staleness invisible # even with this assert enabled. - agreed_stale = getattr(self.cuda_graph_runner, - "agreed_ragged_bucket", None) + agreed_stale = self.cuda_graph_runner.agreed_ragged_bucket assert agreed_stale is None, ( f"uniform step carries a stale published ragged bucket " f"{agreed_stale}; a failed fit exited without clearing it") # Uniform step: assert the converse, that nothing ragged leaked # through. A half-published layout is the dangerous state. - assert getattr(spec_metadata, "total_verify_tokens", None) is None, ( + assert spec_metadata.total_verify_tokens is None, ( "uniform batch left a stale total_verify_tokens on the spec " f"metadata: {spec_metadata.total_verify_tokens}") assert getattr(attn_metadata, "ragged_verify_lens", None) is None, ( @@ -4089,7 +4054,7 @@ def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, from_requests = sum(verify_lens) from_layout = attn_metadata.num_tokens - attn_metadata.num_ctx_tokens from_spec = int(spec_metadata.total_verify_tokens) - agreed = getattr(self.cuda_graph_runner, "agreed_ragged_bucket", None) + agreed = self.cuda_graph_runner.agreed_ragged_bucket assert from_requests == from_layout, ( f"ragged layout mismatch: requests carry {from_requests} generation " @@ -4143,23 +4108,14 @@ def _pinned_host(self, key: str, values, dtype) -> torch.Tensor: PyTorch does not extend the lifetime of an async H2D source -- that is the caller's job -- so the DMA can read a freed page. - It survives most of the time, which is why this went unnoticed: it - needs the allocator to actually hand the page back before the copy - drains. Raising the number of ragged steps per run from 3 to 35 was - enough to turn it from never-seen into a run-killer, surfacing as - `CUDA error: an illegal memory access` on one rank and `unspecified - launch failure` on its peers, hundreds of frames from the copy. - CUDA_LAUNCH_BLOCKING=1 makes the copy synchronous and hides it - completely, which is how it was identified. + Intermittent by nature: the fault needs the allocator to reuse the + page before the queued copy drains (CUDA_LAUNCH_BLOCKING=1 hides it). Buffers are keyed by name and grown monotonically, so a steady state allocates nothing. """ values = list(values) - cache = getattr(self, "_pinned_host_cache", None) - if cache is None: - cache = self._pinned_host_cache = {} - self._pinned_host_events = {} + cache = self._pinned_host_cache # WAR guard, the counterpart of the use-after-free this helper fixed: # a persistent buffer can be REWRITTEN while its previous non_blocking # H2D is still queued behind an in-flight forward, so the device copies @@ -7238,9 +7194,8 @@ def forward(self, self._record_dspark_graph_use( replayed=can_run_graph, shape=None if can_run_graph else - (getattr(self.cuda_graph_runner, "last_miss_reason", None) - or "unknown", - getattr(self.cuda_graph_runner, "last_miss_key", None))) + (self.cuda_graph_runner.last_miss_reason or "unknown", + self.cuda_graph_runner.last_miss_key)) if can_run_graph: attn_metadata = maybe_attn_metadata spec_metadata = maybe_spec_metadata diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ab1fd46046fc..27bc630bd04f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2447,18 +2447,9 @@ def _maybe_arm_dspark_sync_debug(self) -> None: def _maybe_assert_dspark_ragged_active(self, stats) -> None: """Run the ragged gate mid-loop, once the step floor is reached. - The same check also runs at executor-loop cleanup, but only there it is - useless as a gate: cleanup happens after every request has completed and - the driver has already returned its results, so the AssertionError is - raised into a shutdown path that pytest never sees. A run whose ragged - path silently did nothing reported ``1 passed`` with eight ranks each - raising in their logs -- the worst possible outcome for a guard, since - the green result is what gets believed. - - Raised from inside the loop instead, the exception propagates out - through the in-flight requests and fails the run. Only once: the - counters keep moving, and re-raising every step would bury the first - (and only informative) failure. + The cleanup-time copy of this check raises into a shutdown path pytest + never sees; raised mid-loop instead, the exception fails the run. Only + once, so the first (and only informative) failure is not buried. """ # Lazy rather than set in __init__: this class has several # construction paths and a gate that silently no-ops because one of @@ -2466,11 +2457,9 @@ def _maybe_assert_dspark_ragged_active(self, stats) -> None: if getattr(self, "_dspark_asserted", False): return min_steps = self._dspark_assert_min_steps - # Counted over steps with a real batch, not all steps. trtllm-bench - # ramps concurrency from zero, so step 64 can be a batch of two, where - # every window is trivially identical and the gate would fail a run - # whose scheduler had done nothing wrong -- observed as "Broadcasting - # event-loop error to 2 pending request(s)". + # Counted over steps with a real batch, not all steps: trtllm-bench + # ramps concurrency from zero, and a tiny early batch has trivially + # identical windows, which would fail an innocent run. if min_steps is None or stats.steps_multi_request < min_steps: return self._dspark_asserted = True @@ -2507,33 +2496,21 @@ def _report_dspark_ragged_stats(self) -> None: """ worker = getattr(self.model_engine, "_get_spec_worker", None) stats = getattr(worker(), "ragged_stats", None) if worker else None - # Drained BEFORE the ragged early-return. The two are unrelated: STS - # collection runs in static mode too, and gating its final flush on - # `ragged_stats.steps_total > 0` silently discarded the tail of every - # static-mode collection run -- up to flush_every steps x batch rows, - # and the shards still looked well-formed. + # Drained BEFORE the ragged early-return: STS collection runs in static + # mode too, so gating the final flush on ragged counters drops the tail + # of every static-mode collection run. recorder = getattr(self.sampler, "sts_recorder", None) if recorder is not None: recorder.flush() - if stats is None or not getattr(stats, "steps_total", 0): + if stats is None or not stats.steps_total: return stats.log_summary(prefix="DSpark ragged verify [final]") - raw = os.environ.get("TLLM_DSPARK_ASSERT_ACTIVE", "") - if not raw or raw in ("0", "false", "False"): - return - # The value is a MINIMUM STEP COUNT, not a boolean. An engine build - # spins up more than one executor -- KV-cache estimation runs a - # throwaway one that completes a couple of decode steps and shuts down - # -- and asserting on that one kills the run before the real work - # starts, reporting "the feature did nothing" about an executor that was - # never asked to do anything. Requiring a floor is what distinguishes - # "the scheduler declined" from "this executor barely ran". - try: - min_steps = 32 if raw in ("1", "true", "True") else int(raw) - except ValueError: - logger.warning( - f"ignoring TLLM_DSPARK_ASSERT_ACTIVE={raw!r}: expected a " - f"minimum decode-step count, or 1 for the default") + # The value is a MINIMUM STEP COUNT, not a boolean: an engine build + # spins up throwaway executors (KV-cache estimation) that run a couple + # of decode steps, and asserting on one of those reports "the feature + # did nothing" about an executor never asked to do anything. + min_steps = self._dspark_assert_min_steps + if min_steps is None: return if stats.steps_total < min_steps: logger.info( @@ -3380,7 +3357,7 @@ def set_dspark_verify_len_pin(self, verify_len: Optional[int]) -> Optional[int]: "this engine runs static (uniform) verification; the " "verify-length pin only affects window-computing modes, so " "pinning here would be accepted and silently never applied") - if (getattr(worker, "sts_recorder", None) is not None + if (worker.sts_recorder is not None and verify_len is not None and int(verify_len) < int(self.model_engine.max_draft_len)): # The collector's own guard refuses to START under a censoring @@ -3431,7 +3408,7 @@ def _dspark_confidence_draft_len(self, # the environment override exists so a run can be switched between # `static` and `compact` without editing the serving config, and every # counter below is keyed off the same decision. - stats = getattr(worker, "ragged_stats", None) + stats = worker.ragged_stats # The sampler is built independently of the worker, so nothing connects # them; do it here, where both are reachable. Acceptance has to be # recorded in the sampler (that is where the accepted count and the @@ -3447,33 +3424,23 @@ def _dspark_confidence_draft_len(self, # keys must come from the allocator that wrote the buffer. if (self.sampler is not None and getattr(self.sampler, "sts_recorder", None) is None): - recorder = getattr(worker, "sts_recorder", None) + recorder = worker.sts_recorder if recorder is not None: self.sampler.sts_recorder = recorder self.sampler.sts_row_for = worker.sts_row_for self.sampler.sts_seq_provider = worker.verified_draft_seq - mode = getattr(worker, "ragged_verify_mode", None) - compute_windows = (mode.computes_windows if mode is not None else - getattr(self.model_engine.spec_config, - "enable_ragged_verify", False)) + mode = worker.ragged_verify_mode + compute_windows = mode.computes_windows is_distributed = (self.dist is not None - and getattr(self.dist, "tp_size", 1) > 1) + and self.dist.tp_size > 1) # --- Phase 1: decide locally, issue nothing --------------------------- # - # goal doc H9. Every reason to give up on a ragged split is rank-local, - # and one of them -- whether the confidence snapshot's copy event has - # landed (`_ready_snapshot`) -- is timing-dependent, so it can differ - # between ranks on the same step for no reason the ranks can see. A - # collective placed after those checks is therefore entered by some - # ranks and skipped by others: the rank that went ragged used to issue - # two (the planner's max, then `peer_stats`), a rank that declined - # issued one, and a rank that declined only at the bucket fit issued - # three. That is a deadlock, and it is a timing-dependent one. - # - # So: nothing below this point may branch on rank-local state before the - # single collective, and nothing after it may issue another one. + # goal doc H9. Every decline reason is rank-local and one + # (_ready_snapshot) is timing-dependent, so a collective placed after + # them deadlocks: nothing below may branch on rank-local state before + # the single collective, and nothing after it may issue another one. ragged_lens = None fallback_reason = None if compute_windows: @@ -3485,14 +3452,9 @@ def _dspark_confidence_draft_len(self, fallback_reason = "planner_declined" else: fallback_reason = "mode_static" - # Nothing to compute: the fallback verifies the full drafted - # block. This used to call decide_draft_len, which picked one tier - # for the whole batch -- a middle ground between "schedule per - # request" and "verify everything". It is gone: it could not act on - # per-request confidence, which is the point of scheduling, and it - # degenerated to the full block whenever acceptance was high, which - # was every workload measured here. The config that reached it is - # rejected in llm_args now, so there are two states, not three. + # Nothing to compute: the fallback verifies the full drafted block. + # (The batch-wide uniform-tier middle ground was removed; llm_args + # now rejects the config that reached it, so two states, not three.) # --- Phase 2: one collective, same shape on every rank, every step ---- # @@ -3614,17 +3576,10 @@ def _dspark_confidence_draft_len(self, if bucket is None: fallback_reason = "no_captured_shape" else: - # No rank can replay a graph this step -- such steps were - # eager long before this feature existed (can_run_cuda_graph - # is num_context_requests == 0), so declining ragged would buy - # no replay and only forfeit the trimming (measured: 46 of the - # first 64 steps, trim_ratio 0.62 -> 0). Run the windows - # eagerly, and do not fit a bucket: fitting sizes pad rows - # `_get_padded_batch` will not append on an ungraphable step. - # The fault once blamed on this path was a ragged graph KEY - # derived over a window-less step replaying stale row maps; - # with the key derivation checking windows, this path - # publishes no bucket and replays nothing. + # No rank can replay a graph this step (context requests + # present); such steps were always eager, so run the windows + # eagerly. Do not fit a bucket: fitting sizes pad rows that + # _get_padded_batch will not append on an ungraphable step. runner = self.model_engine.cuda_graph_runner runner.agreed_ragged_bucket = None runner.ragged_pad_verify_len = 0 @@ -3667,14 +3622,10 @@ def _dspark_confidence_draft_len(self, stats.record_step(num_gen_requests=num_gen_requests, verify_lens=fitted, bucket=bucket, - # Only when the fit delivered: bucket is non-None - # exactly then. Reading the attribute on every - # step booked the LAST fit's padded row count - # into padded_bs_hist on eager-ragged and - # cap-accept steps -- phantom entries on exactly - # the steps the histogram exists to explain, and - # clearing at fit entry cannot help the branches - # that never enter the fit. + # Only when the fit delivered (bucket non-None): + # reading it every step books the LAST fit's row + # count into padded_bs_hist on eager-ragged and + # cap-accept steps, which never enter the fit. padded_bs=(getattr(self.model_engine, "_dspark_last_padded_bs", None) if bucket is not None else None), @@ -3724,12 +3675,8 @@ def _dspark_confidence_draft_len(self, # The stamp buffer says which draft pass last wrote each slot; the # sequence bump marks "everything drafted so far" so the planner # can histogram how stale each gathered row was at decision time. - staged_seq = None - bump = getattr(worker, "bump_draft_seq", None) - if bump is not None: - staged_seq = bump() - stamps_fn = getattr(worker, "confidence_stamp_buffer", None) - stamps = stamps_fn() if stamps_fn is not None else None + staged_seq = worker.bump_draft_seq() + stamps = worker.confidence_stamp_buffer() planner.stage_confidence(buffer, stamps=stamps, staged_seq=staged_seq) # The STS recorder keeps its own ring of these snapshots: its @@ -3737,7 +3684,7 @@ def _dspark_confidence_draft_len(self, # pass that DRAFTED the block being labeled, not with whatever the # live buffer holds by then. Same relay, same stamps, staged in # the same breath as the planner's copy. - recorder = getattr(worker, "sts_recorder", None) + recorder = worker.sts_recorder if recorder is not None: recorder.stage_snapshot(device_logits=buffer, device_stamps=stamps, diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 75d7738c2b05..5c1f11aec890 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -490,9 +490,9 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> f"every mid-shelf total upward. Re-emit the table with a " f"current profiler if trimming behaves oddly.") check_table_fingerprint(payload=raw, live={ - "tp": int(getattr(self.mapping, "tp_size", 0) or 0), - "ep": int(getattr(self.mapping, "moe_ep_size", 0) or 0), - "attention_dp": bool(getattr(self.mapping, "enable_attention_dp", False)), + "tp": int(self.mapping.tp_size), + "ep": int(self.mapping.moe_ep_size), + "attention_dp": bool(self.mapping.enable_attention_dp), "block": int(block_size), "max_batch_size": int(max_batch), # moe_backend deliberately NOT in the live dict yet: the @@ -583,8 +583,8 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> from .dspark_observability import (DSparkRaggedStats, RaggedVerifyMode, read_ragged_verify_mode) - configured = (RaggedVerifyMode.COMPACT if getattr( - self.spec_config, "enable_ragged_verify", False) else + configured = (RaggedVerifyMode.COMPACT + if self.spec_config.enable_ragged_verify else RaggedVerifyMode.STATIC) self.ragged_verify_mode = read_ragged_verify_mode(default=configured) # STS calibration data collection. Off unless the env var is set; when @@ -595,10 +595,10 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> # next line uses it the same way. self.sts_recorder = make_recorder_from_env( block_size=int(block_size), - rank=int(getattr(getattr(self, "mapping", None), "rank", 0) or 0), + rank=int(self.mapping.rank), has_cost_table=bool( self.spec_config.confidence_sps_table_path), - ragged_mode=getattr(self.ragged_verify_mode, "value", None)) + ragged_mode=self.ragged_verify_mode.value) self.ragged_stats = DSparkRaggedStats(mode=self.ragged_verify_mode, max_draft_len=block_size) @@ -724,7 +724,7 @@ def _draft_gen_block_batched( # Note the ``or K``: SpecMetadata.runtime_draft_len defaults to 0, not # None, so an ``is None`` check would leave a stride of 1 on any path # that has not populated it yet. - runtime_draft_len = getattr(spec_metadata, "runtime_draft_len", 0) or K + runtime_draft_len = spec_metadata.runtime_draft_len or K Kp1 = int(runtime_draft_len) + 1 device = accepted_tokens.device @@ -754,8 +754,8 @@ def _draft_gen_block_batched( # tokens, so the offsets come from the same qo_indptr the input layout # was packed with. Striding by Kp1 there would have request r read from # request r-1's tail -- no error, just a silent acceptance collapse. - qo_indptr = getattr(spec_metadata, "qo_indptr", None) - if getattr(spec_metadata, "verify_lens", None) is not None and qo_indptr is not None: + qo_indptr = spec_metadata.qo_indptr + if spec_metadata.verify_lens is not None and qo_indptr is not None: base = gen_start + qo_indptr[:num_gens].to(device=device, dtype=torch.long) else: arange_g = torch.arange(num_gens, device=device) @@ -985,7 +985,7 @@ def _forward_impl( # cap-accept only: what each request's window threw away, written by # `apply_accept_caps` above. Rides the sampler's existing copy rather # than a read of its own, so the measurement costs no synchronization. - cap_trim = getattr(spec_metadata, "accept_cap_trim", None) + cap_trim = spec_metadata.accept_cap_trim if cap_trim is not None and spec_metadata.accept_caps is not None: outputs["cap_trim_lens"] = cap_trim[:batch_size] return outputs diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index 56a33175ef77..06a9e9077b4d 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -380,12 +380,8 @@ def decide_verify_lens( input layout is built per request (so it goes ragged) while the spec metadata sees a ``None`` and stays uniform. """ - # Counted before any decline, so the four fallback counters below have a - # denominator. Without it they are absolute counts against an unknown - # total, which is how a run reported `planner_declined: 2` alongside an - # all-zero planner block: nothing said how many decisions there had - # been, so nothing said whether zero fallbacks meant "healthy" or - # "never asked". + # Counted before any decline, so the fallback counters below have a + # denominator rather than being absolute counts against an unknown total. self.stats["decisions"] += 1 if num_gen_requests <= 0: @@ -396,18 +392,10 @@ def decide_verify_lens( return None if self._forced_verify_len is not None: - # Ahead of the cost-table gate on purpose: the profiler runs without - # a table (that is the artifact it is producing), so a pin placed - # after this gate would never fire and the run would look entirely - # normal while measuring the untrimmed block. - # - # It sets ``lens`` rather than returning, so the pinned path still - # runs the cross-rank agreement and the length assert below. An - # early return skipped both -- harmless while the pin equals the - # block (every rank agrees trivially and the packed buffers are - # full-width anyway), which is why arm B ran clean, but it left the - # short-window case taking a different route through the same code - # than any real decision ever does. + # Ahead of the cost-table gate on purpose: the profiler runs + # without a table, so a pin placed after it would never fire. Sets + # `lens` rather than returning, so the pinned path still runs the + # cross-rank agreement and the length assert below. self.stats["forced_steps"] += 1 lens = [int(self._forced_verify_len)] * int(num_gen_requests) else: diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 1a7c1e550b96..3ddbd5abf1bb 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1087,7 +1087,7 @@ def apply_accept_caps(num_accepted_tokens: torch.Tensor, num_contexts: int, return caps = caps[:gen.shape[0]].to(gen.dtype) capped = torch.minimum(gen, caps) - trim = getattr(spec_metadata, "accept_cap_trim", None) + trim = spec_metadata.accept_cap_trim if trim is not None: # Zero the context slots too. The buffer is persistent and the batch # composition moves every step, so a slot left unwritten would carry @@ -1785,14 +1785,13 @@ def _rejection_buffers_valid(self, draft_tokens, draft_len, stored_vocab, # rows one-hot, and clamps the accepted count back to `verify_lens`) # unreachable. Greedy GSM8K cannot see any of this. logits_rows = logits.shape[0] if logits.dim() > 1 else 1 - if getattr(spec_metadata, "is_ragged_verify", False): + 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 = getattr(spec_metadata, "total_verify_tokens", - None) + total_verify_tokens = spec_metadata.total_verify_tokens if (total_verify_tokens is None - or getattr(spec_metadata, "verify_lens", None) is None - or getattr(spec_metadata, "qo_indptr", None) 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: diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py index 2981693ba222..45b55e394567 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py @@ -894,7 +894,7 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): # empirically: reverting the stride fix leaves all twelve cases green. What this # does cover is the RoPE, KV-write, sparse-MLA and compressor paths at widths # other than one, which was previously nothing at all. -@pytest.mark.parametrize("generation_seq_len_q", [1, 2, 4, 6]) +@pytest.mark.parametrize("generation_seq_len_q", [1, 2, 6]) def test_deepseek_v4_sparse_mla( context_lengths: List[int], num_generation_steps: int, generation_seq_len_q: int ): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py index cee639cb2adf..5be568d1257f 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py @@ -70,15 +70,6 @@ def test_no_caps_is_a_no_op(): assert counts.tolist() == [3, 1, 4] -def test_counts_inside_their_windows_are_untouched(): - """The draft died before the cut, so the cut cost nothing.""" - counts = torch.tensor([1, 3, 2], dtype=torch.int32) - meta = _Meta(caps=[6, 4, 6]) - apply_accept_caps(counts, 0, meta) - assert counts.tolist() == [1, 3, 2] - assert meta.accept_cap_trim[:3].tolist() == [0, 0, 0] - - def test_counts_beyond_their_windows_are_clamped_and_the_loss_accumulated(): counts = torch.tensor([6, 5, 2], dtype=torch.int32) meta = _Meta(caps=[3, 5, 6]) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py index a4071fdee9ac..9787ba617aae 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py @@ -256,18 +256,6 @@ def test_allocation_is_a_prefix_and_respects_bounds(): assert int((lens - cfg.min_verify_len).sum()) <= 20 -def test_allocation_prefers_the_most_confident_requests(): - cfg = _cfg(min_verify_len=1) - conf = torch.tensor( - [ - [1.0, 0.99, 0.99, 0.99, 0.99, 0.99, 0.99], # confident - [1.0, 0.10, 0.10, 0.10, 0.10, 0.10, 0.10], # not - ] - ) - lens = schedule_verify_lens_topk(survival=compute_survival(conf), budget=6, cfg=cfg) - assert lens[0] > lens[1] - - def test_zero_budget_gives_everyone_the_floor(): cfg = _cfg(min_verify_len=2) surv = compute_survival(torch.rand(5, BLOCK)) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py index a586c44a8c9b..e403bbacc816 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py @@ -192,13 +192,6 @@ def test_ragged_batch_keeps_each_request_positions_contiguous(): assert got["pos_indices"] == [0, 0, 0, 0, 1, 2, 2] -def test_flat_token_total_matches_the_sum_of_windows(): - reqs = [_Req(3), _Req(0), _Req(1)] - got = _build_overlap_layout(reqs, [0, 0, 0], 6) - assert len(got["position_ids"]) == sum(got["tokens_each"]) - assert len(got["pos_indices"]) == sum(got["tokens_each"]) - - # -------------------------------------------------------------------------- # kv_len_offsets: the host estimate and the correction must use one count # -------------------------------------------------------------------------- @@ -247,45 +240,3 @@ def test_batch_wide_correction_under_ragged_is_off_by_the_window_difference(): deltas = [w - r for w, r in zip(wrong, right)] assert deltas == [4 - n, 1 - n, 2 - n] assert all(d != 0 for d in deltas) - - -def test_uniform_batch_is_unaffected_by_the_pairing_rule(): - reqs = [_Req(), _Req(), _Req()] - past = [100, 250, 7] - accepted = [2, 1, 2] - n = 4 - - per_request = _corrected_kv_lens(reqs, past, accepted, n, [n] * 3) - assert per_request == [p + a for p, a in zip(past, accepted)] - - -# -------------------------------------------------------------------------- -# Derived totals the callers compute from the per-request counts -# -------------------------------------------------------------------------- - - -@pytest.mark.parametrize( - "verify_lens", - [ - [None, None, None], - [3, 3, 3], - [3, 0, 1], - [0, 0, 0], - ], -) -def test_draft_token_total_is_tokens_minus_one_per_request(verify_lens): - # previous_batch_draft_tokens = total_num_tokens - num_requests, which has - # to agree with summing (tokens - 1) request by request. - reqs = [_Req(v) for v in verify_lens] - tokens = [get_request_tokens_per_gen_step(r, 4) for r in reqs] - total_num_tokens = sum(tokens) - assert total_num_tokens - len(reqs) == sum(t - 1 for t in tokens) - - -def test_uniform_totals_match_the_multiplicative_formulas(): - # The formulas the uniform path used before this change. - n, num_requests = 4, 5 - reqs = [_Req() for _ in range(num_requests)] - tokens = [get_request_tokens_per_gen_step(r, n) for r in reqs] - assert sum(tokens) == num_requests * n - assert sum(tokens) - num_requests == num_requests * (n - 1) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py index bc0c6476655f..ab7846123d5d 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py @@ -157,16 +157,6 @@ def test_fill_bucket_adds_pad_rows_to_reach_the_captured_batch_size(): assert min(lens[2:]) >= 1, "a pad row with 0 tokens breaks qo_indptr slicing" -def test_fill_bucket_prefers_real_rows_over_pad_rows(): - lay = _layout([1, 1], bucket=10) - filled = lay.fill_bucket(max_verify_len=BLOCK, padded_bs=4) - lens = filled.verify_lens.tolist() - assert sum(lens) == 10 - # 8 spare beyond the baseline (1+1 real + 1+1 pad): real rows take theirs - # first, so both real rows outrank both pad rows. - assert min(lens[:2]) >= max(lens[2:]) - - 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 @@ -521,13 +511,6 @@ def _fit(real_token_lens, bucket, padded_bs, max_verify_len): return layout.fill_bucket(max_verify_len=max_verify_len, padded_bs=padded_bs) -def test_fit_hits_the_bucket_exactly(): - # 3 real requests wanting 2+3+1=6 tokens, padded to 4 rows, bucket 12. - got = _fit([2, 3, 1], bucket=12, padded_bs=4, max_verify_len=6) - assert int(got.verify_lens.sum()) == 12 - assert got.bs == 4 - - def test_fit_spends_slack_on_real_requests_before_pad_rows(): got = _fit([2, 2], bucket=8, padded_bs=4, max_verify_len=6) lens = got.verify_lens.tolist() @@ -537,13 +520,6 @@ def test_fit_spends_slack_on_real_requests_before_pad_rows(): assert lens[0] + lens[1] == 6 -def test_fit_pad_rows_are_never_empty(): - # An empty range breaks qo_indptr's per-row slicing. Baseline here is - # 2*4 real + 6 pad rows = 14, so the bucket has to be at least that. - got = _fit([4, 4], bucket=16, padded_bs=8, max_verify_len=4) - assert all(v >= 1 for v in got.verify_lens.tolist()) - - def test_fit_rejects_a_bucket_that_cannot_hold_the_pad_rows(): # The same batch against a bucket that leaves no room for the pad rows: # 8 real tokens + 6 pad rows needs 14, not 10. Raising here is the point -- @@ -552,12 +528,6 @@ def test_fit_rejects_a_bucket_that_cannot_hold_the_pad_rows(): _fit([4, 4], bucket=10, padded_bs=8, max_verify_len=4) -def test_fit_marks_which_rows_are_real(): - got = _fit([2, 3], bucket=10, padded_bs=4, max_verify_len=6) - assert got.num_real_requests == 2 - assert got.num_pad_requests == 2 - - 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 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py index 4979f20e6df9..e44f3809ea88 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py @@ -263,6 +263,7 @@ def test_fit_recovers_relative_structure_exactly(): every cell must be reproduced to the last bit.""" cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) fit = fit_additive_cost_model(cells) + assert fit.warnings == () for cell in cells: predicted = fit.intercept_ms[cell.batch_size] + fit.theta_ms[cell.total_verify_tokens] @@ -300,11 +301,6 @@ def noise(batch_size, verify_len): assert "bs=32, L=3" in report -def test_fit_stays_quiet_when_the_model_holds(): - cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) - assert fit_additive_cost_model(cells).warnings == () - - def test_fit_clamps_a_dip_upward(): """A non-monotone theta would offer the planner a cheaper *longer* verify.""" dipped = _staircase([(0, 1.0), (64, 5.0), (128, 3.0), (256, 8.0)]) @@ -589,13 +585,14 @@ def test_token_budget_is_scaled_when_acceptance_is_not_pinned(): def test_validate_rejects_a_sequence_that_would_be_evicted(): """An evicted request drains the batch mid-cell and shrinks the shape.""" with pytest.raises(SweepGeometryError, match="max-seq-len"): - _sweep(input_len=4064).validate() + _sweep(input_len=4064, block_follows_verify_len=True).validate() def test_validate_rejects_a_decode_step_wider_than_max_num_tokens(): """A split decode step never has the shape its cell is filed under.""" with pytest.raises(SweepGeometryError, match="widest decode step"): - _sweep(batch_sizes=[8, 4096], max_num_tokens=1024).validate() + _sweep(batch_sizes=[8, 4096], max_num_tokens=1024, max_seq_len=8192, + block_follows_verify_len=True).validate() def test_validate_accepts_a_sane_sweep(): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 2bb9275c4ec5..609dfa3910d1 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -98,6 +98,7 @@ def _make_worker(enable_confidence_scheduling=False, verify_len_tiers=None): max_draft_len=5, spec_dec_mode=SpeculativeDecodingMode.DSPARK, enable_confidence_scheduling=enable_confidence_scheduling, + enable_ragged_verify=False, verify_len_tiers=verify_len_tiers or [1, 3, 5], confidence_sps_table_path=None, confidence_sts_path=None, @@ -488,7 +489,7 @@ def forward_batched(main_hidden, bonus, start_pos, **kwargs): return dm, seen -@pytest.mark.parametrize("runtime_draft_len", [5, 3, 1]) +@pytest.mark.parametrize("runtime_draft_len", [3, 1]) def test_gen_draft_gathers_hidden_with_the_runtime_stride(runtime_draft_len): """The gather stride must follow runtime_draft_len, not max_draft_len. @@ -638,11 +639,3 @@ def test_confidence_buffer_absent_when_scheduling_is_off(): worker._lazy_init(dm, _make_metadata(max_num_requests=8)) assert worker._confidence_logits is None assert worker.verify_planner is None - - -def test_unscored_confidence_slots_read_as_verify_all(): - """A slot never scored must not make the scheduler trim on garbage.""" - worker = _make_worker(enable_confidence_scheduling=True) - dm, _ = _stride_probe_draft_model() - worker._lazy_init(dm, _make_metadata(max_num_requests=8)) - assert torch.all(torch.sigmoid(worker._confidence_logits) > 0.999) diff --git a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py index e7908b91c002..a5084293c3d9 100644 --- a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py +++ b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py @@ -154,6 +154,7 @@ def test_graph_key_uses_the_agreed_bucket_not_a_fresh_sum(): class _R: spec_config = type("S", (), {"enable_ragged_verify": True})() + agreed_ragged_bucket = None def _batch(*verify_lens): return types.SimpleNamespace(generation_requests=[ diff --git a/tests/unittest/_torch/speculative/test_dspark_planner_trims.py b/tests/unittest/_torch/speculative/test_dspark_planner_trims.py index 7b7a8a64ab35..bcb12eecb37c 100644 --- a/tests/unittest/_torch/speculative/test_dspark_planner_trims.py +++ b/tests/unittest/_torch/speculative/test_dspark_planner_trims.py @@ -23,7 +23,6 @@ code path ships untested. """ -import pytest import torch from tensorrt_llm._torch.speculative.dspark_schedule import ( @@ -62,20 +61,6 @@ def test_topk_hands_longer_windows_to_more_confident_requests(): assert sum(lens) - len(lens) * 1 <= 6 -def test_a_budget_below_the_full_block_produces_differing_windows(): - """Ragged shape arises from the planner, not from an override.""" - survival = torch.tensor([ - [0.99, 0.98, 0.97, 0.96, 0.95], - [0.60, 0.30, 0.10, 0.02, 0.00], - [0.95, 0.90, 0.20, 0.05, 0.00], - [0.10, 0.01, 0.00, 0.00, 0.00], - ]) - lens = schedule_verify_lens_topk(survival=survival, budget=5, - cfg=_Cfg()).tolist() - assert len(set(lens)) >= 2, ( - f"expected a ragged split from a constrained budget, got {lens}") - - def test_a_full_budget_degenerates_to_the_uniform_full_window(): """The no-trim case must stay exactly uniform. @@ -88,12 +73,3 @@ def test_a_full_budget_degenerates_to_the_uniform_full_window(): lens = schedule_verify_lens_topk(survival=survival, budget=budget, cfg=_Cfg()).tolist() assert lens == [max_len] * num_reqs - - -@pytest.mark.parametrize("budget", [0, 1, 3, 8, 20]) -def test_windows_stay_within_bounds_for_any_budget(budget): - survival = torch.rand(6, 5).sort(dim=1, descending=True).values - lens = schedule_verify_lens_topk(survival=survival, budget=budget, - cfg=_Cfg()).tolist() - assert all(1 <= v <= 5 for v in lens), lens - assert sum(lens) - len(lens) <= budget diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py index a71b9f1e6640..f398fa1de3c9 100644 --- a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py +++ b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py @@ -91,6 +91,7 @@ def _valid_state(): return types.SimpleNamespace( draft_probs=torch.empty((R, K, V), device="cuda"), batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), + is_ragged_verify=False, ) @@ -209,6 +210,7 @@ def _dispatch_meta(**over): base = dict( use_rejection_sampling=True, is_all_greedy_sample=False, + is_ragged_verify=False, draft_probs_vocab_size=V, draft_probs_last_dim=V, batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), diff --git a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py index bddba0454744..4b26ffffe0bb 100644 --- a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py +++ b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py @@ -14,9 +14,6 @@ bypassing it. See tests/unittest/_torch/speculative/test_dspark_planner_trims.py. """ -import os -from unittest.mock import patch - import pytest from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig @@ -31,19 +28,16 @@ def _cfg(**kwargs): return DSparkDecodingConfig(**base) -@patch.dict(os.environ, {}, clear=False) def test_ragged_without_a_cost_table_is_rejected(): with pytest.raises(ValueError, match="requires a profiled cost table"): _cfg() -@patch.dict(os.environ, {}, clear=False) def test_ragged_with_an_sps_table_is_accepted(): cfg = _cfg(confidence_sps_table_path="/tmp/table.json") assert cfg.enable_ragged_verify -@patch.dict(os.environ, {}, clear=False) def test_confidence_scheduling_without_ragged_is_rejected(): """There is no uniform tier ladder to fall back to. @@ -62,14 +56,3 @@ def test_confidence_scheduling_without_ragged_is_rejected(): speculative_model="/nonexistent/model", enable_confidence_scheduling=True, enable_ragged_verify=False) - - -@patch.dict(os.environ, {}, clear=False) -def test_scheduling_off_needs_no_table_and_verifies_everything(): - """The fallback is static verification, which has nothing to profile.""" - cfg = DSparkDecodingConfig(max_draft_len=5, - speculative_model="/nonexistent/model", - enable_confidence_scheduling=False, - enable_ragged_verify=False) - assert not cfg.enable_confidence_scheduling - assert not cfg.enable_ragged_verify From 3afeb5ecbf5bb1187b5f5bdf2a8e55801bd1b115 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 5 Aug 2026 06:27:45 -0700 Subject: [PATCH 094/177] [None][chore] Deep-slim the DSpark tests; drop the process docs Tests: the suite re-proved the same facts many times over. Every cut was adversarially reviewed against the production modules -- a proposal survived only if no bug exists that the deleted test alone would catch: - near-duplicate bodies collapsed into @pytest.mark.parametrize tables; - facet tests sharing one fixture merged into single tests whose merged bodies pin STRONGER exact values than the originals did; - tests that exercised test-local reimplementations of production logic (helpers like _build_overlap_layout / _corrected_kv_lens) deleted outright: they stay green when production diverges, guarding nothing; - tautologies and trivial-accessor checks removed. Rejected by review and kept: the dp_size/attention-DP rank accounting guard, the non-default freqs_cis parametrization, the STS ring-pairing regression tests, and the stamp-lag histogram accumulation cases. Docs: the four process documents (design goal, task prompt, P0 report, IMA handoff) were working artifacts of this branch's development, not upstream material; deleted, with the four code comments that pointed at them made self-contained. The feature documentation remains at docs/source/features/dspark-confidence-scheduling.md. Validation: python3 -m py_compile on every touched file; every DSpark-related unit test file (plus the KV block-offset overlap-race regressions) run end to end on B300: 331 passed, 0 failed. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_confidence_schedule_goal.md | 399 -------- docs/dspark_p0_report.md | 577 ----------- docs/dspark_p0_task_prompt.md | 933 ------------------ docs/dspark_ragged_ima_handoff.md | 452 --------- .../_torch/attention_backend/sparse/dsa.py | 4 +- .../_torch/pyexecutor/sampler/sampler.py | 2 +- .../dsv4_sparse_mla_presentation.py | 8 +- .../hw_agnostic/test_dspark_attention.py | 49 +- .../hw_agnostic/test_dspark_cap_accept.py | 96 +- .../test_dspark_confidence_schedule.py | 182 +--- .../test_dspark_cost_interpolation.py | 40 +- .../hw_agnostic/test_dspark_cuda_graph.py | 71 +- .../hw_agnostic/test_dspark_draft.py | 11 +- .../hw_agnostic/test_dspark_heads.py | 38 +- .../test_dspark_iter_log_samples.py | 92 +- .../hw_agnostic/test_dspark_overlap_layout.py | 194 +--- .../hw_agnostic/test_dspark_padded_ceiling.py | 32 +- .../hw_agnostic/test_dspark_ragged.py | 361 +++---- .../test_dspark_reconcile_guards.py | 37 +- .../hw_agnostic/test_dspark_sps_profiler.py | 133 +-- .../test_dspark_sts_collection_guards.py | 33 +- .../test_dspark_tier_aligned_budget.py | 52 +- .../hw_agnostic/test_dspark_verify_len_pin.py | 29 +- .../hw_agnostic/test_dspark_worker.py | 144 +-- .../test_dspark_graph_miss_attribution.py | 8 +- .../speculative/test_force_accepted_tokens.py | 97 +- .../speculative/test_group_all_greedy_sync.py | 37 +- .../test_rejection_buffers_guard.py | 53 +- .../test_dspark_ragged_requires_table.py | 8 +- 29 files changed, 533 insertions(+), 3639 deletions(-) delete mode 100644 docs/dspark_confidence_schedule_goal.md delete mode 100644 docs/dspark_p0_report.md delete mode 100644 docs/dspark_p0_task_prompt.md delete mode 100644 docs/dspark_ragged_ima_handoff.md diff --git a/docs/dspark_confidence_schedule_goal.md b/docs/dspark_confidence_schedule_goal.md deleted file mode 100644 index 735a5263afab..000000000000 --- a/docs/dspark_confidence_schedule_goal.md +++ /dev/null @@ -1,399 +0,0 @@ -# DSpark Confidence-Head 调度在 DeepSeek-V4-Pro 上的落地：目标文档 - -> 状态：scoping / 目标定义，**尚未定稿实现方案** -> 关联：PR NVIDIA/TensorRT-LLM#17056，分支 `pr-17056` -> 目标模型：DeepSeek-V4-Pro-DSpark（target = DSv4 Pro + DSA 稀疏注意力；drafter = DSpark block drafter + confidence head） -> `max_draft_len=5` → `block_size K=5` → 均匀 verify 窗口 `next_n = 6` -> 硬件：8×Blackwell（B200/B300，SM100），DEP8 = TP8 + `moe_expert_parallel_size=8` + `enable_attention_dp=True` - -本文档中所有断言均标注了 `file:line`。凡是**没有**通过读代码确认的推断，一律显式标注「待验证」。 - ---- - -## 1. 背景与现状 - -### 1.1 PR #17056 已经做了什么 - -两个 commit： - -| commit | 内容 | -|---|---| -| `830afe4e58` | `[feat] Add DSpark confidence-head verification scheduling` | -| `8723506608` | `[fix] Fix host syncs, slot aliasing and cost model` | - -核心数据流（已读码确认）： - -1. **Confidence head 产出**：`DSparkConfidenceHead.forward` 返回 `[G, K]` fp32 raw logits，按 batch 位置排列（`tensorrt_llm/_torch/models/dspark/heads.py:258-268`），在 draft 的 captured graph 内计算（`tensorrt_llm/_torch/models/dspark/draft.py:108-121`）。 -2. **batch → slot 重索引**：`DSparkWorker` 立刻把它写进按 slot 索引的常驻 device buffer `self._confidence_logits[slots] = confidence`（`tensorrt_llm/_torch/speculative/dspark.py:637-638`），buffer 形状 `[max_batch+2, K]`，其中 `max_batch+1` 行是永不写入的 neutral 行（`dspark.py:285-300`）。 -3. **滞后一步的 host 读取**：`PyExecutor._dspark_confidence_draft_len` 先读上一步 stage 好的 async D2H 快照，再 stage 新的一份，**全程不 sync**（`tensorrt_llm/_torch/pyexecutor/py_executor.py:3189`；staging 见 `tensorrt_llm/_torch/speculative/dspark_verify.py:99-134`）。event 没落地就退化成 `max_tier`（verify 全量）。 -4. **语义**：`py_verify_len` 计的是**每请求的 draft 位置数**，token 窗口是 `1 + py_verify_len`（`tensorrt_llm/_torch/pyexecutor/llm_request.py:1283-1286`）。 - -两个配置开关（`tensorrt_llm/llmapi/llm_args.py:2805` / `:2849`，默认均为 `False`，`status='prototype'`）： - -| 开关 | 行为 | -|---|---| -| `enable_confidence_scheduling` | **均匀路径**：`planner.decide_draft_len` 给整个 batch 选一个 tier，把每个请求的 `py_verify_len` 显式清成 `None`（`py_executor.py:3283`），复用既有的 dynamic-draft-len pad/truncate 机制。CUDA graph 变成 `batch_size × tiers` 的叉积（`model_engine.py:1828-1856`）。 | -| `enable_ragged_verify` | **ragged 路径**：`decide_verify_lens` 给出每请求窗口，`fit_ragged_verify_lens` 把 batch token 总数吸附到某个 captured bucket，`runtime_draft_len` 钉死在**最高 tier**（因为 block 永远整块 draft），graph key 增加一条 token-count 轴。依赖 `enable_confidence_scheduling`（validator 见 `llm_args.py:2863-2871`）。 | - -**默认路径（两个都 False）完全不受影响**，这一点已确认：`verify_len_tiers` 退化成 `[max_draft_len]`（`llm_args.py:2899-2905`），`py_verify_len` 从不设置，`expand_per_gen_token` 走标量 `repeat_interleave` 分支（`dsa.py:693-697`），与 PR 之前的定长扩展 bit-identical。 - -### 1.2 已经做对的部分 - -PR 的 host 侧管线质量是高的，不要推翻重做： - -- `ragged_verify_lens_cuda` / `ragged_qo_indptr_cuda` 一次性预分配、原地写入（`model_engine.py:715-720`、`:3802-3813`），地址稳定，captured graph 能看到新值。 -- 每个 `repeat_interleave` 都传了 `output_size=`，避免 device→host sync（`dspark_ragged.py:411-429`、`dsa.py:693-706`）。 -- staging 用 pinned memory + `non_blocking=True`。 -- `choose_ragged_capture_shape` 对 `peer_stats` 做跨 rank 归约，保证各 rank 落在同一个 `(padded_bs, bucket)`（`dspark_ragged.py:359-382`）。 -- raggedness 活在**内容**里、shape 由 graph key 固定 —— 这是对的构造。 - -### 1.3 工作区当前状态（**与 investigation 描述不同，必须先对齐**） - -`git status` 显示 **11 个已跟踪文件被修改**、2 个新增文件，均**未提交**。这部分工作已经超出了 PR HEAD： - -``` - M cpp/tensorrt_llm/kernels/IndexerTopK.h - M cpp/tensorrt_llm/kernels/indexerTopK.cu - M cpp/tensorrt_llm/thop/IndexerTopKOp.cpp - M cpp/tensorrt_llm/kernels/compressorKernels/compressorKernels.{cu,h} - M cpp/tensorrt_llm/thop/compressorOp.cpp - M tensorrt_llm/_torch/attention_backend/sparse/dsa.py - M tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/{compressor.py,deepseek_v4.py} - M tensorrt_llm/_torch/pyexecutor/py_executor.py - M tensorrt_llm/_torch/speculative/dspark.py -?? tensorrt_llm/_torch/speculative/dspark_observability.py -?? tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py -``` - -已完成的内容（读码确认）： - -| 项 | 状态 | 锚点 | -|---|---|---| -| `indexer_topk_decode` 新增 `Tensor? row_kv_lens=None` | **已实现** | `IndexerTopKOp.cpp:258`（schema）、`indexerTopK.cu:665-690`（kernel 分支）、`IndexerTopKOp.cpp:57-66`（ragged 时**替换**而非放宽 TORCH_CHECK）、`indexerTopK.cu:918-922`（与 GVR `preIdx` 互斥的 TLLM_CHECK） | -| Python 侧路由 ragged → C++ top-k，`next_n` 传死值 1 | **已实现** | `dsa.py:3056-3081` | -| `compressor_paged_kv_compress` 新增 `Tensor? new_tokens_per_seq=None` | **已实现** | `compressorOp.cpp:153`（schema）、`compressorKernels.cu:340-348`（`nn` 取代 `NEXT_N`，`NEXT_N` 降级为编译期上界）、`:425`（Phase 2）；Phase 1 由 `token_idx < kv_len` 守卫（`:369`）、Phase 3 由 `num_compressions` 守卫（`:440-442`），二者对 `nn < NEXT_N` 均正确 | -| `num_gen_tokens_per_seq` ragged 时取 **max** 而非 mean | **已实现** | `deepseek_v4.py:748-769` | -| `gen_new_tokens_per_seq` 每请求向量 + graph-stable buffer | **已实现** | `deepseek_v4.py:337-352` | -| ragged verify mode 开关（static / cap-accept / compact）+ 统计 | **已全部实现**（`cap-accept` 见 `spec_sampler_base._apply_verify_cap` + `py_verify_cap`） | `dspark.py:457-496`、`dspark_observability.py` | -| C++ ragged top-k 单测 | **已新增** | `tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py` | - -**重要副作用**：`num_gen_tokens_per_seq` 改成取 max 之后，compressor 输出 buffer 的 sizing 公式 `num_generations * ceil(w/cr)`（`deepseek_v4.py:753-757`、`:766-768`）重新成为**合法上界**（因为 `w_r ≤ w_max` ⇒ `Σ_r ceil(w_r/cr) ≤ n·ceil(w_max/cr)`）。这消除了「越界写 `kv_comp`」这一类内存破坏风险。取 mean 时该公式**不是**上界，是真实的 OOB 写。 - -### 1.4 已知的**未验证**部分 - -**唯一一次端到端尝试是失败的。** `.dspark-logs/gsm8k.log`（07/30 23:08，早于当前工作区改动）： - -``` -DSpark confidence scheduling: capturing 102 graphs (34 batch sizes x 3 draft-length tiers [1, 3, 5]) -Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=5 ← 完成 -Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=3 ← never returns -TEST_EXIT=124 -``` - -`hang-stack-299685.txt` 停在 `thop.attention`（`fmha/fallback.py:181`）→ `trtllm.py:1775` → `deepseek_v4.py:1564` → `mla.py:1132` → `cuda_graph_runner.py:499 capture`。 - -注意 hang 前紧邻的日志是 `Attention workspace size is not enough, increase the size from 0 bytes to 168493568 bytes` —— **capture 期间的 workspace 重分配**是一个独立的候选根因，必须与 stride 假说区分开（见 §7）。 - -同时，**这次 run 其实什么都没验证到**：配置里 `enable_ragged_verify` 未开，且没给 `confidence_sps_table_path`，planner 拿到 `SpsCostTable.flat()` 后 `_decide_local` 无条件返回 `max_tier`（`dspark_verify.py:268-270`），日志确认 `profiled_cost_table=False`。也就是说该 run 的行为与 baseline 的唯一差别是：多一次 confidence matmul、多一次 D2H staging、以及 3× 的 captured graph。 - ---- - -## 2. 目标 / 非目标 - -### 2.1 目标（可测） - -| # | 目标 | 验收方式 | -|---|---|---| -| **G1** | DSv4-Pro 上 confidence-head 调度可用：`enable_confidence_scheduling=True` 能完成 warmup + capture 且不 hang | 8×B200 上 GSM8K 测试跑到 evaluation 阶段并返回分数 | -| **G2** | 端到端正确性：调度**不改变输出分布** | temperature=0 下，`static` 与开启调度后逐 token 输出**完全一致**（比 GSM8K 分数敏感度高 3 个数量级） | -| **G3** | GSM8K 精度达标 | `score >= 96.0`（`tests/.../references/gsm8k.yaml:158-163` 记录的实测值为 96.475） | -| **G4** | **CUDA graph 兼容** | 稳态 decode 步骤 100% 命中 graph replay；capture 期间无 device→host sync、无新分配；graph key 与实际 token layout 逐字段一致（断言化，见 §6） | -| **G5** | ADP/TP8 下各 rank 形状与集合通信一致 | 各 rank 的 `(padded_bs, bucket)` 与 collectives 调用次数逐步相同；无 hang | -| **G6** | 默认路径零回归 | 两个开关关闭时与 `main` bit-identical（现已成立，需保持） | - -### 2.2 隐含目标（工作本身要求，但用户未明说） - -- **G7**：为「feature 到底有没有在工作」提供可观测性。当前 flat cost table 下 planner 拒绝 trim，run 看起来完全正常但收益为零 —— 必须能从日志/统计区分「调度生效」与「调度退化」。`dspark_observability.py` 是这个方向，需要落地。 -- **G8**：产出两个**目前在 checkpoint 和 repo 里都不存在**的 JSON side-file：STS 校准表、profiled SPS cost table。没有它们，G2/G3 之外的任何**性能**结论都是空的。 - -### 2.3 非目标（明确排除） - -| 排除项 | 理由 | -|---|---| -| 改变 acceptance 规则本身 | 调度只决定「送多少 draft token 给 target」，verify 逻辑不动。这是 G2 可以用逐 token 相等来验证的前提。 | -| Disaggregated serving | 生成侧首步 `py_batch_idx is None` 会命中一条已知有问题的分支（§3.3 K1），单独处理。 | -| 其他模型（DSv4-Flash、非 DSA 模型） | DSv4-Flash-DSpark 权重虽在 `llm_models_root()` 下，但 repo 里零测试；作为后续的小规模 proxy，不在本轮范围。 | -| Pipeline parallelism / LoRA | 与 vLLM #47808 声明的限制一致，本轮不支持。 | -| `enable_heuristic_topk` / `use_cute_dsl_topk` 与 ragged 共存 | GVR 的 `preIdx` 按请求索引且经 `next_n` 展开，CuTe-DSL 把 `next_n` 作为 JIT compile key。二者与 ragged 结构性互斥，**必须显式 config-reject 而非静默降级**。工作区已在 kernel 层加了 TLLM_CHECK（`indexerTopK.cu:918-922`），Python 侧的 config 校验还欠。 | -| 非贪婪采样下的 rejection sampling 等价性 | ragged 下 fail-closed guard 会关掉 rejection sampling（§3.4）。本轮先保证贪婪路径正确，非贪婪单列。 | -| 输出 logprobs | 与 vLLM 同样的限制。 | - ---- - -## 3. 核心问题：kernel 层的 per-request verify len 不兼容 - -### 3.1 问题的本质 - -一旦每个请求有自己的 verify 窗口，**所有从标量 `next_n` 反推 `(request, offset)` 的 kernel 都会失效**。典型形式是： - -``` -request = rowIdx / next_n -offset = rowIdx % next_n -causal = seqLens[request] - next_n + offset + 1 -``` - -而 `next_n` 的来源是 `num_gen_tokens // num_generations`（`dsa.py:2800`、`deepseek_v4.py:748-769`）—— ragged 下这是一个**整数平均值**，没有任何语义。 - -**最危险的一点：失败是静默的，不是报错的。** `ModelEngine.ragged_verify_token_buckets` 返回 `padded_bs * (t+1)`（`model_engine.py:3625`），token 总数**永远**是行数的整数倍，所以形如 `seq_lens.size(0) * next_n == numRows` 的 TORCH_CHECK 会**通过**，然后 kernel 把行归错请求、按错误深度做因果裁剪。这正是 vLLM 在 #47808 之前踩过的坑（他们在 `mla/indexer.py` 里留了明确注释说 `q.view(batch, next_n, ...)` 对非均匀 decode 长度会静默算错）。 - -DSpark 论文 §5.3 对 DeepSeek-V4 的结论与此一致：变长路由「在物理执行层引入了严峻挑战」，解法是把所有 token 展平、用一个 marker tensor 承载序列内结构，并且**「在 DeepSeek-V4 架构上，只有 index-attention 和 compress kernel 需要修改」**。这句话既是我们工作范围的上界，也是「主 sparse-MLA 路径不用动」的外部佐证。 - -### 3.2 DSv4-Pro DSpark 生成路径 kernel 逐项审计 - -按前向顺序。「ragged 容忍」列：✅ 天然支持 / ⚠️ 需 host 侧改动 / ❌ kernel 层阻断 / 🔧 工作区已修。 - -| # | Kernel / Op | Query 轴形状契约 | Ragged? | 阻断点与锚点 | -|---|---|---|---|---| -| **K1** | `mla_rope_generation` → `applyMLARopeAndAssignQKVKernelGeneration` | `TLLM_CHECK(acc_q_len % batch_size == 0)` 后 `seq_len = acc_q_len / batch_size`，**单一标量** | ❌ | `mlaKernels.cu:1161-1163`。派生出 `batch_idx = tok/seq_len`、RoPE `position_id = kv_len[b] - seq_len + local`、**KV 写入偏移**（两处）、`seqQOffset`。窗口 `[6,3]` → `9%2≠0` 直接 abort；`[6,2]` → 8 tokens、`seq_len=4`、check 通过但**每个请求的 K/V 写进错误的行**。这是持久性缓存污染。 | -| **K2** | trtllm-gen sparse MLA generation（`thop.attention` → `AttentionOp::mlaGeneration`） | `mMaxSeqLenQ = acc_q_len / batch_beam`；`mSumOfSeqLensQ = batch_beam * mMaxSeqLenQ` | ❌ | `attentionOp.cpp:1161,1167`。`cumSeqLensQPtr` 字段存在（`fmhaRunnerParams.h:279`）但**在 generation 分支从不赋值**（全 repo 仅 `fmhaDispatcher.cpp:213` 上下文相、`xqaDispatcher.cpp:475` 赋值）。上游 `thop/attentionOp.cpp:920-922` 还有 `TLLM_CHECK(num_tokens % num_seqs == 0, "seq_len should be same for all generation requests")`。**这是最深的阻断点。** | -| **K3** | `fp8_paged_mqa_logits` / `fp8_fp4_paged_mqa_logits`（DeepGEMM） | expanded 路径下 `q.view(-1,1,...)`、`context_lens = kv_lens_expanded[:n].view(-1,1)`、`block_table_expanded` ⇒ `next_n == 1`，**每 token 一行** | ✅ | `dsa.py:2831-2838`。形状上**已经是 ragged 布局**。关键事实：`max_draft_len=5 → next_n=6`，DeepGEMM 在 SM100 只支持 `next_n ∈ {1,2,4}`，所以 **DSv4-Pro 今天本来就跑在这条展平路径上** —— ragged 不是新 regime，它就是 regime。但注意：该 kernel 输出**未裁剪的全长 logits**（`clean_logits` 默认 false），窗口内因果裁剪**全部委托给 top-k**。 | -| **K4** | `indexer_topk_decode`（C++） | 原：`seqLens[rowIdx/next_n]`、`seq_len - next_n + rowIdx%next_n + 1` | 🔧 **工作区已修** | 新增 `Tensor? row_kv_lens`（`IndexerTopKOp.cpp:258`），kernel 分支见 `indexerTopK.cu:665-690`。ragged 时**替换**而非放宽 TORCH_CHECK（`IndexerTopKOp.cpp:57-66`）—— 这个选择是对的，因为放宽后仍会被 `padded_bs*(t+1)` 的整除性静默骗过。split-work 的 merge pass 也正确转发了指针（`indexerTopK.cu:1073-1076`，merge 阶段 `rowEnd` 被覆盖但仍需转发以避免 uniform 分支越界索引 `seqLens`）。 | -| **K5** | GVR / heuristic top-k（C++ `heuristicTopKDecode.cu` + CuTe-DSL `gvr_topk_decode.py`） | `preIdx + (rowIdx/next_n)*stride`；DSL 里 `next_n` 是 `const_expr` 且是 JIT cache key | ❌ **须 reject** | `heuristicTopKDecode.cu:57-65,103`；`IndexerTopKOp.cpp:83-84` 断言 `preIdx.size(0)*next_n == numRows`；`gvr_topk_decode.py:3733-3734,4954`。`order_row` 是**按请求**的 argsort（`dsa.py:1010-1017`），假定每请求占固定连续行数。默认关闭，但是调优部署会打开的旋钮。工作区已在 kernel 层加互斥 check。 | -| **K6** | CuTe-DSL radix / filtered top-k | `next_n` 是 `cutlass.const_expr` + compile key；每行长度 `seqlen[row//next_n] - next_n + row%next_n + 1` | ❌ **已绕开** | `single_pass_multi_cta_radix_topk.py:1062,1124-1125`。DSv4 默认不可达（`dsa.py:867` 在 `compress_ratio>1 且 next_n>1` 时早退，DSv4 的 indexer cr=4）。工作区把 ragged 无条件路由到 C++ kernel（`dsa.py:3056-3081`），进一步保证不可达。 | -| **K7** | CuTe-DSL paged-MQA atom-split `_pick_dsl_expand` | `expand_factor * atom == next_n`，`atom` 是 kernel constexpr | ❌ **结构性不可表达，已 force-off** | `dsa.py:211-263`、ragged 时 `use_dsl=False`（`dsa.py:1566-1574`）。**DSv4-Pro 上性能代价为零**：`use_cute_dsl_paged_mqa_logits` 默认 False（`llm_args.py:888-892`），默认本来就走 DeepGEMM + expanded。 | -| **K8** | `Indexer.prepare_scheduler_metadata`（DeepGEMM schedule） | 硬编码 `num_tokens = num_generations * (1 + max_draft_tokens)` | ⚠️ | `dsa.py:2329-2330`。另两处展开点已改走 `expand_per_gen_token`，唯独这里没改。**这也是均匀 tier 阶梯的 bug**（见 §3.3 H1）。 | -| **K9** | `compressor_paged_kv_compress` | 原：`NEXT_N` 是 C++ **模板常量**，`sp = kv_len - NEXT_N` | 🔧 **工作区已修** | 新增 `Tensor? new_tokens_per_seq`（`compressorOp.cpp:153`）。`nn` 取代 `NEXT_N` 做起点与 Phase-2 计数（`compressorKernels.cu:348,425`），`NEXT_N` 降级为编译期**上界**以保住 uniform 路径的全展开。Phase 1 由 `token_idx < kv_len` 守卫、Phase 3 由 `num_compressions` 守卫，二者对 `nn < NEXT_N` 均正确。`in_off`/`out_off` 本来就是精确前缀和，无需改。 | -| **K10** | DSv4 compressed-token 计数 / `kv_comp` 分配 | `num_generations * ceil(num_gen_tokens_per_seq / cr)` | 🔧 **工作区已修** | `deepseek_v4.py:748-769`。改取 **max** 后该式重新成为合法上界，消除 OOB 写。`_compute_gen_compressed_position_ids` 用 `searchsorted(cu_new_comp, ...)` 回推请求（`deepseek_v4.py:982-990`），对任何**过分配**都仍然正确。 | -| **K11** | `sparse_mla_topk_lens` / `swa_local_indices` / `compressed_local_indices` | 严格逐 token，从 `token_positions` 构造 | ✅ | `deepseek_v4.py:557-648`。**且是 overlap-safe 的**：`on_update_kv_lens` 会用修正后的 device `kv_lens` **重新推导** `token_positions` 并重跑 `prepare_for_deepseek_v4_indices`（`deepseek_v4.py:885-895`）。这条是「token-major 下丢掉 `mIsCausalSpecDecodingGen` 是安全的」的**唯一承重论据**，它是可验证的，不是假设。 | -| **K12** | DSA indexer K-cache slot mapping、`req_idx_per_token`、`convert_req_index_to_global`、`mla_rope_append_paged_kv_assign_q` | 逐 token / indptr 驱动 | ✅ | `dsa.py:1757-1769`、`:2148-2178`、`:3287-3335` | -| **K13** | `spec_decoding_generation_lengths` / packed mask（dense XQA 路径） | C++ 侧**真支持**每请求向量（`gptKernels.cu:140`），但**所有 Python 写入点都是标量 `fill_`** | ⚠️ 但 DSv4 不可达 | `trtllm.py:1189-1191`。SM100 + linear tree 下 `is_spec_decoding_enabled` 被强制 False（`trtllm.py:999-1001`），整块分配跳过；且 MLA 路径走 `mlaGeneration`，**从不读**该字段。**建议加显式 assert 而不是依赖 SM 判断。** | -| **K14** | MoE / dense MLP / allreduce / ADP `num_tokens` allgather | token-flat | ✅ | `attn_metadata.num_tokens` 就是 `_seq_lens.sum().item()`（`attention_backend/interface.py:206-215`），是真实打包计数。`fused_moe/` 与 `distributed/` 全树 grep 无 `next_n` / draft len 耦合。 | -| **K15** | DSpark drafter 自身的 forward | `[G, block]`，与 target 布局解耦 | ✅ | draft 的 attention 用 worker 自有的 `_kv_windows` 滚动 buffer，不碰 paged KV metadata。 | - -### 3.3 Host 侧阻断点（与 kernel 同等重要，且**独立于方案选择**） - -这些不修，任何 kernel 方案都不可能产出正确结果。 - -| # | 问题 | 锚点 | 严重性 | -|---|---|---|---| -| **H0** | **`disable_overlap_scheduler=True` 时，ragged token 布局根本不会被产出。** `_prepare_tp_inputs` 的非 overlap 分支条件是 `next_draft_tokens_device is None or request.is_dummy or request.py_batch_idx is None`（`model_engine.py:4606`），该分支 `sequence_lengths.append(1 + num_draft_tokens)`，**完全不看 `py_verify_len`**（`:4616,4625`）。而 `next_draft_tokens_device` 仅当 `new_tensors_device` 是 `SampleStateTensorsSpec` 时非空（`model_engine.py:4270-4277`），非 overlap 循环调用 `self._forward_step(scheduled_batch)` 不传该参数（`py_executor.py:4276`，默认 `None`，`py_executor.py:6425`）。**精度测试用的就是 `disable_overlap_scheduler=True`**（`test_llm_api_pytorch.py:4036`）。 | 上述 | **P0 / prerequisite #0**。这不是 investigation 描述的「padding 行 / 首步边缘情况」，而是在目标配置下**每个请求每一步**都命中。后果：`spec_metadata.verify_lens` / `qo_indptr` / `total_verify_tokens` 和 graph key 都在描述一个 ragged 布局，而真实 `input_ids`/`seq_lens` 是均匀的 `1+top_tier`；`_accept_draft_tokens` 随后按 `total_verify_tokens` 切片（`interface.py:1436-1441`），**逐请求 token 错位**。 | -| **H1** | **DSA 展开 stride 用的是静态 max，不是运行时 tier —— 均匀 tier 阶梯本身就是坏的。** `expand_per_gen_token` 非 ragged 分支用 `stride = 1 + self.max_draft_tokens`（`dsa.py:696`），`gen_token_repeat_list` 同（`dsa.py:662`）。`self.max_draft_tokens` 只在 `update_spec_dec_param` 赋值（`dsa.py:1401`），而 DSpark 属于 `is_parallel_draft()`（`interface.py:320-324`，`is_dspark()` 在内），故取 `original_max_total_draft_tokens = tokens_per_gen_step - 1 = 5`（`model_engine.py:342`）—— **静态值**。于是 tier 3 时 `kv_lens_expanded` / `block_table_expanded` 仍按 stride 6 铺，而 `Indexer.forward` 按 `q_decode.shape[0] = bs*4` 切片（`dsa.py:2831-2838`），请求 r>0 拿到 r-1 的 kv_len 和 block table。 | 上述 + `dsa.py:2329-2330` | **P0**。这是**所有三个方案共同需要的修复**，约 3 行 Python（`stride = spec_metadata.runtime_tokens_per_gen_step`），且**独立于 ragged 与否**。它对 07/30 那次 tier-3 capture hang 是一个完整自洽的解释：tier 5 时 768==768 通过，tier 3 时 schedule 描述 768 行而 logits 只有 512 行。**注意：`max_draft_tokens` 同时是 buffer 容量（`dsa.py:1310,1321,1342,1355`），修复时容量必须保持静态 max，只让 stride 变，否则会在 capture 期间触发重分配 —— 那是第二类 hang。** | -| **H2** | **`attn_metadata.ragged_verify_lens` 在 `attn_metadata.prepare()` **之后**才发布。** `prepare()` 在 `model_engine.py:5326`，`_attach_ragged_verify_layout`（唯一写入者）在 `:5412`，同一函数内相隔 86 行。 | 上述 | **P0**。后果：`prepare_for_spec_decode` 里的 `use_expanded_buffers_for_mtp` 判定、`kv_lens_expanded_*` 与 `block_table_expanded` 的铺设**全部用上一步的窗口**。因为 `maybe_get_cuda_graph` 返回的是 per-key 的 graph 常驻 metadata 对象（`cuda_graph_runner.py:387-407`），这个「上一步」其实是「**上一次该 key 跑的时候**」。修复：把 attach 移到 `prepare()` 之前（它只需要 `generation_requests`，已在作用域内）。 | -| **H3** | **工作区新增的 `_prepare_ragged_row_kv_lens` 有 device 不匹配，首个 ragged step 必抛异常。** `row_kv_lens.add_(expanded.to(torch.int32))`（`dsa.py:1695`）：`row_kv_lens` 是 `row_kv_lens_cuda` 的切片（CUDA），而 `expanded` 来自 `expand_per_gen_token(kv_lens[...])`，`kv_lens = cached_token_lens + seq_lens_kv` 两项**都在 CPU 上**（`dsa.py:790-801`；`cached_token_lens` 显式 `device='cpu'`）。`.to(torch.int32)` 不改变 device。同函数 docstring 写「kv_len 留在 device 上」，与自身输入不符。 | `dsa.py:1673-1695` | **P0（新引入）**。修法也更省：把逐行修正量直接折进 `row_kv_lens_host`，做**一次** H2D 拷贝 —— 这正是相邻的 `kv_lens_expanded_host` 已经在用的模式（`dsa.py:1590-1592`）。 | -| **H4** | **`row_kv_lens_cuda` 构建时机错 + 从不被 `on_update_kv_lens` 刷新。** (a) 它在 `prepare_for_spec_decode` 内构建（`dsa.py:1608`），受 H2 影响读到陈旧窗口；(b) overlap 场景下 `_preprocess_inputs` 会在 device 上把 `previous_kv_lens_offsets_cuda` 加进 `kv_lens_cuda` 再调 `on_update_kv_lens`（`model_engine.py:2921-2937`），该函数重建了 `kv_lens_expanded_cuda`、`scheduler_metadata_buffer_expanded`、slot mapping —— **但没有 `row_kv_lens_cuda`**。于是同一次 op 调用里，`gen_kv_lens_cuda`（已修正）与 `row_kv_lens`（未修正）对同一个量的说法不一致。 | `dsa.py:1608`、`dsa.py:902-977` | **P0（新引入）** | -| **H5** | **CUDA graph capture 的 token 轴来自 `key[1]`，不看 bucket。** `num_tokens_for_capture = batch_size * max_beam_width * (key[1] + 1)`（`cuda_graph_runner.py:446-453`），而 ragged 下 `key[1]` 被钉死在 top tier（`model_engine.py:1840-1842`）。 | 上述 | **P0**。结合 H0，同一 bs 下的三个 graph 都是按 `bs*(top_tier+1)` 的布局 capture 的 —— **key 里的 bucket 是虚构的**。H0 与 H5 必须**同时**修，否则 graph 只会更糟。 | -| **H6** | **overlap 下 KV rewind 读的是被下一步覆盖过的 `py_verify_len`。** `SpecSamplerBase` 已经把 batch-wide 的 `runtime_draft_len` 快照进 `SampleStateSpec`（`spec_sampler_base.py:218`）正是为了规避这个 hazard，但 `_verified_len` 却**实时**读请求上的 `py_verify_len`（`spec_sampler_base.py:195-196`）。而第 N 步的 `_handle_dynamic_draft_len`（`py_executor.py:4652`）先于第 N-1 步的 `_update_requests`（`py_executor.py:4746`）执行。 | 上述 | **P1**。修法一致：把 `py_verify_len` 同样快照进 `SampleStateSpec`。**方向修正**：investigation 声称的「泄漏 3 个 KV entry」对 KVCacheManager v1 不成立 —— `update_resources` 之后还有补偿项 `extra_rewind = _kv_reserve_draft_tokens - (py_rewind_len + accepted)`（`resource_manager.py:1051-1063`），`py_rewind_len` 在合法区间内代数抵消；真实误差只在 clamp 处出现，方向是**过度 rewind**（丢已接受的 KV）。但在 `kv_cache_manager_v2` 里补偿被 `if self.is_draft:` 门控（`kv_cache_manager_v2.py:3538-3546`），**target manager 无补偿**，那里 investigation 的场景是真的。**须先确认 DSv4-Pro 实际走哪个 manager**「待验证」。 | -| **H7** | **ragged batch 上 rejection sampling 被 fail-closed guard 静默关掉。** `_rejection_buffers_valid` 要求 `logits_rows >= num_contexts + num_gens*(draft_len+1)`（`interface.py:1707-1710`），而 ragged 下 target 只发 `num_contexts + total_verify_tokens` 行，**当调度真的在 trim 时该式必然不成立**。于是永远退回 `_sample_and_accept_draft_tokens_base`（严格 token 相等）。 | 上述 | **P1**。副作用：PR 在 rejection 路径里精心加的 ragged 分支（`interface.py:1793-1806`）成了死代码。GSM8K temperature=0 **检测不到**这一点。 | -| **H8** | **`fit_ragged_verify_lens` 假定 CUDA graph padding 已开。** 它只检查 `runner.enabled`（`model_engine.py:3644-3648`）就按 `padded_bs` 行数规划并预留 `n_pad * pad_len` 个 token，而 `_get_padded_batch` 在 `not self.padding_enabled` 时直接返回 0（`cuda_graph_runner.py:576-579`），且 `BaseCudaGraphConfig.enable_padding` **默认 False**（`llm_args.py:180-184`）。 | 上述 | **P1**。两种后果：(a) padding 关闭时每个 ragged step 静默掉出 graph replay；(b) **混合 ctx+gen batch** 无论如何都无法跑 graph，真实请求却被按「要给 pad 行留位置」缩小了窗口 —— 系统性的吞吐损失，且 `runner.ragged_pad_verify_len` 会作为陈旧状态残留。须在 config 层校验 `enable_padding=True`。 | -| **H9** | **rank-local 早退门控了条件性的 `tp_allgather`，是 hang 风险。** `decide_verify_lens` 的 `all_rank_max` 在早退之后（`dspark_verify.py:234-241,256-259`），`peer_stats` allgather 被 `if ragged_lens is not None` 门控（`py_executor.py:3258-3261`），fallback 分支又调 `decide_draft_len` 再发一次。早退条件是 rank-local 且**时序相关**（`_copy_event.query()`，`dspark_verify.py:132-133`）。 | 上述 | **P1**。走 ragged 的 rank 发 2 次集合通信，fallback 的 rank 发 1 或 3 次 → 死锁。修法：让每个 rank **无条件**发出相同序列的集合通信，再基于归约结果分支。对比之下**均匀路径天然没有这个问题**：`decide_draft_len` 的 `all_rank_max` 无条件执行（`dspark_verify.py:194-201`）。 | -| **H10** | **DSpark drafter 用固定 stride 取 target hidden state。** `Kp1 = runtime_draft_len + 1`；`base = gen_start + arange_g * Kp1`（`dspark.py:600,624`）。ragged 下 `runtime_draft_len` 被钉在 top tier，而真实偏移是 `spec_metadata.qo_indptr[r]`。`dspark.py` 全文无 `qo_indptr` 引用。 | 上述 | **P1**。静默：只是接受率崩塌，不报错（`interim_base` 还被 clamp 到 `captured.shape[0]-1`）。注意 `8723506608` 已把它从 `max_draft_len+1` 修成 `runtime_draft_len+1`，**均匀路径是对的**，只有 ragged 错。 | - -### 3.4 一条重要的语义澄清 - -`indexer_topk_decode` 是**整条 indexer 链路上唯一施加窗口内因果裁剪的地方**。因为 K3 的 paged-MQA-logits GEMM 在 expanded 布局下 `next_n==1`，每行的 `q_pos = ctx_len - 1`，且一个请求的所有行共享同一个 `ctx_len`，所以它输出**全长未裁剪**的 logits。这意味着 `row_kv_lens` 算错**不是噪声，是让一个 verify token 看到它本不该看到的 KV**（这些 KV 确实已经写进去了，因为整个窗口在 attention 之前就已 append）—— 属于信息泄漏级别的正确性错误，不是精度抖动。 - ---- - -## 4. CUDA graph 兼容性要求 - -### 4.1 不变量（replay 期间必须恒定的东西） - -| 类别 | 内容 | -|---|---| -| **可以原地改内容** | 所有 capture 时地址固定的预分配 buffer：`shared_static_tensors['input_ids'/'position_ids']`（`cuda_graph_runner.py:172-184`）、per-key 常驻的 `attn_metadata` 及其 `get_memory_buffers` 持有的一切、`ragged_verify_lens_cuda` / `ragged_qo_indptr_cuda`（`model_engine.py:715-720`）、新增的 `row_kv_lens_cuda`（`dsa.py:1321-1345`）、`gen_new_tokens_per_seq_cuda`（`deepseek_v4.py:337-346`）、DSpark worker 的 `captured_hidden_states` / `_kv_windows` / `_confidence_logits`。 | -| **必须恒定（capture 时冻结）** | 每一个 tensor **shape**，以及每一个用于切片或分支的 **host 标量**：`spec_metadata.total_verify_tokens`、`metadata.is_ragged_verify`、`use_expanded_buffers_for_mtp`、`next_n`、`num_gen_tokens_per_seq`、`attn_metadata.num_tokens` / `num_ctx_tokens` / `num_generations`、`all_rank_num_tokens`、`sum(ragged_verify_lens)`。 | -| **合法** | graph 私有内存池内的分配（`scatter_ragged_to_padded`、`count_accepted_ragged` 等，`dspark_ragged.py:432-512`）。 | - -这条清单解释了为什么 §3 里的多数 bug 是**静默**的：`num_gen_tokens_per_seq` 之流在 capture 时被烤进 graph，replay 时既不会更新也不会报错。同理，`dsa.py` 里那些 `assert` 只在 capture 时执行过一次，replay 时永远不跑。 - -**一个确认是安全、不要「顺手修」的点**：`spec_metadata.total_verify_tokens` 虽然是 captured region 内读取的 host int（`interface.py:1439-1441`），但它**等于 bucket，即 `key` 的一个分量**，因此是 graph key 的函数，构造上安全。 - -### 4.2 Graph 数量与内存代价（DSv4-Pro / dep8 实测） - -来自 `.dspark-logs/gsm8k.log` 的**实测**数字（不是估算）： - -``` -capturing 102 graphs (34 batch sizes x 3 draft-length tiers [1, 3, 5]) -``` - -| 项 | 值 | -|---|---| -| baseline graph 数 | 34 | -| confidence scheduling 后 | **102**（3×） | -| tier ladder | `[1, ceil(K/2), K] = [1,3,5]`（`llm_args.py:2883-2897`） | -| 每 graph metadata 成本 | ~10–23 MB（PR 自述，`model_engine.py:1825-1826`） | -| 额外 68 个 graph | **~0.7–1.6 GB / rank** | -| dep8 聚合 | **~5–13 GB** 的 KV 容量 | - -这笔钱**直接从 KV pool 里出**：captured graph 的字节落在 `extra_cost = total_used_bytes - torch_used_bytes` 项里，`_util.py:990-1015` 把它折进 `peak_memory`，而 `peak_memory` 决定 `kv_cache_max_memory`（该处日志行本身就点名了 CUDA graphs）。 - -**没有**第二次翻倍：`needs_non_greedy_capture` 门控在 `use_one_engine()` 上，DSPARK 不在其中（`interface.py:307-309` vs `:322-323`）；`DeepSeekV4SparseAttentionConfig.needs_separate_short_long_cuda_graphs()` 返回 False（`llm_args.py:1126-1128`）。 - -Warmup 会把整个 ladder 跑两遍（`model_engine.py:1220-1227`），即 204 次前向。 - -**这是个真实的张力**：feature 的收益按其自身 docstring 只在高并发出现（`llm_args.py:2812-2816`），而高并发恰恰是 KV 容量最吃紧的时候。缓解手段是把 tier ladder 按 batch size 门控（低于 SPS 拐点时只留 `[max_tier]`），因为 planner 的 argmax 在拐点以下本来也会选 max_tier。 - -### 4.3 剩余的同步点 - -| 位置 | 问题 | -|---|---| -| `model_engine.py:3910-3919` | `_update_target_input_tensors` 的 ragged 分支 `torch.repeat_interleave(previous_slots, tokens_per_request_device)` **没传 `output_size`**，两个参数都在 device 上 → torch 把 cumsum 读回 host 定尺寸。这正是 `dspark_ragged.py:414-424` 文档里明说不能干的事。在 capture 之外，所以不会让 capture 失败，但每个 ragged step 都在 `_apply_incremental_update_target` 里插一次 D2H sync，把 overlap scheduler 的收益吃掉。host 侧已有总数，直接传即可（旁边 `model_engine.py:5101-5109` 就做对了）。 | -| `heads.py:280-283` | `apply_sts` 在 CPU 快照上做校准，会把 CUDA 上的 `sts_temperatures` `.to('cpu')` 到**可分页**内存 —— 同步拷贝。每个 decode step 一次，且**只在提供了 profiled cost table 时触发**（否则 `is_flat` 提前返回），也就是恰好在真正要用的配置里出现。加载时缓存一份 host 副本即可消除。 | -| `py_executor.py:3232-3236` | ADP 下每 decode step 多一次阻塞式 host collective（`all_rank_max`），PR 自己的注释已承认。 | - -### 4.4 Attention-DP 形状一致性 - -`cuda_graph_runner.py:362-372` 的 ADP 资格门只比较 **`is_all_gen_only` 与 batch size 相等**，**从不比较 token 数**。token 数走另一条 allgather（`model_engine.py:2987-2998`）进 `all_rank_num_tokens`，且因为是 host list 在 forward 内读取，会被**烤进 captured graph** —— replay 时用的是 capture 时的那份，与本步 allgather 结果无关。 - -顺利路径下这是自洽的（`choose_ragged_capture_shape` 归约 `peer_stats` 使各 rank 落在同一 `(padded_bs, bucket)`，且 `fill_bucket` 宁可抛异常也不返回短 batch，`dspark_ragged.py:37-48`）。**暴露面在于任一 rank 退回均匀路径**：它产出 5 元组 key 和 `bs*(chosen+1)` 个 token，而同伴产出 6 元组 key 和 bucket 个 token；batch size 仍相等所以 ADP 门通过，一个 rank replay 满是集合通信的 graph、另一个 eager 跑不同 token 数，而冻结的 `all_rank_num_tokens` 使这个不匹配在 replay 时**无法被察觉**。 - -**要求：ADP 门必须同时比较 bucket。** - -### 4.5 一个尚未有人排查的 capture 期风险 - -`deepseek_v4.py` 里有三个 `@maybe_compile(dynamic=True, options={"max-autotune": True})` 的函数 —— `_prepare_deepseek_v4_indices_compiled`（`:586-589`）、`_compute_gen_compressed_position_ids`（`:955-957`）、`_compute_compressed_mask`（`:996-998`）—— 全部经 `on_update_kv_lens`（`:854-895`）在 `_preprocess_inputs`（`model_engine.py:2891-2895`）内被调用，**而那正是 `CUDAGraphRunner.capture` 录制的区域**。`_compute_gen_compressed_position_ids` 把 `num_gen_tokens_per_seq` 当 Python int 用（`:979-983`），构成 guard → **每个 tier 一次 Inductor 编译，且发生在 graph capture 期间**。 - -`max-autotune` 需要同步来 benchmark 候选 kernel。confidence ladder 把 capture 期编译次数变成 3 倍。这是 07/30 那次 hang 的**另一个独立候选根因**，必须与 H1 的 stride 假说区分开（见 §7）。「待验证」。 - ---- - -## 5. 方案对比 - -三个候选方案。注意 D1 与 D2 **不是独立选项** —— 二者都把成败押在同一个无法从源码证明的 trtllm-gen 行为上（K2 的 cubin 是预编译的）。真正的选择是 **token-major vs tier-uniform**。 - -| 维度 | **D1 Token-flat（`next_n ≡ 1`）** | **D2 Token-major + 3 处签名改动** | **D3 TUV 分层均匀** | -|---|---|---|---| -| **机制** | 把 batch 的 generation 半边呈现成 `num_gen_tokens` 个「长度为 1 的序列」，每 token 带自己的 kv_len / block table / cu_seqlens。所有 `rowIdx/next_n` 退化为恒等式。 | 引入每 token 行向量（`row_req_idx` / `row_kv_len` / `row_idx_end`），给 3 个 kernel 加**可选**参数；attention 仍用 token-major 呈现规避 cubin 改动。 | 不教 kernel ragged，而是把 verify 长度量化到 tier ladder，**保证任何一次 launch 只看到一个 `next_n`**。Stage A：整步一个 tier。 | -| **kernel 改动** | **0**（宣称）。compressor 无法展平（Phase 3 归约的是 Phase 1 同 block 写的状态，`compressorKernels.cu:355-400` vs `:440-530`），靠 host 侧按 tier 分组多次 launch + scratch 行绕开。 | 2 个 op 加可选参数（**工作区已完成**：K4、K9）。attention 不改。 | **0**。四个 uniform kernel 在 `num_gen_tokens = n·(K+1)` 时**逐一精确正确**（已逐条读码确认：`indexerTopK.cu:667-668`、`mlaKernels.cu:1161-1163`、`compressorKernels.cu:294,340`、`attentionOp.cpp:1161`）。 | -| **graph 数量** | 102（bs × tiers，token 轴入 key） | 102（同上） | 102 → 按 bs 门控可降到 **~54** | -| **预期收益** | 完整 Θ = τ·SPS(B) | 完整 Θ | 拿满 **SPS(B)** 项；损失 **τ** 项。D3 自估双峰 batch 下 τ 差距 ~14%（**未实测**） | -| **风险** | **高**。核心假设（K2 在 `mMaxSeqLenQ=1` 下正确且不显著变慢）无法读码证明。compressor 的 tier 分组 + nulled 行 + scratch 是全案最难 review 的部分。 | **高**（同一 K2 假设）。但 kernel 改动是加性可选参数，非 ragged 路径可证明不变。 | **中**。无 K2 假设。但 τ 留在桌上，且 Stage B（两 tier 分子 batch）依赖一个未验证的 FMHA semaphore 复用假设（`attentionOp.cpp:1126-1128`）。 | -| **工作量** | XL | L（**已完成约 40%**） | M | -| **被 D1/D2 漏掉的隐性成本** | token-major 要求 `num_seqs = host_context_lengths.size(0)`（`thop/attentionOp.cpp:1325`、`dsv3RopeOp.cpp:139`），故 `prompt_lens_cpu/cuda_runtime`、`kv_lens_runtime`、`host_request_types_runtime`、**主 KV 的 `kv_cache_block_offsets`**（与 indexer 的 `block_table_expanded` 是不同 buffer）都要按 token 膨胀；`max_num_requests` 变成 `max_batch*(1+max_draft_len)`，连带 `reserveSemaphoreArray`（`thop/attentionOp.cpp:407`）、FMHA/FlashMLA workspace（`common/attentionOp.cpp:925-955`）一起放大。4096 上下文下展开后的主 block table 约 400 KB，128K 上下文下约 12 MB，**每个 decode step 重新 H2D**。D3 完全不付这笔钱。 | 同左 | — | - -### 5.1 推荐 - -**分三阶段推进，先把「独立于方案选择」的东西修干净并拿到真实测量，再定 kernel 方向。** - -理由： - -1. **§3.3 的 H0–H5 全部是 host 侧问题，且三个方案都需要。** 尤其 H1（3 行 Python）就能完整解释唯一一次实测 hang，且**独立于 ragged 与否**。在它修好、`enable_confidence_scheduling` 的均匀路径能跑通之前，讨论 D1/D2/D3 是没有依据的。 -2. **决定 D1/D2 vs D3 的那个数字，谁都没测过**：trtllm-gen sparse-MLA generation 在 `(batch=B, s_q=tier+1)` 与 `(batch=B*(tier+1), s_q=1)` 两种呈现下、相同总 token 与相同逐 token 稀疏索引集时的吞吐。这决定 token-major 是否值得它的爆炸半径。它读不出来（`selectMlaGenerationKernel` 在预编译 cubin 里），但**一张 B200 上直接调 `thop.attention` 就能测**，不需要模型权重。 -3. **工作区已经沿 D2 走了一半**（K4、K9 已实现且设计正确 —— 可选参数、非 ragged 路径 bit-identical、TORCH_CHECK 替换而非放宽）。这部分不该丢，它们对任何 ragged 方案都是必需的，而且对 D3 无害（不传参数即退化）。 -4. **D3 有一个 D1/D2 没有的结构性优势**：`decide_draft_len` 的 `tp_allgather` 是**无条件**的（`dspark_verify.py:194-201`），且各 rank token 数恒为 `bs*(K+1)`，H9 那类 collectives 发散的 hang 在构造上不存在。同理 H7（rejection sampling）在均匀路径下是等式成立，不会被关掉。 - -**阶段划分：** - -| 阶段 | 内容 | 交付 | -|---|---|---| -| **P0：地基（不涉及方案选择）** | H1（stride 用 runtime tier，容量保持静态 max）、H2（attach 移到 prepare 之前）、H3（`row_kv_lens` host 侧折叠 + 单次 H2D）、H4（`on_update_kv_lens` 补刷 `row_kv_lens`）、H0+H5（dummy 分支用 `get_request_tokens_per_gen_step`；`capture()` 吃 bucket —— **必须同时改**）、K8（`prepare_scheduler_metadata` 走 `expand_per_gen_token`）、§4.3 两处 sync、H8 config 校验、K5/K6/K13 显式 config-reject + assert | **均匀 tier 路径**在 8×B200 上跑通 GSM8K；hang 消失 | -| **P1：测量与定向** | 单卡 B200 上的三个微基准（见 §6.4）；产出 STS 校准表与 SPS cost table（G8） | 用数据定 token-major vs tier-uniform | -| **P2：ragged 落地** | 按 P1 结论推进 D2（已完成 K4/K9）或 D3；无论哪条都要修 H6、H7、H9、H10 | ragged 路径 GSM8K + 逐 token 等价 | - -**若 P1 显示 token-major 显著变慢，D3 Stage A 就是 ship 目标** —— 它拿满 SPS 项（vLLM 报的 1.71×/2.19× 主要来自跨过吞吐拐点，而非预算分配的精细度），风险和工作量都低得多，且能把 graph 数从 102 压到 ~54。 - ---- - -## 6. 验收标准 - -> **所有测试都必须在 Slurm 分配的 GPU 节点上跑。当前节点不可用于任何 GPU 测试。** -> 用 `computelab-node-allocator` / `gpu-test-runner` skill 分配（B200/B300 分区，≥140 GB/GPU，≥4 h 预约）。 - -### 6.1 GSM8K 精度门（G3） - -```bash -export LLM_MODELS_ROOT=/scratch.trt_llm_data/llm-models # 容器内路径 -cd /code/tensorrt_llm && python -m pytest \ - "tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling" -vs -``` - -`MODEL_PATH` 在 **class body**（import 时）求值（`test_llm_api_pytorch.py:3964`），所以 `LLM_MODELS_ROOT` 必须在 pytest 启动**前**导出，否则整个 module 收集失败。无需 `mpirun` —— 单 pytest 进程通过 mpi4py.futures 起 8 worker。 - -**这个门本身不够**：GSM8K 的统计阈值是 `96.0 - 3.203 = 92.797`，允许 **42/1319** 题翻转；额外的 `score >= 96.0` 断言对实测 96.475 只有 **6.3 题**余量（1 题 = 0.0758 分），而 GSM8K 在 p=0.96 处的二项标准差约 0.54 分。它既太松（错几个 token 的 bug 能过）又太紧（干净 run 可能假失败）。 - -另外必须注意：该测试**没有被任何 CI stage 引用**（全 `tests/integration/` grep `confidence_scheduling` 只命中测试文件自身），且其 baseline 兄弟 `test_gsm8k_dep8_megamoe_deepgemm` 已被 `waives.txt:184` 以 nvbugs/6506920 跳过 —— **当前没有活的 baseline 可比**。 - -### 6.2 更锐利的不变量测试（G2，这才是真正的门） - -| 测试 | 内容 | 为什么锐利 | -|---|---|---| -| **A1 逐 token 等价** | 同一 engine 配置、同一批 prompt，`static`（verify 全量）与开启调度后，temperature=0 下断言每个 prompt 的 `token_ids` **完全相等** | 贪婪投机解码定义上是 lossless，任何分歧都是 bug。把 6 题的容忍度变成 1 token 的容忍度。**注意**：跨不同 graph 形状的严格 bitwise 相等 TRT-LLM 并不保证（MoE/attention 规约顺序会变），所以诚实的形式是允许极少数 prompt 分歧并设硬上限，或用 `cap-accept` 模式（下条）。 | -| **A2 `cap-accept` 差分** | 实现 SGLang 的三模式（`static` / `cap-accept` / `compact`）。`cap-accept` 跑**均匀 kernel 路径**（`next_n=6`）但只提交每请求窗口内的 token，输出应与 `compact` **逐 token 相同** | `cap-accept` vs `compact` 的差异 ⇒ **必然是 ragged kernel bug**；`static` vs `cap-accept` 的差异 ⇒ **必然是调度/接受逻辑 bug**。把 pass/fail 变成可定位。**已实现**：窗口写在独立属性 `py_verify_cap` 上（不碰 `py_verify_len`），布局侧因此看不见窗口、自动走整块；截断在设备端 acceptance 内部（`interface.apply_accept_caps`，经独立字段 `SpecMetadata.accept_caps` 触发，不借用 `is_ragged_verify` 以免连带走进 draft 缓冲的打包分支）。 | -| **A3 布局一致性断言** | host 侧断言：`sum(1 + py_verify_len over padded generation_requests) == attn_metadata.num_tokens - num_ctx_tokens == spec_metadata.total_verify_tokens == key[bucket]` | 这一条就能同时抓住 H0、H5、H8 | -| **A4 时序断言** | 断言 `attn_metadata.ragged_verify_lens` 在 `prepare()` 执行的那一刻是**本步**的（而非上一步） | 抓 H2、H4 | -| **A5 接受长度分布** | 记录并断言每请求接受长度直方图，以及 `sum(accepted)` 与接受 kernel 报告一致 | 抓「token 落到了错误请求的槽位但整批仍生成合法文本」这一类 | -| **A6 调度确实生效** | 断言 `planner.stats` 的 `fallback_flat_cost` / `fallback_no_snapshot` / `fallback_no_confidence` / `fallback_short_snapshot` 计数（`dspark_verify.py:86-93`）显示 trim 路径被走到 | 否则一个静默退化的 run 会以「正确的分数、错误的理由」通过 | -| **A7 capture 期无 sync/无分配** | capture 期间断言无 device→host sync、无新分配 | G4 | - -### 6.3 必须存在的单元测试 - -| 测试 | 现状 | -|---|---| -| C++ ragged top-k 差分：均匀 batch 用 ragged 形式表达（`row_kv_lens` 填成 uniform 算术会产出的值）必须与 `row_kv_lens=None` **bitwise 相同** | **已新增** `tests/unittest/_torch/thop/parallel/test_indexer_topk_ragged.py`（需在 GPU 节点上实跑验证） | -| compressor ragged 差分：`new_tokens_per_seq` 全填 `next_n` 必须与传 `None` bitwise 相同 | **欠缺** | -| `test_deepseek_v4_sparse_mla.py` 按每请求 q 长度参数化 | **欠缺**。该文件已构造真实的 `DeepseekV4TrtllmAttention` + cache manager（单卡 B200，`l0_b200.yml:80`，TIMEOUT 60），但其参考实现在 `:415` 和 `:473` **硬编码** `seq_len_q = fused_q.shape[0] // num_requests` —— 那正是被测假设本身。**这是性价比最高的一个测试**，单卡分钟级就能抓住 tier-3 那类 hang。 | -| `_prepare_tp_inputs` 在混合 `py_verify_len` + pad 行的 batch 上，断言 `sum(attn_metadata.seq_lens) == _ragged_verify_bucket(batch)` | **欠缺**。现有 220 个 hw_agnostic 测试全部停在纯 Python 布局算术；`test_dspark_ragged.py:515-521` 的 `_fit` 是 `fit_ragged_verify_lens` 的**手写镜像**而非真函数，所以 §3.3 的每一个 bug 对 CI 都不可见。 | - -### 6.4 定向微基准（P1，单卡 B200，分钟级） - -| # | 内容 | 用途 | -|---|---|---| -| **B1** | 修完 H1 后重跑均匀 tier capture，确认 hang 消失 | 区分 H1 stride 假说 vs §4.5 torch.compile 假说。**跑时带 `TORCH_LOGS=recompiles`。** | -| **B2** | 直接调 `thop.attention`，比较 `(batch=B, s_q=tier+1)` 与 `(batch=B*(tier+1), s_q=1)` | **决定 D1/D2 vs D3 的那个数字** | -| **B3** | expanded 布局下 paged-MQA-logits 的吞吐 | 展平路径的 KV 重读代价（每请求 6 次而非 2 次）在关键路径上，必须实测 | - -### 6.5 GSM8K 的并发要求 - -vLLM #47808 在**并发 16 和 64** 两个点上报了 0.945 / 0.951。原因很实际：**调度只在高并发下才真的 trim**，单一低并发的 GSM8K run 是空跑。本项目的验收也必须至少覆盖两个并发点。 - ---- - -## 7. 风险与未知 - -### 7.1 必须在定稿前测量的 - -| # | 未知 | 为什么阻塞 | 怎么解 | -|---|---|---|---| -| **U1** | trtllm-gen sparse-MLA generation 在 `mMaxSeqLenQ=1` + 6× batch 下 vs `s_q=6` 的吞吐 | **决定 token-major 是否可行**。DSv4-Pro 在 DEP8 下每 rank 保有全部 128 个 Q head，BMM1 的 M 维从 768 降到 128；由于稀疏索引集逐行不同，本来也没有跨 6 行共享 KV tile 的机会 —— 但这是**推断，不是测量**。 | B2 | -| **U2** | 07/30 hang 的真实根因 | H1 的 stride 不匹配是一个完整自洽的解释（tier 5 时 768==768 通过，tier 3 时 768 vs 512），但日志里紧邻 hang 的是 **attention workspace 从 0 扩到 168 MB** 的警告，`model_engine.py:1789-1806` 明确记录 capture 期 workspace 重分配是真实 hazard，而 Case 2b **没有** Case 2 那样的 `max_spec_graph` 规避（`model_engine.py:1810-1813`）。§4.5 的 torch.compile / max-autotune 是第三个候选。 | B1 + `TORCH_LOGS=recompiles`；必要时逐个消除 | -| **U3** | ragged 相对 tier-uniform 在固定 token 预算下的 τ 收益 | 决定 D3 是否「够好」。D3 自估双峰 batch ~14%、同质 batch ~0%，**没有任何真实 confidence head 的数据支撑** | 需要先有 STS 表（U4） | -| **U4** | STS 校准表 | `sts_temperatures` 初始化为全 1（`heads.py:249-258`），checkpoint 的 index 里**没有**任何 sts/temperature/calibration 张量，repo 里也没有该 JSON。论文 §3.2.1 明确是**每个位置一个温度**的 1D 网格搜索。未校准的逐位置偏差在累积乘积下**几何放大**。 | 必须先 profile 产出 | -| **U5** | SPS cost table | 没有它 `SpsCostTable.flat()` ⇒ `is_flat` ⇒ `_decide_local` 无条件返回 `max_tier`（`dspark_verify.py:268-270`）⇒ **feature 在构造上无法 trim**。任何在此之前的 A/B 只测到了「多一次 matmul + 3× graph」。 | 必须先 profile 产出 | -| **U6** | DSv4-Pro 实际走 KVCacheManager v1 还是 v2 | H6 的**症状和严重性按 manager 不同**：v1 有补偿项、误差方向是过度 rewind；v2 的 target 路径无补偿（`kv_cache_manager_v2.py:3538-3551`），investigation 的「泄漏」场景才成立。说错会把 debug 引向反方向。 | 读配置/加日志确认 | -| **U7** | token-major 下 buffer 膨胀的真实代价 | `num_seqs` 从 B 变成 ~6B，牵连 semaphore array、FMHA workspace、**主 KV block table 每步重新 H2D**。128K 上下文下约 12 MB/step。 | 若走 token-major，需容量审计 + 实测 | - -### 7.2 已知但可接受 / 需显式处理 - -- **性能故事必须诚实**：**draft pass 是不可裁剪的** —— block 永远整块 draft（`runtime_draft_len` 在 ragged 下钉在 top tier），只有 verification 被裁。所以收益是纯粹的 target 侧 SPS 效应，**只在高并发出现**，小 batch 下可能持平甚至因为 3× graph 占用 KV 而**净负**。vLLM 报的 1.71×(c=128) / 2.19×(c=256) 的对照组是「固定 7-token verify」，而后者在 c=256 时比不投机还慢 33% —— 那是被修复的病态，不是通用加速比。 -- **`derive_verify_len_tiers` 实际是死代码**：`_build_verify_planner` 只在 `verify_len_tiers` 为假值时才调它（`dspark.py:433-450`），而该属性总是返回非空 list（`llm_args.py:2899-2905`）。这**其实是安全的**（planner 的 tier 必须与 `_get_graphs_to_capture` capture 的一致），但 `dspark.py:450` 的日志永不触发，reviewer 会误以为 ladder 是 profile 派生的。要么删分支、要么让派生结果同时喂给 capture 列表。 -- **Case 2b 丢了 `draft_len=0` 的 graph**：它提前 return（`model_engine.py:1848,1856`），从不到达 `should_capture_no_spec` 逻辑（`:1867-1873`），而 tier ladder 最小值是 1。一旦 speculation gate 因低接受率关掉投机（`py_executor.py:3317-3322` 设 `runtime_draft_len = 0`），key 变成从未 capture 过的 `(bs, 0, ...)`，**之后每个 decode step 都 eager** —— 大 MoE 模型上是巨大的静默性能悬崖。 -- **PR 的外部引用有误**，会削弱 reviewer 信任：算法在论文 **§3.2.2**（Algorithm 1）而非 §5；SGLang **没有** `dspark_schedule.py`（其真实布局是 `speculative/ragged_verify.py` + `speculative/dspark_components/*`）；`align_verify_tokens_to_graph_tier` 全网零命中（SGLang 的实际原语是 `round_up_grid()` / `compute_target_verify_graph_key()`）。 -- **论文 §5.2 支持 PR 的一处偏离**：Algorithm 1 的 greedy first-descent 被论文自己否定（「真实硬件容量 SPS(B) 是离散的、阶梯状退化的」），生产中「移除 early-stopping break，启用无约束全局搜索」—— 这正是 PR 采用 global argmax 的理由，可以引用。但论文用的是 **t-2** 快照（为兼容 Zero-Overhead Scheduling 并形成「因果屏障」），PR 用的是 t-1。**t-1 是否足够早取决于 TRT-LLM 的 overlap scheduler，不取决于论文**「待验证」。 - -### 7.3 上游可复用的东西（尚未采用） - -TRT-LLM pin 的 DeepGEMM commit（`245dc5d6`，`cpp/build/_deps/deepgemm-subbuild/CMakeLists.txt:30`）**已经**在 `get_paged_mqa_logits_metadata` / `fp8_paged_mqa_logits` / `fp8_fp4_paged_mqa_logits` 上暴露了可选的 `indices` 参数（varlen 分支由 `indices.has_value()` 门控，断言 `dim()==1 && size(0)==batch_size`、int32、contiguous）—— 这就是论文说的「marker tensor」。vLLM #47808 用的正是它（`indices=decode_metadata.indices`）。 - -TRT-LLM 的 `dsa.py` **从不传 `indices`**，而是用既有的 "expanded buffers" 展平达到同样形状。二者数值等价，差别在于 expanded 路径要物化一个 `[num_tokens, max_blocks_per_seq]` 的 block table，而 varlen 只需一个 1-D int32 marker。**限制**：DeepGEMM 的 varlen 分支断言 `arch_major == 10 and next_n == 1` —— 仅 SM100，且必须是展平形式。这条路值得作为一个**有意的决定**评估，而不是默认忽略。 - ---- - -## 8. 开放问题（需要你拍板） - -1. **token-major 还是 tier-uniform？** 这是全文最大的分叉。建议先跑 B2 微基准再定 —— 但如果你出于其他原因（例如已经决定要与 SGLang/vLLM 的实现保持形态一致）倾向某一边，可以省掉这一步。 - -2. **`enable_ragged_verify` 是否是本轮必须交付的？** 若可以接受先只交付**均匀 tier 阶梯**（D3 Stage A），P0 阶段结束就能 ship 一个正确且 CUDA-graph 兼容的版本，把 ragged 留作 follow-up。这会显著降低风险，代价是 τ 项（未量化）。 - -3. **工作区那 11 个未提交的改动如何处置？** K4/K9 的实现质量是好的（可选参数、非 ragged bit-identical、TORCH_CHECK 替换而非放宽），但 H3/H4 是新引入的 P0 bug。建议：先修 H3/H4 再提交，**不要**在带着必抛异常的代码的状态下推进其他工作。 - -4. ~~**`cap-accept` 模式要不要实现？**~~ **已实现**（§6.2 A2）。它把 GSM8K 从「pass/fail 硬币」变成可定位诊断，并产出 `compact` 结构上拿不到的 `cap_trim_tokens`（裁剪的真实接受代价）。 - -5. **STS 表和 SPS 表由谁、按什么流程产出？** 这是 U4/U5，也是 G8。在它们存在之前，任何**性能**结论都是空的，而且 feature 在构造上无法 trim。需要一个 profiling 脚本 + 产物的存放约定（checkpoint 内？repo 内 `examples/configs/`？）。 - -6. **tier ladder 要不要按 batch size 门控？** 低于 SPS 拐点时 planner 本来也会选 `max_tier`，门控后 graph 数从 102 降到 ~54，换回 ~0.4–0.8 GB/rank 的 KV。需要一个共享 helper 让 `_get_graphs_to_capture`（`model_engine.py:1841-1856`）与 planner 的 `allowed_lens`（`dspark_verify.py:280`）**不可能漂移**。 - -7. **非贪婪采样（H7）的处理时机？** 现在 ragged batch 上 rejection sampling 被 fail-closed guard 静默关掉，改变了 temperature>0 的采样分布。要么本轮修 guard（放宽成 `num_contexts + total_verify_tokens`），要么显式在 config 层拒绝「ragged + 非贪婪」组合。**静默降级是最差的选项。** - -8. **KVCacheManager v1 还是 v2？**（U6）影响 H6 的修复优先级和症状描述。 - -9. **要不要采用 DeepGEMM 的 `indices` varlen 参数？**（§7.3）能去掉一个大 buffer 并与 vLLM 对齐，但把该路径钉死在 SM100。 diff --git a/docs/dspark_p0_report.md b/docs/dspark_p0_report.md deleted file mode 100644 index 4a7e17690c02..000000000000 --- a/docs/dspark_p0_report.md +++ /dev/null @@ -1,577 +0,0 @@ -# DSpark confidence-head 调度 P0 阶段：执行报告 - -> # ⚠️ 本文是中途快照,主要结论已被推翻 -> -> **当前事实以 `docs/dspark_p0_task_prompt.md` 第 10 节为准。** 本文写于 -> token-major 实现之前,其核心判断已不成立: -> -> | 本文的说法 | 现在的事实 | -> |---|---| -> | 「ragged 被 K1 挡住,从未真正跑起来过」(§0.4) | K1/K2 已由 token-major 呈现解除,**纯 Python,零 C++ 改动** | -> | 「ragged 在 DSv4 上结构上不可能端到端跑通」 | 已跑通:planner 自主产出 4 种窗口、211/318 步 ragged、裁剪 60.6% | -> | 「吞吐今天无法测量」 | 已测:此 checkpoint 无收益,因接受率 ~90% 时 planner 正确地不裁 | -> | 引用 `TLLM_DSPARK_FORCE_VERIFY_LENS` | 该开关已从代码中删除,改用合成陡峭 cost table 驱动真实 planner | -> -> 仍然有效的部分:§B2 的 token-major microbenchmark(0.93–1.00×)、对 K1/K2 -> 检查本身的代码分析、以及各处 `file:line` 锚点。 - - -> 分支：`dspark-p0`（基于 PR #17056 head `d8897046`，PR 分支未改动） -> 硬件：`bia` 集群，8× B300 SXM6（sm103），x86_64 -> 构建：`/lustre/fsw/coreai_comparch_trtllm/laliao/build/dspark-p0/`（`-a '103-real;'`，rc23） -> 日期：2026-08-01 -> -> 本文沿用 goal doc 的规范：凡未经读码或实测确认的推断，显式标注「待验证」。 - ---- - -## 0. 先说三个与任务书不同的结论 - -### 0.1 goal doc 的 11 个 host 阻断点里，7 个在 PR head 上已经修好了 - -任务书 §1(b) 说「H3 和 H4 现在是已经进了 PR head 的 bug」。**这一条不成立。** - -证据： - -``` -$ git log --oneline -S "row_kv_lens.add_(expanded" --all -- .../dsa.py -(空) -``` - -那个「CUDA tensor 加 CPU tensor」的形式**从未进入过任何 commit**。goal doc §1.3 描述的是提交*之前*的工作区快照；等它变成 `e1c28eeae2` 时，H3/H4 已经按 goal doc 建议的方式修好了： - -| 项 | 现状 | 锚点 | -|---|---|---| -| **H0** | 已修 | `model_engine.py:4670-4671` 非 overlap 分支已走 `get_request_tokens_per_gen_step` | -| **H2** | 已修 | `model_engine.py:5410-5412` `_publish_ragged_verify_lens` 在 `prepare()` **之前** | -| **H3** | 已修 | `dsa.py:1697-1757` 全在 host 上组装 + 单次 pinned H2D | -| **H4** | 已修 | `dsa.py:902` → `refresh_ragged_row_kv_lens`（`dsa.py:1759-1785`） | -| **H5** | 已修 | `cuda_graph_runner.py:462-464` `num_tokens_for_capture` 取 `key[5]` | -| **H6** | 已修 | `SampleStateSpec.verify_lens_snapshot`（`spec_sampler_base.py:63,246,340-354`） | -| **H10** | 已修 | `dspark.py:629-634` ragged 时用 `qo_indptr` | - -全部由 `e1c28eeae2` 引入（`git log -S` 逐个确认）。 - -### 0.2 B2 的结论是「token-major 不慢」，ragged 路线不被否决 - -见 §1。这与 goal doc §7.1 U1 担心的方向相反。 - -### 0.3 参数化后的 `test_deepseek_v4_sparse_mla` **抓不到** H1 - -goal doc §6.3 称它是「性价比最高的一个测试」。参数化之后它确实覆盖了宽度 1/2/4/6（此前只有 1），但**实测确认它对 H1 不敏感**：把 stride 修复回退成修复前的行为，12 个 case 依然全绿。 - -原因：该测试对 `compress_ratio=4` 的层**直接传入 `topk_indices`**（`deepseek_v4.py:1674-1676` 走 `forward_args.topk_indices`），于是 `Indexer.forward` 根本不执行 —— 而 `kv_lens_expanded` / `block_table_expanded` / DeepGEMM schedule 正是在那里被消费的。 - -这一点已写进测试文件本身，避免下一个人被绿色的测试误导。 - ---- - -## 0.4 ~~⛔ 最重要的结论：ragged 在 DSv4 上被 K1 挡住，从未真正跑起来过~~ - -> **【已推翻】** K1/K2 已由 token-major 呈现解除(纯 Python,`cpp/` 零改动), -> ragged 已端到端跑通。以下分析对「K1 检查本身做了什么」仍然准确,但结论 -> 「结构上不可能跑通」不成立。 - -端到端跑到真实 decode 之后，第一件撞上的事是： - -``` -RuntimeError: [TensorRT-LLM][ERROR] Assertion failed: - MLA can only support input sequences with the same sequence length. - (cpp/tensorrt_llm/kernels/mlaKernels.cu:1161) -``` - -调用栈毫无歧义： - -``` -modeling_deepseekv4.py:2509 forward - -> mla.py:1981 forward_impl_with_deepseek_v4 - -> mla.py:2147 forward_generation_sparse_mla - -> mla.py:2549 forward_absorption_generation - -> self.mqa.mla_rope_generation(...) - -> trtllm.py:2038 torch.ops.trtllm.mla_rope_generation - -> mlaKernels.cu:1161 -``` - -kernel 侧就是 goal doc §3.2 **K1** 逐字描述的那个检查： - -```cpp -TLLM_CHECK_WITH_INFO(params.acc_q_len % params.batch_size == 0, - "MLA can only support input sequences with the same sequence length."); -auto seq_len = params.acc_q_len / params.batch_size; -``` - -### 为什么这条是决定性的 - -commit `e1c28eeae2` 的标题是 **"Make DSpark ragged verification work on the DSv4 kernel path"**。 -但 `git show --stat` 显示它只动了两个 kernel： - -| 文件 | goal doc 编号 | -|---|---| -| `indexerTopK.cu` / `IndexerTopK.h` / `IndexerTopKOp.cpp` | **K4** 🔧 | -| `compressorKernels.cu/.h` / `compressorOp.cpp` | **K9** 🔧 | - -**`mlaKernels.cu`（K1）和 `attentionOp.cpp`（K2）从未被触碰** —— 而 goal doc §3.2 把这两个都标为 ❌ -kernel 层阻断点，且 K1 排在前面。所以 ragged 路径在 DSv4 上**结构上就不可能端到端跑通**， -这个 commit 的标题超出了它实际交付的范围。 - -### 危险的那一半 - -goal doc K1 还指出了比 abort 更坏的模式：窗口 `[6,2]` → 8 tokens / batch 2 → `seq_len=4`， -**检查通过**，但每个请求的 K/V 写进错误的行 —— 持久性 KV 缓存污染，静默。 -本轮撞到的是响亮的那一半（整除性不成立）；另一半在别的窗口组合下可达。 - -### 这对 D1/D2 意味着什么 - -goal doc §5 说得很清楚：D1 和 D2 **都**靠把 generation 半边呈现成 -「`num_gen_tokens` 个长度为 1 的序列」来让 K1/K2 的整除检查平凡通过。 -**那个呈现变更从未实现。** - -而 §1 的 B2 实测说这个呈现**不慢**（0.93–1.00×）。所以路线是通的、代价是可接受的， -只是**活还没干**。这正是 P0 之后该做的第一件事，而不是调参或找配置。 - -### 对「吞吐有没有提升」这个问题的直接回答 - -**今天无法测量。** 不是因为缺 SPS 表，而是因为 ragged 在 DSv4 的 MLA generation 路径上 -根本跑不起来。任何声称测到了 ragged 吞吐的数字，要么没真正走 ragged（planner 静默退化）， -要么就是撞了上面那个静默的 KV 污染分支。 - -## 1. 交付物 0：B2 实测数字（最先交） - -**问题**（goal doc §7.1 U1）：trtllm-gen sparse-MLA generation 在 `(batch=B, s_q=tier+1)` 与 -`(batch=B*(tier+1), s_q=1)` 两种呈现下，相同总 token、相同逐 token 稀疏索引集时的吞吐差多少。 -D1 和 D2 都靠把 generation 半边呈现成「长度为 1 的序列」来让 -`thop/attentionOp.cpp` 的 `num_tokens % num_seqs == 0` 平凡通过，二者押的是同一个赌注。 - -**脚本**：`tests/microbenchmarks/dsv4_sparse_mla_presentation.py`（单卡，不需要模型权重，分钟级） - -**结果**（B300 sm103，context_len=4096，20 iters，中位数）： - -| B | tier | next_n | tokens | batch-major (ms) | token-major (ms) | slowdown | -|---|---|---|---|---|---|---| -| 4 | 1 | 2 | 8 | 0.148 | 0.146 | **0.99×** | -| 4 | 3 | 4 | 16 | 0.142 | 0.137 | **0.97×** | -| 4 | 5 | 6 | 24 | 0.150 | 0.140 | **0.93×** | -| 16 | 1 | 2 | 32 | 0.147 | 0.139 | **0.95×** | -| 16 | 3 | 4 | 64 | 0.154 | 0.146 | **0.95×** | -| 16 | 5 | 6 | 96 | 0.145 | 0.145 | **1.00×** | -| 64 | 1 | 2 | 128 | 0.148 | 0.147 | **0.99×** | - -另一次独立测量：B=8, tier=5 → **0.88×**。 - -**结论：token-major 没有变慢，范围 0.93×–1.00×（甚至略快）。** -按任务书 §3 表三的判据（「若 token-major 显著变慢或不正确：立刻停下来报告，改推 D3 Stage A」）， -**不触发兜底条件，ragged 路线在吞吐上是可行的。** - -这与 goal doc §7.1 U1 的推断方向一致（DEP8 下每 rank 保有全部 Q head，BMM1 的 M 维从 -`128·(tier+1)` 降到 128；且稀疏索引集逐行不同，本来就没有跨窗口共享 KV tile 的机会可损失）， -但现在是**测量**而不是推断。 - -**未覆盖**：B≥64 且 tier≥3 的点。这些点卡在 benchmark 自己的 KV pool sizing 上 -（token-major 呈现有 `next_n` 倍的 sequence，整块分配的 rounding 也就付 `next_n` 倍）， -不是 kernel 结果。已把 pool 按整块向上取整并加一块 slack，仍未覆盖到 B=128。 -考虑到 7 个点跨 B=4..64 × 全部三个 tier 全部落在 0.93–1.00×，方向性结论是稳的； -但**「高并发下是否仍然如此」严格说仍待验证**，而高并发恰恰是这个 feature 唯一有收益的区间。 - -**范围说明**：该脚本计时的是整个 attention layer，包含 indexer（其呈现也随之改变： -batch-major 走 expanded 一 token 一行的路径，token-major 走 `next_n == 1` 的 strided 路径）。 -这是有意的 —— 那正是 ragged step 真实要付的代价，并且它顺带覆盖了 B3。 -纯 attention 的差值无法在不上 kernel-level profiler 的前提下分离出来。 - ---- - -## 1.5 θ 是真实且可观的 —— 一个我先前判断错误的更正 - -L=1 与 L=3 两组样本齐了之后，第一次可以把 `theta(M)` 从 `alpha(bs)` 里分离出来： -固定 `bs` 时 `T(bs,L=3) - T(bs,L=1)` 让 `bias` 和 `alpha(bs)` **精确抵消**， -剩下的就是 `theta(4bs) - theta(2bs)` —— 唯一能被裁剪减少的那一项。 - -| bs | tokens | Δθ (ms) | 占该步耗时 | -|---|---|---|---| -| 8 | 16→32 | +3.75 | +15.0% | -| 16 | 32→64 | +5.26 | +18.3% | -| 32 | 64→128 | +9.30 | +27.5% | -| 64 | 128→256 | +7.42 | +17.4% | -| **128** | **256→512** | **+18.33** | **+37.0%** | - -**我先前的判断是错的。** 我根据 B2（sparse-MLA layer 在 8→128 token 上基本持平， -0.137–0.150 ms）推断 θ 没什么可裁的空间。错在把**一个 attention layer** 外推成**整个 decode step**： -该 layer 确实对 token 数不敏感，但整步的开销由 **MoE** 主导，而 MoE 的代价随 token 数增长 -（每个 token 都要路由到专家）。所以 θ 既真实又很大，且**随 batch 增大而增大** —— -这正好符合该 feature docstring 自称的「收益只在高并发出现」。 - -对照可裁剪门槛（profiler 源码自述：相邻 tier 之间需 >15% 才值得裁）， -**bs=128 的 37% 有充裕余量**。 - -**上界性质**：profiler 通过缩小 `block_size` 来扫 L，draft pass 也跟着缩， -所以测得的 Δθ 里混入了 draft 代价，是**上界**。这使得「小值」是可靠的否定结论， -而「大值」需要进一步确认。要测无泄漏的版本，就是保持 `block_size=5` 不动、 -改用 `TLLM_DSPARK_FORCE_VERIFY_LENS=` 强制窗口 —— **【该开关已删除】**,现由合成陡峭 cost table 驱动真实 planner -（见 §2 的 commit）。这是下一步该做的测量。 - -## 2. 交付物 1：修复清单 - -每项一个 commit，commit message 里引用 goal doc 编号。 - -| commit | 项 | 内容 | -|---|---|---| -| `cb5cb0bb92` | **H1 + K8** | DSA 展开 stride 走 runtime tier | -| `44a81930ca` | **§4.3 sync ×2** | `repeat_interleave` 补 `output_size=`；`apply_sts` 缓存 host 副本 | -| `f99c34a498` | **H8 / K5 / K6 / K13** | 显式 config-reject，不静默降级 | -| `b3a5fef92a` | **H9 + §4.4** | 集合通信收敛为恒定一次；ADP 资格门比 bucket | -| `3b52d19fcd` | **H7** | rejection sampling 的 fail-closed guard 用 `total_verify_tokens` | - -### 2.1 H1 + K8：这是 tier-3 hang 的自洽解释 - -`DSAtrtllmAttentionMetadata.max_draft_tokens` 是**静态上限**，不是本步窗口。 -对 parallel-draft 模式（DSpark 属于其一），`forward()` 传给 `update_spec_dec_param` 的是 -`original_max_total_draft_tokens`（`model_engine.py:6686-6688`）， -即 `tokens_per_gen_step - 1`，engine 生命周期内**永不改变**。 - -于是： - -- **producer**（`prepare_for_spec_decode`）按 stride `1 + max_draft_tokens = 6` 铺 - `kv_lens_expanded` / `block_table_expanded`； -- **consumer**（`Indexer.forward`, `dsa.py:2967`）按 `num_gen_tokens // num_generations` 切片。 - -tier 3 时 consumer 取 4。3 个请求、宽度 4 的例子： - -``` -正确: idx 0-3 → kv0, 4-7 → kv1, 8-11 → kv2 -修前: idx 0-5 → kv0, 6-11 → kv1, (kv2 从不出现) -``` - -**每个 r>0 的请求都拿到邻居的 kv_len 和 block table。** - -K8（`dsa.py:2496-2497`）是同一个错位，但落在一个会 **hang 而不是算错**的地方： -`prepare_scheduler_metadata` 按 `num_generations * (1 + max_draft_tokens)` 行建 DeepGEMM -schedule，而 paged-MQA-logits 只拿到 `num_gen_tokens` 行 —— schedule 承诺了没人会产出的 tile。 - -这与唯一一次实测失败**逐字吻合**：`draft_len=5` 时 warmup 完成（768 == 768）， -`draft_len=3` 时永不返回（768 vs 512）。 - -**修法**：新增 `runtime_tokens_per_gen_step` 与 `gen_token_stride`，所有 *stride* 走它， -所有 *capacity* 保持静态上限（expanded buffer 必须按最大值分配，否则较短的 tier 会在 -capture 期触发重分配 —— 那是第二类 hang）。未发布（0）时回落到 `1 + max_draft_tokens`， -所以任何不缩短 draft 长度的路径 bit-identical。 - -发布点在 `_prepare_tp_inputs` 而不是 `update_spec_dec_param`：后者作用于 base metadata， -而 `prepare()` 跑在之后取的 per-key CUDA-graph 副本上，对 base 的属性赋值对副本不可见。 - -`use_expanded_buffers_for_mtp` **有意**保持按静态上限判定，使 layout regime 不随 tier 变化。 - -### 2.2 H9：这不只是「多一次集合通信」，是时序相关的死锁 - -`_dspark_confidence_draft_len` 里每个 rank 发出的集合通信次数不同： - -| rank 走到哪 | 次数 | -|---|---| -| 完整 ragged | 2（planner 的 `all_rank_max` + `peer_stats`） | -| 在 planner 处放弃 | 1（uniform `decide_draft_len`） | -| 到 bucket fit 才放弃 | 3 | - -而**所有放弃的理由都是 rank-local 的**，其中一条 —— confidence 快照的 copy event -是否落地（`_ready_snapshot` → `_copy_event.query()`）—— **纯粹是时序**。 -同一步上两个 rank 可以因为谁都观察不到的原因走进不同分支。 - -**修法**：拆成三段 —— (1) 本地决策，不发任何集合通信；(2) 发**一次**定宽 allgather； -(3) 基于归约结果分支，之后不再发。payload 顺带把 `peer_stats` 带上，所以常见路径从 -两次降到**一次**。只有当**所有** rank 都能走 ragged 时才走 ragged；有一个反对就整组退回 -uniform —— 那是唯一一种 token layout 仍然一致的结局。 - -### 2.3 §4.4：ADP 资格门 - -`fit_ragged_verify_lens` 在达成一致*之后*仍可能在单个 rank 上失败（它的 pad 行拟合用的是 -本地请求数），所以 ADP 门必须兜住。它此前只比 `is_all_gen_only` 和 batch size，**从不比 token 数**。 -而这个不匹配在 replay 时**无法被察觉**：`all_rank_num_tokens` 是 forward 内读取的 host list， -replay 的 graph 用的是 capture 时那份。现在同时比 ragged token bucket 与 draft length。 - -### 2.4 H7：让精心写的 ragged 分支不再是死代码 - -`_rejection_buffers_valid` 要求 `num_contexts + num_gens * (draft_len + 1)` 行， -而 ragged 下 target 只发 `num_contexts + total_verify_tokens` —— 差值恰好是调度 trim 掉的量。 -**guard 恰好在 feature 起作用时失败**，然后 fail-closed 静默退回严格接受，改变了 -temperature > 0 的采样分布。 - -rejection 路径的 ragged 分支本身是**完整的**（按 `total_verify_tokens` 打包、scatter 进 -padded 矩形、pad 行填 one-hot 使其不可能被接受、最后把 accepted 数 clamp 回 `verify_lens`）， -所以该修的是行数而不是 config —— 直接拒绝「ragged + 非贪婪」会丢掉能用的代码。 - -goal doc 说得对：**贪婪 GSM8K 观察不到这一点**，所以只能从代码论证。 - ---- - -## 3. 交付物 2：测试 - -| 测试 | 状态 | -|---|---| -| `test_deepseek_v4_sparse_mla` 按每请求 q 长度参数化（宽度 1/2/4/6） | **12/12 通过**，但见 §0.3 的覆盖限制 | -| compressor ragged 差分（`test_compressor_ragged.py`，新增） | **4/4 通过** | -| C++ ragged top-k（`test_indexer_topk_ragged.py`，PR 已有） | **通过** | -| DSpark hw_agnostic 全套（confidence_schedule / ragged / worker / overlap_layout / heads / cuda_graph） | **通过** | -| A3 布局一致性断言（`TLLM_DSPARK_ASSERT_LAYOUT=1`） | 已加；e2e run 全程启用，**57 个 graph 的两轮 warmup/capture 无一触发** | -| A4 时序断言 | 已加（无条件，代价是一次 host 求和） | -| A6 planner 计数并入 ragged stats summary | 已加 | - -上述 DSpark 相关单测合并跑：**202 passed**（B300，13 秒）。 - -> 注意：`tests/unittest/_torch/speculative/hw_agnostic/` 整目录跑会有大量失败， -> 但全部是环境问题 —— 该目录里有需要 llama-3.1 / Qwen3 / EAGLE3 / DeepSeek-V3-Lite -> 权重的用例，而本轮的 `LLM_MODELS_ROOT` 只软链了 DeepSeek-V4-Pro-DSpark； -> 另有若干用例会自己起 `MpiPoolSession`，在已有 MPI world 里 spawn 失败。 -> 与本轮改动无关。 - -### 3.1 参数化过程中抓到的两个测试自身的 bug - -1. `_allocate_kv_cache_for_generation` 每步只扩一个 token —— 只有在 - `generation_seq_len_q == 1` 时才对。宽度 4 时第二步一旦跨页就越界（illegal access， - 不是干净的失败）。 -2. `_build_compressed_topk_indices` 按**每请求**一行构建，而 kernel 要的是**每 query token** - 一行。这正是被测试的那个「均匀窗口」假设本身。 - -### 3.2 compressor 差分测试的设计 - -第一个测试（uniform 填充 == 传 `None`，bitwise）是让那个可选参数**可以被接受**的论据。 -但它有一个结构性盲点：对于一个**完全忽略** `new_tokens_per_seq` 的 kernel，uniform 向量 -也是正确答案。所以第二个测试用真正不同的计数，与**单请求** uniform 调用逐个比对 —— -batch 为 1 时「uniform」无歧义，且与 ragged 调用不共享代码路径。 - ---- - -## 4. 交付物 4：hang 的真实根因 —— **H1/K8，已实测确认** - -goal doc §7 列了三个独立候选。**本轮 e2e run 把它们分开了。** - -### 4.1 实测：hang 消失 - -07/30 那次失败（goal doc §1.4）： - -``` -DSpark confidence scheduling: capturing 102 graphs (34 batch sizes x 3 draft-length tiers [1, 3, 5]) -Attention workspace size is not enough, increase the size from 0 bytes to 168493568 bytes -Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=5 ← 完成 -Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=3 ← never returns -TEST_EXIT=124 -``` - -本轮（修复 H1/K8 后，overlap scheduler 开启，DEP8）： - -``` -06:36:23 DSpark confidence scheduling: capturing 57 graphs (19 batch sizes x 3 draft-length tiers [1, 3, 5]) -06:36:23 Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=5 - Attention workspace size is not enough, increase the size from 0 bytes to 168493568 bytes ← 同一个数字 -06:36:24 Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=3 ← ~1 秒完成 -06:36:29 Run generation-only CUDA graph warmup (greedy) for batch size=128, draft_len=1 -06:36:35 ... batch size=120, draft_len=5 - ... 继续走完全部 19 个 batch size × 3 个 tier -``` - -（graph 数从 102 变 57 是因为本轮配置的 batch-size 桶只有 19 个，与 hang 无关。） - -### 4.2 修完 hang 之后立刻暴露出来的第二个 bug（新发现） - -hang 消失之后，run 走到了**从来没有人到达过的地方**，于是撞上了下一个问题： - -``` -Run generation-only CUDA graph capture (advanced sampling) for batch size=1, draft_len=1 ← 57 个 graph 全部 capture 完成 -Run warmup with 4096 tokens, include 0 generation tokens -Run warmup with 128 tokens, include 128 generation tokens -RuntimeError: shape '[128, 1]' is invalid for input of size 640 -``` - -`640 = 128 × 5`。根因：capture 循环给每个 graph 盖上 -`self.runtime_draft_len = draft_len`（`model_engine.py:2056`），**结束后不还原** -（`:2081-2082` 只还原了 `enable_spec_decode`）。 - -只有一个 draft length 要 capture 时这无害 —— 留下的值就是原本的值。 -但 **tier ladder 结束在它最小的 tier 上**，于是 engine 带着 `runtime_draft_len=1` -走出 capture，而紧接着的通用 token warmup 按完整 draft 长度造 dummy 请求。 -acceptance 把 draft 重塑成 `[num_gens, runtime_draft_len]`，就成了 640 → `[128, 1]`。 - -这个泄漏**早于本 feature**（随 dynamic draft length 一起进来的）， -但只有多 tier 阶梯才会让泄漏值与正确值不同。已修（commit `e59c8101f2`）。 - -**这条本身就是 P0 阶段的一个论据**：capture 跑通不等于 feature 跑通， -而在 hang 被修掉之前，这个 bug 在物理上不可能被观察到。 - -### 4.3 三个候选的判定 - -| 候选 | 判定 | 依据 | -|---|---|---| -| **1. H1 stride 不匹配** | **✅ 根因** | 唯一被修的东西；修完之后 `batch=128, draft_len=3` 从「永不返回」变成 ~1 秒完成。且它的机制与症状逐字吻合：tier 5 时 768==768 通过，tier 3 时 schedule 描述 768 行而 logits 只有 512 行 —— K8 那一处**机制上就是 hang 而不是算错**（DeepGEMM schedule 承诺了没人会产出的 tile，kernel 等在那里）。 | -| **2. capture 期 attention workspace 重分配** | **❌ 排除（是共存现象，不是原因）** | 同一条警告、**同一个字节数 168493568**、同一个位置（紧邻 `draft_len=5` warmup）在本轮**照样出现**，而 capture 正常走完。所以它不是 hang 的原因。 | -| **3. torch.compile / max-autotune 在 capture 区内** | **❌ 排除为根因** | tier 5 / 3 / 1 各自在 1–6 秒内完成。若 Inductor 在 capture 区内编译并 hang，这里就走不过去。Inductor 编译大概率确实发生了（每个 tier 一次 guard），但它会完成。**未用 `TORCH_LOGS=recompiles` 单独确认编译次数，这一点仍标「待验证」。** | - -**结论：hang 的真实根因是 H1（连带 K8）—— DSA 展开 stride 用静态 max 而非 runtime tier。** -候选 2 是无害的共存现象；候选 3 未被证明有害。 - -## 4.4 host-device sync 验证（overlap scheduler 不受影响） - -**静态审计**（逐行核实 `dspark_ragged.py` / `dspark_verify.py` / `dspark.py` / `heads.py`）： - -| 命中 | 判定 | -|---|---| -| `dspark_verify.py` 两处 `.numpy()` | ✅ survival 来自 `_host_buffer`（**pinned CPU**）；planner 全程在 host 快照上算 | -| `dspark_ragged.py` 两处 `.tolist()` | ✅ `verify_lens` 是 CPU 张量 | -| `dspark_ragged.py` `row_ids_from_lens` | ✅ `output_size` 是**必填参数**，docstring 说明这里 sync 是 "outright illegal" | -| `dspark.py:634` `qo_indptr.to(device)` | ✅ 同 device，no-op | -| `dspark.py:551` `chunk_positions[0].item()` | ⚠️ 真 sync，但在 **prefill** 路径且由 base DSpark `147404864d` (#15808) 引入，**不属本 PR** | - -加上本轮已修的两处（`repeat_interleave` 补 `output_size`、`apply_sts` 缓存 host 副本）， -**decode 路径上 confidence/ragged 没有新增 sync**。 - -静态分析不是证明，所以加了 `TLLM_DSPARK_SYNC_DEBUG=`：在 n 个 DSpark decode 迭代后 -arm `torch.cuda.set_sync_debug_mode("warn")`。**故意晚 arm** —— warmup 与 graph capture -本就合法地同步，一开始就开会把信号埋掉。用 `warn` 不用 `error`，一次 run 能看到全部违规者。 -**实测结果待补。** - -## 4.5 kernel 改动兼容性 - -K4 / K9 是 PR 自带的（`e1c28eeae2`），本轮验证其「纯加性」性质： - -| 判据 | 证据 | -|---|---| -| (a) 不传新参数时 uniform 路径不变 | `indexerTopK.cu:680-688` 的 `else` 分支是**原公式逐字保留**；`compressorKernels.cu:348` 在指针为空时 `nn == NEXT_N`，`sp = kv_len - NEXT_N` 即原式，且 `NEXT_N` 仍是模板常量（保住全展开） | -| (a') bitwise 差分实测 | compressor：uniform 填充 == 传 `None`，**4/4 通过**；indexer top-k：`test_ragged_with_equal_lens_matches_uniform`，在 202 passed 内 | -| (b) ragged 结果正确 | compressor 与**单请求 uniform** 逐个比对（避开「kernel 直接忽略参数也能过」的盲点），通过 | -| (c) 互斥组合显式拒绝 | kernel `TLLM_CHECK`（`indexerTopK.cu:921`）+ 本轮新增的 Python config reject | -| (d) G6 两开关关闭零回归 | **穷举证明**：新 stride 与原 `1 + max_draft_tokens` 对 MTP / DSpark / PARD / DFlash 在 K=1,3,5,8 上全部相等。唯一变化发生在 `runtime_draft_len < max`（动态 draft 长度），那是修 bug | - -## 4.6 新增的可观测性：trimming 到底有没有代价 - -SGLang 博客（https://www.lmsys.org/blog/2026-07-06-dspark-sglang）的核心指标是 -`accept_length / block_accept_length`（报 0.88–0.97）。TRT-LLM 缺的是**分母** —— -`ceiling_tokens` 数的是「不裁会**提交**多少 token」，回答「省了多少算力」， -永远回答不了「毁了多少接受率」。后果：planner 过度裁剪只表现为端到端接受长度缓慢下降， -与「workload 变难」不可区分。 - -**没有照抄 SGLang 的估计器**，而是直接测那件事：被裁到窗口 `w` 的请求若**接受了整个 `w`**， -说明剪的时候 draft 还活着，这一刀确实有代价；接受数 `< w` 的本来就会死，剪了不亏。 -`trim_regret_rate` = 被裁请求中前者的占比 —— **精确、无偏、不需要模型**。 -（从未被裁的请求外推是**构造性有偏**的：哪些不裁正是 planner 挑的。） - -记录点在 sampler，因为那是唯一一处「接受数」与「它对应的窗口」同时在 host 侧且属于同一步 -的地方 —— 活的 `py_verify_len` 那时已翻到下一步（正是 H6 快照存在的原因）。 - -## 5. 交付物 3：e2e 实测 - -### 5.1 GSM8K = **96.2092** (± 0.526) - -配置：uniform tier ladder（`enable_confidence_scheduling=True`）、overlap scheduler 开启、 -DEP8、MEGAMOE_DEEPGEMM、B300 × 8。 - -| | | -|---|---| -| 实测 | **96.2092**（flexible-extract 与 strict-match 一致） | -| 记录参考 | 96.475 | -| 差 | 0.27，**在 1 个标准误（0.526）之内** —— 统计上不可区分 | - -这印证了 goal doc §6.1 对该门的批评：`>= ref_accuracy` 对实测值只有 6.3 题余量， -而 GSM8K 在 p=0.96 处的二项标准误约 0.54 分 —— **既太松又太紧，干净的 run 会假失败**。 - -> pytest 报 fail，但**精度断言从未执行到**：失败在 -> `tensorrt_llm/evaluate/lm_eval.py:657` → `numpy.mean` 对字符串 dtype 求平均， -> 是本轮拼装环境里 lm_eval 0.4.12 与容器 numpy 2.1.0 的版本不兼容，与模型和本轮改动无关。 -> lm_eval 已完整跑完 1319 题并输出了分数。 - -### 5.2 这个分数是**基线**，不是调度生效后的分数 - -随附计数器把这一点讲得毫不含糊： - -``` -mode: static, steps_total: 114, steps_no_windows: 114 -planner: {decisions: 114, fallback_flat_cost: 112} -requests_trimmed: 0, trim_ratio: 0.0, trim_regret_rate: 0.0 -accept_len: 3.3009 (requests_scored: 3436) -graph_replays: 520, graph_eager: 9 -``` - -`fallback_flat_cost: 112/114` —— 该 build 尚未接上 SPS 表（开关是启动之后才补的）， -所以 planner 每步都拒绝裁剪。**这次 run 等价于 baseline DSpark，只是多捕获了 3× graph。** - -**这正是可观测性存在的理由**：一个「分数正确、日志正常、但 feature 什么都没干」的 run， -靠计数器一眼看穿，不需要推测。 - -其他可用数字： - -| 量 | 值 | 含义 | -|---|---|---| -| `accept_len` | **3.3009** | DSpark block drafter 在 GSM8K 上的真实接受长度（5 个 draft 位置接受 3.30，66%）—— τ 项的基准 | -| `graph_replays / (replays+eager)` | **520/529 = 98.3%** | G4：稳态 decode 基本全部命中 graph replay | -| `trim_regret_rate` | 0.0 | 与 `requests_trimmed: 0` 一致 | - -### 5.3 ragged 路径：**跑不起来**（见 §0.4） - -死于 K1（`mlaKernels.cu:1161`），31:51。 - -### 5.4 复现环境的坑（供下一个人参考） - -1. 容器无 `pytest` / `lm_eval` / `mako` / `oyaml` / `parameterized` / `sacrebleu`。 -2. **不要用 `pip install --target` 装带依赖的包**：`PYTHONPATH` 永远排在 site-packages 前面， - 一次 `pip install lm_eval` 就把 `numpy` / `transformers` / `datasets` / `pandas` / `pyarrow` - 等 **50 个包**装进去并 shadow 掉容器版本，症状是 - `No module named 'torch._C'` 或 `No module named 'transformers.utils.chat_template_utils'`。 - 做法：装完立刻把与容器重名的目录全部隔离。 -3. `srun --ntasks=8 pytest` 会让 LLM API 再去 `MPI_Comm_spawn`，报 `MPI_ERR_SPAWN`； - **必须套 `trtllm-llmapi-launch`**。 -4. 被 kill 的 srun step **不会回收 MPI worker**，它们继续占每卡 ~180 GiB， - 下一次 run 会在 executor 创建时 OOM 并「怪自己」。启动前必须清理（`tmp/cleanup_node.sh`）。 -5. SPS profiler 需要 `TORCHINDUCTOR_COMPILE_THREADS=1`，否则 Inductor 每遇新 shape - 就 fork 一个编译 worker（父进程常驻上百 GB × 8 rank）→ 宿主 OOM（exit 137）。 -6. **在花 25 分钟建 engine 之前，先廉价验证整条 import 链**。本轮为此损失了三次 25 分钟。 - -## 6. 交付物 5：对 goal doc §8 剩余待拍板问题的建议 - -### Q4：`cap-accept` 要不要实现 —— **建议实现，但优先级低于 Q5** - -任务书倾向实现，理由是把 pass/fail 硬币变成可定位诊断。这个理由成立，但要注意 -**【已删除】** ~~PR head 已经有了一个更便宜的等价物~~：`TLLM_DSPARK_FORCE_VERIFY_LENS=1` -（`dspark_observability.py:131-162`）从 tier ladder 里按批内位置轮转发窗口， -确定性、每个 rank 一致、且不需要 cost table。它已经能把 -「ragged packing 对不对」与「planner 要不要 trim」分开 —— 而这正是 A2 想要的那一半。 - -`cap-accept` 额外提供的是另一半：跑**均匀 kernel 路径**但只提交窗口内 token， -于是 `cap-accept ≠ compact` ⇒ 必然是 ragged kernel bug。这个价值是真的， -但在 forced-lens 已经存在的前提下没有那么紧急。 - -### Q5：STS 表和 SPS cost table —— **这是当前最高优先级的缺口** - -不是「性能结论是空的」那么简单：**没有它 feature 在构造上无法 trim**， -于是任何 e2e run 都只测到「多一次 matmul + 3× graph」。本轮的 run 就正好落在这个状态 -（`profiled_cost_table=False`）。 - -建议：先做 SPS cost table 的 profiling 脚本（比 STS 简单得多 —— -它只是「总 verified token 数 → decode step 耗时」的一维曲线，用现成的 -`dspark_sps_profiler.py` 加一个 driver 即可），存到 `examples/configs/`， -按 (模型, 并行度, GPU) 命名。STS 表需要 1-D 网格搜索，成本高，排在其后。 - -### Q6：tier ladder 按 batch size 门控 —— **建议做，但要一个共享 helper** - -102 → ~54 个 graph，换回 ~0.4–0.8 GB/rank 的 KV。风险在于 -`_get_graphs_to_capture` 与 planner 的 `allowed_lens` 一旦漂移， -planner 就会选一个没有 captured graph 的 tier，**静默掉出 graph replay** —— -那比多占的显存贵得多。所以必须让两边读同一个函数，而不是各写一遍。 - -### Q7：非贪婪采样 —— **已按「修 guard」处理**（见 §2.4），不再需要拍板 - -理由：ragged 分支本来就是完整的，静默降级是三个选项里最差的那个。 - -### Q9：DeepGEMM 的 `indices` varlen 参数 —— **建议本轮不动，但记录为明确的后续项** - -B2 的结果（§1）说明 expanded 路径**没有**吞吐问题，所以采用 `indices` 的收益 -主要是省掉 `[num_tokens, max_blocks_per_seq]` 那个 block table 的物化， -而不是速度。代价是把该路径钉死在 SM100（DeepGEMM varlen 分支断言 -`arch_major == 10 and next_n == 1`）。在 B2 已经证明 expanded 不慢的前提下， -这笔交易不划算 —— 除非 block table 的 H2D 在 128K 上下文下成为瓶颈（U7，未测）。 - ---- - -## 7. 还没做的 - -1. **e2e 分数**（§5）—— 均匀 tier 与 ragged 两条路径都没有分数。 -2. **B1 的对照实验**：`TORCH_LOGS=recompiles`，用来排除 §4 的候选 3。 -3. **两个并发点**（goal doc §6.5）—— 一个都没跑。 -4. **G2 逐 token 等价**、**A2 `cap-accept` 差分**。 -5. **B2 在 B≥64 且 tier≥3 的点**（§1 末尾）。 -6. ~~`scripts/generate_llm_args_golden_manifest.py`~~ —— **已确认**：在容器内重跑该脚本后 - `tensorrt_llm/usage/llm_args_golden_manifest.json` 无 diff。本轮只加了 validator 和一个 - 私有方法，没有新增用户可见字段，所以不需要 telemetry/privacy CODEOWNER 审批。 diff --git a/docs/dspark_p0_task_prompt.md b/docs/dspark_p0_task_prompt.md deleted file mode 100644 index a799a05c75f2..000000000000 --- a/docs/dspark_p0_task_prompt.md +++ /dev/null @@ -1,933 +0,0 @@ -# 任务：让 DSpark confidence-head 调度在 DeepSeek-V4-Pro 上跑通（P0 阶段） - -> 本文是给 agent 的任务书。全部路径、commit、权重张量名均已在 `bia` 集群 login -> 节点上实地核实过（2026-08-01）。 - ---- - -## 0. 先读这个（不要跳过，也不要从零设计方案） - -本任务的 **唯一 source of truth** 是 PR NVIDIA/TensorRT-LLM#17056 里的 -`docs/dspark_confidence_schedule_goal.md`（399 行中文目标文档）。 - -它已经完成了全部前期分析：H0–H10 十个 host 侧阻断点、K1–K15 kernel 逐项审计、 -D1/D2/D3 三个方案对比、G1–G8 验收标准、§7 的 U1–U7 未知项、§8 的 9 个待拍板 -问题。**每一条断言都带 `file:line` 锚点。** - -你的工作是**执行它的 P0 阶段**，不是重新做分析。任何与该文档冲突的结论，必须先用 -`file:line` 证明文档错了，再改。 - -### 拉取方式 - -本地 checkout：`/lustre/fsw/coreai_comparch_trtllm/laliao/repos/trtllm-dsv4/TensorRT-LLM` -（`origin` = `github.com/lancelly/TensorRT-LLM`，`upstream` = `github.com/NVIDIA/TensorRT-LLM`） - -```bash -cd /lustre/fsw/coreai_comparch_trtllm/laliao/repos/trtllm-dsv4/TensorRT-LLM -git fetch upstream refs/pull/17056/head:pr-17056 -git checkout pr-17056 # head = d8897046653331bca51242bf068b052568fca4ce -git show pr-17056:docs/dspark_confidence_schedule_goal.md -``` - -PR 的 4 个 commit： - -| commit | 内容 | -|---|---| -| `830afe4e58` | feat: Add DSpark confidence-head verification scheduling（26 files, +4683） | -| `8723506608` | fix: Fix host syncs, slot aliasing and cost model（13 files, +641） | -| `e1c28eeae2` | feat: Make DSpark ragged verification work on the DSv4 kernel path（21 files, +4172） | -| `d889704665` | doc: Scope the ragged-verify measurements to x86 DGX B300 | - ---- - -## 1. 三个必须先纠正的认知 - -### (a) confidence head 已经实现了 - -不要从零写。代码在： - -- `tensorrt_llm/_torch/models/dspark/heads.py`（`DSparkConfidenceHead`） -- `tensorrt_llm/_torch/models/dspark/{draft,attention}.py`、`modeling_dspark.py` -- `tensorrt_llm/_torch/speculative/dspark{,_planner,_verify,_ragged,_schedule,_observability,_sps_profiler}.py` -- 开关：`llm_args.py` 的 `enable_confidence_scheduling` / `enable_ragged_verify` - （均默认 `False`，`status='prototype'`） - -任务是**让它正确、且 CUDA-graph 兼容地跑起来**，不是新增功能。 - -### (b) goal doc §1.3 描述的「11 个未提交的工作区改动」现在已经是 commit `e1c28eeae2` - -文件清单与该 commit 逐项吻合。所以文档里标为「新引入的 P0 bug」的 **H3** 和 **H4** -现在是**已经进了 PR head 的 bug**： - -- **H3**：`dsa.py:1673-1695`，`row_kv_lens.add_(expanded.to(torch.int32))` —— CUDA - tensor 加 CPU tensor（`kv_lens = cached_token_lens + seq_lens_kv` 两项都在 CPU）， - 首个 ragged step 必抛异常。修法：把逐行修正量折进 `row_kv_lens_host` 做**一次** - H2D，复用相邻 `kv_lens_expanded_host` 已有模式（`dsa.py:1590-1592`）。 -- **H4**：`row_kv_lens_cuda` 从不被 `on_update_kv_lens` 刷新。 - -另外注意：本地 checkout 的工作区是干净的，goal doc §1.3 描述的是**另一个**工作区 -的状态，不要去找那 11 个「未提交」的文件。 - -### (c) PR 的外部引用有误（goal doc §7.2 已指出） - -找参考实现时用正确的地址： - -- **SGLang 没有 `dspark_schedule.py`**。真实布局是 - `python/sglang/srt/speculative/ragged_verify.py` + - `python/sglang/srt/speculative/dspark_components/*`。 - 三模式 `static` / `cap-accept` / `compact` 出自这里。 -- **vLLM 是 PR vllm-project/vllm#47808**。关键参考点： - - `vllm/v1/attention/backends/mla/indexer.py` 里关于 `q.view(batch, next_n, ...)` - 对非均匀 decode 长度**静默算错**的注释 - - 它如何使用 DeepGEMM 的 varlen `indices` 参数（TRT-LLM 目前用 "expanded - buffers" 展平达到同样形状，从不传 `indices`；见 goal doc §7.3） -- `align_verify_tokens_to_graph_tier` 这个符号全网不存在。SGLang 的实际原语是 - `round_up_grid()` / `compute_target_verify_graph_key()`。 - ---- - -## 2. 目标配置 - -- **模型**：DeepSeek-V4-Pro-DSpark - （target = DSv4 Pro + DSA 稀疏注意力；drafter = DSpark block drafter + confidence head） -- **并行**：DEP8 = TP8 + `moe_expert_parallel_size=8` + `enable_attention_dp=True` -- **投机**：`max_draft_len=5` → block K=5 → 均匀 verify 窗口 `next_n=6` -- **硬件**：8×Blackwell（B200/B300，SM100） - -### 2.1 权重（已实地确认存在且可读） - -``` -/lustre/share/coreai_dlalgo_ci/artifacts/model/deepseek-ai_deepseek-v4-pro-dspark/hf/hf-cd80265_orig -``` - -832 GB，66 shards，149782 个 tensor。路径 world-readable。 - -**决定性证据 —— confidence head 权重就在里面：** - -``` -mtp.2.confidence_head.proj.weight BF16 [1, 7680] (model-00066-of-00066.safetensors) -``` - -checkpoint 结构：`mtp.0` / `mtp.1` / `mtp.2` 三个 DSpark block，只有 `mtp.2` 带 -confidence head。config 关键字段： - -| key | 值 | -|---|---| -| `dspark_block_size` | **5**（与 `max_draft_len=5` 对上） | -| `dspark_target_layer_ids` | `[58, 59, 60]` | -| `dspark_markov_rank` | 512 | -| `dspark_noise_token_id` | 128799 | -| `expert_dtype` | fp4 | -| `index_topk` | **1024**（普通 dsv4-pro 是 512） | -| `num_nextn_predict_layers` | 1 | - -**全 index grep `sts|temperature|calib` 零命中** —— 印证 goal doc 的 U4：STS 校准表 -既不在 checkpoint 也不在 repo，必须自己 profile 产出。 - -### 2.2 必须先建软链 - -测试里是 `MODEL_PATH = f"{llm_models_root()}/DeepSeek-V4-Pro-DSpark"`，而实际目录名 -是 `hf-cd80265_orig` 且嵌了两层： - -```bash -mkdir -p ~/llm-models -ln -s /lustre/share/coreai_dlalgo_ci/artifacts/model/deepseek-ai_deepseek-v4-pro-dspark/hf/hf-cd80265_orig \ - ~/llm-models/DeepSeek-V4-Pro-DSpark -export LLM_MODELS_ROOT=~/llm-models -``` - -`MODEL_PATH` 在 class body（import 时）求值，`LLM_MODELS_ROOT` 必须在 pytest 启动 -**前**导出，否则整个 module 收集失败。 - -容器里 `~` 可能不是同一个 home，建议软链建在 `/lustre/fsw/coreai_comparch_trtllm/laliao/llm-models/` -下并用绝对路径导出。 - -### 2.3 用不到但容易认错的几个目录 - -| 路径 | 是什么 | 为什么不是你要的 | -|---|---|---| -| `/lustre/share/coreai_comparch_trtllm/dsv4-pro` | 806 GB，纯 DSv4-Pro | 只有 `mtp.0`，**无 confidence head** | -| `/lustre/share/coreai_comparch_trtllm/dsv4-flash` | 149 GB | DSv4-Flash，非本轮范围 | -| `/lustre/share/coreai_comparch_aarwlt/hf_repos/deepseek-ai/DeepSeek-V4-{Pro,Flash}` | HF 原始副本 | 同上，无 DSpark | -| `/lustre/share/coreai_comparch_trtllm/llm-models/` | **空目录** | 不要指向这里 | - -### 2.4 ⚠️ 与现有测试的冲突（必须先解决） - -现有测试 `TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling` -（`tests/integration/defs/accuracy/test_llm_api_pytorch.py`）用的是 -**`disable_overlap_scheduler=True`**。 - -但本任务**要求 overlap scheduler 开启**。这不是配置项调换，而是把 goal doc 的 -**H0 从边缘情况变成第一优先级**： - -> **H0**：`disable_overlap_scheduler=True` 时，ragged token 布局**根本不会被产出**。 -> `_prepare_tp_inputs` 的非 overlap 分支（`model_engine.py:4606`）条件含 -> `next_draft_tokens_device is None`，该分支 `sequence_lengths.append(1 + num_draft_tokens)`， -> **完全不看 `py_verify_len`**（`:4616,4625`）。后果：`spec_metadata.verify_lens` / -> `qo_indptr` / `total_verify_tokens` 和 graph key 都在描述一个 ragged 布局，而真实 -> `input_ids`/`seq_lens` 是均匀的 `1+top_tier`，`_accept_draft_tokens` 随后按 -> `total_verify_tokens` 切片（`interface.py:1436-1441`）→ **逐请求 token 错位**。 - -goal doc §3.3 明确要求 **H0 与 H5 必须同时修**（H5：CUDA graph capture 的 token 轴 -来自 `key[1]` 而不看 bucket，`cuda_graph_runner.py:446-453`），否则 graph 只会更糟。 - -**默认做法**：新增一个 `disable_overlap_scheduler=False` 的 test case，保留原 case -作对照。若要直接改原 case，先确认。 - -### 2.5 已经拍板的两个决策（goal doc §8 的 Q2 / Q8 已关闭） - -goal doc §8 把这两条列为待拍板，**现在已定，不要再当开放问题讨论**。 - -#### 决策一：**ragged path 是第一优先级**（关闭 Q2） - -Q2 原本问的是：本轮交付「均匀 tier 阶梯」还是「ragged」。两者的区别是—— - -| | **均匀 tier 阶梯（D3 Stage A）** | **ragged（D1/D2）** | -|---|---|---| -| 语义 | 每个 decode step **整个 batch 用同一个 verify 长度**，从 tier ladder `[1,3,5]` 里选一个 | **每个请求各自的 verify 窗口**，同一步里可以是 `[6,3,5,2,...]` | -| kernel | 0 改动。四个 uniform kernel 在 `num_gen_tokens = n·(K+1)` 时逐一精确正确 | 需要 kernel 支持每 token 行向量；K4/K9 已完成 | -| 收益 | 拿满 SPS(B) 项，**损失 τ 项** | 完整 Θ = τ·SPS(B) | -| 风险 | 中。无 K2 假设 | **高**。押在 K2（trtllm-gen sparse-MLA generation）这个读码无法证明的行为上 | - -> **后续（见 §10.9）**：均匀 tier 阶梯不只是「没选它」，而是**已从代码里整个删除**。 -> 它在实测中退化成「验证整块」，却仍占着一条独立代码路径。下表保留作决策记录。 - -**已定：做 ragged。** 因此： - -1. goal doc §5.1 里划给 **P2 的 H6 / H7 / H9 / H10 全部上提到 P0**（见 §3 表二）。 -2. **K2 从「以后再说」变成本轮最大风险**。它的确切阻断点已核实： - `cpp/tensorrt_llm/thop/attentionOp.cpp:920-921` - ```cpp - TLLM_CHECK_WITH_INFO(num_tokens % num_seqs == 0, - "seq_len should be same for all generation requests, num_tokens=%d, num_seqs=%d", - num_tokens, num_seqs); - int32_t const input_seq_length = num_tokens / num_seqs; - ``` - 加上 `mMaxSeqLenQ = acc_q_len / batch_beam`（`attentionOp.cpp:1161,1167`），以及 - `cumSeqLensQPtr` 字段虽存在（`fmhaRunnerParams.h:279`）但**在 generation 分支从不 - 赋值**。D1 和 D2 都靠把 generation 半边**呈现成 `num_gen_tokens` 个「长度为 1 的 - 序列」**来让这个检查平凡通过（`seq_len=1`），二者押的是同一个赌注。 -3. **因此 B2 微基准从 P1 上提为 P0 的 gating 测量**，而且要**最先跑**：它单卡、 - 分钟级、**不需要模型权重**，却决定整条 ragged 路线可不可行。见 §3 表三。 - -#### 决策二：**DSv4-Pro 走 KVCacheManager v2**（关闭 Q8 / U6） - -已确认。后果： - -- **H6 上提为 P0**（原 P1）。 -- 修法不变且与 manager 版本无关：`SpecSamplerBase._verified_len` **实时**读请求上的 - `py_verify_len`（`spec_sampler_base.py:195-196`），而第 N 步的 - `_handle_dynamic_draft_len`（`py_executor.py:4652`）先于第 N-1 步的 - `_update_requests`（`py_executor.py:4746`）执行 → 读到被下一步覆盖过的值。 - `runtime_draft_len` 已经为此快照进 `SampleStateSpec`（`spec_sampler_base.py:218`）， - **把 `py_verify_len` 一并快照进去即可**。 -- ⚠️ **一处需要你自己核实的措辞**：goal doc 说 v2 的补偿项被 `if self.is_draft:` - 门控、「target manager 无补偿」所以会泄漏 KV。但该分支紧邻的注释明确写着 - *"Target managers do not allocate this reserve slack."* - （`kv_cache_manager_v2.py:3538-3546`）——即 target 本来就没分配这块 slack，不补偿 - 是**正确的**。所以「泄漏」这个结论还差一步论证。 - **`py_verify_len` 读到陈旧值这个 bug 本身独立成立**，先按这个修；「是否额外泄漏」 - 作为一个独立问题实测确认，标注「待验证」。 - ---- - -## 3. 执行范围：P0 阶段 - -因为 ragged 是第一优先级（§2.5 决策一），P0 的范围比 goal doc §5.1 原本划的要大： -它 = 原 P0（地基）+ 原 P2 的 ragged 必需项 + 原 P1 的 B2 测量。 - -**执行顺序有讲究，不要并行乱做：** - -``` -表三 B2 微基准（单卡，无需权重） ──┐ - ├──► 表一 地基（H0–H5, K8, sync, H8, K5/K6/K13） - │ └──► 均匀 tier 路径跑通 = 中间里程碑 - └──► 表二 ragged 必需项（H6, H7, H9, H10） - └──► ragged 路径跑通 = 交付 -``` - -- **B2 最先跑**，它可能直接否掉整条 ragged 路线（§2.5 决策一第 3 点）。 -- **表一必须先于表二**。表一全部是「无论 ragged 与否都要修」的东西，其中 H1 只有 - 约 3 行 Python 却能完整解释唯一一次实测 hang。在均匀 tier 路径跑通之前调 ragged， - 等于在两层未知上 debug。 -- **均匀 tier 路径不是可以砍掉的中间产物**：G2 的逐 token 等价、A2 的 `cap-accept` - 差分都要拿它当参照系（§6）。 - -### 表一：地基（全部 host 侧，独立于 kernel 方案选择） - -| 项 | 内容 | 锚点 | -|---|---|---| -| **H1** | DSA 展开 stride 用 runtime tier 而非静态 max。**约 3 行 Python**（`stride = spec_metadata.runtime_tokens_per_gen_step`）。⚠️ **buffer 容量必须保持静态 max**，只让 stride 变，否则 capture 期重分配 → 第二类 hang | `dsa.py:662,696,1401`、`model_engine.py:342` | -| **H2** | `_attach_ragged_verify_layout` 移到 `attn_metadata.prepare()` **之前**（它只需要 `generation_requests`，已在作用域内） | `model_engine.py:5326` vs `:5412` | -| **H3** | `row_kv_lens` device 不匹配。修法：折进 `row_kv_lens_host` + 单次 H2D | `dsa.py:1673-1695` | -| **H4** | `on_update_kv_lens` 补刷 `row_kv_lens_cuda` | `dsa.py:1608`、`:902-977` | -| **H0+H5** | **必须同时改**。dummy 分支用 `get_request_tokens_per_gen_step`；`capture()` 吃 bucket | `model_engine.py:4606`、`cuda_graph_runner.py:446-453` | -| **K8** | `prepare_scheduler_metadata` 改走 `expand_per_gen_token`（现在硬编码 `num_generations * (1 + max_draft_tokens)`） | `dsa.py:2329-2330` | -| **sync** | `_update_target_input_tensors` 的 `repeat_interleave` 补 `output_size=`（旁边 `model_engine.py:5101-5109` 已做对）；`apply_sts` 缓存 host 副本 | `model_engine.py:3910-3919`、`heads.py:280-283` | -| **H8** | config 层校验 `cuda_graph_config.enable_padding=True`（默认 False，`fit_ragged_verify_lens` 却假定已开） | `model_engine.py:3644-3648`、`llm_args.py:180-184` | -| **K5/K6/K13** | 显式 **config-reject** 而非静默降级：`enable_heuristic_topk` / `use_cute_dsl_topk` 与 ragged 结构性互斥（kernel 层已有 TLLM_CHECK，Python config 校验欠缺）；K13 加显式 assert 而非依赖 SM 判断 | `indexerTopK.cu:918-922`、`trtllm.py:999-1001` | - -**表一交付（中间里程碑）**：均匀 tier 路径（`enable_confidence_scheduling=True`、 -`enable_ragged_verify=False`）在 8×B200/B300 上跑通 GSM8K，hang 消失。 - -### 表二：ragged 必需项（原 goal doc P2，因决策一上提到 P0） - -| 项 | 内容 | 锚点 | -|---|---|---| -| **H6** | `py_verify_len` 快照进 `SampleStateSpec`，不要实时读。**KVCacheManager v2 已确认**，见 §2.5 决策二（含一处需自行核实的措辞） | `spec_sampler_base.py:195-196,218`、`py_executor.py:4652,4746`、`kv_cache_manager_v2.py:3538-3546` | -| **H7** | ragged batch 上 rejection sampling 被 fail-closed guard **静默关掉**：`_rejection_buffers_valid` 要求 `logits_rows >= num_contexts + num_gens*(draft_len+1)`，而 ragged 下 target 只发 `num_contexts + total_verify_tokens` 行 —— **调度真的在 trim 时该式必然不成立**。于是 PR 在 rejection 路径里精心加的 ragged 分支（`interface.py:1793-1806`）成了死代码。**GSM8K temperature=0 检测不到这一点。** 要么放宽 guard 成 `num_contexts + total_verify_tokens`，要么在 config 层显式拒绝「ragged + 非贪婪」。**静默降级是最差的选项** | `interface.py:1707-1710,1793-1806` | -| **H9** | **rank-local 早退门控了条件性的 `tp_allgather`，是 hang 风险。** `decide_verify_lens` 的 `all_rank_max` 在早退之后，`peer_stats` allgather 被 `if ragged_lens is not None` 门控，fallback 分支又发一次 → 走 ragged 的 rank 发 2 次集合通信、fallback 的 rank 发 1 或 3 次 → **死锁**。且早退条件是 rank-local 且**时序相关**（`_copy_event.query()`）。修法：每个 rank **无条件**发出相同序列的集合通信，再基于归约结果分支 | `dspark_verify.py:234-241,256-259,132-133`、`py_executor.py:3258-3261` | -| **H10** | DSpark drafter 用固定 stride 取 target hidden state：`base = gen_start + arange_g * (runtime_draft_len + 1)`，而 ragged 下真实偏移是 `spec_metadata.qo_indptr[r]`（`dspark.py` 全文无 `qo_indptr` 引用）。**静默：只是接受率崩塌，不报错**。注意 `8723506608` 已把它从 `max_draft_len+1` 修成 `runtime_draft_len+1`，**均匀路径是对的**，只有 ragged 错 | `dspark.py:600,624` | -| **§4.4** | **ADP 资格门必须同时比较 bucket**。现在只比 `is_all_gen_only` 和 batch size，从不比 token 数；`all_rank_num_tokens` 是 host list、会被**烤进 captured graph**，所以「任一 rank 退回均匀路径」造成的不匹配在 replay 时**无法被察觉** | `cuda_graph_runner.py:362-372`、`model_engine.py:2987-2998` | - -**表二交付**：ragged 路径（两个开关都开）GSM8K + 逐 token 等价（§6）。 - -### 表三：B2 微基准（最先跑，单卡，无需模型权重） - -**这是决定整条 ragged 路线可不可行的那个数字**（goal doc §7.1 U1）： - -> 直接调 `thop.attention`，在**相同总 token、相同逐 token 稀疏索引集**下，比较 -> `(batch=B, s_q=tier+1)` 与 `(batch=B*(tier+1), s_q=1)` 两种呈现的**正确性与吞吐**。 - -- 读不出来（`selectMlaGenerationKernel` 在预编译 cubin 里），但一张 B200/B300 上 - 直接测就行，分钟级。 -- 关注点：DSv4-Pro 在 DEP8 下每 rank 保有全部 128 个 Q head，BMM1 的 M 维从 768 降到 - 128；由于稀疏索引集逐行不同，本来也没有跨 6 行共享 KV tile 的机会 —— 但这是 - **推断，不是测量**。 -- **若 token-major 显著变慢或不正确**：立刻停下来报告，不要硬推 ragged。 - goal doc §5.1 的兜底是 D3 Stage A（均匀 tier 阶梯），它拿满 SPS(B) 项、无 K2 假设、 - 还能把 graph 数从 102 压到 ~54。 - -顺带把 goal doc §6.4 的另两个也跑了（都是分钟级）： - -- **B1**：修完 H1 后重跑均匀 tier capture，确认 hang 消失。**带 `TORCH_LOGS=recompiles`** - —— 这是区分 §7 三个 hang 候选根因的最短路径。 -- **B3**：expanded 布局下 paged-MQA-logits 的吞吐（展平路径每请求重读 6 次 KV 而非 - 2 次，在关键路径上，必须实测）。 - -### 仍然不做的 - -- D1 vs D2 的最终选型：等 B2 出数再定。工作区已沿 D2 走了一半（K4/K9 已实现且设计 - 正确 —— 可选参数、非 ragged 路径 bit-identical、TORCH_CHECK 替换而非放宽）， - **这部分不要丢**，它对任何 ragged 方案都是必需的，对 D3 也无害（不传参数即退化）。 -- STS 校准表与 SPS cost table 的产出流程（goal doc G8 / U4 / U5）—— 见 §8 交付物第 5 条。 - ---- - -## 4. 编译 - -**不要手工敲 `build_wheel.py`。**用仓库里现成的脚本： - -```bash -/lustre/fsw/coreai_comparch_trtllm/laliao/repos/trtllm-dsv4/TensorRT-LLM/tmp/compile.sh -``` - -它做的事（`tmp/compile.slurm`）： - -- `sbatch --partition=batch --nodes=1`，`--time=4:00:00`， - `--account=coreai_comparch_trtllm` -- 容器镜像从 `jenkins/current_image_tags.properties` 的 `LLM_DOCKER_IMAGE` 读， - URL 里的 `urm.nvidia.com/` 被换成 `urm.nvidia.com#` 供 enroot 使用 -- mount：`/lustre:/lustre` + repo → `/code/tensorrt_llm`，workdir `/code/tensorrt_llm` -- 构建目录按 **git short hash** 命名：`/lustre/fsw/coreai_comparch_trtllm/laliao/build//` -- 实际命令：`git clean -fdx cpp/` → 删 `tensorrt_llm/flash_mla` → - `python scripts/build_wheel.py --clean -G Ninja -a '103-real;'` → `pip install` -- 产物：`build//trtllm.sqsh`（`--container-save`）+ `build//tensorrt_llm-*.whl` -- 日志：`tmp/logs/compile_.{log,err}` - -**注意事项：** - -1. **`-a '103-real;'` 只编 SM103（B300）。** 如果你拿到的是 B200（SM100），必须改成 - `'100-real;'` 或 `'100-real;103-real;'`，否则 kernel 加载会失败。 -2. **构建目录以 git hash 为键，并且会先 `rm -r` 掉旧的。** 两个同 hash 的 compile - job 并发跑会互相清目录 —— 见 §5 的竞态提交规则。 -3. Python-only 改动（本任务 P0 清单里除 K5/K6/K13 的 C++ assert 外**全部是 Python**） - 不需要重编。直接用已有的 `trtllm.sqsh` + mount repo 到 `/code/tensorrt_llm` 即可， - 因为 `PYTHONPATH=/code/tensorrt_llm` 会让改动立即生效（见 `tmp/serve_*.slurm` 的写法）。 - **先判断要不要编，再编。** -4. 跑测试时用 `CONT="/lustre/fsw/coreai_comparch_trtllm/laliao/build//trtllm.sqsh"`， - 参照 `tmp/serve_dsv4_bs8_relaxed_constraint.slurm` 的 srun 写法 - （`--mpi=pmix -N 1 --ntasks=8 --ntasks-per-node=8`）。 - ---- - -## 5. 申请节点：所有账号一起拍，没排到的取消 - -### 5.1 集群事实（实测） - -| 项 | 值 | -|---|---| -| 集群 | `bia`（login-bia02.bia.clusters.nvidia.com），Slurm 25.11.2 | -| 节点 | 128 × B300，`bia0001-0128`，x86_64，2 socket × 64 core（256 线程），~2.0 TiB 内存 | -| **每节点 8 GPU** | 由 `tmp/serve_*.slurm` 的 `--ntasks-per-node=8` 佐证；Slurm **未配 GRES**，不能按卡申请，`OverSubscribe=EXCLUSIVE` 整节点独占 | -| 分区 | `batch`（默认，默认 2h / 最大 **5h**，PriorityTier 2）
`backfill`（默认 2h / 最大 **8h**，PriorityTier 1，**`PreemptMode=CANCEL` 会被直接取消**，GraceTime 600s） | -| 你的账号 | `coreai_comparch_trtllm`、`coreai_comparch_aarwlt` | -| 配额 | 单用户 **104 节点**（QOS `MaxTRESPU`），最多 1000 个排队作业 | -| 当前负载 | 112/128 alloc、0 idle、303 个 pending —— **要排队** | - -DEP8 需要 8 GPU，**1 个节点即可**。 - -### 5.2 竞态提交 → 先跑的赢，其余取消 - -四个组合全拍上：`{coreai_comparch_trtllm, coreai_comparch_aarwlt} × {batch, backfill}`。 -`sbatch` 命令行的 `--account` / `--partition` **会覆盖**脚本内的 `#SBATCH` 指令。 - -**推荐做法：竞态的是「分配」，不是「负载」。** 先抢一个空 allocation，抢到之后再往 -里塞真正的活。这样不存在两个 job 同时开跑把对方的 build 目录清掉的问题。 - -```bash -#!/bin/bash -# race_alloc.sh —— 抢一个 1 节点 allocation，先到先得，其余取消 -set -uo pipefail - -ACCOUNTS=(coreai_comparch_trtllm coreai_comparch_aarwlt) -PARTITIONS=(batch backfill) -TIME_LIMIT=5:00:00 # batch 上限 5h；backfill 可以到 8h -LOGDIR=/lustre/fsw/coreai_comparch_trtllm/laliao/repos/trtllm-dsv4/TensorRT-LLM/tmp/logs -mkdir -p "$LOGDIR" - -JOBS=() -for A in "${ACCOUNTS[@]}"; do - for P in "${PARTITIONS[@]}"; do - T="$TIME_LIMIT" - [[ "$P" == backfill ]] && T=8:00:00 - JID=$(sbatch --parsable \ - --account="$A" --partition="$P" \ - --nodes=1 --exclusive --time="$T" \ - --job-name="race-${A}-${P}" \ - --output="$LOGDIR/race_%j.log" --error="$LOGDIR/race_%j.err" \ - --wrap="sleep infinity") - echo "submitted $JID account=$A partition=$P time=$T" - JOBS+=("$JID") - done -done - -# 轮询，谁先 RUNNING 谁赢 -WINNER="" -while [[ -z "$WINNER" ]]; do - for J in "${JOBS[@]}"; do - ST=$(squeue -h -j "$J" -o "%T" 2>/dev/null) - if [[ "$ST" == RUNNING ]]; then WINNER="$J"; break; fi - done - [[ -z "$WINNER" ]] && sleep 10 -done - -NODE=$(squeue -h -j "$WINNER" -o "%N") -echo "WINNER=$WINNER on $NODE" - -# 只取消这批里没排到的，不要动别人/别的作业 -for J in "${JOBS[@]}"; do - [[ "$J" == "$WINNER" ]] && continue - scancel "$J" && echo "cancelled $J" -done - -echo "export WINNER_JOBID=$WINNER" -``` - -拿到 allocation 之后，把活塞进去（`--overlap` 是必须的，因为 `sleep infinity` 那个 -step 已经占着资源）： - -```bash -srun --jobid=$WINNER_JOBID --overlap \ - --container-image="$CONT" \ - --container-mounts="/lustre:/lustre,$REPO:/code/tensorrt_llm" \ - --container-workdir=/code/tensorrt_llm \ - -N1 --ntasks=1 --pty bash # 交互式排查 - -# 或者跑测试（8 卡） -srun --jobid=$WINNER_JOBID --overlap --mpi=pmix \ - --container-image="$CONT" \ - --container-mounts="/lustre:/lustre,$REPO:/code/tensorrt_llm" \ - --container-workdir=/code/tensorrt_llm \ - -N1 --ntasks=8 --ntasks-per-node=8 \ - bash -c "..." -``` - -活干完后**记得 `scancel $WINNER_JOBID`**，`sleep infinity` 不会自己退。 - -### 5.3 直接竞态提交真实负载（更简单，但有坑） - -如果你就想直接把 `compile.sh` 那种脚本拍四份： - -```bash -for A in coreai_comparch_trtllm coreai_comparch_aarwlt; do - for P in batch backfill; do - sbatch --parsable --account="$A" --partition="$P" \ - --job-name="compile-${A}-${P}" tmp/compile.slurm - done -done -``` - -⚠️ **`compile.slurm` 的 `BUILD_DIR` 只按 git hash 命名，且开头就 `rm -r`。** 四个 job -在同一 commit 上跑，谁先起来谁就把别人的目录删了。所以： - -- 轮询间隔要短（≤5s），**一发现有 job 离开 PENDING 就立刻取消其余的** -- 如果不幸有两个同时起来了，取消其一之后**必须重跑赢家**，不能假定它的 build 是完整的 -- 更稳的做法就是 §5.2：竞态 allocation，然后只在赢家里跑一次编译 - -### 5.4 收尾 - -```bash -squeue -u $USER # 看还剩什么 -scancel [ ...] # 精确取消 -``` - -**不要用 `scancel -u $USER`** 无差别清场 —— 可能会杀掉正在跑的别的活。 - ---- - -## 6. 验收（goal doc §6） - -> **所有测试都必须在 Slurm 分配的 GPU 节点上跑。login node 不可用于任何 GPU 测试。** - -**必须做的，不只是 GSM8K：** - -- **G3 GSM8K**：`score >= 96.0`。但 goal doc §6.1 明确说这个门**不够** —— 统计阈值 - 允许 42/1319 题翻转，而 `>=96.0` 对实测 96.475 只有 6.3 题余量，既太松又太紧。 - **且当前没有活的 baseline 可比**（兄弟测试 `test_gsm8k_dep8_megamoe_deepgemm` 已被 - `waives.txt:184` 以 nvbugs/6506920 跳过）。 -- **G2 逐 token 等价（这才是真正的门）**：temperature=0 下，`static` 与开启调度后 - 逐 token 输出一致。贪婪投机解码定义上 lossless，任何分歧都是 bug。 - ⚠️ 诚实的形式：跨不同 graph 形状的严格 bitwise 相等 TRT-LLM 并不保证（MoE/attention - 规约顺序会变），所以允许极少数 prompt 分歧并设硬上限。 -- **A2 `cap-accept` 差分（ragged 成为第一优先级后，这条从可选变成强烈建议）**： - 实现 SGLang 的三模式 `static` / `cap-accept` / `compact`。`cap-accept` 跑**均匀 - kernel 路径**（`next_n=6`）但只提交每请求窗口内的 token，输出应与 `compact` - **逐 token 相同**。于是： - - `cap-accept` ≠ `compact` ⇒ **必然是 ragged kernel bug** - - `static` ≠ `cap-accept` ⇒ **必然是调度/接受逻辑 bug** - - 这是把「GSM8K 掉分了但不知道为什么」变成可定位诊断的唯一手段。 - `dspark_observability.py` 已有 mode 枚举，`cap-accept` 目前显式 `NotImplementedError` - （`dspark.py:471-478`）。 -- **A3 布局一致性断言**（一条同时抓 H0/H5/H8）： - ``` - sum(1 + py_verify_len over padded generation_requests) - == attn_metadata.num_tokens - num_ctx_tokens - == spec_metadata.total_verify_tokens - == key[bucket] - ``` -- **A4 时序断言**（抓 H2/H4）：断言 `attn_metadata.ragged_verify_lens` 在 `prepare()` - 那一刻是**本步**的（而非上一步）。 -- **A6 调度确实生效**（抓静默退化）：断言 `planner.stats` 的 `fallback_flat_cost` / - `fallback_no_snapshot` / `fallback_no_confidence` / `fallback_short_snapshot` - （`dspark_verify.py:86-93`）显示 trim 路径被走到。 - **没有 profiled SPS cost table 时 `_decide_local` 无条件返回 `max_tier` - （`dspark_verify.py:268-270`），feature 在构造上无法 trim** —— 一个看起来完全正常的 - run 收益可能是零。 -- **G4 CUDA graph**：稳态 decode 100% 命中 replay；capture 期无 D2H sync、无新分配； - graph key 与实际 token layout 逐字段一致。 -- **G5 ADP 一致性**：goal doc §4.4 要求 **ADP 资格门必须同时比较 bucket**（现在只比 - `is_all_gen_only` 和 batch size，从不比 token 数）。 -- **G6 零回归**：两个开关关闭时与 `main` bit-identical。 -- **单测**：`tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py` - 按每请求 q 长度参数化 —— goal doc §6.3 称这是**性价比最高的一个测试**，单卡分钟级 - 就能抓住 tier-3 那类 hang（该文件 `:415`/`:473` 现在硬编码 - `seq_len_q = fused_q.shape[0] // num_requests`，那正是被测假设本身）。 - -集成测试命令（无需 `mpirun`，单 pytest 进程经 mpi4py.futures 起 8 worker）： - -```bash -export LLM_MODELS_ROOT=/lustre/fsw/coreai_comparch_trtllm/laliao/llm-models -python -m pytest "tests/integration/defs/accuracy/test_llm_api_pytorch.py::\ -TestDeepSeekV4ProDSpark::test_gsm8k_dep8_megamoe_deepgemm_confidence_scheduling" -vs -``` - -**验收必须覆盖至少两个并发点**（goal doc §6.5）：调度只在高并发下才真的 trim，单一 -低并发 run 是空跑。vLLM #47808 报的是并发 16 和 64。 - ---- - -## 7. 已知 hang（必须区分三个候选根因，不要认定其中之一） - -唯一一次端到端尝试失败在 tier-3 的 CUDA graph capture： - -``` -capturing 102 graphs (34 batch sizes x 3 draft-length tiers [1, 3, 5]) -warmup batch=128, draft_len=5 ← 完成 -warmup batch=128, draft_len=3 ← never returns TEST_EXIT=124 -``` - -三个**独立**候选（goal doc §7.1 U2）： - -1. **H1 stride 不匹配** —— tier 5 时 768==768 通过，tier 3 时 schedule 描述 768 行而 - logits 只有 512 行。完整自洽。 -2. **capture 期 attention workspace 重分配** —— hang 前紧邻的日志是 - `Attention workspace ... increase from 0 bytes to 168493568 bytes`； - `model_engine.py:1789-1806` 明确记录这是真实 hazard，且 Case 2b **没有** Case 2 的 - `max_spec_graph` 规避（`:1810-1813`）。 -3. **torch.compile / max-autotune 在 capture 区内**（goal doc §4.5）—— - `deepseek_v4.py` 三个 `@maybe_compile(dynamic=True, max-autotune)` 函数经 - `on_update_kv_lens` 在 `CUDAGraphRunner.capture` 录制区内被调用； - `_compute_gen_compressed_position_ids` 把 `num_gen_tokens_per_seq` 当 Python int 用 - （`:979-983`）构成 guard → **每个 tier 一次 Inductor 编译，发生在 graph capture - 期间**，而 max-autotune 需要同步 benchmark。 - -**先修 H1 再重跑，并带 `TORCH_LOGS=recompiles`** —— 这是区分三者的最短路径。 - -goal doc 引用的 `.dspark-logs/gsm8k.log` 和 `hang-stack-299685.txt` 不在本 repo 里， -是写文档那位的工作区产物。不要去找，重新复现即可。 - ---- - -## 8. 交付物 - -0. **B2 的实测数字**（最先交，可能改变后续全部计划）：`(batch=B, s_q=tier+1)` vs - `(batch=B*(tier+1), s_q=1)` 的正确性与吞吐对比。若 token-major 不可行，立刻报告 - 并改推 D3 Stage A。 -1. 一个基于 `pr-17056` 的分支，表一 + 表二逐项修完，每项 commit message 里引用 - goal doc 的编号（H0/H1/.../K8） -2. A3/A4/A6 断言 + 参数化的 `test_deepseek_v4_sparse_mla.py` 单测；compressor ragged - 差分单测（`new_tokens_per_seq` 全填 `next_n` 必须与传 `None` bitwise 相同 —— goal - doc §6.3 标为「欠缺」） -3. GSM8K 在 **overlap scheduler + DEP8 + ragged** 下的实测分数，**以及 planner.stats - 的 fallback 计数**（证明调度真的生效，而不是静默退化成 max_tier）；均匀 tier 路径 - 的同口径分数作为对照 -4. 一份结论：hang 的真实根因是 §7 三个候选中的哪个（附证据） -5. **对 goal doc §8 剩余的待拍板问题给出你的推荐**（Q2 和 Q8 已关闭，见 §2.5），尤其： - - **Q4**：`cap-accept` 模式要不要实现。**建议实现** —— ragged 成为第一优先级之后 - 它的价值更高了：`cap-accept` vs `compact` 的差异 ⇒ 必然是 ragged kernel bug； - `static` vs `cap-accept` 的差异 ⇒ 必然是调度/接受逻辑 bug。把 pass/fail 硬币变成 - 可定位诊断。成本不高（复用已有的 `torch.minimum(..., verify_lens)` clamp， - `interface.py:1473-1480`、`:1889-1896`）。当前是显式 `NotImplementedError` - （`dspark.py:471-478`） - - **Q5**：STS 表和 SPS cost table 由谁、按什么流程产出、存哪（checkpoint 内？ - repo `examples/configs/`？）。**在它们存在之前，任何性能结论都是空的，而且 - feature 在构造上无法 trim** - - **Q6**：tier ladder 要不要按 batch size 门控（102 → ~54 个 graph，换回 - ~0.4–0.8 GB/rank 的 KV） - - **Q7**：非贪婪采样的处理时机 —— 见表二 H7 - - **Q9**：要不要采用 DeepGEMM 的 `indices` varlen 参数（goal doc §7.3）。**ragged - 成为第一优先级之后这条值得重新评估**：vLLM #47808 用的正是它，能去掉 - `[num_tokens, max_blocks_per_seq]` 那个大 block table，代价是把该路径钉死在 SM100 - （DeepGEMM varlen 分支断言 `arch_major == 10 and next_n == 1`） - ---- - -## 9. 工作方式 - -- 每个改动都要能指出它对应 goal doc 的哪一条，以及 `file:line` -- 凡是没读码确认的推断，显式标注「待验证」—— 沿用 goal doc 的规范 -- **不要静默降级**。goal doc 反复强调：不支持的组合要显式 config-reject，静默降级是 - 最差的选项 -- 遇到与 goal doc 冲突的发现，先给证据再改结论，不要默默按自己的理解走 -- 提交遵循仓库 `AGENTS.md`：`git commit -s`（DCO），PR 标题格式 - `[JIRA/NVBUG/None][type] description`，不要加 AI co-author - ---- - -## 10. P0 执行结果（2026-08-02 收尾） - -本节由执行 agent 追加，记录实际跑出来的结论。**与前文冲突处以本节为准**，因为前文 -是任务下达时的认知，本节是实测。 - -> **硬件与配置口径:全部实测都在 x86 + NVIDIA B300 SXM6 AC(275040 MiB / 卡, -> 驱动 580.167.08)、8 卡单节点、DEP8(TP8 + EP8 + attention-DP)上取得。MoE -> backend 为 MEGAMOE_DEEPGEMM 或 TRTLLM,`max_draft_len=5`,tier 阶梯 `{1,3,5}`。 -> -> **换到 GB300(或任何 Grace-Blackwell 一体机)结论可能不同**,以下几组明确依赖 -> 硬件: -> -> * **显存(10.4)**:「ragged 的 CUDA graph 开销 ≈0」是在 267 GiB HBM、 -> `free_gpu_memory_fraction=0.5` 下测的。GB300 的显存容量、NVLink 拓扑和 -> C2C 一致性内存会改变 graph 池与 KV 的相对占比。 -> * **裁剪是否划算(10.3)**:planner 的 argmax 比较「每秒接受 token 数」,分母 -> 是实测的 `T(bs, M)`。不同的算力/带宽比会改变 `theta(M)` 的斜率,从而移动 -> 「裁剪开始划算」的接受率阈值。本文结论是「p≈0.90 时任何批次档位都不该裁」, -> **这个阈值不可跨硬件搬运**。 -> * **SPS cost table 必须在目标硬件上重采**。`sps_real_final.json` 只对 B300 -> 有效,在目标机上重跑 `dspark_sps_profiler` 即可。 -> * **每 rank 批次 = 并发 / DP 数**,本文按 DP=8 计算。 -> -> 与硬件**无关**的部分:10.5 的机制验证(每请求不同 verify len 能否正确运行)、 -> 10.6 的六个 bug 及其修复(都是 host 侧逻辑错误)、10.7 的方法论。这些是正确性 -> 问题,不是性能问题。 - - -### 10.1 核心结论：K1/K2 已解除，纯 Python，不需要编译 - -ragged 原本卡在 `mla_rope_generation`（`mlaKernels.cu:1161`）和 FMHA dispatch -（`thop/attentionOp.cpp:920`）都要求 `acc_q_len % batch_size == 0`——per-request -窗口从构造上违反它。解法是把 generation 半边按 **token-major** 呈现：每个 -generation token 作为一行长度为 1 的序列，`seq_len == 1` 整除任何行数。 - -`git diff pr-17056..HEAD -- cpp/` **为空**。op 的 `mMaxNumRequests` 只是 JIT -warmup 提示，per-step 请求数来自 batch，`getFmhaMultiCtasKvScratchSize` 按 SM 数 -而非 batch 算——从 Python 侧传入静态行数上限即可。 - -### 10.2 精度（G3/G4/G5） - -DEP8 + attention-DP + overlap scheduler + CUDA graph，forced ragged 窗口： - -| MoE backend | exact_match (flexible) | exact_match (strict) | -|---|---|---| -| TRTLLM | 96.4367 ± 0.5106 | 96.5125 ± 0.5053 | -| MEGAMOE_DEEPGEMM | 断言通过（分数未捕获，该轮早于 `-s` 开关） | - -uniform 基线 96.2092。`steps_ragged 81/129`，5 种窗口，`trim_ratio 0.23`。 - -### 10.3 吞吐（G6）：**这个 checkpoint 上没有收益空间** - -planner 在三种工况下都正确地选择不裁剪： - -| 工况 | 每 rank 批次 | accept_len | planner | -|---|---|---|---| -| GSM8K | 小 | 3.30 | 不裁 | -| 合成随机 token | 4–8 | 4.09 | 不裁 | -| 合成随机 token | **64** | 3.73（p≈0.90） | **不裁**，零 fallback | - -DSpark block drafter 在此模型上每 token 接受率约 **90%**。拟合表在该接受率下 -于所有批次档位都判定"验满整块最优"（判据：`bs=64, p≥0.80 → L=5`）。 -**这不是功能失效，是没有机会。** 要看到收益需要接受率显著更低的模型或负载。 - -吞吐实测（128 请求，isl1024/osl512，并发 64，KV 泄漏修复后）： - -| 配置 | tok/s | -|---|---| -| baseline（confidence head 关） | 7032.37 | -| ragged 窗口=5（不裁剪） | 7319.47 | -| rotating（confidence 无关地裁 29%） | 5994.39 | - -**唯一站得住的吞吐结论是 KV 泄漏修复值 +6.5%**（6869.73 → 7319.47，同配置）。 -ragged 相对 baseline 的 +4.1% 不作为结论：单次测量，且从未表征吞吐的 -run-to-run 方差。rotating 的 5994.39 **不得**当作本功能的代价——它是与 -confidence 无关的盲裁，是最坏情况且方向与真实策略相反。 - -### 10.4 显存代价:实测约等于零 - -**本节的前一版是错的,已按实测重写。** 前一版声称「CUDA graph 张数 3×、 -attention workspace 9×、ragged 让显存悬崖来得更早」,其中只有张数是对的, -而张数并不等于显存。 - -实测(同节点、同数据集、同并发,只切换 baseline / ragged 的 yaml,用 -`torch.cuda.memory_snapshot()` 求 CUDA graph 私有池的**绝对值**): - -| max_batch_size | uniform 图数 / 池 | ragged 图数 / 池 | 差 | -|---|---|---|---| -| 128 | 19 / **8.746 GiB** | 57 / **8.740 GiB** | **−0.006** | -| 256 | 21 / **16.527 GiB** | 63 / **16.588 GiB** | **+0.061** | -| 512 | 25 / **33.14 GiB** | 75 / — | — | - -**图数确实是 3×(等于 tier 数),但显存差落在噪声里。** CUDA graph 共享内存池, -多捕的形状复用同样的缓冲区,张数没有转化成显存代价。 - -真正决定显存的是 `max_batch_size` 本身:128 → 256 → 512 对应 -8.7 → 16.5 → 33.1 GiB,与 ragged 无关。 - -前一版还声称「ragged 让 bs=512 OOM」。补做对照后:同配置下 uniform 跑通, -而 ragged 那次是在 **25.10 GiB** 处挂的,**低于 uniform 自己 33.14 GiB 的峰值** -—— 一个更费显存的配置不会在更低点失败。那次的日志里有别的进程占着显存,最可能 -是节点残留(本次实验中已两次遇到:一次是 relaunch 过快导致两个进程各占 134 GiB)。 -一次不可复现、且与其余三个档位矛盾的观察,不足以支撑结论。 - -`attention workspace 68 MiB → 604 MiB` 那条同样**从未做过 uniform 对照**, -在补做之前不应引用。 - -**方法论**:前一版三条断言里,唯一做过对照的那条(张数)是对的,另外两条都是 -单次观察 + 事后编的机制解释。而且第一次测显存用的是「捕获前后取差值」,在多轮 -捕获共用分配器的情况下得出 `+13.9 / +5.5 / −0.8 / +0.7 GiB` 这种自相矛盾的 -数——增量法本身不成立;换成私有池绝对值后三次独立测量复现到小数点后三位。 - -### 10.5 每请求不同的 verify len:已验证 - -这是 P0 的核心能力,也是 K1/K2 那个原始 blocker 要解决的东西。**已由 planner -自主产出,无任何外部强制**(`TLLM_DSPARK_FORCE_VERIFY_LENS` 已从代码中删除): - -``` -steps_ragged: 211/318 同一步内窗口不相同的步数 -distinct_verify_lens: 4 {1: 27375, 2: 5831, 3: 12, 5: 66} -trim_ratio: 0.6064 裁掉 60.6% 的 token 上限 -KV 溢出: 0 OOM: 0 -``` - -分配链路是真实的:confidence head → 每请求 survival 概率 → 成本表算出预算 → -top-k 按 survival 高低分配窗口。置信度高的拿长窗口,容易早被拒的拿短窗口。 - -**这证明了机制,不证明策略有收益。** 该轮用的是刻意造陡的合成成本表 -(`tests/microbenchmarks/dspark_make_steep_sps_table.py`,输出带 `SYNTHETIC: true`), -planner 在对着虚构的成本曲线优化,所以那轮的吞吐数字无意义。真实表下 planner 在 -所有批次档位都选满窗口(见 10.3)。 - -**为什么必须造这张表**:此 checkpoint 接受率约 90%,planner 永远不会自己走到 -裁剪路径。而下面三个 bug **只在真裁剪时可达**——没有这张表,它们会带着绿色的 -测试合入。 - -### 10.6 修掉的真 bug - -1. **KV 泄漏(ragged 独有)**。预留侧每步加 `1 + max(draft_len, reserve)`, - 回退侧只退 `py_rewind_len = verified_len − accepted`;`kv_cache_manager_v2` - 把补偿项 gate 在 `if self.is_draft` 后面,于是 target manager 每步漏 - `draft_len − verified_len`,直到序列需要的 block 超过 `max_seq_len`,报 - `User-provided base page indices is too short`。uniform 下该项恒为 0。 - **定位方式**:窗口=5 vs 窗口轮转、其余全同的一次对照,而非读代码。 - -2. **token-major block table 布局**。展开表用了 `[rows, max_blocks]`,而 op 读 - `[num_pools, num_seqs, 2, max_blocks]`,且 gather 折掉了 `num_seqs` 轴 → - 8 卡 device-side OOB。 - -3. **`host_request_types` / `prompt_lens_cuda` 未行化**。op 按 `num_seqs`(=896) - 索引 128 长的数组。**主动审计 12 个 per-sequence 参数找出来的,不是等它崩。** - -4. **bucket fit 假设 padding 必然发生**。fit 按 padded_bs 推导 bucket 网格并把 - 余量分给尚不存在的 pad 行,但 `_get_padded_batch` 可以拒绝 padding。拒绝时 - 网格对应的行数不存在,填出的总数落在网格外:`graph_miss_keys - {'(193, 5, False, False, True, 449)': 16}`,捕获档位有 192 和 256 却没有 193。 - 修法:fit 前先问 `runner.will_pad_to`,不能 pad 就退回 uniform。 - -5. **跨 rank bucket 不一致**。bucket 被独立推导两次:fit 从 allgather 的 - `peer_stats` 算,graph key 又遍历 `generation_requests`(padding 追加**之后**) - 重新求和。两条独立推导却要求跨 rank 精确相等,一致性是碰巧的。实测 - `peer_shape_mismatch` 不裁剪时 2/265 步、裁剪时 16/318 步,每次让 8 个 rank - 一起掉出 replay。修法:fit 发布商定值,key 直接读——按构造必然一致。 - -6. **lm-eval 分数选取**(与 DSpark 无关的既有 bug)。把 `samples`(1319×100) - 混进平均,报出 **44030.98 并通过了 96.0 的门槛**。按 `metric,filter` 键结构修正。 - -### 10.7 踩过的坑(值得写进流程) - -**探针失效比功能失效更常见。** 本次至少七次「失败」中,探针本身有问题的占多数: - -| 现象 | 真因 | -|---|---| -| `graph_miss_shapes` 为空 | 模块版本偏斜,默认参数把「没记录」伪装成「没发生」 | -| sync 探针零警告 | 探针根本没 arm,我差点报通过 | -| gate 在 ramp 期误杀三轮实验 | floor 数的是总步数,而并发从 0 爬升 | -| 显存增量 `+13.9/+5.5/−0.8/+0.7` | 多轮捕获共用分配器,增量法不成立 | -| 造表后 ragged 仍不亮(两次) | 固定开销压倒可裁剪项;网格覆盖不到工作区间 | -| 预览说会裁、实跑不裁 | 我重推了 planner 的 argmax,与真实实现不符 | - -由此得到三条: - -* **沉默不等于通过。** 任何「没有告警」的结论,必须先证明探针确实生效。 -* **能调用就别重推。** 造表预览改成直接 `import` 真实的 - `budget_argmax_over_uniform_lens` 后一次就对,并能回测历史结果;测试里也用真的 - `DSparkScheduleConfig` 而非 mock。 -* **夹具必须能自检。** 造表工具现在会预测 planner 的选择,并在「任何档位都不裁」 - 时直接 WARNING——失败在生成的那一秒暴露,而不是 20 分钟后。 - -**没有对照的归因一律不成立。** 本次有四条结论因补对照而推翻或修正: - -| 断言 | 结局 | -|---|---| -| ragged graph 显存 3× | 实测 ≈0(8.746 vs 8.740 GiB) | -| ragged 让 bs=512 OOM | 不开 ragged 同样配置也曾跑通;且 ragged 在更低的 25.10 GiB 就挂,与「更费显存」矛盾 | -| `peer_shape_mismatch` 是 bug | 是那道门在正常拦截(但推导重复确实该修) | -| G2 token 等价性失败 | 两次**相同** static run 之间也 15/16 发散——噪声底 = 信号,该测试在此部署上问不出东西 | - -**局部证据不能支撑全局断言。** 我曾报「ragged 未引入新的 host-device sync」, -依据是审计了自己新增的 token-major 代码(那部分确实干净),但动态探针在 -`dspark.py:117/516/551` 和 `deepseek_v4.py:891` 抓到 32 次真同步。静态审计只能 -覆盖「我写的代码」,动态探针才覆盖「这条路径实际执行的一切」。 - -**在功能没满负荷运行时做的验证,适用范围比看起来窄。** 上面 4 和 5 两个 bug 都 -只在真裁剪时可达;我此前核对过「两条 padded_bs 规则一致」,但那是在不裁剪的前提 -下做的,总数恒定完全掩盖了行数这一维。 - -### 10.8 未完成 - -* **动态 sync 证明**(唯一未定论项)。ragged 侧已测得 arm 后 32 次真同步,位置见 - 上;**对照(confidence head 关、其余逐字相同)尚未跑完**,因此「这些是 ragged - 引入的还是 DSpark 既有的」暂无结论。在此之前不应引用任何 sync 相关结论。 -* **`peer_not_gen_only` 类 graph miss**:32/318 步,属正常连续批处理(uniform - 路径同样存在),非缺陷,但可作为后续调度优化的方向。 -* **吞吐收益需要更低接受率的模型或负载**。这是 checkpoint 的属性,非本 PR 可解。 - -### 10.9 删除均匀 tier 阶梯:三状态收敛成两状态 - -调度原本有三种状态,中间那种不值得保留: - -| 配置 | 语义 | 结局 | -|---|---|---| -| `scheduling=False` | 验证整个 drafted block | **保留**,唯一兜底 | -| `scheduling=True, ragged=False` | 整批取同一个 verify 长度(tier `{1,3,5}` 里选) | **已删除** | -| `scheduling=True, ragged=True` | 每请求各自的窗口 | **保留** | - -删除理由有两条,都是实测出来的: - -1. **它做不到调度的本意**。整批一个长度,就无法把长窗口给高置信请求、短窗口给 - 坍塌的请求——而那正是 confidence head 存在的理由。 -2. **它在本 checkpoint 上恒等于「验证整块」**。阶梯只有 `{1,3,5}` 三档,接受率 - ~90% 时 argmax 永远落在 5,即整块。也就是说:一条独立代码路径(进 graph key、 - 进跨 rank allgather、进 planner 的 fallback 计数),运行效果和 `scheduling=False` - 完全一样。 - -SGLang 也没有对应物:它的 scheduler 一律经 top-k 产出每请求窗口,"均匀"只作为 -verify-all 的退化情形出现。 - -删掉的东西:`decide_draft_len` 及其 `_decide_local`、`snap_to_tier`(唯一调用者 -就是 uniform 分支)、`py_executor` 里的 uniform 分支、跨 rank allgather 里的 tier -字段(5→4 个 int)、以及覆盖该路径的 12 个测试和一个 rank 模拟 helper。共 -−350 行。兜底现在直接是 `runtime_draft_len = planner.max_tier`,不需要跨 rank -协商,所以这一步也少了一次归约。 - -`llm_args` 现在**直接拒绝** `scheduling=True + ragged=False`,而不是把它悄悄当成 -别的东西。关掉特性用 `enable_confidence_scheduling=False`。 - -> **历史数据是否受影响**:§10.2 的 uniform 基线 96.2092 仍然有效,但理由不是 -> 「它没用中间态」——我没有回去核对那一轮的 flag(§2 表一把中间里程碑定义成 -> `scheduling=True + ragged=False`,所以有可能就是它)。理由是 §10.3 已实测 -> planner 在三种工况下**一次都没裁剪**:中间态在本 checkpoint 上必然选到 tier 5, -> 即整块。因此无论那一轮用的是哪套 flag,96.2092 描述的都是「验证整块」,而那条 -> 路径原封不动地留着。 -> -> §10.8 的 sync 对照写的是「confidence head 关」= `scheduling=False`,配置依然合法, -> 不需要重新设计。 - -### 10.10 实现 `cap-accept`(A2 / Q4 收口) - -三模式现已全部可用,`dspark.py` 里的 `NotImplementedError` 已移除。 - -`cap-accept` 算出每请求窗口并**按窗口提交**,但**不缩小交给 target 的 token 轴**。 -它把 `compact` 混在一起的两件事拆开: - - cap-accept 输出 ≠ static 输出 ⇒ 调度策略有问题 - compact 输出 ≠ cap-accept 输出 ⇒ 布局压缩有问题 - -**实现是纯增量的**:窗口写在独立属性 `py_verify_cap` 上,**不碰 `py_verify_len`**。 -于是布局侧(`model_engine` / `cuda_graph_runner` / `dsa.py`)完全看不见窗口、自动 -走整块;`_verified_len` 因为 `py_verify_len is None` 自动返回整块,**KV rewind 随之 -正确**(`rewind = 整块 − capped`)。这样就不必在「为 compact 调了三个 bug 才对」的 -记账代码里再加一组条件分支——§10.6 那个 KV leak 正是这类记账错误。 - -**截断必须在设备端 acceptance 内部**(`interface.apply_accept_caps`)。第一版写在主机侧 -`update_requests`,是错的——这条值得单独记: - -`num_accepted_tokens` 在**同一次 forward 的后续、设备上还要被消费**。`dspark.py:606-652` -里 drafter 用它挑下一块的条件隐状态、回填滚动 KV 窗口、并推进持久解码位置 -(`_ctx_len += nacc`)。主机同步后再截断,设备已按未截断的数推进过 drafter 状态,而主机 -只提交 capped 前缀——输出因验证权威仍无损,但 drafter 状态永久漂移,这个模式也就不再 -是它存在意义所在的可信参照。SGLang 截在同一位置(`dspark_accept.py:138`)。 - -**且不能借 `is_ragged_verify` 触发已有的 clamp**:那个标志同时意味着「draft 缓冲已按 -窗口打包」(`_padded_gen_draft_tokens` 切 `draft_tokens[:total−num_gens]`)。cap-accept -的缓冲是普通矩形,借用会把矩形当 ragged 解包,一个请求的 draft 记到另一个头上,**全程 -不报错**。所以新增独立字段 `SpecMetadata.accept_caps`。 - -**它产出 `compact` 结构上拿不到的数**:`cap_trim_tokens` = 被丢弃的、本可接受的位置 -数,即裁剪的真实接受代价。`compact` 下那些位置根本没算,SGLang 为此专门养了一个 -~600 行的统计估计器;我们此前只有 `trim_regret_rate` 这个下界。 - -两个埋在计数器里的坑,都已处理并有测试: - -1. **不能给它记上算力节省**。`record_step` 会按 `sum(1+v)` 推算 delivered,那会让 - cap-accept 显示出一个它刻意不去拿的 trim_ratio——而那正是判断功能是否生效的 - 头号指标。新增显式 `delivered` 参数,由调用方传整块。 -2. **`assert_ragged_active(require_trim=True)` 会误杀它**。它的 `trim_ratio` 恒为 0 - 是设计如此。该判据现在对不裁 token 轴的模式豁免,否则这道门在它唯一该跑的模式 - 里恒不可满足。 - -代价是**不省任何算力**,所以它是诊断模式,永远不是服务配置。 - -损失**逐请求**记录,不是只记总量:总量分不出「每个请求丢一点」和「少数请求被砍烂」, -两者 `cap_trim_tokens` 相同而结论相反。因此另记 `requests_cap_trimmed` / -`cap_trim_max` / `cap_trim_hist` / `cap_trim_concentration`。链路是 -`apply_accept_caps` 写持久 `[bs]` 缓冲 → worker 返回 `cap_trim_lens` → 采样器 store → -随**已有的** host 拷贝回来,**不额外引入同步**。 - -这里的真风险是**陈旧缓冲**:缓冲持久且按 slot 索引,不产出 caps 的步必须往本批次 -slot **写零**,否则 slot 回收给新请求后会报上一个占用者的损失——今晚已经踩过两次的 -同一类。`_process_outputs` 显式 `zeros_like`,有测试卡住。 - -覆盖:`tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py` -(63 项,含 clamp 的全枚举不变式、context slot 清零、同一总量的两种分布可区分、 -三个计数器坑的回归)。dspark 全套 **335 passed**。**尚未 e2e 跑过。** diff --git a/docs/dspark_ragged_ima_handoff.md b/docs/dspark_ragged_ima_handoff.md deleted file mode 100644 index 9ab9a6ec5bd0..000000000000 --- a/docs/dspark_ragged_ima_handoff.md +++ /dev/null @@ -1,452 +0,0 @@ -# DSpark ragged verify — handoff - -Written 2026-08-04. Supersedes nothing; read alongside -`docs/dspark_p0_task_prompt.md` (the original task; §10 is the results record). - -Branch `dspark-p0`, fork `origin` = `github.com/lancelly/TensorRT-LLM`, -PR branch `origin/confidence_head` (NVIDIA/TensorRT-LLM#17056). - ---- - -## 0. The one thing to read first - -**The CUDA illegal memory access this document hands off has since been root-caused and FIXED** — see §3's dated entries: a ragged graph key derived over a window-less step replayed the previous ragged step's row state; fixed at the key derivation (`6a59b01e53`, with the WAR-staging hardening in `c7e6d78ac2`/`74f1a82d06`/`5c38441a00`), validated 131 steps / zero IMA with ragged fully active. The rest of this section is preserved as the record of how the framing went wrong: I spent most of a day believing the fault was confined to a path I had just added (running ragged on steps with no CUDA graph); the final run disproved that with the path gated OFF. Everything in §4 was designed under that wrong hypothesis — the falsifications stand as facts, but re-derive any future candidate list from "ragged trimming is unsafe", not "eager is unsafe". - ---- - -## 1. Goal - -Make DSpark confidence-head verification scheduling work, and be -CUDA-graph-compatible, on DeepSeek-V4-Pro-DSpark at DEP8 (TP8 + EP8 + -attention-DP, 8x B300), `max_draft_len=5`, tiers `{1,3,5}`, overlap scheduler ON. - -Acceptance, all required: -1. Real checkpoint, not a stub. -2. GSM8K no worse than baseline. Scheduling changes only how many drafted - positions are submitted, never which token is emitted — **accuracy loss is a - bug, not a tradeoff**. -3. attention-DP verified (not just TP). -4. Ragged demonstrably *live*, not silently falling back. -5. Kernels compatible (indexer top-k, DSv4 compressor). -6. CUDA graphs actually replaying, not forced eager. -7. Overlap scheduler ON — that is the shipping configuration. - -Requirement 4 exists because **every failure mode of this feature is silent**: -the planner declining, a flat cost model, a batch missing its captured shape, a -partially-windowed batch — each still produces correct output and baseline -accuracy. GSM8K passing proves nothing on its own. - ---- - -## 2. Where it stands - -### Works, with evidence - -| | Evidence | -|---|---| -| K1/K2 unblocked, **zero C++ changes** | token-major presentation: each generation token becomes a length-1 row, so `seq_len == 1` divides any row count and `num_tokens % num_seqs == 0` passes trivially | -| GSM8K on the ragged path | TRTLLM **96.4367** (uniform baseline 96.2092) | -| cap-accept implemented | **96.096** TRTLLM + MEGAMOE_DEEPGEMM both pass | -| static reference | **96.171** | -| compact | **96.323** — but see the caveat in §3 | -| Uniform tier ladder removed | three states collapsed to two, −350 lines | -| Capture set no longer builds dead tiers | 57 graphs → 19 on the non-ragged branch | -| Planner counters readable | `decisions` was a dead counter; now increments | - -### Open - -| # | Item | -|---|---| -| ~~IMA~~ | **Closed** — root cause and fix in §3 (2026-08-04 ~16:30 entry); validated 131 steps, zero IMA/hang/tripwire, ragged fully active. | -| #19 | Verify budget is quantised to the tier ladder. Measured cost up to **73%** of the planner's own objective. Not implemented. | -| #11 | Sync control never completed. **Do not cite any sync-related conclusion.** | -| #20 | Three-way differential incomplete (compact leg needs a non-blocking run) | - ---- - -## 3. The IMA - -### Symptom - -One rank reports `CUDA error: an illegal memory access was encountered`, peers -report `unspecified launch failure`, at decode step ~75–128 of GSM8K on DEP8 -with the steep synthetic cost table. Surfaces at `dspark.py:117`, but CUDA -errors are asynchronous — that is the next sync point, not the origin. Ranks -that throw stop issuing collectives, so the survivors deadlock and the hang -detector kills the job at 300s. **The hang is a consequence; the IMA is the -cause.** - -### The only configuration that completes - -`CUDA_LAUNCH_BLOCKING=1` — 128 steps, GSM8K 96.323 (`compact_ima.log`). That is -where the 96.323 in §2 comes from, so treat it as "the numbers are right when -the race is suppressed", not as a clean pass. - -### Scope — read this carefully - -The last run of the session (`compact_gated.log`) used the **default** -configuration with the eager-ragged path gated off: - -``` -fallbacks: {'no_graph_this_step': 46, ...} <- gating worked -steps_ragged: 43 <- all of these were GRAPH steps -IMA: 1, hang: 5 -``` - -So the fault occurs on the ordinary graph ragged path — `fit_ragged_verify_lens`, -a captured bucket, a replayed graph. It is not specific to any path added on -this branch. - -Corollary: `commit 4f81cc5534` ("Gate the eager-ragged path off until its IMA is -understood") does **not** make the branch safe. Its stated justification is -wrong. Keep the gate or drop it on its own merits, but do not rely on it. - -### Root cause (best-evidenced theory; coarse-guard fix in validation 2026-08-04) - -GPU coredump-on-exception (`COREDUMP=1` in the runner) ended the guessing that -§6 used to recommend compute-sanitizer for. Three independent runs produced -corefiles naming the same kernel: -`applyMLARopeAndAssignQKVKernelGeneration`, -Warp MMU Fault, and disassembly at the trigger PC shows the faulting -instruction is the 128-bit **KV-cache store** (`STG.E.128`) whose address is -computed from a block offset — i.e. the kernel was handed a block table / KV -extent pair that points outside mapped memory. - -The mechanism is a **write-after-read race on persistent pinned staging**: the -base class already fixed exactly this bug for `kv_cache_block_offsets` (nvbug -6293536 — see the long comment in `resource_manager.py -copy_batch_block_offsets` and `_stage_block_offsets_for_copy`): under the -overlap scheduler the CPU runs an iteration ahead, so an in-place refill of a -pinned host buffer can clobber the source of the previous iteration's -still-queued `non_blocking=True` H2D. The base class's fix was to snapshot into -a fresh pinned buffer per copy. **`dsa.py` re-introduces the unfixed pattern -wholesale**: `prepare_for_mla_rope_append` (four indptr buffers — direct -inputs of the faulting rope kernel), `prepare_for_indexer_k_cache` -(`host_indexer_k_cache_block_offsets`), `prepare_for_indices_conversion` -(`host_req_idx_per_token`), `kv_lens_expanded_host`, and the ragged row maps. -Ragged is what turns the latent race into an MMU fault: uniform decode's -next-step values differ by +1 token, but ragged verify's per-row extents and -row→request mapping change shape step to step, so a one-step-stale mix of -extents and block table walks off the allocated blocks. - -Two narrower fixes failed before the mechanism was fully mapped: per-buffer -events on the ragged row maps only (fault persisted — the indptr staging was -still exposed), and skipping the copies during graph capture (a no-op: -`_forward_step` is captured on **pre-built inputs**, `prepare()` never runs -inside the capture region — verified at `model_engine.py:7322`, the captured -body is `_preprocess_inputs` + `model_forward` only). - -The fix has two layers, because the racy buffers have two different *writers*: - -1. **Buffers rewritten inside `prepare()`** (the DSA indptr/indexer/row-map - staging): one coarse guard in `DSAtrtllmAttentionMetadata.prepare` — - synchronize on an event recorded at the end of the previous `prepare()` - before any pinned buffer is rewritten, record it after this step's copies - are enqueued. Steady-state cost is zero. Trap discovered while validating: - `DeepseekV4TrtllmAttentionMetadata` overrode `prepare()` and skipped - straight to the grandparent, silently bypassing the guard — the V4 override - is now `_prepare_impl` so the wrapper covers it. Any future subclass must - override `_prepare_impl`, never `prepare`. -2. **`host_kv_cache_block_offsets` under KVCacheManagerV2** (the run config - uses `use_kv_cache_manager_v2=True`): its rows are live page-index buffers - the allocator rewrites at *scheduling* time — outside any prepare-level - fence — while the previous step's H2D from the same buffer can still be - queued behind an in-flight forward. Neither of V2's two copy branches had - V1's nvbug-6293536 snapshot fix. A prepare-scope guard provably does not - cover this: a run with layer-1 alone (fix5) still faulted at step ≥64. - Fixed by `_stage_host_block_offsets_for_copy` in `kv_cache_manager_v2.py` - — snapshot to a fresh pinned buffer per copy, V1 parity. This is the layer - that matches the coredump: the block table is what the faulting store's - address is computed from, and ragged trimming is what makes roster churn - (block free/reassign at scheduling time) frequent enough to hit the window. - -Validation note: default-gated runs now die at step ~67 on the runner's own -activity gate (`trim_ratio` ≈ −0.02) because the `4f81cc5534` gate declines -eager-ragged on no-graph steps — a config artifact, not the IMA. Use -`EAGER_RAGGED=1` runs for validation; they exercise BOTH eager-ragged and -graph-ragged (bucketed) steps in one run. - -### 2026-08-04 ~07:00 — the tripwire verdict: VALUE corruption, confirmed live - -The two staging-fix layers above were *necessary but not sufficient*: with both -in, 4 identical runs gave 1 pass / 3 faults. The remaining hypothesis space was -split by a zero-sync device tripwire (`TLLM_DSPARK_VALIDATE_ROWS`, -`torch._assert_async` on the exact (extent, block id) pair the rope KV-append -computes its store address from — placed after the prepare-time gather and -inside the capture body after the extent rebuild) run under -`CUDA_DEVICE_WAITS_ON_EXCEPTION`: - -**The tripwire FIRED** (`CUDA_EXCEPTION_12, Warp Assert` in -`_assert_async_cuda_kernel`, frozen live on the faulting rank) — the value -pair is out of bounds *before* the rope kernel consumes it. This kills the -VMM/page-mapping family (suspend/resume, free-vs-replay unmap) outright and -reduces the problem to: **who writes a bad extent or block id, between two -validated points**. The launch-blocking algebra had already excluded -deterministic value bugs and graph-internal races (a replay is ONE launch — -CUDA_LAUNCH_BLOCKING does not serialize inside it, yet 79 bucketed steps -passed under it), so the writer is a host-side actor racing in-flight device -work on ragged-only row state. - -Forensics protocol that finally worked (traps everywhere else): -- Attach from the host via `enroot exec cuda-gdb -p ` — a sibling - pyxis container cannot ptrace, the host has no cuda-gdb, and gdb through - `/proc//root` cannot map the target's libraries. -- NEVER wrap an attached cuda-gdb in `timeout` — killing it mid-handshake - kills the frozen inferior (lost one scene to this). -- Do not `p *(@parameter ...)` on the assert kernel — cuda-gdb 13.2 crashes - (out_of_range in print_cuda_exception_string) and the crash cascade tears - the run down. -- The faulting rank is named by the peers' `timed out waiting for completion - flag from rank R` printfs, which flush minutes after the freeze. - -The tripwire now also STASHES (site, rows, bad row, extent, block id, bound, -capacity, ok, req idx, kv_len) into a plain-cudaMalloc int64[10] whose address -is logged at allocation (`DSPARK VALIDATE debug buffer @ 0x…`) — after a -freeze, `x/10gd ` in the attached debugger reads the verdict directly. -site 1 = prepare-time gather produced the bad pair (bad inputs from the host -side); site 2 = it appeared only at the in-graph rebuild (a concurrent writer -touched row state between prepare and replay). - -### 2026-08-04 ~09:15 — ROOT CAUSE, read off two frozen scenes - -Two independent stash readouts, identical signature: - -| scene | bucket | row | ext | block id | req | kv_lens_cuda | drift vs allocation | -|---|---|---|---|---|---|---|---| -| i / rank6 | 256 | 60 | 1221 | −1 | 10 | 1226 | ≈ +70 | -| j / rank2 | 96 | 36 | 926 | −1 | 6 | 931 | ≈ +35 | - -Both **site 2** with a clean site 1 in the same step. Since both sites use the -same formula on the same table, site1-clean + site2-trip forces one -conclusion: `kv_lens_cuda[slot]` moved between the prepare H2D reset and the -in-graph rebuild — and the only mutation in that interval is the captured -overlap correction `kv_lens_cuda += previous_kv_lens_offsets_cuda`. Its legal -values are bounded by one verify window (≤7); the observed drift is tens of -tokens. - -The unbounded source: `previous_kv_lens_offsets` is built as -`new_tokens_lens_device[previous_slots] − tokens_per_previous_request` -(`model_engine.py`, ragged branch), and `new_tokens_lens` is the **spec -sampler's slot-indexed store, allocated with `torch.empty`** and rewritten -each step only for the slots actually sampled. A slot the sampler never wrote -reads as arbitrary bytes; one garbage offset pushes the device KV length past -the host's block allocation, the in-graph extent rebuild follows it, and the -rope KV-append store walks into a −1 block entry → unmapped address → Warp -MMU Fault. - -Every historical constraint checks out: **uniform never faults** because a -stable roster rewrites every slot every step (even a stale read is bounded); -**ragged faults** because trimming churns completions and the fault window -(step ~64–90) is exactly when batch concurrency reaches new peaks, i.e. when -slots are read before their first write; **CUDA_LAUNCH_BLOCKING passes** -because it serializes the sampler's store writes ahead of the next step's -staging gather; the staging-fix layers never touched it because the offsets -gather is device-side, not a pinned-host copy. - -Fix (`sampler.py int_tensor` + `model_engine.py` ragged offsets): -1. Slot-indexed sampler stores are **zeros, not empty** — an unwritten slot - reads 0. -2. The ragged offsets gather is **clamped to ±(window)** — a stale value is - then at worst one window of KV length, never an out-of-bounds address. - -Note the residual: a *stale-but-clamped* offset is still a small KV-length -error on roster-churn steps (silent, bounded). The clean long-term fix is an -explicit ordering guarantee (event) between the sampler's store writes and -the next step's previous-batch gather; not attempted yet — measure first. - -### 2026-08-04 ~16:30 — THE ROOT CAUSE, self-described by the stamp tripwire - -The kv-lens-drift reading above was a *symptom* misread (the apparent drift -was a stale row→request map gathering the new roster's kv_lens). The final -mechanism, caught by a prepare-tick/staged-tick stamp pair raising on the -step itself: - -``` -DSpark ragged staleness: prepare tick 559, staged tick 558, skip reason None -``` - -`skip reason None` = the token-major staging chain was NEVER INVOKED that -step: `_ragged_token_lens` returned None (a generation request without a -window), so `_attach_ragged_verify_layout` set `ragged_verify_lens = None` -and prepare skipped every row-map/table refresh — while the graph key still -carried the fit's `agreed_ragged_bucket`, so the step REPLAYED a -ragged-bucket graph over the PREVIOUS ragged step's row state: stale -row→request mapping, stale extents, stale per-row block tables. Wherever the -retired roster's allocation was shorter than a current extent, the rope -KV-append read a −1 block entry and stored to an unmapped address. - -`_ragged_verify_bucket`'s docstring always promised "None unless *every* -generation request carries a window"; the code returned the published bucket -without checking. **Fix: make the code match the docstring** — a step whose -windows and bucket disagree derives the uniform key or falls back to eager -(`no_captured_shape`), never a ragged replay it staged no rows for. - -Why every earlier observation fit a race: the divergence between the fit's -bucket and the requests' windows is timing/roster dependent (planner or peer -declining after the fit, completions between fit and key derivation), so the -fault was intermittent, followed the first captured buckets by a few dozen -steps, never touched uniform (no bucket key), and vanished under -CUDA_LAUNCH_BLOCKING's pacing. Every WAR/staging fix left it untouched -because nothing was racing — the replay was reading last step's rows by -design once the staging chain was skipped. - -First validation (fix11_a): **131 steps, zero IMA/hang/tripwire**, ragged -fully active (`steps_ragged` 38, `trim_ratio` 0.5375, 527 replays), the -consistency decline fired only 6 times (`no_captured_shape: 6`). GSM8K -95.944 vs the 96.0 gate — within one σ (≈0.54 at n=1319) of the historical -96.1–96.4 band; two more samples running to separate noise from a real -regression (suspect list if real: the offsets clamp, the trim-schedule -shift from declined steps). - -The investigation scaffolding (zero-sync device tripwire + forensic stash, -prepare/staged tick stamps, instance markers, pre-replay staleness assert) -was stripped from the tree once the fix was validated — this document and -`tmp/autopsy_recipe.md` (frozen-scene attach protocol, kernel param-space -offsets, the never-wrap-cuda-gdb-in-timeout rule) are the durable record. -Rebuilding any piece takes an hour with this doc; keeping them all in-tree -was judged not worth the noise. - -### Reproduce - -```bash -export LLM_MODELS_ROOT=/lustre/fsw/coreai_comparch_trtllm/laliao/llm-models -SPS_TABLE=/lustre/fsw/coreai_comparch_trtllm/laliao/dspark-runs/steep_sps.json \ - tmp/run_dspark_e2e.sh ragged TRTLLM -``` - -A steep table is **required**: at this checkpoint's ~90% acceptance the planner -correctly declines to trim against the real table, so nothing is exercised. -`tests/microbenchmarks/dspark_make_steep_sps_table.py` generates one; it carries -`SYNTHETIC: true` and **no throughput number taken with it means anything**. - -Useful switches, all wired through `tmp/run_dspark_e2e.sh`: - -| env | effect | -|---|---| -| `LAUNCH_BLOCKING=1` | `CUDA_LAUNCH_BLOCKING=1` — the only thing that passes | -| `NO_OVERLAP=1` | `disable_overlap_scheduler=True` | -| `ASSERT_ROWS=1` | row-map staleness assertion (skipped during capture) | -| `SANITIZER=1` | compute-sanitizer — see §5, not currently viable | - ---- - -## 4. Falsified hypotheses - -Each by its own run unless marked. **Caveat**: several were designed to -discriminate "eager vs graph", which §3 shows was the wrong axis. The -observations stand; their relevance may not. - -By experiment: -- stale `ragged_pad_verify_len` (`compact_confirm`) -- temporary pinned buffers as async H2D sources (`compact_pinned`) -- ragged on a mixed batch (`compact_genonly`; `steps_ragged` rose 35→69 and it - still faulted) -- cross-step overlap (`compact_nooverlap`) -- bucket alignment and the fit (`compact_fitnokey`: full fit, only the graph key - dropped) -- per-step buffer allocation (`compact_eagerpool`: a persistent pool made it - **worse**, 18 faults) -- window raggedness (`compact_eageruni`: identical windows, 12 faults) -- **trimming itself** (`compact_eagerfull`: every window set to `max_draft_len`, - making the step semantically identical to static — still faulted) -- a barrier at the scheduling point (`compact_syncfinal`, `compact_wait`) -- token-major row-map staleness (`compact_assert2`: IMA reproduced 15x, the - assertion never fired, and `steps_ragged: 74` confirms it had every chance to) - -By reading: -- DeepGEMM 32-token alignment — `scheduler_metadata_buffer_expanded` is - `(num_sms+1, 2)`, independent of token count -- row-buffer overflow — `rows <= cap` is asserted (`dsa.py:2024`) -- buffer-pool use-after-free — eager metadata has `cuda_graph_buffers=None`, so - `get_empty` returns `torch.zeros` and never touches the pool -- metadata rebuilt per step — `_set_up_attn_metadata` caches `self.attn_metadata` -- MoE aux-stream overlap — `with_multi_stream(True)` appears only in - `cuda_graph_runner` capture/replay, so multi-stream is OFF on eager steps - ---- - -## 5. Traps — I fell into each of these - -**"It stopped without error" is not "it passed."** Three times a run exited -early for an unrelated reason and I read the absence of a fault as a fix: - -- `compact_eagersync` reported IMA 0 — it died at step 67 on the - `assert_ragged_active` gate, before the step range where the fault occurs. I - wrote a task claiming the fix was found. Re-run to completion: IMA 3. -- `compact_trtllm4` (conservative fallback) showed no fault at step 64 and I - concluded the eager path was to blame. It never completed either. This is what - produced the wrong framing in §0. -- My own row-map assertion did a `.max()` device-to-host read inside the CUDA - graph capture body, raising `cudaErrorStreamCaptureInvalidated` and killing the - engine before any real step ran. That run's "IMA 0" meant nothing. - -**Always check the run reached the failure window** (step ≥ ~128, or a -non-zero `steps_ragged` in the periodic summary) before believing a negative. - -**Post-hoc mechanism is not evidence.** Twice I saw a symptom, found code that -could explain it, and reported it as the cause without checking: -- the hang: I blamed rank-local `enable_spec_decode`; `git diff` showed this - branch never touched it. The real trigger was that ranks 5 and 6 *threw* (my - own layout assertion) and stopped issuing collectives. **The ranks that did - not print a stack were the informative ones.** -- the workspace: I claimed `68 MiB → 604 MiB` was the token-major cost and - changed `predicted_tokens_per_seq` accordingly. Measured: byte-for-byte - identical. The 604 MiB is the *prefill* workspace at `max_num_tokens`, and - `getWorkspaceSize` returns `max(context, generation)`, so the generation term - is invisible. Change reverted. - -**Run the full test selection.** Removing the uniform tier ladder broke 4 tests -in `tests/unittest/llmapi/test_llm_args.py` and 1 integration test. I reported -"all green" from a sweep that never executed either file. - -**Capture is a systematic blind spot.** Three bugs this session had the same -shape — runtime path updated, capture path not: - -| | runtime | capture | -|---|---|---| -| capture set keyed on config flag | by mode | by flag | -| dead tier ladder | draft_len fixed | still built 3x graphs | -| `agreed_ragged_bucket` | set by the fit | set by nobody | - -Unit tests cannot reach capture; it needs a real pass. - -**attention-DP turns a rank-local raise into a cluster deadlock.** A diagnostic -assertion that fires on 2 of 8 ranks costs 300 seconds and produces a hang -report that names nothing. Consider reducing across ranks before raising. - -**compute-sanitizer is not currently viable** — it instruments the ~40-minute -weight load; the log grew zero bytes in 55 minutes. Restricting it to a single -rank is untried and is the obvious next move. - ---- - -## 6. Next steps - -1. #19 (budget quantisation — the largest measured gap, up to 73% of the planner's own objective), then #11 (sync control; until it completes, do not cite any sync conclusion) and #20 (the three-way static / cap-accept / compact differential). All IMA work has landed: coarse prepare() WAR guard with the `_prepare_impl` subclass rule (`c7e6d78ac2`), KVCacheManagerV2 snapshot staging (`74f1a82d06`), bounded/zeroed overlap staging inputs (`5c38441a00`), window-consistent ragged graph key (`6a59b01e53`), and the eager-ragged gate reverted on its merits (`35a1ac4c4a`). - -### Do not - -- Do not resolve by declining ragged whenever a graph is unavailable. Those steps - were eager before this feature existed (`can_run_cuda_graph` is - `num_context_requests == 0`), so declining buys no graph replay — it only - forfeits trimming, measured at 46 of the first 64 steps, `trim_ratio` 0.62→0. -- Do not cite `docs/dspark_p0_task_prompt.md` §10 sync numbers; the control never - ran. -- Do not trust a throughput number taken with the synthetic steep table. - ---- - -## 8. Environment - -- 8x B300 SXM6 (SM103, 275 GB/card), x86_64, node pool `bia*`, partition `batch` -- **Queue on both SLURM accounts in parallel** — `coreai_comparch_trtllm` and - `coreai_comparch_aarwlt`. Their fairshare is independent and differs - enormously: on 2026-08-03 one sat PENDING for over two hours with reason - `Priority` while the other, submitted to the same full partition minutes - later, was RUNNING in 30 seconds. -- Weights: `~/llm-models/DeepSeek-V4-Pro-DSpark` (symlink; `LLM_MODELS_ROOT` - must be exported *before* pytest starts — `MODEL_PATH` is evaluated at import) -- Run logs: `/lustre/fsw/coreai_comparch_trtllm/laliao/dspark-runs/` -- Every number in this document is x86 DGX B300. **GB300 is aarch64 with - NVLink-C2C unified memory — all timing, memory and load-path figures must be - re-measured there.** The structural conclusions transfer. diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 45a24661bddd..50568dbd0621 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -1654,8 +1654,8 @@ def update_spec_dec_param( capture_graph = self.is_cuda_graph # GROW-ONLY ratchet, never an equality check: captured graphs bake the # buffer addresses in, and the old `!=` re-created the buffers whenever - # max_draft_len moved per step -- the graph-path ragged IMA (see - # docs/dspark_ragged_ima_handoff.md). Consumers slice [:rows]/[:width]. + # max_draft_len moved per step -- the graph-path ragged IMA. Consumers + # slice [:rows]/[:width]. prev_cap = getattr(self, "_expanded_alloc_draft_cap", -1) ratchet = max(self._draft_sizing_cap, prev_cap) if ratchet > prev_cap: diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 3d1ce5eec0df..54b424c95b6b 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -807,7 +807,7 @@ def int_tensor(shape: tuple[int, ...], device: str = "cuda") -> torch.Tensor: # yet must read as 0 -- with empty it reads arbitrary bytes, and one such # read (previous_kv_lens_offsets) feeds the captured kv_lens correction, # where a garbage value walks the KV append past the request's allocated - # blocks (the DSpark ragged IMA; see docs/dspark_ragged_ima_handoff.md). + # blocks (the DSpark ragged IMA). return torch.zeros(shape, dtype=torch.int, device=device) diff --git a/tests/microbenchmarks/dsv4_sparse_mla_presentation.py b/tests/microbenchmarks/dsv4_sparse_mla_presentation.py index db82711230c9..4ab64673ad92 100644 --- a/tests/microbenchmarks/dsv4_sparse_mla_presentation.py +++ b/tests/microbenchmarks/dsv4_sparse_mla_presentation.py @@ -2,8 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """B2: batch-major vs token-major presentation of DeepSeek-V4 sparse-MLA generation. -This is the measurement the DSpark ragged-verification design hinges on -(``docs/dspark_confidence_schedule_goal.md`` §6.4 B2 / §7.1 U1). +This is the measurement the DSpark ragged-verification design hinges on. The question ------------ @@ -41,9 +40,8 @@ How to read the result ---------------------- -``docs/dspark_p0_task_prompt.md`` §3: if token-major is materially slower or -incorrect, stop and fall back to the uniform tier ladder (D3 Stage A), which -needs no kernel change at all. The prior is that DeepSeek-V4-Pro under DEP8 +If token-major is materially slower or incorrect, the fallback is the uniform +tier ladder, which needs no kernel change at all. The prior is that DeepSeek-V4-Pro under DEP8 keeps all 128 Q heads per rank, so BMM1's M dimension falls from ``128 * (tier+1)`` to ``128`` -- and because the sparse index set differs per row there was never any KV tile sharing across the window to lose. That is an diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py index 6ed4f1cba050..b8cce346cd75 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py @@ -255,9 +255,8 @@ def _loop_reference(q, kv, attn_sink, topk_idxs, scale): return out -@pytest.mark.parametrize("seed", [0, 1, 2]) -def test_sparse_attn_matches_loop_reference(seed): - torch.manual_seed(seed) +def test_sparse_attn_matches_loop_reference(): + torch.manual_seed(0) b, m, h, d, n, topk = 2, 5, 3, 16, 40, 12 q = torch.randn(b, m, h, d) kv = torch.randn(b, n, d) @@ -291,20 +290,6 @@ def test_sparse_attn_no_sink_matches_sdpa(): torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) -def test_sparse_attn_sink_reduces_mass(): - """A finite sink must strictly shrink the attention output magnitude vs an - infinitely-negative (disabled) sink, because it adds denominator mass only.""" - torch.manual_seed(0) - b, m, h, d, topk = 1, 2, 2, 16, 6 - q = torch.randn(b, m, h, d) - kv = torch.randn(b, topk, d) - idx = torch.arange(topk).view(1, 1, -1).expand(b, m, topk).int() - scale = d**-0.5 - no_sink = dspark_sparse_attn(q, kv, torch.full((h,), float("-inf")), idx, scale) - with_sink = dspark_sparse_attn(q, kv, torch.zeros(h), idx, scale) - assert with_sink.abs().sum() < no_sink.abs().sum() - - def test_sparse_attn_masked_indices_excluded(): """An index of -1 must be excluded exactly (equiv. to dropping that column).""" torch.manual_seed(0) @@ -324,9 +309,7 @@ def test_sparse_attn_masked_indices_excluded(): assert not torch.allclose(got_full, got_masked, rtol=1e-3, atol=1e-3) -@pytest.mark.parametrize( - "start_pos,window,block", [(1, 128, 5), (3, 128, 5), (10, 4, 5), (200, 128, 6)] -) +@pytest.mark.parametrize("start_pos,window,block", [(3, 128, 5), (10, 4, 5)]) def test_get_dspark_topk_idxs_matches_reference(start_pos, window, block): bsz = 3 got = get_dspark_topk_idxs(window, bsz, block, start_pos) @@ -340,11 +323,6 @@ def test_get_dspark_topk_idxs_matches_reference(start_pos, window, block): torch.testing.assert_close(got[bi, mi], ref_row) -def test_get_dspark_topk_idxs_requires_generation(): - with pytest.raises(AssertionError): - get_dspark_topk_idxs(128, 1, 5, 0) - - @pytest.mark.parametrize("rope_head_dim,seqlen", [(64, 16), (64, 1), (128, 8)]) def test_precompute_freqs_cis_matches_reference(rope_head_dim, seqlen): got = precompute_dspark_freqs_cis(rope_head_dim, seqlen, rope_theta=10000.0) @@ -361,16 +339,8 @@ def test_apply_rotary_matches_reference(ndim): got = apply_dspark_rotary(x, fc) ref = _ref_apply_rotary_emb(x, fc) torch.testing.assert_close(got, ref) - - -@pytest.mark.parametrize("ndim", [3, 4]) -def test_apply_rotary_inverse_roundtrip(ndim): - """De-rotation (inverse) must undo the forward rotation (property test).""" - torch.manual_seed(1) - b, s, h, rd = 2, 6, 3, 64 - x = torch.randn(b, s, h, rd) if ndim == 4 else torch.randn(b, s, rd) - fc = precompute_dspark_freqs_cis(rd, s) - roundtrip = apply_dspark_rotary(apply_dspark_rotary(x, fc), fc, inverse=True) + # De-rotation (inverse) must undo the forward rotation (property test). + roundtrip = apply_dspark_rotary(got, fc, inverse=True) torch.testing.assert_close(roundtrip, x, rtol=1e-5, atol=1e-5) @@ -440,15 +410,10 @@ def _run(g): def test_attention_forward_shape_and_determinism(): g = _make_attn_inputs() + before = g["kv_cache0"].clone() o = _run(g) assert tuple(o.shape) == (g["b"], g["block"], g["dim"]) assert torch.isfinite(o.float()).all() torch.testing.assert_close(o, _run(g)) # deterministic - - -def test_attention_forward_does_not_mutate_kv_cache(): - """The rolling window write must be functional (cache cloned, not mutated).""" - g = _make_attn_inputs() - before = g["kv_cache0"].clone() - _run(g) + # The rolling window write must be functional (cache cloned, not mutated). torch.testing.assert_close(g["kv_cache0"], before) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py index 5be568d1257f..909fce895ed4 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py @@ -71,26 +71,21 @@ def test_no_caps_is_a_no_op(): def test_counts_beyond_their_windows_are_clamped_and_the_loss_accumulated(): - counts = torch.tensor([6, 5, 2], dtype=torch.int32) - meta = _Meta(caps=[3, 5, 6]) - apply_accept_caps(counts, 0, meta) - assert counts.tolist() == [3, 5, 2] - # Only the first request lost anything: 6 -> 3. And we can say WHICH. - assert meta.accept_cap_trim[:3].tolist() == [3, 0, 0] - - -def test_each_step_overwrites_the_previous_step_s_losses(): - """Per step, not cumulative: the buffer is what THIS step lost. + """Clamp plus per-step (not cumulative) loss recording. - The host accumulates across steps. A buffer that added would double-count - every request still in the batch next step. + The host accumulates across steps, so the buffer must hold what THIS step + lost: a buffer that added would double-count every request still in the + batch next step. Repeating the same step must therefore leave the same + losses, not multiples of them. """ - meta = _Meta(caps=[2, 2]) + meta = _Meta(caps=[3, 5, 6]) for _ in range(3): - counts = torch.tensor([5, 5], dtype=torch.int32) + counts = torch.tensor([6, 5, 2], dtype=torch.int32) apply_accept_caps(counts, 0, meta) - assert counts.tolist() == [2, 2] - assert meta.accept_cap_trim[:2].tolist() == [3, 3] + # Request 2's cap (6) is wider than its count (2): never raised. + assert counts.tolist() == [3, 5, 2] + # Only the first request lost anything: 6 -> 3. And we can say WHICH. + assert meta.accept_cap_trim[:3].tolist() == [3, 0, 0] def test_context_requests_are_left_alone(): @@ -102,11 +97,9 @@ def test_context_requests_are_left_alone(): # Context slots are zeroed, not left stale. assert meta.accept_cap_trim[:4].tolist() == [0, 0, 4, 3] - -def test_a_generation_free_batch_does_not_crash(): + # Degenerate endpoint: a generation-free batch must not crash or clamp. counts = torch.tensor([4, 4], dtype=torch.int32) - meta = _Meta(caps=[2, 2]) - apply_accept_caps(counts, 2, meta) + apply_accept_caps(counts, 2, _Meta(caps=[2, 2])) assert counts.tolist() == [4, 4] @@ -132,14 +125,6 @@ def test_nothing_is_lost_or_invented(accepted, window): assert committed <= window + 1 -def test_the_cap_never_raises_a_count(): - """A window wider than the block must not invent acceptance.""" - counts = torch.tensor([2], dtype=torch.int32) - meta = _Meta(caps=[99]) - apply_accept_caps(counts, 0, meta) - assert counts.tolist() == [2] - - # --- what the counters must say --------------------------------------------- @@ -158,6 +143,15 @@ def test_cap_trim_is_recorded_per_request(): assert (stats.accept_len + stats.accept_loss_per_request) == pytest.approx(3.0) + # And the summary carries the measurement. + summary = stats.summary() + assert summary["mode"] == "cap-accept" + assert summary["cap_trim_tokens"] == 3 + assert summary["accept_loss_per_request"] == pytest.approx(1.5) + for key in ("requests_cap_trimmed", "cap_trim_max", "cap_trim_hist", + "cap_trim_concentration"): + assert key in summary, f"{key} missing from the summary" + def test_the_same_total_loss_is_distinguished_by_its_shape(): """The reason this is per request at all. @@ -188,18 +182,13 @@ def test_the_same_total_loss_is_distinguished_by_its_shape(): assert concentrated.cap_trim_hist == {20: 1} -def test_concentration_is_zero_when_nothing_was_measured(): - stats = _stats(mode=RaggedVerifyMode.COMPACT) - stats.record_acceptance(accepted=3, window=5) - assert stats.cap_trim_concentration == 0.0 - assert stats.cap_trim_max == 0 - - def test_the_loss_is_absent_rather_than_zero_when_not_measured(): """`compact` cannot produce this number, and must not appear to.""" stats = _stats(mode=RaggedVerifyMode.COMPACT) stats.record_acceptance(accepted=3, window=3) assert stats.cap_trim_tokens == 0 + assert stats.cap_trim_concentration == 0.0 + assert stats.cap_trim_max == 0 # The bound that IS available under compact still works. assert stats.trimmed_hit_ceiling == 1 @@ -210,6 +199,12 @@ def test_cap_accept_must_not_be_credited_with_a_compute_saving(): It submits the full block by construction. Deriving delivered tokens from the windows would report a trim ratio for a run that trimmed no compute at all -- and that ratio is the headline number for whether the feature works. + + For the same reason the active gate must not demand a saving here: + `require_trim` asks "did delivered tokens come in under the no-trim + ceiling", cap-accept answers no by construction, and keying the gate on + the ratio alone would fail every cap-accept run -- the reference run + against which the others are judged. """ stats = _stats() num_reqs = 4 @@ -225,6 +220,8 @@ def test_cap_accept_must_not_be_credited_with_a_compute_saving(): # The schedule is still visible -- that is the whole point of the mode. assert stats.window_tokens == sum(1 + v for v in (5, 3, 2, 1)) assert stats.steps_ragged == 1 + # ...and the gate stays satisfiable in the mode it exists to validate. + stats.assert_ragged_active(require_trim=True) def test_compact_still_derives_delivered_from_the_bucket(): @@ -252,23 +249,6 @@ def test_cap_accept_computes_windows_but_does_not_trim_the_token_axis(): assert not RaggedVerifyMode.STATIC.trims_submitted_tokens -def test_the_active_gate_does_not_demand_a_saving_cap_accept_never_makes(): - """The gate must stay satisfiable in the mode it exists to validate. - - `require_trim` asks "did delivered tokens come in under the no-trim - ceiling". cap-accept answers no by construction, so keying the gate on the - ratio alone would fail every cap-accept run -- the reference run against - which the others are judged. - """ - stats = _stats() - stats.record_step(num_gen_requests=4, - verify_lens=[5, 3, 2, 1], - bucket=None, - delivered=4 * (1 + MAX_DRAFT_LEN)) - assert stats.trim_ratio == 0.0 - stats.assert_ragged_active(require_trim=True) - - def test_the_active_gate_still_demands_a_saving_from_compact(): """The regression the clause above must not introduce.""" stats = _stats(mode=RaggedVerifyMode.COMPACT) @@ -280,18 +260,6 @@ def test_the_active_gate_still_demands_a_saving_from_compact(): stats.assert_ragged_active(require_trim=True) -def test_the_summary_carries_the_measurement(): - stats = _stats() - stats.record_acceptance(accepted=2, window=2, cap_trim=3) - summary = stats.summary() - assert summary["mode"] == "cap-accept" - assert summary["cap_trim_tokens"] == 3 - assert summary["accept_loss_per_request"] == pytest.approx(3.0) - for key in ("requests_cap_trimmed", "cap_trim_max", "cap_trim_hist", - "cap_trim_concentration"): - assert key in summary, f"{key} missing from the summary" - - def test_a_slot_left_unwritten_cannot_inherit_the_previous_occupant_s_loss(): """The stale-buffer hazard, guarded at the one place it can arise. diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py index 9787ba617aae..f1ff0934ec04 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py @@ -55,19 +55,6 @@ def test_survival_is_cumulative_product(): assert torch.allclose(surv[0, :3], torch.tensor([0.9, 0.72, 0.36]), atol=1e-6) -def test_survival_is_non_increasing(): - """The allocator's prefix property depends on this holding for any input.""" - torch.manual_seed(0) - conf = torch.rand(16, BLOCK) - surv = compute_survival(conf) - assert torch.all(surv[:, 1:] <= surv[:, :-1] + 1e-7) - - -def test_survival_rejects_wrong_rank(): - with pytest.raises(ValueError, match=r"\[bs, K\]"): - compute_survival(torch.rand(BLOCK)) - - # -------------------------------------------------------------------------- # budget argmax # -------------------------------------------------------------------------- @@ -95,37 +82,6 @@ def _brute_force_budget(survival, num_gen, table, min_len=1, max_len=None): return best_n -def test_tau_counts_the_always_verified_floor_positions(): - """The floor is verified for free, but its yield still belongs in tau. - - argmax((C + f(n)) / g(n)) depends on C, so dropping the floor's contribution - silently moves the optimum. Two batches that differ only in how good their - floor positions are must be able to choose different budgets. - """ - # The riser sits just past the floor (bs*(min_len+1) = 8 tokens), so both - # shelves are reachable and the constant term can actually decide between - # them. With the riser outside the reachable range this test would still - # pass but prove nothing. - table = SpsCostTable(token_counts=(0, 10), step_time_ms=(1.0, 3.0), fixed_overhead_ms=1.0) - bs = 4 - good_floor = np.cumprod(np.full((bs, BLOCK), 0.99), axis=1) - bad_floor = good_floor.copy() - bad_floor[:, 0] = 0.01 # a hopeless first position drags the whole row down - bad_floor = np.cumprod( - np.concatenate([bad_floor[:, :1], np.full((bs, BLOCK - 1), 0.99)], axis=1), axis=1 - ) - budgets = [] - for surv in (good_floor, bad_floor): - got = compute_verify_token_budget(survival=surv, num_gen_requests=bs, cost_table=table) - assert got == _brute_force_budget(surv, bs, table) - budgets.append(got) - assert budgets[0] != budgets[1], ( - "the two batches differ only in their floor positions' survival, so a " - "planner that ignored the floor's contribution to tau would give them " - "the same budget -- keep the riser inside the reachable token range" - ) - - def test_budget_matches_brute_force_on_staircase(): """Theta is not unimodal on a staircase; a greedy first-descent scan is wrong.""" rng = np.random.default_rng(7) @@ -143,44 +99,6 @@ def test_budget_matches_brute_force_on_staircase(): assert got == want, f"trial {trial}: bs={bs} got={got} want={want}" -def test_greedy_first_descent_would_be_wrong(): - """Pin the specific failure mode the global argmax exists to avoid. - - Theta rises along a shelf, drops at a riser, then rises again along the next - shelf. A greedy loop that stops at the first non-improvement (as the paper's - pseudocode does) parks at the end of the *first* shelf; here the second - shelf's optimum is almost 3x better. - """ - bs = 4 - # A steep measured segment right past the floor (bs*(min_len+1) = 8 - # tokens), then a nearly-flat one out to the end of the range. Under - # interpolation the kink between segment slopes is what makes Theta - # non-unimodal: the first candidates buy tokens on the steep slope and - # Theta drops, then the flat segment makes every further token nearly - # free and Theta climbs past the starting point. - table = SpsCostTable(token_counts=(0, 8, 9, 40), - step_time_ms=(0.5, 0.5, 4.0, 4.2), - fixed_overhead_ms=1.0) - # Uniformly high confidence: every extra candidate adds nearly a full - # token, so tau grows ~7x across the range while cost less than doubles - # past the wall. - surv = np.cumprod(np.full((bs, BLOCK), 0.995), axis=1) - n_star = compute_verify_token_budget(survival=surv, num_gen_requests=bs, cost_table=table) - - cand = np.sort(surv[:, 1:].reshape(-1))[::-1] - - def theta(n): - return (bs + cand[:n].sum()) / table.step_time(2 * bs + n, bs) - - greedy = 0 - while greedy + 1 <= cand.size and theta(greedy + 1) > theta(greedy): - greedy += 1 - - assert theta(n_star) >= theta(greedy) - assert n_star > greedy, "expected the cost kink to defeat a first-descent scan" - assert theta(n_star) > 1.5 * theta(greedy), "the gap should be large, not marginal" - - def test_flat_cost_table_degenerates_to_verify_all(): """A flat cost model makes every token free -> spend the whole budget.""" table = SpsCostTable.flat(2.0) @@ -188,19 +106,7 @@ def test_flat_cost_table_degenerates_to_verify_all(): surv = np.cumprod(np.full((6, BLOCK), 0.9), axis=1) budget = compute_verify_token_budget(survival=surv, num_gen_requests=6, cost_table=table) assert budget == 6 * (BLOCK - 1) - - -def test_low_confidence_yields_small_budget(): - table = SpsCostTable(token_counts=(0, 8, 16, 24, 32), step_time_ms=(2.0, 3.0, 4.0, 5.0, 6.0)) - hi = np.cumprod(np.full((8, BLOCK), 0.98), axis=1) - lo = np.cumprod(np.full((8, BLOCK), 0.30), axis=1) - b_hi = compute_verify_token_budget(survival=hi, num_gen_requests=8, cost_table=table) - b_lo = compute_verify_token_budget(survival=lo, num_gen_requests=8, cost_table=table) - assert b_lo < b_hi - - -def test_budget_is_zero_for_empty_batch(): - table = SpsCostTable.flat() + # An empty batch has no candidates to buy, whatever the table says. assert ( compute_verify_token_budget( survival=np.zeros((0, BLOCK)), num_gen_requests=0, cost_table=table @@ -254,20 +160,19 @@ def test_allocation_is_a_prefix_and_respects_bounds(): assert torch.all(lens <= cfg.resolved_max_verify_len) # Budget is spent above the floor and never exceeded. assert int((lens - cfg.min_verify_len).sum()) <= 20 - - -def test_zero_budget_gives_everyone_the_floor(): - cfg = _cfg(min_verify_len=2) - surv = compute_survival(torch.rand(5, BLOCK)) - lens = schedule_verify_lens_topk(survival=surv, budget=0, cfg=cfg) + # Zero budget degenerates to the floor for everyone. + cfg_floor = _cfg(min_verify_len=2) + lens = schedule_verify_lens_topk( + survival=compute_survival(torch.rand(5, BLOCK)), budget=0, cfg=cfg_floor) assert torch.equal(lens, torch.full((5,), 2, dtype=torch.int32)) - - -def test_huge_budget_saturates_at_the_cap(): - cfg = _cfg(min_verify_len=1, max_verify_len=4) - surv = compute_survival(torch.full((5, BLOCK), 0.99)) - lens = schedule_verify_lens_topk(survival=surv, budget=10**6, cfg=cfg) + # A huge budget saturates at the cap. + cfg_cap = _cfg(min_verify_len=1, max_verify_len=4) + lens = schedule_verify_lens_topk( + survival=compute_survival(torch.full((5, BLOCK), 0.99)), budget=10**6, cfg=cfg_cap) assert torch.equal(lens, torch.full((5,), 4, dtype=torch.int32)) + # An empty batch allocates nothing. + assert schedule_verify_lens_topk( + survival=torch.zeros(0, BLOCK), budget=5, cfg=_cfg()).numel() == 0 def test_survival_eps_excludes_hopeless_positions(): @@ -301,12 +206,6 @@ def test_ties_break_toward_earlier_positions(): assert int(lens.max()) - int(lens.min()) <= 1 -def test_empty_batch_returns_empty(): - cfg = _cfg() - lens = schedule_verify_lens_topk(survival=torch.zeros(0, BLOCK), budget=5, cfg=cfg) - assert lens.numel() == 0 - - def test_schedule_config_rejects_bad_bounds(): with pytest.raises(ValueError, match="min_verify_len must be >= 1"): DSparkScheduleConfig(block_size=BLOCK, min_verify_len=0) @@ -314,6 +213,8 @@ def test_schedule_config_rejects_bad_bounds(): DSparkScheduleConfig(block_size=0) with pytest.raises(ValueError, match="< min_verify_len"): DSparkScheduleConfig(block_size=BLOCK, min_verify_len=5, max_verify_len=3) + with pytest.raises(ValueError, match=r"\[bs, K\]"): + compute_survival(torch.rand(BLOCK)) # -------------------------------------------------------------------------- @@ -368,6 +269,13 @@ def theta(length): return tau / table.step_time(8 * (length + 1), 8) assert theta(chosen) == max(theta(v) for v in (1, 3, 5, 7)) + # With nothing runnable, fall back to the floor. + assert ( + budget_argmax_over_uniform_lens( + survival=surv, num_gen_requests=8, cost_table=table, allowed_lens=[] + ) + == 1 + ) def test_derived_tiers_lose_nothing_versus_continuous_k(): @@ -432,25 +340,6 @@ def test_derived_tiers_are_a_function_of_batch_size(): ) == [1, BLOCK] -def test_flat_cost_table_derives_only_the_endpoints(): - """No shelves to find -> nothing to derive beyond floor and full block.""" - tiers = derive_verify_len_tiers( - cost_table=SpsCostTable.flat(), num_requests=8, block_size=BLOCK, max_tiers=99 - ) - assert tiers == [1, BLOCK] - - -def test_discrete_argmax_falls_back_when_nothing_is_allowed(): - table = SpsCostTable.flat() - surv = np.cumprod(np.full((4, BLOCK), 0.9), axis=1) - assert ( - budget_argmax_over_uniform_lens( - survival=surv, num_gen_requests=4, cost_table=table, allowed_lens=[] - ) - == 1 - ) - - # -------------------------------------------------------------------------- # host-side planner: relay, fallbacks, cross-rank agreement # -------------------------------------------------------------------------- @@ -468,13 +357,6 @@ def _planner(**kw): return DSparkVerifyPlanner(**kw) -def test_planner_returns_one_verify_len_per_request(): - p = _planner() - p._host_buffer, p._copy_event, p._snapshot_valid = torch.rand(6, BLOCK), None, True - lens = p.decide_verify_lens(num_gen_requests=6) - assert lens is not None and len(lens) == 6 - - def test_planner_reads_confidence_by_row_not_by_batch_position(): """The snapshot is slot-indexed; ``rows`` is what re-associates it. @@ -573,21 +455,15 @@ def test_trim_regret_counts_only_drafts_alive_at_the_cut(): assert stats.trimmed_hit_ceiling == 1 assert stats.trim_regret_rate == 0.5 assert stats.accept_len == pytest.approx(2.0) - - -def test_trim_regret_is_zero_when_nothing_was_trimmed(): - stats = _stats() - for accepted in (0, 3, 5): - stats.record_acceptance(accepted=accepted, window=5) - assert stats.requests_trimmed == 0 - assert stats.trim_regret_rate == 0.0 - assert stats.summary()["accept_len"] == pytest.approx(8 / 3, abs=1e-4) - - -def test_acceptance_metrics_appear_in_summary(): - stats = _stats() - stats.record_acceptance(accepted=3, window=3) summary = stats.summary() for key in ("accept_len", "requests_scored", "requests_trimmed", "trim_regret_rate"): assert key in summary, f"{key} missing from the summary" + + # Nothing trimmed: the regret rate is 0/0-safe and stays 0. + untrimmed = _stats() + for accepted in (0, 3, 5): + untrimmed.record_acceptance(accepted=accepted, window=5) + assert untrimmed.requests_trimmed == 0 + assert untrimmed.trim_regret_rate == 0.0 + assert untrimmed.summary()["accept_len"] == pytest.approx(8 / 3, abs=1e-4) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py index 32bc50364039..acbf6eba383f 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py @@ -16,11 +16,9 @@ interpolates both theta(M) and alpha(bs); this port must match. """ -import numpy as np import pytest -from tensorrt_llm._torch.speculative.dspark_planner import ( - SpsCostTable, compute_verify_token_budget) +from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable # The certified GB300 table's shape, abbreviated to the two breakpoints whose # gap hid the bug. @@ -31,13 +29,14 @@ batch_overhead_ms=(11.462, 19.343)) -def test_theta_is_exact_on_breakpoints(): +def test_theta_interpolates_between_breakpoints(): + """Exact on breakpoints, linear between them. + + The regression case: 1512 tokens must cost ~the 1536 price, not the 768 + one. + """ assert TABLE.step_time(768, 256) == pytest.approx(25.244 + 19.343 + 35.61) assert TABLE.step_time(1536, 256) == pytest.approx(25.244 + 19.343 + 105.64) - - -def test_theta_interpolates_between_breakpoints(): - """The regression case: 1512 tokens must cost ~the 1536 price, not the 768 one.""" got = TABLE.step_time(1512, 252) theta = 35.61 + (105.64 - 35.61) * (1512 - 768) / (1536 - 768) alpha = 11.462 + (19.343 - 11.462) * (252 - 128) / (256 - 128) @@ -60,28 +59,3 @@ def test_theta_clamps_outside_the_measured_range(): fixed_overhead_ms=25.244) assert plain.step_time(100, 0) == pytest.approx(25.244 + 23.85) assert plain.step_time(4096, 0) == pytest.approx(25.244 + 105.64) - - -def test_step_times_vector_matches_scalar(): - tokens = np.asarray([512, 900, 1512, 1536]) - vec = TABLE.step_times(tokens, 200) - for tok, v in zip(tokens.tolist(), vec.tolist()): - assert v == pytest.approx(TABLE.step_time(tok, 200)) - - -def test_full_batch_argmax_now_trims(): - """The live failure in miniature: GSM8K pooled survivals at a near-full batch. - - tau5/tau2 ~ 1.48 sits below the true cost ratio ~2.0 but above the floored - one (1.19). With interpolation the argmax must stop buying the full block; - the whole campaign's zero-trim runs came down to this comparison. - """ - bs = 252 - survival = np.tile([0.771, 0.693, 0.502, 0.360, 0.315], (bs, 1)) - budget = compute_verify_token_budget(survival=survival, - num_gen_requests=bs, - cost_table=TABLE, - min_verify_len=1, - allowed_lens=[1, 2, 5]) - rung = 1 + budget // bs - assert rung < 5, f"argmax still buys the full block (budget={budget})" diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py index a799d4b6c174..9b1a27bc24ff 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py @@ -153,13 +153,21 @@ def test_batched_attention_matches_scalar_per_request(seed): def test_batched_attention_persist_writes_through_window(): - """persist=True writes main_kv into the shared window at start_pos%window.""" + """persist gates the window write: False mutates nothing, True writes + main_kv into exactly the start_pos%window row of each request.""" g = _make_batched_inputs(seed=3) G, win = g["G"], g["window"] start_pos = torch.tensor(g["start_positions"], dtype=torch.long) slots = torch.arange(G, dtype=torch.long) cache = g["kv_cache"].clone() before = cache.clone() + + dspark_attention_forward_batched( + g["x"], g["main_x"], start_pos, cache, slots, persist=False, **_attn_kwargs(g) + ) + # persist=False must not mutate the caller's window (functional). + torch.testing.assert_close(cache, before) + dspark_attention_forward_batched( g["x"], g["main_x"], start_pos, cache, slots, persist=True, **_attn_kwargs(g) ) @@ -171,14 +179,6 @@ def test_batched_attention_persist_writes_through_window(): assert torch.equal(changed, expected), f"req {i}: wrong window row written" -def test_batched_attention_no_persist_keeps_window(): - """persist=False must not mutate the caller's window (functional).""" - g = _make_batched_inputs(seed=4) - before = g["kv_cache"].clone() - _batched(g, persist=False) - torch.testing.assert_close(g["kv_cache"], before) - - @pytest.mark.parametrize("start_positions", [(1, 3, 20), (5, 5, 5), (2, 7, 200)]) def test_batched_topk_matches_scalar(start_positions): """Fixed-size masked batched topk == scalar topk per request (valid slots).""" @@ -389,20 +389,34 @@ def test_apply_sts_accepts_host_logits_from_a_device_head(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") def test_ragged_scatter_is_cuda_graph_capturable(): + """scatter_ragged_to_padded and count_accepted_ragged capture and replay. + + Both run inside the target's captured graph; the scatter is the op that + needs ``repeat_interleave`` told its output size up front, and the accept + count must never gain a host sync either. + """ from tensorrt_llm._torch.speculative.dspark_ragged import ( build_qo_indptr, + count_accepted_ragged, scatter_ragged_to_padded, ) lens = torch.tensor([3, 1, 2], dtype=torch.int32, device="cuda") indptr = build_qo_indptr(lens) flat = torch.tensor([10, 11, 12, 20, 30, 31], dtype=torch.int32, device="cuda") + draft = torch.tensor([[1, 2, 3, 9], [4, 9, 9, 9], [5, 6, 9, 9]], device="cuda") + target = torch.tensor([[1, 2, 7, 9], [4, 9, 9, 9], [5, 6, 9, 9]], device="cuda") out = torch.empty(3, 4, dtype=torch.int32, device="cuda") + accepted = torch.empty(3, dtype=torch.int64, device="cuda") def run(): - return scatter_ragged_to_padded( + padded = scatter_ragged_to_padded( flat, verify_lens=lens, qo_indptr=indptr, max_len=4, pad_value=-1 ) + counts = count_accepted_ragged( + draft_tokens=draft, target_tokens=target, verify_lens=lens + ) + return padded, counts s = torch.cuda.Stream() s.wait_stream(torch.cuda.current_stream()) @@ -413,44 +427,19 @@ def run(): g = torch.cuda.CUDAGraph() with torch.cuda.graph(g): - out.copy_(run()) + padded, counts = run() + out.copy_(padded) + accepted.copy_(counts) g.replay() torch.cuda.synchronize() assert out.tolist() == [[10, 11, 12, -1], [20, -1, -1, -1], [30, 31, -1, -1]] + # req 0 matches 2 then diverges; req 1 matches its single position; req 2 + # matches both. Column 3 is padding for every row and must never count. + assert accepted.tolist() == [2, 1, 2] # New packed contents must flow through on replay. flat.copy_(torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32, device="cuda")) g.replay() torch.cuda.synchronize() assert out.tolist() == [[1, 2, 3, -1], [4, -1, -1, -1], [5, 6, -1, -1]] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") -def test_ragged_accept_count_is_cuda_graph_capturable(): - from tensorrt_llm._torch.speculative.dspark_ragged import count_accepted_ragged - - lens = torch.tensor([3, 1, 2], dtype=torch.int32, device="cuda") - draft = torch.tensor([[1, 2, 3, 9], [4, 9, 9, 9], [5, 6, 9, 9]], device="cuda") - target = torch.tensor([[1, 2, 7, 9], [4, 9, 9, 9], [5, 6, 9, 9]], device="cuda") - out = torch.empty(3, dtype=torch.int64, device="cuda") - - def run(): - return count_accepted_ragged(draft_tokens=draft, target_tokens=target, verify_lens=lens) - - s = torch.cuda.Stream() - s.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(s): - for _ in range(3): - run() - torch.cuda.current_stream().wait_stream(s) - - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): - out.copy_(run()) - - g.replay() - torch.cuda.synchronize() - # req 0 matches 2 then diverges; req 1 matches its single position; req 2 - # matches both. Column 3 is padding for every row and must never count. - assert out.tolist() == [2, 1, 2] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py index 80c8c3bb7884..d4270eb836ee 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py @@ -74,7 +74,8 @@ def _propose_with_confidence(conf, markov, *, return_confidence): def test_dspark_propose_scores_without_shortening_the_block(): - """The block is always proposed in full; confidence only scores it.""" + """The block is always proposed in full; confidence only scores it, and + only when asked (return_confidence is the run-constant opt-in).""" torch.manual_seed(1) markov = build_markov_head( markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID @@ -93,13 +94,7 @@ def test_dspark_propose_scores_without_shortening_the_block(): assert torch.all(confidence < 0.0) assert torch.all(conf.apply_sts(confidence) < 0.5) - -def test_dspark_propose_confidence_is_opt_in(): - torch.manual_seed(2) - markov = build_markov_head( - markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID - ).eval() - conf = DSparkConfidenceHead(hidden_size=HID, block_size=BLK).eval() + # Opt-in: with the flag off the head is not consulted at all. _, confidence = _propose_with_confidence(conf, markov, return_confidence=False) assert confidence is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py index e6c121d0f857..c8de13d02ff9 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py @@ -99,15 +99,13 @@ def test_confidence_head_emits_raw_logits(): assert conf.dtype == torch.float32 -def test_apply_sts_is_identity_by_default(): +def test_apply_sts_uses_per_position_temperatures(): head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + # Default (all-ones) calibration is a plain sigmoid, bounded in [0, 1]. raw = torch.tensor([[2.0, 0.0, -2.0, 1.0, -1.0]]) assert torch.allclose(head.apply_sts(raw), torch.sigmoid(raw), atol=1e-6) assert torch.all((head.apply_sts(raw) >= 0.0) & (head.apply_sts(raw) <= 1.0)) - -def test_apply_sts_uses_per_position_temperatures(): - head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) temps = torch.tensor([1.0, 2.0, 4.0, 1.0, 1.0]) head.load_sts_temperatures(temps) raw = torch.full((1, BLK), 4.0) @@ -118,26 +116,19 @@ def test_apply_sts_uses_per_position_temperatures(): def test_load_sts_temperatures_updates_in_place(): - """Rebinding the buffer would be invisible to an already-captured graph.""" - head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) - before = head.sts_temperatures - head.load_sts_temperatures(torch.full((BLK,), 2.0)) - assert head.sts_temperatures is before - assert torch.allclose(before, torch.full((BLK,), 2.0)) - - -def test_load_sts_temperatures_validates(): + """Rebinding the buffer would be invisible to an already-captured graph; + and it is calibration, not a checkpoint weight, so it must stay out of + state_dict and refuse malformed tables.""" head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + assert "sts_temperatures" not in head.state_dict() with pytest.raises(ValueError, match="one per block position"): head.load_sts_temperatures(torch.ones(BLK + 1)) with pytest.raises(ValueError, match="strictly positive"): head.load_sts_temperatures(torch.zeros(BLK)) - - -def test_sts_buffer_is_not_persistent(): - """It is calibration, not a checkpoint weight -- it must not affect loading.""" - head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) - assert "sts_temperatures" not in head.state_dict() + before = head.sts_temperatures + head.load_sts_temperatures(torch.full((BLK,), 2.0)) + assert head.sts_temperatures is before + assert torch.allclose(before, torch.full((BLK,), 2.0)) def test_confidence_head_load_weights_rejects_a_dropped_bias(): @@ -150,12 +141,11 @@ def test_confidence_head_load_weights_rejects_a_dropped_bias(): with pytest.raises(ValueError, match="bias=True"): head.load_weights([{**good, "proj.bias": torch.zeros(1)}]) - -def test_confidence_head_with_bias_accepts_bias_weights(): - head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK, bias=True) + # Constructed with bias=True, the same checkpoint bias is accepted. + biased = DSparkConfidenceHead(hidden_size=HID, block_size=BLK, bias=True) w = {"proj.weight": torch.randn(1, HID), "proj.bias": torch.randn(1)} - head.load_weights([w]) - assert torch.allclose(head.proj.bias, w["proj.bias"]) + biased.load_weights([w]) + assert torch.allclose(biased.proj.bias, w["proj.bias"]) def test_confidence_head_with_markov_concat_dim(): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py index f476cf3a938c..e2b9f049807b 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py @@ -60,54 +60,48 @@ def _write(lines, name="gen.log"): def test_verify_length_comes_from_the_padded_token_total(): - """The label is read out of the log, never inferred from step times.""" + """The label is read out of the log, never inferred from step times. + + Steps the ladder cannot explain are skipped, not mislabelled: filing a + step under a shape it did not run is the mislabelling that once shipped + eight fictional cells with near-zero residuals. + """ log = _write([ _line(1, 1, 62098.0, 6), # executor's first step: dropped _line(2, 244, 79.2, 768), # 244 rows padded to 256 -> rung-2 _line(3, 242, 117.1, 1536), # padded to 256 -> full block _line(4, 128, 61.0, 384), # exactly 128 rows -> rung-2 + _line(5, 256, 300.0, 2048), # L = 7, past the block: skipped + _line(6, 300, 400.0, 1800), # more rows than the ladder holds: skipped ]) + samples = samples_from_iter_log([log], max_draft_len=5, + padded_batch_sizes=LADDER) got = sorted((s.batch_size, s.verify_len, s.step_time_ms) - for s in samples_from_iter_log( - [log], max_draft_len=5, padded_batch_sizes=LADDER)) + for s in samples) assert got == [(128, 2, 61.0), (256, 2, 79.2), (256, 5, 117.1)] + # The table is indexed by the padded token total, which each sample exposes. + assert sorted(s.total_verify_tokens for s in samples) == [384, 768, 1536] def test_the_padded_width_is_what_gets_recorded(): """A 244-row step cost what 256 rows cost; the table is indexed by that. - Filing it under 244 would put the same shape in as many cells as the - roster had sizes, and alpha(bs) would be fitted against a row count no - step ever ran. + The ladder is also what removes the ambiguity: 768 tokens over 128 rows + is L=5 at 128, L=3 at 192, or L=2 at 256 -- without the ladder all three + are admissible and the step is dropped; with it, padding rounds 128 up to + 128 and the answer is L=5 (the first version of this collector kept 312 + of 5682 production steps for lack of this). And a step whose token total + the padded width cannot explain (_get_padded_batch can bail) is dropped: + nothing in the log says padding declined, so that is the only honest move. """ assert resolve_padded_shape(num_rows=244, num_generation_tokens=768, max_draft_len=5, max_batch_size=256, padded_batch_sizes=LADDER) == (256, 2) - - -def test_the_ladder_is_what_removes_the_ambiguity(): - """768 tokens over 128 rows is L=5 at 128, L=3 at 192, or L=2 at 256. - - Without the ladder all three are admissible and the step is dropped; with - it, padding rounds 128 up to 128 and the answer is L=5. Nearly every - low-occupancy step is in this position, which is why the ladder matters: - the first version of this collector kept 312 of 5682 production steps. - """ - ambiguous = resolve_padded_shape(num_rows=128, num_generation_tokens=768, - max_draft_len=5, max_batch_size=256) - assert ambiguous is None - resolved = resolve_padded_shape(num_rows=128, num_generation_tokens=768, - max_draft_len=5, max_batch_size=256, - padded_batch_sizes=LADDER) - assert resolved == (128, 5) - - -def test_a_step_whose_padding_declined_is_skipped(): - """_get_padded_batch can bail; then the width the ladder implies is wrong. - - Nothing in the log says padding declined, so the only honest move is to - drop a step whose token total the padded width cannot explain. - """ + assert resolve_padded_shape(num_rows=128, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256) is None + assert resolve_padded_shape(num_rows=128, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) == (128, 5) assert resolve_padded_shape(num_rows=200, num_generation_tokens=600, max_draft_len=5, max_batch_size=256, padded_batch_sizes=LADDER) is None @@ -130,33 +124,15 @@ def test_the_first_step_of_each_executor_is_dropped(): assert times == [79.0, 79.4] -def test_steps_beyond_the_ladder_are_skipped_not_mislabelled(): - """Filing a step under a shape it did not run is the mislabelling that - once shipped eight fictional cells with near-zero residuals.""" - log = _write([ - _line(2, 244, 79.2, 768), # padded to 256 -> kept - _line(3, 256, 300.0, 2048), # L = 7, past the block - _line(4, 300, 400.0, 1800), # more rows than the ladder holds - ]) - got = [(s.batch_size, s.verify_len) for s in - samples_from_iter_log([log], max_draft_len=5, - padded_batch_sizes=LADDER)] - assert got == [(256, 2)] - - def test_only_the_requested_rank_is_read(): - """Every rank logs; counting all of them would multiply every cell.""" - log = _write([ - _line(2, 256, 79.2, 768, rank="0"), - _line(2, 256, 79.9, 768, rank="3"), - ]) - assert len(samples_from_iter_log([log], max_draft_len=5)) == 1 - + """Every rank logs; counting all of them would multiply every cell. -def test_outlier_steps_are_bounded_out(): + Outlier steps are bounded out on the same pass. + """ log = _write([ - _line(2, 256, 79.2, 768), - _line(3, 256, 5000.0, 768), + _line(2, 256, 79.2, 768, rank="0"), + _line(2, 256, 79.9, 768, rank="3"), # another rank's copy: dropped + _line(3, 256, 5000.0, 768), # outlier past max_step_ms ]) times = [s.step_time_ms for s in samples_from_iter_log([log], max_draft_len=5, max_step_ms=1000.0)] @@ -174,12 +150,6 @@ def test_several_logs_are_pooled(): assert got == [2, 5] -def test_samples_expose_the_token_total_the_table_is_indexed_by(): - log = _write([_line(1, 1, 100.0, 6), _line(2, 256, 79.2, 768)]) - sample = samples_from_iter_log([log], max_draft_len=5)[0] - assert sample.total_verify_tokens == 768 - - def test_a_log_without_the_states_dict_yields_nothing(): """enable_iter_perf_stats off: fail loudly upstream, not with a bad table.""" path = pathlib.Path(tempfile.mkdtemp()) / "bare.log" diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py index e403bbacc816..9619cb9eeaa0 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py @@ -16,18 +16,13 @@ The overlap path builds this iteration's inputs from the *previous* iteration's device tensors, and does it with a fixed stride: every request is assumed to -contribute ``runtime_tokens_per_gen_step`` tokens. Two places break silently -when that assumption is dropped: - - * the flat gathers out of ``new_tokens`` / ``next_draft_tokens``, and the - ``position_ids`` / ``gather_ids`` / ``seq_lens`` the host appends -- a wrong - count shifts every following request, so RoPE phases and attention windows - go wrong without any shape error; and - * ``kv_len_offsets``, which is subtracted from a host-side estimate that was - built with the *same* count. If only one of the two is made per-request the - KV length is off by exactly their difference. - -These tests pin the arithmetic of both, and -- most importantly -- that a +contribute ``runtime_tokens_per_gen_step`` tokens. The flat gathers out of +``new_tokens`` / ``next_draft_tokens`` and the ``position_ids`` / +``gather_ids`` / ``seq_lens`` the host appends break silently when that +assumption is dropped: a wrong count shifts every following request, so RoPE +phases and attention windows go wrong without any shape error. + +These tests pin that arithmetic, and -- most importantly -- that a uniform batch reproduces the strided layout exactly, since that is the path every non-DSpark model keeps taking. """ @@ -52,25 +47,29 @@ def __init__(self, verify_len=None): # -------------------------------------------------------------------------- -def test_missing_attribute_yields_the_batch_wide_count(): - # Every non-DSpark request lands here: the attribute does not exist at all. - assert get_request_tokens_per_gen_step(_Req(), 4) == 4 - - -def test_none_verify_len_yields_the_batch_wide_count(): +_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() - req.py_verify_len = None - assert get_request_tokens_per_gen_step(req, 4) == 4 - - -def test_verify_len_is_a_draft_length_not_a_token_count(): - # py_verify_len counts draft positions, matching runtime_draft_len; the - # accepted token makes it one more. - assert get_request_tokens_per_gen_step(_Req(2), 6) == 3 - - -def test_zero_verify_len_still_carries_the_accepted_token(): - assert get_request_tokens_per_gen_step(_Req(0), 6) == 1 + if verify_len is not _MISSING: + req.py_verify_len = verify_len + assert get_request_tokens_per_gen_step(req, runtime_tokens) == expected # -------------------------------------------------------------------------- @@ -100,143 +99,14 @@ def test_ragged_counts_pack_each_request_window(): # slot 2 columns 0..2, then slot 0 column 0 assert gathered.tolist() == [10, 11, 12, 0] - -def test_zero_count_contributes_nothing(): - rows, cols = ragged_gather_index_lists([1, 2, 3], [2, 0, 1]) - assert rows == [1, 1, 3] - assert cols == [0, 1, 0] - - -def test_empty_batch(): + # 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]) - - -def test_negative_count_raises(): with pytest.raises(ValueError, match="negative gather count"): ragged_gather_index_lists([0], [-1]) - - -def test_total_gathered_equals_the_sum_of_counts(): - counts = [4, 1, 3, 2] - rows, cols = ragged_gather_index_lists([0, 1, 2, 3], counts) - assert len(rows) == len(cols) == sum(counts) - - -# -------------------------------------------------------------------------- -# Host-side layout arithmetic (mirrors the extend-request loop) -# -------------------------------------------------------------------------- - - -def _build_overlap_layout(requests, past_seens, runtime_tokens_per_gen_step): - """Reproduce the per-request bookkeeping the overlap branch emits. - - Mirrors _prepare_tp_inputs' previous-batch branch: one entry per request in - seq_lens/draft_lens/cached, and ``tokens`` flat entries per request in - position_ids/gather_ids/previous_pos_indices. - """ - seq_lens, draft_lens, cached, tokens_each = [], [], [], [] - position_ids, gather_ids, pos_indices = [], [], [] - for slot, (req, past_seen) in enumerate(zip(requests, past_seens)): - tokens = get_request_tokens_per_gen_step(req, runtime_tokens_per_gen_step) - seq_lens.append(tokens) - draft_lens.append(tokens - 1) - tokens_each.append(tokens) - gather_ids.extend(range(len(position_ids), len(position_ids) + tokens)) - position_ids.extend(range(past_seen, past_seen + tokens)) - pos_indices.extend([slot] * tokens) - cached.append(past_seen + tokens) - return { - "seq_lens": seq_lens, - "draft_lens": draft_lens, - "cached": cached, - "tokens_each": tokens_each, - "position_ids": position_ids, - "gather_ids": gather_ids, - "pos_indices": pos_indices, - } - - -def test_uniform_batch_reproduces_the_fixed_stride_layout(): - n = 4 - reqs = [_Req(), _Req(), _Req()] - past = [100, 250, 7] - got = _build_overlap_layout(reqs, past, n) - - assert got["seq_lens"] == [n] * 3 - assert got["draft_lens"] == [n - 1] * 3 - assert len(got["position_ids"]) == 3 * n - assert got["pos_indices"] == [0] * n + [1] * n + [2] * n - assert got["gather_ids"] == list(range(3 * n)) - for slot, base in enumerate(past): - window = got["position_ids"][slot * n : (slot + 1) * n] - assert window == list(range(base, base + n)) - - -def test_ragged_batch_keeps_each_request_positions_contiguous(): - # The bug this pins: emitting a fixed count per request shifts every later - # request's position_ids, which is a silent RoPE-phase error. - reqs = [_Req(3), _Req(0), _Req(1)] - past = [100, 250, 7] - got = _build_overlap_layout(reqs, past, 6) - - assert got["tokens_each"] == [4, 1, 2] - assert got["seq_lens"] == [4, 1, 2] - assert got["draft_lens"] == [3, 0, 1] - assert got["position_ids"] == [100, 101, 102, 103, 250, 7, 8] - assert got["gather_ids"] == list(range(7)) - assert got["pos_indices"] == [0, 0, 0, 0, 1, 2, 2] - - -# -------------------------------------------------------------------------- -# kv_len_offsets: the host estimate and the correction must use one count -# -------------------------------------------------------------------------- - - -def _corrected_kv_lens( - requests, past_seens, accepted, runtime_tokens_per_gen_step, correction_tokens -): - """Host estimate (past_seen + tokens) plus the device correction. - - ``correction_tokens`` is what kv_len_offsets subtracts; passing the - batch-wide value while the estimate used a per-request one is exactly the - A2 bug. - """ - out = [] - for req, past_seen, acc, corr in zip(requests, past_seens, accepted, correction_tokens): - tokens = get_request_tokens_per_gen_step(req, runtime_tokens_per_gen_step) - estimate = past_seen + tokens - out.append(estimate + (acc - corr)) - return out - - -def test_paired_counts_cancel_to_past_seen_plus_accepted(): - reqs = [_Req(3), _Req(0), _Req(1)] - past = [100, 250, 7] - accepted = [2, 1, 2] - tokens_each = [get_request_tokens_per_gen_step(r, 6) for r in reqs] - - got = _corrected_kv_lens(reqs, past, accepted, 6, tokens_each) - - # The token count drops out entirely -- that is what makes the correction - # correct regardless of which window the scheduler picked. - assert got == [p + a for p, a in zip(past, accepted)] - - -def test_batch_wide_correction_under_ragged_is_off_by_the_window_difference(): - reqs = [_Req(3), _Req(0), _Req(1)] - past = [100, 250, 7] - accepted = [2, 1, 2] - n = 6 - - wrong = _corrected_kv_lens(reqs, past, accepted, n, [n] * 3) - right = [p + a for p, a in zip(past, accepted)] - - # Off by (per-request tokens - batch-wide tokens) for every request. - deltas = [w - r for w, r in zip(wrong, right)] - assert deltas == [4 - n, 1 - n, 2 - n] - assert all(d != 0 for d in deltas) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py index bde9c5e320ce..94892660fbaf 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py @@ -40,20 +40,10 @@ def test_padding_rows_do_not_make_the_trim_look_negative(): f"the ratio is 0 -- got {summary['trim_ratio']}, which came from " f"comparing {summary['delivered_tokens']} padded-row tokens against a " f"ceiling counted over {3} real requests") - - -def test_narrow_rank_is_not_charged_for_its_widest_peer(): - """Under attention-DP the padding comes from the widest peer, not locally. - - This is the larger of the two terms: a rank holding 2 requests while its - widest peer holds 7 submits the peer's rows. Booked against a two-request - ceiling that is a ratio of -3.0, i.e. the metric claims the feature made the - step four times more expensive. - """ - stats = _stats(block_size=5) - stats.record_step(num_gen_requests=2, verify_lens=[5, 5], - bucket=8 * 6, padded_bs=8) - assert stats.summary()["trim_ratio"] == pytest.approx(0.0) + # The padded row count itself must land in the summary: the executor once + # passed padded_bs=None unconditionally, and an always-empty padded_bs_hist + # is indistinguishable from "padding never happened". + assert summary["padded_bs_hist"] == {4: 1} def test_a_real_trim_still_reads_as_a_saving(): @@ -82,17 +72,3 @@ def test_paths_without_a_bucket_keep_the_local_row_base(): explicit.record_step(num_gen_requests=3, verify_lens=[2, 3, 1], bucket=8 * 6, padded_bs=8, delivered=3 * 6) assert explicit.summary()["trim_ratio"] == pytest.approx(0.0) - - -def test_padded_bs_reaches_the_histogram(): - """``padded_bs_hist`` was empty in every summary ever logged. - - The executor passed ``padded_bs=None`` unconditionally, so the padded row - count -- the thing that explains the negative ratio -- was never recorded - anywhere. An always-empty histogram is indistinguishable from "padding never - happened", which is how this survived. - """ - stats = _stats(block_size=5) - stats.record_step(num_gen_requests=3, verify_lens=[5, 5, 5], - bucket=4 * 6, padded_bs=4) - assert stats.summary()["padded_bs_hist"] == {4: 1} diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py index ab7846123d5d..ef65a9d25d29 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py @@ -49,30 +49,15 @@ def _layout(lens, bucket=None, buckets=None): # -------------------------------------------------------------------------- -def test_qo_indptr_is_the_exclusive_prefix_sum(): - lens = torch.tensor([3, 1, 4, 2], dtype=torch.int32) - assert build_qo_indptr(lens).tolist() == [0, 3, 4, 8, 10] - - -def test_qo_indptr_slices_recover_each_request(): - """Request r must own exactly [qo_indptr[r], qo_indptr[r+1]).""" - lens = [3, 1, 4, 2] - lay = _layout(lens, bucket=10) - flat = torch.arange(sum(lens)) - for r, n in enumerate(lens): - lo, hi = int(lay.qo_indptr[r]), int(lay.qo_indptr[r + 1]) - assert hi - lo == n - assert flat[lo:hi].numel() == n - assert int(lay.qo_indptr[-1]) == sum(lens) - - -def test_extend_start_loc_matches_indptr_head(): - lay = _layout([2, 5, 1], bucket=8) - assert torch.equal(lay.extend_start_loc, lay.qo_indptr[:-1]) - - -def test_empty_batch_indptr(): - assert build_qo_indptr(torch.zeros(0, dtype=torch.int32)).tolist() == [0] +@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(): @@ -106,6 +91,8 @@ def test_layout_derives_the_bucket_from_the_host_total(): torch.tensor([3, 1, 4], dtype=torch.int32), buckets=[8, 16, 32], total_verify_tokens=8 ) assert lay.graph_num_tokens == 8 + # extend_start_loc is the offset form some attention backends want directly. + assert torch.equal(lay.extend_start_loc, lay.qo_indptr[:-1]) def test_layout_refuses_to_sync_for_the_total(): @@ -118,12 +105,19 @@ def test_layout_refuses_to_sync_for_the_total(): # -------------------------------------------------------------------------- -def test_fill_bucket_converts_padding_into_real_verification(): +@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([2, 1, 1], bucket=12) # 4 used, 8 spare + lay = _layout(lens, bucket=bucket) filled = lay.fill_bucket(max_verify_len=BLOCK) - assert filled.total_verify_tokens == 12 - assert int(filled.verify_lens.sum()) == 12 + assert filled.total_verify_tokens == bucket + assert int(filled.verify_lens.sum()) == bucket filled.validate(max_verify_len=BLOCK, exact_fill=True) @@ -136,14 +130,8 @@ def test_fill_bucket_spreads_over_real_requests_first(): assert max(lens) - min(lens) <= 1, f"unbalanced: {lens}" -def test_fill_bucket_is_a_noop_when_exact(): - lay = _layout([3, 2, 3], bucket=8) - assert lay.fill_bucket(max_verify_len=BLOCK) is not None - assert int(lay.fill_bucket(max_verify_len=BLOCK).verify_lens.sum()) == 8 - - def test_fill_bucket_adds_pad_rows_to_reach_the_captured_batch_size(): - """Row count must hit padded_bs; the leftovers land on the pad rows.""" + """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 @@ -151,10 +139,11 @@ def test_fill_bucket_adds_pad_rows_to_reach_the_captured_batch_size(): assert filled.num_pad_requests == 2 assert int(filled.verify_lens.sum()) == 12 filled.validate(max_verify_len=BLOCK, exact_fill=True) - # Real rows are saturated before any slack is wasted on pad rows. + # 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[0] >= 3 and lens[1] >= 2 - assert min(lens[2:]) >= 1, "a pad row with 0 tokens breaks qo_indptr slicing" + assert lens[2:] == [1, 1] + assert lens[0] + lens[1] == 10 def test_fill_bucket_raises_when_the_bucket_cannot_be_reached(): @@ -164,10 +153,20 @@ def test_fill_bucket_raises_when_the_bucket_cannot_be_reached(): lay.fill_bucket(max_verify_len=BLOCK) -def test_fill_bucket_raises_when_the_bucket_is_too_small(): - lay = _layout([5, 5], bucket=8) # 10 real tokens already > 8 +@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) + lay.fill_bucket(max_verify_len=BLOCK, padded_bs=padded_bs) def test_fill_bucket_rejects_a_padded_bs_below_the_real_count(): @@ -222,8 +221,10 @@ def test_validate_rejects_a_stale_indptr(): # -------------------------------------------------------------------------- -def test_scheduler_output_packs_into_a_valid_layout(): - """The whole point: heterogeneous confidence must yield heterogeneous lengths.""" +def test_ragged_beats_uniform_at_equal_budget(): + """Pins the reason ragged exists: heterogeneous confidence must yield + heterogeneous lengths that pack into a valid layout and out-earn the best + uniform split at the same token spend.""" conf = torch.tensor( [ [0.98, 0.97, 0.96, 0.95, 0.94, 0.93, 0.92], # deep, worth verifying @@ -232,10 +233,14 @@ def test_scheduler_output_packs_into_a_valid_layout(): [0.60, 0.45, 0.35, 0.28, 0.22, 0.18, 0.12], ] ) + surv = compute_survival(conf) cfg = DSparkScheduleConfig(block_size=BLOCK, min_verify_len=1) - lens = schedule_verify_lens_topk(survival=compute_survival(conf), budget=8, cfg=cfg) + budget, bs = 8, conf.shape[0] + lens = schedule_verify_lens_topk(survival=surv, budget=budget, cfg=cfg) assert lens[0] > lens[1], "the confident request should verify deeper" + + # The lengths pack into a valid captured layout. total = int(lens.sum()) lay = RaggedVerifyLayout.from_verify_lens( lens, buckets=[8, 16, 32, 64], total_verify_tokens=total @@ -243,31 +248,13 @@ def test_scheduler_output_packs_into_a_valid_layout(): lay.validate(max_verify_len=BLOCK) assert lay.graph_num_tokens >= total - -def test_ragged_beats_uniform_at_equal_budget(): - """Pins the reason ragged exists: same tokens verified, more expected yield.""" - conf = torch.tensor( - [ - [0.98, 0.97, 0.96, 0.95, 0.94, 0.93, 0.92], - [0.55, 0.40, 0.30, 0.25, 0.20, 0.15, 0.10], - [0.90, 0.85, 0.70, 0.55, 0.45, 0.35, 0.25], - [0.60, 0.45, 0.35, 0.28, 0.22, 0.18, 0.12], - ] - ) - surv = compute_survival(conf) - cfg = DSparkScheduleConfig(block_size=BLOCK, min_verify_len=1) - budget, bs = 8, conf.shape[0] - - ragged_lens = schedule_verify_lens_topk(survival=surv, budget=budget, cfg=cfg) - tau_ragged = sum(float(surv[r, : int(ragged_lens[r])].sum()) for r in range(bs)) - - # Best uniform length that fits the same total token spend. + # Same tokens verified, more expected yield than the best uniform length. + tau_ragged = sum(float(surv[r, : int(lens[r])].sum()) for r in range(bs)) budget_tokens = bs * cfg.min_verify_len + budget tau_uniform = max( float(surv[:, :L].sum()) for L in range(1, BLOCK + 1) if bs * L <= budget_tokens ) - - assert int(ragged_lens.sum()) <= budget_tokens + assert int(lens.sum()) <= budget_tokens assert tau_ragged > tau_uniform, ( f"ragged {tau_ragged:.2f} should beat uniform {tau_uniform:.2f} at the same token spend" ) @@ -306,50 +293,24 @@ def test_scatter_unpacks_a_flat_batch_into_rows(): assert out.tolist() == [[10, 11, 12, -1], [20, -1, -1, -1], [30, 31, -1, -1]] -def test_scatter_round_trips_through_qo_indptr(): - torch.manual_seed(4) - lens_list = [4, 1, 7, 2, 5] - lens = torch.tensor(lens_list, dtype=torch.int32) - lay = _layout(lens_list, bucket=32) - flat = torch.randint(0, 1000, (sum(lens_list),)) - out = scatter_ragged_to_padded(flat, verify_lens=lens, qo_indptr=lay.qo_indptr, max_len=BLOCK) - for r, n in enumerate(lens_list): - lo = int(lay.qo_indptr[r]) - assert out[r, :n].tolist() == flat[lo : lo + n].tolist() - - -def test_accept_count_stops_at_the_first_mismatch(): - draft = torch.tensor([[1, 2, 3, 4]]) - target = torch.tensor([[1, 2, 9, 4]]) - lens = torch.tensor([4]) - assert count_accepted_ragged( - draft_tokens=draft, target_tokens=target, verify_lens=lens - ).tolist() == [2] - - -def test_accept_count_respects_each_request_window(): - """The core ragged invariant: a request only gets credit inside its window.""" - draft = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]]) - target = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]]) - # Row 0 verified all 4; row 1 verified only 2 even though everything matches. - lens = torch.tensor([4, 2]) - got = count_accepted_ragged(draft_tokens=draft, target_tokens=target, verify_lens=lens) - assert got.tolist() == [4, 2] - - -def test_accept_count_ignores_stale_padding_beyond_the_window(): - """Stale slots must not be able to credit an acceptance the target never made. - - Positions past verify_len were never sent to the target. If they happen to - hold matching values from a previous iteration and are not masked, the - cumulative product runs straight through them and reports more accepted - tokens than were actually verified -- wrong output, no error. - """ - draft = torch.tensor([[1, 2, 3, 4]]) - target = torch.tensor([[1, 2, 3, 4]]) # everything "matches", incl. stale tail - lens = torch.tensor([2]) # but only 2 were really verified - got = count_accepted_ragged(draft_tokens=draft, target_tokens=target, verify_lens=lens) - assert got.tolist() == [2], "unmasked padding leaked into the accept count" +@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(): @@ -391,26 +352,20 @@ def _shape(num_real, total, peers=None): ) -def test_shape_rounds_both_axes_up_to_captured_values(): - got = _shape(num_real=5, total=20) - assert got.padded_bs == 8 # 5 -> 8 - # slack = 20 - 5 = 15; needed = 8 + 15 = 23 -> 32 - assert got.bucket == 32 - - -def test_bucket_covers_the_padded_rows_floor(): - # Every row verifies at least its bonus position, so the bucket can never - # be smaller than the row count. - got = _shape(num_real=3, total=3) - assert got.padded_bs == 4 - assert got.bucket >= got.padded_bs - - -def test_exact_captured_values_are_not_rounded_further(): - # 8 requests, slack 8 -> needed 16, both already captured. - got = _shape(num_real=8, total=16) - assert got.padded_bs == 8 - assert got.bucket == 16 +@pytest.mark.parametrize( + "num_real, total, want_bs, want_bucket", + [ + # 5 -> 8 rows; slack = 20 - 5 = 15; needed = 8 + 15 = 23 -> 32. + (5, 20, 8, 32), + # exact captured values are not rounded further: 8 requests with + # slack 8 -> needed 16, both axes already captured. + (8, 16, 8, 16), + ], +) +def test_shape_rounds_both_axes_up_to_captured_values(num_real, total, want_bs, want_bucket): + got = _shape(num_real=num_real, total=total) + assert got.padded_bs == want_bs + assert got.bucket == want_bucket def test_every_rank_computes_the_same_shape(): @@ -421,37 +376,37 @@ def test_every_rank_computes_the_same_shape(): assert len(shapes) == 1 -def test_shape_is_driven_by_the_widest_rank_not_the_local_one(): - # This rank is small (3 requests, 9 tokens) but must still size for the - # widest peer: rows from max real (6 -> 8), slack from max (30 - 6 = 24). - peers = [(3, 9), (6, 30), (5, 11)] - got = _shape(num_real=3, total=9, peers=peers) - assert got.padded_bs == 8 - assert got.bucket == 32 # 8 + 24 = 32, exactly a captured bucket - - -def test_widest_axes_can_come_from_different_ranks(): - # rank A has the most rows, rank B the most slack. - peers = [(7, 8), (2, 20)] - got = _shape(num_real=7, total=8, peers=peers) - assert got.padded_bs == 8 # from rank A - # slack: A = 1, B = 18 -> needed = 8 + 18 = 26 -> 32 - assert got.bucket == 32 - - -def test_shape_raises_when_the_batch_exceeds_captured_rows(): - with pytest.raises(ValueError, match="exceed the largest captured batch"): - _shape(num_real=17, total=17) - - -def test_shape_raises_when_no_token_bucket_fits(): - with pytest.raises(ValueError, match="exceeds the largest captured bucket"): - _shape(num_real=16, total=1000) - - -def test_shape_rejects_fewer_tokens_than_requests(): - with pytest.raises(ValueError, match="cannot be below"): - _shape(num_real=5, total=4) +@pytest.mark.parametrize( + "num_real, total, peers, want_bs, want_bucket", + [ + # This rank is small (3 requests, 9 tokens) but must size for the + # widest peer: rows from max real (6 -> 8), slack from the max + # (30 - 6 = 24) -> 8 + 24 = 32, exactly a captured bucket. + (3, 9, [(3, 9), (6, 30), (5, 11)], 8, 32), + # The widest axes can come from different ranks: rank A has the most + # rows (7 -> 8), rank B the most slack (18) -> 8 + 18 = 26 -> 32. + (7, 8, [(7, 8), (2, 20)], 8, 32), + ], +) +def test_shape_is_driven_by_the_widest_rank_not_the_local_one( + num_real, total, peers, want_bs, want_bucket +): + got = _shape(num_real=num_real, total=total, peers=peers) + assert got.padded_bs == want_bs + assert got.bucket == want_bucket + + +@pytest.mark.parametrize( + "num_real, total, match", + [ + (17, 17, "exceed the largest captured batch"), # rows overflow + (16, 1000, "exceeds the largest captured bucket"), # tokens overflow + (5, 4, "cannot be below"), # fewer tokens than requests + ], +) +def test_shape_raises_when_the_batch_exceeds_captured_rows(num_real, total, match): + with pytest.raises(ValueError, match=match): + _shape(num_real=num_real, total=total) def test_shape_requires_bs_buckets(): @@ -461,16 +416,6 @@ def test_shape_requires_bs_buckets(): ) -def test_uniform_batch_shape_matches_the_one_dimensional_rule(): - # With a uniform K the token total is bs * (K + 1), so the joint rule must - # not pick anything larger than the uniform path would have needed. - k_plus_1 = 4 - for num_real in (1, 3, 8): - got = _shape(num_real=num_real, total=num_real * k_plus_1) - assert got.padded_bs >= num_real - assert got.bucket >= num_real * k_plus_1 - - # --------------------------------------------------------------------------- # fill_padded_rows_onehot # --------------------------------------------------------------------------- @@ -489,12 +434,11 @@ def test_padding_rows_become_a_valid_distribution(): assert torch.allclose(probs[1, 1].sum(), torch.tensor(1.0)) assert torch.allclose(probs[1, 2].sum(), torch.tensor(1.0)) - -def test_full_length_rows_are_left_alone(): - probs = torch.full((2, 3, 4), 0.25) - before = probs.clone() - fill_padded_rows_onehot(probs, verify_lens=torch.tensor([3, 3])) - assert torch.equal(probs, before) + # 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) # --------------------------------------------------------------------------- @@ -511,42 +455,11 @@ def _fit(real_token_lens, bucket, padded_bs, max_verify_len): return layout.fill_bucket(max_verify_len=max_verify_len, padded_bs=padded_bs) -def test_fit_spends_slack_on_real_requests_before_pad_rows(): - got = _fit([2, 2], bucket=8, padded_bs=4, max_verify_len=6) - lens = got.verify_lens.tolist() - # Two real rows grew; the two pad rows stay at the 1-token minimum. - assert sum(lens) == 8 - assert lens[2] == 1 and lens[3] == 1 - assert lens[0] + lens[1] == 6 - - -def test_fit_rejects_a_bucket_that_cannot_hold_the_pad_rows(): - # The same batch against a bucket that leaves no room for the pad rows: - # 8 real tokens + 6 pad rows needs 14, not 10. Raising here is the point -- - # a short batch would desync attention from the MoE. - with pytest.raises(ValueError, match="too small"): - _fit([4, 4], bucket=10, padded_bs=8, max_verify_len=4) - - 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 -def test_capture_buckets_match_the_uniform_token_totals(): - # The ragged ladder is one bucket per tier at each batch size, which is - # exactly what the uniform ladder would have produced -- so enabling ragged - # does not grow the captured graph count. - tiers = [1, 3, 5] - for bs in (1, 4, 16): - buckets = [bs * (t + 1) for t in tiers] - assert buckets == sorted(buckets) - assert len(buckets) == len(tiers) - # Every bucket is reachable: bs rows at >= 1 token, <= max tier + 1. - for b in buckets: - assert bs <= b <= bs * (tiers[-1] + 1) - - def test_warmup_split_reaches_any_captured_bucket(): # Capture has to hit the bucket exactly or the graph is keyed on a total no # runtime batch produces, and every ragged step silently runs eager. @@ -582,28 +495,24 @@ def _pad_len(*, bucket, n_real, n_pad, max_verify_len, floor=None): return pad_len -def test_pad_rows_take_the_minimum_when_real_rows_can_absorb_the_bucket(): - # 3 real rows capped at 6 hold up to 18; bucket 24 minus 5 pad rows leaves - # 19 -- one too many, so the pad rows have to grow. Contrast with below. - padded_bs, bucket, max_verify_len = 8, 18, 6 - real = [2, 3, 2] - n_pad = padded_bs - len(real) - pad_len = _pad_len(bucket=bucket, n_real=len(real), n_pad=n_pad, max_verify_len=max_verify_len) - assert pad_len == 1 - - real_target = bucket - n_pad * pad_len - got = _fit(real, bucket=real_target, padded_bs=len(real), max_verify_len=max_verify_len) - assert int(got.verify_lens.sum()) + n_pad * pad_len == bucket - - -def test_pad_rows_grow_when_the_real_rows_cannot_absorb_the_bucket(): - # This is the case that used to have no solution: 3 real rows cap at 18, - # but bucket 24 with 5 one-token pad rows would demand 19 from them. - padded_bs, bucket, max_verify_len = 8, 24, 6 +@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) pad_len = _pad_len(bucket=bucket, n_real=len(real), n_pad=n_pad, max_verify_len=max_verify_len) - assert pad_len == 2, "pad rows must absorb what the real rows cannot" + assert pad_len == expected_pad_len real_target = bucket - n_pad * pad_len assert sum(real) <= real_target <= len(real) * max_verify_len diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py index 8817f6f9a170..78d882c0ec0f 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py @@ -28,10 +28,6 @@ def _payload(engine): "_meta": {"engine": engine} if engine is not None else {}} -def test_matching_fingerprint_loads_silently(): - check_table_fingerprint(payload=_payload(dict(LIVE)), live=dict(LIVE)) - - def test_mismatched_engine_is_refused(): """The 23%-apart case: same cell, different max_batch_size.""" wrong = dict(LIVE, max_batch_size=64) @@ -39,26 +35,23 @@ def test_mismatched_engine_is_refused(): check_table_fingerprint(payload=_payload(wrong), live=dict(LIVE)) -def test_comparison_is_case_insensitive(): - """String facts must not fail on spelling case (MEGAMOE vs MegaMoe).""" - fp = dict(LIVE, moe_backend="megamoe_cutedsl") - live = dict(LIVE, moe_backend="MEGAMOE_CUTEDSL") - check_table_fingerprint(payload=_payload(fp), live=live) - +def test_fingerprints_that_must_load(): + """Every non-contradicted fingerprint loads; only a contradicted fact refuses. -def test_extra_table_facts_do_not_block_loading(): - """Facts the consumer cannot see (image hash, geometry) are informational. - - They must neither raise nor be silently swallowed -- the function logs - them -- but loading proceeds: refusing on unverifiable keys would make - every fingerprint addition a breaking change. + Matching facts load silently; string facts must not fail on spelling case + (MEGAMOE vs MegaMoe); facts the consumer cannot see (image hash, geometry) + are informational and logged, not refused -- refusing on unverifiable keys + would make every fingerprint addition a breaking change; and an old table + with no fingerprint warns rather than breaking existing deployments. """ - fp = dict(LIVE, image="faf2c60935", geometry="constant_block") - check_table_fingerprint(payload=_payload(fp), live=dict(LIVE)) - - -def test_unfingerprinted_table_still_loads(): - """Backward compatibility: old tables warn, they do not break.""" + check_table_fingerprint(payload=_payload(dict(LIVE)), live=dict(LIVE)) + check_table_fingerprint( + payload=_payload(dict(LIVE, moe_backend="megamoe_cutedsl")), + live=dict(LIVE, moe_backend="MEGAMOE_CUTEDSL")) + check_table_fingerprint( + payload=_payload(dict(LIVE, image="faf2c60935", + geometry="constant_block")), + live=dict(LIVE)) check_table_fingerprint(payload=_payload(None), live=dict(LIVE)) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py index e44f3809ea88..4e4718e80dbc 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py @@ -88,13 +88,10 @@ def _stats_row(*, iteration, num_gen, rank=0, num_ctx=0, host_ms=10.0, gpu_ms=8. # -------------------------------------------------------------------------- -def test_aligned_steps_single_rank(): - rows = [_stats_row(iteration=i, num_gen=8, host_ms=10.0 + i) for i in range(3)] - assert aligned_steps_from_stats(rows) == [(0, 8, 10.0), (1, 8, 11.0), (2, 8, 12.0)] - - def test_aligned_steps_drops_prefill_iterations(): """A step with context requests is a different shape, not a cheap decode.""" + steady = [_stats_row(iteration=i, num_gen=8, host_ms=10.0 + i) for i in range(3)] + assert aligned_steps_from_stats(steady) == [(0, 8, 10.0), (1, 8, 11.0), (2, 8, 12.0)] rows = [ _stats_row(iteration=0, num_gen=0, num_ctx=8, host_ms=90.0), _stats_row(iteration=1, num_gen=8, num_ctx=2, host_ms=40.0), @@ -104,7 +101,12 @@ def test_aligned_steps_drops_prefill_iterations(): def test_aligned_steps_requires_every_dp_rank(): - """A missing rank means the iteration was not observed batch-wide.""" + """Multi-rank alignment: complete, agreeing, and timed by rank 0. + + A missing rank means the iteration was not observed batch-wide; ranks at + different batch sizes are not one (bs, M) point; and ADP fans out + rank-local counters with rank-0's clock, so timing must never be averaged. + """ rows = [ _stats_row(iteration=0, num_gen=8, rank=0), _stats_row(iteration=0, num_gen=8, rank=1), @@ -112,23 +114,17 @@ def test_aligned_steps_requires_every_dp_rank(): ] assert aligned_steps_from_stats(rows, expected_ranks=2) == [(0, 8, 10.0)] - -def test_aligned_steps_requires_ranks_to_agree_on_batch_size(): - """Ranks at different batch sizes are not one (bs, M) point.""" - rows = [ + disagreeing = [ _stats_row(iteration=0, num_gen=8, rank=0), _stats_row(iteration=0, num_gen=7, rank=1), ] - assert aligned_steps_from_stats(rows, expected_ranks=2) == [] + assert aligned_steps_from_stats(disagreeing, expected_ranks=2) == [] - -def test_aligned_steps_takes_rank0_timing(): - """ADP fans out rank-local counters with rank-0's clock; do not average.""" - rows = [ + rank0_not_first = [ _stats_row(iteration=0, num_gen=8, rank=1, host_ms=99.0), _stats_row(iteration=0, num_gen=8, rank=0, host_ms=10.0), ] - assert aligned_steps_from_stats(rows, expected_ranks=2) == [(0, 8, 10.0)] + assert aligned_steps_from_stats(rank0_not_first, expected_ranks=2) == [(0, 8, 10.0)] def test_aligned_steps_skips_missing_timing_field(): @@ -171,12 +167,6 @@ def test_summarize_reports_median_not_mean(): assert cells[0].p90_ms > cells[0].step_time_ms -def test_summarize_refuses_short_cell(): - """A silently dropped cell can disconnect the grid, not just widen error bars.""" - with pytest.raises(InsufficientSamplesError, match="steady samples"): - summarize_cells(_samples(8, 3, [10.0] * 5), warmup_steps=4, min_samples=8) - - def test_summarize_drops_thin_cells_outside_the_requested_grid(): """A deployment log carries ramp/drain occupancy nobody asked to measure. @@ -184,7 +174,9 @@ def test_summarize_drops_thin_cells_outside_the_requested_grid(): requested (bs, L) died at the very end on (bs=2, L=1) with one sample -- the admission ramp of the first cell's client, filed as a cell and then held to the requested-grid bar. Incidental cells must be dropped when - thin, kept when thick, and a thin REQUESTED cell must still refuse. + thin, kept when thick, and a thin REQUESTED cell must still refuse -- + with or without a requested grid, since a silently dropped cell can + disconnect the grid, not just widen error bars. """ grid = [(8, 3)] requested = _samples(8, 3, [10.0] * 8) @@ -196,6 +188,9 @@ def test_summarize_drops_thin_cells_outside_the_requested_grid(): with pytest.raises(InsufficientSamplesError, match="steady samples"): summarize_cells(_samples(8, 3, [10.0]) + thick_extra, warmup_steps=0, min_samples=4, expected_cells=grid) + # No grid given means every observed cell was requested: same refusal. + with pytest.raises(InsufficientSamplesError, match="steady samples"): + summarize_cells(_samples(8, 3, [10.0] * 5), warmup_steps=4, min_samples=8) def test_summarize_orders_by_iteration_before_cutting_warmup(): @@ -313,13 +308,6 @@ def test_fit_clamps_a_dip_upward(): assert any("monoton" in w for w in fit.warnings) -def test_fit_keeps_theta_positive(): - """SpsCostTable rejects a non-positive step time, so the shift is bounded.""" - cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) - fit = fit_additive_cost_model(cells) - assert min(fit.theta_ms.values()) > 0.0 - - def test_fit_refuses_a_disconnected_sweep(): """Two islands each carry their own unknown offset; splicing them lies.""" cells = _synthetic_cells([8], [1, 2], alpha=ALPHA.get, theta=THETA) @@ -328,13 +316,6 @@ def test_fit_refuses_a_disconnected_sweep(): fit_additive_cost_model(cells) -def test_fit_refuses_a_single_length_sweep(): - """One verify length per batch size confounds alpha and theta completely.""" - cells = _synthetic_cells(BATCH_SIZES, [5], alpha=ALPHA.get, theta=THETA) - with pytest.raises(SweepGeometryError): - fit_additive_cost_model(cells) - - # -------------------------------------------------------------------------- # staircase shaping # -------------------------------------------------------------------------- @@ -395,17 +376,21 @@ def _payload(): return build_cost_table_payload(cells, riser_tolerance=0.01, max_breakpoints=8), cells -def test_payload_has_exactly_the_loader_keys(): - payload, _ = _payload() +def test_round_trip_reproduces_measured_cells(): + """The whole point: what the planner reads back must be what was measured. + + The payload must carry exactly the keys ``dspark.py::_build_verify_planner`` + reads (everything else under ``_meta``, where the loader ignores it), + survive a JSON file round trip into a real ``SpsCostTable``, and price + every measured cell back to within one shelf width -- the emitted curve is + a floor staircase by construction, so a cell between two risers is + deliberately priced at the lower shelf. + """ + payload, cells = _payload() assert LOADER_KEYS.issubset(payload) - # Everything else must live under _meta, where the loader ignores it. assert set(payload) - LOADER_KEYS == {"_meta"} - -def test_payload_survives_a_json_file_round_trip(): - payload, _ = _payload() - reloaded = json.loads(json.dumps(payload)) - table = load_cost_table(reloaded) + table = load_cost_table(json.loads(json.dumps(payload))) assert isinstance(table, SpsCostTable) assert not table.is_flat # The loader builds SpsCostTable directly, so the profiler's own validation @@ -413,16 +398,6 @@ def test_payload_survives_a_json_file_round_trip(): assert list(table.token_counts) == payload["token_counts"] assert len(table.batch_sizes) == len(table.batch_overhead_ms) - -def test_round_trip_reproduces_measured_cells(): - """The whole point: what the planner reads back must be what was measured. - - The tolerance is one shelf width, because the emitted curve is a floor - staircase by construction -- a cell between two risers is deliberately - priced at the lower shelf. - """ - payload, cells = _payload() - table = load_cost_table(json.loads(json.dumps(payload))) for cell in cells: predicted = table.step_time(cell.total_verify_tokens, cell.batch_size) assert predicted == pytest.approx(cell.step_time_ms, abs=1e-4) @@ -442,15 +417,12 @@ def test_emitted_terms_are_not_double_counted(): def test_batch_overhead_grows_with_batch_size(): - payload, _ = _payload() + payload, cells = _payload() overheads = payload["batch_overhead_ms"] assert overheads == sorted(overheads) assert payload["batch_sizes"] == sorted(BATCH_SIZES) assert payload["fixed_overhead_ms"] >= 0.0 - - -def test_meta_carries_the_tier_advice_and_the_cells(): - payload, cells = _payload() + # Provenance rides along under _meta: the encoding marker and the cells. assert payload["_meta"]["encoding"] == "decomposed" assert len(payload["_meta"]["cells"]) == len(cells) @@ -508,34 +480,16 @@ def test_refuses_an_inert_table(): """Non-flat, but no tier pair is worth trading yield for. This is the failure mode that does *not* trip ``fallback_flat_cost``, so - nothing downstream would report it. + nothing downstream would report it. ``allow_inert=True`` is the explicit + escape hatch -- and it must still say that nothing trims. """ + kwargs = dict(batch_sizes=[32], tiers=[1, 3, 5], acceptance_rates=[0.9, 0.95]) with pytest.raises(InertCostTableError, match="inert"): - check_table_is_informative( - INERT_PAYLOAD, - batch_sizes=[32], - tiers=[1, 3, 5], - acceptance_rates=[0.9, 0.95], - ) - - -def test_allow_inert_escape_hatch(): - diagnostics = check_table_is_informative( - INERT_PAYLOAD, - batch_sizes=[32], - tiers=[1, 3, 5], - acceptance_rates=[0.9, 0.95], - allow_inert=True, - ) + check_table_is_informative(INERT_PAYLOAD, **kwargs) + diagnostics = check_table_is_informative(INERT_PAYLOAD, allow_inert=True, **kwargs) assert diagnostics["trims_somewhere"] is False -def test_refuses_a_single_entry_ladder(): - payload, _ = _payload() - with pytest.raises(InertCostTableError, match="single entry"): - check_table_is_informative(payload, batch_sizes=BATCH_SIZES, tiers=[5]) - - def test_informative_table_passes_and_reports_where_trimming_wins(): payload, _ = _payload() diagnostics = check_table_is_informative( @@ -546,6 +500,9 @@ def test_informative_table_passes_and_reports_where_trimming_wins(): ) assert diagnostics["trims_somewhere"] is True assert diagnostics["best_step_time_spread"] > 0.02 + # Even a good table cannot rescue a one-rung ladder: nothing to choose. + with pytest.raises(InertCostTableError, match="single entry"): + check_table_is_informative(payload, batch_sizes=BATCH_SIZES, tiers=[5]) # -------------------------------------------------------------------------- @@ -570,14 +527,10 @@ def _sweep(**overrides) -> SweepConfig: def test_token_budget_is_the_step_count_when_acceptance_is_pinned(): - """Pinned, a step commits exactly one token, so the two are the same number.""" - config = _sweep() - assert config.max_tokens_for(5) == config.warmup_steps + config.measure_steps + 16 - - -def test_token_budget_is_scaled_when_acceptance_is_not_pinned(): - """Unpinned, a step can commit verify_len + 1 tokens, so the budget must cover it.""" + """Pinned, a step commits exactly one token, so the two are the same + number; unpinned it can commit verify_len + 1, so the budget must scale.""" pinned = _sweep() + assert pinned.max_tokens_for(5) == pinned.warmup_steps + pinned.measure_steps + 16 loose = _sweep(pin_acceptance=0.0) assert loose.max_tokens_for(5) == pinned.max_tokens_for(5) * 6 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py index 6eec0a98e46a..1324ad0c7684 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py @@ -68,17 +68,20 @@ def test_a_loaded_cost_table_without_a_pin_is_refused(): def test_a_pinned_window_makes_the_table_harmless(): - """Pinning to the full block makes the label uncensored by construction.""" + """The uncensored regimes must all build a recorder. + + Pinning to the full block makes the label uncensored by construction; + without a table the planner cannot trim, so nothing censors anything; + static mode contributes no padding rows. + """ with collecting(pin=5): rec = make_recorder_from_env(block_size=5, has_cost_table=True) assert rec is not None and rec.block_size == 5 - - -def test_no_table_needs_no_pin(): - """Without a table the planner cannot trim, so nothing censors anything.""" with collecting(): assert make_recorder_from_env(block_size=5, has_cost_table=False) is not None + assert make_recorder_from_env(block_size=5, + ragged_mode="static") is not None def test_compact_mode_is_refused(): @@ -92,12 +95,6 @@ def test_compact_mode_is_refused(): make_recorder_from_env(block_size=5, ragged_mode="compact") -def test_static_mode_is_allowed(): - with collecting(): - assert make_recorder_from_env(block_size=5, - ragged_mode="static") is not None - - class _Head: pass @@ -125,20 +122,14 @@ def __init__(self, head): self.dspark_model = _Bare(head) -def test_head_resolves_on_the_bare_model(): +def test_head_resolves_across_both_layouts_and_absence_is_none(): from tensorrt_llm._torch.speculative.dspark_sts import resolve_confidence_head h = _Head() assert resolve_confidence_head(_Bare(h)) is h - - -def test_head_resolves_through_the_forcausallm_wrapper(): - from tensorrt_llm._torch.speculative.dspark_sts import resolve_confidence_head - h = _Head() + # The wrapper layout is the one the chained getattr missed: resolving to + # None here silently disabled calibration for every confidence-scheduled + # run of the first campaign. assert resolve_confidence_head(_Wrapper(h)) is h - - -def test_absent_head_is_none_not_an_error(): - from tensorrt_llm._torch.speculative.dspark_sts import resolve_confidence_head assert resolve_confidence_head(_Bare(None)) is None assert resolve_confidence_head(_Wrapper(None)) is None assert resolve_confidence_head(object()) is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py index 2c0a83f58351..2160e1d0ba49 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py @@ -54,40 +54,22 @@ def test_restricted_answer_is_always_realisable(tiers): f"budget {n} implies tier {(n // bs) + 1}, not in {tiers}") -def test_agrees_with_the_uniform_scorer_at_a_degenerate_batch(): - """With one row the two taus coincide, so the rungs must match. - - A single request makes the uniform allocation and the top-k allocation the - same set, which is the one case where any disagreement would be a plain bug - rather than the intended improvement. - """ - table = _table() - rng = np.random.default_rng(7) - tiers = [1, 2, 5] - for _ in range(100): - survival = np.sort(rng.random((1, 5)), axis=1)[:, ::-1] - n = compute_verify_token_budget( - survival=survival, num_gen_requests=1, cost_table=table, - min_verify_len=1, allowed_lens=tiers) - uniform = budget_argmax_over_uniform_lens( - survival=survival, num_gen_requests=1, cost_table=table, - allowed_lens=tiers, min_verify_len=1) - assert (n // 1) + 1 == uniform - - def test_restricted_never_scores_worse_than_the_uniform_choice(): """Same grid, better numerator: the chosen rung's realised theta wins. Both routes evaluate the same discrete cost points, so the restricted form cannot be worse under the tau that the step actually collects. This is the property that makes the swap safe -- if it ever failed, the change would be - trading a modelling error for a real one. + trading a modelling error for a real one. At bs=1 the uniform and top-k + allocations are the same set, so there the two scorers must agree on the + rung outright: any disagreement is a plain bug, not the intended + improvement. """ rng = np.random.default_rng(31337) table = _table() tiers = [1, 2, 5] for _ in range(300): - bs = int(rng.integers(2, 17)) + bs = int(rng.integers(1, 17)) survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] def realised_theta(tier: int) -> float: @@ -108,23 +90,7 @@ def realised_theta(tier: int) -> float: assert realised_theta(chosen) >= realised_theta(uniform) - 1e-12, ( f"restricted picked tier {chosen} whose realised theta is below " f"the uniform scorer's tier {uniform}") - - -def test_unrestricted_form_is_unchanged(): - """Omitting allowed_lens must still answer the unconstrained question. - - The existing tests use that form as the reference, and it is still the right - thing to ask when sizing a ladder rather than executing one. - """ - table = _table() - rng = np.random.default_rng(99) - survival = np.sort(rng.random((6, 5)), axis=1)[:, ::-1] - free = compute_verify_token_budget( - survival=survival, num_gen_requests=6, cost_table=table, - min_verify_len=1) - assert 0 <= free <= 6 * 4 - tied = compute_verify_token_budget( - survival=survival, num_gen_requests=6, cost_table=table, - min_verify_len=1, allowed_lens=[1, 2, 5]) - # The restricted answer is one of the grid points; the free one need not be. - assert tied in {0, 6, 24} + if bs == 1: + assert chosen == uniform, ( + "at bs=1 the uniform and top-k allocations are the same set; " + "the two scorers must pick the same rung") diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py index 20f892fc5872..04b8d5c0dce6 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py @@ -33,17 +33,14 @@ def _planner(tiers=(1, 2, 5)): return DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=list(tiers)) -def test_queued_pin_is_not_yet_in_force(): - """Queuing must not change this rank's behaviour before the group agrees.""" +def test_adopting_the_agreed_value_applies_it(): + """Queuing must not change this rank's behaviour before the group agrees; + adopting the agreed value applies it and consumes the pending request.""" planner = _planner() assert planner.request_verify_len_pin(2) == 2 + # Queued, not yet in force. assert planner._forced_verify_len is None assert planner.pending_verify_len_pin() == 2 - - -def test_adopting_the_agreed_value_applies_it(): - planner = _planner() - planner.request_verify_len_pin(2) planner.adopt_verify_len_pin(2) assert planner._forced_verify_len == 2 # Consumed: the next step must not re-broadcast a pin nobody asked for. @@ -58,17 +55,13 @@ def test_a_rank_that_queued_nothing_still_adopts_the_group_value(): assert peer._forced_verify_len == 5 -def test_minus_one_means_nobody_asked_and_leaves_the_pin_alone(): +def test_zero_clears_the_pin(): + """Wire protocol: -1 means nobody asked (leave the pin alone); clearing + needs its own value, 0, because None cannot travel in an int payload.""" planner = _planner() planner.adopt_verify_len_pin(2) planner.adopt_verify_len_pin(-1) assert planner._forced_verify_len == 2 - - -def test_zero_clears_the_pin(): - """Clearing needs its own wire value: None cannot travel in an int payload.""" - planner = _planner() - planner.adopt_verify_len_pin(2) assert planner.request_verify_len_pin(None) is None assert planner.pending_verify_len_pin() == 0 planner.adopt_verify_len_pin(0) @@ -80,14 +73,10 @@ def test_an_uncaptured_length_is_refused_at_the_call_site(): planner = _planner(tiers=(1, 2, 5)) with pytest.raises(ValueError, match="captured tier ladder"): planner.request_verify_len_pin(3) - assert planner.pending_verify_len_pin() == -1 - assert planner._forced_verify_len is None - - -def test_a_length_outside_the_bounds_is_refused(): - planner = _planner() with pytest.raises(ValueError, match="outside"): planner.request_verify_len_pin(9) + assert planner.pending_verify_len_pin() == -1 + assert planner._forced_verify_len is None def test_pinned_steps_hand_every_request_the_same_window(): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py index 609dfa3910d1..ba186f4d942e 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_worker.py @@ -58,39 +58,29 @@ def test_metadata_buffer_and_layer_lookup(): # sorted, O(1) membership assert meta.is_layer_capture(58) and meta.is_layer_capture(60) assert not meta.is_layer_capture(0) and not meta.is_layer_capture(61) + # prepare() publishes contiguous batch indices for the current batch. + meta.request_ids = [7, 3, 5] + meta.prepare() + assert meta.batch_indices_cuda[:3].tolist() == [0, 1, 2] def test_metadata_capture_plain_hidden(): - """A [num_tokens, hidden] capture is stored at the layer's slice as-is.""" + """Capture routing: a [num_tokens, hidden] capture is stored at the layer's + slice as-is; a flattened mHC residual [N, hc_mult*hidden] is reduced by mean + over the hc streams; an unlisted layer writes nothing.""" meta = _make_metadata() + meta.captured_hidden_states.zero_() hs = torch.randn(4, HIDDEN, device="cuda", dtype=torch.bfloat16) meta.maybe_capture_hidden_states(59, hs) # layer 59 -> capture index 1 - got = meta.get_hidden_states(4) - assert torch.equal(got[:, HIDDEN : 2 * HIDDEN], hs) - - -def test_metadata_capture_hc_mean_reduction(): - """A flattened mHC residual [N, hc_mult*hidden] is reduced by mean over hc.""" - meta = _make_metadata() mhc = torch.randn(4, HC_MULT * HIDDEN, device="cuda", dtype=torch.bfloat16) meta.maybe_capture_hidden_states(58, mhc) # layer 58 -> capture index 0 - expected = mhc.reshape(4, HC_MULT, HIDDEN).mean(dim=1) + meta.maybe_capture_hidden_states(10, torch.randn(4, HIDDEN, device="cuda")) # unlisted got = meta.get_hidden_states(4) + assert torch.equal(got[:, HIDDEN : 2 * HIDDEN], hs) + expected = mhc.reshape(4, HC_MULT, HIDDEN).mean(dim=1) assert torch.equal(got[:, 0:HIDDEN], expected) - - -def test_metadata_no_capture_for_unlisted_layer(): - meta = _make_metadata() - meta.captured_hidden_states.zero_() - meta.maybe_capture_hidden_states(10, torch.randn(4, HIDDEN, device="cuda")) - assert torch.count_nonzero(meta.get_hidden_states(4)) == 0 - - -def test_metadata_prepare_batch_indices(): - meta = _make_metadata() - meta.request_ids = [7, 3, 5] - meta.prepare() - assert meta.batch_indices_cuda[:3].tolist() == [0, 1, 2] + # The unlisted layer landed nowhere: the remaining slice stays zero. + assert torch.count_nonzero(got[:, 2 * HIDDEN :]) == 0 def _make_worker(enable_confidence_scheduling=False, verify_len_tiers=None): @@ -108,12 +98,6 @@ def _make_worker(enable_confidence_scheduling=False, verify_len_tiers=None): return DSparkWorker(cfg, Mapping()) -def test_worker_reads_confidence_flag_once(): - """The flag must be run-constant: a per-step read could diverge the graph.""" - assert _make_worker().return_confidence is False - assert _make_worker(enable_confidence_scheduling=True).return_confidence is True - - def _fake_draft_model(num_stages=3, window_size=128, head_dim=64): return types.SimpleNamespace( num_stages=num_stages, @@ -143,14 +127,11 @@ def test_worker_lazy_init_window_buffers(): worker._lazy_init(dm, meta) assert id(worker._kv_windows) == buf_id - -def test_worker_rejects_mismatched_block_size(): - worker = _make_worker() - draft_model = _fake_draft_model() - draft_model.block_size = 4 - + # A draft model whose block_size disagrees with the worker is refused. + bad = _fake_draft_model() + bad.block_size = 4 with pytest.raises(ValueError, match="block_size must equal worker max_draft_len"): - worker._lazy_init(draft_model, _make_metadata()) + _make_worker()._lazy_init(bad, _make_metadata()) def test_worker_slot_assignment_and_reset(): @@ -235,7 +216,8 @@ def write_context_windows(self, hidden, positions, windows): def test_prepare_builds_batch_to_slot_on_batched_path(): - """prepare() mirrors the host slot map into _batch_to_slot (default batched path).""" + """prepare() mirrors the host slot map into _batch_to_slot (default batched + path) and returns a dropped request's slot to the free pool.""" worker = _make_worker() meta = _make_metadata(max_num_requests=4) worker._lazy_init(_fake_draft_model(), meta) # batched is the default @@ -244,24 +226,13 @@ def test_prepare_builds_batch_to_slot_on_batched_path(): # Assign slots for two requests (as the prefill path would). sa = worker._assign_slot(100, reset=True) sb = worker._assign_slot(101, reset=True) + worker._ctx_len[sa] = 17 meta.request_ids = [101, 100] meta.prepare() # Mirror reflects request-order -> slot. assert worker._batch_to_slot[:2].tolist() == [sb, sa] - -def test_prepare_frees_stale_slots_on_batched_path(): - """A request that drops out of the batch returns its slot to the free pool.""" - worker = _make_worker() - meta = _make_metadata(max_num_requests=4) - worker._lazy_init(_fake_draft_model(), meta) - meta._dspark_worker = worker - - sa = worker._assign_slot(100, reset=True) - worker._assign_slot(101, reset=True) - worker._ctx_len[sa] = 17 - # Only request 101 survives; 100's slot must be freed + cleared. meta.request_ids = [101] meta.prepare() @@ -489,7 +460,7 @@ def forward_batched(main_hidden, bonus, start_pos, **kwargs): return dm, seen -@pytest.mark.parametrize("runtime_draft_len", [3, 1]) +@pytest.mark.parametrize("runtime_draft_len", [3, None]) def test_gen_draft_gathers_hidden_with_the_runtime_stride(runtime_draft_len): """The gather stride must follow runtime_draft_len, not max_draft_len. @@ -497,6 +468,8 @@ def test_gen_draft_gathers_hidden_with_the_runtime_stride(runtime_draft_len): Striding by ``max_draft_len + 1`` instead walks into the *next* request's hidden states as soon as the scheduler trims the draft length -- the draft then conditions on another request's context, with no error anywhere. + ``runtime_draft_len=None`` pins the fallback: metadata without a runtime + length keeps the pre-existing max_draft_len stride. """ num_gens, hidden = 4, HIDDEN * NCAP worker = _make_worker() @@ -506,7 +479,7 @@ def test_gen_draft_gathers_hidden_with_the_runtime_stride(runtime_draft_len): worker._batch_to_slot[:num_gens] = torch.arange(num_gens, device="cuda") # Row r of request g is tagged with g so a cross-request read is visible. - stride = runtime_draft_len + 1 + stride = (runtime_draft_len or worker.max_draft_len) + 1 captured = torch.zeros(num_gens * stride, hidden, device="cuda") for g in range(num_gens): captured[g * stride : (g + 1) * stride] = float(g) @@ -538,82 +511,39 @@ def test_gen_draft_gathers_hidden_with_the_runtime_stride(runtime_draft_len): ) -def test_gen_draft_stride_falls_back_to_max_draft_len(): - """Metadata without runtime_draft_len keeps the pre-existing behavior.""" - num_gens, hidden = 3, HIDDEN * NCAP - worker = _make_worker() - dm, seen = _stride_probe_draft_model() - meta = _make_metadata(max_num_requests=8, max_num_tokens=256) - worker._lazy_init(dm, meta) - worker._batch_to_slot[:num_gens] = torch.arange(num_gens, device="cuda") - - stride = worker.max_draft_len + 1 - captured = torch.zeros(num_gens * stride, hidden, device="cuda") - for g in range(num_gens): - captured[g * stride : (g + 1) * stride] = float(g) - - meta.runtime_draft_len = None - meta.get_hidden_states = lambda _n: captured - attn = types.SimpleNamespace(num_ctx_tokens=0, num_seqs=num_gens, num_contexts=0) - nacc = torch.ones(num_gens, dtype=torch.int32, device="cuda") - accepted = torch.zeros(num_gens, worker.max_draft_len + 1, dtype=torch.int32, device="cuda") - - worker._draft_gen_block_batched( - dm, - meta, - attn, - accepted, - nacc, - num_contexts=0, - batch_size=num_gens, - total_target_tokens=captured.shape[0], - ) - assert seen["main_hidden"][:, 0].tolist() == [float(g) for g in range(num_gens)] - - def test_confidence_rows_are_looked_up_by_slot(): """The confidence buffer is slot-indexed; a batch rarely owns slots 0..G-1. The snapshot the scheduler reads is one iteration old, and joins/departures reshuffle the batch in between, so batch position is not a usable key. ``confidence_row_for`` re-associates each request with the row its own draft - wrote. + wrote. A request with no slot yet falls back to a dedicated, permanently + neutral row ("verify the whole block") past every real slot -- the scratch + row is not safe for that, since padded and unknown requests write through it. """ worker = _make_worker(enable_confidence_scheduling=True) dm, _ = _stride_probe_draft_model() worker._lazy_init(dm, _make_metadata(max_num_requests=8)) assert worker._confidence_logits is not None - # Tag every row with its index, then hand three requests explicit slots. - for s in range(worker._confidence_logits.shape[0]): + # Unscored request: the permanently-neutral row, past every real slot AND + # past the scratch row, and it reads as ~certain acceptance. + neutral = worker.confidence_row_for(12345) + assert neutral == worker._neutral_conf_row + assert neutral > worker._scratch_slot + buf = worker.staged_confidence_buffer() + assert torch.sigmoid(buf[neutral]).min().item() > 0.999 + + # Tag every real row with its index, then hand three requests explicit slots. + for s in range(worker._scratch_slot): worker._confidence_logits[s] = float(s) worker._req_to_slot = {100: 5, 101: 2, 102: 7} rows = [worker.confidence_row_for(r) for r in (100, 101, 102)] assert rows == [5, 2, 7] - buf = worker.staged_confidence_buffer() assert buf[rows][:, 0].tolist() == [5.0, 2.0, 7.0] -def test_unscored_request_maps_to_a_permanently_neutral_row(): - """A request with no slot yet must read as "verify the whole block". - - The scratch row is not safe for this -- padded and unknown requests write - through it -- so there is a dedicated row past every real slot that nothing - ever writes. - """ - worker = _make_worker(enable_confidence_scheduling=True) - dm, _ = _stride_probe_draft_model() - worker._lazy_init(dm, _make_metadata(max_num_requests=8)) - - neutral = worker.confidence_row_for(12345) - assert neutral == worker._neutral_conf_row - # Past every real slot AND past the scratch row. - assert neutral > worker._scratch_slot - buf = worker.staged_confidence_buffer() - assert torch.sigmoid(buf[neutral]).min().item() > 0.999 - - def test_a_freshly_assigned_slot_starts_neutral(): """Otherwise a new request's first decision trims on a dead draft's scores. diff --git a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py index a5084293c3d9..eff4f9884e5e 100644 --- a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py +++ b/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py @@ -39,7 +39,9 @@ def _stats(**reasons): def test_continuous_batching_misses_are_not_blamed_on_ragged(): - """A rank with context requests is normal, not a ragged defect.""" + """A rank with context requests is normal, not a ragged defect -- and a + fully clean run (no misses at all) passes the same gate.""" + _stats().assert_ragged_active(require_trim=True) _stats(peer_not_gen_only=9).assert_ragged_active(require_trim=True) @@ -76,10 +78,6 @@ def test_unattributed_misses_are_a_broken_probe_not_a_clean_run(): stats.assert_ragged_active(require_trim=True) -def test_clean_run_passes(): - _stats().assert_ragged_active(require_trim=True) - - def test_ramp_up_steps_do_not_count_toward_the_gate_floor(): """A batch of one is uniform by construction, not a scheduler failure. diff --git a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py index fbfaca610386..144de6a6c174 100644 --- a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py +++ b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py @@ -80,32 +80,21 @@ def _require_cuda(): @pytest.mark.parametrize( - "env_value, expected", + "env_value, expected_float, expected_int", [ - ("0", 0.0), - ("3", 3.0), - ("2.6", 2.6), - ("0.5", 0.5), - ("not-a-number", 0.0), + ("0", 0.0, 0), + ("3", 3.0, 3), + ("2.6", 2.6, 0), + ("0.5", 0.5, 0), + ("not-a-number", 0.0, 0), ], ) -def test_get_force_num_accepted_tokens_float(env_value, expected): +def test_get_force_num_accepted_tokens_float(env_value, expected_float, expected_int): + """The float helper parses fractional rates; the int helper (used by the + 2-model path) must keep its original behavior on the same values.""" with patch.dict(os.environ, {FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR: env_value}): - assert get_force_num_accepted_tokens_float() == pytest.approx(expected) - - -@pytest.mark.parametrize( - "env_value, expected", - [ - ("0", 0), - ("3", 3), - ("2.6", 0), - ], -) -def test_get_force_num_accepted_tokens_int_unchanged(env_value, expected): - """The int helper must keep its original behavior (used by 2-model).""" - with patch.dict(os.environ, {FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR: env_value}): - assert get_force_num_accepted_tokens() == expected + assert get_force_num_accepted_tokens_float() == pytest.approx(expected_float) + assert get_force_num_accepted_tokens() == expected_int # ---------------- _apply_force_accepted_tokens semantics -------------------- @@ -123,57 +112,50 @@ def test_zero_value_is_noop(): assert worker._force_accept_rng_counter is None -@pytest.mark.parametrize("value", [1.0, 2.0, 4.0]) -def test_integer_value_matches_legacy_behavior(value): +@pytest.mark.parametrize( + "value, runtime_draft_len", + [ + (1.0, 4), + (2.0, 4), + (4.0, 4), + # 10 draft tokens requested but only 2 available -> capped at 3. + (10.0, 2), + # Fractional but no room: max_total = 3, base_total = min(10, 3) = 3, + # so ``base_total < max_total`` is False -> all 3, no RNG draw. + (9.5, 2), + ], +) +def test_integer_value_matches_legacy_behavior(value, runtime_draft_len): _require_cuda() worker = _make_worker(value) out = worker._apply_force_accepted_tokens( - _make_input(batch_size=8), num_contexts=0, runtime_draft_len=4 + _make_input(batch_size=8), num_contexts=0, runtime_draft_len=runtime_draft_len ) - expected = min(int(value) + 1, 4 + 1) + expected = min(int(value) + 1, runtime_draft_len + 1) assert torch.all(out == expected) - # Pure-integer path must not allocate the RNG pool. + # The deterministic (non-random) path must not allocate the RNG pool. assert worker._force_accept_rng_pool is None -def test_integer_value_caps_at_runtime_draft_len_plus_one(): - _require_cuda() - worker = _make_worker(10.0) - out = worker._apply_force_accepted_tokens( - _make_input(batch_size=4), num_contexts=0, runtime_draft_len=2 - ) - # 10 draft tokens requested but only 2 available → all 3 (= 2 + target). - assert torch.all(out == 3) - - -def test_fractional_only_emits_two_values(): - _require_cuda() - worker = _make_worker(2.6) - seen = set() - for _ in range(50): - out = worker._apply_force_accepted_tokens( - _make_input(batch_size=64), num_contexts=0, runtime_draft_len=4 - ) - seen.update(out.unique().tolist()) - # Either 2 draft + target = 3, or 3 draft + target = 4. - assert seen == {3, 4} - - def test_fractional_distribution_matches_target_probability(): - """Average over many draws should converge to the target fraction.""" + """Fractional draws emit only base/base+1 and average to the target rate.""" _require_cuda() worker = _make_worker(2.6) target_frac = 0.6 n_iters = 200 batch = 64 + seen = set() extra_count = 0 total = 0 for _ in range(n_iters): out = worker._apply_force_accepted_tokens( _make_input(batch_size=batch), num_contexts=0, runtime_draft_len=4 ) + seen.update(out.unique().tolist()) extra_count += int((out == 4).sum().item()) total += batch + # Either 2 draft + target = 3, or 3 draft + target = 4 -- never anything else. + assert seen == {3, 4} measured = extra_count / total # Pool-based RNG is deterministic; the empirical mean over ~13k draws # is tight against 0.6. @@ -182,19 +164,6 @@ def test_fractional_distribution_matches_target_probability(): ) -def test_fractional_capped_when_no_room(): - """If ``int_part + 1`` already saturates the cap, no extra is granted.""" - _require_cuda() - worker = _make_worker(9.5) - out = worker._apply_force_accepted_tokens( - _make_input(batch_size=8), num_contexts=0, runtime_draft_len=2 - ) - # max_total = runtime_draft_len + 1 = 3, base_total = min(10, 3) = 3. - # frac_part > 0 but ``base_total < max_total`` is False → all 3, no RNG. - assert torch.all(out == 3) - assert worker._force_accept_rng_pool is None - - def test_num_contexts_offset_is_respected(): """Context (prefill) rows must be left untouched by the override.""" _require_cuda() 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 bdcc3fd77992..d43882ec9a83 100644 --- a/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py +++ b/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py @@ -55,12 +55,22 @@ def test_local_value_used_when_no_group_sync(): def test_group_override_pulls_greedy_rank_onto_advanced_path(): - # This rank's batch is all-greedy, but another rank in the LM-head-TP - # group has a sampling request: the group AND (False) must win so the - # whole group takes the advanced path together. - meta = _fake_meta(group_all_greedy_sample=False) - _scan(meta, [_fake_request(), _fake_request()]) - assert meta.is_all_greedy_sample is False + # (group value, capture force) -> final flag. This rank's batch is + # all-greedy in every case; the group AND (False) must win so the whole + # LM-head-TP group takes the advanced path together, group-wide True keeps + # greedy, and the warmup capture force composes without changing the + # outcome (the synced value is derived from capture-forced locals). + cases = [ + (False, False, False), + (True, False, True), + (False, True, False), + ] + for group_value, force_capture, expected in cases: + meta = _fake_meta(group_all_greedy_sample=group_value, + force_capture=force_capture) + _scan(meta, [_fake_request(), _fake_request()]) + assert meta.is_all_greedy_sample is expected, (group_value, + force_capture) def test_group_override_survives_rescan(): @@ -71,18 +81,3 @@ def test_group_override_survives_rescan(): for _ in range(3): _scan(meta, [_fake_request()]) assert meta.is_all_greedy_sample is False - - -def test_group_override_true_keeps_greedy(): - meta = _fake_meta(group_all_greedy_sample=True) - _scan(meta, [_fake_request()]) - assert meta.is_all_greedy_sample is True - - -def test_capture_override_composes_with_group_sync(): - # Warmup forces the advanced variant to capture its CUDA graph; the group - # value is derived from capture-forced locals (all False), so the final - # flag stays False regardless of composition order. - meta = _fake_meta(group_all_greedy_sample=False, force_capture=True) - _scan(meta, [_fake_request()]) - assert meta.is_all_greedy_sample is False diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py index f398fa1de3c9..82efadf6679f 100644 --- a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py +++ b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py @@ -58,8 +58,6 @@ def test_prepare_buffers_allocates_when_enabled(): assert m.full_draft_probs is None assert m.draft_probs_vocab_size == V - -def test_prepare_buffers_allocates_full_draft_probs_on_vocab_mismatch(): # Distinct draft vocab: full_draft_probs (d2t-expanded) is allocated. m = _alloc_meta(draft_vocab_size=V - 1) SpecMetadata.prepare_rejection_sampling_buffers(m) @@ -237,38 +235,33 @@ def _rej(logits, draft_tokens, draft_probs, num_contexts, batch_size, spec_metad return w, calls -def test_accept_dispatch_fails_closed_on_malformed_buffers(): - w, calls = _make_worker() - # Buffers are malformed (draft_probs None) -> _rejection_buffers_valid must - # return False, so acceptance falls back to base and the rejection kernel is - # skipped. - meta = _dispatch_meta(draft_probs=None) - draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") - logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") - out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) - assert out == ("base", None) - assert calls["base"] == 1 and calls["rejection"] == 0 - - -def test_accept_dispatch_routes_to_rejection_on_valid_state(): - w, calls = _make_worker() - meta = _dispatch_meta() # valid buffers - draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") - logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") - out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) - assert out == ("rejection", None) - assert calls["rejection"] == 1 and calls["base"] == 0 - - -def test_accept_dispatch_base_when_all_greedy(): +@pytest.mark.parametrize( + "meta_over, expected_path", + [ + # Malformed buffers (draft_probs None): _rejection_buffers_valid must + # return False, so acceptance fails closed to base and the rejection + # kernel is skipped. + (dict(draft_probs=None), "base"), + # Valid buffers: the rejection path runs. + (dict(), "rejection"), + # All-greedy batch: rejection is bypassed regardless of buffers. + (dict(is_all_greedy_sample=True), "base"), + ], + ids=[ + "fails_closed_on_malformed_buffers", + "routes_to_rejection_on_valid_state", + "base_when_all_greedy", + ], +) +def test_accept_dispatch(meta_over, expected_path): w, calls = _make_worker() - # All-greedy batch: rejection is bypassed regardless of buffers. - meta = _dispatch_meta(is_all_greedy_sample=True) + meta = _dispatch_meta(**meta_over) draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) - assert out == ("base", None) - assert calls["base"] == 1 and calls["rejection"] == 0 + assert out == (expected_path, None) + other = "rejection" if expected_path == "base" else "base" + assert calls[expected_path] == 1 and calls[other] == 0 if __name__ == "__main__": diff --git a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py index 4b26ffffe0bb..898699b69e91 100644 --- a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py +++ b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py @@ -31,11 +31,9 @@ def _cfg(**kwargs): def test_ragged_without_a_cost_table_is_rejected(): with pytest.raises(ValueError, match="requires a profiled cost table"): _cfg() - - -def test_ragged_with_an_sps_table_is_accepted(): - cfg = _cfg(confidence_sps_table_path="/tmp/table.json") - assert cfg.enable_ragged_verify + # The same gate must open once a table is named -- a gate that rejects + # everything would pass the raise above. + assert _cfg(confidence_sps_table_path="/tmp/table.json").enable_ragged_verify def test_confidence_scheduling_without_ragged_is_rejected(): From 7327a93c163c488b08e49bc43d998ea22da43832 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 5 Aug 2026 07:28:58 -0700 Subject: [PATCH 095/177] [None][chore] Test only what ships: drop tool and observability test suites Scope the unit suite to production correctness. Deleted outright, on the principle that dev tooling and diagnostics do not need their own tests: - the SPS profiler suite (the profiler is a measurement tool; its output is reviewed operationally) and the SGLang budget-parity metric harness; - observability plumbing: snapshot instruments, iter-log sampling, reconcile guards, and the STS calibration-collection suites; - measurement-support features: the runtime verify-length pin and synthetic forced acceptance; - config-validation and cost-model micro-suites (interpolation, tier budget quantisation). Preserved by moving, not deleting: - the two confidence-follows-budget scheduling guards from test_dspark_planner_trims.py into test_dspark_confidence_schedule.py (they were the only remaining guard of that property); - the three graph-key regressions (padding-ladder fit, key reads the fitted bucket -- the ragged-IMA family -- and capture publishes its bucket) from the miss-attribution file, which is renamed test_dspark_graph_key.py with its observability half removed. One comment in dspark_planner.py pointed at a deleted test; made self-contained. Remaining DSpark suite after the cut: 223 tests, all passing on B300. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/speculative/dspark_planner.py | 3 +- .../test_dspark_budget_parity_sglang.py | 232 ------- .../test_dspark_confidence_schedule.py | 40 ++ .../test_dspark_cost_interpolation.py | 61 -- .../test_dspark_iter_log_samples.py | 159 ----- .../test_dspark_reconcile_guards.py | 99 --- .../test_dspark_snapshot_instruments.py | 110 ---- .../hw_agnostic/test_dspark_sps_profiler.py | 620 ------------------ .../test_dspark_sts_collection_guards.py | 135 ---- .../test_dspark_sts_ring_pairing.py | 168 ----- .../test_dspark_tier_aligned_budget.py | 96 --- .../hw_agnostic/test_dspark_verify_len_pin.py | 97 --- ...ttribution.py => test_dspark_graph_key.py} | 104 +-- .../speculative/test_dspark_planner_trims.py | 75 --- .../speculative/test_force_accepted_tokens.py | 363 ---------- .../test_dspark_ragged_requires_table.py | 56 -- 16 files changed, 48 insertions(+), 2370 deletions(-) delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py rename tests/unittest/_torch/speculative/{test_dspark_graph_miss_attribution.py => test_dspark_graph_key.py} (58%) delete mode 100644 tests/unittest/_torch/speculative/test_dspark_planner_trims.py delete mode 100644 tests/unittest/_torch/speculative/test_force_accepted_tokens.py delete mode 100644 tests/unittest/llmapi/test_dspark_ragged_requires_table.py diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py index 494b7354f3c7..6afd63b2e039 100644 --- a/tensorrt_llm/_torch/speculative/dspark_planner.py +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -120,8 +120,7 @@ class SpsCostTable: breakpoint's price, making tokens look free right before a riser, and the argmax then over-spends -- measured live as the planner buying the full block on ~95% of decisions because its cost ratio between tiers collapsed - below the survival ratio (see test_dspark_cost_interpolation for the - pinned numbers). A table that wants a flat shelf must MEASURE the shelf: + below the survival ratio. A table that wants a flat shelf must MEASURE the shelf: two breakpoints with equal values. Queries outside the measured range clamp to the end values; the caller is expected to bound the budget so the high clamp does not happen silently. diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py deleted file mode 100644 index ac989a423428..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py +++ /dev/null @@ -1,232 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Decision-level parity with the SGLang budget planner this was ported from. - -A throughput gain alone does not validate the port: the wrong economics can -still move the needle in the right direction. This pins the DECISION function -against a line-for-line reference copy of SGLang's -``compute_verify_token_budget`` + ``_additive_step_time_tensor`` -(python/sglang/srt/speculative/dspark_components/dspark_planner.py), so a -divergence in either port shows up as an argmax mismatch here rather than as -an unexplained fleet-level regression. - -The two planners parameterize the same optimization differently: - -- SGLang admits EVERY position (including position 0) as a budget candidate - and its floor is one bonus token per request: total M = bs + n. -- TensorRT-LLM grants positions below ``min_verify_len`` for free (their - survival rides in the tau base) and only positions at or beyond it are - schedulable: total M = bs*(min_verify_len+1) + m. - -For ``min_verify_len=1`` these describe the same curve shifted by ``bs``: -survival is non-increasing along a row, so the top-(bs+m) SGLang candidates -are exactly the bs position-0 entries plus the top-m deeper entries, giving -``tau_sgl(bs+m) == tau_trt(m)`` and ``M_sgl(bs+m) == M_trt(m)``. The argmaxes -must therefore agree up to that shift whenever SGLang admits at least the -full position-0 set. Known, deliberate divergences are pinned separately: -SGLang filters candidates by ``survival_eps`` inside the budget function -(TensorRT-LLM applies eps at allocation time instead), and SGLang can choose -sub-floor budgets (windows of zero draft positions) that TensorRT-LLM's -``min_verify_len=1`` floor excludes. -""" - -import numpy as np -import pytest -import torch - -from tensorrt_llm._torch.speculative.dspark_planner import ( - SpsCostTable, compute_verify_token_budget) - -# -------------------------------------------------------------------------- -# Reference implementation, copied from SGLang dspark_components/dspark_planner -# (additive-table path) with only cosmetic renames. Do not "fix" this copy: -# its job is to be what SGLang runs. -# -------------------------------------------------------------------------- - - -def _sgl_interp_clamped(xs, ys, x: float) -> float: - xs = torch.tensor(xs, dtype=torch.float64) - ys = torch.tensor(ys, dtype=torch.float64) - x_t = torch.tensor(float(x), dtype=torch.float64).clamp_(xs[0], xs[-1]) - hi = torch.bucketize(x_t, xs, right=True).clamp_(1, xs.numel() - 1) - lo = hi - 1 - span = (xs[hi] - xs[lo]).clamp_(min=1e-9) - frac = (x_t - xs[lo]) / span - return float(ys[lo] + frac * (ys[hi] - ys[lo])) - - -def _sgl_additive_step_time(table, num_requests: int, num_budgets: int): - floor = table["bias_seconds"] + _sgl_interp_clamped( - table["bs_probes"], table["alpha_seconds"], float(num_requests)) - m_probes = torch.tensor(table["m_probes"], dtype=torch.float64) - theta_vals = torch.tensor(table["theta_seconds"], dtype=torch.float64) - m = (num_requests + torch.arange(num_budgets, dtype=torch.float64)).clamp_( - min=float(table["m_probes"][0]), max=float(table["m_probes"][-1])) - hi = torch.bucketize(m, m_probes, right=True).clamp_(1, m_probes.numel() - 1) - lo = hi - 1 - span = (m_probes[hi] - m_probes[lo]).clamp_(min=1e-9) - frac = (m - m_probes[lo]) / span - theta_at_m = theta_vals[lo] + frac * (theta_vals[hi] - theta_vals[lo]) - return floor + theta_at_m - - -def _sgl_compute_verify_token_budget(*, history_survival_probs, table, - max_verify_len, survival_eps): - num_requests = history_survival_probs.shape[0] - candidates = history_survival_probs[:, :max_verify_len].flatten() - candidates = candidates[candidates >= survival_eps].to(torch.float64) - candidates_sorted = torch.sort(candidates, descending=True).values - prefix_sum = torch.cumsum(candidates_sorted, dim=0) - tau_star = num_requests + torch.cat( - [torch.zeros(1, dtype=torch.float64), prefix_sum]) - step_time = _sgl_additive_step_time(table, int(num_requests), - int(tau_star.numel())) - theta = tau_star / step_time - return int(torch.argmax(theta)) - - -# -------------------------------------------------------------------------- -# One certified-table shape shared by both sides. TensorRT-LLM stores ms, -# SGLang stores seconds; the argmax is scale-invariant but the tables are -# built with the honest factor anyway. -# -------------------------------------------------------------------------- - -M_PROBES = (64, 96, 128, 192, 384, 512, 768, 1536) -THETA_MS = (9.68, 11.08, 11.95, 13.95, 22.09, 23.85, 35.61, 105.64) -BIAS_MS = 25.244443 -BS_PROBES = (32, 64, 128, 256) -ALPHA_MS = (0.0, 2.177477, 11.462212, 19.342661) -BLOCK = 5 - -TRT_TABLE = SpsCostTable(token_counts=M_PROBES, step_time_ms=THETA_MS, - fixed_overhead_ms=BIAS_MS, batch_sizes=BS_PROBES, - batch_overhead_ms=ALPHA_MS) -SGL_TABLE = { - "m_probes": M_PROBES, - "theta_seconds": tuple(v / 1e3 for v in THETA_MS), - "bias_seconds": BIAS_MS / 1e3, - "bs_probes": BS_PROBES, - "alpha_seconds": tuple(v / 1e3 for v in ALPHA_MS), -} - - -def _inversion_free_survival(rng, bs): - """Survivals where every position-0 beats every deeper candidate. - - The exact-parity theorem needs this: SGLang admits position 0 as a budget - candidate while TensorRT-LLM grants it in the floor, so the two describe - the same curve (shifted by ``bs``) only when the top-(bs+m) SGLang set is - exactly "all of position 0, plus the top-m deeper entries". A weak row's - position-0 falling below a strong row's position-1 breaks that -- see - test_starving_weak_rows_is_a_known_gap for what SGLang does with it. - """ - conf0 = rng.uniform(0.85, 0.99, size=(bs, 1)) - deeper = rng.uniform(0.25, 0.80, size=(bs, BLOCK - 1)) - return np.cumprod(np.concatenate([conf0, deeper], axis=1), axis=1) - - -@pytest.mark.parametrize("bs", [1, 3, 8, 64, 128, 252, 256]) -def test_budget_argmax_matches_sglang_up_to_the_floor_shift(bs): - """Same survival, same table, eps off, no inversions: shift-exact parity.""" - rng = np.random.default_rng(17 + bs) - for _ in range(5): - surv = _inversion_free_survival(rng, bs) - sgl_n = _sgl_compute_verify_token_budget( - history_survival_probs=torch.tensor(surv, dtype=torch.float32), - table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.0) - trt_m = compute_verify_token_budget( - survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, - min_verify_len=1) - if sgl_n >= bs: - assert trt_m == sgl_n - bs, ( - f"bs={bs}: SGLang admits {sgl_n} candidates " - f"(= floor {bs} + {sgl_n - bs}), TRT-LLM budget {trt_m}") - else: - # SGLang went below the one-draft floor; the closest budget - # TensorRT-LLM can express under min_verify_len=1 is zero. - assert trt_m == 0 - - -def test_live_pooled_survivals_park_on_the_breakpoint_together(): - """The campaign's operating point, both consumers, both argmax modes. - - With interpolated segments the continuous optimum parks where the theta - slope jumps: buying position-2 survivals (0.502/token) beats the - 512..768 slope (0.046 ms/token) but not the 768..1536 one (0.091), so - both planners stop at exactly M = 768 -- TensorRT-LLM at budget - 768 - 2*252 = 264, SGLang at 768 - 252 = 516. Tier alignment then rounds - down to the rung-2 budget the executor can capture. - """ - bs = 252 - surv = np.tile([0.771, 0.693, 0.502, 0.360, 0.315], (bs, 1)) - sgl_n = _sgl_compute_verify_token_budget( - history_survival_probs=torch.tensor(surv, dtype=torch.float32), - table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.0) - trt_m = compute_verify_token_budget( - survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, - min_verify_len=1) - assert trt_m == 768 - 2 * bs - assert sgl_n - bs == trt_m - tiered = compute_verify_token_budget( - survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, - min_verify_len=1, allowed_lens=[1, 2, 5]) - assert tiered == bs * 1 # rung-2: one scheduled position past the floor - - -def test_starving_weak_rows_is_a_known_gap(): - """SGLang can verify NOTHING for a hopeless request; TRT-LLM cannot. - - SGLang's budget admits position-0 candidates individually, so a row whose - draft is dead on arrival can be starved to a zero-length window and its - would-be tokens spent on a strong row's deep positions. TensorRT-LLM's - ``min_verify_len=1`` floor buys position 0 for every request - unconditionally. On a batch that mixes both kinds, SGLang's achievable - Theta is therefore strictly better -- pinned here so a future - min_verify_len=0 port is a deliberate decision with a test to flip. - """ - # Half the batch is strong all the way down, half is dead on arrival. - # The batch must be big enough that its token range sits on a genuinely - # rising part of the theta curve: below the first breakpoint the clamp - # makes cost flat, every candidate is free, and both planners trivially - # buy everything at the same theta -- the two-row version of this test - # sat in that regime and proved nothing. - bs = 64 - strong = np.tile([0.99, 0.98, 0.97, 0.96, 0.95], (bs // 2, 1)) - dead = np.tile([0.01, 0.008, 0.006, 0.004, 0.002], (bs // 2, 1)) - surv = np.concatenate([strong, dead], axis=0) - sgl_n = _sgl_compute_verify_token_budget( - history_survival_probs=torch.tensor(surv, dtype=torch.float32), - table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.0) - trt_m = compute_verify_token_budget( - survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, - min_verify_len=1) - - def theta_at(tau, tokens): - return tau / float(TRT_TABLE.step_times(np.asarray([tokens]), bs)[0]) - - cand = np.sort(surv.reshape(-1))[::-1] - sgl_theta = theta_at(bs + cand[:sgl_n].sum(), bs + sgl_n) - deeper = np.sort(surv[:, 1:].reshape(-1))[::-1] - trt_theta = theta_at(bs + surv[:, 0].sum() + deeper[:trt_m].sum(), - 2 * bs + trt_m) - assert sgl_theta > trt_theta - - -def test_eps_filtering_is_a_known_divergence(): - """SGLang drops sub-eps candidates inside the budget; TRT-LLM does not. - - TensorRT-LLM applies survival_eps at allocation time - (schedule_verify_lens_topk's eligibility mask), so its BUDGET can count a - candidate the allocator then refuses. This pins the divergence so a - future unification is a deliberate change, not an accident: with a heavy - sub-eps tail the SGLang budget must not exceed TRT-LLM's shifted one. - """ - bs = 16 - surv = np.tile([0.9, 0.4, 0.008, 0.004, 0.002], (bs, 1)) - sgl_n = _sgl_compute_verify_token_budget( - history_survival_probs=torch.tensor(surv, dtype=torch.float32), - table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.01) - trt_m = compute_verify_token_budget( - survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, - min_verify_len=1) - assert sgl_n - bs <= trt_m diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py index f1ff0934ec04..ffad376de6e4 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py @@ -467,3 +467,43 @@ def test_trim_regret_counts_only_drafts_alive_at_the_cut(): assert untrimmed.requests_trimmed == 0 assert untrimmed.trim_regret_rate == 0.0 assert untrimmed.summary()["accept_len"] == pytest.approx(8 / 3, abs=1e-4) + + +# -------------------------------------------------------------------------- +# top-k allocation follows confidence (moved from test_dspark_planner_trims.py) +# -------------------------------------------------------------------------- + + +def _block5_cfg() -> DSparkScheduleConfig: + return DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + + +def test_topk_hands_longer_windows_to_more_confident_requests(): + """The budget must follow survival, not batch position.""" + # survival[r, k] = P(the first k+1 drafted tokens all get accepted) + survival = torch.tensor([ + [0.99, 0.98, 0.97, 0.96, 0.95], # confident + [0.90, 0.70, 0.50, 0.30, 0.10], # middling + [0.20, 0.04, 0.01, 0.00, 0.00], # collapses + ]) + lens = schedule_verify_lens_topk(survival=survival, budget=6, + cfg=_block5_cfg()).tolist() + + assert len(lens) == 3 + assert lens[0] > lens[2], ( + f"confident request got {lens[0]}, collapsing request got {lens[2]}; " + f"the budget must follow survival, not batch position") + assert lens[0] >= lens[1] >= lens[2] + # Every request keeps at least the floor, and the budget is respected. + assert min(lens) >= 1 + assert sum(lens) - len(lens) * 1 <= 6 + + +def test_a_full_budget_degenerates_to_the_uniform_full_window(): + """The no-trim case must stay exactly uniform: ragged must cost nothing.""" + num_reqs, max_len = 4, 5 + survival = torch.full((num_reqs, max_len), 0.99) + budget = num_reqs * (max_len - 1) + lens = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=_block5_cfg()).tolist() + assert lens == [max_len] * num_reqs diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py deleted file mode 100644 index acbf6eba383f..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py +++ /dev/null @@ -1,61 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Cost lookup must interpolate between breakpoints, not floor onto them. - -The floor lookup this replaces silently disabled the whole feature: with -measured points at 768 and 1536, every total in 769..1535 was billed at the -768 price, so a full block at bs=252 (1512 tokens, true cost ~148 ms) was -priced 72.3 ms. The planner's cost ratio T(6bs)/T(3bs) collapsed from a true -~2.0 to 1.19 -- below the GSM8K survival ratio tau5/tau2 ~ 1.48 -- so the -argmax bought the full block on ~95% of decisions, the bucket landed on -padded_bs*6 on every step of every run, and trimming never happened. The only -exception ever observed was bs=256 exactly, where the tier totals land on -breakpoints and floor happens to be exact. - -The consumer this was ported from (SGLang ``_additive_step_time_tensor``) -interpolates both theta(M) and alpha(bs); this port must match. -""" - -import pytest - -from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable - -# The certified GB300 table's shape, abbreviated to the two breakpoints whose -# gap hid the bug. -TABLE = SpsCostTable(token_counts=(512, 768, 1536), - step_time_ms=(23.85, 35.61, 105.64), - fixed_overhead_ms=25.244, - batch_sizes=(128, 256), - batch_overhead_ms=(11.462, 19.343)) - - -def test_theta_interpolates_between_breakpoints(): - """Exact on breakpoints, linear between them. - - The regression case: 1512 tokens must cost ~the 1536 price, not the 768 - one. - """ - assert TABLE.step_time(768, 256) == pytest.approx(25.244 + 19.343 + 35.61) - assert TABLE.step_time(1536, 256) == pytest.approx(25.244 + 19.343 + 105.64) - got = TABLE.step_time(1512, 252) - theta = 35.61 + (105.64 - 35.61) * (1512 - 768) / (1536 - 768) - alpha = 11.462 + (19.343 - 11.462) * (252 - 128) / (256 - 128) - assert got == pytest.approx(25.244 + alpha + theta) - # The floor price it used to return -- and must never return again. - assert got > 130.0 - - -def test_alpha_interpolates_and_clamps(): - assert TABLE.batch_overhead(128) == pytest.approx(11.462) - assert TABLE.batch_overhead(192) == pytest.approx((11.462 + 19.343) / 2) - assert TABLE.batch_overhead(64) == pytest.approx(11.462) # clamp low - assert TABLE.batch_overhead(512) == pytest.approx(19.343) # clamp high - - -def test_theta_clamps_outside_the_measured_range(): - """Theta clamps to the end values; a table without a batch axis adds no alpha.""" - plain = SpsCostTable(token_counts=(512, 768, 1536), - step_time_ms=(23.85, 35.61, 105.64), - fixed_overhead_ms=25.244) - assert plain.step_time(100, 0) == pytest.approx(25.244 + 23.85) - assert plain.step_time(4096, 0) == pytest.approx(25.244 + 105.64) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py deleted file mode 100644 index e2b9f049807b..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py +++ /dev/null @@ -1,159 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Reading cost-table samples out of production iteration logs. - -The in-process sweep measures an engine the profiler built itself, and when -that engine's configuration drifts from the deployment the table describes a -machine nobody is running -- a swept (bs=256, M=1536) cell read 150.2 ms -against 118.8 ms measured live at the same shape, with nothing in the -fingerprint to record the difference. ``print_iter_log`` already emits the row -count, the token total and the host step time for every step of the real -deployment, so a table can be fitted from traffic that actually ran. - -What must hold: the token total identifies the verify length exactly (no -clustering, no guessing), steps that cannot be filed honestly are dropped -rather than mislabelled, and the first step of each executor instance -- whose -host time contains the wait for the first request, measured at 62 s live -- is -never treated as a measurement. -""" - -import importlib.util -import pathlib -import sys -import tempfile - -_SPEC = importlib.util.spec_from_file_location( - "dspark_sps_profiler", - pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / - "dspark_sps_profiler.py") -_PROFILER = importlib.util.module_from_spec(_SPEC) -sys.modules[_SPEC.name] = _PROFILER # @dataclass needs the module registered -_SPEC.loader.exec_module(_PROFILER) - -samples_from_iter_log = _PROFILER.samples_from_iter_log -resolve_padded_shape = _PROFILER.resolve_padded_shape - -# A deployment's captured ladder, abbreviated. Padding rounds a step's rows up -# to the smallest entry at or above them, which is what turns a padded token -# total into exactly one verify length. -LADDER = [1, 2, 4, 8, 16, 32, 64, 128, 192, 256] - - -def _line(iteration, bs, step_ms, gen_tokens, rank="0"): - return (f"{rank}: [08/04/2026-19:08:13] [TRT-LLM] [I] [_torch][RANK 0] " - f"iter = {iteration}, global_rank = 0, rank = 0, " - f"num_scheduled_requests = {bs}, kv_cache_util = 0.125, " - f"currank_total_requests = 1/8, host_step_time = {step_ms}ms, " - f"prev_device_step_time = {step_ms}ms, " - f"timestamp = 2026-08-04 19:08:13, " - f"states = {{'num_ctx_requests': 0, 'num_ctx_tokens': 0, " - f"'num_generation_tokens': {gen_tokens}, " - f"'cached_kv_tokens': 63821}}") - - -def _write(lines, name="gen.log"): - # tempfile rather than pytest's tmp_path: these tests also run under the - # fixture-less in-container runner used for the GPU images. - path = pathlib.Path(tempfile.mkdtemp()) / name - path.write_text("\n".join(lines) + "\n") - return str(path) - - -def test_verify_length_comes_from_the_padded_token_total(): - """The label is read out of the log, never inferred from step times. - - Steps the ladder cannot explain are skipped, not mislabelled: filing a - step under a shape it did not run is the mislabelling that once shipped - eight fictional cells with near-zero residuals. - """ - log = _write([ - _line(1, 1, 62098.0, 6), # executor's first step: dropped - _line(2, 244, 79.2, 768), # 244 rows padded to 256 -> rung-2 - _line(3, 242, 117.1, 1536), # padded to 256 -> full block - _line(4, 128, 61.0, 384), # exactly 128 rows -> rung-2 - _line(5, 256, 300.0, 2048), # L = 7, past the block: skipped - _line(6, 300, 400.0, 1800), # more rows than the ladder holds: skipped - ]) - samples = samples_from_iter_log([log], max_draft_len=5, - padded_batch_sizes=LADDER) - got = sorted((s.batch_size, s.verify_len, s.step_time_ms) - for s in samples) - assert got == [(128, 2, 61.0), (256, 2, 79.2), (256, 5, 117.1)] - # The table is indexed by the padded token total, which each sample exposes. - assert sorted(s.total_verify_tokens for s in samples) == [384, 768, 1536] - - -def test_the_padded_width_is_what_gets_recorded(): - """A 244-row step cost what 256 rows cost; the table is indexed by that. - - The ladder is also what removes the ambiguity: 768 tokens over 128 rows - is L=5 at 128, L=3 at 192, or L=2 at 256 -- without the ladder all three - are admissible and the step is dropped; with it, padding rounds 128 up to - 128 and the answer is L=5 (the first version of this collector kept 312 - of 5682 production steps for lack of this). And a step whose token total - the padded width cannot explain (_get_padded_batch can bail) is dropped: - nothing in the log says padding declined, so that is the only honest move. - """ - assert resolve_padded_shape(num_rows=244, num_generation_tokens=768, - max_draft_len=5, max_batch_size=256, - padded_batch_sizes=LADDER) == (256, 2) - assert resolve_padded_shape(num_rows=128, num_generation_tokens=768, - max_draft_len=5, max_batch_size=256) is None - assert resolve_padded_shape(num_rows=128, num_generation_tokens=768, - max_draft_len=5, max_batch_size=256, - padded_batch_sizes=LADDER) == (128, 5) - assert resolve_padded_shape(num_rows=200, num_generation_tokens=600, - max_draft_len=5, max_batch_size=256, - padded_batch_sizes=LADDER) is None - - -def test_the_first_step_of_each_executor_is_dropped(): - """A build spins up more than one executor, and each numbers from 1. - - The first step of each carries the wait for its first request inside - host_step_time (62 s on a live server), which would otherwise enter the - fit as a cell measurement. - """ - log = _write([ - _line(1, 1, 200.0, 6), # KV-estimation executor - _line(2, 256, 79.0, 768), - _line(1, 1, 62098.0, 6), # real executor starts over - _line(2, 256, 79.4, 768), - ]) - times = [s.step_time_ms for s in samples_from_iter_log([log], max_draft_len=5)] - assert times == [79.0, 79.4] - - -def test_only_the_requested_rank_is_read(): - """Every rank logs; counting all of them would multiply every cell. - - Outlier steps are bounded out on the same pass. - """ - log = _write([ - _line(2, 256, 79.2, 768, rank="0"), - _line(2, 256, 79.9, 768, rank="3"), # another rank's copy: dropped - _line(3, 256, 5000.0, 768), # outlier past max_step_ms - ]) - times = [s.step_time_ms for s in - samples_from_iter_log([log], max_draft_len=5, max_step_ms=1000.0)] - assert times == [79.2] - - -def test_several_logs_are_pooled(): - """One arm gives one rung; the table needs the ladder, so pool the runs.""" - a = _write([_line(1, 1, 100.0, 6), _line(2, 256, 117.1, 1536)], - name="a.log") - c = _write([_line(1, 1, 100.0, 6), _line(2, 256, 79.2, 768)], - name="c.log") - got = sorted(s.verify_len for s in - samples_from_iter_log([a, c], max_draft_len=5)) - assert got == [2, 5] - - -def test_a_log_without_the_states_dict_yields_nothing(): - """enable_iter_perf_stats off: fail loudly upstream, not with a bad table.""" - path = pathlib.Path(tempfile.mkdtemp()) / "bare.log" - path.write_text( - "0: [TRT-LLM] iter = 2, num_scheduled_requests = 256, " - "host_step_time = 79.2ms\n") - assert samples_from_iter_log([str(path)], max_draft_len=5) == [] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py deleted file mode 100644 index 78d882c0ec0f..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py +++ /dev/null @@ -1,99 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The two reconciliation guards: engine fingerprint and price prediction. - -Both exist because this pipeline's characteristic failure is a PREDICTION -recorded as a MEASUREMENT with nothing positioned to contradict it: a cost -table measured on the wrong MoE kernel almost planned MegaMoE steps, the same -nominal cell measured 23% apart across two engine configs, and a planner -argmax'd against a table whose labels described shapes that never ran. The -fingerprint refuses the wrong table at load; the recorded prediction gives -every run a number that reality (hostStepTimeMS) can contradict. -""" - -import numpy as np -import pytest - -from tensorrt_llm._torch.speculative.dspark_planner import ( - SpsCostTable, check_table_fingerprint) -from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig -from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner - -LIVE = {"tp": 8, "ep": 8, "attention_dp": True, "block": 5, - "max_batch_size": 256} - - -def _payload(engine): - return {"token_counts": [512, 1536], "step_time_ms": [60.0, 150.0], - "_meta": {"engine": engine} if engine is not None else {}} - - -def test_mismatched_engine_is_refused(): - """The 23%-apart case: same cell, different max_batch_size.""" - wrong = dict(LIVE, max_batch_size=64) - with pytest.raises(ValueError, match="different engine configuration"): - check_table_fingerprint(payload=_payload(wrong), live=dict(LIVE)) - - -def test_fingerprints_that_must_load(): - """Every non-contradicted fingerprint loads; only a contradicted fact refuses. - - Matching facts load silently; string facts must not fail on spelling case - (MEGAMOE vs MegaMoe); facts the consumer cannot see (image hash, geometry) - are informational and logged, not refused -- refusing on unverifiable keys - would make every fingerprint addition a breaking change; and an old table - with no fingerprint warns rather than breaking existing deployments. - """ - check_table_fingerprint(payload=_payload(dict(LIVE)), live=dict(LIVE)) - check_table_fingerprint( - payload=_payload(dict(LIVE, moe_backend="megamoe_cutedsl")), - live=dict(LIVE, moe_backend="MEGAMOE_CUTEDSL")) - check_table_fingerprint( - payload=_payload(dict(LIVE, image="faf2c60935", - geometry="constant_block")), - live=dict(LIVE)) - check_table_fingerprint(payload=_payload(None), live=dict(LIVE)) - - -def test_planner_records_the_price_it_paid(): - """Every budget decision must leave behind the step time it assumed. - - Without this the run has no number reality can contradict: the argmax - consults the table and the table's opinion evaporates. With it, the - [final] planner block carries predicted_ms_sum / predicted_steps, and a - systematic gap against measured hostStepTimeMS is the one signal that - catches a wrong table, a wrong lookup, and an engine mismatch alike. - """ - table = SpsCostTable(token_counts=(0, 512, 768, 1536), - step_time_ms=(5.0, 68.4, 80.2, 150.5), - fixed_overhead_ms=1.0) - cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) - planner = DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=[1, 2, 5]) - - bs = 8 - rng = np.random.default_rng(11) - survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() - - import torch - # Only the snapshot is faked; calibration stays the constructor's default - # sigmoid (setting it to None post-construction would bypass the - # `apply_calibration or torch.sigmoid` fallback and crash the decide path). - planner._gather_rows = lambda **_: torch.tensor(survival, - dtype=torch.float32) - - lens = planner.decide_verify_lens(num_gen_requests=bs, - reduce_across_ranks=False) - assert lens is not None - - steps = planner.stats.get("predicted_steps", 0) - total = planner.stats.get("predicted_ms_sum", 0.0) - assert steps == 1 and total > 0.0 - - # The recorded price must be the table's own answer for the tokens the - # decision implies -- anything else and the reconciliation compares two - # unrelated numbers. - budget = sum(lens) - bs * cfg.min_verify_len - tokens = bs * (cfg.min_verify_len + 1) + budget - expected = float(table.step_times(np.asarray([tokens]), bs)[0]) - assert planner.last_predicted_step_ms == pytest.approx(expected) - assert total == pytest.approx(expected) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py deleted file mode 100644 index 85517f5c4c30..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py +++ /dev/null @@ -1,110 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The two snapshot instruments: stamp-lag histogram and neutral-row count. - -They exist because the planner's inputs have never been measured: the argmax -eats a slot-indexed snapshot staged one step earlier, and nothing recorded how -old each gathered row was (relay staleness) or how many rows were still the -+30 neutral fill (unknowns entering as certainties). The live anomaly they -instrument: 445/476 full-batch steps chose the full block while calibrated -pooled survivals favour rung-2 by 36% -- so the per-step snapshot content must -differ from the pooled shards, and these counters say in which way. - -Both are measurement-only. A failure inside them must never kill a step. -""" - -import numpy as np -import pytest -import torch - -from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable -from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig -from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner - - -def _planner(): - table = SpsCostTable(token_counts=(0, 512, 768, 1536), - step_time_ms=(5.0, 68.4, 80.2, 150.5), - fixed_overhead_ms=1.0) - cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) - return DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=[1, 2, 5]) - - -def test_neutral_rows_are_counted_and_real_rows_are_not(): - planner = _planner() - selected = torch.zeros((8, 5)) - selected[3] = 30.0 # the neutral fill: sigmoid ~ 1.0 at any temperature - planner._note_snapshot_stats(selected, rows=None) - assert planner.stats["snap_rows"] == 8 - assert planner.stats["snap_neutral_rows"] == 1 - - -def test_lag_histogram_reads_staged_seq_minus_stamp(): - planner = _planner() - planner._host_stamps = torch.tensor([7, 6, 5, 0], dtype=torch.int32) - planner._staged_seq = 7 - selected = torch.zeros((4, 5)) - planner._note_snapshot_stats(selected, rows=[0, 1, 2, 3]) - hist = planner.stats["stamp_lag_hist"] - # Stamp 0 is "never drafted": lag 7 in truth, and it must land visibly at - # the histogram's far end rather than vanish -- these are exactly the rows - # the neutral counter should agree about. - assert hist == {0: 1, 1: 1, 2: 1, 7: 1} - - -def test_large_and_negative_lags_clamp_instead_of_exploding_the_dict(): - planner = _planner() - planner._host_stamps = torch.tensor([0, 500], dtype=torch.int32) - planner._staged_seq = 400 - planner._note_snapshot_stats(torch.zeros((2, 5)), rows=[0, 1]) - assert planner.stats["stamp_lag_hist"] == {8: 1, -2: 1} - - -def test_counts_accumulate_across_steps(): - planner = _planner() - planner._host_stamps = torch.tensor([3, 3], dtype=torch.int32) - planner._staged_seq = 4 - for _ in range(3): - planner._note_snapshot_stats(torch.full((2, 5), 30.0), rows=[0, 1]) - assert planner.stats["snap_rows"] == 6 - assert planner.stats["snap_neutral_rows"] == 6 - assert planner.stats["stamp_lag_hist"] == {1: 6} - - -def test_instrument_failure_is_counted_not_raised(): - planner = _planner() - planner._host_stamps = torch.tensor([1], dtype=torch.int32) - planner._staged_seq = 1 - # Rows that cannot index the stamp buffer must clamp, and anything worse - # must be swallowed into the error counter: instruments never kill steps. - planner._note_snapshot_stats(torch.zeros((1, 5)), rows=[10_000]) - assert planner.stats.get("snap_stats_errors", 0) == 0 - planner._note_snapshot_stats(object(), rows=None) # type: ignore[arg-type] - assert planner.stats["snap_stats_errors"] == 1 - - -def test_decide_path_feeds_the_instruments(): - """The counters must fill from the real decide path, not only direct calls.""" - planner = _planner() - bs = 8 - rng = np.random.default_rng(3) - survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() - planner._gather_rows = lambda **_: torch.tensor(survival, - dtype=torch.float32) - lens = planner.decide_verify_lens(num_gen_requests=bs, - reduce_across_ranks=False) - assert lens is not None - assert planner.stats["snap_rows"] == bs - assert planner.stats["snap_neutral_rows"] == 0 - # The local-decision instruments must fill on the same path: one rung per - # decision, one window per request, and the two must agree -- the windows - # are the top-k realisation of exactly that rung's budget. - rung_hist = planner.stats["local_rung_hist"] - assert sum(rung_hist.values()) == 1 - (rung,) = rung_hist - assert rung in (1, 2, 5) - len_hist = planner.stats["local_len_hist"] - assert sum(len_hist.values()) == bs - budget = bs * (rung - planner.cfg.min_verify_len) - extra = sum((k - planner.cfg.min_verify_len) * v for k, v in len_hist.items()) - assert extra <= budget diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py deleted file mode 100644 index 4e4718e80dbc..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py +++ /dev/null @@ -1,620 +0,0 @@ -# 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. -"""Unit tests for the DSpark SPS cost-table profiler (CPU only, no model load). - -Everything here exercises the pure half of the profiler: turning ``get_stats`` -rows into aligned samples, summarizing cells, separating ``alpha(bs)`` from -``theta(M)``, and assembling / refusing the emitted JSON. The sweep driver needs -a GPU and a checkpoint and is deliberately not covered. - -The load-bearing case is :func:`test_round_trip_reproduces_measured_cells`: it -takes a synthetic machine with a known cost curve, runs it through the whole -fit-and-emit pipeline, and reloads the JSON exactly the way -``dspark.py::_build_verify_planner`` does. If the profiler and the loader ever -disagree about units, keys or the meaning of a token count, that test is where -it shows up. -""" - -import importlib.util -import json -import pathlib -import sys - -import numpy as np -import pytest - -from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable, total_verify_tokens -# The profiler is a development tool under tests/microbenchmarks, not part of -# the shipped package: nothing on the serving path imports it, and keeping a -# 1.8k-line sweep driver inside tensorrt_llm/ would ship it to every user. -# Load it by path so this test follows the file rather than an import root. -_SPEC = importlib.util.spec_from_file_location( - "dspark_sps_profiler", - pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / - "dspark_sps_profiler.py") -_PROFILER = importlib.util.module_from_spec(_SPEC) -# Registered BEFORE exec: @dataclass resolves annotations through -# sys.modules[cls.__module__].__dict__, so a module executed without being -# registered raises "'NoneType' object has no attribute '__dict__'" on its -# first dataclass. -sys.modules[_SPEC.name] = _PROFILER -_SPEC.loader.exec_module(_PROFILER) - -globals().update({name: getattr(_PROFILER, name) for name in ('CellStat', 'FlatCostTableError', 'InertCostTableError', 'InsufficientSamplesError', 'StepSample', 'SweepConfig', 'SweepGeometryError', 'aligned_steps_from_stats', 'build_cost_table_payload', 'check_table_is_informative', 'compress_to_risers', 'fit_additive_cost_model', 'load_cost_table', 'profitability_probe', 'running_max', 'summarize_cells')}) - -#: The exact set of keys ``dspark.py::_build_verify_planner`` reads. Anything -#: else in the file is ignored by the loader, so the profiler must not rely on -#: it -- and must not omit any of these. -LOADER_KEYS = { - "token_counts", - "step_time_ms", - "fixed_overhead_ms", - "batch_sizes", - "batch_overhead_ms", -} - - -def _stats_row(*, iteration, num_gen, rank=0, num_ctx=0, host_ms=10.0, gpu_ms=8.0): - """One /metrics-shaped iteration record.""" - row = { - "iter": iteration, - "attentionDpRank": rank, - "inflightBatchingStats": { - "numContextRequests": num_ctx, - "numGenRequests": num_gen, - }, - } - if host_ms is not None: - row["hostStepTimeMS"] = host_ms - if gpu_ms is not None: - row["gpuForwardTimeMS"] = gpu_ms - return row - - -# -------------------------------------------------------------------------- -# aligned step extraction -# -------------------------------------------------------------------------- - - -def test_aligned_steps_drops_prefill_iterations(): - """A step with context requests is a different shape, not a cheap decode.""" - steady = [_stats_row(iteration=i, num_gen=8, host_ms=10.0 + i) for i in range(3)] - assert aligned_steps_from_stats(steady) == [(0, 8, 10.0), (1, 8, 11.0), (2, 8, 12.0)] - rows = [ - _stats_row(iteration=0, num_gen=0, num_ctx=8, host_ms=90.0), - _stats_row(iteration=1, num_gen=8, num_ctx=2, host_ms=40.0), - _stats_row(iteration=2, num_gen=8, host_ms=10.0), - ] - assert aligned_steps_from_stats(rows) == [(2, 8, 10.0)] - - -def test_aligned_steps_requires_every_dp_rank(): - """Multi-rank alignment: complete, agreeing, and timed by rank 0. - - A missing rank means the iteration was not observed batch-wide; ranks at - different batch sizes are not one (bs, M) point; and ADP fans out - rank-local counters with rank-0's clock, so timing must never be averaged. - """ - rows = [ - _stats_row(iteration=0, num_gen=8, rank=0), - _stats_row(iteration=0, num_gen=8, rank=1), - _stats_row(iteration=1, num_gen=8, rank=0), - ] - assert aligned_steps_from_stats(rows, expected_ranks=2) == [(0, 8, 10.0)] - - disagreeing = [ - _stats_row(iteration=0, num_gen=8, rank=0), - _stats_row(iteration=0, num_gen=7, rank=1), - ] - assert aligned_steps_from_stats(disagreeing, expected_ranks=2) == [] - - rank0_not_first = [ - _stats_row(iteration=0, num_gen=8, rank=1, host_ms=99.0), - _stats_row(iteration=0, num_gen=8, rank=0, host_ms=10.0), - ] - assert aligned_steps_from_stats(rank0_not_first, expected_ranks=2) == [(0, 8, 10.0)] - - -def test_aligned_steps_skips_missing_timing_field(): - """gpuForwardTimeMS is read without a sync and is absent on some steps.""" - rows = [ - _stats_row(iteration=0, num_gen=8, gpu_ms=None), - _stats_row(iteration=1, num_gen=8, gpu_ms=7.5), - ] - assert aligned_steps_from_stats(rows, timing_key="gpuForwardTimeMS") == [(1, 8, 7.5)] - - -# -------------------------------------------------------------------------- -# warmup + median -# -------------------------------------------------------------------------- - - -def _samples(batch_size, verify_len, times): - return [ - StepSample(batch_size=batch_size, verify_len=verify_len, step_time_ms=t, iteration=i) - for i, t in enumerate(times) - ] - - -def test_summarize_discards_warmup(): - warm = [500.0] * 4 - steady = [10.0] * 8 - cells = summarize_cells(_samples(8, 3, warm + steady), warmup_steps=4, min_samples=8) - assert len(cells) == 1 - assert cells[0].step_time_ms == pytest.approx(10.0) - assert cells[0].num_samples == 8 - - -def test_summarize_reports_median_not_mean(): - """The right tail is real and must not move the point.""" - times = [10.0, 10.0, 10.0, 10.0, 10.0, 1000.0] - cells = summarize_cells(_samples(8, 3, times), warmup_steps=0, min_samples=1) - assert cells[0].step_time_ms == pytest.approx(10.0) - assert cells[0].step_time_ms != pytest.approx(float(np.mean(times))) - # ...but the tail is still reported, so a pathological cell is visible. - assert cells[0].p90_ms > cells[0].step_time_ms - - -def test_summarize_drops_thin_cells_outside_the_requested_grid(): - """A deployment log carries ramp/drain occupancy nobody asked to measure. - - Job 2585129: a five-rung server sweep with 984-step cells at every - requested (bs, L) died at the very end on (bs=2, L=1) with one sample -- - the admission ramp of the first cell's client, filed as a cell and then - held to the requested-grid bar. Incidental cells must be dropped when - thin, kept when thick, and a thin REQUESTED cell must still refuse -- - with or without a requested grid, since a silently dropped cell can - disconnect the grid, not just widen error bars. - """ - grid = [(8, 3)] - requested = _samples(8, 3, [10.0] * 8) - ramp = _samples(2, 3, [30.0]) # thin, incidental: drop - thick_extra = _samples(4, 3, [20.0] * 8) # thick, incidental: keep - cells = summarize_cells(requested + ramp + thick_extra, warmup_steps=0, - min_samples=4, expected_cells=grid) - assert [(c.batch_size, c.verify_len) for c in cells] == [(4, 3), (8, 3)] - with pytest.raises(InsufficientSamplesError, match="steady samples"): - summarize_cells(_samples(8, 3, [10.0]) + thick_extra, warmup_steps=0, - min_samples=4, expected_cells=grid) - # No grid given means every observed cell was requested: same refusal. - with pytest.raises(InsufficientSamplesError, match="steady samples"): - summarize_cells(_samples(8, 3, [10.0] * 5), warmup_steps=4, min_samples=8) - - -def test_summarize_orders_by_iteration_before_cutting_warmup(): - """Stats arrive per queue, not necessarily in iteration order.""" - out_of_order = [ - StepSample(batch_size=8, verify_len=1, step_time_ms=t, iteration=i) - for i, t in sorted(enumerate([500.0, 500.0, 10.0, 10.0]), key=lambda p: -p[0]) - ] - cells = summarize_cells(out_of_order, warmup_steps=2, min_samples=2) - assert cells[0].step_time_ms == pytest.approx(10.0) - - -def test_step_sample_json_round_trip(): - sample = StepSample(batch_size=16, verify_len=3, step_time_ms=12.5, iteration=7) - assert StepSample.from_json(json.loads(json.dumps(sample.to_json()))) == sample - assert sample.total_verify_tokens == total_verify_tokens(16, 3) == 64 - - -# -------------------------------------------------------------------------- -# additive fit -# -------------------------------------------------------------------------- - - -def _synthetic_cells(batch_sizes, verify_lens, *, alpha, theta, noise=None): - """Cells generated by an exactly additive machine ``alpha(bs) + theta(M)``.""" - cells = [] - for batch_size in batch_sizes: - for verify_len in verify_lens: - tokens = total_verify_tokens(batch_size, verify_len) - value = alpha(batch_size) + theta(tokens) - if noise is not None: - value += noise(batch_size, verify_len) - cells.append( - CellStat( - batch_size=batch_size, - verify_len=verify_len, - step_time_ms=value, - num_samples=32, - p10_ms=value, - p90_ms=value, - ) - ) - return cells - - -def _staircase(edges_and_costs): - def theta(tokens): - value = edges_and_costs[0][1] - for edge, cost in edges_and_costs: - if tokens >= edge: - value = cost - return value - - return theta - - -BATCH_SIZES = [8, 16, 32, 64] -VERIFY_LENS = [1, 2, 3, 4, 5] -ALPHA = {8: 20.0, 16: 22.0, 32: 26.0, 64: 34.0} -THETA = _staircase([(0, 1.0), (64, 2.0), (128, 4.0), (256, 8.0)]) - - -def test_fit_recovers_relative_structure_exactly(): - """The fit is only identified up to one shared constant -- pin that, then - every cell must be reproduced to the last bit.""" - cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) - fit = fit_additive_cost_model(cells) - assert fit.warnings == () - - for cell in cells: - predicted = fit.intercept_ms[cell.batch_size] + fit.theta_ms[cell.total_verify_tokens] - assert predicted == pytest.approx(cell.step_time_ms, abs=1e-6) - assert fit.max_rel_residual == pytest.approx(0.0, abs=1e-9) - - # The batch axis keeps its measured *differences*, which is what alpha means. - for smaller, larger in zip(BATCH_SIZES, BATCH_SIZES[1:]): - measured = ALPHA[larger] - ALPHA[smaller] - fitted = fit.intercept_ms[larger] - fit.intercept_ms[smaller] - assert fitted == pytest.approx(measured, abs=1e-6) - - -def test_fit_surfaces_a_bad_cell_rather_than_hiding_it(): - """A cell the additive model cannot explain is evidence, not noise. - - Dropping it would produce a table that fits beautifully and predicts badly, - and would delete the one signal that says ``theta`` is not a function of - ``M`` alone. - """ - - def noise(batch_size, verify_len): - return 500.0 if (batch_size, verify_len) == (32, 3) else 0.0 - - cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA, noise=noise) - fit = fit_additive_cost_model(cells) - - assert fit.max_abs_residual_ms > 100.0 - assert fit.max_rel_residual > 1.0 - report = " ".join(fit.warnings) - assert "additive model misses" in report - # The least-squares solve spreads the bad cell over everything sharing its - # batch size or token count, so the report names a neighbourhood -- the bad - # cell is in it, but so are innocent ones, which is why it is worded that way. - assert "bs=32, L=3" in report - - -def test_fit_clamps_a_dip_upward(): - """A non-monotone theta would offer the planner a cheaper *longer* verify.""" - dipped = _staircase([(0, 1.0), (64, 5.0), (128, 3.0), (256, 8.0)]) - cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=dipped) - fit = fit_additive_cost_model(cells) - - tokens = sorted(fit.theta_ms) - values = [fit.theta_ms[m] for m in tokens] - assert values == running_max(values) - assert any("monoton" in w for w in fit.warnings) - - -def test_fit_refuses_a_disconnected_sweep(): - """Two islands each carry their own unknown offset; splicing them lies.""" - cells = _synthetic_cells([8], [1, 2], alpha=ALPHA.get, theta=THETA) - cells += _synthetic_cells([100], [1, 2], alpha=lambda _: 40.0, theta=THETA) - with pytest.raises(SweepGeometryError, match="disconnected"): - fit_additive_cost_model(cells) - - -# -------------------------------------------------------------------------- -# staircase shaping -# -------------------------------------------------------------------------- - - -def test_running_max_is_monotone(): - assert running_max([1.0, 3.0, 2.0, 2.5, 9.0]) == [1.0, 3.0, 3.0, 3.0, 9.0] - - -def test_compress_keeps_risers_and_closes_their_shelves(): - """A shelf must survive as a breakpoint PAIR under the interp consumer. - - Keeping only [16, 64] would turn the measured flat 16..48 into one long - rising segment: mid-shelf totals get billed part of the riser -- the - over-trim mirror of the floor-lookup bug. The shelf's last point (48) - rides along so interpolation stays flat where the hardware was flat. - """ - tokens = [16, 32, 48, 64, 80] - times = [1.0, 1.002, 1.004, 2.0, 2.001] - kept_tokens, kept_times = compress_to_risers( - tokens, times, min_riser_ms=0.02, max_breakpoints=8 - ) - assert kept_tokens == [16, 48, 64] - assert kept_times == [1.0, 1.004, 2.0] - - -def test_truncation_drops_the_point_cheapest_to_interpolate(): - """Over budget, remove the interior point interpolation best reconstructs. - - Here 64 goes: its neighbors predict it within 0.27ms, while removing 32 - would ramp the near-flat 16..32 stretch into the 64 riser and misprice it - by 0.9ms. The previous rule (drop the smallest incoming jump first) chose - 32 -- by construction it always chose the shelf-side point, re-billing - mid-shelf totals upward, the exact over-charge the pair encoding of - shelves exists to prevent. - """ - tokens = [16, 32, 64, 128] - times = [1.0, 1.1, 4.0, 9.0] - kept_tokens, _ = compress_to_risers(tokens, times, min_riser_ms=0.0, max_breakpoints=3) - assert kept_tokens == [16, 32, 128] - - -def test_compress_never_drops_the_base_shelf(): - """Index 0 is what every below-range query clamps onto.""" - tokens = [16, 32, 64] - times = [1.0, 100.0, 101.0] - kept_tokens, _ = compress_to_risers(tokens, times, min_riser_ms=0.0, max_breakpoints=2) - assert kept_tokens[0] == 16 - - -# -------------------------------------------------------------------------- -# payload assembly and JSON round-trip through SpsCostTable -# -------------------------------------------------------------------------- - - -def _payload(): - cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) - return build_cost_table_payload(cells, riser_tolerance=0.01, max_breakpoints=8), cells - - -def test_round_trip_reproduces_measured_cells(): - """The whole point: what the planner reads back must be what was measured. - - The payload must carry exactly the keys ``dspark.py::_build_verify_planner`` - reads (everything else under ``_meta``, where the loader ignores it), - survive a JSON file round trip into a real ``SpsCostTable``, and price - every measured cell back to within one shelf width -- the emitted curve is - a floor staircase by construction, so a cell between two risers is - deliberately priced at the lower shelf. - """ - payload, cells = _payload() - assert LOADER_KEYS.issubset(payload) - assert set(payload) - LOADER_KEYS == {"_meta"} - - table = load_cost_table(json.loads(json.dumps(payload))) - assert isinstance(table, SpsCostTable) - assert not table.is_flat - # The loader builds SpsCostTable directly, so the profiler's own validation - # must be the real constructor rather than a copy of its rules. - assert list(table.token_counts) == payload["token_counts"] - assert len(table.batch_sizes) == len(table.batch_overhead_ms) - - for cell in cells: - predicted = table.step_time(cell.total_verify_tokens, cell.batch_size) - assert predicted == pytest.approx(cell.step_time_ms, abs=1e-4) - - -def test_emitted_terms_are_not_double_counted(): - """step_time_ms is theta only; the floor lives in the two overhead fields.""" - payload, cells = _payload() - table = load_cost_table(payload) - smallest = min(BATCH_SIZES) - floor = payload["fixed_overhead_ms"] + table.batch_overhead(smallest) - # The floor is the M-independent part, i.e. the intercept of the measured - # row at the smallest batch size -- not the whole step time. - row = sorted((c for c in cells if c.batch_size == smallest), key=lambda c: c.verify_len) - assert floor < row[0].step_time_ms - assert floor == pytest.approx(ALPHA[smallest], abs=0.5) - - -def test_batch_overhead_grows_with_batch_size(): - payload, cells = _payload() - overheads = payload["batch_overhead_ms"] - assert overheads == sorted(overheads) - assert payload["batch_sizes"] == sorted(BATCH_SIZES) - assert payload["fixed_overhead_ms"] >= 0.0 - # Provenance rides along under _meta: the encoding marker and the cells. - assert payload["_meta"]["encoding"] == "decomposed" - assert len(payload["_meta"]["cells"]) == len(cells) - - -# -------------------------------------------------------------------------- -# refusals -# -------------------------------------------------------------------------- - - -def test_refuses_a_flat_curve(): - """A machine whose step cost ignores M must not produce a table at all.""" - cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=lambda _: 3.0) - payload = build_cost_table_payload(cells) - assert load_cost_table(payload).is_flat - with pytest.raises(FlatCostTableError, match="is_flat"): - check_table_is_informative(payload, batch_sizes=BATCH_SIZES, tiers=[1, 3, 5]) - - -def test_refuses_a_curve_that_only_moves_by_noise(): - """is_flat is exact float equality, so it cannot catch this one.""" - cells = _synthetic_cells( - BATCH_SIZES, - VERIFY_LENS, - alpha=ALPHA.get, - theta=lambda tokens: 3.0 + 1e-4 * tokens, - ) - payload = build_cost_table_payload(cells, riser_tolerance=0.0, max_breakpoints=8) - assert not load_cost_table(payload).is_flat - with pytest.raises(FlatCostTableError, match="noise floor"): - check_table_is_informative( - payload, batch_sizes=BATCH_SIZES, tiers=[1, 3, 5], min_step_time_spread=0.02 - ) - - -# A real but shallow riser: at bs=32 tier 3 (M=128) sits below it and tier 5 -# (M=192) above, so the curve is genuinely non-flat and clears the noise floor -# (2.5%), yet the step only gets 2.5% more expensive while the acceptance yield -# given up between those tiers is ~36% at p=0.9. The argmax therefore never -# leaves max_tier -- non-flat and inert at the same time. -# -# The shelf is ENCODED as a breakpoint pair: under the interpolating consumer -# a bare [0, 160] would ramp from 2.0 to 2.5 across the whole range, dropping -# the ladder spread below the noise floor and turning this fixture into the -# flat-noise case instead of the inert one it exists to pin. -INERT_PAYLOAD = { - "token_counts": [0, 159, 160], - "step_time_ms": [2.0, 2.0, 2.5], - "fixed_overhead_ms": 18.0, - "batch_sizes": [32], - "batch_overhead_ms": [0.0], -} - - -def test_refuses_an_inert_table(): - """Non-flat, but no tier pair is worth trading yield for. - - This is the failure mode that does *not* trip ``fallback_flat_cost``, so - nothing downstream would report it. ``allow_inert=True`` is the explicit - escape hatch -- and it must still say that nothing trims. - """ - kwargs = dict(batch_sizes=[32], tiers=[1, 3, 5], acceptance_rates=[0.9, 0.95]) - with pytest.raises(InertCostTableError, match="inert"): - check_table_is_informative(INERT_PAYLOAD, **kwargs) - diagnostics = check_table_is_informative(INERT_PAYLOAD, allow_inert=True, **kwargs) - assert diagnostics["trims_somewhere"] is False - - -def test_informative_table_passes_and_reports_where_trimming_wins(): - payload, _ = _payload() - diagnostics = check_table_is_informative( - payload, - batch_sizes=BATCH_SIZES, - tiers=[1, 3, 5], - acceptance_rates=[0.3, 0.6, 0.9], - ) - assert diagnostics["trims_somewhere"] is True - assert diagnostics["best_step_time_spread"] > 0.02 - # Even a good table cannot rescue a one-rung ladder: nothing to choose. - with pytest.raises(InertCostTableError, match="single entry"): - check_table_is_informative(payload, batch_sizes=BATCH_SIZES, tiers=[5]) - - -# -------------------------------------------------------------------------- -# sweep geometry guards (pure; the sweep itself needs a GPU) -# -------------------------------------------------------------------------- - - -def _sweep(**overrides) -> SweepConfig: - kwargs = dict( - model="/models/x", - speculative_model="/models/x", - batch_sizes=[8, 16], - verify_lens=[1, 3, 5], - input_len=1024, - warmup_steps=8, - measure_steps=32, - max_seq_len=4096, - max_num_tokens=8192, - ) - kwargs.update(overrides) - return SweepConfig(**kwargs) - - -def test_token_budget_is_the_step_count_when_acceptance_is_pinned(): - """Pinned, a step commits exactly one token, so the two are the same - number; unpinned it can commit verify_len + 1, so the budget must scale.""" - pinned = _sweep() - assert pinned.max_tokens_for(5) == pinned.warmup_steps + pinned.measure_steps + 16 - loose = _sweep(pin_acceptance=0.0) - assert loose.max_tokens_for(5) == pinned.max_tokens_for(5) * 6 - - -def test_validate_rejects_a_sequence_that_would_be_evicted(): - """An evicted request drains the batch mid-cell and shrinks the shape.""" - with pytest.raises(SweepGeometryError, match="max-seq-len"): - _sweep(input_len=4064, block_follows_verify_len=True).validate() - - -def test_validate_rejects_a_decode_step_wider_than_max_num_tokens(): - """A split decode step never has the shape its cell is filed under.""" - with pytest.raises(SweepGeometryError, match="widest decode step"): - _sweep(batch_sizes=[8, 4096], max_num_tokens=1024, max_seq_len=8192, - block_follows_verify_len=True).validate() - - -def test_validate_accepts_a_sane_sweep(): - """Sane means the drafted block matches what each cell verifies.""" - _sweep(block_follows_verify_len=True).validate() - _sweep(verify_lens=[5], max_draft_len=5).validate() - - -def test_validate_refuses_constant_block_verifying_below_it(): - """Building an engine that drafts 5 and verifies 3 runs the regime that - trips the Repeat.cu output_size assert -- the sweep would crash mid-cell - after minutes of engine build, or worse, measure a crash-adjacent shape. - """ - with pytest.raises(SweepGeometryError, match="Repeat.cu"): - _sweep().validate() - - -def test_dp_size_only_counts_ranks_under_attention_dp(): - """Non-ADP publishes one stats row per iteration however wide TP is.""" - assert _sweep(tp_size=8).dp_size == 1 - assert _sweep(tp_size=8, enable_attention_dp=True).dp_size == 8 - - -# -------------------------------------------------------------------------- -# profitability probe (runs the real planner argmax) -# -------------------------------------------------------------------------- - - -def test_probe_trims_at_low_acceptance_and_not_at_high(): - """Trimming is a trade: it only wins when the drafts are unlikely to survive.""" - payload = { - # A riser between bs*(3+1)=128 and bs*(5+1)=192 at bs=32, so tiers 3 and - # 5 sit on different shelves and the argmax has something to choose: - # T(3) = 3 ms, T(5) = 4 ms, a 33% step for a yield gain that is worth it - # only when the drafts survive. - "token_counts": [0, 160], - "step_time_ms": [2.0, 3.0], - "fixed_overhead_ms": 0.0, - "batch_sizes": [32], - "batch_overhead_ms": [1.0], - } - probes = profitability_probe( - payload, batch_size=32, tiers=[1, 3, 5], acceptance_rates=[0.05, 0.99] - ) - by_rate = {p["acceptance_rate"]: p for p in probes} - assert by_rate[0.05]["chosen_verify_len"] < 5 - assert by_rate[0.05]["trims"] is True - assert by_rate[0.99]["chosen_verify_len"] == 5 - - -def test_probe_uses_the_bonus_token_convention(): - """A tier's cost must be looked up at bs*(L+1), never bs*L. - - The two differ by a whole batch, which is easily a shelf's width; this pins - the profiler to the same convention as ``total_verify_tokens``. - """ - payload = { - # Priced so that tier 5 (M = 8*6 = 48) is expensive and tier 3 - # (M = 8*4 = 32) is cheap. Under the wrong bs*L convention tier 5 would - # be M = 40 and land on the cheap shelf instead. The shelf is a - # breakpoint pair so the interpolating consumer keeps it flat. - "token_counts": [0, 43, 44], - "step_time_ms": [1.0, 1.0, 50.0], - "fixed_overhead_ms": 0.0, - "batch_sizes": [8], - "batch_overhead_ms": [0.0], - } - table = load_cost_table(payload) - assert table.step_time(total_verify_tokens(8, 5), 8) == pytest.approx(50.0) - assert table.step_time(total_verify_tokens(8, 3), 8) == pytest.approx(1.0) - probes = profitability_probe(payload, batch_size=8, tiers=[1, 3, 5], acceptance_rates=[0.9]) - assert probes[0]["chosen_verify_len"] == 3 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py deleted file mode 100644 index 1324ad0c7684..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py +++ /dev/null @@ -1,135 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""A collection run must refuse the regimes that fit the scheduler, not the head. - -Both regimes below produce shards that look perfectly well-formed. The run -completes, the row counts are plausible, the fitter converges, and the emitted -temperatures are wrong -- there is no downstream check that can tell. That is -why these are construction-time failures rather than warnings. - -Measured on this branch: collecting while the planner trimmed gave -``ece_before = [.244 .319 .361 .389 .394]``; pinning the window to the full -block, same model and same fitter, gave ``[.157 .208 .270 .321 .344]``. The -difference is entirely the censored label. - -SGLang guards the first regime next to its own collector -(``dspark_planner.py:100-108``) and names the second in the sibling -``ConfidenceMetricsProbe`` -- "padded verify rows corrupt the per-position -prefix label" (``dspark_observability.py:688-697``). -""" - -import contextlib -import os -import tempfile - -import pytest - -from tensorrt_llm._torch.speculative.dspark_sts import (STS_COLLECT_ENV, - make_recorder_from_env) - -_PIN = "TLLM_DSPARK_FORCE_VERIFY_LEN" - - -@contextlib.contextmanager -def collecting(pin=None): - """Set the collection env for the block, restore it after. - - A context manager rather than a pytest fixture: these run inside the - container, which ships no pytest, under a minimal shim that does not - implement fixtures or tmp_path. - """ - saved = {k: os.environ.get(k) for k in (STS_COLLECT_ENV, _PIN)} - try: - with tempfile.TemporaryDirectory() as tmp: - os.environ[STS_COLLECT_ENV] = os.path.join(tmp, "collect") - if pin is None: - os.environ.pop(_PIN, None) - else: - os.environ[_PIN] = str(pin) - yield - finally: - for k, v in saved.items(): - os.environ.pop(k, None) if v is None else os.environ.__setitem__(k, v) - - -def test_collection_off_builds_nothing(): - """The serving default must be untouched by any of this.""" - os.environ.pop(STS_COLLECT_ENV, None) - assert make_recorder_from_env(block_size=5) is None - assert make_recorder_from_env(block_size=5, has_cost_table=True, - ragged_mode="compact") is None - - -def test_a_loaded_cost_table_without_a_pin_is_refused(): - """A table means the planner trims, and trimming censors the label.""" - with collecting(): - with pytest.raises(ValueError, match="censors the acceptance label"): - make_recorder_from_env(block_size=5, has_cost_table=True) - - -def test_a_pinned_window_makes_the_table_harmless(): - """The uncensored regimes must all build a recorder. - - Pinning to the full block makes the label uncensored by construction; - without a table the planner cannot trim, so nothing censors anything; - static mode contributes no padding rows. - """ - with collecting(pin=5): - rec = make_recorder_from_env(block_size=5, has_cost_table=True) - assert rec is not None and rec.block_size == 5 - with collecting(): - assert make_recorder_from_env(block_size=5, - has_cost_table=False) is not None - assert make_recorder_from_env(block_size=5, - ragged_mode="static") is not None - - -def test_compact_mode_is_refused(): - """Padding rows contribute prefix labels that measure nothing. - - Refused even with the pin set: the pin fixes the *window*, while this is - about rows that belong to no real request at all. - """ - with collecting(pin=5): - with pytest.raises(ValueError, match="Padded verify rows"): - make_recorder_from_env(block_size=5, ragged_mode="compact") - - -class _Head: - pass - - -class _Bare: - """DSparkDraftModel layout: stages under .mtp_layers.""" - - def __init__(self, head): - stage = type("Stage", (), {})() - if head is not None: - stage.confidence_head = head - self.mtp_layers = [object(), object(), stage] - - -class _Wrapper: - """DSparkForCausalLM layout: bare model under .dspark_model. - - This is the layout the chained getattr missed: the wrapper has no - mtp_layers of its own, so `getattr(model, "mtp_layers", [None])[-1]` - resolved to None and silently disabled calibration for every - confidence-scheduled run of the first campaign. - """ - - def __init__(self, head): - self.dspark_model = _Bare(head) - - -def test_head_resolves_across_both_layouts_and_absence_is_none(): - from tensorrt_llm._torch.speculative.dspark_sts import resolve_confidence_head - h = _Head() - assert resolve_confidence_head(_Bare(h)) is h - # The wrapper layout is the one the chained getattr missed: resolving to - # None here silently disabled calibration for every confidence-scheduled - # run of the first campaign. - assert resolve_confidence_head(_Wrapper(h)) is h - assert resolve_confidence_head(_Bare(None)) is None - assert resolve_confidence_head(_Wrapper(None)) is None - assert resolve_confidence_head(object()) is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py deleted file mode 100644 index 23f2467ce68c..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py +++ /dev/null @@ -1,168 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""STS pairing must join on draft-pass identity, never on arrival order. - -The recorder's label (accepted count) and its feature (confidence logits) are -produced by different steps: the block verified at step ``i`` was drafted at -step ``i-1``, and under the overlap scheduler the sampler consumes the label -while pass ``i+1`` has already overwritten the shared confidence buffer. The -single mutable stash this ring replaced had NO execution order under which -the pair was correct -- measured on job 2562577, 82.2% of rows paired -confidence(t+1) with accepted(t), correlation 0.14 -- and the shards looked -perfectly healthy. - -These tests pin the join contract: a pair is appended only when the ring -snapshot's per-row stamp equals the pass the label verifies, everything else -is a *counted* decline, and the shard carries provenance the fitter can -refuse. No CUDA: the recorder itself is device-agnostic (CPU tensors skip the -pinned-copy event), which is exactly what lets this run in pre-merge CI. -""" - -import importlib.util -import os -import pathlib -import sys -import tempfile - -import pytest -import torch - -from tensorrt_llm._torch.speculative.dspark_sts import DSparkStsRecorder - -BLOCK = 5 -ROWS = 8 - - -def _recorder(tmp, flush_every=1000): - return DSparkStsRecorder(path_stem=os.path.join(tmp, "shard"), - block_size=BLOCK, rank=0, - flush_every=flush_every) - - -def _stage(rec, seq, written_rows, fill): - """Stage one pass's snapshot: `written_rows` stamped `seq`, rest stale. - - Mirrors the producer: the in-graph scatter stamps only the rows the pass - actually drafted; every other row keeps whatever stamp it had (0 for - never-written, matching the worker's zero-initialized stamp buffer). - """ - logits = torch.full((ROWS, BLOCK), float(fill)) - stamps = torch.zeros(ROWS, dtype=torch.int32) - stamps[list(written_rows)] = seq - rec.stage_snapshot(device_logits=logits, device_stamps=stamps, - staged_seq=seq) - - -def test_pairs_by_pass_identity_not_arrival_order(): - """A label for pass 1 gets pass-1 logits even after pass 2 was staged.""" - with tempfile.TemporaryDirectory() as tmp: - rec = _recorder(tmp) - _stage(rec, seq=1, written_rows=[2], fill=1.0) - _stage(rec, seq=2, written_rows=[2], fill=2.0) # overwrite arrives first - rec.record(row=2, accepted=3, target_seq=1) - rec.record(row=2, accepted=1, target_seq=2) - assert rec.stats["recorded"] == 2 - assert rec._logits[0][0].item() == 1.0 # pass 1 content, not latest - assert rec._logits[1][0].item() == 2.0 - - -def test_row_not_written_for_that_pass_is_declined(): - """A row the pass never drafted must not contribute a fabricated pair.""" - with tempfile.TemporaryDirectory() as tmp: - rec = _recorder(tmp) - _stage(rec, seq=1, written_rows=[2], fill=1.0) - rec.record(row=3, accepted=2, target_seq=1) # row 3 stamp is 0, not 1 - assert rec.stats["recorded"] == 0 - assert rec.stats["stale_stamp"] == 1 - - -def test_ring_eviction_declines_instead_of_mispairing(): - """A label older than the ring depth is dropped, never matched to the - - evicting pass's content: the slot's stamps say who actually wrote it. - """ - with tempfile.TemporaryDirectory() as tmp: - rec = _recorder(tmp) - for seq in range(1, 2 + DSparkStsRecorder._RING_DEPTH): - _stage(rec, seq=seq, written_rows=[2], fill=float(seq)) - rec.record(row=2, accepted=3, target_seq=1) # slot now holds seq 5 - assert rec.stats["recorded"] == 0 - assert rec.stats["stale_stamp"] == 1 - - -def test_every_giving_up_path_is_counted(): - """Silent declines are how the mispaired fit shipped; none may be silent.""" - with tempfile.TemporaryDirectory() as tmp: - rec = _recorder(tmp) - rec.record(row=2, accepted=1, target_seq=None) # no pass id - rec.record(row=2, accepted=1, target_seq=7) # ring slot never staged - _stage(rec, seq=1, written_rows=[2], fill=1.0) - rec.record(row=None, accepted=1, target_seq=1) # unresolvable request - rec.record(row=ROWS + 3, accepted=1, target_seq=1) - assert rec.stats == {"recorded": 0, "no_snapshot": 2, - "stale_stamp": 0, "row_out_of_range": 1, - "no_row": 1, "snapshots_staged": 1} - - -def test_staging_without_stamps_is_a_noop(): - """No stamps means no way to verify identity -- refuse to guess.""" - with tempfile.TemporaryDirectory() as tmp: - rec = _recorder(tmp) - rec.stage_snapshot(device_logits=torch.zeros(ROWS, BLOCK), - device_stamps=None, staged_seq=1) - rec.stage_snapshot(device_logits=torch.zeros(ROWS, BLOCK), - device_stamps=torch.zeros(ROWS, dtype=torch.int32), - staged_seq=None) - assert rec.stats["snapshots_staged"] == 0 - rec.record(row=2, accepted=1, target_seq=1) - assert rec.stats["no_snapshot"] == 1 - - -def test_shard_carries_provenance_and_prefix_semantics(): - with tempfile.TemporaryDirectory() as tmp: - rec = _recorder(tmp) - _stage(rec, seq=1, written_rows=[2], fill=1.5) - rec.record(row=2, accepted=3, target_seq=1) - rec.flush() - blob = torch.load(os.path.join(tmp, "shard.r0.0.pt"), - map_location="cpu") - assert blob["meta"]["pairing"] == "draft_seq_ring" - assert blob["meta"]["stats"]["recorded"] == 1 - assert blob["logits"].shape == (1, BLOCK) - # accepted=3 -> positions 0..2 were ALL accepted: the prefix event - # whose probability survival[k] estimates. - assert blob["prefix_mask"].tolist() == [[1.0, 1.0, 1.0, 0.0, 0.0]] - - -# --------------------------------------------------------------------------- -# The fitter must refuse shards from the pre-ring recorders: their pairs are -# mislabeled in ways the tensors themselves cannot reveal. -# --------------------------------------------------------------------------- - -_SPEC = importlib.util.spec_from_file_location( - "dspark_fit_sts", - pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / - "dspark_fit_sts.py") -_FITTER = importlib.util.module_from_spec(_SPEC) -sys.modules[_SPEC.name] = _FITTER -_SPEC.loader.exec_module(_FITTER) - - -def test_fitter_refuses_shards_without_ring_provenance(): - with tempfile.TemporaryDirectory() as tmp: - torch.save({"logits": torch.zeros(4, BLOCK), - "prefix_mask": torch.zeros(4, BLOCK)}, - os.path.join(tmp, "old.r0.0.pt")) - with pytest.raises(SystemExit, match="draft_seq_ring"): - _FITTER.load_shards(os.path.join(tmp, "*.pt")) - - -def test_fitter_accepts_ring_shards(): - with tempfile.TemporaryDirectory() as tmp: - rec = _recorder(tmp) - _stage(rec, seq=1, written_rows=[2, 3], fill=0.5) - rec.record(row=2, accepted=2, target_seq=1) - rec.record(row=3, accepted=5, target_seq=1) - rec.flush() - logits, mask = _FITTER.load_shards(os.path.join(tmp, "*.pt")) - assert logits.shape == (2, BLOCK) and mask.shape == (2, BLOCK) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py deleted file mode 100644 index 2160e1d0ba49..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py +++ /dev/null @@ -1,96 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The budget must be scored with the tau the step actually collects. - -``decide_verify_lens`` picks a rung, converts it to ``B = n*(L - min)`` tokens, -and hands B to ``schedule_verify_lens_topk``, which spends it on the B largest -survivals in the batch. But the rung was chosen by -``budget_argmax_over_uniform_lens``, which scores rung L as if every request took -columns ``[min, L)``. Top-k is by construction the largest sum over all size-B -subsets, so the realised tau is >= the scored one at an identical token cost -- -and the gap is not the same at every rung, so the argmax can land on the wrong -one. - -The fix is not to let the planner choose an arbitrary budget: the executor can -only run ``bs*(t+1)`` tokens for a captured tier, and any other total rounds UP, -which would push a budget deliberately chosen below a cost riser back over it. -So the ladder stays the grid and only the numerator changes. -""" - -import numpy as np -import pytest - -from tensorrt_llm._torch.speculative.dspark_planner import ( - SpsCostTable, budget_argmax_over_uniform_lens, compute_verify_token_budget) - - -def _table() -> SpsCostTable: - """A staircase with genuine risers, so trimming is worth something. - - Breakpoints are spaced so that neither tier is trivially optimal: a flat - table degenerates to verify-all and would make every assertion here vacuous. - """ - token_counts = tuple(range(0, 400, 16)) - step_time_ms = tuple(4.0 + 0.6 * (tok // 96) + 0.004 * tok - for tok in token_counts) - return SpsCostTable(token_counts=token_counts, step_time_ms=step_time_ms, - fixed_overhead_ms=1.0) - - -@pytest.mark.parametrize("tiers", [[1, 2, 5], [1, 3, 5], [1, 5]]) -def test_restricted_answer_is_always_realisable(tiers): - """Every returned budget corresponds to a rung the executor captured.""" - rng = np.random.default_rng(20260803) - table = _table() - for _ in range(200): - bs = int(rng.integers(2, 17)) - survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] - n = compute_verify_token_budget( - survival=survival, num_gen_requests=bs, cost_table=table, - min_verify_len=1, allowed_lens=tiers) - assert n % bs == 0, ( - f"budget {n} for {bs} requests is not n*(t-min) for any tier") - assert (n // bs) + 1 in tiers, ( - f"budget {n} implies tier {(n // bs) + 1}, not in {tiers}") - - -def test_restricted_never_scores_worse_than_the_uniform_choice(): - """Same grid, better numerator: the chosen rung's realised theta wins. - - Both routes evaluate the same discrete cost points, so the restricted form - cannot be worse under the tau that the step actually collects. This is the - property that makes the swap safe -- if it ever failed, the change would be - trading a modelling error for a real one. At bs=1 the uniform and top-k - allocations are the same set, so there the two scorers must agree on the - rung outright: any disagreement is a plain bug, not the intended - improvement. - """ - rng = np.random.default_rng(31337) - table = _table() - tiers = [1, 2, 5] - for _ in range(300): - bs = int(rng.integers(1, 17)) - survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] - - def realised_theta(tier: int) -> float: - budget = bs * (tier - 1) - cand = np.sort(survival[:, 1:].reshape(-1))[::-1] - tau = float(bs) + float(survival[:, :1].sum()) + float( - cand[:budget].sum()) - tokens = np.array([bs * (tier + 1)]) - return tau / float(table.step_times(tokens, bs)[0]) - - n = compute_verify_token_budget( - survival=survival, num_gen_requests=bs, cost_table=table, - min_verify_len=1, allowed_lens=tiers) - chosen = (n // bs) + 1 - uniform = budget_argmax_over_uniform_lens( - survival=survival, num_gen_requests=bs, cost_table=table, - allowed_lens=tiers, min_verify_len=1) - assert realised_theta(chosen) >= realised_theta(uniform) - 1e-12, ( - f"restricted picked tier {chosen} whose realised theta is below " - f"the uniform scorer's tier {uniform}") - if bs == 1: - assert chosen == uniform, ( - "at bs=1 the uniform and top-k allocations are the same set; " - "the two scorers must pick the same rung") diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py deleted file mode 100644 index 04b8d5c0dce6..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py +++ /dev/null @@ -1,97 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The runtime verify-length pin: queue, agree, adopt. - -A cost-table sweep has to hold every request at one verify length while it -measures a cell, or the cell's label describes a shape that never ran (this -branch shipped a sweep where eight of twelve cells were mislabelled and the -fit dutifully concluded the curve was flat). The environment variable that -does this today is read once at planner construction, so walking the ladder -means rebuilding the engine per length -- and on a pre-spawned world it -cannot reach the ranks that are already running. - -The runtime pin fixes both, but it must never let one rank pin while its -peers do not: the ranks would derive different token totals, and the -attention-DP shape gate would drop the whole group out of graph replay for as -long as they disagreed. So the pin is queued locally, carried by the step's -existing allgather, and adopted from that agreed payload -- which is what -these tests pin down. -""" - -import pytest - -from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable -from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig -from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner - - -def _planner(tiers=(1, 2, 5)): - table = SpsCostTable(token_counts=(0, 512, 768, 1536), - step_time_ms=(5.0, 68.4, 80.2, 150.5), - fixed_overhead_ms=1.0) - cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) - return DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=list(tiers)) - - -def test_adopting_the_agreed_value_applies_it(): - """Queuing must not change this rank's behaviour before the group agrees; - adopting the agreed value applies it and consumes the pending request.""" - planner = _planner() - assert planner.request_verify_len_pin(2) == 2 - # Queued, not yet in force. - assert planner._forced_verify_len is None - assert planner.pending_verify_len_pin() == 2 - planner.adopt_verify_len_pin(2) - assert planner._forced_verify_len == 2 - # Consumed: the next step must not re-broadcast a pin nobody asked for. - assert planner.pending_verify_len_pin() == -1 - - -def test_a_rank_that_queued_nothing_still_adopts_the_group_value(): - """The endpoint lands on ONE rank; the others learn it from the payload.""" - peer = _planner() - assert peer.pending_verify_len_pin() == -1 - peer.adopt_verify_len_pin(5) - assert peer._forced_verify_len == 5 - - -def test_zero_clears_the_pin(): - """Wire protocol: -1 means nobody asked (leave the pin alone); clearing - needs its own value, 0, because None cannot travel in an int payload.""" - planner = _planner() - planner.adopt_verify_len_pin(2) - planner.adopt_verify_len_pin(-1) - assert planner._forced_verify_len == 2 - assert planner.request_verify_len_pin(None) is None - assert planner.pending_verify_len_pin() == 0 - planner.adopt_verify_len_pin(0) - assert planner._forced_verify_len is None - - -def test_an_uncaptured_length_is_refused_at_the_call_site(): - """Rejected when requested, not on some later step, and never queued.""" - planner = _planner(tiers=(1, 2, 5)) - with pytest.raises(ValueError, match="captured tier ladder"): - planner.request_verify_len_pin(3) - with pytest.raises(ValueError, match="outside"): - planner.request_verify_len_pin(9) - assert planner.pending_verify_len_pin() == -1 - assert planner._forced_verify_len is None - - -def test_pinned_steps_hand_every_request_the_same_window(): - """The whole point: the shape a sweep cell can honestly label.""" - import numpy as np - import torch - - planner = _planner() - bs = 8 - rng = np.random.default_rng(5) - survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() - planner._gather_rows = lambda **_: torch.tensor(survival, - dtype=torch.float32) - planner.adopt_verify_len_pin(2) - lens = planner.decide_verify_lens(num_gen_requests=bs, - reduce_across_ranks=False) - assert lens == [2] * bs - assert planner.stats["forced_steps"] == 1 diff --git a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py b/tests/unittest/_torch/speculative/test_dspark_graph_key.py similarity index 58% rename from tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py rename to tests/unittest/_torch/speculative/test_dspark_graph_key.py index eff4f9884e5e..df9ba57412c2 100644 --- a/tests/unittest/_torch/speculative/test_dspark_graph_miss_attribution.py +++ b/tests/unittest/_torch/speculative/test_dspark_graph_key.py @@ -1,103 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Which CUDA-graph misses the ragged gate is allowed to blame on ragged. - -A miss can mean opposite things and the counter alone cannot tell them apart. -``peer_not_gen_only`` is ordinary continuous batching -- some rank had context -requests, which drops the whole attention-DP world out of replay and happens on -the uniform path too. ``key_not_captured`` means the batch was graphable, every -rank agreed, and the shape was still missing: the token bucket grid is wrong. -``peer_shape_mismatch`` means one rank chose a different ragged bucket, which is -the attention-DP divergence this feature exists to avoid. - -Gating on the raw count made the check unsatisfiable on any real workload; not -gating at all let a real bucket bug through. These tests pin the split. -""" - -import pytest - -from tensorrt_llm._torch.speculative.dspark_observability import ( - DSparkRaggedStats, RaggedVerifyMode) - - -def _stats(**reasons): - """Stats that look like a healthy ragged run apart from the given misses.""" - stats = DSparkRaggedStats(mode=RaggedVerifyMode.COMPACT, max_draft_len=5) - stats.steps_total = 100 - stats.steps_ragged = 80 - # distinct_verify_lens is derived from this histogram, so populate the - # histogram rather than the property. - stats.verify_len_hist.update({1: 10, 3: 10, 5: 10}) - stats.ceiling_tokens = 1000 - stats.delivered_tokens = 800 # trim_ratio 0.2 > 0, so require_trim passes - stats.window_tokens = 800 - stats.graph_replays = 500 - for reason, count in reasons.items(): - stats.graph_eager += count - stats.graph_miss_reasons[reason] = count - return stats - - -def test_continuous_batching_misses_are_not_blamed_on_ragged(): - """A rank with context requests is normal, not a ragged defect -- and a - fully clean run (no misses at all) passes the same gate.""" - _stats().assert_ragged_active(require_trim=True) - _stats(peer_not_gen_only=9).assert_ragged_active(require_trim=True) - - -@pytest.mark.parametrize( - "reason", ["key_not_captured", "bs_not_supported", "peer_shape_mismatch"]) -def test_ragged_attributable_misses_raise(reason): - with pytest.raises(AssertionError) as excinfo: - _stats(**{reason: 1}).assert_ragged_active(require_trim=True) - assert reason in str(excinfo.value) - - -def test_mixed_misses_raise_and_name_only_the_attributable_ones(): - with pytest.raises(AssertionError) as excinfo: - _stats(peer_not_gen_only=20, - key_not_captured=2).assert_ragged_active(require_trim=True) - message = str(excinfo.value) - assert "key_not_captured" in message - # The tolerated reason must not appear in the *blame* dict, or the message - # sends the reader to the wrong place. It still appears in the summary. - blame = message.split("(", 1)[1].split(")", 1)[0] - assert "peer_not_gen_only" not in blame - - -def test_unattributed_misses_are_a_broken_probe_not_a_clean_run(): - """graph_eager > 0 with no reasons recorded must not read as success. - - This is what actually happened: a build skew left the reason unrecorded and - the summary showed misses with an empty reason map. Silently passing there - would make a broken probe indistinguishable from a healthy run. - """ - stats = _stats() - stats.graph_eager = 9 # counted, but no reason recorded - with pytest.raises(AssertionError, match="cannot be attributed"): - stats.assert_ragged_active(require_trim=True) - +"""The ragged CUDA-graph key: fit, publication, and the padding ladder. -def test_ramp_up_steps_do_not_count_toward_the_gate_floor(): - """A batch of one is uniform by construction, not a scheduler failure. - - trtllm-bench ramps concurrency from zero, so a step-count floor of 64 can be - reached while only two requests are in flight. Every window is then - trivially identical, and gating on total steps failed a run whose scheduler - had done nothing wrong -- observed as "Broadcasting event-loop error to 2 - pending request(s)". The floor counts steps whose batch could actually be - ragged. - """ - stats = DSparkRaggedStats(mode=RaggedVerifyMode.COMPACT, max_draft_len=5) - for _ in range(64): - stats.record_step(num_gen_requests=1, verify_lens=[5]) - assert stats.steps_total == 64 - assert stats.steps_multi_request == 0, ( - "single-request steps must not count toward the gate floor") - - for _ in range(64): - stats.record_step(num_gen_requests=4, verify_lens=[1, 3, 5, 5]) - assert stats.steps_multi_request == 64 +Three regression guards for one invariant -- the graph key's token axis must +describe the batch the graph will actually replay over. Each test pins a bug +that shipped: a fit against unreachable pad rows, a key derived by re-summing +the batch instead of reading the fitted value (the ragged-IMA family), and a +capture pass that shaped a batch without publishing its bucket. +""" class _Runner: diff --git a/tests/unittest/_torch/speculative/test_dspark_planner_trims.py b/tests/unittest/_torch/speculative/test_dspark_planner_trims.py deleted file mode 100644 index bcb12eecb37c..000000000000 --- a/tests/unittest/_torch/speculative/test_dspark_planner_trims.py +++ /dev/null @@ -1,75 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Make the planner itself produce per-request windows, with a crafted table. - -This replaces TLLM_DSPARK_FORCE_VERIFY_LENS. That knob assigned windows by -*batch position*, which is orthogonal to confidence: it reproduced the ragged -*shape* but never the ragged *policy*, so it could not tell a correct -confidence-driven assignment from an arbitrary one. It also had to sit in -production code, ahead of the cost-table gate, to be reachable at all. - -A cost table does the same job through the real path. The planner trims when -the marginal cost of extra verify tokens outweighs their expected acceptance, -so a table whose theta(M) rises steeply makes it trim at any acceptance rate -- -and the budget then flows through the same confidence-ordered top-k, which -hands different windows to requests with different survival. That exercises -what actually ships. - -Why this matters beyond tidiness: the KV leak fixed in this branch (KV reserved -for the full drafted block but only the verified window rewound) is reachable -only when verified_len < draft_len. On this checkpoint the planner declines to -trim at every batch size, because draft acceptance is ~90% and the fitted table -says verifying the full block wins -- so without a way to make it trim, that -code path ships untested. -""" - -import torch - -from tensorrt_llm._torch.speculative.dspark_schedule import ( - DSparkScheduleConfig, schedule_verify_lens_topk) - - -def _Cfg(min_verify_len=1, max_verify_len=5): - """The real config, so the test cannot drift from what ships.""" - return DSparkScheduleConfig(block_size=max_verify_len, - min_verify_len=min_verify_len, - max_verify_len=max_verify_len) - - -def test_topk_hands_longer_windows_to_more_confident_requests(): - """The property the position-rotating knob could never check. - - Request 0 is confident throughout, request 2 collapses immediately. A - correct scheduler spends the budget where survival is highest. - """ - # survival[r, k] = P(the first k+1 drafted tokens all get accepted) - survival = torch.tensor([ - [0.99, 0.98, 0.97, 0.96, 0.95], # confident - [0.90, 0.70, 0.50, 0.30, 0.10], # middling - [0.20, 0.04, 0.01, 0.00, 0.00], # collapses - ]) - lens = schedule_verify_lens_topk(survival=survival, budget=6, - cfg=_Cfg()).tolist() - - assert len(lens) == 3 - assert lens[0] > lens[2], ( - f"confident request got {lens[0]}, collapsing request got {lens[2]}; " - f"the budget must follow survival, not batch position") - assert lens[0] >= lens[1] >= lens[2] - # Every request keeps at least the floor, and the budget is respected. - assert min(lens) >= 1 - assert sum(lens) - len(lens) * 1 <= 6 - - -def test_a_full_budget_degenerates_to_the_uniform_full_window(): - """The no-trim case must stay exactly uniform. - - This is what the planner correctly chooses on this checkpoint, and it is - the case where ragged must cost nothing. - """ - num_reqs, max_len = 4, 5 - survival = torch.full((num_reqs, max_len), 0.99) - budget = num_reqs * (max_len - 1) - lens = schedule_verify_lens_topk(survival=survival, budget=budget, - cfg=_Cfg()).tolist() - assert lens == [max_len] * num_reqs diff --git a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py deleted file mode 100644 index 144de6a6c174..000000000000 --- a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py +++ /dev/null @@ -1,363 +0,0 @@ -# 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. -"""Unit tests for the synthetic forced acceptance rate in 1-model spec dec. - -These tests exercise ``SpecWorkerBase._apply_force_accepted_tokens``, which -honors the ``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` environment -variable. The implementation must be: - - - Faithful to fractional rates: integer part always accepts, fractional - part acts as the per-iteration probability of accepting one extra draft - token. - - Tensor-parallel deterministic: identical seed/counter state across - ranks so all ranks agree on per-request accepted counts (otherwise - downstream collectives expecting matching shapes deadlock). - - CUDA-graph compatible: pure-tensor ops with statically-known shapes, - captureable and replayable. -""" - -import os -from typing import Optional -from unittest.mock import patch - -import pytest -import torch - -from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata -from tensorrt_llm._torch.speculative.interface import ( - FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR, - SpecWorkerBase, - get_force_num_accepted_tokens, - get_force_num_accepted_tokens_float, -) - - -class _StubSpecWorker(SpecWorkerBase): - """Concrete ``SpecWorkerBase`` that stubs out the abstract API. - - Used purely to drive ``_apply_force_accepted_tokens`` in isolation. - """ - - @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: - worker = _StubSpecWorker() - if value is not None: - worker.force_num_accepted_tokens = value - return worker - - -def _make_input(batch_size: int, device: str = "cuda") -> torch.Tensor: - # Initial counts: 1 (target token, no draft accepts yet). The forced - # override only ever rewrites entries from ``num_contexts:`` onward. - return torch.ones(batch_size, dtype=torch.int32, device=device) - - -def _require_cuda(): - if not torch.cuda.is_available(): - pytest.skip("CUDA required for this test.") - - -# ---------------- env-var parsing helpers ----------------------------------- - - -@pytest.mark.parametrize( - "env_value, expected_float, expected_int", - [ - ("0", 0.0, 0), - ("3", 3.0, 3), - ("2.6", 2.6, 0), - ("0.5", 0.5, 0), - ("not-a-number", 0.0, 0), - ], -) -def test_get_force_num_accepted_tokens_float(env_value, expected_float, expected_int): - """The float helper parses fractional rates; the int helper (used by the - 2-model path) must keep its original behavior on the same values.""" - with patch.dict(os.environ, {FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR: env_value}): - assert get_force_num_accepted_tokens_float() == pytest.approx(expected_float) - assert get_force_num_accepted_tokens() == expected_int - - -# ---------------- _apply_force_accepted_tokens semantics -------------------- - - -def test_zero_value_is_noop(): - _require_cuda() - worker = _make_worker(0.0) - base = _make_input(batch_size=4) - before = base.clone() - out = worker._apply_force_accepted_tokens(base, num_contexts=0, runtime_draft_len=4) - assert torch.equal(out, before) - # No RNG state should have been touched in the early-exit path. - assert worker._force_accept_rng_pool is None - assert worker._force_accept_rng_counter is None - - -@pytest.mark.parametrize( - "value, runtime_draft_len", - [ - (1.0, 4), - (2.0, 4), - (4.0, 4), - # 10 draft tokens requested but only 2 available -> capped at 3. - (10.0, 2), - # Fractional but no room: max_total = 3, base_total = min(10, 3) = 3, - # so ``base_total < max_total`` is False -> all 3, no RNG draw. - (9.5, 2), - ], -) -def test_integer_value_matches_legacy_behavior(value, runtime_draft_len): - _require_cuda() - worker = _make_worker(value) - out = worker._apply_force_accepted_tokens( - _make_input(batch_size=8), num_contexts=0, runtime_draft_len=runtime_draft_len - ) - expected = min(int(value) + 1, runtime_draft_len + 1) - assert torch.all(out == expected) - # The deterministic (non-random) path must not allocate the RNG pool. - assert worker._force_accept_rng_pool is None - - -def test_fractional_distribution_matches_target_probability(): - """Fractional draws emit only base/base+1 and average to the target rate.""" - _require_cuda() - worker = _make_worker(2.6) - target_frac = 0.6 - n_iters = 200 - batch = 64 - seen = set() - extra_count = 0 - total = 0 - for _ in range(n_iters): - out = worker._apply_force_accepted_tokens( - _make_input(batch_size=batch), num_contexts=0, runtime_draft_len=4 - ) - seen.update(out.unique().tolist()) - extra_count += int((out == 4).sum().item()) - total += batch - # Either 2 draft + target = 3, or 3 draft + target = 4 -- never anything else. - assert seen == {3, 4} - measured = extra_count / total - # Pool-based RNG is deterministic; the empirical mean over ~13k draws - # is tight against 0.6. - assert abs(measured - target_frac) < 0.03, ( - f"measured fraction {measured:.4f} differs from target {target_frac}" - ) - - -def test_num_contexts_offset_is_respected(): - """Context (prefill) rows must be left untouched by the override.""" - _require_cuda() - worker = _make_worker(2.6) - batch_size = 8 - num_contexts = 3 - base = _make_input(batch_size=batch_size) - sentinel = torch.iinfo(base.dtype).max - base[:num_contexts] = sentinel - out = worker._apply_force_accepted_tokens(base, num_contexts=num_contexts, runtime_draft_len=4) - assert torch.all(out[:num_contexts] == sentinel) - assert torch.all((out[num_contexts:] == 3) | (out[num_contexts:] == 4)) - - -# ---------------- TP determinism -------------------------------------------- - - -def test_tp_determinism_across_independent_workers(): - """ - Two independently-instantiated workers (simulating two TP ranks) must - produce bit-identical accepted-token counts when fed the same call - sequence. This is the property whose absence used to hang TP>1. - """ - _require_cuda() - rank0 = _make_worker(2.6) - rank1 = _make_worker(2.6) - # Use a non-power-of-two batch to make accidental shape coincidences - # less likely to mask divergence. - for _ in range(32): - out0 = rank0._apply_force_accepted_tokens( - _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 - ) - out1 = rank1._apply_force_accepted_tokens( - _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 - ) - assert torch.equal(out0, out1) - - -def test_tp_determinism_survives_default_generator_drift(): - """ - Even if the default CUDA generator state diverges between ranks (e.g. - because some other component consumed random numbers), the synthetic - AR draws must still match. This is the regression guard for the - original bug, where ``torch.rand`` against the default generator gave - each rank a different sequence. - """ - _require_cuda() - rank0 = _make_worker(2.6) - rank1 = _make_worker(2.6) - - def _advance_default_generator(seed: int): - # Seed the default CUDA generator on the same device differently - # for each "rank" and consume some randoms, simulating drift. - torch.cuda.manual_seed(seed) - _ = torch.rand(128, device="cuda") - - _advance_default_generator(seed=11) - out0 = rank0._apply_force_accepted_tokens( - _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 - ) - _advance_default_generator(seed=999) - out1 = rank1._apply_force_accepted_tokens( - _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 - ) - assert torch.equal(out0, out1) - - -# ---------------- CUDA graph capture/replay --------------------------------- - - -def test_cuda_graph_capture_and_replay_match_eager(): - """ - Capturing the override into a CUDA graph and replaying it must produce - the exact same per-replay output sequence as running it eagerly. - """ - _require_cuda() - - n_iters = 8 - batch_size = 16 - runtime_draft_len = 4 - value = 2.6 - - # 1) Eager reference. - eager_worker = _make_worker(value) - eager_outputs = [] - for _ in range(n_iters): - out = eager_worker._apply_force_accepted_tokens( - _make_input(batch_size=batch_size), num_contexts=0, runtime_draft_len=runtime_draft_len - ) - eager_outputs.append(out.clone()) - - # 2) Captured-graph run, aligned to the eager reference. - graph_worker = _make_worker(value) - static_input = _make_input(batch_size=batch_size) - # Eager warmup: forces lazy RNG state allocation OUTSIDE capture. - graph_worker._apply_force_accepted_tokens( - static_input, num_contexts=0, runtime_draft_len=runtime_draft_len - ) - # Realign the device-side counter so the captured graph's first replay - # advances 0 → 1, matching the eager loop's first iteration. - graph_worker._force_accept_rng_counter.zero_() - - static_input.fill_(1) - static_output = torch.empty_like(static_input) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_worker._apply_force_accepted_tokens( - static_input, num_contexts=0, runtime_draft_len=runtime_draft_len - ) - # ``_apply_force_accepted_tokens`` mutates ``static_input`` in place; - # mirror it into ``static_output`` so we can snapshot per replay. - static_output.copy_(static_input) - - graph_outputs = [] - for _ in range(n_iters): - # Reset input each replay so the override semantics match eager - # (which also starts from all-1s every iteration). - static_input.fill_(1) - graph.replay() - torch.cuda.synchronize() - graph_outputs.append(static_output.clone()) - - for i, (eager, graphed) in enumerate(zip(eager_outputs, graph_outputs)): - assert torch.equal(eager, graphed), ( - f"Iteration {i}: eager={eager.tolist()} vs graphed={graphed.tolist()}" - ) - - -def test_cuda_graph_replay_advances_rng_state(): - """ - Across replays, the captured counter must advance, producing different - (but deterministic) draws. A constant output across replays would mean - the counter increment was not captured. - """ - _require_cuda() - - worker = _make_worker(0.5) - batch_size = 64 - runtime_draft_len = 4 - - static_input = _make_input(batch_size=batch_size) - worker._apply_force_accepted_tokens( - static_input, num_contexts=0, runtime_draft_len=runtime_draft_len - ) - worker._force_accept_rng_counter.zero_() - - static_input.fill_(1) - static_output = torch.empty_like(static_input) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - worker._apply_force_accepted_tokens( - static_input, num_contexts=0, runtime_draft_len=runtime_draft_len - ) - static_output.copy_(static_input) - - snapshots = [] - for _ in range(8): - static_input.fill_(1) - graph.replay() - torch.cuda.synchronize() - snapshots.append(static_output.clone()) - - # At least two replays must differ; otherwise the captured counter is - # not advancing (the original bug-class). - assert any(not torch.equal(snapshots[0], s) for s in snapshots[1:]), ( - "Captured graph produced identical outputs on every replay — " - "RNG counter is not advancing inside the captured graph." - ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) - - -class _FailingSpecWorker(_StubSpecWorker): - """Stub whose forward saves spec-dec metadata state and then fails.""" - - def _forward_impl(self, *args: object, **kwargs: object) -> None: - attn_metadata = kwargs["attn_metadata"] - attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") - raise RuntimeError("simulated draft failure") - - -def test_forward_restores_spec_dec_state_on_failure() -> None: - """A failure between prepare_for_spec_dec and restore must not leak saved - attn-metadata state; SpecWorkerBase.forward restores it in its cleanup - (https://nvbugs/6442074).""" - _require_cuda() - attn_metadata = AttentionMetadata(max_num_requests=2, max_num_tokens=16) - attn_metadata.seq_lens = torch.ones(2, dtype=torch.int32) - worker = _FailingSpecWorker() - for _ in range(2): - # The second iteration would trip the pairing assert inside - # prepare_for_spec_dec if the first failure had leaked saved state. - with pytest.raises(RuntimeError, match="simulated draft failure"): - worker(attn_metadata=attn_metadata, spec_metadata=None) - assert not attn_metadata.has_spec_dec_saved_state diff --git a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py deleted file mode 100644 index 898699b69e91..000000000000 --- a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py +++ /dev/null @@ -1,56 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Ragged verification must refuse to start without a profiled cost table. - -Without one the planner's budget degenerates to verify-all -- correctly, since -a flat cost model makes every extra verify token look free -- so the ragged -path runs and changes nothing. As a runtime warning that produced 27-minute -runs which "succeeded" having done nothing, detectable only from a counter -(``steps_ragged: 0``) that nobody reads unless they already suspect the -problem. Config errors belong at construction. - -There is no exception: a crafted (non-flat) cost table is how a test makes the -planner trim, which exercises the real confidence-driven path rather than -bypassing it. See tests/unittest/_torch/speculative/test_dspark_planner_trims.py. -""" - -import pytest - -from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig - - -def _cfg(**kwargs): - base = dict(max_draft_len=5, - speculative_model="/nonexistent/model", - enable_confidence_scheduling=True, - enable_ragged_verify=True) - base.update(kwargs) - return DSparkDecodingConfig(**base) - - -def test_ragged_without_a_cost_table_is_rejected(): - with pytest.raises(ValueError, match="requires a profiled cost table"): - _cfg() - # The same gate must open once a table is named -- a gate that rejects - # everything would pass the raise above. - assert _cfg(confidence_sps_table_path="/tmp/table.json").enable_ragged_verify - - -def test_confidence_scheduling_without_ragged_is_rejected(): - """There is no uniform tier ladder to fall back to. - - Scheduling used to have a middle state: pick one verify length for the - whole batch from the captured tiers. It could not act on per-request - confidence -- which is the entire point of a confidence head -- and it - degenerated to the full block whenever acceptance was high, i.e. on every - workload measured on this checkpoint. It is gone, so this combination now - names a path that does not exist and must not silently mean something else. - - The two remaining states are: schedule per request (ragged), or verify the - whole drafted block (confidence scheduling off). - """ - with pytest.raises(ValueError, match="enable_ragged_verify"): - DSparkDecodingConfig(max_draft_len=5, - speculative_model="/nonexistent/model", - enable_confidence_scheduling=True, - enable_ragged_verify=False) From b7e091d5a7d18bc43172454196db93bd6f260c74 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 5 Aug 2026 08:35:55 -0700 Subject: [PATCH 096/177] [None][fix] Choose the ragged bucket every rank can decompose The ADP shape gate compares each rank's ragged token bucket, and the bucket used to be chosen in two stages with different scopes: the bucket VALUE came from the allgathered peer stats (identical everywhere), but whether a rank could decompose it into its real rows plus the shared-window pad rows was checked locally, against rank-local row and token counts. Pad rows all carry one shared window (a single dummy object), so decomposability is a divisibility question, and it is sharpest exactly where nothing else varies: under a full-block pin, real rows have zero slack, and (bucket - floor) must divide evenly among this rank's pad rows. A rank that failed declined the fit alone and published no bucket while its peers published one; the gate saw the divergence and took the whole group eager. Measured on job 2586075 (pinned): 70 of 962 steps, at ~150 ms per miss -- most of the arm's 14% regression -- with runs of consecutive misses wherever ranks straddled a bucket boundary during ramp and drain. choose_ragged_capture_shape now takes the verify-window cap and picks the smallest captured bucket that every rank in the payload can decompose, walking past candidates any rank cannot realise; when none works it raises for every rank identically, and the group falls back to uniform together -- which replays the uniform graph instead of splitting keys. Accept versus decline is thereby a pure function of the allgathered payload, the same principle the runtime pin adoption already follows. The local decomposition checks in fit_ragged_verify_lens remain as safety nets and now warn instead of silently declining, because a hit there means the two arithmetics diverged. Cost: on the steps that used to diverge, at most one bucket step of extra verification, spent on real tokens by fill_bucket. Ordinary trimmed batches keep the exact bucket they got before (pinned by test_slacky_batch_keeps_the_smallest_bucket). 294 hw-agnostic DSpark unit tests green, including three new group-agreement regressions (the job-2586075 divergence miniaturised, a group-consistent decline on a truncated grid, and the no-growth guarantee). Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 16 +++- .../_torch/speculative/dspark_ragged.py | 58 ++++++++++++- .../hw_agnostic/test_dspark_ragged.py | 85 +++++++++++++++++++ 3 files changed, 154 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index e1d1c404f098..9a3ded9b5553 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3863,7 +3863,13 @@ def fit_ragged_verify_lens( total_verify_tokens=sum(token_lens), bs_buckets=bs_buckets, token_buckets=buckets, - peer_stats=peer_stats) + peer_stats=peer_stats, + # With the window cap the chooser also enforces per-rank + # decomposability from the allgathered stats, so accept vs + # decline is identical on every rank -- the local checks + # below become unreachable safety nets instead of the + # group-splitting decision points they used to be. + max_verify_len=max_verify_len) except ValueError as exc: logger.debug(f"DSpark ragged shape selection failed, falling back " f"to uniform scheduling: {exc}") @@ -3895,7 +3901,11 @@ def fit_ragged_verify_lens( lo = max(1, -(-(bucket - real_capacity) // n_pad)) # ceil division hi = min(max_verify_len, (bucket - floor_tokens) // n_pad) if lo > hi: - logger.debug( + # Unreachable when the chooser ran with max_verify_len (it + # verified this rank's decomposition from the payload); a hit + # here means the two arithmetics diverged -- worth a warning, + # not a silent debug line. + logger.warning( f"DSpark ragged: bucket {bucket} admits no pad-row window " f"for {n_real} real ({floor_tokens}..{real_capacity} " f"tokens) + {n_pad} pad rows; falling back to uniform " @@ -3904,7 +3914,7 @@ def fit_ragged_verify_lens( pad_len = lo real_target = bucket - n_pad * pad_len if real_target < floor_tokens or real_target > real_capacity: - logger.debug( + logger.warning( f"DSpark ragged: bucket {bucket} leaves {real_target} tokens " f"for {n_real} real requests ({n_pad} pad rows at {pad_len}), " f"outside [{floor_tokens}, {real_capacity}]; falling back " diff --git a/tensorrt_llm/_torch/speculative/dspark_ragged.py b/tensorrt_llm/_torch/speculative/dspark_ragged.py index 0fdb42e58b98..c977e14cb8fe 100644 --- a/tensorrt_llm/_torch/speculative/dspark_ragged.py +++ b/tensorrt_llm/_torch/speculative/dspark_ragged.py @@ -321,6 +321,32 @@ class RaggedCaptureShape: bucket: int +def _bucket_decomposable(bucket: int, rows: int, tokens: int, padded_bs: int, + max_verify_len: int) -> bool: + """Can a rank with ``rows`` real requests (``tokens`` floor) realise ``bucket``? + + Mirrors the pad-row arithmetic in ``fit_ragged_verify_lens`` exactly: the + ``padded_bs - rows`` pad rows all carry ONE shared window in + ``[1, max_verify_len]`` (they are a single shared dummy object), and the + real rows absorb the remainder within ``[tokens, rows * max_verify_len]``. + + This is computed here, from the allgathered ``(rows, tokens)``, so that + every rank evaluates every OTHER rank's feasibility too: a bucket that + rank 3 cannot decompose used to make rank 3 alone decline the fit while + its peers published the bucket -- a graph-key divergence the ADP shape + gate then turned into a whole-group eager step. Measured on job 2586075 + (pinned full block, where floor == capacity and the divisibility trap is + sharpest): 70 of 962 steps went eager on exactly this. + """ + n_pad = padded_bs - rows + floor, cap = tokens, rows * max_verify_len + if n_pad == 0: + return floor <= bucket <= cap + lo = max(1, -(-(bucket - cap) // n_pad)) + hi = min(max_verify_len, (bucket - floor) // n_pad) + return lo <= hi + + def choose_ragged_capture_shape( *, num_real_requests: int, @@ -328,6 +354,7 @@ def choose_ragged_capture_shape( bs_buckets: Sequence[int], token_buckets: Sequence[int], peer_stats: Optional[Sequence[Sequence[int]]] = None, + max_verify_len: Optional[int] = None, ) -> RaggedCaptureShape: """Pick the captured ``(padded_bs, bucket)`` for a ragged batch. @@ -355,6 +382,16 @@ def choose_ragged_capture_shape( Raises when no captured bucket fits, rather than clamping: a batch with no graph runs eager, which costs far more than the trimming saves, so the caller has to see it and shrink the batch instead. + + With ``max_verify_len`` given, the answer is additionally required to be + DECOMPOSABLE BY EVERY RANK (see ``_bucket_decomposable``): the smallest + such bucket is chosen, walking past candidates any rank cannot realise. + This makes accept/decline a pure function of the allgathered payload -- + all ranks pick the same bucket or all raise together -- where previously + the per-rank decomposition check ran locally and a skewed rank declined + alone, splitting the group's graph keys. The cost is at most one bucket + step of extra verification on the steps that used to diverge, spent on + real tokens by ``fill_bucket``. """ if not bs_buckets: raise ValueError("choose_ragged_capture_shape requires bs_buckets") @@ -381,8 +418,25 @@ def choose_ragged_capture_shape( padded_bs = sorted_bs[bisect.bisect_left(sorted_bs, max_real)] needed = padded_bs + max_slack - bucket = round_up_to_bucket(needed, sorted({int(t) for t in token_buckets})) - return RaggedCaptureShape(padded_bs=padded_bs, bucket=bucket) + candidates = sorted({int(t) for t in token_buckets}) + if max_verify_len is None: + bucket = round_up_to_bucket(needed, candidates) + return RaggedCaptureShape(padded_bs=padded_bs, bucket=bucket) + # Group-consistent choice: smallest captured bucket every rank can + # decompose. The floor requirement (bucket >= tokens_r + pad rows at one + # token each, i.e. >= needed at the widest rank) is subsumed by the + # decomposability test, so no separate >= needed cut is applied. + for cand in candidates: + if all( + _bucket_decomposable(cand, int(peer[0]), int(peer[1]), + padded_bs, int(max_verify_len)) + for peer in stats): + return RaggedCaptureShape(padded_bs=padded_bs, bucket=cand) + raise ValueError( + f"no captured bucket (grid {candidates}) is decomposable by every " + f"rank (peer stats {list(stats)}, padded_bs {padded_bs}, " + f"max_verify_len {max_verify_len}); the group falls back to uniform " + f"together rather than splitting its graph keys") def ragged_gather_index_lists( diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py index ef65a9d25d29..6b18e41b943d 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py @@ -416,6 +416,91 @@ def test_shape_requires_bs_buckets(): ) +# --------------------------------------------------------------------------- +# group-consistent bucket choice (max_verify_len given) +# +# The regression these pin: the chooser used to pick the smallest bucket >= +# needed and leave "can MY rank decompose it into real + pad rows" to a LOCAL +# check in fit_ragged_verify_lens. Pad rows all share one window (single +# shared dummy), so decomposability is a divisibility question -- sharpest +# under a full-block pin where real rows have zero slack -- and a rank that +# failed it declined ALONE while its peers published the bucket. The ADP +# shape gate then took the whole group eager: 70/962 steps on job 2586075. +# --------------------------------------------------------------------------- + +MAXLEN_BUCKETS = [16, 24, 32, 40, 48] # padded_bs 8, tiers 1..5 -> 8*(t+1) + + +def _shape6(num_real, total, peers): + return choose_ragged_capture_shape( + num_real_requests=num_real, + total_verify_tokens=total, + bs_buckets=BS_BUCKETS, + token_buckets=MAXLEN_BUCKETS, + peer_stats=peers, + max_verify_len=6, + ) + + +def test_pinned_ranks_agree_on_a_bucket_every_rank_can_decompose(): + """The job-2586075 divergence, miniaturised. + + Both ranks are pinned to the full block (floor == capacity). needed = + 8 + 25 = 33 rounds to 40, which rank (5, 30) cannot decompose (3 pad rows, + (40-30)/3 is not an integer window) while rank (4, 24) can (4 rows at + window 4 exactly) -- the old code split the group here. The chooser must + now walk past 40 and land BOTH ranks on 48. + """ + peers = [(5, 30), (4, 24)] + shapes = {_shape6(num_real=r, total=t, peers=peers) for r, t in peers} + assert len(shapes) == 1 + assert shapes.pop().bucket == 48 + + +def test_slacky_batch_keeps_the_smallest_bucket(): + """Feasibility filtering must not grow the bucket when real rows have + slack to absorb the remainder -- the ordinary trimmed batch keeps the + exact answer the ungated chooser gave.""" + peers = [(3, 9), (6, 30), (5, 11)] + got = _shape6(num_real=3, total=9, peers=peers) + ungated = choose_ragged_capture_shape( + num_real_requests=3, total_verify_tokens=9, + bs_buckets=BS_BUCKETS, token_buckets=MAXLEN_BUCKETS, + peer_stats=peers) + assert got.bucket == ungated.bucket == 32 + assert got.padded_bs == ungated.padded_bs == 8 + + +def test_group_declines_together_when_nothing_decomposes(): + """When no captured bucket works for every rank, every rank must see the + same ValueError -- a group-consistent decline replays the uniform graph, + where a lone decline used to split the graph keys. + + A full-block-pinned rank can always decompose the FULL bucket (pad rows + at the max window), so total infeasibility requires a grid whose top tier + is missing -- which capture produces for real when max_num_tokens + truncates the ladder at the largest batch size. + """ + truncated = [16, 24, 32, 40] # top tier 48 not captured + peers = [(5, 30), (4, 24)] # pinned: 40 fails (5,30), 32 fails it too + for r, t in peers: + with pytest.raises(ValueError, match="decomposable by every rank"): + choose_ragged_capture_shape( + num_real_requests=r, total_verify_tokens=t, + bs_buckets=BS_BUCKETS, token_buckets=truncated, + peer_stats=peers, max_verify_len=6) + + +def test_uniform_batch_shape_matches_the_one_dimensional_rule(): + # With a uniform K the token total is bs * (K + 1), so the joint rule must + # not pick anything larger than the uniform path would have needed. + k_plus_1 = 4 + for num_real in (1, 3, 8): + got = _shape(num_real=num_real, total=num_real * k_plus_1) + assert got.padded_bs >= num_real + assert got.bucket >= num_real * k_plus_1 + + # --------------------------------------------------------------------------- # fill_padded_rows_onehot # --------------------------------------------------------------------------- From 22f93300bffa9e6bad9fa9d2881a94fba4d4230c Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 5 Aug 2026 08:37:11 -0700 Subject: [PATCH 097/177] [None][fix] Server sweeps under attention DP promise no bs grid The requested-grid check assumed a steady client concurrency pins the observed batch size. Without attention DP it does; with it the router spreads the load and each rank's occupancy is emergent -- 1024 requests over 8 ranks oscillates around 128 per rank and pads to whatever ladder entry the drift lands on (job 2585368: requested 128, measured at 192; only the cap-saturated 256 landed exactly). The pin guarantees L; bs is whatever the deployment did. Under ADP, require nothing, keep what is thick, and let the fit's grid-connectivity check catch a real hole. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- tests/microbenchmarks/dspark_sps_profiler.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/tests/microbenchmarks/dspark_sps_profiler.py b/tests/microbenchmarks/dspark_sps_profiler.py index 7f1a9a4589d5..43691fc441da 100644 --- a/tests/microbenchmarks/dspark_sps_profiler.py +++ b/tests/microbenchmarks/dspark_sps_profiler.py @@ -2196,11 +2196,19 @@ def main(argv: Optional[Sequence[str]] = None) -> int: print(f"[dspark-sps] {len(samples)} steps from {args.base_url}", file=sys.stderr) provenance = {"source": "server_sweep", "base_url": args.base_url} - # The requested concurrencies are ladder entries, so a steady cell's - # padded batch size lands exactly on them; everything else in the - # server's log is admission ramp or drain-down. - expected_cells = frozenset( + # Which cells the sweep can PROMISE depends on who controls the batch + # size. Without attention DP the server schedules one global batch, so + # a steady client concurrency IS the row count and the requested grid + # is enforceable. Under attention DP the router spreads the load and + # each rank's occupancy is emergent -- 1024 requests over 8 ranks + # oscillates around 128/rank and pads to whatever ladder entry the + # drift lands on (job 2585368: requested 128, measured at 192; only + # the cap-saturated 256 landed exactly). The pin guarantees L; bs is + # whatever the deployment did. Require nothing, keep what is thick, + # and let the fit's grid-connectivity check catch a real hole. + expected_cells = (frozenset( (int(b), int(l)) for b in batch_sizes for l in verify_lens) + if not args.enable_attention_dp else frozenset()) else: if not args.model: raise SystemExit("--model is required unless --fit-only is used") From 53644827cf30d7e8b75de1f371e31097d64ec9ad Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 5 Aug 2026 08:54:36 -0700 Subject: [PATCH 098/177] [None][chore] Strip war-story comments, residual getattr, and debug env knobs Comments and docstrings across the DSpark production modules cut to the live contracts: ordering constraints, capture-safety rules, cross-rank invariants, shape/unit conventions, and fail-closed semantics. Bug archaeology, measurements from past runs, design-alternative narratives, and cross-file line references are gone -- git history holds them. Remaining defensive getattr calls inlined where the attribute is unconditionally declared (dataclass fields, __init__ assignments, or the spec-sampler base-class defaults); the genuinely polymorphic guards (e.g. the spec-worker probe on the engine, the lazily attached _dspark_worker back-reference) stay as getattr. Debug environment knobs removed with their scaffolding: TLLM_DSPARK_SYNC_DEBUG (sync-detection arming for an investigation that ended) and TLLM_DSPARK_ASSERT_ROWS (a device-sync row-map assert whose failure mode the layout assert covers). The five knobs with real consumers stay: FORCE_VERIFY_LEN and RAGGED_VERIFY_MODE (the profiler and the cap-accept mode's only entry points), STS_COLLECT_PATH (calibration collection), ASSERT_ACTIVE (the ragged-liveness gate), and ASSERT_LAYOUT (opt-in layout cross-check). No behavior change on any default path. Every touched file byte-compiles; the DSpark suite passes end to end on B300 (223 tests). Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa.py | 430 +++++---------- .../_torch/models/dspark/attention.py | 163 ++---- tensorrt_llm/_torch/models/dspark/heads.py | 77 +-- .../_torch/pyexecutor/cuda_graph_runner.py | 124 ++--- .../_torch/pyexecutor/model_engine.py | 522 +++++------------- tensorrt_llm/_torch/pyexecutor/py_executor.py | 334 +++-------- .../_torch/pyexecutor/sampler/sampler.py | 11 +- tensorrt_llm/_torch/speculative/dspark.py | 474 ++++++---------- .../speculative/dspark_observability.py | 271 +++------ .../_torch/speculative/dspark_planner.py | 245 +++----- .../_torch/speculative/dspark_ragged.py | 216 ++------ .../_torch/speculative/dspark_schedule.py | 89 +-- tensorrt_llm/_torch/speculative/dspark_sts.py | 143 ++--- .../_torch/speculative/dspark_verify.py | 227 +++----- tensorrt_llm/_torch/speculative/interface.py | 114 ++-- 15 files changed, 1019 insertions(+), 2421 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 50568dbd0621..4107f0aa2945 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -657,29 +657,20 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): # Number of compressed KV tokens for context requests num_ctx_kv_tokens: int = 0 gen_indexer_kv_lens_cuda_runtime: Optional[torch.Tensor] = None - # Ragged verification: how many query tokens each generation request - # contributes this step, host-side and in batch order. None (the default, - # and every non-DSpark path) means the uniform `1 + max_draft_tokens` per - # request, which is what the strided expansions below assume. Kept on the - # host because the expansions are built host-side and copied in; reading - # the repeats off a device tensor would sync. + # Ragged verification: query tokens per generation request this step, + # host-side and in batch order. None (the default) means the uniform + # `1 + max_draft_tokens` the strided expansions assume. Host-side because + # reading repeats off a device tensor would sync. ragged_verify_lens: Optional[List[int]] = None - # Query tokens per generation request on a *uniform* batch this step, i.e. - # `1 + runtime_draft_len`. Zero means "not published", and every expansion - # below falls back to the static `1 + max_draft_tokens` -- which is what - # every caller that never shortens its draft length wants, and keeps them - # bit-identical. + # Query tokens per generation request on a *uniform* batch this step + # (`1 + runtime_draft_len`). Zero means "not published"; expansions then + # fall back to the static `1 + max_draft_tokens`. # - # This is deliberately NOT `max_draft_tokens`. That field is the static - # ceiling: it sizes every expanded buffer (see `create_expanded_buffers`) - # and must stay at the maximum so a shorter step never reallocates inside a - # CUDA-graph capture. The *stride* has to follow the runtime tier instead, - # because `Indexer.forward` slices the very same buffers by the true token - # count (`num_gen_tokens // num_generations`, see the decode branch). Left - # at the ceiling, a step that verifies fewer positions writes each - # request's kv_len/block-table row `1 + max_draft_tokens` apart and then - # reads them back `1 + runtime_draft_len` apart, so every request past the - # first silently picks up a neighbour's rows. + # Deliberately NOT `max_draft_tokens`: that ceiling sizes the expanded + # buffers and must never change under CUDA-graph capture, while the + # *stride* must follow the runtime tier -- `Indexer.forward` slices the + # same buffers by the true token count, so a ceiling stride makes every + # request past the first read a neighbour's rows. runtime_tokens_per_gen_step: int = 0 @property @@ -703,14 +694,9 @@ def gen_token_stride(self) -> int: def gen_token_repeat_list(self) -> List[int]: """Per-generation-request query-token counts, on the host. - Uniform batches get `gen_token_stride` for every request, so the - callers below stay bit-identical whenever the runtime draft length is - the static maximum; ragged batches get the scheduler's per-request - windows. - - Host-side on purpose: the expanded token count derived from these feeds - buffer slicing, and reducing it off a device tensor would be a - device->host sync on every step of the *uniform* path too. + Uniform batches get `gen_token_stride` per request; ragged batches get + the scheduler's windows. Host-side on purpose: the total feeds buffer + slicing, and reducing it off a device tensor would sync every step. """ if self.ragged_verify_lens is None: return [self.gen_token_stride] * self.num_generations @@ -727,22 +713,10 @@ def expand_per_gen_token(self, dim: int = 0) -> Tuple[torch.Tensor, int]: """Repeat `values` once per query token, returning `(expanded, num_tokens)`. - Uniform batches take the scalar `repeat_interleave` -- no repeats - tensor, no host->device copy, and no output-size question at all, so - this is bit-identical and cost-identical to the fixed-stride expansion - it replaced. - - Ragged batches need a per-request repeats vector, which costs two things - the scalar form does not, and both are paid for here: - - * `repeat_interleave` with a tensor `repeats` reads their cumulative sum - back to the host to size its output unless `output_size` is given, so - the total is computed on the host and passed explicitly; - * materializing `repeats` on the device is a host->device copy, and from - *pageable* memory CUDA synchronizes the stream before starting it. The - staging tensor is therefore pinned and the copy issued - `non_blocking=True`, matching how the rest of the engine stages - host-built index vectors. + Uniform batches take the scalar `repeat_interleave` (no repeats tensor, + no copy). Ragged batches must pass `output_size` (a tensor-repeats + interleave otherwise reads the cumsum back to the host) and stage the + repeats through pinned memory so the H2D copy is non-blocking. """ repeats = self.gen_token_repeat_list() num_tokens = sum(repeats) @@ -1039,10 +1013,8 @@ def on_update_kv_lens(self): self.scheduler_metadata_buffer_full_next_n.copy_( scheduler_metadata_buffer_full_next_n, non_blocking=True) if self.use_expanded_buffers_for_mtp: - # One row per query token; the repeat count is per-request - # under ragged verification and a constant otherwise. Both the - # host-side total and the sync-free expansion live in - # ``expand_per_gen_token``. + # One row per query token; per-request repeats under ragged + # verification (see ``expand_per_gen_token``). expanded, num_tokens = self.expand_per_gen_token( gen_indexer_kv_lens) self.kv_lens_expanded_cuda[:num_tokens] = expanded @@ -1382,14 +1354,12 @@ def _draft_sizing_cap(self) -> int: """Static draft-token ceiling for buffer SIZING (not per-step logic). Every expanded buffer must be shaped from one process-wide ceiling, - never from the per-step ``max_draft_tokens``: the pooled backing - memory is shared between the base metadata and every CUDA-graph - clone, and two instances viewing the same bytes with different row - counts interleave their rows differently -- a prepare() on one - layout silently corrupts every captured graph reading the other - (the DSpark ragged IMA). ``update_spec_dec_param`` records the - config-wide ``max_total_draft_tokens`` here; before the first call - the mutable value is all there is. + never the per-step ``max_draft_tokens``: the pooled backing memory is + shared between the base metadata and every CUDA-graph clone, and two + instances viewing the same bytes with different row counts corrupt + each other. ``update_spec_dec_param`` records the config-wide + ``max_total_draft_tokens`` here; before the first call the mutable + value is all there is. """ cap = getattr(self, "_draft_alloc_cap", None) if cap is None: @@ -1483,34 +1453,23 @@ def create_expanded_buffers(self, capture_graph=False): self._ragged_num_rows = 0 # ------------------------------------------------------------------ - # Token-major presentation of the generation half, for the MLA rope and - # sparse-MLA generation ops (goal doc K1 / K2). - # - # Both of those ops learn their sequence count from a tensor *shape* -- - # `num_seqs = host_context_lengths.size(0)` in dsv3RopeOp.cpp and - # thop/attentionOp.cpp -- and then require the query tokens to divide - # evenly across it. A ragged batch does not divide, so the assert - # "MLA can only support input sequences with the same sequence length" - # fires. Handing those two ops one row per generation *token* makes the - # division trivial (seq_len == 1) and satisfies the assert rather than - # weakening it; no kernel changes. + # Token-major presentation of the generation half, for the MLA rope + # and sparse-MLA generation ops. Both derive their sequence count from + # `host_context_lengths.size(0)` and require the query tokens to + # divide evenly across it; a ragged batch does not, so hand them one + # row per generation token (seq_len == 1). # - # These are a PARALLEL set of views, never the metadata's own fields. - # `num_seqs` is read as a request count in roughly twenty places -- - # spec workers sizing acceptance rectangles, the KV-length correction in - # _preprocess_inputs, the indexer's own metadata -- so redefining it - # would be a global re-typing rather than a presentation change. - # - # Rows are laid out prefix-preserving: the context rows first, unchanged, - # then one row per generation token. That keeps `num_contexts`, - # `num_ctx_tokens` and the ops' `seq_offset = num_contexts` all correct - # without touching them. + # These are a PARALLEL set of views, never the metadata's own fields: + # `num_seqs` is read as a request count throughout and must not be + # redefined. Rows are prefix-preserving -- context rows first, + # unchanged, then one row per generation token -- keeping + # `num_contexts`, `num_ctx_tokens` and the ops' `seq_offset` correct. # # Capacity is the static ceiling and must never be recomputed from a # runtime count: these live in the process-wide buffer pool keyed by - # `cache_name`, so a later reallocation would move addresses already - # baked into captured graphs. `2 + max_draft_tokens` per sequence covers - # one context row plus the widest generation window. + # `cache_name`, so reallocation would move addresses baked into + # captured graphs. `2 + max_draft_tokens` per sequence covers one + # context row plus the widest generation window. attn_rows_cap = self.max_num_sequences * (2 + self._draft_sizing_cap) self._attn_num_rows = 0 self.attn_row_kv_lens_cuda = self.get_empty( @@ -1546,14 +1505,11 @@ def create_expanded_buffers(self, capture_graph=False): pin_memory=prefer_pinned()) self.attn_row_req_idx_host = torch.zeros( (attn_rows_cap, ), dtype=torch.long, pin_memory=prefer_pinned()) - # host_context_lengths: the shape oracle both ops key off. - # The op indexes request_types over its whole num_seqs -- which under - # this presentation is the ROW count -- in - # `for (idx = num_contexts; idx < num_seqs; idx++) - # TLLM_CHECK(request_types[idx] == kGENERATION)` - # (thop/attentionOp.cpp). Handing it the batch-major array would read - # past its end. Generation is 1, context 0; prefill the whole capacity - # as generation so a step only has to stamp its context prefix. + # host_context_lengths: the shape oracle both ops key off. The op + # reads request_types over its whole num_seqs -- the ROW count under + # this presentation -- so the array must be row-sized. Generation is + # 1, context 0; prefill the whole capacity as generation so a step + # only stamps its context prefix. self.attn_row_request_types_host = torch.ones( (attn_rows_cap, ), dtype=torch.int32, pin_memory=prefer_pinned()) @@ -1571,16 +1527,10 @@ def create_expanded_buffers(self, capture_graph=False): pin_memory=prefer_pinned()) # The main KV block table, one row per query token. Distinct from # `block_table_expanded` below, which carries *indexer* pool indices at - # the indexer's own compression. On the DSv4 sparse generation path the - # only consumer that walks this per row is the KV write inside the rope - # kernel; the attention op overrides its page index pointer with the - # sparse top-k indices, which are already per token. - # Same [num_pools, num_seqs, 2, max_blocks_per_seq] layout the attention - # op reads, expanded on the *sequence* axis only: this tensor is - # substituted for `kv_cache_block_offsets` wholesale while the ops are - # presented token-major, so its rank and axis order have to match what - # they already expect. The pool and K/V axes are carried through - # unchanged -- only the row axis is what token-major changes. + # the indexer's own compression. Same [num_pools, num_seqs, 2, + # max_blocks_per_seq] layout the attention op reads, expanded on the + # sequence axis only: it substitutes for `kv_cache_block_offsets` + # wholesale, so rank and axis order must match. num_attention_op_pools = getattr(self.kv_cache_manager, "num_attention_op_pools", self.kv_cache_manager.num_pools) @@ -1653,9 +1603,9 @@ def update_spec_dec_param( getattr(self, "_draft_alloc_cap", 0) or 0) capture_graph = self.is_cuda_graph # GROW-ONLY ratchet, never an equality check: captured graphs bake the - # buffer addresses in, and the old `!=` re-created the buffers whenever - # max_draft_len moved per step -- the graph-path ragged IMA. Consumers - # slice [:rows]/[:width]. + # buffer addresses in, so re-creating buffers when max_draft_len moves + # per step corrupts every captured graph. Consumers slice + # [:rows]/[:width]. prev_cap = getattr(self, "_expanded_alloc_draft_cap", -1) ratchet = max(self._draft_sizing_cap, prev_cap) if ratchet > prev_cap: @@ -1675,12 +1625,10 @@ def update_spec_dec_param( dtype=torch.float32, capture_graph=capture_graph, ) - # The Radix-split-work scratch (radix_aux_*) is sized the same way - # (num_seqs * (1 + max_draft_tokens) rows) and is allocated - # unconditionally, so it must be resized here too -- otherwise the - # cpp indexer_topk_decode op aborts once MTP raises max_draft_tokens - # ("radix_aux_* must hold at least num_rows*blocks_per_row*index_topk - # elements"). + # radix_aux_* is sized by the same generation-row count and + # allocated unconditionally, so it must be resized here too or the + # cpp indexer_topk_decode op aborts once MTP raises + # max_draft_tokens. self._create_radix_aux_buffers(capture_graph=capture_graph) def _get_pool_block_indices(self) -> torch.Tensor: @@ -1824,41 +1772,27 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): use_dsl = self.sparse_metadata_params.use_cute_dsl_paged_mqa_logits if not self.is_ragged_verify: # A uniform step must not inherit a previous ragged step's row - # count: refresh_ragged_row_kv_lens and ragged_row_kv_lens both key - # off it, and this covers every uniform path, not just the one that - # skips the expanded buffers. + # count: refresh_ragged_row_kv_lens and ragged_row_kv_lens key + # off it. self._ragged_num_rows = 0 if self.is_ragged_verify: - # The DSL paged-MQA path reshapes (num_gen, next_n) -> (num_gen * - # factor, atom) and broadcasts a single 1-D kv_lens per request, so - # it is structurally uniform: there is no next_n that describes a - # ragged batch. The expanded path already materializes one row per - # token, which is exactly the ragged layout, so route there. + # The DSL paged-MQA path is structurally uniform (one broadcast + # kv_len per request, factor * atom == next_n); route ragged + # batches to the expanded one-row-per-token path instead. use_dsl = False - # goal doc K13. The dense XQA spec-dec metadata - # (`spec_decoding_generation_lengths`, packed mask) is written by - # scalar `fill_` at every Python call site, so it can only describe a - # uniform window -- even though the C++ side would accept a - # per-request vector. Today that block is unreachable here twice - # over: `update_spec_dec_param` forces `is_spec_decoding_enabled` - # False on the Blackwell trtllm-gen kernels for a linear tree, and - # MLA generation reads `mlaGeneration` state instead. Both are - # incidental. Assert the property that actually matters so a future - # SM, or a change to that gate, fails here rather than silently - # verifying every request at the batch-wide window. + # The dense XQA spec-dec metadata (spec_decoding_generation_lengths, + # packed mask) is filled with a single scalar window and cannot + # describe per-request verify lengths; assert it stays off so a + # future change fails here rather than silently verifying every + # request at the batch-wide window. assert not self.is_spec_decoding_enabled, ( "ragged verification requires the dense spec-decoding metadata " "to be off: spec_decoding_generation_lengths and the packed " "mask are filled with a single scalar window, which cannot " "describe per-request verify lengths") - # Deliberately keyed on the static ceiling rather than this step's tier. - # Whether DeepGEMM supports a given next_n natively is a per-step - # question, so a shorter tier could sometimes take the cheaper strided - # path -- but that would make the *layout regime* a function of the tier, - # and the expanded path is the general one (one kernel row per query - # token) that is correct at every tier and required for ragged batches - # anyway. Holding the regime fixed keeps the tiers, and the ragged path - # built on top of them, comparable to each other. + # Keyed on the static ceiling, not this step's tier, so the layout + # regime never changes per step; the expanded path is correct at every + # tier and required for ragged batches anyway. self.use_expanded_buffers_for_mtp = ( not use_dsl and (self.is_ragged_verify or @@ -1945,27 +1879,13 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): def _prepare_ragged_row_kv_lens(self, kv_lens: torch.Tensor) -> None: """Populate ``row_kv_lens_cuda``: how far back each query row may attend. - The decode top-k kernel normally reconstructs this from ``next_n``:: - - actual_kv_len = seq_len - next_n + (row % next_n) + 1 - - which is only meaningful while every request contributes the same number - of rows. Under ragged verification neither ``row // next_n`` (which - request) nor ``row % next_n`` (which position in its window) holds, so - the extent is materialized per row instead. - - The generalization is exact: for a request verifying ``v`` positions, - row ``o`` of its window may attend to ``kv_len - v + o + 1`` tokens. - Substituting ``v = next_n`` recovers the formula above, which is why the - uniform path can keep passing ``nullptr`` and stay bit-identical. - - Sync-free by construction: ``kv_lens`` is already a host tensor here - (:meth:`prepare` builds it from ``num_cached_tokens_per_seq``, which - lives on the CPU), so the whole vector is assembled on the host and - crosses once, through pinned memory so the copy does not serialize the - stream. Reading it on the device instead would mean a D2H sync, and - combining the two halves on the device would mean adding a CPU tensor to - a CUDA one. + The uniform kernel reconstructs this from ``next_n``, which is + meaningless on a ragged batch, so materialize it per row: row ``o`` of + a request verifying ``v`` positions attends to ``kv_len - v + o + 1`` + tokens (``v = next_n`` recovers the uniform formula, so that path + keeps passing ``nullptr``). Assembled on the host -- ``kv_lens`` is + already a CPU tensor -- and copied once through pinned memory; + building it on device would sync. """ verify_lens = self.ragged_verify_lens if not verify_lens: @@ -1974,22 +1894,14 @@ def _prepare_ragged_row_kv_lens(self, kv_lens: torch.Tensor) -> None: assert len(gen_kv_lens) == len(verify_lens), ( f"ragged verify lengths {len(verify_lens)} != generation requests " f"{len(gen_kv_lens)}") - # goal doc A4. `seq_lens` is rebuilt from this step's requests, so its - # generation half is by definition current; `ragged_verify_lens` is - # published separately and used to be published *after* prepare(), which - # meant reading the previous step's split -- and under CUDA graphs - # "previous step" really meant "whenever this key last ran". Checking - # the two against each other tests the property (these windows describe - # this step) rather than the mechanism (publish happens before prepare). + # `ragged_verify_lens` is published separately from this step's token + # layout; check they describe the same step. expected_gen_tokens = self.num_tokens - self.num_ctx_tokens assert sum(verify_lens) == expected_gen_tokens, ( f"ragged verify windows sum to {sum(verify_lens)} but this step's " f"token layout has {expected_gen_tokens} generation tokens; the " f"windows belong to a different step than seq_lens does") - # Row o of a request verifying v positions sees kv_len - v + o + 1 - # tokens: the last row sees the request's full KV, each earlier row one - # token less. rows: List[int] = [] for kv_len, verify_len in zip(gen_kv_lens, verify_lens): base = int(kv_len) - int(verify_len) + 1 @@ -2004,10 +1916,8 @@ def _prepare_ragged_row_kv_lens(self, kv_lens: torch.Tensor) -> None: num_tokens = len(rows) self._ragged_num_rows = num_tokens # The write-after-read hazard on these pinned buffers (rewritten each - # ragged step, copied with non_blocking=True under the overlap - # scheduler's run-ahead) is fenced wholesale by the event guard in - # prepare(): it waits for the previous prepare's copies to execute - # before any pinned buffer is touched again. + # ragged step, copied non_blocking under the overlap scheduler) is + # fenced by the event guard in prepare(). self.row_kv_lens_host[:num_tokens].copy_( torch.tensor(rows, dtype=torch.int32)) self.row_kv_lens_cuda[:num_tokens].copy_( @@ -2024,25 +1934,18 @@ def _prepare_ragged_row_kv_lens(self, kv_lens: torch.Tensor) -> None: def refresh_ragged_row_kv_lens(self) -> None: """Rebuild the per-row extent from the corrected ``kv_lens_cuda``. - ``prepare()`` builds the extent from the host-side kv lengths, but under - the overlap scheduler ``_preprocess_inputs`` corrects ``kv_lens_cuda`` - afterwards to account for how many draft tokens were actually accepted. - Leaving the extent at its prepare-time value would make the top-k kernel - and ``gen_kv_lens_cuda`` -- both arguments of the same op call -- disagree - about the same quantity. - - Stays on the device: the correction term was staged at prepare time and - does not move, so this is one gather plus one add, with no sync. + Under the overlap scheduler ``_preprocess_inputs`` corrects + ``kv_lens_cuda`` after prepare(); the extent must follow or the top-k + kernel's arguments disagree. Device-only: one gather plus one add, no + sync. """ num_tokens = self._ragged_num_rows if not self.is_ragged_verify or num_tokens <= 0: return gen_kv_lens = self.kv_lens_cuda[self.num_contexts:self.num_seqs] - # Gather, not repeat_interleave: this runs inside the capture body, and - # a tensor-repeats interleave both allocates and (without output_size) - # reads the cumulative sum back to the host. The row->request map and - # the correction were staged at prepare() time and are pinned by the - # graph key, so all that is left is one gather plus one add. + # Gather, not repeat_interleave: this runs inside the capture body, + # where a tensor-repeats interleave would allocate and sync. The map + # and correction were staged at prepare() time. row_kv_lens = self.row_kv_lens_cuda[:num_tokens] torch.index_select(gen_kv_lens, 0, self.row_req_idx_cuda[:num_tokens], @@ -2052,19 +1955,11 @@ def refresh_ragged_row_kv_lens(self) -> None: def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: """Build the one-row-per-generation-token view for the MLA gen ops. - See the buffer comments in :meth:`create_expanded_buffers` for why this - exists and why it is a parallel view rather than a redefinition of - ``num_seqs``. - - Rows are ``[contexts unchanged] ++ [one per generation token]``. Row - ``o`` of a request verifying ``v`` positions attends to - ``kv_len - v + o + 1`` tokens, the same generalization the decode top-k - uses; substituting ``v = next_n`` recovers the uniform formula, which is - why the uniform path can keep its single row per request. - - Assembled entirely on the host -- ``kv_lens`` is already a CPU tensor - here -- and crossed once through pinned memory. Reading it on the device - would be a D2H sync on a path the overlap scheduler depends on. + See the buffer comments in :meth:`create_expanded_buffers`. Rows are + ``[contexts unchanged] ++ [one per generation token]``; row ``o`` of a + request verifying ``v`` positions attends to ``kv_len - v + o + 1`` + tokens. Assembled on the host and copied once through pinned memory; + building it on device would sync. """ verify_lens = self.ragged_verify_lens ngt = self._ragged_num_rows @@ -2090,14 +1985,9 @@ def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: self.attn_row_prompt_lens_cpu[:nc].copy_( self.prompt_lens_cpu[:nc].to(torch.int32)) - # Reuse what _prepare_ragged_row_kv_lens just built rather than walking - # the windows again: it produced exactly these per-row extents, - # corrections and request indices one call earlier. Rebuilding them - # would double the only O(num_gen_tokens) host work on the step. - # Same WAR guard as _prepare_ragged_row_kv_lens above, same reason: - # attn_row_req_idx_host holds whole-batch request indices, and a - # rewrite racing last step's queued H2D turns into an out-of-bounds - # gather in refresh_token_major_gen_rows. + # Reuse the per-row extents/corrections/request indices + # _prepare_ragged_row_kv_lens just built. The pinned-buffer WAR hazard + # is fenced by the event guard in prepare(). self.attn_row_kv_lens_host[nc:rows].copy_( self.row_kv_lens_host[:ngt]) self.attn_row_kv_correction_host[nc:rows].copy_( @@ -2106,9 +1996,8 @@ def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: # rows are indexed over the whole batch, so shift past the contexts. torch.add(self.row_req_idx_host[:ngt], nc, out=self.attn_row_req_idx_host[nc:rows]) - # host_context_lengths is only a shape oracle for the generation half - # (the ops read it for num_seqs and, for contexts, for the q length); - # one token per row. + # One query token per generation row; contexts keep their real q + # length. self.attn_row_prompt_lens_cpu[nc:rows].fill_(1) # Capacity is prefilled with kGENERATION, so only the context prefix # needs stamping, and nc is small. @@ -2126,53 +2015,18 @@ def _prepare_token_major_gen_rows(self, kv_lens: torch.Tensor) -> None: self.attn_row_prompt_lens_cpu[:rows], non_blocking=True) self._attn_num_rows = rows - # The block table is expanded on the DEVICE, from the copy that is - # already there. Building it host-side and copying would add ~96 KiB of - # H2D per decode step (768 rows x max_blocks_per_seq int32) plus a host - # gather, both on the critical path the overlap scheduler exists to keep - # clear. One index_select costs neither. self.refresh_token_major_block_table() def refresh_token_major_gen_rows(self) -> None: """Rebuild the per-row KV extent from the corrected ``kv_lens_cuda``. - Mirrors :meth:`refresh_ragged_row_kv_lens` and for the same reason: the - overlap scheduler corrects ``kv_lens_cuda`` on device after prepare(), - inside the CUDA-graph capture body, so a prepare-time extent would make - the rope kernel and the attention op disagree about the same quantity. - - One gather plus one add. No allocation, no repeat_interleave, no D2H -- - the row->request map and the correction were staged at prepare() time - and their addresses are pinned by the graph key. + Mirrors :meth:`refresh_ragged_row_kv_lens`: the overlap scheduler + corrects ``kv_lens_cuda`` on device after prepare(), inside the + capture body. One gather plus one add; no allocation, no D2H. """ rows = self._attn_num_rows if not self.is_ragged_verify or rows <= 0: return - # TLLM_DSPARK_ASSERT_ROWS: name the failure instead of taking an - # illegal memory access several hundred frames later. - # - # The docstring above says the row->request map's addresses "are - # pinned by the graph key". That is true on the graph path and false - # on an eager one, where the same metadata object serves ragged and - # non-ragged steps in turn. If `_attn_num_rows` and - # `attn_row_req_idx_cuda` are stale relative to this step's - # `num_seqs` -- 64 requests last step, 3 now -- the gather below reads - # request 63 out of a length-3 tensor, which is exactly the shape of - # the fault seen on the eager-ragged path. - # - # Costs a device-to-host read, so it is opt-in. - if (os.environ.get("TLLM_DSPARK_ASSERT_ROWS") - and not torch.cuda.is_current_stream_capturing()): - # The .max() is a device-to-host read; during capture it would - # raise cudaErrorStreamCaptureInvalidated. The eager path is where - # the interesting case lives anyway. - worst = int(self.attn_row_req_idx_cuda[:rows].max()) - assert worst < self.num_seqs, ( - f"token-major row map is stale: {rows} rows index up to " - f"request {worst} but this step has only {self.num_seqs} " - f"sequences. The map was staged by a previous prepare() and " - f"never restaged; the gather that follows would read out of " - f"bounds.") torch.index_select(self.kv_lens_cuda[:self.num_seqs], 0, self.attn_row_req_idx_cuda[:rows], out=self.attn_row_kv_lens_cuda[:rows]) @@ -2182,23 +2036,16 @@ def refresh_token_major_gen_rows(self) -> None: def refresh_token_major_block_table(self) -> None: """Expand the main KV block table to one row per query token, on device. - Every row of a request sees that request's blocks -- a verify window - never spans a request -- so this is a plain gather by the row->request - map. Done on the device because the source is already resident; doing it - host-side would mean a per-step H2D of the whole table. - - Not refreshed from ``on_update_kv_lens``: block IDs do not move within a - step, the same reason ``block_table_expanded`` is built once per - prepare(). + A verify window never spans a request, so this is a gather by the + row->request map on the already-resident table. Not refreshed from + ``on_update_kv_lens``: block IDs do not move within a step. """ rows = self._attn_num_rows if rows <= 0 or self.kv_cache_block_offsets is None: return src = self.kv_cache_block_offsets - # Gather along the sequence axis (1) and leave the pool and K/V axes - # alone. Collapsing leading axes instead would strip num_seqs itself and - # index a length-2 K/V axis with request ids, which the gather kernel - # catches as an out-of-bounds index rather than silently mis-reading. + # Gather along the sequence axis (1); the pool and K/V axes pass + # through unchanged. pools = min(src.shape[0], self.attn_row_block_offsets.shape[0]) width = min(src.shape[-1], self.attn_row_block_offsets.shape[-1]) torch.index_select( @@ -2208,11 +2055,9 @@ def refresh_token_major_block_table(self) -> None: def token_major_gen_view(self) -> Optional["TokenMajorGenView"]: """The token-major presentation for this step, or None. - None on every uniform batch, which is what keeps those paths - bit-identical: at ``s_q == 1`` trtllm-gen selects different kernels - (`mIsCausalSpecDecodingGen` is gated on ``mMaxSeqLenQ > 1``), so - presenting a uniform batch token-major would be functionally equivalent - but not bitwise so. + Must be None on every uniform batch: at ``s_q == 1`` trtllm-gen + selects different kernels, so presenting a uniform batch token-major + would not stay bit-identical. """ rows = self._attn_num_rows if not self.is_ragged_verify or rows <= 0: @@ -2920,12 +2765,10 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): metadata.kv_lens_cuda_runtime[num_contexts:num_contexts + num_generations]) metadata.gen_indexer_kv_lens_cuda_runtime = gen_seq_lens - # The 2D DeepGEMM metadata API encodes next_n into the schedule, and - # `Indexer.forward` launches the strided path with the *runtime* - # next_n (`num_gen_tokens // num_generations`). Building the - # schedule at the buffer's full static width instead would describe - # more query positions than the kernel is given whenever a shorter - # verify tier is in play. + # The 2D DeepGEMM metadata API encodes next_n into the schedule; + # build it at the *runtime* next_n `Indexer.forward` launches + # with, or a shorter verify tier describes more query positions + # than the kernel is given. next_n_cap = min(metadata.gen_token_stride, metadata.kv_lens_cuda_2d.shape[1]) metadata.kv_lens_cuda_2d[:num_generations, :next_n_cap].copy_( @@ -2943,16 +2786,10 @@ def prepare_scheduler_metadata(metadata: DSAtrtllmAttentionMetadata): scheduler_metadata_buffer_full_next_n, non_blocking=True) else: # Expand schedule metadata buffer (only generation). The DeepGEMM - # API requires 2D; each expanded token becomes a (1,) row. - # - # The row count has to be the one the expansions actually populated, - # which is what `gen_token_repeat_list` reports: the runtime tier on - # a uniform batch, the scheduler's windows on a ragged one. Deriving - # it from `1 + max_draft_tokens` instead described `num_generations * - # (1 + max_draft_tokens)` rows while `Indexer.forward` handed the - # kernel only `num_gen_tokens` of them -- the schedule and the logits - # then disagree about how much work exists, which is a hang rather - # than an error because the kernel waits on tiles nobody produces. + # API requires 2D; each expanded token becomes a (1,) row. The row + # count must be what the expansions actually populated + # (`gen_token_repeat_list`); a mismatch makes the kernel wait on + # tiles nobody produces (hang, not an error). num_tokens = sum(metadata.gen_token_repeat_list()) scheduler_metadata_buffer_expanded = get_paged_mqa_logits_metadata( metadata.kv_lens_expanded_cuda[:num_tokens].view(-1, 1), @@ -3411,10 +3248,8 @@ def sparse_attn_indexer( gen_seq_lens = metadata.seq_lens[num_contexts:num_contexts + num_generations] if not metadata.is_ragged_verify: - # Ragged verification makes unequal decode lengths the whole - # point; the expanded path below handles them natively (one - # kernel row per query token), so this uniformity check only - # applies to the strided paths. + # Uniformity only holds on the strided paths; ragged batches + # take the expanded (one row per token) path. max_decode_len = gen_seq_lens.max().item() min_decode_len = gen_seq_lens.min().item() assert max_decode_len == min_decode_len, "max_decode_len != min_decode_len, we need padding" @@ -3463,9 +3298,8 @@ def sparse_attn_indexer( block_table = metadata.block_table_expanded[:num_tokens] scheduler_metadata_buffer = metadata.scheduler_metadata_buffer_expanded - # Ragged verification packs a different number of query tokens per - # request, so the only invariant left is that the expanded buffers - # cover exactly the tokens handed to the kernel. + # The expanded buffers must cover exactly the tokens handed to the + # kernel. if metadata.is_ragged_verify: assert num_gen_tokens == sum(metadata.ragged_verify_lens), ( f"ragged gen tokens {num_gen_tokens} != " @@ -3482,11 +3316,9 @@ def sparse_attn_indexer( self.layer_idx) indexer_max_seq_len = metadata.get_indexer_max_seq_len() - # The DSL path broadcasts one kv_len per request across all next_n - # positions and requires factor * atom == next_n, neither of which - # a ragged batch can satisfy. prepare_for_spec_decode already - # routed such batches to the expanded buffers; keep the forward - # consistent with that decision. + # A ragged batch cannot satisfy the DSL path's uniform kv_len + # broadcast / factor * atom == next_n; prepare_for_spec_decode + # routed it to the expanded buffers, keep the forward consistent. if self.use_cute_dsl_paged_mqa_logits and not metadata.is_ragged_verify: # DSL kernel design: 1 atom per q (atom = real next_n positions), # kNumNextNAtoms = 1 for any real next_n. The matching schedule @@ -3606,13 +3438,10 @@ def sparse_attn_indexer( metadata.heuristic_scratch_values[ :num_gen_tokens] - # Ragged batches carry a per-row causal extent instead of - # deriving one from next_n. Both CuTe DSL top-k kernels still - # take next_n as a compile-time-ish scalar and reconstruct the - # row->request map from it, and GVR additionally indexes its - # previous-step hint the same way, so neither can express a - # ragged batch. Route to the C++ kernel, which now accepts the - # extent directly. + # The CuTe DSL top-k kernels reconstruct the row->request map + # from a scalar next_n and cannot express a ragged batch; + # route to the C++ kernel, which accepts the per-row extent + # directly. ragged_row_kv_lens = metadata.ragged_row_kv_lens(num_gen_tokens) if ragged_row_kv_lens is not None: torch.ops.trtllm.indexer_topk_decode( @@ -3620,9 +3449,8 @@ def sparse_attn_indexer( gen_kv_lens_cuda, topk_indices_buffer[num_ctx_tokens:num_ctx_tokens + num_gen_tokens, :], - # Dead on this path: the kernel reads the extent out of - # row_kv_lens rather than rebuilding it. Pass 1 so a - # stale next_n cannot quietly come back into use. + # next_n is dead on this path (the kernel reads + # row_kv_lens); pass 1 so a stale value cannot be used. 1, self.index_topk, pre_idx=None, diff --git a/tensorrt_llm/_torch/models/dspark/attention.py b/tensorrt_llm/_torch/models/dspark/attention.py index 7f63a40edfbe..0908ec90d05b 100644 --- a/tensorrt_llm/_torch/models/dspark/attention.py +++ b/tensorrt_llm/_torch/models/dspark/attention.py @@ -15,20 +15,13 @@ # # The DSpark captured-context attention primitives are ported from DeepSeek's # DeepSpec reference ``inference/kernel.py`` (``sparse_attn``) and -# ``inference/model.py`` (``get_dspark_topk_idxs``). The reference computes these -# with a TileLang kernel; this is a functional-first pure-PyTorch port with the -# same math (index-gather + online softmax + a learnable attention sink that -# contributes only to the softmax denominator). +# ``inference/model.py`` (``get_dspark_topk_idxs``); pure-PyTorch, same math. """DSpark draft captured-context attention primitives (hardware-agnostic). -The DSpark draft uses *dense* sliding-window MLA (``compress_ratio == 0``): the -query comes from the block's draft tokens, while the keys/values are gathered -from a small per-request set of positions (a sliding window of the projected -captured context plus the current block's own positions). Two primitives capture -the parts that differ from the standard MLA path: - -* :func:`get_dspark_topk_idxs` — the (window-context + block) position list. -* :func:`dspark_sparse_attn` — index-gathered attention with an attention sink. +Dense sliding-window MLA (``compress_ratio == 0``) for the draft block: queries +come from the block's draft tokens; keys/values are gathered from a rolling +window of projected captured context plus the block's own positions, with an +attention-sink softmax. """ from functools import lru_cache @@ -54,11 +47,8 @@ def precompute_dspark_freqs_cis( rope_theta: float = 10000.0, device: torch.device | str = "cpu", ) -> torch.Tensor: - """Plain (non-YaRN) RoPE complex exponentials for the DSpark draft. - - The dense draft attention (``compress_ratio == 0``) disables YaRN and uses the - base ``rope_theta`` (DeepSpec ``precompute_freqs_cis`` with - ``original_seq_len == 0``). + """Plain (non-YaRN) RoPE table: the dense draft disables YaRN and uses the + base ``rope_theta``. Returns: complex64 tensor ``[seqlen, rope_head_dim // 2]``. @@ -75,13 +65,11 @@ def precompute_dspark_freqs_cis( def apply_dspark_rotary( x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False ) -> torch.Tensor: - """Apply (or, with ``inverse``, de-apply) rotary embeddings, DeepSpec-style. + """Apply (or, with ``inverse``, de-apply) rotary embeddings. - Functional (non-in-place) port of DeepSpec ``apply_rotary_emb``: treats the - last dim as adjacent (re, im) pairs, rotates by ``freqs_cis`` indexed along the - sequence axis, and conjugates for the inverse (de-rotation applied to the - attention output). ``x`` is the rope-dim slice only: ``[b, s, rd]`` (3D) or - ``[b, s, h, rd]`` (4D), with ``freqs_cis`` of shape ``[s, rd // 2]``. + Last dim is adjacent (re, im) pairs; ``inverse`` conjugates. ``x`` is the + rope-dim slice only: ``[b, s, rd]`` or ``[b, s, h, rd]``, with ``freqs_cis`` + of shape ``[s, rd // 2]``. """ orig_dtype = x.dtype xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) @@ -98,13 +86,10 @@ def apply_dspark_rotary( def apply_dspark_rotary_batched( x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False ) -> torch.Tensor: - """Per-row (batched) variant of :func:`apply_dspark_rotary`. - - Identical math, but ``freqs_cis`` carries a leading batch axis so each row of - ``x`` is rotated by its own per-request phases (the generation draft runs each - request at a different absolute ``start_pos``). ``x`` is the rope-dim slice - only: ``[G, s, rd]`` (3D) or ``[G, s, h, rd]`` (4D), with ``freqs_cis`` of shape - ``[G, s, rd // 2]``. + """Per-row variant of :func:`apply_dspark_rotary`: ``freqs_cis`` carries a + leading batch axis (one set of phases per request). ``x`` is the rope-dim + slice only: ``[G, s, rd]`` or ``[G, s, h, rd]``, with ``freqs_cis`` of + shape ``[G, s, rd // 2]``. """ orig_dtype = x.dtype xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) @@ -138,19 +123,10 @@ def get_dspark_topk_idxs( ) -> torch.Tensor: """Per-query attended-position indices for the DSpark draft block. - Mirrors DeepSpec ``get_dspark_topk_idxs``: every one of the ``block_size`` - query positions attends to the same set — the ``min(window_size, start_pos+1)`` - most-recent context positions in the rolling KV window, then the - ``block_size`` positions of the current block (stored at offset - ``window_size`` in the concatenated KV). Note this is *non-causal* within the - block (every position sees every block position), matching the reference. - - Args: - window_size: sliding-window length of the captured-context KV cache. - bsz: batch size. - block_size: number of draft positions per request. - start_pos: absolute decode position (must be > 0); bounds the context. - device: device for the returned index tensor. + Every query position attends to the same set: the + ``min(window_size, start_pos+1)`` most-recent context positions in the + rolling KV window, then the ``block_size`` block positions stored at offset + ``window_size``. *Non-causal* within the block, matching the reference. Returns: int32 tensor ``[bsz, block_size, topk]`` with @@ -166,26 +142,13 @@ def get_dspark_topk_idxs_batched( block_size: int, start_pos: torch.Tensor, ) -> torch.Tensor: - """Sync-free, fixed-size (CUDA-graph-safe) batched ``get_dspark_topk_idxs``. - - Unlike the scalar :func:`get_dspark_topk_idxs` (whose ``topk`` width - ``min(window_size, start_pos+1) + block_size`` depends on the host int - ``start_pos``), this always returns the **fixed** width ``window_size + - block_size`` and masks the unfilled context slots with ``-1``. The masked - slots are excluded by :func:`dspark_sparse_attn` exactly as if they were - absent, so the result is numerically identical to gathering only the - ``min(window_size, start_pos+1)`` valid context positions — but the shape no - longer depends on the data, which is what CUDA-graph capture requires. - - Every query position attends to the same set: context window slots - ``0..window_size-1`` (slot ``c`` valid iff ``c <= start_pos[g]``, i.e. it has - been written) followed by the ``block_size`` block positions at offset - ``window_size`` (always valid). + """Sync-free, fixed-shape (CUDA-graph-safe) batched ``get_dspark_topk_idxs``. - Args: - window_size: sliding-window length of the captured-context KV cache. - block_size: number of draft positions per request. - start_pos: ``[G]`` int tensor of per-request absolute decode positions. + Always returns the fixed width ``window_size + block_size``, masking + unfilled context slots with ``-1`` (excluded by :func:`dspark_sparse_attn`), + so the result matches the scalar path while the shape stays data-independent + for graph capture. Context slot ``c`` is valid iff ``c <= start_pos[g]``; + ``start_pos`` is a ``[G]`` int tensor of absolute decode positions. Returns: int32 tensor ``[G, block_size, window_size + block_size]``. @@ -193,8 +156,6 @@ def get_dspark_topk_idxs_batched( device = start_pos.device g = start_pos.shape[0] ctx_cols = torch.arange(window_size, device=device) # [win] - # Context slot c holds a written key iff c <= start_pos (slots 0..start_pos - # filled; for start_pos >= window_size-1 the whole rolling window is filled). valid = ctx_cols.unsqueeze(0) <= start_pos.unsqueeze(1) # [G, win] ctx_idx = torch.where( valid, ctx_cols.unsqueeze(0).expand(g, -1), torch.full_like(valid, -1, dtype=torch.long) @@ -214,12 +175,10 @@ def dspark_sparse_attn( ) -> torch.Tensor: """Index-gathered multi-query attention with an attention sink. - Functional-first port of the DeepSpec ``sparse_attn`` TileLang kernel. For - each ``(batch, query, head)`` it gathers the ``topk`` KV rows named by - ``topk_idxs`` (an index of ``-1`` masks that slot), computes a scaled - dot-product softmax over them, and adds a per-head learnable *sink* logit that - participates only in the softmax denominator (i.e. an "attend-to-nothing" - option with a zero value vector). KV is shared across query heads (MQA). + Gathers the ``topk`` KV rows named by ``topk_idxs`` (``-1`` masks a slot), + softmaxes over them, and adds a per-head learnable *sink* logit that + participates only in the softmax denominator (an "attend-to-nothing" option + with a zero value vector). KV is shared across query heads (MQA). Args: q: ``[b, m, h, d]`` query (``m`` = block_size, ``h`` = heads). @@ -309,21 +268,15 @@ def dspark_attention_forward( ) -> torch.Tensor: """Captured-context DSpark draft attention (generation path, ``start_pos > 0``). - Functional port of DeepSpec ``DSparkAttention.forward`` for the dense - (``compress_ratio == 0``) draft: low-rank Q (``wq_a`` -> ``q_norm`` -> ``wq_b``) - with a per-head RMS + RoPE, MQA K/V from ``wkv`` (shared across heads), keys - gathered from a rolling captured-context window (``kv_cache``, into which the - projected ``main_x`` context is written at ``start_pos % window_size``) plus the - block's own positions, attention-sink softmax, inverse-RoPE on the output, and a - grouped low-rank O projection (``wo_a`` einsum + ``wo_b``). - - Weights are plain tensors for ``F.linear`` (the caller supplies the loaded / - dequantized projection weights); ``wo_a`` is the raw grouped weight matrix - ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]``. ``kv_cache`` is - ``[b, window_size, head_dim]`` and is updated functionally (cloned). + Port of DeepSpec ``DSparkAttention.forward`` for the dense + (``compress_ratio == 0``) draft. Weights are plain (already dequantized) + tensors for ``F.linear``; ``wo_a`` is the raw grouped weight + ``[n_groups * o_lora_rank, n_heads * head_dim // n_groups]``. ``kv_cache`` + is ``[b, window_size, head_dim]``; the projected ``main_x`` context is + written at ``start_pos % window_size`` (into a clone unless ``persist``). Returns: - ``[b, block_size, dim]`` attention output (residual stream contribution). + ``[b, block_size, dim]`` attention output. """ assert start_pos > 0, "DSpark draft attention runs at generation (start_pos > 0)" b, block, _ = x.shape @@ -347,10 +300,8 @@ def dspark_attention_forward( kv = _rmsnorm(F.linear(x, wkv), kv_norm_w, eps) # [b, block, head_dim] kv = _rope_last_dims(kv, rd, blk_freqs) - # Write the context K/V into the rolling window, then attend over - # [window context | block] with the sink. ``persist=True`` writes through - # to the caller's buffer (cross-step decode, worker-owned window); the - # default clones so single-shot callers (golden / unit tests) stay pure. + # persist=True writes through to the caller's window (cross-step decode); + # the default clones so single-shot callers stay pure. cache = kv_cache if persist else kv_cache.clone() cache[:, start_pos % window_size] = main_kv.squeeze(1) kv_full = torch.cat([cache, kv], dim=1) # [b, window + block, head_dim] @@ -391,35 +342,25 @@ def dspark_attention_forward_batched( freqs_cis: torch.Tensor, persist: bool = False, ) -> torch.Tensor: - """Batched, CUDA-graph-safe captured-context DSpark draft attention. - - Numerically identical, per request, to :func:`dspark_attention_forward`, but - free of host syncs and data-dependent shapes so it can be captured into a CUDA - graph (the one-engine drafter runs inside the target's graph). The differences - from the scalar path are purely mechanical: - - * ``start_pos`` is a ``[G]`` int tensor (one absolute decode position per gen - request) instead of a python int; RoPE phases are *gathered* per request from - the fixed ``freqs_cis`` table rather than sliced. - * the rolling-window context K/V is written/read through the ``slots`` index - into a shared ``kv_cache`` (``persist=True`` writes through to the caller's - worker-owned buffer; otherwise a clone is used), instead of mutating a - per-request cache in place. - * the attended-position list has the fixed width ``window_size + block_size`` - with ``-1`` masking (see :func:`get_dspark_topk_idxs_batched`). + """Batched, CUDA-graph-safe :func:`dspark_attention_forward`. + + Numerically identical per request to the scalar path, but free of host + syncs and data-dependent shapes so it can be captured into a CUDA graph + (the one-engine drafter runs inside the target's graph). Args: x: ``[G, block, dim]`` block layer input (per gen request). main_x: ``[G, 1, hidden]`` projected captured context. start_pos: ``[G]`` int tensor of absolute decode positions (> 0). - kv_cache: ``[N, window_size, head_dim]`` rolling captured-context windows - (``N`` rows indexed by ``slots``; ``N == G`` for single-shot callers). + kv_cache: ``[N, window_size, head_dim]`` rolling captured-context + windows, rows indexed by ``slots`` (``persist=True`` writes + through to the caller's buffer; otherwise a clone is used). slots: ``[G]`` int tensor mapping each request to its ``kv_cache`` row. - freqs_cis: ``[maxlen, rope_head_dim // 2]`` precomputed plain-RoPE table; - must satisfy ``maxlen > start_pos.max() + block_size``. + freqs_cis: ``[maxlen, rope_head_dim // 2]`` plain-RoPE table; must + satisfy ``maxlen > start_pos.max() + block_size``. Returns: - ``[G, block, dim]`` attention output (residual stream contribution). + ``[G, block, dim]`` attention output. """ g, block, _ = x.shape rd = rope_head_dim @@ -442,10 +383,8 @@ def dspark_attention_forward_batched( kv = _rmsnorm(F.linear(x, wkv), kv_norm_w, eps) # [G, block, head_dim] kv = _rope_last_dims_batched(kv, rd, blk_freqs) - # Write the context K/V into the rolling window at slot start_pos%window_size, - # then attend over [window context | block]. ``persist=True`` writes through to - # the worker-owned buffer (cross-step decode); otherwise clone so single-shot - # callers stay pure. Indexed scatter/gather by (slots, slot_pos) is graph-safe. + # persist=True writes through to the worker-owned window; the default clones. + # Indexed scatter/gather by (slots, slot_pos) is graph-safe. write_target = kv_cache if persist else kv_cache.clone() slot_pos = start_pos % window_size # [G] write_target[slots, slot_pos] = main_kv.squeeze(1).to(write_target.dtype) diff --git a/tensorrt_llm/_torch/models/dspark/heads.py b/tensorrt_llm/_torch/models/dspark/heads.py index 62c3d0612143..e2010c0f1436 100644 --- a/tensorrt_llm/_torch/models/dspark/heads.py +++ b/tensorrt_llm/_torch/models/dspark/heads.py @@ -15,23 +15,13 @@ # # The DSpark Markov/RNN/confidence-head math is ported from DeepSeek's DeepSpec # reference implementation (https://github.com/deepseek-ai/DeepSpec, MIT License). -"""DSpark draft-network heads (pure-torch, framework-agnostic). - -These modules implement the *sequential refinement* and *acceptance-confidence* -parts of DeepSeek's DSpark speculative-decoding draft network: - - - Markov head: a low-rank token-bigram logit bias ``logits_k += W2(W1[t_{k-1}])`` - applied autoregressively across the ``block_size`` draft positions (the cheap - "sequential" half of DSpark's "semi-parallel" drafting). RNN variant carries - a GRU-style recurrent state across positions. - - Confidence head: predicts a per-position *conditional* acceptance probability - (``P(accept_k | accept_1..k-1)``); the cumulative product over positions is the - prefix-survival probability the verification scheduler budgets against. It only - decides *how many* tokens are verified, never *whether* one is accepted, so - scheduling stays lossless. See ``_torch/speculative/dspark_schedule.py``. - -This file deliberately depends on ``torch`` only so it can be unit-tested in -isolation (token-for-token) against the DeepSpec reference. +"""DSpark draft-network heads (pure-torch so they unit-test against DeepSpec). + +Markov/RNN heads apply an autoregressive per-position logit bias across the +draft block. The confidence head predicts per-position *conditional* acceptance +probabilities (``P(accept_k | accept_1..k-1)``) whose cumulative product the +verification scheduler budgets against; it decides *how many* tokens are +verified, never *whether* one is accepted, so scheduling stays lossless. """ from typing import Optional @@ -215,16 +205,9 @@ def build_markov_head( class DSparkConfidenceHead(nn.Module): """Per-position acceptance-confidence predictor (DeepSpec AcceptRatePredictor). - Input features are the backbone hidden state, optionally concatenated with the - Markov head's previous-token embedding. Output is a single *raw logit* per - position; callers turn it into a probability with :meth:`apply_sts`, which - folds in the sequential-temperature-scaling (STS) calibration. - - Calibration matters here because the scheduler consumes the *cumulative - product* of per-position probabilities: a per-position bias is amplified - geometrically along the block, so an uncalibrated head systematically - over-estimates prefix survival. ``sts_temperatures`` defaults to all-ones, - which makes :meth:`apply_sts` a plain sigmoid (identity calibration). + Emits one raw logit per position; :meth:`apply_sts` turns it into a + probability, folding in the sequential-temperature-scaling (STS) + calibration (``sts_temperatures`` defaults to all-ones: plain sigmoid). """ def __init__( @@ -239,16 +222,11 @@ def __init__( super().__init__() self.with_markov = bool(with_markov) input_dim = int(hidden_size) + (int(markov_rank) if with_markov else 0) - # The checkpoint stores ``proj`` as a bias-free bf16 weight, but the - # confidence score is computed in fp32 (mirrors the DeepSpec reference - # ``Linear(input_dim, 1, dtype=torch.float32)`` with the fp32 matmul). + # Checkpoint weight is bf16, but the score matmul is fp32 (DeepSpec parity). self.proj = nn.Linear(input_dim, 1, bias=bool(bias), dtype=torch.float32) - # Per-position temperatures, broadcast against ``[batch, block]``. Shape - # ``(block_size,)`` when known, else ``(1,)``; both broadcast correctly. - # ``persistent=False`` keeps it out of ``state_dict`` so it never - # participates in checkpoint loading. It is sized at construction and - # only ever updated in place, so a captured CUDA graph keeps seeing the - # same storage (rebinding the attribute would be invisible to the graph). + # Per-position temperatures, broadcast against ``[batch, block]``. + # Update in place only: a captured CUDA graph holds this storage; + # rebinding the attribute would be invisible to the graph. self.register_buffer( "sts_temperatures", torch.ones(max(int(block_size), 1), dtype=torch.float32), @@ -273,20 +251,10 @@ def forward( def apply_sts(self, confidence_logits: torch.Tensor) -> torch.Tensor: """Raw logits -> calibrated per-position acceptance probabilities in [0, 1]. - Works on host tensors too. The verification planner stages confidence to - pinned CPU memory and calibrates there, so it calls this with CPU logits - while the head (and therefore ``sts_temperatures``) lives on the device; - without the transfer that is a device-mismatch ``RuntimeError`` on the - first step that reaches calibration. The branch is host-side and the - devices match on the in-graph path, so nothing extra is captured. - - The host copy is cached rather than re-fetched. Transferring the table - per call is a device->host copy into pageable memory, which serializes - the stream, and it lands on the planner's path once per decode step -- - and only once a profiled cost table exists, i.e. exactly in the - configuration where the feature is supposed to pay for itself. The table - is written only through :meth:`load_sts_temperatures`, which refreshes - this cache, so it cannot go stale. + Must accept CPU logits: the verification planner calibrates on pinned + host memory while the head lives on the device. The device-mismatch + branch is host-side and devices match on the in-graph path, so nothing + extra is captured. """ temperatures = self.sts_temperatures if temperatures.device != confidence_logits.device: @@ -320,13 +288,8 @@ def load_sts_temperatures(self, temperatures: torch.Tensor) -> None: self._sts_temperatures_host = None def load_weights(self, weights: list) -> None: - """Strict loader that refuses to silently drop a checkpoint bias. - - The generic DeepseekV4 loader copies by iterating ``named_parameters()``, - so a ``proj.bias`` present in the checkpoint but absent from the module - would be dropped without any error -- shifting every confidence score. - Fail loudly instead. - """ + """Strict loader: fail loudly on checkpoint keys the module lacks (a + silently dropped ``proj.bias`` would shift every confidence score).""" (module_weights,) = weights expected = dict(self.named_parameters()) unexpected = [k for k in module_weights if k not in expected] diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index af56bdd14d6c..998c59b9c6d5 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -274,17 +274,10 @@ def get_graph_key( set(draft_len_list)) == 1, "All draft lengths must be the same" key = (batch_size, draft_len, False, short_seq_len_mode, is_all_greedy_sample) - # Ragged verification breaks the link between the row count and the - # token count: the block is still drafted in full (so draft_len is - # unchanged) but the number of positions actually verified varies, - # and that flat token count is baked into the captured graph. Add it - # as a second axis. The element is absent for every uniform batch, - # so those keys stay exactly what they were. - # - # Derived from the batch rather than plumbed in, because this runs - # inside pad_batch: the padding rows are already present and must be - # counted, and the scheduler has already fitted the total onto a - # captured bucket (see ModelEngine.fit_ragged_verify_lens). + # Ragged verification bakes the flat verified-token total into the + # captured graph, so it becomes a second key axis; the element is + # absent for uniform batches, whose keys stay exactly what they + # were. verify_bucket = self._ragged_verify_bucket(batch) if verify_bucket is not None: key = key + (verify_bucket, ) @@ -293,45 +286,34 @@ def get_graph_key( def _ragged_verify_bucket(self, batch: ScheduledRequests) -> Optional[int]: """Total verified token count when the batch uses ragged windows. - None unless *every* generation request carries a window, so a batch the - scheduler declined to make ragged -- or any non-DSpark batch -- keeps - the original key. - - Config-gated before touching the batch: this runs on every graph-key - build for every model, and a per-request scan there is pure overhead for - the ~all of them that will never be ragged. + None unless *every* generation request carries a window; None keeps + the original (uniform) key. Config-gated before touching the batch: + this runs on every graph-key build for every model. """ # Keyed on the MODE: `cap-accept` shares this config flag but never # shrinks the token axis, so it must keep the ordinary uniform key. from ..speculative.dspark_observability import trims_submitted_tokens if not trims_submitted_tokens(self.spec_config): return None - # Enforce the docstring: None unless EVERY generation request carries a - # window. Windows can vanish between the fit and key derivation (planner - # or peer declined, roster changed); a ragged key over a window-less - # batch replays stale token-major row maps -- the ragged IMA. + # Windows can vanish between the fit and key derivation; a ragged key + # over a window-less batch replays stale token-major row maps (IMA). if any( getattr(request, "py_verify_len", None) is None for request in batch.generation_requests): return None - # The bucket the fit agreed on, not a fresh sum over the batch. Summing - # here walks generation_requests *after* _get_padded_batch has appended - # padding, so it is a second, independent derivation of a quantity every - # attention-DP rank must match exactly -- agreement was incidental. The - # fitted value is computed from allgathered peer stats by rules every - # rank runs identically, so it agrees by construction. None means this - # step is not ragged, which keeps the key exactly what it was. + # Return the fitted bucket, never a fresh sum over the batch: every + # attention-DP rank must key on the same value, and the fit's + # allgathered arithmetic agrees by construction where a local sum only + # agrees by accident. return self.agreed_ragged_bucket @staticmethod def _local_draft_len(batch: ScheduledRequests) -> int: """This rank's draft length, for the attention-DP shape comparison. - Deliberately tolerant where :meth:`get_graph_key` asserts: this runs - before the batch is known to be graphable, and a batch whose requests - disagree has no single draft length to compare. Reporting the maximum - makes such a batch differ from a well-formed peer, which is the - conservative answer -- the gate then declines the graph. + Tolerant where :meth:`get_graph_key` asserts: a malformed batch + reports its max, differs from well-formed peers, and the gate + declines the graph. """ return max( (len(request.py_draft_tokens) @@ -385,13 +367,8 @@ def maybe_get_cuda_graph( - The spec_metadata for the graph, if applicable. - The key for the graph, if applicable. """ - # Why the last call declined, for the DSpark ragged counters. A miss is - # counted at the call site, which cannot see which of the several - # independent reasons applied -- and they mean opposite things: a peer - # shape mismatch is one rank declining to go ragged and dragging the - # whole world eager, while an uncaptured key is the token bucket grid - # being wrong. Attributing one to the other sends the fix to the wrong - # place. + # Why the last call declined, for the DSpark ragged counters; the call + # site that counts the miss cannot see the reason. self.last_miss_reason = None self.last_miss_key = None @@ -403,16 +380,11 @@ def maybe_get_cuda_graph( can_run_cuda_graph = batch.can_run_cuda_graph batch_size = batch.batch_size if self.enabled and self.config.enable_attention_dp and self.config.mapping.tp_size > 1: - # goal doc §4.4. Batch size alone is not enough once the token - # count stops tracking the row count, and replay freezes - # all_rank_num_tokens, so a mismatch is invisible at replay time: - # compare every key component a rank can decide for itself. - # - # `draft_len` is agreed by the DSpark planner's single collective, - # but it is also settable rank-locally by the acceptance-rate - # speculation gate, so it is compared here too rather than assumed. - # -1 stands in for "uniform batch" / "not applicable" so the payload - # is plain ints on every rank and compares elementwise. + # Compare every key component a rank can decide for itself: replay + # freezes all_rank_num_tokens, so a shape mismatch is invisible at + # replay time. draft_len is settable rank-locally (acceptance-rate + # gate), so it is compared rather than assumed. -1 = "not + # applicable" so the payload is plain ints on every rank. local_bucket = self._ragged_verify_bucket( batch) if can_run_cuda_graph else None all_can_graph_batch = self.config.dist.tp_allgather([ @@ -427,10 +399,6 @@ def maybe_get_cuda_graph( for peer in all_can_graph_batch) if not is_all_gen_only or not all_shapes_equal: - # Deliberately distinguished: "some rank had contexts" is - # ordinary continuous batching, while "shapes differ" means a - # rank chose a different ragged bucket and is the one worth - # chasing. self.last_miss_reason = ("peer_not_gen_only" if not is_all_gen_only else "peer_shape_mismatch") @@ -459,8 +427,6 @@ def maybe_get_cuda_graph( # Graph doesn't exist yet. If on-the-fly capture is not allowed, # fall back to eager so the caller doesn't need a separate check. if not self._capture_allowed: - # THE reason that means the bucket grid is wrong: the batch was - # graphable and every rank agreed, but this key was never captured. self.last_miss_reason = "key_not_captured" self.last_miss_key = key return None, None, None @@ -523,10 +489,8 @@ def capture(self, # We're forced to do this because we cannot reallocate inputs over many graph runs. max_draft_len = key[1] token_per_request = max_draft_len + 1 - # Ragged keys carry the flat token total as a sixth element (see - # _get_graph_key); that is the replay width, NOT bs * (draft_len + 1): - # the block is drafted in full while the verified total shrinks (the - # two only coincide at the widest bucket, which is what warmup sees). + # Ragged keys carry the flat token total as a sixth element; that is + # the replay width, NOT bs * (draft_len + 1). num_tokens_for_capture = (key[5] if len(key) > 5 else batch_size * self.max_beam_width * token_per_request) @@ -637,15 +601,11 @@ def replay(self, key: KeyType, def will_pad_to(self, padded_bs: int, num_requests: int) -> bool: """Whether a batch of ``num_requests`` can actually reach ``padded_bs``. - The ragged fit derives its token-bucket grid from an assumed padded row - count and hands the rounding slack to pad rows; if padding declines, - the fitted total lands in no captured bucket and every rank runs eager. - - Mirrors the size guards in :meth:`_get_padded_batch`. It deliberately - does not try to predict the resource-dependent bails (KV capacity for - the dummy request, encoder-decoder state): those are checked when the - dummy is created, and being wrong in the optimistic direction there - leaves the pre-existing behaviour rather than making it worse. + Mirrors the SIZE guards in :meth:`_get_padded_batch` only; the + resource-dependent bails (dummy KV capacity, encoder-decoder state) + are deliberately not predicted. The ragged fit sizes its bucket grid + from the padded row count, so a wrong yes strands the fitted total in + no captured bucket. """ if padded_bs <= 0 or num_requests <= 0: return False @@ -756,27 +716,21 @@ def _get_padded_batch(self, batch: ScheduledRequests, self.padding_dummy_requests[runtime_draft_len] = dummy_request padding_dummy_request = self.padding_dummy_requests[runtime_draft_len] - # Under ragged verification the padding rows are part of the token - # budget the scheduler fitted onto a captured bucket. The scheduler - # decided their (uniform) window when it chose the bucket -- it has to, - # because these rows are one shared object and cannot carry different - # values -- so read it back rather than assume. Cleared otherwise: the - # dummy is cached across steps and would carry a stale window into a - # uniform batch. + # The pad rows' (uniform) window was decided by the fit when it chose + # the bucket (they are one shared object), so read it back. Cleared + # otherwise: the dummy is cached across steps and would carry a stale + # window into a uniform batch. if getattr(self.spec_config, "enable_ragged_verify", False): ragged = any( getattr(request, "py_verify_len", None) is not None for request in batch.generation_requests) padding_dummy_request.py_verify_len = (self.ragged_pad_verify_len if ragged else None) - # cap-accept needs the same treatment for the same reason, and - # fails differently without it: `_attach_accept_caps` drops the - # whole batch's caps if any request lacks one, so a single - # capless dummy would degrade the mode to `static` -- silently, - # and only on the steps that happen to need padding. The dummy's - # own acceptance is discarded, so give it the full block, i.e. no - # cap at all. Cleared otherwise: the object is cached across steps - # and would carry a stale cap into a static batch. + # Same for cap-accept: `_attach_accept_caps` drops the whole + # batch's caps if any request lacks one, so a capless dummy would + # silently degrade the mode to `static`. Its acceptance is + # discarded, so give it the full block; cleared otherwise against + # staleness. cap_accept = any( getattr(request, "py_verify_cap", None) is not None for request in batch.generation_requests) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 9a3ded9b5553..2bc99b84f031 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -259,11 +259,9 @@ def _filter_cuda_graph_seq_lens(cuda_graph_seq_lens: list[int], _DEEP_GEMM_PDL_CONFIGURED = False -# goal doc A3/A4. Ragged-verification layout bugs are silent by construction -- -# the token totals stay integer multiples of the row count, so every -# divisibility check inside the kernels still passes and rows are merely -# attributed to the wrong requests. These assertions make that class of bug -# loud, at the cost of walking the batch on the hot path, so they are opt-in. +# Ragged layout bugs are silent (rows get attributed to the wrong requests +# while every kernel-side divisibility check still passes); these assertions +# make them loud at the cost of walking the batch, so they are opt-in. _DSPARK_ASSERT_LAYOUT = os.environ.get("TLLM_DSPARK_ASSERT_LAYOUT", "0") in ("1", "true", "True") @@ -610,13 +608,8 @@ def __init__( self.is_warmup = False self.previous_request_ids = [] # Per-request verify windows seen by the last full _prepare_tp_inputs - # pass. All None without ragged verification, so the incremental-update - # gate below never trips for uniform speculation. + # pass; all None without ragged verification. self.previous_verify_lens = [] - # NOTE: the captured token bucket and the padding rows' window live on - # the CUDA graph runner (``ragged_pad_verify_len``), which is what reads - # them; the bucket itself is re-derived from the batch in - # ``CUDAGraphRunner._ragged_verify_bucket`` so it cannot go stale. self.has_previous_device_draft = False self.previous_accepted_tokens_cuda = torch.empty((self.batch_size, ), dtype=torch.int, @@ -715,22 +708,17 @@ def __init__( self.previous_batch_indices_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') - # Ragged verification windows. These are read *inside* the captured - # graph (one-engine acceptance runs in the target's graph), so they must - # live at a stable address and be written in place -- rebinding - # spec_metadata.verify_lens to a freshly allocated tensor each step - # would leave the graph replaying against the first step's buffer. + # Ragged verification windows, read *inside* the captured graph: they + # must live at a stable address and be written in place. self.ragged_verify_lens_cuda = torch.empty((self.batch_size + 1, ), dtype=torch.int, device='cuda') self.ragged_qo_indptr_cuda = torch.empty((self.batch_size + 2, ), dtype=torch.int, device='cuda') - # cap-accept: per-request acceptance ceilings for an UNTRIMMED batch, - # and the per-request positions those ceilings discarded. Same - # stable-address requirement as the two above -- both are read and - # written inside the captured graph. The losses ride back to the host - # on the sampler's existing copy, so nothing here ever syncs. + # cap-accept: per-request acceptance ceilings for an untrimmed batch + # and the positions they discarded. Same stable-address requirement as + # above; nothing here ever syncs. self.accept_caps_cuda = torch.empty((self.batch_size + 1, ), dtype=torch.int, device='cuda') @@ -1839,38 +1827,25 @@ def _get_graphs_to_capture( f"Capturing {len(graphs)} graphs: {graphs}") return graphs - # Case 2b: DSpark confidence-scheduled verification. - # The drafted block is always full -- scheduling trims *verification*, - # not drafting -- so draft_len is pinned to the top tier in every mode. - # What varies is which axis the tiers land on: `compact` moves them to - # the token count (a cross product with the batch sizes), while - # `cap-accept` and the static fallback keep the ordinary one graph per - # batch size. - # - # Captured graphs are the expensive part of this feature: each costs - # ~10-23 MB of metadata outside the shared pool, taken directly out of - # KV cache (the KV pool is sized from what is left after capture). So a - # tier that no mode can select is not free, it is KV cache spent on a - # graph that will never be replayed. + # Case 2b: DSpark confidence-scheduled verification. The drafted block + # is always full, so draft_len is pinned to the top tier in every + # mode; `compact` moves the tiers to the token axis, `cap-accept` and + # the static fallback keep one graph per batch size. Captured graphs + # cost KV cache, so capture no tier a mode cannot select. if (self.spec_config is not None and getattr( self.spec_config, "enable_confidence_scheduling", False)): tiers = sorted({ int(t) for t in self.spec_config.verify_len_tiers if int(t) >= 0 }) - # The MODE, not the config flag. `cap-accept` is configured with - # enable_ragged_verify=True as well, but submits the FULL block, so - # capturing the ragged token ladder for it builds synthetic batches - # at token totals its real steps never produce -- which fails - # during capture, inside MoE, long before any of it is explicable. + # The MODE, not the config flag: `cap-accept` also sets + # enable_ragged_verify=True but submits the FULL block, so it must + # not capture the ragged token ladder. from ..speculative.dspark_observability import \ trims_submitted_tokens if trims_submitted_tokens(self.spec_config): - # Ragged always drafts the full block -- only verification is - # trimmed -- so draft_len is pinned to the top tier and the - # ladder moves to the *token* axis: one bucket per tier, i.e. - # exactly the token totals the uniform ladder would produce. - # Same graph count, different second axis. + # One token bucket per tier -- exactly the totals the uniform + # ladder would produce, so the graph count does not grow. top_tier = tiers[-1] graphs = [(bs, top_tier, bs * (t + 1)) for bs in cuda_graph_batch_sizes for t in tiers] @@ -1881,9 +1856,7 @@ def _get_graphs_to_capture( f"{top_tier}).") return graphs # One draft length, not the ladder: nothing varies draft_len at - # runtime any more (cap-accept pins runtime_draft_len to - # planner.max_tier; compact takes the ragged branch above), so - # lower tiers would be captured graphs paid for out of KV cache. + # runtime, so lower tiers would be captured graphs never replayed. top_tier = tiers[-1] graphs = [(bs, top_tier) for bs in cuda_graph_batch_sizes] logger.info( @@ -2048,9 +2021,6 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: # in this pass uses the non-greedy key; populate's override # below will keep it False on every subsequent iteration. spec_metadata.is_all_greedy_sample = False - # Count graphs and report private-pool memory (in the finally - # block below): ragged keys multiply the capture set, and OOM - # messages report private-pool bytes without attributing them. _graph_mem_before = (torch.cuda.memory_reserved(), torch.cuda.memory_allocated()) _graph_count = 0 @@ -2101,16 +2071,13 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: finally: if force_non_greedy and spec_metadata is not None: spec_metadata._force_non_greedy_for_capture = False - # The warmup batches published a bucket per captured shape; the - # first real step must not inherit the last one. `fit_ragged_ - # verify_lens` clears it on entry, but a step that never reaches - # the fit -- the planner declining, or a non-ragged mode -- would - # otherwise key on a bucket from capture. + # Warmup published a bucket per captured shape; the first real + # step must not inherit the last one (a step that never reaches + # the fit would otherwise key on a capture-time bucket). self.cuda_graph_runner.agreed_ragged_bucket = None - # Absolute private-pool size, not a delta: capture passes share - # the allocator, so per-pass deltas mislead. CUDA graphs live in - # private pools (segment_pool_id != (0, 0)) -- the same quantity - # PyTorch's OOM message reports -- so sum those directly. + # Absolute private-pool size, not a per-pass delta (passes + # share the allocator); CUDA graphs live in private pools + # (segment_pool_id != (0, 0)). _pool_bytes = 0 _pool_ids = set() for _seg in torch.cuda.memory_snapshot(): @@ -3116,14 +3083,10 @@ def _set_spec_metadata_all_rank_num_tokens( # sequence contributes one token per iteration). spec_metadata.all_rank_num_tokens = spec_all_rank_num_tokens spec_metadata.all_rank_num_seqs = all_rank_num_seqs - # DSpark can draft only after the target processes the current bonus token, - # because it consumes captured target-layer hidden states for that token. - # Prefill computes hidden states for prompt tokens; the first generated token - # is sampled from the last prompt logits and has not itself passed through the - # target layers. Thus context requests seed the rolling window but do not run - # the draft. On mixed steps, num_seqs therefore over-counts the draft MoE - # workload; gen-only per-rank counts keep the FUSED_COMM (DeepGEMM MegaMoE) - # chunk loop identical across EP ranks. + # DSpark drafts only for generation requests (it needs the bonus + # token's target hidden states), so on mixed steps num_seqs + # over-counts the draft MoE workload; gen-only per-rank counts keep + # the FUSED_COMM chunk loop identical across EP ranks. if all_rank_num_gens is not None: spec_metadata.all_rank_num_gens = all_rank_num_gens if (spec_metadata.spec_dec_mode.is_mtp_eagle_one_model() @@ -3407,11 +3370,9 @@ def _can_use_incremental_update( if self.previous_request_ids != request_ids: return False - # This path reuses the attention/spec metadata built by the last full - # prepare, which is only valid while every request's sequence length is - # unchanged. Uniform speculation guarantees that; ragged verification - # re-picks each window every step, so fall back to a full prepare - # whenever the windows moved. + # The incremental path is only valid while every request's sequence + # length is unchanged; ragged verification re-picks windows every + # step, so fall back to a full prepare whenever they moved. verify_lens = [ getattr(request, "py_verify_len", None) for request in scheduled_requests.generation_requests @@ -3661,12 +3622,10 @@ def _set_warmup_ragged_windows(self, batch, verify_bucket: int, draft_len: int) -> None: """Give a capture batch per-request windows summing to ``verify_bucket``. - A captured ragged graph is keyed on its token total, so the warmup batch - has to hit that total exactly -- otherwise the graph is captured under a - key no runtime batch will ever produce, and every ragged step silently - falls back to eager. Any split with the right sum captures the right - shape (the raggedness lives in the *contents* of qo_indptr, not the - shape), so reuse the same spreading rule the runtime fit uses. + The warmup batch must hit the bucket total exactly, or the graph is + keyed under a total no runtime batch produces and every ragged step + silently falls back to eager. Any split with the right sum captures + the right shape (the raggedness lives in the contents of qo_indptr). """ requests = batch.generation_requests if not requests: @@ -3681,34 +3640,26 @@ def _set_warmup_ragged_windows(self, batch, verify_bucket: int, filled = layout.fill_bucket(max_verify_len=1 + int(draft_len)) for request, tokens in zip(requests, filled.verify_lens.tolist()): request.py_verify_len = int(tokens) - 1 - # Publish it to the runner as well: the graph key carries this value - # and the layout assert cross-checks it against the batch walk; the - # runtime fit never runs during capture, so nothing else sets it here. + # Publish to the runner too: the graph key carries this value, and the + # runtime fit never runs during capture. self.cuda_graph_runner.agreed_ragged_bucket = int(verify_bucket) def ragged_verify_token_buckets(self, padded_bs: int) -> List[int]: """Captured token totals for a ragged batch of ``padded_bs`` rows. - One bucket per verify-length tier, i.e. exactly the token counts the - uniform ladder would have produced at this batch size. A ragged batch - rounds its own total up to one of these, so the captured graph set does - not grow beyond the uniform one on the token axis. + One bucket per verify-length tier -- exactly the totals the uniform + ladder would produce, so the captured graph set does not grow. """ if self.spec_config is None: return [] tiers = sorted({ int(t) - for t in getattr(self.spec_config, "verify_len_tiers", None) or [] + for t in self.spec_config.verify_len_tiers }) return [int(padded_bs) * (t + 1) for t in tiers] def _record_dspark_graph_use(self, *, replayed: bool, shape=None) -> None: - """Count graph replay vs eager for the DSpark ragged stats, if enabled. - - Deliberately tolerant: this runs on the hot path for every model, and a - non-DSpark engine has no worker, no planner and no stats object. Guard - rather than branch on config so the cost is one attribute lookup. - """ + """Count graph replay vs eager for the DSpark ragged stats, if enabled.""" stats = self._dspark_ragged_stats if stats is not None: stats.record_graph(replayed=replayed, shape=shape) @@ -3726,49 +3677,30 @@ def fit_ragged_verify_lens( peer_stats: Optional[List[List[int]]] = None) -> Optional[int]: """Round this batch's verify total onto a captured bucket, in place. - The graph key gains a token axis under ragged verification: two batches - with the same row count and the same maximum draft length can still - carry very different token totals, and the flat token count is baked - into the captured graph. So the total has to land on a captured value - *before* the key is built, and the slack is spent on real verification - rather than thrown away (see ``RaggedVerifyLayout.fill_bucket``). - - Writes the fitted window back onto each request's ``py_verify_len`` and - returns the bucket, or None when no captured shape fits -- in which - case the caller should fall back to uniform scheduling rather than drop - out of graph replay. + The flat token count is a CUDA-graph key axis, so the total must land + on a captured value *before* the key is built; the rounding slack is + spent on real verification (``RaggedVerifyLayout.fill_bucket``). + Writes the fitted window onto each request's ``py_verify_len`` and + returns the bucket, or None when no captured shape fits (the caller + falls back to uniform scheduling). ``peer_stats`` is every attention-DP rank's ``[num_requests, - total_tokens, all_can_graph]``, including this rank's. The first two - are required under ADP because ``_get_padded_batch`` pads to the - *maximum* batch size across ranks: sizing the bucket from the local - count alone makes each rank pick a different shape, which is the - graph-key divergence this whole path exists to avoid. - - The third is the group's answer to "will that maximum be taken at - all", identical on every rank. ``_get_padded_batch`` only takes it when - every rank can graph the step; one rank holding a context request makes - each keep its own row count instead. It is gathered rather than - predicted because the graph runner's own allgather, which decides this, - runs after the fit. + total_tokens, all_can_graph]``, including this rank's: the bucket must + be sized from the cross-rank maximum or ranks pick different shapes, + and ``all_can_graph`` is the group's answer to whether that maximum + will be taken at all. """ runner = self.cuda_graph_runner - # Cleared on entry so that every early return leaves no stale - # value behind: the graph key reads this, and a bucket carried - # over from a previous step would key into the wrong graph. The - # padded row count follows the same rule -- record_step books it - # into padded_bs_hist on every step, and a value surviving from the - # last successful fit would pollute exactly the fallback steps the - # histogram exists to distinguish. + # Cleared on entry so every early return leaves no stale value: a + # bucket carried over from a previous step would key into the wrong + # graph. The padded row count follows the same rule. runner.agreed_ragged_bucket = None self._dspark_last_padded_bs = None if not runner.enabled or not verify_lens: return None if len(verify_lens) != len(generation_requests): - # Partially windowing a batch is worse than not windowing it: the - # token layout is built per request and would go ragged, while - # _attach_ragged_verify_layout sees the missing windows and leaves - # the spec metadata uniform. Refuse and let the caller fall back. + # A partially windowed batch would go ragged in the token layout + # while the spec metadata stayed uniform; refuse instead. logger.debug( f"DSpark ragged: got {len(verify_lens)} verify lengths for " f"{len(generation_requests)} generation requests; falling back " @@ -3776,7 +3708,7 @@ def fit_ragged_verify_lens( return None tiers = sorted({ int(t) - for t in getattr(self.spec_config, "verify_len_tiers", None) or [] + for t in self.spec_config.verify_len_tiers }) if not tiers: return None @@ -3789,33 +3721,17 @@ def fit_ragged_verify_lens( bs_buckets = sorted({int(b) for b in runner.supported_batch_sizes}) if not bs_buckets: return None - # The token bucket grid depends on the row count, which itself depends - # on the widest rank, so resolve the rows first and then the tokens. - # Whether _get_padded_batch will take the cross-rank row maximum at all - # is a group decision, gathered with the peer stats rather than guessed - # here: it only does so when EVERY rank can graph this step. One rank - # holding a context request makes each rank keep its own row count, and - # a fit that assumed the maximum would budget tokens for pad rows that - # never arrive. Third element of a peer stat; absent means single-rank, - # where the local count IS the maximum. + # Resolve rows first, then tokens: the bucket grid depends on the + # widest rank's row count. `all_can_graph` (third peer-stat element; + # absent means single-rank) says whether _get_padded_batch will take + # the cross-rank row maximum at all -- only when EVERY rank can graph. all_can_graph = (bool(peer_stats[0][2]) if peer_stats and len(peer_stats[0]) > 2 else True) if not all_can_graph: - # Some rank holds a context request, so no rank will replay a graph - # this step -- and everything below exists to land the batch on a - # captured one. Two things then go wrong at once if we continue: - # _get_padded_batch declines the cross-rank row maximum AND pads the - # TOTAL batch size (context + generation) while this fit counts - # generation rows only, so its budget goes to rows that never - # appear. Observed as `requests carry 18 generation tokens but the - # CUDA-graph key's token axis is 24` -- 3 generation rows fitted as - # 4 -- which throws inside the forward and, since a rank that throws - # stops issuing collectives, deadlocks every other rank behind it. - # - # An eager ragged step would still save real compute, so this is - # conservative rather than necessary; making it safe means teaching - # the bucket machinery that a step may have no graph at all, which - # is a larger change than the saving justifies on a mixed batch. + # No rank will replay a graph this step, and _get_padded_batch + # pads differently than this fit assumes on such steps, so the + # budget would go to rows that never appear. Decline (conservative: + # an eager ragged step would still save compute). logger.debug( "DSpark ragged: some rank cannot run a CUDA graph this step, " "so no captured bucket applies; falling back to uniform") @@ -3826,28 +3742,11 @@ def fit_ragged_verify_lens( padded_bs = runner._round_up_batch_size(widest_rows) if padded_bs == 0: return None - # The whole fit is built on the assumption that the batch will be padded - # up to `padded_bs`: the bucket grid is derived from it and fill_bucket - # hands the rounding slack to rows that do not exist yet. But padding is - # allowed to decline -- _get_padded_batch bails when the dummy request - # has no KV capacity, when padding would exceed the configured batch - # size, and on a few other resource checks. When it does, the rows never - # appear, the grid was derived for a row count that is not realised, and - # the filled total lands in no captured bucket: observed as - # graph_miss_keys {'(193, 5, False, False, True, 449)': 16} on a run - # where the ladder held 192 and 256 but not 193, costing every rank a - # replay on 16 of 318 steps. - # - # Only reachable with real trimming. Without it the total is always - # padded_bs * (max_verify_len + 1), which is a captured bucket by - # construction, so no amount of uniform testing surfaces this. - # - # Decline to go ragged rather than fit against rows that will not exist. - # Uniform always has a captured graph, so the fallback is cheap and the - # step still runs. - # `will_pad_to` mirrors only the SIZE guards; the group's graph answer - # is `all_can_graph` above, and the resource bails stay unmodeled -- an - # optimistic miss desyncs the layout from the graph key under ADP. + # The fit assumes the batch will actually be padded to `padded_bs`; + # padding may decline, and then the fitted total lands in no captured + # bucket. Decline to go ragged rather than fit against rows that will + # not exist (`will_pad_to` mirrors only the SIZE guards; the group's + # graph answer is `all_can_graph` above). if not runner.will_pad_to(padded_bs, len(token_lens)): logger.debug( f"DSpark ragged: padding to {padded_bs} rows is not available " @@ -3876,17 +3775,11 @@ def fit_ragged_verify_lens( return None padded_bs, bucket = shape.padded_bs, shape.bucket - # The CUDA-graph padding rows are appended later by _get_padded_batch, - # which appends a *single shared dummy object*, so every pad row - # necessarily carries the same window. Their contribution therefore has - # to be decided here, not left to the fill: a total the padding step - # then contradicts would key into a graph that was never captured. - # - # Pad rows take the minimum (one token) so real requests get the slack, - # and only grow when the real rows cannot absorb the bucket on their - # own -- SGLang spreads the leftover across pad rows for the same - # reason, just at per-row granularity it can afford because it builds - # the pad block on device. + # Pad rows are appended later as a *single shared dummy object*, so + # they all carry the same window and their contribution must be + # decided here, not left to the fill. They take the minimum (one + # token) so real requests get the slack, growing only when the real + # rows cannot absorb the bucket. n_real = len(token_lens) n_pad = padded_bs - n_real real_capacity = n_real * max_verify_len @@ -3929,23 +3822,16 @@ def fit_ragged_verify_lens( try: filled = layout.fill_bucket(max_verify_len=max_verify_len) except ValueError as exc: - # The guards are about caller consistency, not impossibility; a - # mismatch means this batch has no captured shape, so degrade to + # A mismatch means this batch has no captured shape; degrade to # uniform rather than run eager. logger.debug(f"DSpark ragged bucket fit failed, falling back to " f"uniform scheduling: {exc}") return None # Published only now, past the last way this fit can fail: a stale - # published bucket on a fallback step is the state family behind the - # ragged IMA, and clearing at entry protects the next call, not the - # remainder of THIS one. The graph key carries `agreed_ragged_bucket` - # verbatim so every attention-DP rank keys on the value they - # allgathered, and the layout assert cross-checks it against the - # batch walk. `_dspark_last_total_tokens` lets a graph miss report the - # shape it submitted; `_dspark_last_padded_bs` is the row count the - # stats need to book a ceiling comparable to `delivered`; - # `ragged_pad_verify_len` is stamped on the shared dummy by + # published bucket on a fallback step is the state behind the ragged + # IMA. `_dspark_last_total_tokens`/`_dspark_last_padded_bs` feed the + # stats; `ragged_pad_verify_len` is stamped on the shared dummy by # _get_padded_batch. self._dspark_last_total_tokens = int(bucket) self._dspark_last_padded_bs = int(padded_bs) @@ -3961,13 +3847,9 @@ def fit_ragged_verify_lens( def _ragged_token_lens(generation_requests) -> Optional[List[int]]: """Each generation request's token window, or None if the batch is uniform. - ``py_verify_len`` counts drafted positions; a request also verifies its - bonus position, so its token window is one wider. - - Returns None unless *every* request carries a window. A partially - windowed batch is worse than an unwindowed one: the flat token layout - would go ragged while the spec metadata stayed uniform, which is silent - token misattribution rather than a crash. + ``py_verify_len`` counts drafted positions; the token window adds the + bonus position. None unless *every* request carries a window -- + partially windowed batches are silent token misattribution. """ verify_lens = [ getattr(request, "py_verify_len", None) @@ -3982,31 +3864,14 @@ def _publish_gen_token_layout(self, attn_metadata, """Hand the attention metadata this step's gen-token layout, before it prepares. - Split out of :meth:`_attach_ragged_verify_layout` because of ordering: - ``attn_metadata.prepare()`` is the *only* consumer of - ``ragged_verify_lens`` -- it drives the DSA expanded-buffer decision, the - per-row causal extents, and DeepSeek-V4's per-request compressor token - counts -- and it runs well before the spec metadata is assembled. - Publishing both together meant ``prepare()`` read the previous step's - split every time, and None on the first ragged step. - - The spec-metadata half stays where it was; it is consumed later and has - no such constraint. - - The uniform stride is published for the same reason and at the same - moment. ``prepare()`` expands kv_lens and the block table once per query - token, and the only value that describes how many tokens a request - actually contributes this step is ``1 + runtime_draft_len``. The - attention metadata's own ``max_draft_tokens`` cannot stand in for it: it - is the static ceiling that sizes those buffers, and for the parallel-draft - modes it is pinned to ``tokens_per_gen_step - 1`` for the engine's whole - lifetime (see the ``update_spec_dec_param`` call in :meth:`forward`), so - it does not move when a shorter tier is chosen. - - Note this has to happen here rather than in ``update_spec_dec_param``: - that runs against the base metadata, while ``prepare()`` runs against the - per-key CUDA-graph copy taken later, and a plain attribute set on the - base is invisible to the copy. + Split out of :meth:`_attach_ragged_verify_layout` for ordering: + ``attn_metadata.prepare()`` is the only consumer of + ``ragged_verify_lens`` and runs well before the spec metadata is + assembled. The uniform stride is published here too, because + ``max_draft_tokens`` is a static buffer-sizing ceiling that does not + move when a shorter tier is chosen. It cannot live in + ``update_spec_dec_param`` either: that runs against the base metadata, + and ``prepare()`` runs against the per-key CUDA-graph copy. """ if attn_metadata is None: return @@ -4019,34 +3884,21 @@ def _publish_gen_token_layout(self, attn_metadata, def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, scheduled_requests) -> None: - """goal doc A3: one identity that catches H0, H5 and H8 at once. - - Four independent descriptions of "how many generation tokens are in this - step" have to agree: + """Assert the four descriptions of this step's generation-token count agree: - sum(1 + py_verify_len) the requests, as the scheduler fitted them - num_tokens - num_ctx_tokens the flat input layout the kernels see - total_verify_tokens what acceptance will slice by - the ragged bucket the CUDA-graph key's token axis + sum(1 + py_verify_len) the requests, as fitted + num_tokens - num_ctx_tokens the flat input layout + total_verify_tokens what acceptance slices by + the ragged bucket the CUDA-graph key's token axis - Each of the bugs this guards against breaks exactly one of these while - leaving the others self-consistent, which is why none of them raise on - their own: the token totals are always an integer multiple of the row - count, so the divisibility checks inside the kernels still pass and the - rows are simply attributed to the wrong requests. - - Off by default -- it walks the batch on the hot path. Enable with - TLLM_DSPARK_ASSERT_LAYOUT=1 for correctness runs. + A mismatch is silent token misattribution, not a crash. Off by + default (walks the batch); enable with TLLM_DSPARK_ASSERT_LAYOUT=1. """ verify_lens = self._ragged_token_lens( scheduled_requests.generation_requests) if verify_lens is None: - # A uniform step must not be carrying a published ragged bucket: - # that is exactly the stale state a failed fit used to leave - # behind, and a later ragged step would key a graph the batch - # does not describe. Checked here because the ragged branch below - # never runs for uniform steps, which made the staleness invisible - # even with this assert enabled. + # A uniform step must not carry a published ragged bucket: a later + # ragged step would key a graph the batch does not describe. agreed_stale = self.cuda_graph_runner.agreed_ragged_bucket assert agreed_stale is None, ( f"uniform step carries a stale published ragged bucket " @@ -4077,22 +3929,17 @@ def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, f"tokens but spec_metadata.total_verify_tokens is {from_spec}; " f"acceptance would slice the target logits at the wrong offsets.") if agreed is not None: - # Two derivations of the token axis: `from_requests` walks this - # batch's windows, `agreed` is the fit's allgathered arithmetic, - # and the graph key carries the latter. If the fit reserved for - # padding rows that were then not appended, the replayed graph is - # a different width than the batch -- and nothing downstream says - # so, because acceptance's output_size was frozen at capture. + # `from_requests` walks this batch's windows; `agreed` is what the + # graph key carries. A mismatch means the replayed graph is a + # different width than the batch, and nothing downstream says so. assert from_requests == int(agreed), ( f"ragged layout mismatch: the batch carries {from_requests} " f"generation tokens but the fit agreed on {agreed}, which is " f"what the CUDA-graph key's token axis says; the replayed " f"graph was captured at a different width.") - # Row axis. Everything above is a token count, so all four agree even - # when the token-major presentation is malformed -- the presentation is - # about how many ROWS the MLA generation ops are shown, which is a - # different quantity and needs its own check. + # Row axis: the token-major presentation can be malformed while every + # token count above agrees, so it needs its own check. view = attn_metadata.token_major_gen_view() if hasattr( attn_metadata, "token_major_gen_view") else None if view is not None: @@ -4111,28 +3958,16 @@ def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, def _pinned_host(self, key: str, values, dtype) -> torch.Tensor: """A persistent pinned staging buffer holding ``values``. - Every async host-to-device copy on the ragged path used to source from - a freshly built ``torch.tensor(..., pin_memory=True)``. Nothing keeps - that temporary alive: the statement ends, Python is free to reclaim the - page, and ``non_blocking=True`` means the copy has only been QUEUED. - PyTorch does not extend the lifetime of an async H2D source -- that is - the caller's job -- so the DMA can read a freed page. - - Intermittent by nature: the fault needs the allocator to reuse the - page before the queued copy drains (CUDA_LAUNCH_BLOCKING=1 hides it). - - Buffers are keyed by name and grown monotonically, so a steady state - allocates nothing. + Async H2D sources must outlive the queued copy -- PyTorch does not + extend the lifetime of an async source, so a temporary tensor can be + reclaimed while the DMA still reads it. Buffers are keyed by name and + grown monotonically. """ values = list(values) cache = self._pinned_host_cache - # WAR guard, the counterpart of the use-after-free this helper fixed: - # a persistent buffer can be REWRITTEN while its previous non_blocking - # H2D is still queued behind an in-flight forward, so the device copies - # this step's values into last step's consumer. Callers record the - # key's event via _pinned_host_record after enqueuing the copy; waiting - # here closes the window, and in the steady state the event completed - # long ago so this costs nothing. + # WAR guard: the buffer must not be rewritten while its previous + # non_blocking H2D is still queued. Callers record the key's event via + # _pinned_host_record after enqueuing; waiting here closes the window. evt = self._pinned_host_events.get(key) if evt is not None: evt.synchronize() @@ -4165,28 +4000,19 @@ def _pinned_host_record(self, *keys: str) -> None: def _attach_accept_caps(self, spec_metadata, generation_requests) -> None: """Publish ``cap-accept`` ceilings, without touching the layout. - ``py_verify_cap`` is the cap-accept twin of ``py_verify_len``: the - scheduler's window for a batch whose token axis was deliberately NOT - trimmed. Keeping them on separate attributes is what makes the mode a - pure addition -- every layout consumer keys off ``py_verify_len``, so - the input packing, the DSA row expansion, the graph key and the KV - rewind all keep seeing an ordinary uniform batch with no changes. - - The cap is in ``num_accepted_tokens`` units (accepted drafts plus the - bonus), matching ``verify_lens``, hence ``1 + py_verify_cap``. + ``py_verify_cap`` is kept separate from ``py_verify_len`` so every + layout consumer (input packing, DSA row expansion, graph key, KV + rewind) keeps seeing an ordinary uniform batch. The cap is in + ``num_accepted_tokens`` units (accepted drafts plus the bonus), hence + ``1 + py_verify_cap``. """ caps = [ getattr(request, "py_verify_cap", None) for request in generation_requests ] if not caps or any(cap is None for cap in caps): - # Partial caps are worse than none: the uncapped requests would - # commit their full block while the rest honoured a window, which - # is neither mode and reproducible by neither. So drop them all -- - # but say so, because the result is a run that looks like a - # successful cap-accept and measured nothing. The CUDA-graph - # padding dummy used to land here on every padded step (it carries - # no cap of its own); it is given one in `pad_batch` now. + # Partial caps are neither mode: drop them all, but warn, because + # the step then silently runs as `static`. if caps and any(cap is not None for cap in caps): if not getattr(self, "_warned_partial_accept_caps", False): self._warned_partial_accept_caps = True @@ -4201,9 +4027,8 @@ def _attach_accept_caps(self, spec_metadata, generation_requests) -> None: return n = len(caps) - # Written through the persistent buffer for the same reason as - # verify_lens: acceptance runs inside the target's captured graph, so a - # freshly allocated tensor here would be invisible to every replay. + # Persistent buffer: acceptance runs inside the captured graph, so a + # fresh tensor here would be invisible to every replay. caps_view = self.accept_caps_cuda[:n] caps_view.copy_(self._pinned_host("accept_caps", [1 + int(cap) for cap in caps], @@ -4217,12 +4042,10 @@ def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, generation_requests) -> None: """Publish this step's per-request verify windows to both metadatas. - Only the DSpark confidence scheduler sets ``py_verify_len``; every other - path leaves the fields None and the uniform strides downstream stay - live. Both metadatas need it and for different reasons: acceptance has - to know each request's window to stop its cumulative product at the - right place, and the DSA indexer has to expand kv_lens/block tables one - row per query token instead of a fixed count per request. + Only the DSpark confidence scheduler sets ``py_verify_len``; every + other path leaves the fields None. Acceptance needs the windows to + slice correctly; the DSA indexer needs them to expand kv_lens/block + tables per query token. """ self._attach_accept_caps(spec_metadata, generation_requests) token_lens = self._ragged_token_lens(generation_requests) @@ -4238,9 +4061,8 @@ def _attach_ragged_verify_layout(self, spec_metadata, attn_metadata, from ..speculative.dspark_ragged import build_qo_indptr n = len(token_lens) - # Write through the persistent buffers rather than binding a fresh - # tensor: a captured graph baked in the address it saw at capture time, - # so a new allocation here would be invisible to every replay. + # Persistent buffers: a captured graph baked in the address it saw at + # capture time, so a fresh tensor would be invisible to every replay. lens_view = self.ragged_verify_lens_cuda[:n] lens_view.copy_(self._pinned_host("ragged_verify_lens", token_lens, torch.int32), @@ -4354,14 +4176,9 @@ def _update_target_input_tensors( dtype=torch.long, pin_memory=prefer_pinned()).to(new_tokens_lens_device.device, non_blocking=True) - # `output_size` is not optional here. With a tensor `repeats`, - # repeat_interleave has to know the output length before it can - # allocate, and without this it gets it by reading the cumulative - # sum back to the host -- a device->host sync on every ragged step, - # inside the overlap scheduler's window, which is precisely what - # overlap exists to avoid. The total is already known on the host: - # the caller derived `total_num_tokens` as - # `sum(tokens_per_extend_request)`. + # `output_size` is not optional: without it repeat_interleave with + # tensor `repeats` reads the cumulative sum back to the host -- a + # device->host sync on every ragged step. previous_pos_indices = torch.repeat_interleave( previous_slots, tokens_per_request_device, @@ -5042,7 +4859,7 @@ def append_cross_attention_state(request: LlmRequest, self.runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 for request in extend_requests: - if getattr(request, "py_needs_onehot_draft_probs", False): + if request.py_needs_onehot_draft_probs: if request.py_seq_slot is not None: padding_gen_slots.append(request.py_seq_slot) request.py_needs_onehot_draft_probs = False # consume once @@ -5057,15 +4874,9 @@ def append_cross_attention_state(request: LlmRequest, if next_draft_tokens_device is None or request.is_dummy or request.py_batch_idx is None: # get token ids, including input token ids and draft token ids. For these dummy requests, # no need to copy the token ids. - # Under ragged verification the block is still drafted in full, - # but only the request's own window is submitted to the target. - # This branch runs whenever the overlap scheduler is off -- i.e. - # on every step of the accuracy configuration -- so taking the - # full drafted length here would leave the flat token layout - # uniform while the spec metadata and the CUDA-graph key both - # believe it is ragged: silent token misattribution rather than - # a crash. Mirrors the overlap branch below, which already - # derives every length from the same per-request value. + # Only the request's own window is submitted to the target; + # taking the full drafted length here would desynchronize the + # flat token layout from the spec metadata and the graph key. num_draft_tokens = get_request_tokens_per_gen_step( request, 1 + get_draft_token_length(request)) - 1 if not (request.is_attention_dp_dummy @@ -5504,13 +5315,9 @@ def previous_seq_slots_device(): if previous_batch_len > 0: previous_slots = previous_seq_slots_device() - # Ragged verification gives each request its own window, so the - # fixed [previous_batch_len, runtime_tokens_per_gen_step] block - # gathers below no longer describe the layout. Detect that here - # and fall back to explicit (row, col) index lists; a uniform - # batch -- every non-DSpark model, and DSpark whenever the - # scheduler picks one window for the whole batch -- keeps the - # cheaper strided path. + # Ragged windows break the fixed strided gathers below; use + # explicit (row, col) index lists, keeping the cheaper strided + # path for uniform batches. is_ragged_gen = any( tokens != runtime_tokens_per_gen_step for tokens in previous_batch_tokens_per_request) @@ -5569,15 +5376,10 @@ def previous_seq_slots_device(): previous_kv_len_offsets = ( new_tokens_lens_device[previous_slots] - tokens_per_previous_request) - # The gather above reads the sampler's slot-indexed store; - # a slot the sampler has not (yet) written for the previous - # step yields a stale or uninitialized count, and the - # captured kv_lens correction then walks the KV append past - # the request's allocated blocks (the ragged IMA). The - # physically possible range is tight -- accepted tokens - # minus this request's window -- so clamp to it; a clamped - # value is at worst one window of KV length, never an OOB - # address. + # A slot the sampler has not yet written yields a stale + # count, and the captured kv_lens correction then walks the + # KV append out of bounds (the ragged IMA); clamp to the + # physically possible range. previous_kv_len_offsets = previous_kv_len_offsets.clamp_( min=-int(self.runtime_draft_len + 1), max=int(self.runtime_draft_len + 1)) @@ -5743,24 +5545,13 @@ def previous_seq_slots_device(): helix_is_inactive_rank=helix_is_inactive_rank, ) - # Ragged verification is the one case where the per-request lengths - # genuinely move between replays of the same graph: the row count and - # the padded token total are both pinned by the key, but how the tokens - # are split across rows is not. Leaving the captured ones(bs)*(1+top) - # in place would make attn_metadata.num_tokens (== seq_lens.sum(), - # interface.py) disagree with the input_ids width on every trimmed step - # -- and that sum is what attention-DP all-gathers and the MoE chunks - # from. The setter already copies in place under CUDA graphs (see its - # comment: "The seqlens can change if we are doing spec decode"), so - # this reallocates nothing as long as the row count is unchanged. - # Publish before the refresh decision below, not after: that decision - # asks whether this is a ragged step, and asking `attn_metadata` before - # this step's windows are on it reads the *previous* step's answer -- - # None on a freshly copied per-key graph metadata. The refresh is then - # skipped and `seq_lens` keeps the uniform `ones(bs) * (1 + - # max_draft_tokens)` that `create_cuda_graph_metadata` seeded it with, - # so the layout claims bs*(top_tier+1) tokens while the windows describe - # the trimmed bucket. + # Under ragged verification the per-request lengths move between + # replays of the same graph (the key pins the row count and token + # total, not the split), so seq_lens must be refreshed or + # attn_metadata.num_tokens disagrees with the input_ids width. Publish + # BEFORE the refresh decision below: it asks the metadata whether this + # step is ragged, and asking before this step's windows are on it + # reads the previous step's answer. self._publish_gen_token_layout(attn_metadata, scheduled_requests.generation_requests) @@ -7189,17 +6980,10 @@ def forward(self, ) can_run_graph = key is not None - # A ragged step that misses its captured shape does not raise -- it - # quietly runs eager, which costs far more than the tokens it - # trimmed. Record it so the miss is a counter rather than something - # only a profiler would reveal. - # - # Only generation-only steps count. A batch carrying context - # requests has no generation graph to miss, so recording it as an - # eager fallback would make `graph_eager` a prefill counter -- and - # `assert_ragged_active` refuses a run with any eager step, so that - # miscount alone makes the check unsatisfiable on every real - # workload. + # A ragged step that misses its captured shape quietly runs eager; + # count it. Only generation-only steps count: a batch with context + # requests has no generation graph to miss, and booking it as + # eager would make `assert_ragged_active` unsatisfiable. if padded_requests.can_run_cuda_graph: self._record_dspark_graph_use( replayed=can_run_graph, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 27bc630bd04f..fe24886b903d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -558,9 +558,6 @@ def __init__( self.peft_cache_config = peft_cache_config self.iter_counter = 0 - # TLLM_DSPARK_SYNC_DEBUG bookkeeping; see _maybe_arm_dspark_sync_debug. - self._dspark_sync_debug_iters = 0 - self._dspark_sync_debug_armed = False # profile config self.profile_start_iters, self.profile_stop_iters = _load_iteration_indexes( PROFILE_START_STOP_ENV_VAR_NAME) @@ -2402,64 +2399,19 @@ def _process_iter_stats( prev_device_step_time_ms=prev_device_step_time_ms, gpu_forward_time_ms=gpu_forward_time_ms) - def _maybe_arm_dspark_sync_debug(self) -> None: - """Turn on CUDA sync detection once decode is in steady state. - - The overlap scheduler's entire value is that the host runs ahead of the - device, so one device->host sync per decode step gives it back. Static - review can show the code *looks* sync-free; only the allocator's own - detector shows that it *is*. - - Armed late on purpose. Warmup, CUDA-graph capture and the first steady - steps legitimately synchronize, so arming at startup would bury the - signal in expected hits. ``TLLM_DSPARK_SYNC_DEBUG=`` arms after ``n`` - DSpark decode iterations (``1`` means "on", which arms after the default - 64). Each subsequent sync emits a Python warning with a stack, which is - what makes the offender identifiable rather than merely countable. - - Diagnostic only: never on by default, and 'warn' rather than 'error' so - a run surfaces every offender instead of dying at the first. - """ - if self._dspark_sync_debug_armed: - return - raw = os.environ.get("TLLM_DSPARK_SYNC_DEBUG", "") - if not raw or raw in ("0", "false", "False"): - return - try: - after = 64 if raw in ("1", "true", "True") else int(raw) - except ValueError: - logger.warning( - f"ignoring TLLM_DSPARK_SYNC_DEBUG={raw!r}: expected an integer " - f"iteration count, or 1 for the default") - self._dspark_sync_debug_armed = True - return - self._dspark_sync_debug_iters += 1 - if self._dspark_sync_debug_iters < after: - return - torch.cuda.set_sync_debug_mode("warn") - self._dspark_sync_debug_armed = True - logger.warning( - f"DSpark: CUDA sync detection armed after " - f"{self._dspark_sync_debug_iters} decode iterations. Any " - f"device->host synchronization from here on emits a warning with a " - f"stack; on the overlap-scheduler path there should be none.") - def _maybe_assert_dspark_ragged_active(self, stats) -> None: """Run the ragged gate mid-loop, once the step floor is reached. - The cleanup-time copy of this check raises into a shutdown path pytest - never sees; raised mid-loop instead, the exception fails the run. Only - once, so the first (and only informative) failure is not buried. + Raises mid-loop (a cleanup-time raise lands in a shutdown path the + test runner never sees); fires at most once. """ - # Lazy rather than set in __init__: this class has several - # construction paths and a gate that silently no-ops because one of - # them missed an attribute is the exact failure being fixed here. + # Deliberately lazy: this class has several construction paths, and a + # missed __init__ attribute would silently disable the gate. if getattr(self, "_dspark_asserted", False): return min_steps = self._dspark_assert_min_steps - # Counted over steps with a real batch, not all steps: trtllm-bench - # ramps concurrency from zero, and a tiny early batch has trivially - # identical windows, which would fail an innocent run. + # Count only multi-request steps: tiny ramp-up batches have trivially + # identical windows and would fail an innocent run. if min_steps is None or stats.steps_multi_request < min_steps: return self._dspark_asserted = True @@ -2479,36 +2431,23 @@ def _dspark_assert_min_steps(self) -> Optional[int]: def _report_dspark_ragged_stats(self) -> None: """Final DSpark scheduling summary, and optionally a hard gate on it. - The counters live in the worker process. Under TP8 + attention-DP the - driver cannot reach them -- ``collective_rpc`` refuses a multi-rank - world -- so a test can assert on a GSM8K score and nothing else. But a - run where the planner declined on every step scores *exactly* the same - as one where it worked, so that assertion is evidence-free: the known - failure mode here is not a wrong answer, it is a feature that quietly - did nothing. - - Two things close that. The summary is logged unconditionally at - shutdown, so the numbers are in the log of every run rather than only in - whatever the 32-step periodic emit happened to catch. And with - ``TLLM_DSPARK_ASSERT_ACTIVE=1`` the worker asserts on its own counters - here -- an exception in the worker does propagate to the driver, which - is the one channel that crosses the MPI boundary reliably. + The counters live in the worker process and are unreachable from the + driver under multi-rank TP, so the summary is always logged at + shutdown, and with ``TLLM_DSPARK_ASSERT_ACTIVE`` the worker asserts on + its own counters (a worker exception does propagate to the driver). """ worker = getattr(self.model_engine, "_get_spec_worker", None) stats = getattr(worker(), "ragged_stats", None) if worker else None - # Drained BEFORE the ragged early-return: STS collection runs in static - # mode too, so gating the final flush on ragged counters drops the tail - # of every static-mode collection run. + # Flush BEFORE the ragged early-return: STS collection also runs in + # static mode. recorder = getattr(self.sampler, "sts_recorder", None) if recorder is not None: recorder.flush() if stats is None or not stats.steps_total: return stats.log_summary(prefix="DSpark ragged verify [final]") - # The value is a MINIMUM STEP COUNT, not a boolean: an engine build - # spins up throwaway executors (KV-cache estimation) that run a couple - # of decode steps, and asserting on one of those reports "the feature - # did nothing" about an executor never asked to do anything. + # A MINIMUM STEP COUNT, not a boolean: throwaway executors (KV-cache + # estimation) run a few decode steps and must not trip the assert. min_steps = self._dspark_assert_min_steps if min_steps is None: return @@ -3318,16 +3257,10 @@ def wait_on_pp_send_handles(self, send_handles, microbatch_id): def set_dspark_verify_len_pin(self, verify_len: Optional[int]) -> Optional[int]: """Queue a DSpark verify-length pin, or ``None`` to clear it. - The pin bypasses the confidence planner and makes every request verify - the same number of positions -- the shape a cost-table sweep needs to - measure a cell it can label. It exists so a sweep can walk the ladder - against a RUNNING server instead of rebuilding the engine per length - (which is what the ``TLLM_DSPARK_FORCE_VERIFY_LEN`` environment - variable requires, and which cannot reach ranks that were spawned - before it was set). - - Validated here and applied on the next decode step, when the step's - existing cross-rank allgather hands the same value to every rank. + The pin bypasses the confidence planner so every request verifies the + same number of positions (for cost-table sweeps against a running + server). Validated here, applied on the next decode step, when the + step's cross-rank allgather hands the same value to every rank. Returns: The validated pin, or None when cleared. @@ -3347,12 +3280,10 @@ def set_dspark_verify_len_pin(self, verify_len: Optional[int]) -> Optional[int]: "scheduling enabled") mode = getattr(worker, "ragged_verify_mode", None) computes = (mode.computes_windows if mode is not None else - getattr(self.model_engine.spec_config, - "enable_ragged_verify", False)) + self.model_engine.spec_config.enable_ragged_verify) if not computes: # Static mode never calls decide_verify_lens, so an accepted pin - # would be adopted, logged as set -- and never influence a step. - # Answering 200 there is the mislabelled-sweep trap all over again. + # would be adopted and silently never applied. raise RuntimeError( "this engine runs static (uniform) verification; the " "verify-length pin only affects window-computing modes, so " @@ -3360,10 +3291,9 @@ def set_dspark_verify_len_pin(self, verify_len: Optional[int]) -> Optional[int]: if (worker.sts_recorder is not None and verify_len is not None and int(verify_len) < int(self.model_engine.max_draft_len)): - # The collector's own guard refuses to START under a censoring - # window; this closes the runtime door that guard cannot see. A - # sub-block pin clips the acceptance label of every sample - # collected after it, and the shards keep looking healthy. + # A sub-block pin censors the acceptance label of every STS sample + # collected after it; the collector's start-time guard cannot see + # a runtime pin. raise RuntimeError( "STS calibration collection is active on this engine; a pin " "below the full block censors the acceptance label for every " @@ -3375,14 +3305,10 @@ def _dspark_confidence_draft_len(self, scheduled_batch: ScheduledRequests) -> int: """Draft length for this iteration, from the DSpark confidence planner. - Runs entirely outside CUDA-graph capture, which is what makes the - device->host relay legal: the worker writes confidence device-side inside - the captured graph, and this hook -- called before the next forward -- - first *reads* the previously staged snapshot, then stages a fresh one. - That gives a one-iteration lag for free, with no synchronization. - - Falls back to the full block on any missing piece, so an engine without a - confidence-capable checkpoint behaves exactly as before. + Runs outside CUDA-graph capture: it reads the confidence snapshot + staged on the previous iteration, then stages a fresh one -- a + one-iteration lag with no synchronization. Falls back to the full + block on any missing piece. """ max_draft_len = self.model_engine.max_draft_len worker = self.model_engine._get_spec_worker() @@ -3390,40 +3316,29 @@ def _dspark_confidence_draft_len(self, if planner is None: return max_draft_len - self._maybe_arm_dspark_sync_debug() - gen_requests = scheduled_batch.generation_requests num_gen_requests = len(gen_requests) - # Confidence is read back by SLOT, not by batch position. The snapshot - # is one iteration old, and joins/departures reshuffle the batch in - # between, so position i is routinely a different request than the one - # that was scored there. A request with no slot yet maps to the neutral - # row, i.e. "verify the whole block". + # Read back by SLOT, not batch position: the snapshot is one iteration + # old and the batch reshuffles in between. Slotless requests map to + # the neutral row ("verify the whole block"). confidence_rows = [ worker.confidence_row_for(request.py_request_id) for request in gen_requests ] - # The mode, not the config flag, decides whether windows are computed: - # the environment override exists so a run can be switched between - # `static` and `compact` without editing the serving config, and every - # counter below is keyed off the same decision. + # The mode, not the config flag, decides whether windows are computed; + # every counter below is keyed off the same decision. stats = worker.ragged_stats - # The sampler is built independently of the worker, so nothing connects - # them; do it here, where both are reachable. Acceptance has to be - # recorded in the sampler (that is where the accepted count and the - # window it was measured against are both host-side and belong to the - # same step), but the counters live on the worker. + # Wire the worker's counters into the sampler here, where both are + # reachable: acceptance is only observable host-side in the sampler. if (stats is not None and self.sampler is not None - and getattr(self.sampler, "acceptance_stats", None) is None): + and self.sampler.acceptance_stats is None): self.sampler.acceptance_stats = stats - # Same wiring, same reason, for STS collection: the logits live on the - # worker and the accepted count is only host-side in the sampler. The - # sampler gets the worker's OWN row resolver and pass counter, never a - # raw buffer: the recorder joins by (draft pass, buffer row), and both - # keys must come from the allocator that wrote the buffer. + # Same wiring for STS. The sampler gets the worker's OWN row resolver + # and pass counter: the recorder joins by (draft pass, buffer row), + # and both keys must come from the allocator that wrote the buffer. if (self.sampler is not None - and getattr(self.sampler, "sts_recorder", None) is None): + and self.sampler.sts_recorder is None): recorder = worker.sts_recorder if recorder is not None: self.sampler.sts_recorder = recorder @@ -3437,8 +3352,7 @@ def _dspark_confidence_draft_len(self, # --- Phase 1: decide locally, issue nothing --------------------------- # - # goal doc H9. Every decline reason is rank-local and one - # (_ready_snapshot) is timing-dependent, so a collective placed after + # Every decline reason is rank-local, so a collective placed after # them deadlocks: nothing below may branch on rank-local state before # the single collective, and nothing after it may issue another one. ragged_lens = None @@ -3453,28 +3367,14 @@ def _dspark_confidence_draft_len(self, else: fallback_reason = "mode_static" # Nothing to compute: the fallback verifies the full drafted block. - # (The batch-wide uniform-tier middle ground was removed; llm_args - # now rejects the config that reached it, so two states, not three.) # --- Phase 2: one collective, same shape on every rank, every step ---- # - # Carries everything the rest of the step needs, so the `peer_stats` - # gather that used to follow is folded in here rather than added to it: - # [wants_ragged, num_rows, total_verify_tokens, max_window, can_graph] - # `_can_queue` already runs `tp_allgather(batch_size)` just before this - # hook and could absorb this too, but that would mean computing the - # local draft length inside a hot path shared by every speculation mode. - # - # `can_graph` rides along because the ragged fit has to know what row - # padding will actually do, and it runs before the graph runner's own - # allgather answers that. `_get_padded_batch` only takes the cross-rank - # row maximum when EVERY rank can graph this step; one rank holding a - # context request makes each rank keep its local count instead, so a fit - # that assumed the maximum budgets tokens for pad rows that are never - # appended. Predicting that from size rules alone is what `will_pad_to` - # tried and got wrong. Carrying the input costs nothing here -- the - # collective is already being issued, and one more int does not change - # its cost -- and it replaces a second derivation with the same data. + # Payload: [wants_ragged, num_rows, total_verify_tokens, max_window, + # can_graph, pending_pin]. `can_graph` rides along because the ragged + # fit runs before the graph runner's own allgather and must know + # whether row padding will happen: _get_padded_batch only takes the + # cross-rank row maximum when EVERY rank can graph this step. local_rows = len(ragged_lens) if ragged_lens is not None else 0 local_tokens = (sum(1 + int(v) for v in ragged_lens) if ragged_lens else 0) @@ -3492,76 +3392,52 @@ def _dspark_confidence_draft_len(self, 1 if scheduled_batch.can_run_cuda_graph else 0, planner.pending_verify_len_pin(), ]) - # A runtime verify-length pin (profiling endpoint) reaches the - # group here rather than through each rank's own environment: the - # endpoint lands on one rank, and a pin applied locally would make - # that rank disagree with its peers for as long as it held. Taking - # the first rank that queued one makes the choice a function of the - # allgathered payload, identical everywhere, effective on the same - # step. -1 means nobody queued anything. + # The pin rides the allgather so every rank adopts the same value + # on the same step (the endpoint lands on one rank). -1 means + # nobody queued anything. agreed_pin = next( (int(payload[5]) for payload in payloads if len(payload) > 5 and int(payload[5]) >= 0), -1) planner.adopt_verify_len_pin(agreed_pin) if not all(int(payload[0]) for payload in payloads): - # Any rank that cannot go ragged takes the whole group with it: - # a mixed step means one rank replays a ragged graph while - # another runs a different token count, and `all_rank_num_tokens` - # is frozen into the captured graph, so the mismatch would not - # even be observable at replay time. + # All-or-nothing across ranks: `all_rank_num_tokens` is frozen + # into the captured graph, so a mixed step's mismatch would + # not even be observable at replay time. if ragged_lens is not None: fallback_reason = "peer_declined" ragged_lens = None else: - # payload[3] (the local max window) stays in the collective so - # its shape is stable, but capping local windows at the global - # max was a no-op by construction (the global max is >= every - # local window) and nothing else consumed it. - # Third element is the group's answer, identical on every rank - # by construction, not each rank's own flag: what the fit needs - # to know is whether the row maximum will be taken at all. + # payload[3] is unread; it stays so the collective keeps a + # stable shape. peer_stats' third element is the GROUP's + # can-graph answer, identical on every rank, not each rank's + # own flag. all_can_graph = all(int(payload[4]) for payload in payloads) peer_stats = [[ int(payload[1]), int(payload[2]), 1 if all_can_graph else 0 ] for payload in payloads] else: - # A single-rank world has no allgather to carry the pin, and - # without this branch a queued pin sat forever: the endpoint - # returned 200, forced_steps stayed 0, and a sweep labelled every - # cell with a length the planner never ran. The local queue IS the - # group's agreement when the group is one rank. + # Single-rank world: the local queue IS the group's agreement. planner.adopt_verify_len_pin(planner.pending_verify_len_pin()) # --- Phase 3: no more collectives ------------------------------------- # - # Land the batch's token total on a captured bucket before the graph key - # is built, spending the rounding slack on real verification. A batch - # with no captured shape gets None back and falls through to uniform, - # which always has a graph -- running ragged without one would cost far - # more than the trimmed tokens save. This can still come out differently - # on different ranks (the pad-row fit is sized from the local request - # count), which is why the attention-DP graph gate compares the token - # bucket and not just the batch size. + # Land the token total on a captured bucket BEFORE the graph key is + # built; no captured shape means fall through to uniform, which always + # has a graph. The fit can still differ across ranks, which is why the + # attention-DP graph gate compares the token bucket, not just the + # batch size. # - # `cap-accept` stops short of this: it wants the schedule's commitment - # policy without the layout that pays for it, so it submits the full - # block and needs no bucket at all. The windows go onto `py_verify_cap` - # rather than `py_verify_len` precisely so the whole token-layout path - # below and downstream (input layout, DSA row expansion, graph key, KV - # rewind) keeps seeing an untrimmed batch -- the mode is then a pure - # addition rather than a second set of conditions inside code that took - # three bugs to get right for `compact`. + # `cap-accept` submits the full block: its windows go on + # `py_verify_cap`, not `py_verify_len`, so the whole token-layout path + # downstream keeps seeing an untrimmed batch. bucket = None ragged_active = False cap_accept = (mode is not None and mode.computes_windows and not mode.trims_submitted_tokens) if cap_accept: - # Stale caps are the same hazard as stale windows, minus the one - # thing that makes stale windows survivable: a mismatched layout - # eventually trips an assert, whereas a stale cap just commits - # against last step's schedule with nothing to contradict it. - # Clear every request first, set second. + # Clear every request first, set second: a stale cap silently + # commits against last step's schedule, with no assert to trip. for request in gen_requests: request.py_verify_cap = None if ragged_lens is not None: @@ -3576,10 +3452,9 @@ def _dspark_confidence_draft_len(self, if bucket is None: fallback_reason = "no_captured_shape" else: - # No rank can replay a graph this step (context requests - # present); such steps were always eager, so run the windows - # eagerly. Do not fit a bucket: fitting sizes pad rows that - # _get_padded_batch will not append on an ungraphable step. + # No rank can graph this step, so run the windows eagerly. Do + # not fit a bucket: the fit sizes pad rows _get_padded_batch + # will not append on an ungraphable step. runner = self.model_engine.cuda_graph_runner runner.agreed_ragged_bucket = None runner.ragged_pad_verify_len = 0 @@ -3587,32 +3462,23 @@ def _dspark_confidence_draft_len(self, request.py_verify_len = int(window) ragged_active = True - # Record what the step actually decided, reading the windows back off - # the requests rather than off `ragged_lens`: fit_ragged_verify_lens - # spends the bucket's rounding slack on real verification, so the - # fitted windows are what the target will see. + # Record what the step actually decided. if stats is not None: - # The stats object lives in the worker process; under TP the caller - # that wants to assert on it is on the far side of an MPI boundary - # and cannot reach it. Emit a periodic summary so the decision is - # recoverable from the log, which is the only shared channel. + # Periodic summary: under TP the counters sit across an MPI + # boundary, so the log is the only shared channel. if stats.steps_total and stats.steps_total % 32 == 0: stats.log_summary(prefix="DSpark ragged verify [periodic]") self._maybe_assert_dspark_ragged_active(stats) fitted = None delivered = None if cap_accept: - # Nothing was fitted, so there is nothing to read back: the - # scheduled windows ARE the decision. The block is submitted - # whole, which `delivered` has to say explicitly -- otherwise - # the windows would be booked as if they had been submitted. + # Nothing was fitted: the scheduled windows ARE the decision, + # and `delivered` must say the block was submitted whole. fitted = [int(v) for v in ragged_lens] if ragged_lens else None delivered = num_gen_requests * (1 + max_draft_len) elif ragged_active: - # Read back off the requests, not off `ragged_lens`: when a - # bucket was fitted its rounding slack was spent on real - # verification, so the fitted windows are what the target sees. - # Without a bucket the two are the same thing. + # Read back off the requests, not `ragged_lens`: the fit spent + # the bucket's rounding slack on real verification. fitted = [ int(request.py_verify_len) for request in gen_requests if getattr(request, "py_verify_len", None) is not None @@ -3622,10 +3488,9 @@ def _dspark_confidence_draft_len(self, stats.record_step(num_gen_requests=num_gen_requests, verify_lens=fitted, bucket=bucket, - # Only when the fit delivered (bucket non-None): - # reading it every step books the LAST fit's row - # count into padded_bs_hist on eager-ragged and - # cap-accept steps, which never enter the fit. + # Only when the fit delivered: otherwise this + # books the LAST fit's row count on steps that + # never entered the fit. padded_bs=(getattr(self.model_engine, "_dspark_last_padded_bs", None) if bucket is not None else None), @@ -3633,37 +3498,24 @@ def _dspark_confidence_draft_len(self, delivered=delivered) if ragged_active and not cap_accept: - # The block is always drafted in full -- only verification is - # trimmed -- so the batch-wide draft length stays at the top tier. - # That keeps the drafted-token buffer and the padded acceptance - # rectangle a fixed width; each request's own (smaller) verify - # window rides along on py_verify_len. + # The block is always drafted in full, so the batch-wide draft + # length stays at the top tier (fixed-width buffers); per-request + # windows ride on py_verify_len. runtime_draft_len = max( int(t) for t in self.model_engine.spec_config.verify_len_tiers) else: - # Clear any window left over from a previous ragged step, otherwise - # a fallback iteration would keep trimming on a stale split. This - # has to cover EVERY generation request: _prepare_tp_inputs builds - # the token layout per request from py_verify_len, but - # _attach_ragged_verify_layout disables the ragged spec metadata as - # soon as one request is missing a window. A batch where only some - # requests still carry a stale window therefore gets a ragged input - # layout with uniform acceptance -- silent token misattribution. + # Clear stale windows on EVERY generation request: a batch where + # only some requests carry one gets a ragged input layout with + # uniform acceptance -- silent token misattribution. for request in gen_requests: request.py_verify_len = None - # And the published fit state with them: a fallback step that - # inherits the previous ragged step's agreed bucket carries a - # claim about a batch it does not describe. The window guard in - # the graph key tolerates that today; the layout assert's uniform - # branch does not, on purpose -- one line of defense where there - # were zero. + # Clear the published fit state too: a stale agreed bucket + # describes a batch that no longer exists. if self.model_engine.cuda_graph_runner is not None: self.model_engine.cuda_graph_runner.agreed_ragged_bucket = None self.model_engine.cuda_graph_runner.ragged_pad_verify_len = 0 - # Verify the full drafted block. This is the only fallback: the - # uniform tier ladder that used to sit between "schedule per - # request" and "verify everything" is gone, so there is nothing to - # agree on across ranks here and no collective is needed. + # The only fallback is the full drafted block; nothing to agree on + # across ranks, so no collective. runtime_draft_len = planner.max_tier # Stage this step's confidence for the *next* decision. Non-blocking; if @@ -3679,11 +3531,9 @@ def _dspark_confidence_draft_len(self, stamps = worker.confidence_stamp_buffer() planner.stage_confidence(buffer, stamps=stamps, staged_seq=staged_seq) - # The STS recorder keeps its own ring of these snapshots: its - # labels arrive one-plus steps later and must be paired with the - # pass that DRAFTED the block being labeled, not with whatever the - # live buffer holds by then. Same relay, same stamps, staged in - # the same breath as the planner's copy. + # The STS recorder rings its own copy: labels arrive steps later + # and must pair with the pass that DRAFTED the block, not whatever + # the live buffer holds by then. recorder = worker.sts_recorder if recorder is not None: recorder.stage_snapshot(device_logits=buffer, diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 54b424c95b6b..c221380490db 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -800,14 +800,9 @@ def add_token( def int_tensor(shape: tuple[int, ...], device: str = "cuda") -> torch.Tensor: - # zeros, not empty: several of these buffers are slot-indexed stores that - # are only written for the slots sampled in a given step, while consumers - # (the overlap scheduler's previous-batch gather in model_engine) index - # them by slots from a *previous* roster. A slot that has not been written - # yet must read as 0 -- with empty it reads arbitrary bytes, and one such - # read (previous_kv_lens_offsets) feeds the captured kv_lens correction, - # where a garbage value walks the KV append past the request's allocated - # blocks (the DSpark ragged IMA). + # zeros, not empty: these buffers are slot-indexed stores whose consumers + # (e.g. the overlap scheduler's previous-batch gather) may read slots the + # current step never wrote; an unwritten slot must read as 0. return torch.zeros(shape, dtype=torch.int, device=device) diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 5c1f11aec890..8116971abfa7 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -13,11 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. # -# DSpark worker / metadata mirror the DFlash plumbing (capture target-layer -# hidden states, accept the previous block with standard verification, draft a -# new block in one backbone forward), adapted to DSpark's draft model which -# produces the whole block (and its confidence-truncated length) inside a single -# ``DSparkDraftModel.forward`` rather than via mask-token cross-attention. +# DSpark speculative decoding: capture target-layer hidden states during the +# target forward, verify the previous block, and draft a new block in a single +# ``DSparkDraftModel.forward``. from collections import deque from dataclasses import dataclass @@ -36,8 +34,8 @@ if TYPE_CHECKING: from ...llmapi.llm_args import DSparkDecodingConfig -# Raw-logit fill for an unscored confidence slot: sigmoid(30) rounds to 1.0 in -# fp32, so an unwritten row schedules as "verify the whole block". +# Unscored-slot fill: sigmoid saturates to 1.0, so an unwritten row schedules +# as "verify the whole block" (fail-safe). _NEUTRAL_CONFIDENCE_LOGIT = NEUTRAL_CONFIDENCE_LOGIT @@ -45,16 +43,10 @@ class DSparkSpecMetadata(SpecMetadata): """Metadata for DSpark speculative decoding. - Captures hidden states from the target model's ``layers_to_capture`` during - the target forward pass. DSpark captures the *mean over the multi-head - (mHC) residual streams* at each captured layer (handled by the target-side - capture hook), concatenated across layers, and feeds them to the draft - model's ``main_proj`` + ``main_norm`` (inside ``DSparkDraftModel.forward``) - as the captured-context attention input (``main_x``). - - Mirrors :class:`DFlashSpecMetadata`; the only DSpark-specific detail is that - the per-layer captured width is the model hidden size (post hc-mean), so the - buffer is ``[max_num_tokens, hidden_size * num_capture_layers]``. + Captures target-layer hidden states (mean over the mHC residual streams, + per-layer width = model hidden size) into a + ``[max_num_tokens, hidden_size * num_capture_layers]`` buffer that feeds + ``DSparkDraftModel.forward``. """ batch_indices_cuda: Optional[torch.Tensor] = None @@ -76,11 +68,9 @@ def __post_init__(self): self.is_spec_dec_tree = False self.is_spec_dec_dynamic_tree = False - # Set up hidden state capture buffer if self.layers_to_capture is not None and len(self.layers_to_capture) > 0: self.layers_to_capture = sorted(list(self.layers_to_capture)) self.num_capture_layers = len(self.layers_to_capture) - # O(1) lookups for is_layer_capture() and maybe_capture_hidden_states() self._capture_layer_set = frozenset(self.layers_to_capture) self._layer_to_idx = {lid: i for i, lid in enumerate(self.layers_to_capture)} self.captured_hidden_states = torch.empty( @@ -105,10 +95,9 @@ def prepare(self): ) self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) - # CUDA-graph-safe path: maintain the request->slot mapping on the host - # (outside the captured region) and mirror it into ``_batch_to_slot`` so the - # captured gen forward can index the rolling windows by tensor. Mirrors - # ``DFlashSpecMetadata.prepare`` (dflash.py:96-113). + # Maintain the request->slot map on the host (outside any captured + # region) and mirror it into ``_batch_to_slot`` so the captured gen + # forward indexes the rolling windows by tensor. worker = getattr(self, "_dspark_worker", None) if worker is not None and worker._win_inited: current = set(self.request_ids) @@ -118,16 +107,11 @@ def prepare(self): worker._ctx_len[slot] = 0 worker._kv_windows[slot].zero_() worker._free_slots.append(slot) - # Assign a persistent rolling-window slot to every real generation - # request that never ran a context/seed forward on this worker. In - # disaggregated serving the prompt is prefilled (and the window - # seeded) on the *context* server, so ``_seed_context_windows`` never - # runs on the generation server and ``_req_to_slot`` stays empty; - # without this, all concurrent gen requests fall through to the shared - # scratch row below and corrupt each other's draft window at batch - # size > 1 (GitHub #16767). Context-prefix entries are left to - # ``_seed_context_windows``; the ADP-idle (id 0) and CUDA-graph - # padding dummies are kept on the scratch row. + # Assign a slot to every real gen request that never ran a + # context/seed forward on this worker (disaggregated serving + # prefills on the context server), so concurrent gen requests do + # not share the scratch row. Context requests are seeded by + # ``_seed_context_windows``; dummies stay on the scratch row. num_contexts = max(0, len(self.request_ids) - self.num_generations) for rid in self.request_ids[num_contexts:]: if ( @@ -136,10 +120,9 @@ def prepare(self): and rid not in worker._req_to_slot ): worker._assign_slot(rid, reset=False) - # Unknown request IDs (synthetic warmup / CUDA-graph padding, ADP idle - # requests, or disagg seed forwards without a real id) map to the - # dedicated throwaway scratch row so they cannot overwrite a live - # request's rolling window (they previously aliased to slot 0). + # Unknown request IDs (warmup / CUDA-graph padding, ADP idle, disagg + # seeds without a real id) map to the throwaway scratch row so they + # cannot overwrite a live request's rolling window. scratch = worker._scratch_slot mapping = torch.tensor( [worker._req_to_slot.get(rid, scratch) for rid in self.request_ids], @@ -155,14 +138,11 @@ def is_layer_capture(self, layer_id: int) -> bool: def maybe_capture_hidden_states( self, layer_id: int, hidden_states: torch.Tensor, residual: Optional[torch.Tensor] = None ) -> None: - """Capture hidden states from a target model layer into the buffer. - - DeepSeek-V4 keeps the multi-head (mHC) residual stream flattened as - ``[num_tokens, hc_mult * hidden]``; DSpark captures the *mean over the hc - streams* (reference ``h.mean(dim=2)`` with ``h`` shaped - ``[*, hc_mult, hidden]``). We reduce here so the V4 decoder layer's - existing capture call is unchanged. A ``[num_tokens, hidden]`` input - (already reduced / non-mHC) is stored as-is. + """Capture a target layer's hidden states into the buffer. + + An mHC-flattened ``[num_tokens, hc_mult * hidden]`` input is reduced by + mean over the hc streams; a ``[num_tokens, hidden]`` input is stored + as-is. """ if self.captured_hidden_states is None: return @@ -170,7 +150,6 @@ def maybe_capture_hidden_states( if i is not None: num_tokens = hidden_states.shape[0] to_save = hidden_states + residual if residual is not None else hidden_states - # mHC residual -> mean over the hc_mult streams. if to_save.shape[-1] != self.hidden_size: hc_mult = to_save.shape[-1] // self.hidden_size to_save = to_save.reshape(num_tokens, hc_mult, self.hidden_size).mean(dim=1) @@ -190,29 +169,12 @@ def get_hidden_states(self, num_tokens: int) -> Optional[torch.Tensor]: class DSparkWorker(SpecWorkerBase): """Worker for DSpark speculative decoding. - DSpark drafts a whole block of ``block_size`` tokens in one backbone forward - (``DSparkDraftModel.forward``): it projects the captured target-layer hidden - states (``main_proj`` + ``main_norm``) into the draft's captured-context - attention, runs the ``num_stages`` DSpark blocks over a rolling captured - window, refines the per-position logits with the Markov head, and (when - ``enable_confidence_scheduling`` is set) scores each position with the - confidence head. The block is always proposed in full; the confidence only - feeds the verification scheduler's budget, never the acceptance decision. - - Unlike DFlash, the draft does NOT use the paged KV cache or mask-token - cross-attention: its attention K/V come from the worker-owned rolling window - of projected captured context (one ``main_kv`` per decode step, per stage). - Acceptance of the previous block goes through the unified - :meth:`SpecWorkerBase.sample_and_accept_draft_tokens` (strict target-verify, - or rejection sampling for a non-greedy batch), so greedy parity with no-spec - is preserved regardless of draft quality. - - The rolling window is kept consistent across the whole decode: it is seeded - from the prompt's captured context at prefill and back-filled with the - intermediate accepted tokens of a multi-accept step (both via - ``DSparkDraftModel.write_context_windows``), in addition to the per-step bonus - write done by the generation path. These affect draft acceptance rate only, - not correctness, which the standard target verify guarantees. + Drafts a whole block of ``block_size`` tokens in one + ``DSparkDraftModel.forward`` over a worker-owned rolling window of + projected captured context (no paged KV cache). Acceptance goes through + :meth:`SpecWorkerBase.sample_and_accept_draft_tokens`, so the window state + affects acceptance rate only, never correctness. Confidence scores feed + only the verification scheduler's budget, never the acceptance decision. Reference: DeepSeek DeepSpec (https://github.com/deepseek-ai/DeepSpec). """ @@ -234,23 +196,19 @@ def __init__( self._ctx_len: Optional[torch.Tensor] = None # [max_batch] abs decode position self._win = 0 - # Slot management. ``_req_to_slot`` (python dict) + ``_free_slots`` are the - # source of truth, updated in prepare()/forward(); ``_batch_to_slot`` is the - # CUDA mirror (request-order -> slot) read by the CUDA-graph-safe batched - # gen path (set on the host in prepare(), so the captured forward indexes - # the rolling windows through a tensor instead of a python dict lookup). + # Slot management: ``_req_to_slot`` + ``_free_slots`` (host) are the + # source of truth; ``_batch_to_slot`` is the CUDA mirror + # (request-order -> slot) the captured gen forward indexes through. self._req_to_slot = {} # request_id -> slot index self._free_slots = deque() # available slot indices self._batch_to_slot: Optional[torch.Tensor] = None # [max_batch] long, cuda - # Index of the throwaway "scratch" window row that absorbs padded / - # unknown request IDs (set in ``_lazy_init`` to ``max_batch``); it is - # never handed out through ``_free_slots``. + # Throwaway window row for padded / unknown request IDs (set to + # ``max_batch`` in ``_lazy_init``); never handed out via ``_free_slots``. self._scratch_slot = 0 - # Confidence-scheduled verification. ``return_confidence`` is read once - # here (never per step) so the captured draft graph cannot diverge. - # ``_confidence_logits`` is the slot-indexed handoff buffer to the - # verification scheduler; see ``_lazy_init`` for its neutral value. + # ``return_confidence`` is read once here (never per step) so the + # captured draft graph cannot diverge. ``_confidence_logits`` is the + # slot-indexed handoff to the verification scheduler. self.return_confidence = bool(spec_config.enable_confidence_scheduling) import os as _os @@ -269,17 +227,14 @@ def __init__( self._draft_seq_cuda: Optional[torch.Tensor] = None # Row index of the permanently-neutral confidence row; see ``_lazy_init``. self._neutral_conf_row = 0 - # Host-side coordinator that turns the confidence snapshot into this - # iteration's draft length. Built in ``_lazy_init`` (it needs the draft - # model's block size and calibration head). + # Host-side planner turning the confidence snapshot into this + # iteration's draft length; built in ``_lazy_init``. self.verify_planner = None - # The generation draft path is the batched, host-sync-free - # ``_draft_gen_block_batched`` + ``DSparkDraftModel.forward_batched`` + - # ``dspark_attention_forward_batched``: it is correct in eager mode AND safe - # to capture into the target's CUDA graph (DSpark is a one-engine drafter — - # its worker forward runs inside that graph, so the draft path MUST be - # capture-safe whenever ``cuda_graph_config`` is set). + # DSpark is a one-engine drafter: the worker forward runs inside the + # target's CUDA graph, so the gen draft path + # (``_draft_gen_block_batched``) must stay host-sync-free and + # capture-safe whenever ``cuda_graph_config`` is set. logger.info( f"DSparkWorker initialized with " @@ -289,11 +244,8 @@ def __init__( def staged_confidence_buffer(self) -> Optional[torch.Tensor]: """The whole slot-indexed confidence buffer, or None when disabled. - The *whole* buffer, not the current batch's rows: the planner reads it - back by slot, so it must be able to look up any request that is still - alive. It is small (``max_num_requests x block`` fp32 -- kilobytes), so - staging all of it costs nothing and removes the need to communicate a - row ordering across the one-iteration lag. + The whole buffer, not the current batch's rows: the planner reads it + back by slot across the one-iteration lag. """ return self._confidence_logits @@ -304,11 +256,9 @@ def confidence_stamp_buffer(self) -> Optional[torch.Tensor]: def bump_draft_seq(self) -> int: """Advance the draft-pass sequence and return the PREVIOUS value. - Called once per step from the executor's host path, before this step's - forward is enqueued: the in-graph stamp scatter then labels this step's - confidence rows with the new sequence, while the returned previous - value is the sequence of the last COMPLETED draft -- i.e. the newest - content a snapshot staged right now can possibly contain. + Must be called once per step from the executor's host path before the + step's forward is enqueued; the returned previous value is the + sequence of the last completed draft. """ prev = self._draft_seq_host self._draft_seq_host += 1 @@ -319,15 +269,9 @@ def bump_draft_seq(self) -> int: def confidence_row_for(self, req_id: int) -> int: """Buffer row holding ``req_id``'s confidence, host-side. - Falls back to the permanently-neutral row for a request that has not - been drafted yet (or was never assigned a slot). Neutral reads as - "certain to be accepted", i.e. verify the full block -- the safe - direction, and the same answer today's static behavior gives. - - Keying by slot rather than by batch position is what makes the - one-iteration-lagged snapshot usable: the batch is reshuffled by joins - and departures between the step that wrote the confidence and the step - that reads it, so position ``i`` is routinely a different request. + Falls back to the permanently-neutral row ("verify the full block" -- + fail-safe). Keyed by slot, not batch position, so the one-iteration- + lagged snapshot survives batch reshuffles. """ if self._confidence_logits is None: return self._neutral_conf_row @@ -336,12 +280,8 @@ def confidence_row_for(self, req_id: int) -> int: def sts_row_for(self, req_id: int) -> Optional[int]: """``confidence_row_for`` for the STS recorder: None, never neutral. - The planner WANTS the neutral fallback (it reads as "verify the full - block", the safe direction). The recorder must not get it: the neutral - row is a real buffer row that no draft ever writes, and pairing a - request's label against it would be a fabricated sample, not a safe - default. An unresolvable request is a dropped sample, counted by the - recorder as ``no_row``. + Pairing a label against the never-written neutral row would fabricate + a sample; an unresolvable request is dropped (counted as ``no_row``). """ if self._confidence_logits is None: return None @@ -350,12 +290,9 @@ def sts_row_for(self, req_id: int) -> Optional[int]: def verified_draft_seq(self) -> Optional[int]: """The draft pass that produced the block the CURRENT step verifies. - Read at sampling time: by then this step's ``bump_draft_seq`` has - already advanced the counter for the draft running alongside the - verification, so the pass being verified is the previous value. The - first pass is numbered 1 (the stamp buffer starts at zeros, so a - sequence of 0 would false-match every never-written row -- including - the neutral one), hence None until a pass numbered >= 1 exists. + Read at sampling time, after this step's ``bump_draft_seq``, so the + verified pass is the previous value. Passes start at 1: a sequence of + 0 would false-match every never-written stamp row. """ return self._draft_seq_host - 1 if self._draft_seq_host > 1 else None @@ -378,23 +315,18 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: self._win = int(draft_model._attn_params["window_size"]) head_dim = int(draft_model._attn_params["head_dim"]) - # Real requests occupy slots ``[0, max_batch)``; one extra "scratch" row - # at index ``max_batch`` absorbs padded / unknown request IDs (CUDA-graph - # padding, ADP idle requests, or disagg seed forwards that arrive without - # a real request id) so they can never overwrite a live request's rolling - # window. Previously such IDs aliased to slot 0 and corrupted whichever - # real request occupied it. The scratch row is never handed out through - # ``_free_slots`` and its contents are throwaway. + # Real requests occupy slots ``[0, max_batch)``; the extra scratch row + # at index ``max_batch`` absorbs padded / unknown request IDs so they + # can never overwrite a live request's rolling window. It is never + # handed out through ``_free_slots``. self._scratch_slot = max_batch num_rows = max_batch + 1 - # CUDA-graph padding requests carry ids in - # ``[CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, CUDA_GRAPH_DUMMY_REQUEST_ID]``, - # while real request ids start at ``max_batch_size`` and grow, so a simple - # floor cleanly separates them. Together with ``ATTENTION_DP_DUMMY_REQUEST_ID`` - # (0) these dummies must route to the scratch row (see ``prepare()``) and - # never consume a real slot. Imported lazily to break the - # dspark -> cuda_graph_runner -> speculative.utils -> dspark import cycle. + # CUDA-graph padding ids sit in + # ``[CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, CUDA_GRAPH_DUMMY_REQUEST_ID]``; + # real ids start at ``max_batch_size``, so a floor separates them. + # Imported lazily to break the dspark -> cuda_graph_runner -> + # speculative.utils -> dspark import cycle. from ..pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID self._graph_dummy_id_floor = CUDA_GRAPH_DUMMY_REQUEST_ID - self.max_draft_len @@ -407,18 +339,10 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: self._ctx_len = torch.zeros(num_rows, dtype=torch.long, device="cuda") self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") if self.return_confidence: - # Neutral fill = a large positive logit, i.e. sigmoid ~ 1.0. A slot - # that has not been scored yet therefore reads as "certain to be - # accepted", which makes the scheduler fall back to verifying the - # full block -- the safe direction (never verifies fewer tokens than - # today's static behavior on stale/unwritten rows). - # - # One row beyond the window buffer: rows ``[0, max_batch)`` are real - # request slots and row ``max_batch`` is the scratch row that padded - # / unknown requests write through, so neither is reliably neutral. - # Row ``max_batch + 1`` is never in ``_batch_to_slot`` and therefore - # never written -- that is the row ``confidence_row_for`` hands out - # for an unscored request. + # Neutral fill = large positive logit (sigmoid ~ 1.0): unscored + # slots schedule as "verify the full block" -- the fail-safe + # direction. The neutral row sits one past the scratch row; it is + # never in ``_batch_to_slot``, so no draft ever writes it. self._neutral_conf_row = num_rows self._confidence_logits = torch.full( (num_rows + 1, block_size), @@ -426,22 +350,18 @@ def _lazy_init(self, draft_model, spec_metadata) -> None: dtype=torch.float32, device="cuda", ) - # Freshness stamps, one per buffer row: the draft-pass sequence - # number of the last write. SGLang carries the same thing - # (_carry_generation) next to its confidence ring; without it there - # is no way to know whether a gathered row is this step's data, a - # stale replay, or the neutral fill -- the planner path has never - # been measured. Written in-graph through the same `slots` scatter - # as the confidence itself; the scalar is bumped outside the graph. + # Per-row freshness stamps: the draft-pass sequence of the last + # write, distinguishing fresh data from a stale replay or the + # neutral fill. Written in-graph through the same ``slots`` scatter + # as the confidence; the scalar is bumped outside the graph. self._confidence_stamp = torch.zeros( (num_rows + 1,), dtype=torch.int32, device="cuda") self._draft_seq_host = 0 - # Bumped by fill_ from the executor's host path, which runs - # OUTSIDE inference mode; a tensor born here (inside the draft - # forward's inference mode) is an "inference tensor" and refuses - # that update -- it killed all 8 ranks on the first decode step. - # The stamp buffer above stays an inference tensor on purpose: - # it is only written in-graph, like _confidence_logits. + # Bumped by ``fill_`` from the executor's host path, outside + # inference mode -- an inference tensor refuses that update, so + # allocate this scalar with inference mode off. The stamp buffer + # above intentionally stays an inference tensor (written in-graph + # only). with torch.inference_mode(False): self._draft_seq_cuda = torch.zeros((), dtype=torch.int32, device="cuda") @@ -475,10 +395,9 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> if table_path: with open(table_path, encoding="utf-8") as f: raw = json.load(f) - # ``bias`` and ``alpha(bs)`` are optional but not cosmetic: Theta is - # a ratio, so the non-trimmable part of the step time moves the - # argmax. Omitting them says "step_time_ms already measures whole - # steps at the deployment's batch size". + # The overhead fields are optional but not cosmetic: omitting them + # asserts ``step_time_ms`` already measures whole steps at the + # deployment's batch size. from .dspark_planner import check_table_fingerprint if (raw.get("_meta") or {}).get("lookup") != "interp": logger.warning( @@ -495,11 +414,9 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> "attention_dp": bool(self.mapping.enable_attention_dp), "block": int(block_size), "max_batch_size": int(max_batch), - # moe_backend deliberately NOT in the live dict yet: the - # MoeConfig.backend -> model_config.moe_backend propagation is - # not textually obvious, and a wrong live value here would - # false-reject a correct table. Until verified, the table's - # moe_backend surfaces on the "check manually" INFO line. + # moe_backend deliberately absent: a wrong live value would + # false-reject a correct table; the table's value surfaces on + # the "check manually" INFO line instead. }) cost_table = SpsCostTable( token_counts=[int(v) for v in raw["token_counts"]], @@ -509,9 +426,8 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> batch_overhead_ms=[float(v) for v in raw.get("batch_overhead_ms", [])], ) else: - # Explicitly flat: the planner treats this as "unprofiled" and keeps - # verifying the full block rather than trimming on a cost model that - # says every extra token is free. + # Flat = "unprofiled": the planner keeps verifying the full block + # rather than trimming on a cost model where every token is free. cost_table = SpsCostTable.flat() logger.warning( "DSpark confidence scheduling is enabled but no " @@ -527,26 +443,16 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> apply_calibration = head.apply_sts if head is not None else None sts_path = self.spec_config.confidence_sts_path if sts_path and head is None: - # Refuse rather than degrade. The old chained getattr defaulted to - # None when the runtime object was the DSparkForCausalLM wrapper - # (which keeps the stages under .dspark_model, not .mtp_layers), so - # every confidence-scheduled run in the first campaign silently - # planned on RAW sigmoid survivals -- wildly overconfident (the - # fitted temperatures are 2.8-10) -- and chose the full block 99.7% - # of the time. The only visible symptom was the ABSENCE of one log - # line. A configured calibration that cannot be attached is a - # wiring bug, never a preference. + # Refuse rather than degrade: a configured calibration that cannot + # be attached is a wiring bug, never a preference. raise ValueError( f"confidence_sts_path is set but no confidence head was found " f"on {type(draft_model).__name__}; calibration cannot be " f"applied, so the planner would schedule on raw sigmoid " f"survivals. This is a model-wiring bug, not a config choice.") if sts_path and head is not None: - # Accepts either key spelling: this repo writes "sts_temperatures", - # SGLang's DSparkStsCalibration writes "temperatures". The vectors - # are interchangeable -- both are fitted against the same cumulative - # product with the same sigmoid -- so rejecting the other spelling - # would refuse a usable table over a name. + # Accepts either key spelling ("sts_temperatures" here, + # "temperatures" in SGLang); the vectors are interchangeable. from .dspark_sts import load_sts_temperatures_from_path temps = load_sts_temperatures_from_path(sts_path) head.load_sts_temperatures(torch.tensor(temps, dtype=torch.float32)) @@ -554,13 +460,11 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> tiers = self.spec_config.verify_len_tiers if not tiers: - # Not configured: derive the ladder from the measured cost curve - # instead of hardcoding one. Each tier is a captured graph whose - # memory comes out of KV cache, so the derivation is capped -- and - # it is only zero-loss at the batch size it was derived for (the - # shelf right edges move with batch size), hence the steady-state - # modal batch size is the right input. A flat table yields just - # [min, block_size], which is the no-scheduling ladder. + # Derive the ladder from the cost curve. Each tier is a captured + # graph whose memory comes out of KV cache, so the derivation is + # capped; tiers are only zero-loss at the batch size they were + # derived for. A flat table yields [min, block_size] + # (no scheduling). from .dspark_planner import derive_verify_len_tiers tiers = derive_verify_len_tiers( @@ -577,9 +481,8 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> apply_calibration=apply_calibration, ) - # The config knob says what was asked for; the environment overrides it - # so a mode can be swapped without editing a serving config. `static` - # unless ragged was explicitly requested, matching the shipped default. + # The env var overrides the config knob so the mode can be swapped + # without editing a serving config. from .dspark_observability import (DSparkRaggedStats, RaggedVerifyMode, read_ragged_verify_mode) @@ -587,12 +490,9 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> if self.spec_config.enable_ragged_verify else RaggedVerifyMode.STATIC) self.ragged_verify_mode = read_ragged_verify_mode(default=configured) - # STS calibration data collection. Off unless the env var is set; when - # on it costs one device->host copy per decode step, which is why it is - # a collection mode rather than something a serving run carries. + # STS collection is env-gated; it costs one device->host copy per + # decode step. from .dspark_sts import make_recorder_from_env - # `block_size` is the local parameter, not an attribute -- the very - # next line uses it the same way. self.sts_recorder = make_recorder_from_env( block_size=int(block_size), rank=int(self.mapping.rank), @@ -602,8 +502,7 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> self.ragged_stats = DSparkRaggedStats(mode=self.ragged_verify_mode, max_draft_len=block_size) - # goal doc A6: the ragged counters say what happened, the planner's say - # why. Sharing the dict (not copying it) keeps the summary live. + # Shared (not copied) so the summary stays live. self.ragged_stats.planner_stats = self.verify_planner.stats logger.info( f"DSpark verify planner: tiers={self.verify_planner.tiers}, " @@ -612,9 +511,6 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> ) if (self.ragged_verify_mode.trims_submitted_tokens and cost_table.is_flat): - # Not a warning about a tuning choice -- in this combination the - # feature cannot do anything at all, and the run would look like a - # successful `compact` run in every log line. logger.warning( "DSpark ragged verify is set to 'compact' but the cost table is " "flat, so the planner's budget degenerates to verify-all and " @@ -641,10 +537,9 @@ def _assign_slot(self, req_id: int, reset: bool) -> int: self._ctx_len[slot] = 0 self._kv_windows[slot].zero_() if self._confidence_logits is not None: - # A recycled slot still holds the *previous* occupant's scores. - # Reset to neutral so the incoming request's first scheduling - # decision verifies the full block instead of inheriting a dead - # draft's low survival. + # A recycled slot still holds the previous occupant's scores; + # reset to neutral so the first scheduling decision verifies + # the full block. self._confidence_logits[slot].fill_(_NEUTRAL_CONFIDENCE_LOGIT) return self._req_to_slot[req_id] @@ -700,30 +595,20 @@ def _draft_gen_block_batched( ) -> torch.Tensor: """CUDA-graph-safe batched gen draft (all gen requests in one forward). - Free of host syncs and data-dependent shapes: per-request quantities - (``nacc``, the bonus, ``main_hidden``, ``start_pos``, the multi-accept - back-fill) are gathered as tensors, slots come from the host-built - ``_batch_to_slot`` mirror, and the backbone runs once via - ``DSparkDraftModel.forward_batched``. Returns the per-position corrected - block logits ``[num_gens, K, vocab]`` (or ``None`` when there is nothing to - draft); the worker feeds them to ``SpecWorkerBase.sample_draft_tokens``. - Confidence truncation stays disabled — the full block is proposed. + Must stay free of host syncs and data-dependent shapes. Returns block + logits ``[num_gens, K, vocab]`` for ``sample_draft_tokens``, or + ``None`` when there is nothing to draft. The full block is always + proposed. """ num_gens = batch_size - num_contexts - # Two different widths meet here, and conflating them corrupts the gather: - # K - the DRAFT block width. Fixed by the checkpoint; the draft - # always proposes a full block regardless of how much of it - # the target was asked to verify. - # Kp1 - the TARGET's per-gen-request token stride in ``captured``. - # That is ``runtime_draft_len + 1``, which shrinks when the - # verification scheduler trims the batch's draft length. - # Using max_draft_len for the stride reads into the *next* request's - # hidden states whenever runtime_draft_len < max_draft_len, silently - # drafting from the wrong context. + # K is the draft block width (fixed by the checkpoint); Kp1 is the + # target's per-gen-request token stride in ``captured`` + # (runtime_draft_len + 1, which shrinks when verification is trimmed). + # Striding by max_draft_len instead reads into the next request's + # hidden states. K = self.max_draft_len - # Note the ``or K``: SpecMetadata.runtime_draft_len defaults to 0, not - # None, so an ``is None`` check would leave a stride of 1 on any path - # that has not populated it yet. + # ``or K``: runtime_draft_len defaults to 0 (not None), so an + # ``is None`` check would leave a stride of 1 on unpopulated paths. runtime_draft_len = spec_metadata.runtime_draft_len or K Kp1 = int(runtime_draft_len) + 1 device = accepted_tokens.device @@ -746,14 +631,10 @@ def _draft_gen_block_batched( accepted_tokens[num_contexts:batch_size].gather(1, gidx.unsqueeze(1)).squeeze(1).long() ) # [G] - # Captured target hidden at the bonus position within each request's Kp1 - # processed tokens. - # Where each request's processed tokens start inside `captured`, which - # is indexed by flat token position. A uniform batch strides by Kp1; - # under ragged verification requests contribute different numbers of - # tokens, so the offsets come from the same qo_indptr the input layout - # was packed with. Striding by Kp1 there would have request r read from - # request r-1's tail -- no error, just a silent acceptance collapse. + # Start offset of each request's processed tokens inside ``captured``. + # Uniform batches stride by Kp1; under ragged verification the offsets + # must come from the same qo_indptr the input layout was packed with, + # or request r silently reads request r-1's tail. qo_indptr = spec_metadata.qo_indptr if spec_metadata.verify_lens is not None and qo_indptr is not None: base = gen_start + qo_indptr[:num_gens].to(device=device, dtype=torch.long) @@ -764,13 +645,11 @@ def _draft_gen_block_batched( # count, so a bad one must not turn into an out-of-bounds read. main_hidden = captured[(base + gidx).clamp(min=0, max=captured.shape[0] - 1)] - # Fixed-size ([G, K]) masked back-fill of the intermediate accepted tokens - # (everything but the bonus) into the rolling window — same frames as the - # eager path (old+1 .. old+nacc-1), with j >= nacc-1 masked out. - # The width stays at the static K even when the runtime stride is smaller: - # ``nacc <= runtime_draft_len + 1`` bounds the *valid* j to within this - # request's own block, so the extra columns are always masked out. Keeping - # the width static also keeps this tensor's shape independent of K. + # Fixed-size [G, K] masked back-fill of the intermediate accepted + # tokens (all but the bonus) into frames old+1 .. old+nacc-1. The + # width stays at static K (capture-safe); ``nacc <= runtime_draft_len + # + 1`` bounds the valid j to this request's own block, so the extra + # columns are always masked out. old = self._ctx_len[slots] # [G] pre-increment decode position j = torch.arange(K, device=device) # [K] interim_valid = j.unsqueeze(0) < (nacc.unsqueeze(1) - 1) # [G, K] @@ -783,14 +662,11 @@ def _draft_gen_block_batched( interim_hidden, interim_pos, slots, interim_valid, self._kv_windows ) - # Advance the decode position by the accepted count; start_pos (= post- - # increment ctx_len) matches the eager path's frame value. + # Advance the decode position by the accepted count; start_pos is the + # post-increment ctx_len. start_pos = old + nacc # [G] self._ctx_len[slots] = start_pos - # Surface the per-position corrected block logits ([num_gens, K, vocab]) - # and let SpecWorkerBase.sample_draft_tokens do the (greedy or rejection) - # sampling + TP gather + draft_probs scatter, rather than argmaxing here. _toks, confidence, block_logits = draft_model.forward_batched( main_hidden, bonus, @@ -802,14 +678,10 @@ def _draft_gen_block_batched( return_logits=True, all_rank_num_tokens=all_rank_num_tokens, ) - # Stash the raw [G, K] confidence logits for the verification scheduler. - # Slot-indexed (not batch-position-indexed) so it survives batch - # reshuffles; written in place so a captured graph keeps the same - # storage, and scattered through ``slots`` (a view into - # ``_batch_to_slot``) so a replay lands on whichever slots the current - # batch occupies. The planner reads rows back by slot via - # ``confidence_row_for``, so no row ordering has to survive the - # one-iteration lag. + # Stash the [G, K] confidence for the verification scheduler. + # Slot-indexed (survives batch reshuffles), written in place (a + # captured graph keeps the same storage), and scattered through + # ``slots`` so a replay lands on the current batch's slots. if confidence is not None: self._confidence_logits[slots] = confidence.detach() # Same slots, same capture region: content-only update, graph-safe. @@ -839,42 +711,34 @@ def _forward_impl( spec_metadata._dspark_worker = self self._execute_guided_decoder_if_present(logits) - # Target-verify acceptance via the unified SpecWorkerBase entry: it - # reshapes the stored draft tokens (default (num_gens, runtime_draft_len) - # hook), then routes to strict or rejection sampling. Greedy parity with - # the previous hand-rolled path is preserved (rejection only engages for a - # non-greedy batch with valid draft_probs). + # Acceptance: strict target-verify, or rejection sampling for a + # non-greedy batch with valid draft_probs. accepted_tokens, num_accepted_tokens = self.sample_and_accept_draft_tokens( logits, attn_metadata, spec_metadata ) total_target_tokens = input_ids.shape[0] - # CUDA-graph warmup guard: the warmup forwards (is_cuda_graph set, stream - # NOT yet capturing) run synthetic gen batches that would otherwise advance - # the persistent rolling-window state. Snapshot and restore it so warmup is - # side-effect-free. (During the capture pass itself the stream IS capturing, - # so we skip the save/restore and let the ops be recorded; real requests - # reset their slot's window+ctx_len at prefill, wiping any capture-time - # mutation.) + # CUDA-graph warmup guard: warmup forwards (is_cuda_graph set, stream + # NOT yet capturing) run synthetic gen batches; snapshot/restore the + # persistent window state so they are side-effect-free. During the + # capture pass the stream IS capturing, so skip the save/restore and + # let the ops be recorded; prefill resets wipe capture-time mutation. is_warmup = ( - getattr(spec_metadata, "is_cuda_graph", False) + spec_metadata.is_cuda_graph and not torch.cuda.is_current_stream_capturing() ) if is_warmup: saved_ctx_len = self._ctx_len.clone() saved_windows = self._kv_windows.clone() - # The confidence buffer is persistent state too: warmup scores - # synthetic hidden states, and those scores would otherwise be the - # first thing the verification planner ever reads. + # The confidence buffer is persistent state too; warmup scores + # must not reach the verification planner. saved_confidence = ( None if self._confidence_logits is None else self._confidence_logits.clone() ) - # Assign / reset window slots for context (prefill) requests and seed each - # request's rolling KV window from its prompt's captured context, so the - # first generation step drafts against real context instead of an all-zero - # window (acceptance-rate only; verified decoding keeps output correct). + # Seed prefill requests' rolling windows from their prompt's captured + # context (affects acceptance rate only). if num_contexts > 0: self._seed_context_windows( draft_model, @@ -884,24 +748,18 @@ def _forward_impl( total_target_tokens, ) - # FUSED_COMM MoE backends (DeepGEMM MegaMoE) synchronize EP ranks with an - # in-kernel phase-flip NVLink barrier that flips on every kernel call, so - # every rank must invoke the draft MoE the same number of times and with - # the same globally-gathered per-rank token list, or the barrier desyncs - # (hang / "unspecified launch failure"). The draft runs over generation - # requests only, each expanded to ``block`` positions, so the per-rank - # draft-MoE token count is ``num_gens * block``. ``all_rank_num_gens`` is - # gathered at metadata-prep time (model_engine, outside any CUDA-graph - # capture region); it is None for non-ADP / single-rank runs, where the - # local ``[num_tokens]`` fallback in ``_forward_stage`` is correct. + # FUSED_COMM MoE backends sync EP ranks with a phase-flip barrier: + # every rank must invoke the draft MoE the same number of times with + # the same globally-gathered per-rank token list (num_gens * block + # each) or the barrier desyncs. ``all_rank_num_gens`` is gathered at + # metadata-prep time, outside any capture region; None for non-ADP / + # single-rank runs (local fallback in ``_forward_stage``). block = int(draft_model.block_size) - all_rank_num_gens = getattr(spec_metadata, "all_rank_num_gens", None) - # A rank with zero local gen requests still has to cross the draft MoE's - # cross-rank barrier, but DeepseekV4MoE's router / shared-expert dense - # GEMMs reject a 0-row input (cuBLAS CUBLAS_STATUS_INVALID_VALUE), so such - # a rank runs a single 1-row dummy through the MoE (like ADP padding). - # Encode that as ``1`` in the globally-shared per-rank token list so every - # rank agrees on the FUSED_COMM chunk count and per-rank slice. + all_rank_num_gens = spec_metadata.all_rank_num_gens + # A rank with zero local gen requests must still cross the draft MoE's + # barrier, but the router / shared-expert GEMMs reject 0-row input, so + # it runs a 1-row dummy -- encoded as ``1`` in the shared per-rank list + # so every rank agrees on the chunk count. all_rank_draft_tokens = ( [max(1, int(g) * block) for g in all_rank_num_gens] if all_rank_num_gens is not None @@ -912,8 +770,6 @@ def _forward_impl( ) if num_gens > 0: - # The batched gen-block draft returns the per-position corrected block - # logits [num_gens, K, vocab] and is CUDA-graph-safe. gen_logits = self._draft_gen_block_batched( draft_model, spec_metadata, @@ -926,25 +782,20 @@ def _forward_impl( all_rank_num_tokens=all_rank_draft_tokens, ) if gen_logits is not None: - # SpecWorkerBase samples the draft tokens (greedy argmax, or - # rejection sampling for a non-greedy batch), performs the TP - # gather, and scatters the proposal distribution into draft_probs. gen_draft_tokens = self.sample_draft_tokens( gen_logits, spec_metadata, batch_size, num_contexts=num_contexts ) - # The context one-hot must match the width the gen scatter just - # published to draft_probs, NOT gen_logits.shape[-1]: under TP the - # draft logits are vocab-sharded and sample_draft_tokens gathers - # them to full vocab before scattering, so the pre-gather shard - # width would leave stale columns and corrupt rejection. + # The context one-hot must match the width published to + # draft_probs, NOT gen_logits.shape[-1]: under TP the draft + # logits are vocab-sharded before sample_draft_tokens gathers + # them to full vocab. gen_vocab = spec_metadata.draft_probs_last_dim else: gen_draft_tokens = torch.zeros((num_gens, K), dtype=torch.int32, device="cuda") gen_vocab = None else: - # No local generation requests: if any peer EP rank has some, we must - # still cross the draft MoE's cross-rank barrier the same number of - # times (zero-token) so a FUSED_COMM phase-flip barrier stays lockstep. + # No local gens: if any peer EP rank has some, still cross the + # draft MoE barrier so the FUSED_COMM phase-flip stays lockstep. if global_has_gen: draft_model.run_moe_lockstep_noop(all_rank_draft_tokens, accepted_tokens.device) gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") @@ -982,9 +833,8 @@ def _forward_impl( "next_draft_tokens": next_draft_tokens, "next_new_tokens": next_new_tokens, } - # cap-accept only: what each request's window threw away, written by - # `apply_accept_caps` above. Rides the sampler's existing copy rather - # than a read of its own, so the measurement costs no synchronization. + # cap-accept only: per-request trim written by ``apply_accept_caps``. + # Rides the sampler's existing copy, so it costs no extra sync. cap_trim = spec_metadata.accept_cap_trim if cap_trim is not None and spec_metadata.accept_caps is not None: outputs["cap_trim_lens"] = cap_trim[:batch_size] diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index e538fd498498..dc8950ae60f2 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -12,50 +12,16 @@ # 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. -"""Verification mode selection and per-step observability for DSpark. - -Why this module exists ----------------------- -Every failure mode of confidence-scheduled verification is *silent*. The -scheduler can decline to trim, the planner can fall back on a flat cost model, -a batch can miss its captured shape and drop out of graph replay, and a -partially-windowed batch can be refused -- and in every one of those cases the -run still produces correct output at the baseline accuracy. A GSM8K pass -therefore proves nothing about whether the feature was ever active. - -So the two questions a reviewer actually has -- *did the ragged path run?* and -*did it trim anything?* -- are made into counters rather than left to -inference, and :meth:`DSparkRaggedStats.assert_ragged_active` turns them into a -test assertion. - -Verification modes ------------------- -Mirrors SGLang's ``SGLANG_RAGGED_VERIFY_MODE``, and for the same reason: it -separates "is the ragged layout computed correctly" from "does the planner -choose to trim", which otherwise fail together and are indistinguishable. - -``static`` - Verify the whole drafted block for every request. The baseline, and what - every non-DSpark path already does. Scheduling is off. - -``cap-accept`` - Compute the per-request windows and carry them through the whole layout, - but still submit the full block to the target and only *commit* the window. - The expensive path is unchanged, so this buys no throughput -- it exists - because its output must be **bit-identical** to ``static``. Any divergence - is a scheduling bug isolated from the kernel-side packing, which is the one - comparison that separates the two. - -``compact`` - Submit only each request's window. The production path, and the only one - that trims tokens. - -Cost-table dependence ---------------------- -``compact`` needs a profiled :class:`~.dspark_planner.SpsCostTable`. Without -one the planner's cost model is flat, every extra verified token looks free, -and the budget degenerates to verify-all -- ragged silently becomes uniform. -:meth:`DSparkRaggedStats.assert_ragged_active` catches exactly that. +"""Verification-mode selection and per-step counters for DSpark ragged verify. + +Failures on this path are silent (output stays correct at baseline accuracy), +so activity is counted rather than inferred, and +:meth:`DSparkRaggedStats.assert_ragged_active` turns the counters into a test +gate. Modes: ``static`` verifies the full block; ``cap-accept`` computes and +commits per-request windows but submits the full block, so its output must +stay bit-identical to ``static``; ``compact`` submits only each request's +window and needs a profiled :class:`~.dspark_planner.SpsCostTable` -- with a +flat cost model the budget degenerates to verify-all. """ import os @@ -71,9 +37,8 @@ "DSparkRaggedStats", ] -#: Overrides ``DSparkDecodingConfig`` so a mode can be selected without editing -#: a test or a serving config. Reading it per call (rather than caching at -#: import) is deliberate: tests flip it with ``monkeypatch.setenv``. +#: Overrides ``DSparkDecodingConfig``. Read per call, not cached at import: +#: tests flip it with ``monkeypatch.setenv``. RAGGED_VERIFY_MODE_ENV = "TLLM_DSPARK_RAGGED_VERIFY_MODE" @@ -86,12 +51,9 @@ class RaggedVerifyMode(str, Enum): @property def computes_windows(self) -> bool: - """Whether per-request windows are computed at all. - - ``cap-accept`` computes them and commits against them without trimming - the submitted batch, so it shares every host-side code path with - ``compact`` except the one that shrinks the token axis. - """ + """Whether per-request windows are computed at all (``cap-accept`` + shares every host-side path with ``compact`` except the token-axis + trim).""" return self is not RaggedVerifyMode.STATIC @property @@ -106,10 +68,8 @@ def read_ragged_verify_mode( """Resolve the verification mode from the environment. Raises: - ValueError: if the variable is set to something that is not a mode. - Silently falling back to ``static`` on a typo would turn "I asked - for compact" into "the feature never ran", which is the one - outcome this module exists to make impossible. + ValueError: on an unrecognized value; a silent fallback to ``static`` + would turn a typo into "the feature never ran". """ value = os.environ.get(RAGGED_VERIFY_MODE_ENV) if value is None or value == "": @@ -125,15 +85,10 @@ def read_ragged_verify_mode( def resolve_ragged_verify_mode(spec_config) -> RaggedVerifyMode: """The mode a config plus the environment actually select. - The config flag alone is not the answer, and using it as one is a bug with - a long fuse: ``enable_ragged_verify`` is what ``cap-accept`` is configured - with too, but cap-accept submits the FULL block. Anything that shapes the - batch -- the CUDA-graph capture set, the graph key -- has to ask whether - the token axis actually shrinks (:attr:`RaggedVerifyMode.trims_submitted_tokens`), - not whether ragged was requested. - - Resolved identically to the DSpark worker's own resolution, and separately - from it, because the capture set is built before any batch exists. + ``enable_ragged_verify`` alone is not the answer: ``cap-accept`` sets it + too but submits the FULL block, so batch-shaping code (capture set, graph + key) must ask :attr:`RaggedVerifyMode.trims_submitted_tokens`. Must stay + identical to the DSpark worker's own resolution. """ configured = (RaggedVerifyMode.COMPACT if getattr( spec_config, "enable_ragged_verify", False) else RaggedVerifyMode.STATIC) @@ -150,27 +105,21 @@ def trims_submitted_tokens(spec_config) -> bool: class DSparkRaggedStats: """Per-step counters for the ragged verification path. - Cheap enough to leave on: one integer histogram update per decode step, no - device work and no synchronization. Nothing here reads a device tensor -- - every value recorded is already a host value by the time the scheduler has - made its decision. + Host-only and cheap enough to leave on: no device reads, no + synchronization. """ def __init__(self, *, mode: RaggedVerifyMode, max_draft_len: int): self.mode = mode self.max_draft_len = int(max_draft_len) - #: The planner's own decision counters, attached by the worker once the - #: planner exists (goal doc A6). They answer a question these counters - #: cannot: *why* a step declined to trim. `fallback_flat_cost` in - #: particular means the run had no profiled SPS table, under which the - #: budget search is constructionally incapable of trimming -- so a run - #: can look entirely healthy here and still have delivered nothing. + #: Planner decision counters, attached by the worker once the planner + #: exists; the only record of WHY a step declined to trim + #: (`fallback_flat_cost` = no profiled SPS table, so no trimming). self.planner_stats: Optional[Dict[str, int]] = None self.steps_total = 0 - #: Steps where per-request windows were produced AND differed from each - #: other. A step where every request got the same window is counted as - #: uniform: it exercises none of the ragged packing. + #: Steps where per-request windows were produced AND differed; a step + #: where every request got the same window counts as uniform. self.steps_ragged = 0 #: Windows were produced, but every request got the same one. self.steps_uniform_windows = 0 @@ -187,13 +136,12 @@ def __init__(self, *, mode: RaggedVerifyMode, max_draft_len: int): self.verify_len_hist: Counter = Counter() self.bucket_hist: Counter = Counter() self.padded_bs_hist: Counter = Counter() - #: why the graph runner declined -> count. The reasons mean opposite - #: things: "peer_shape_mismatch" is one rank declining to go ragged and - #: dragging the world eager, "key_not_captured" is the bucket grid being - #: wrong, "peer_not_gen_only" is ordinary continuous batching. - #: Steps whose batch held at least two generation requests. The gate - #: counts these rather than all steps: trtllm-bench ramps concurrency, - #: so step 64 can still be a batch of two. + #: why the graph runner declined -> count. "peer_shape_mismatch" = one + #: rank declined ragged and dragged the world eager; "key_not_captured" + #: = wrong bucket grid; "peer_not_gen_only" = ordinary continuous + #: batching. + #: Steps whose batch held at least two generation requests; the gate's + #: floor counts these, not all steps (concurrency ramps up). self.steps_multi_request = 0 self.graph_miss_reasons: Dict[str, int] = {} #: the actual graph key, kept only for the reasons where it localizes @@ -201,41 +149,25 @@ def __init__(self, *, mode: RaggedVerifyMode, max_draft_len: int): self.graph_miss_shapes: Dict[str, int] = {} #: Acceptance, split by whether the request's window was shortened. - #: This is the term the trim/ceiling ratio needs and the one TRT-LLM - #: did not have: `ceiling_tokens` above counts tokens *submitted* with - #: no trim, which answers "how much compute did we save" and never - #: "how much acceptance did we destroy". self.accepted_tokens = 0 self.requests_scored = 0 #: Requests whose window was shorter than the full block. self.requests_trimmed = 0 - #: Trimmed requests that accepted their ENTIRE window. Those are the - #: ones where the draft was still alive at the cut, so trimming - #: certainly cost acceptance -- a trimmed request that accepted fewer - #: than its window would have died on its own and cost nothing. This is - #: an exact, unbiased count, which is why it is preferred here over - #: estimating the uncapped acceptance length: no model, no extrapolation - #: from the untrimmed subpopulation (which the planner selected, so it - #: is not a fair sample of the trimmed one). + #: Trimmed requests that accepted their ENTIRE window: the draft was + #: still alive at the cut, so the trim certainly cost acceptance (one + #: that accepted fewer would have died on its own and cost nothing). self.trimmed_hit_ceiling = 0 - #: Positions the target accepted and the window then discarded -- - #: the exact acceptance the schedule cost. Only ``cap-accept`` can - #: populate this, because only there are those positions scored at all; - #: it stays 0 under ``compact``, where the honest answer is that the - #: number is unknowable without re-running. + #: Positions the target accepted and the window then discarded. Only + #: ``cap-accept`` scores these; under ``compact`` 0 means "unknowable", + #: not "no loss". self.cap_trim_tokens = 0 - #: The distribution, not just the total. "Every request gave up a - #: little" and "a few were cut to pieces" sum to the same - #: cap_trim_tokens and mean opposite things about the planner: the - #: first is a tuning knob, the second is the scheduler concentrating - #: its whole cost on a subpopulation it selected. A total cannot tell - #: them apart, which is the reason this is kept per request at all. + #: Kept per request because a broad mild trim and a concentrated one + #: sum to the same total and mean opposite things about the planner. self.requests_cap_trimmed = 0 self.cap_trim_max = 0 self.cap_trim_hist: Counter = Counter() - #: Graph replay is the whole point of landing on a bucket; an eager - #: step costs far more than the tokens it saved. + #: An eager step costs far more than the tokens the trim saved. self.graph_replays = 0 self.graph_eager = 0 @@ -265,35 +197,20 @@ def record_step( delivered: tokens actually handed to the target, when that is not derivable from ``verify_lens`` and ``bucket``. Required by ``cap-accept``, which schedules windows but submits the full - block: deriving it there would report the windows as if they - had been submitted, i.e. credit the mode with a compute saving - it deliberately does not take, and ``trim_ratio`` would then - describe a run that never happened. + block. """ self.steps_total += 1 - # Steps where raggedness is even definable. A window can only differ - # *between* requests, so a batch of one is uniform by construction -- - # counting it toward the gate's floor makes the gate fire during ramp-up - # and blame the scheduler for a batch that had not filled yet. + # A batch of one is uniform by construction; counting it toward the + # gate's floor would make the gate fire during ramp-up. if int(num_gen_requests) >= 2: self.steps_multi_request += 1 - # What a no-trim step would have submitted: every row sends its bonus - # token plus the full block. - # - # ROWS, not requests. When the step ran on a fitted bucket, `delivered` - # below books `bucket = padded_bs * (tier + 1)`, and `padded_bs` is the - # CUDA-graph-rounded maximum row count ACROSS RANKS -- so a rank holding - # fewer requests than its widest peer still submits the peer's rows. - # Booking those rows as delivered against a ceiling counted over this - # rank's real requests made trim_ratio negative in four separate runs - # (jobs 2560992 -0.001, 2561519 -0.0438, 2563581 -0.0082), and the - # assert_ragged_active gate then reported "nothing was saved", which is - # the opposite of what a negative ratio means. - # - # Only that one case rebases. Every other path books `delivered` over - # the real rows -- the no-window path and the cap-accept path each add - # exactly n*(1+K), and the bucket-less ragged path adds sum(1+v) <= - # n*(1+K) -- so widening their ceiling would understate the trim. + # Ceiling = what a no-trim step would have submitted: each ROW sends + # its bonus token plus the full block. ROWS, not requests: on a fitted + # bucket `delivered` books `padded_bs * (tier + 1)` where `padded_bs` + # is the CUDA-graph-rounded max row count ACROSS RANKS, so the ceiling + # must count the same rows or trim_ratio goes negative. Only that case + # rebases; every other path books `delivered` over the real rows, and + # widening their ceiling would understate the trim. ceiling_rows = int(num_gen_requests) if bucket and delivered is None and padded_bs: ceiling_rows = max(ceiling_rows, int(padded_bs)) @@ -336,17 +253,9 @@ def record_graph(self, *, replayed: bool, shape) -> None: """Record whether this step replayed a CUDA graph or ran eager. ``shape`` is the ``(num_rows, total_verify_tokens)`` the step actually - submitted, or None on a replay. It is kept only for misses, and only - the distinct ones: a miss means some shape was not captured, and - knowing *which* is the whole diagnosis. Counting misses without - recording the shape leaves only guessing or re-running with a profiler. - - Required rather than defaulted, deliberately. A run reported - ``graph_eager: 9`` alongside an empty ``graph_miss_shapes`` because the - caller was loaded from a build that predated this argument -- with a - default that reads as "no misses had shapes", which is indistinguishable - from "the probe works and found nothing". Without a default it is a - TypeError instead. + submitted, or None on a replay; kept only for misses. Deliberately + required, not defaulted: a default would make "probe not wired" + indistinguishable from "no misses had shapes". """ if replayed: self.graph_replays += 1 @@ -423,11 +332,9 @@ def accept_len(self) -> float: def accept_loss_per_request(self) -> float: """Mean accepted positions the schedule discarded, per request. - The exact counterpart to :attr:`accept_len`: add the two and you get - the acceptance a no-trim run would have achieved on the same drafts. - Meaningful only under ``cap-accept``; 0.0 elsewhere, where it means - "not measured", not "no loss" -- ``trim_regret_rate`` is the bound - available under ``compact``. + :attr:`accept_len` plus this is the no-trim acceptance on the same + drafts. Meaningful only under ``cap-accept``; elsewhere 0.0 means "not + measured", not "no loss". """ if not self.requests_scored: return 0.0 @@ -437,10 +344,8 @@ def accept_loss_per_request(self) -> float: def cap_trim_concentration(self) -> float: """Share of scored requests that lost anything at all. ``cap-accept``. - Read it together with :attr:`accept_loss_per_request`. The same mean - loss at 0.9 concentration is a broad, mild trim; at 0.02 it is the - scheduler paying for its throughput out of two percent of the traffic. - Only the second is a problem, and the mean alone cannot see it. + Distinguishes a broad mild trim from the same mean loss concentrated + on a few requests; only the second is a problem. """ if not self.requests_scored: return 0.0 @@ -450,11 +355,8 @@ def cap_trim_concentration(self) -> float: def trim_regret_rate(self) -> float: """Share of trimmed requests whose draft was still alive at the cut. - The operational reading: 0.0 means every trim was free -- those drafts - would have died anyway -- and the scheduler is taking compute back for - nothing. Climbing toward 1.0 means trimming is buying its throughput by - throwing away acceptance, which is the regression a delivered-only - acceptance metric cannot distinguish from a harder workload. + 0.0 means every trim was free (those drafts would have died anyway); + toward 1.0 the trims are buying throughput with acceptance. """ if not self.requests_trimmed: return 0.0 @@ -502,19 +404,16 @@ def log_summary(self, *, prefix: str = "DSpark ragged verify") -> None: def assert_ragged_active(self, *, require_trim: bool = True) -> None: """Fail loudly if the ragged path never actually ran. - This is the check that makes an accuracy run meaningful. Passing GSM8K - with ``compact`` selected proves nothing on its own -- the scheduler - may have declined on every step and the run may have been a plain - uniform baseline wearing a different config. + An accuracy pass alone proves nothing: the scheduler may have declined + every step. Args: - require_trim: also require that the delivered token count came in - under the no-trim ceiling. Leave it on unless deliberately - measuring a workload where every draft survives. + require_trim: also require that delivered tokens came in under the + no-trim ceiling. Leave on unless deliberately measuring a + workload where every draft survives. Raises: - AssertionError: with the full counter summary, so the reason is in - the failure rather than in a log the CI dropped. + AssertionError: with the full counter summary. """ summary = self.summary() if self.steps_total == 0: @@ -531,12 +430,9 @@ def assert_ragged_active(self, *, require_trim: bool = True) -> None: raise AssertionError( f"DSpark ragged verify handed out only one distinct window " f"size, so the batch was uniform in all but name: {summary}") - # The two counter families have to move together. A run once reported - # per-step scheduling activity here beside an all-zero planner block, - # because `decisions` was never incremented -- which made the fallback - # counters unreadable, and those are the only thing that says WHY a - # step declined to trim. Zero decisions against non-zero windowed steps - # means the planner's stats are not wired to this object at all. + # Zero planner decisions against non-zero windowed steps means the + # planner's stats are not wired to this object -- and its fallback + # counters are the only record of WHY a step declined to trim. if self.planner_stats is not None: decisions = int(self.planner_stats.get("decisions", 0)) windowed = self.steps_ragged + self.steps_uniform_windows @@ -546,26 +442,19 @@ def assert_ragged_active(self, *, require_trim: bool = True) -> None: f"but the planner recorded 0 decisions, so its fallback " f"counters cannot be read -- and they are the only record " f"of why a step declined to trim: {summary}") - # `cap-accept` submits the full block on purpose -- it measures the - # commitment policy, not a compute saving -- so a zero trim ratio is - # the expected outcome there, not a silent no-op. Requiring one would - # make the gate unsatisfiable in the one mode whose entire job is to be - # the trusted reference the other modes are compared against. + # `cap-accept` submits the full block on purpose, so a zero trim ratio + # is expected there, not a silent no-op. if (require_trim and self.mode.trims_submitted_tokens and self.trim_ratio <= 0.0): raise AssertionError( f"DSpark ragged verify delivered as many tokens as a no-trim " f"run, so nothing was saved: {summary}") - # Only the misses ragged is responsible for. A batch some rank filled - # with context requests falls out of replay on the uniform path too -- - # the uniform GSM8K run sat at ~1.7% eager for exactly that reason -- - # so failing on it makes this check unsatisfiable on any workload with - # continuous batching, which is all of them. The three below are - # different: "key_not_captured" and "bs_not_supported" mean the batch - # was graphable, every rank agreed, and the shape was still missing -- - # i.e. the token bucket grid is wrong. "peer_shape_mismatch" means one - # rank chose a different ragged bucket and dragged the whole world - # eager, which is the attention-DP divergence this feature has to avoid. + # Fail only on misses ragged is responsible for: context-filled batches + # fall out of replay on the uniform path too, and blaming those would + # make the check unsatisfiable under continuous batching. + # "key_not_captured"/"bs_not_supported" mean the batch was graphable + # and the bucket grid is wrong; "peer_shape_mismatch" means one rank + # chose a different ragged bucket and dragged the whole world eager. blamed = {"key_not_captured", "bs_not_supported", "peer_shape_mismatch"} attributable = { reason: n diff --git a/tensorrt_llm/_torch/speculative/dspark_planner.py b/tensorrt_llm/_torch/speculative/dspark_planner.py index 6afd63b2e039..17b4a100108c 100644 --- a/tensorrt_llm/_torch/speculative/dspark_planner.py +++ b/tensorrt_llm/_torch/speculative/dspark_planner.py @@ -14,51 +14,20 @@ # limitations under the License. """DSpark verification-budget planner: how many tokens are worth verifying. -The scheduler maximizes *system* goodput, not per-request speed:: +The scheduler maximizes *system* goodput:: tau(n) = num_gen_requests + sum of the n largest prefix-survival values Theta(n) = tau(n) / step_cost_ms(num_gen_requests + n) budget = argmax_n Theta(n) -``tau(n)`` is the expected number of tokens the step emits if it verifies ``n`` -tokens beyond the per-request floor: every generating request contributes one -bonus token for free, and each admitted draft position contributes its survival -probability. Dividing by the step cost turns that into expected tokens per -millisecond. - -Two things are easy to get wrong here: - -**Take the global argmax, not the first descent.** The marginal cost of a verify -token is a staircase -- flat while a token slots into the current kernel waves, -then a jump when it starts a new one -- so ``Theta`` is not unimodal. A greedy -loop that stops at the first non-improvement parks at the end of the first -shelf and leaves most of the win behind. Both SGLang and vLLM PR #47808 scan the -whole curve; so do we, and :mod:`tests` pins it against a brute-force scan. - -**Cost is per step, not per token.** ``SpsCostTable`` therefore stores measured -step time against the *total* verified token count, interpolated between -breakpoints (see the class docstring for why the old floor lookup was wrong). - -**Total verified tokens includes the bonus position.** A request that verifies -``L`` drafted positions submits ``L + 1`` tokens to the target -- the bonus -token it already holds, plus the ``L`` drafts. So a batch of ``bs`` requests at -uniform length ``L`` costs ``step_time(bs * (L + 1))``, not ``step_time(bs * L)``. -The two differ by a whole ``bs``, which is easily a shelf's width. - -**The cost is two-dimensional, and only one dimension is trimmable.** SGLang -publishes ``T(bs, K) = bias + alpha(bs) + theta(M)``: a fixed cost, a -batch-size-dependent cost (the draft pass, weight movement -- untouched by -trimming), and the target's verify-token cost ``theta(M)``, which is the only -term a smaller ``M`` reduces. ``Theta = tau / T`` is a *ratio*, so the -non-trimmable terms are not a constant that cancels -- getting them wrong moves -the argmax directly. They therefore have to be supplied, not guessed: -:class:`SpsCostTable` carries ``fixed_overhead_ms`` for ``bias`` and an optional -``batch_sizes`` / ``batch_overhead_ms`` staircase for ``alpha(bs)``. - -The planner runs on the host: its output selects a CUDA graph, and graph -selection is a host decision in TensorRT-LLM. Keeping it here (rather than on -device) is deliberate -- a device-side budget would have to be copied back -before the batch could be shaped, which is exactly the sync we are avoiding. +Invariants: take the GLOBAL argmax (marginal cost is a staircase, so ``Theta`` +is not unimodal and a first-descent stop parks on the first shelf); cost is +per step, indexed by the TOTAL verified token count including the bonus +position (a request verifying ``L`` drafts submits ``L + 1`` tokens); and +``Theta`` is a ratio, so the non-trimmable ``bias + alpha(bs)`` terms do not +cancel -- they must be supplied, not guessed (see :class:`SpsCostTable`). The +planner runs on the host because its output selects a CUDA graph, a host +decision. """ import math @@ -94,36 +63,21 @@ class SpsCostTable: """Measured decode step time as a function of total verified tokens. Attributes: - token_counts: strictly increasing total-token breakpoints. "Total" means - what :func:`total_verify_tokens` computes -- bonus tokens included. + token_counts: strictly increasing total-token breakpoints, in + :func:`total_verify_tokens` units (bonus tokens included). step_time_ms: measured ``theta(M)`` at each breakpoint. fixed_overhead_ms: the ``bias`` term -- per-step cost that trimming - cannot touch (sampling, bookkeeping, launch latency), if it is not - already inside ``step_time_ms``. Defaults to 0.0 on the assumption - that a profiled table measured whole steps; set it only when the - table isolates the verify term. + cannot touch -- if not already inside ``step_time_ms``. + Understating the non-trimmable terms makes the planner over-trim: + ``tau / T`` is a ratio, so they do not cancel. batch_sizes / batch_overhead_ms: optional ``alpha(bs)`` breakpoints -- - the batch-size-dependent, *non-trimmable* part of the step (the - draft pass, weight movement). Interpolated on ``batch_sizes``. - - Getting the non-trimmable terms wrong is not harmless. The planner maximizes - ``tau / T``, a ratio, so ``bias + alpha(bs)`` does not cancel: understating - it makes every verified token look proportionally more expensive than it is - and the planner over-trims. On a large-MoE deployment ``alpha(bs)`` is tens - of milliseconds while ``theta(M)`` is a few, so a hardcoded guess here would - dominate the decision. - - Lookup is a clamped LINEAR INTERPOLATION between breakpoints, matching the - additive-table consumer this was ported from (SGLang's - ``_additive_step_time_tensor``). A floor lookup here is wrong on a sparse - table: it bills every total below the next breakpoint at the previous - breakpoint's price, making tokens look free right before a riser, and the - argmax then over-spends -- measured live as the planner buying the full - block on ~95% of decisions because its cost ratio between tiers collapsed - below the survival ratio. A table that wants a flat shelf must MEASURE the shelf: - two breakpoints with equal values. Queries outside the measured range - clamp to the end values; the caller is expected to bound the budget so the - high clamp does not happen silently. + the batch-size-dependent, non-trimmable part of the step. + + Lookup is a clamped linear interpolation between breakpoints; a flat shelf + must be MEASURED as two breakpoints with equal values (a floor lookup makes + tokens look free right before a riser and the argmax over-spends). Queries + outside the measured range clamp to the end values; callers must bound the + budget so the high clamp does not happen silently. """ token_counts: Sequence[int] @@ -182,41 +136,27 @@ def step_times(self, num_tokens: np.ndarray, num_requests: int = 0) -> np.ndarra @classmethod def flat(cls, step_time_ms: float = 1.0) -> "SpsCostTable": - """A cost model with no staircase at all. - - Useful as an explicit "I have not profiled this machine" marker. With a - flat cost every extra token is free, so the planner will always spend the - entire budget -- i.e. it degenerates to verify-all. Callers should warn - rather than silently ship this. - """ + """Explicit "not profiled" marker: every extra token looks free, so + the planner degenerates to verify-all. Callers should warn rather than + silently ship this.""" return cls(token_counts=(0,), step_time_ms=(float(step_time_ms),)) @property def is_flat(self) -> bool: - """Whether the *trimmable* term is constant. - - Only ``theta(M)`` matters here: ``bias`` and ``alpha(bs)`` are the same - for every candidate length, so a table that varies only in those cannot - distinguish a cheap verified token from an expensive one either. - """ + """Whether the trimmable term ``theta(M)`` is constant. ``bias`` and + ``alpha(bs)`` are identical for every candidate length, so variation + there cannot distinguish lengths.""" return len(set(self.step_time_ms)) <= 1 def check_table_fingerprint(*, payload: dict, live: dict) -> None: """Refuse a cost table profiled on a different engine than this one. - The table's numbers are only meaningful for the engine configuration they - were measured on: the same (bs=64, M=384) cell measured 49.5 ms on a - max_batch_size=256 engine and 61.0 ms on a max_batch_size=64 engine -- 23% - apart against a 0.04-1.2% run-to-run noise floor -- and a table taken on - CUTLASS MoE kernels was nearly loaded to plan MegaMoE steps. Neither - mistake produces an error anywhere downstream; the planner just trims to - the wrong economics. - - Only the keys PRESENT in both dicts are compared, so a table may carry - facts the consumer cannot see (image hash, geometry) for humans to check, - and an old table with no fingerprint loads with a warning rather than - breaking existing deployments. + A mismatched table produces no error downstream -- the planner just trims + to the wrong economics. Only the keys PRESENT in both dicts are compared, + so a table may carry facts the consumer cannot see (for humans to check), + and a table with no fingerprint loads with a warning rather than breaking + existing deployments. """ fp = (payload.get("_meta") or {}).get("engine") or payload.get("engine") if not fp: @@ -275,20 +215,11 @@ def compute_verify_token_budget( cost_table: measured step cost; see :class:`SpsCostTable`. min_verify_len / max_verify_len: per-request bounds, mirroring :class:`~.dspark_schedule.DSparkScheduleConfig`. - allowed_lens: restrict the answer to budgets that a tier in this ladder - can actually realise, i.e. ``n in {bs*(t - min_verify_len)}``. Pass - it whenever the result will be executed, and omit it only to ask - what an unconstrained scheduler would have wanted. - - The unrestricted answer is not runnable. The captured graphs are - ``bs*(t+1)`` tokens for ``t`` in the ladder - (``model_engine.ragged_verify_token_buckets``), and any other total - is rounded UP to one of them -- so a budget chosen strictly inside a - cost shelf, which is the only reason to choose one, is pushed back - over the riser it was avoiding. The planner would then buy tokens - priced on the cheap shelf and run on the expensive one. Realising a - fine budget honestly would need a graph per token count, which is - tens of MB of metadata each, taken out of KV cache. + allowed_lens: restrict the answer to budgets a captured tier can + realise (``n in {bs*(t - min_verify_len)}``). Pass it whenever the + result will be executed: any non-tier total is rounded UP to a + captured bucket, pushing a budget chosen inside a cost shelf back + over the riser it was avoiding. Returns: ``n`` maximizing ``tau(n) / step_cost(floor_tokens + n)``, in ``[0, N]`` @@ -308,30 +239,21 @@ def compute_verify_token_budget( if candidates.size == 0: return 0 - # tau(n) = expected tokens emitted when n candidates are admitted: - # * one bonus token per generating request, emitted regardless; - # * the floor positions, which every request verifies for free -- their - # survival is part of the yield even though no budget buys them. Leaving - # this out shifts the whole curve and can move the argmax, since - # argmax (C + f(n)) / g(n) depends on C. - # * the n admitted candidates, best-first. + # tau(n): one bonus token per request, plus the floor positions' survival + # (unbought but part of the yield -- omitting the constant moves the argmax + # of a ratio), plus the n admitted candidates best-first. base = float(num_gen_requests) + float(surv[:, :min_verify_len].sum()) tau = base + np.concatenate(([0.0], np.cumsum(candidates))) - # Tokens actually submitted to the target at the floor: every request's - # bonus position plus its floor draft positions. Dropping the bonus term - # under-reports the batch by a whole ``bs``, which is easily a shelf wide. + # Includes the bonus position: dropping it under-reports the batch by a + # whole ``bs``. floor_tokens = total_verify_tokens(bs, min_verify_len) tokens = floor_tokens + np.arange(tau.size) theta = tau / cost_table.step_times(tokens, num_gen_requests) if allowed_lens is None: return int(np.argmax(theta)) - # Tier-aligned: same theta curve, evaluated only where the executor can - # land. tokens(n) at n = bs*(t - min_verify_len) is bs*(min+1) + bs*(t-min) - # = bs*(t+1), which is exactly total_verify_tokens(bs, t) -- the same cost - # grid budget_argmax_over_uniform_lens uses. The only thing that changes is - # the numerator: tau here is the sum of the top n survivals, which is what - # schedule_verify_lens_topk actually collects, rather than the sum of a - # uniform column, which nothing executes. + # Tier-aligned: the same theta curve, evaluated only at totals the executor + # can land on -- n = bs*(t - min_verify_len) gives tokens = bs*(t+1) + # = total_verify_tokens(bs, t). idx = sorted({ int(bs) * (int(t) - int(min_verify_len)) for t in allowed_lens @@ -353,48 +275,22 @@ def derive_verify_len_tiers( ) -> List[int]: """Derive the verify lengths worth capturing a CUDA graph for. - Restricting K to a handful of captured values sounds like it must cost - something versus a freely-chosen K. **At a fixed batch size, on a table - with genuine shelves, it does not**: - - Within one cost shelf the step time is constant while ``tau`` strictly - increases with each admitted token, so ``Theta = tau / cost`` is strictly - increasing across the shelf. The optimum therefore always sits at a - shelf's *right edge* -- never in its interior. For that batch size, a - tier set of the shelf right edges contains the continuous optimum - exactly. - - Since the consumer became an interpolation, "shelf" means a measured flat - segment (adjacent breakpoints with equal values), not the gap between two - sparse samples. On a table whose segments genuinely rise -- like the live - GB300 one -- the optimum can sit inside a segment and the right-edge set - is a slope-change heuristic rather than an exact cover; the residual loss - has to be measured. (The old floor consumer made every gap LOOK like a - shelf, which both made this theorem vacuously "exact" and priced a - 1512-token step at the 768-token price, disabling trimming entirely.) - - **The zero-loss property does not survive across batch sizes.** The cost - shelves live in token space, so a shelf's right edge in *length* space is - ``(breakpoint - 1) // batch_size - 1`` -- a function of the batch size. A - deployment captures graphs for many batch sizes but can only capture one - ladder, so a single tier set cannot sit on every batch size's right edges. - Deriving tiers is more robust than hardcoding a pair like ``{1, 5}``, but it - is still a lossy discretization once the batch size moves; the residual loss - has to be measured, not assumed away. Callers that capture one ladder should - derive it at the steady-state modal batch size and check the worst-case loss - across the others. - - This returns the right-edge length for each shelf the given batch size can - reach, always including ``min_verify_len`` and the full block. + Within a measured flat cost shelf (adjacent breakpoints with equal values) + ``Theta = tau / cost`` strictly increases, so at a fixed batch size the + optimum sits on a shelf's *right edge*, never in its interior; this + returns those edges, always including ``min_verify_len`` and the full + block. On segments that genuinely rise the right-edge set is a + slope-change heuristic, and the property does not survive across batch + sizes (shelves live in token space) -- derive at the modal batch size and + measure the residual loss at the others. Args: - num_requests: batch size the tiers are derived for. Tiers are a function - of it (total tokens = ``num_requests * length``), so a deployment - that captures several batch sizes derives a set per batch size. - max_tiers: cap on how many lengths to return. Each extra tier is another - captured graph, and captured graphs consume memory that would - otherwise be KV cache -- so this is a real budget, not a formality. - The full block and ``min_verify_len`` are kept first. + num_requests: batch size the tiers are derived for; tiers are a + function of it, so a deployment that captures several batch sizes + derives a set per batch size. + max_tiers: cap on how many lengths to return. Each extra tier is + another captured graph, which consumes memory that would otherwise + be KV cache. The full block and ``min_verify_len`` are kept first. """ lo, hi = int(min_verify_len), int(block_size) if num_requests <= 0 or hi < lo: @@ -402,10 +298,9 @@ def derive_verify_len_tiers( edges = set() for breakpoint_tokens in cost_table.token_counts[1:]: - # Largest length whose total token count still sits on the shelf below - # this breakpoint. Total tokens for length L is num_requests * (L + 1) - # (see total_verify_tokens), so invert that, not num_requests * L -- - # otherwise the derived edge sits one shelf too far right. + # Largest length whose total (num_requests * (L + 1), see + # total_verify_tokens) still sits below this breakpoint; inverting + # num_requests * L instead puts the edge one shelf too far right. length = (int(breakpoint_tokens) - 1) // int(num_requests) - 1 if lo <= length <= hi: edges.add(length) @@ -416,8 +311,7 @@ def derive_verify_len_tiers( tiers = sorted(edges) if len(tiers) <= max_tiers: return tiers - # Keep the endpoints, then spread the remaining slots over the interior - # edges rather than clustering them at one end of the curve. + # Keep the endpoints; spread the remaining slots over the interior edges. interior = tiers[1:-1] step = len(interior) / float(max_tiers - 2) if max_tiers > 2 else 0 keep = [interior[int(i * step)] for i in range(max_tiers - 2)] if max_tiers > 2 else [] @@ -434,11 +328,9 @@ def budget_argmax_over_uniform_lens( ) -> int: """Pick the best verify length from a discrete set (e.g. captured graph tiers). - TensorRT-LLM can only run one draft length per step, and only lengths that - have a captured CUDA graph. Optimizing over the continuous budget and then - rounding is not the same as optimizing over what is actually runnable: the - rounded-down point can sit on the far side of a cost riser. This evaluates - ``Theta`` directly at each achievable length and returns the best one. + Evaluates ``Theta`` directly at each runnable length: optimizing the + continuous budget and rounding is not equivalent, since the rounded point + can land on the far side of a cost riser. Args: allowed_lens: verify lengths with a captured graph. Values below @@ -459,9 +351,8 @@ def budget_argmax_over_uniform_lens( surv = np.asarray(survival, dtype=np.float64) best_len, best_theta = lens[0], -np.inf for length in lens: - # Every request verifies ``length`` positions: the expected yield is the - # bonus token plus the survival of *every* verified position (including - # the floor ones), summed over the batch. + # Expected yield: the bonus token plus survival of every verified + # position, floor included. tau = float(num_gen_requests) + float(surv[:, :length].sum()) theta = tau / cost_table.step_time(total_verify_tokens(bs, length), num_gen_requests) if theta > best_theta: diff --git a/tensorrt_llm/_torch/speculative/dspark_ragged.py b/tensorrt_llm/_torch/speculative/dspark_ragged.py index c977e14cb8fe..255fcfae6b04 100644 --- a/tensorrt_llm/_torch/speculative/dspark_ragged.py +++ b/tensorrt_llm/_torch/speculative/dspark_ragged.py @@ -12,54 +12,16 @@ # 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. - -Uniform-K scheduling picks one verify length for the whole batch. That leaves -real budget on the table, because requests differ a lot in how deep their draft -survives: giving a request with survival ``0.55, 0.22, 0.07`` the same window as -one sitting at ``0.98, 0.95, 0.91`` wastes verification on the first and starves -the second. Ragged verification gives each request its own length. - -The packing follows the shape every variable-length attention path already uses: -a flat token axis plus an index pointer. - - verify_lens [bs] per-request number of verified positions - qo_indptr [bs + 1] exclusive prefix sum; request r owns - [qo_indptr[r], qo_indptr[r + 1]) - extend_start_loc [bs] == qo_indptr[:-1], the offset form some - attention backends want directly - -Two invariants make the rest of the system safe: - -* ``verify_lens >= 1`` for every request. Position 0 carries the bonus/anchor - token; a request that verifies nothing makes no progress at all and would - stall forever. -* ``sum(verify_lens) == graph_num_tokens`` after :meth:`fill_bucket`. The batch - is padded up to a captured CUDA-graph token bucket, so the packed shape is one - of a small fixed set even though the per-request split is arbitrary. That is - what lets a graph capture a ragged batch: the *shape* is constant, the - raggedness lives in the *contents* of ``verify_lens`` / ``qo_indptr``. - -Why the token count has to land exactly on the bucket ------------------------------------------------------ -``seq_lens.sum()`` becomes ``attn_metadata.num_tokens``, which is what gets -all-gathered across attention-DP ranks and drives the MoE's chunk count. If the -packed token count under-reports the rows actually in the batch, attention and -the MoE disagree about how many tokens are in flight -- and nothing raises. - -Reclaiming the bucket padding ------------------------------ -Rounding up to a bucket would normally waste the difference. :meth:`fill_bucket` -spends it in two stages: real requests first (up to ``max_verify_len``, so those -tokens verify real draft positions the step is already paying for), then pad rows -added to reach the captured batch size. Same idea as SGLang's -``align_verify_tokens_to_graph_tier``. - -The real requests cannot always absorb everything -- a request verifies at most -``max_verify_len`` positions, so the batch's capacity is -``padded_bs * max_verify_len``. When the bucket exceeds what the real requests -can take, the remainder must go to pad rows; when it exceeds even the padded -capacity, :meth:`fill_bucket` raises rather than returning a short batch. +"""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. """ import bisect @@ -85,10 +47,8 @@ def round_up_to_bucket(total: int, buckets: Sequence[int]) -> int: """Smallest captured bucket that fits ``total``. - Raises rather than clamping: a batch that exceeds the largest captured - bucket has no graph, and silently running it eager costs far more than the - trimming saves. Callers must reject such a batch *before* selecting a - layout (see :func:`exceeds_captured_buckets`). + Raises rather than clamping; callers must reject an oversized batch before + selecting a layout (see :func:`exceeds_captured_buckets`). """ if not buckets: raise ValueError("round_up_to_bucket requires a non-empty bucket list") @@ -159,10 +119,9 @@ def from_verify_lens( ) -> "RaggedVerifyLayout": """Pack ``[bs]`` per-request lengths, rounding up to a captured bucket. - ``graph_num_tokens`` may be given directly (the usual path -- the caller - already knows the bucket because it also had to pick the graph), or - derived from ``buckets`` plus a host-side ``total_verify_tokens``. - Deriving it from the device tensor would need a sync. + Give ``graph_num_tokens`` directly, or ``buckets`` plus a host-side + ``total_verify_tokens``; deriving the total from the device tensor + would sync. """ lens = verify_lens.to(torch.int32) if graph_num_tokens is None: @@ -186,28 +145,12 @@ def from_verify_lens( def fill_bucket( self, *, max_verify_len: int, padded_bs: Optional[int] = None ) -> "RaggedVerifyLayout": - """Pad the batch to the captured shape, spending the slack on real work. - - Two things have to line up before the batch can run: - - * the row count must equal the captured batch size (``padded_bs``), and - * the token count must equal the captured bucket - (``graph_num_tokens``) -- because ``seq_lens.sum()`` becomes - ``attn_metadata.num_tokens``, which is what gets all-gathered and - drives the MoE's chunk count. If it under-reports, attention and MoE - disagree about how many tokens are in flight. - - So the slack is distributed in two stages: - - 1. **Real requests first**, up to ``max_verify_len``. The step is - already paying for the bucket, so these tokens verify real draft - positions instead of nothing. - 2. **Pad rows** take whatever is left. They exist only to reach - ``padded_bs``; every pad row needs at least one token because an - empty range breaks ``qo_indptr``'s per-row slicing. - - Raises when the bucket cannot be hit exactly -- silently returning a - short batch is what would desynchronize attention from the MoE. + """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") @@ -238,11 +181,8 @@ def fill_bucket( lens = real + [1] * n_pad spare = self.graph_num_tokens - baseline - # Round-robin over real rows first -- spreading keeps the extra - # verification where survival is still plausible for several requests - # instead of pushing one request deep into its low-survival tail. Only - # once every real row is saturated does the remainder go to pad rows, - # where it is genuinely wasted. + # 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)): @@ -268,15 +208,12 @@ def validate( max_verify_len: Optional[int] = None, exact_fill: bool = False, ) -> None: - """Assert the invariants the rest of the pipeline relies on. + """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. Check it on anything - heading for a captured graph: a short batch does not raise anywhere - downstream, it just makes ``seq_lens.sum()`` disagree with the row count - the graph was captured for. - - Host-syncing, so this is for tests and debug builds -- not the hot path. + 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: @@ -358,40 +295,24 @@ def choose_ragged_capture_shape( ) -> RaggedCaptureShape: """Pick the captured ``(padded_bs, bucket)`` for a ragged batch. - Uniform speculation needs one number -- the batch size -- because the token - count follows from it (``bs * (K + 1)``). Ragged verification breaks that - link: two batches with the same row count can hold very different token - totals, so the captured shape is genuinely two-dimensional and both axes - have to be agreed on before the graph is looked up. - Under attention DP every rank must land on the *same* shape or the ranks - replay different graphs and their collectives diverge. Ranks do not have - equal batches, so agreement cannot come from each rank rounding its own - numbers; it comes from rounding a reduction. Pass every rank's + replay different graphs and their collectives diverge, so agreement comes + from rounding a reduction: pass every rank's ``(num_real_requests, total_verify_tokens)`` as ``peer_stats`` (including - this rank's own) and each rank computes an identical answer: + this rank's own) and each rank computes an identical padded_bs = round_up_bs(max_r num_real_r) - bucket = round_up_bucket(padded_bs + max_r slack_r) + bucket = round_up_bucket(padded_bs + max_r (total_r - num_real_r)) where ``slack_r = total_r - num_real_r`` is the drafted-position count on - rank ``r``. The ``padded_bs +`` term is the floor every row contributes - (each request verifies at least its bonus position), so the bucket is large - enough for the widest rank's drafts on top of the widest rank's rows. - - Raises when no captured bucket fits, rather than clamping: a batch with no - graph runs eager, which costs far more than the trimming saves, so the - caller has to see it and shrink the batch instead. + rank ``r`` and the ``padded_bs +`` term is the one-token-per-row floor. + Raises when no captured bucket fits; the caller must shrink the batch. With ``max_verify_len`` given, the answer is additionally required to be DECOMPOSABLE BY EVERY RANK (see ``_bucket_decomposable``): the smallest - such bucket is chosen, walking past candidates any rank cannot realise. - This makes accept/decline a pure function of the allgathered payload -- - all ranks pick the same bucket or all raise together -- where previously - the per-rank decomposition check ran locally and a skewed rank declined - alone, splitting the group's graph keys. The cost is at most one bucket - step of extra verification on the steps that used to diverge, spent on - real tokens by ``fill_bucket``. + such bucket is chosen, walking past candidates any rank cannot realise, + so accept/decline is a pure function of the allgathered payload -- all + ranks pick the same bucket or all raise together. """ if not bs_buckets: raise ValueError("choose_ragged_capture_shape requires bs_buckets") @@ -403,9 +324,7 @@ def choose_ragged_capture_shape( ) stats = list(peer_stats) if peer_stats else [(num_real_requests, total_verify_tokens)] - # Tolerate extra trailing fields: the caller's payload has grown before - # (the group's graph-capability answer rides along with the row and token - # counts) and only the first two are this function's business. + # Peer entries may carry extra trailing fields; only the first two are read. max_real = max(int(peer[0]) for peer in stats) max_slack = max(int(peer[1]) - int(peer[0]) for peer in stats) @@ -445,11 +364,9 @@ def ragged_gather_index_lists( """Row/column index lists for gathering a ragged block from a ``[num_slots, max_width]`` tensor. - ``slots[i]`` contributes ``counts[i]`` entries taken from columns - ``0..counts[i]-1``, concatenated in batch order. This is the ragged - replacement for the ``tensor[slots, :width]`` strided gather the overlap - scheduler uses when every request verifies the same number of positions; - with a constant ``counts`` the two produce the identical flat sequence. + ``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( @@ -468,16 +385,9 @@ def ragged_gather_index_lists( def row_ids_from_lens(verify_lens: torch.Tensor, *, total: int) -> torch.Tensor: """``[bs]`` lengths -> ``[total]`` owning-request id for each packed token. - ``repeat_interleave`` with a *tensor* of repeats is what makes the packing - ragged; the uniform paths use a scalar repeat and get a fixed stride. - - ``total`` (``sum(verify_lens)``) is mandatory rather than derived. With a - device-resident ``repeats`` and no ``output_size``, torch reads the - cumulative sum back to the host to size its output -- a sync that is - outright *illegal* here, because DSpark is a one-engine drafter whose - acceptance runs inside the target's captured CUDA graph. The callers all - know the total on the host already (``spec_metadata.total_verify_tokens``, - or the packed buffer's own leading dimension). + ``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), @@ -496,15 +406,9 @@ def scatter_ragged_to_padded( ) -> torch.Tensor: """Unpack a flat ragged batch into ``[bs, max_len, ...]`` with padding. - The verify path is packed (one flat token axis) because that is what the - attention kernels want, but the acceptance logic is naturally rectangular - (compare draft vs target position by position). This bridges the two - without a host sync: token ``t`` belongs to request ``row_ids[t]`` at column - ``t - qo_indptr[row_ids[t]]``. - - ``flat.shape[0]`` is the packed token total, so it is exactly the - ``output_size`` :func:`row_ids_from_lens` needs to stay sync-free (and - therefore capturable). + 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]) @@ -521,14 +425,10 @@ def fill_padded_rows_onehot( ) -> torch.Tensor: """Make ``[bs, max_len, vocab]`` padding rows a valid one-hot distribution. - Scattering a ragged batch into a rectangle leaves all-zero probability rows - past each request's window. Those rows are never read -- the accepted count - is clamped to the window -- but a rejection-sampling kernel still walks them - when every real position was accepted, and sampling from an all-zero - distribution can index out of bounds. Putting all the mass on token 0 keeps - the kernel well-defined without affecting any position that is read. - - Modifies ``probs`` in place and returns it. + Padding rows are never read, but a rejection-sampling kernel still walks + them and sampling an all-zero row can index out of bounds; all mass goes + to token 0, which affects no position that is read. Modifies ``probs`` in + place and returns it. """ positions = torch.arange(probs.shape[1], device=probs.device) invalid = positions.unsqueeze(0) >= verify_lens.to(positions.dtype).unsqueeze(1) @@ -544,18 +444,14 @@ def count_accepted_ragged( ) -> 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 at the same - positions. - verify_lens: ``[bs]`` how many positions each request actually verified. - - Positions at or beyond a request's ``verify_len`` were never sent to the - target, so whatever sits there is stale padding. Masking them to "no match" - before the cumulative product makes the run of matches stop at the window - edge -- without the mask a stale slot could compare equal and credit an - acceptance the target never made, which is a silent correctness bug rather - than a throughput one. + 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( diff --git a/tensorrt_llm/_torch/speculative/dspark_schedule.py b/tensorrt_llm/_torch/speculative/dspark_schedule.py index eaba8be7ad3e..b025f82c3382 100644 --- a/tensorrt_llm/_torch/speculative/dspark_schedule.py +++ b/tensorrt_llm/_torch/speculative/dspark_schedule.py @@ -12,33 +12,17 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""DSpark confidence-scheduled verification: survival + per-request allocation. - -Given the draft's calibrated per-position acceptance probabilities, decide how -many of the proposed tokens are worth sending to the target for verification. - -The algorithm (DSpark, arXiv:2607.05147 §5; cross-checked against SGLang's -``dspark_schedule.py`` and vLLM PR #47808's ``adaptive_verification.py``): - - 1. ``survival[r][j] = prod(conf[r][:j+1])`` -- the probability that draft - position ``j`` of request ``r`` is reached *and* accepted. The confidence - head predicts a *conditional* acceptance probability, which is exactly what - makes this cumulative product the right prefix statistic. - 2. Rank every ``(request, position)`` candidate by survival, globally across - the batch, and admit the best ``budget`` of them. - -Step 2 needs no explicit prefix constraint: ``survival`` is non-increasing along -each request's positions, so any global top-k is automatically a per-request -*prefix*. That is why the allocation only has to *count* admitted candidates per -row rather than track which ones they were. - -Nothing here decides acceptance -- only how many tokens get verified -- so -scheduling is lossless with respect to the target distribution. - -Determinism note: ties are broken by ``(position, request)``, never by value. -Two TP ranks that compute bitwise-different confidences must still choose the -same verify lengths, or their batch shapes (and therefore their collectives) -diverge. See :func:`schedule_verify_lens_topk`. +"""DSpark confidence-scheduled verification (arXiv:2607.05147 §5). + +``survival[r][j] = prod(conf[r][:j+1])`` ranks every ``(request, position)`` +candidate globally; the best ``budget`` are admitted. Survival is +non-increasing along each request's positions, so any global top-k is +automatically a per-request prefix and the allocation only counts admitted +candidates per row. Only the number of verified tokens is decided here, never +acceptance, so scheduling is lossless w.r.t. the target distribution. Ties are +broken by ``(position, request)``, never by value: ranks with +bitwise-different confidences must still choose identical lengths or their +batch shapes (and collectives) diverge. """ from dataclasses import dataclass @@ -52,12 +36,10 @@ ] -# Raw logit written into confidence rows that carry no measurement (slot-free -# fill, never-drafted slots): sigmoid(30/T) ~ 1.0 at any fitted temperature, so -# an unknown row is treated as "certainly accept" -- optimistic by design, the -# same convention as SGLang's ones-fill for stale rows. Shared from here -# because both the worker (writes it) and the planner (counts it) need it and -# they cannot import each other. +# Raw logit for confidence rows that carry no measurement: sigmoid(30/T) ~ 1.0 +# at any fitted temperature, so an unknown row is treated as "certainly accept" +# -- optimistic by design. Lives here because the worker (writes it) and the +# planner (counts it) cannot import each other. NEUTRAL_CONFIDENCE_LOGIT = 30.0 @@ -66,18 +48,14 @@ class DSparkScheduleConfig: """Bounds for the per-request verify length. Attributes: - block_size: draft block length ``K`` (the number of speculative tokens - the draft proposes per step). - min_verify_len: floor on every request's verify length. Must stay >= 1: - position 0 carries the bonus/anchor token, and a request that - verifies nothing makes no progress at all. Budget is allocated + block_size: draft block length ``K``. + min_verify_len: floor on every request's verify length; must stay >= 1 + (position 0 carries the bonus/anchor token). Budget is allocated *above* this floor. max_verify_len: cap on every request's verify length; 0 means - ``block_size``. Never exceeds ``block_size`` -- a request cannot - verify more tokens than the draft proposed. - survival_eps: candidates whose survival falls below this are dropped - before ranking. Purely a numerical floor, not a tuning knob: it - keeps hopeless positions out of the ordering. + ``block_size``; never exceeds ``block_size``. + survival_eps: numerical floor (not a tuning knob); candidates below it + are dropped before ranking. """ block_size: int @@ -143,10 +121,9 @@ def schedule_verify_lens_topk( ``[min_verify_len, resolved_max_verify_len]``, summing to at most ``bs * min_verify_len + budget``. - Determinism: ranking uses a value-independent tie-break, so equal survivals - always resolve the same way regardless of float noise or backend. Positions - are ordered before requests, which biases ties toward *earlier* positions -- - the ones every downstream prefix depends on. + Determinism: ranking uses a value-independent tie-break (position before + request), so equal survivals resolve identically regardless of float noise + or backend. """ if survival.dim() != 2: raise ValueError(f"survival must be [bs, K], got shape {tuple(survival.shape)}") @@ -164,22 +141,16 @@ def schedule_verify_lens_topk( if bs == 0 or schedulable <= 0 or budget <= 0: return verify_lens - # Only positions above the floor are up for grabs: positions [0, floor) are - # already granted to every request, so scoring them would let a request win - # budget for something it is getting for free. + # Positions [0, floor) are already granted to every request; only those + # above compete for budget. candidates = survival[:, floor : floor + schedulable].to(torch.float32) flat = candidates.reshape(-1) eligible = flat >= cfg.survival_eps - # Rank by survival, then break ties deterministically. Sorting the tie-break - # key first and using a *stable* sort on the values makes the final order a - # pure function of (survival values, position, request) -- no dependence on - # the sort implementation's handling of equal keys. - # - # Ineligible candidates are scored -1 so they sort below every eligible one; - # the ``& eligible`` below then drops them even when ``budget`` is larger - # than the eligible count. That keeps the whole path free of any device->host - # sync (no ``.item()``), so it stays usable inside a captured graph. + # Sort the tie-break key first, then stable-sort on the scores: the final + # order is a pure function of (survival, position, request). Ineligible + # candidates score -1 and are masked by ``& eligible`` below; no + # ``.item()``/host sync anywhere, so the path stays capture-safe. scores = torch.where(eligible, flat, torch.full_like(flat, -1.0)) tie_break = torch.arange(schedulable, device=device).repeat(bs) * bs + torch.arange( bs, device=device diff --git a/tensorrt_llm/_torch/speculative/dspark_sts.py b/tensorrt_llm/_torch/speculative/dspark_sts.py index f7c2424db524..c910a5cc0c2d 100644 --- a/tensorrt_llm/_torch/speculative/dspark_sts.py +++ b/tensorrt_llm/_torch/speculative/dspark_sts.py @@ -14,29 +14,12 @@ # limitations under the License. """STS calibration: collecting the data, and reading the fitted table. -The confidence head emits a raw logit per drafted position. ``apply_sts`` turns -it into a probability with a *per-position* temperature:: - - confidence[r][j] = sigmoid(logit[r][j] / T[j]) - survival[r][j] = prod_{i <= j} confidence[r][i] - -``T`` defaults to all ones, which makes that a plain sigmoid -- i.e. no -calibration at all. That default is not neutral, because the planner does not -consume ``confidence``; it consumes the **cumulative product**. A per-position -bias of x compounds to x^(j+1) by position j, so a head that is 5% over-confident -per position over-states the depth-5 survival by more than 30%, and the budget -argmax divides by a cost while multiplying that error into its numerator. - -So the table has to be fitted against the survival, not against the per-position -probability -- which is what makes the fitting objective (see -``tests/microbenchmarks/dspark_fit_sts.py``) a cumulative-product ECE rather than -a per-position one. The recorder here produces exactly the pairs that objective -needs. - -Collection is opt-in via ``TLLM_DSPARK_STS_COLLECT_PATH`` and is **not** a -serving path: it reads the confidence buffer back to the host once per decode -step. That is a real synchronization, deliberately accepted because a -calibration run is not a throughput run. +``apply_sts`` computes ``confidence[r][j] = sigmoid(logit[r][j] / T[j])``. The +planner consumes the cumulative product ``survival[r][j]``, in which +per-position calibration error compounds geometrically, so the table must be +fitted against the survival (see ``tests/microbenchmarks/dspark_fit_sts.py``). +Collection is opt-in via ``TLLM_DSPARK_STS_COLLECT_PATH`` and is not a serving +path: it reads the confidence buffer back to the host once per decode step. """ import json @@ -77,9 +60,8 @@ class DSparkStsCalibration: num_samples: int = 0 def to_json(self) -> dict: - # Both key spellings are written. ``sts_temperatures`` is what this repo - # has always read; ``temperatures`` is what SGLang's DSparkStsCalibration - # emits, so a table produced here is loadable there and vice versa. + # Both key spellings on purpose: this repo reads ``sts_temperatures``, + # SGLang uses ``temperatures``; tables stay interchangeable. return { "sts_temperatures": list(self.temperatures), "temperatures": list(self.temperatures), @@ -93,10 +75,8 @@ def to_json(self) -> dict: def load_sts_temperatures_from_path(path: str) -> List[float]: """Read a temperature vector, accepting either spelling of the key. - This repo writes ``sts_temperatures``; SGLang's ``DSparkStsCalibration`` - writes ``temperatures``. The vectors are interchangeable -- both are fitted - against the cumulative product with the same sigmoid -- so refusing the - other spelling would reject a usable table for no reason. + This repo writes ``sts_temperatures``; SGLang writes ``temperatures``. + The vectors are interchangeable. """ with open(path, encoding="utf-8") as handle: payload = json.load(handle) @@ -118,10 +98,8 @@ def resolve_confidence_head(draft_model): Two layouts exist: the bare ``DSparkDraftModel`` (stages under ``.mtp_layers``, head on the last stage) and the ``DSparkForCausalLM`` - wrapper (bare model under ``.dspark_model``). A chained ``getattr`` with - soft defaults conflated "wrapper" with "no head" and silently disabled - calibration for an entire measurement campaign; this helper exists so the - resolution is written once and unit-tested against both layouts. + wrapper (bare model under ``.dspark_model``). Written once and unit-tested + against both; returns None when no head is found. """ inner = getattr(draft_model, "dspark_model", draft_model) stages = getattr(inner, "mtp_layers", None) @@ -133,30 +111,18 @@ def resolve_confidence_head(draft_model): class DSparkStsRecorder: """Pair each drafted block's logits with how much of it was accepted. - The two halves live in different places and on different clocks: the - logits sit in the worker's persistent slot-indexed device buffer, written - by draft pass ``i``; the accepted count arrives in the sampler while pass - ``i+1`` (or ``i+2``, under the overlap scheduler) has already overwritten - that buffer. The join is therefore made on TWO axes, neither of which is - wall-clock arrival: - - * **time** -- ``stage_snapshot`` keeps a small ring of host copies keyed - by the worker's draft sequence counter, and ``record`` selects the ring - entry whose per-row stamp equals the pass that drafted the block being - verified; - * **row** -- the caller resolves the request's buffer row through the - worker's own allocator (``confidence_row_for``), never through - ``py_seq_slot``, which belongs to a different allocator and only - coincides until the first request completes. - - A pair that cannot be made exactly is dropped and counted in ``stats`` - rather than approximated: a mispaired shard is strictly worse than a - smaller one, and looks identical downstream. + The confidence buffer is overwritten by later draft passes before the + accepted count arrives, so the join is made on two axes, never wall-clock + arrival: by draft-pass stamp (``stage_snapshot`` keeps a ring of host + copies keyed by the worker's draft sequence counter) and by buffer row + resolved through the worker's own allocator (``confidence_row_for`` -- + never ``py_seq_slot``, which belongs to a different allocator). A pair + that cannot be made exactly is dropped and counted in ``stats``, never + approximated. """ - #: Ring depth. The pairing lag is one draft pass (two under the overlap - #: scheduler); four snapshots is comfortably past both without letting a - #: wrong target_seq silently match ancient content. + #: Pairing lag is one draft pass (two under the overlap scheduler); four + #: clears both without letting a wrong target_seq match ancient content. _RING_DEPTH = 4 def __init__(self, *, path_stem: str, block_size: int, rank: int = 0, @@ -170,19 +136,12 @@ def __init__(self, *, path_stem: str, block_size: int, rank: int = 0, self._shard = 0 self._steps_since_flush = 0 # Snapshot ring, indexed by staged_seq % _RING_DEPTH. Each entry is - # (host_logits, host_stamps, cuda_event, staged_seq): a copy of the - # whole slot-indexed confidence buffer plus the draft-pass stamp of - # every row, taken on the host path each step -- the same relay the - # planner trusts. `record` selects the entry whose stamp for the - # request's row equals the pass that drafted the block being verified, - # so the pairing is by draft-pass identity, not by wall-clock arrival: - # the single mutable stash this replaces was overwritten by the very - # forward that produced the label, and NO execution order paired it - # correctly (off by one plain, off by two under overlap). + # (host_logits, host_stamps, cuda_event, staged_seq). `record` selects + # the entry whose stamp for the request's row equals the pass that + # drafted the block being verified: pairing is by draft-pass identity, + # never wall-clock arrival. self._ring: List[Optional[tuple]] = [None] * self._RING_DEPTH - # Every decline is counted: this feature's failures are all silent, - # and an empty or thin shard with healthy-looking stats was exactly - # how the mispaired fit shipped. + # Every decline is counted; a thin shard must explain itself. self.stats: dict = {"recorded": 0, "no_snapshot": 0, "stale_stamp": 0, "row_out_of_range": 0, "no_row": 0, "snapshots_staged": 0} @@ -192,10 +151,9 @@ def stage_snapshot(self, *, device_logits: torch.Tensor, staged_seq: Optional[int]) -> None: """Snapshot the confidence buffer + stamps, keyed by draft pass. - Called from the executor's host path once per step (next to the - planner's own staging), never from inside the captured graph: Python - in a captured region runs at capture time only, which is how the old - stash silently replayed stale rows on every graphed step. + Must be called from the executor's host path once per step, never from + inside the captured graph: Python in a captured region runs at capture + time only. """ if device_stamps is None or staged_seq is None: return @@ -228,17 +186,12 @@ def record(self, *, row: Optional[int], accepted: int, Args: row: the request's row in the confidence buffer, resolved through - the worker's own allocator (``confidence_row_for``). The - previous join key -- ``py_seq_slot`` -- came from a DIFFERENT - allocator that only coincides until the first request - completes; every pair after roster churn was request A's - label against request B's logits, and the shards looked - healthy. + the worker's own allocator (``confidence_row_for``), never + ``py_seq_slot``, which belongs to a different allocator. accepted: drafted positions accepted, excluding the bonus token. target_seq: the draft pass that produced the block this label verifies, snapshotted into the SampleState at sampling time - (the ``verify_lens_snapshot`` pattern; reading it live at - update time rewinds wrong under the overlap scheduler). + (reading it live rewinds wrong under the overlap scheduler). """ if row is None: self.stats["no_row"] += 1 @@ -257,10 +210,8 @@ def record(self, *, row: Optional[int], accepted: int, if event is not None: event.synchronize() if int(stamps_host[int(row)]) != int(target_seq): - # The buffer content for this row is from a different draft pass - # than the one this label verifies -- the ring slot was reused, or - # the row was never written for that pass. Refusing IS the fix: - # appending it anyway is the mispairing this design replaces. + # This row's content is from a different draft pass than the label + # verifies; refusing is the point -- appending would mispair. self.stats["stale_stamp"] += 1 return self._logits.append(logits_host[int(row)].clone()) @@ -288,10 +239,8 @@ def flush(self) -> None: path = Path(f"{self.path_stem}.r{self.rank}.{self._shard}.pt") path.parent.mkdir(parents=True, exist_ok=True) - # The pairing marker lets the fitter reject shards from the pre-ring - # recorder, whose pairs were mislabeled in ways the tensors themselves - # cannot reveal. The running decline counters ride along so a thin - # shard explains itself. + # The pairing marker lets the fitter reject mispaired legacy shards; + # the decline counters ride along so a thin shard explains itself. payload = { "logits": logits, "prefix_mask": prefix_mask, @@ -316,13 +265,8 @@ def make_recorder_from_env(*, block_size: int, rank: int = 0, ) -> Optional[DSparkStsRecorder]: """Build a recorder iff ``TLLM_DSPARK_STS_COLLECT_PATH`` is set. - Refuses the two regimes in which the collected pairs describe the scheduler - rather than the head. Both are silent failures otherwise: the run completes, - the shards look normal, and the fitted temperatures are wrong in a way no - downstream check can see. SGLang guards the first of these next to its own - collector (dspark_planner.py:100-108) and names the second in the sibling - ConfidenceMetricsProbe ("padded verify rows corrupt the per-position prefix - label", dspark_observability.py:688-697). + Refuses the two regimes in which the collected pairs describe the + scheduler rather than the head; both would otherwise fail silently. """ stem = os.environ.get(STS_COLLECT_ENV, "").strip() if not stem: @@ -336,11 +280,6 @@ def make_recorder_from_env(*, block_size: int, rank: int = 0, except ValueError: forced_full = False if has_cost_table and not forced_full: - # A loaded table means the planner trims, and a trimmed window clips - # `py_num_accepted_draft_tokens`: a request granted 2 positions can show - # at most 2 accepted, however well the head predicted. Fitting that - # calibrates the scheduler. Pinning the window to the full block makes - # the label uncensored by construction. raise ValueError( f"{STS_COLLECT_ENV} is set while a cost table is loaded. The planner " f"will trim, and a trimmed window censors the acceptance label -- the " @@ -350,10 +289,6 @@ def make_recorder_from_env(*, block_size: int, rank: int = 0, f"block still censors the label and is refused).") if (ragged_mode or "").lower() == "compact": - # Compact packs the token axis and fills the bucket slack with padding - # rows. Those rows carry positions that were never drafted for a real - # request, so the per-position prefix label they contribute is not a - # measurement of anything. raise ValueError( f"{STS_COLLECT_ENV} is set with ragged verify mode 'compact'. Padded " f"verify rows corrupt the per-position prefix label. Collect in " diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index 06a9e9077b4d..bde6001718bc 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -14,30 +14,13 @@ # limitations under the License. """Host-side coordinator turning DSpark confidence into a per-iteration draft length. -Three constraints shape this component, and all three point the same way: - -1. **The decision must be on the host.** TensorRT-LLM picks a decode CUDA graph - from ``(batch_size, draft_len, ...)`` -- every input to that choice is a host - value. A device-resident budget would have to be copied back before the batch - could be shaped, which is the synchronization we are trying to avoid. -2. **The step must land on a captured graph.** A shape with no captured graph - does not raise; ``maybe_get_cuda_graph`` returns ``None`` and the step - silently runs eager, which costs far more than the tokens saved. The drafted - block is always ``max_tier`` long, so what varies is the *verified* token - total, and the layout rounds that up to a captured verify bucket before the - batch is shaped. -3. **Every rank must decide identically.** The graph key is not part of the - attention-DP consistency allgather, so two ranks that shape the batch - differently select different graphs -- one replays, one falls back to eager - -- and their collectives diverge. The caller therefore reduces the ragged - decision across ranks (a single allgather in ``py_executor``) before any rank - acts on it; :meth:`decide_verify_lens` itself is rank-local by default. - -Confidence for step ``t`` only becomes readable on the host at ``t+1`` (or ``t+2`` -with the overlap scheduler), so the decision always runs on a lagged snapshot. -That is fine -- it is a throughput heuristic, not a correctness input -- but the -snapshot has to be *taken* without blocking, and a snapshot that is not ready yet -must degrade to verifying everything rather than to a stale guess. +Contracts: the decision is made on the host (CUDA-graph selection is a host +value); the verified token total must round up to a captured verify bucket, or +the step silently runs eager; and every rank must shape the batch identically +-- the caller reduces the ragged decision across ranks before any rank acts on +it, while :meth:`decide_verify_lens` itself is rank-local by default. +Confidence is read from a lagged, non-blocking snapshot; a snapshot that is +not ready degrades to verifying everything, never to a stale guess. """ import os @@ -52,11 +35,8 @@ from .dspark_schedule import (NEUTRAL_CONFIDENCE_LOGIT, DSparkScheduleConfig, compute_survival, schedule_verify_lens_topk) -#: Pin every request's verify window to this length, bypassing the planner. -#: Set by the SPS profiler to move M at a fixed draft length; also used to -#: hold the full block during STS collection, where a trimmed window would -#: clip the acceptance label and fit a temperature to the scheduler instead -#: of to the head. Not a serving knob. +#: Profiling-only pin (SPS profiling / STS collection): fixes every request's +#: verify window at this length, bypassing the planner. Not a serving knob. FORCE_VERIFY_LEN_ENV = "TLLM_DSPARK_FORCE_VERIFY_LEN" __all__ = ["DSparkVerifyPlanner"] @@ -92,32 +72,16 @@ def __init__( self.apply_calibration = apply_calibration or torch.sigmoid self.all_rank_max = all_rank_max - # Profiling pin. The SPS profiler has to hold M at a known value while it - # times steps, and it cannot use the planner's own decision to do that: - # the planner needs a cost table, and the cost table is what the profiler - # is producing. SGLang solves this with `dspark_force_budget_frac` pushed - # to a running server; TRT-LLM has no runtime control plane at - # world_size > 1, so this is an env var read ONCE here -- constant per - # process, zero cost on the step path, and forwarded to every worker at - # spawn, so all ranks agree by construction rather than by a collective. - # - # Absolute length, not SGLang's fraction. `schedule_verify_lens_topk` - # drops candidates below survival_eps, so a fractional budget yields - # `granted <= budget` and `choose_ragged_capture_shape` then rounds the - # total DOWN onto a smaller captured bucket -- the cell would be timed at - # one M and filed under another. With an absolute L every request gets - # exactly L, so M = bs * (L + 1) is the captured bucket by construction. + # Profiling pin, read once at init: constant per process and forwarded + # to every worker at spawn, so all ranks agree without a collective. + # An absolute length (not a fraction) so M = bs * (L + 1) lands exactly + # on a captured bucket. self._forced_verify_len = self._read_forced_verify_len() - # A pin requested at runtime (profiling endpoint) but not yet in force, - # held as its WIRE value because three states have to survive the trip: - # -1 nothing queued, 0 clear the pin, >0 pin to that length. Storing - # Optional[int] collapsed "clear" into "nothing queued", so a clear - # never reached the peers. min_verify_len is validated >= 1, so 0 can - # never collide with a real length. - # It is adopted through the step's existing cross-rank allgather so - # every rank starts pinning on the SAME step. Setting it locally would - # diverge the shape gate for as long as the ranks disagreed -- the same - # failure the env-var pin has on a pre-spawned world. + # Runtime pin queued but not yet in force, held as its wire value: + # -1 nothing queued, 0 clear the pin, >0 pin to that length + # (min_verify_len >= 1, so 0 is unambiguous). Adopted through the + # step's cross-rank allgather so every rank starts pinning on the + # same step; applying it locally would diverge the shape gate. self._pending_verify_len_pin: int = -1 self._host_buffer: Optional[torch.Tensor] = None @@ -125,8 +89,7 @@ def __init__( self._host_stamps: Optional[torch.Tensor] = None self._staged_seq: Optional[int] = None self._snapshot_valid = False - # Observability: silent degradation is the main failure mode here, so - # count every path that gives up on trimming. + # Count every path that declines to trim; degradation here is silent. self.stats = { "decisions": 0, "fallback_no_snapshot": 0, @@ -148,14 +111,10 @@ def stage_confidence(self, confidence_logits: torch.Tensor, ``confidence_logits`` is the worker's whole slot-indexed buffer, staged in full so :meth:`decide_verify_lens` can look up any live request by - slot. It is kilobytes (``max_num_requests x block`` fp32), so there is - nothing to gain from staging a subset -- and a subset would need a row - ordering to survive the one-iteration lag, which is exactly what slot - keying exists to avoid. - - Never synchronizes: the copy is issued on the current stream and an event - is recorded. :meth:`decide_verify_lens` polls that event and simply does - not use the snapshot if it has not landed. + slot across the one-iteration lag. Never synchronizes: the copy is + issued on the current stream and an event is recorded; + :meth:`decide_verify_lens` polls the event and skips a snapshot that + has not landed. """ if confidence_logits is None or confidence_logits.shape[0] == 0: return @@ -206,10 +165,6 @@ def _read_forced_verify_len(self) -> Optional[int]: raise ValueError( f"{self._FORCE_ENV}={value} outside [{lo}, {hi}]") if value not in self.tiers: - # A length with no captured graph makes every step run eager, and an - # eager step is orders of magnitude off the graph-replay cost the - # profiler is trying to measure -- the curve would be noise wearing - # the shape of a cost model. Refuse rather than produce it. raise ValueError( f"{self._FORCE_ENV}={value} is not in the captured tier ladder " f"{self.tiers}; every step would fall out of CUDA-graph replay " @@ -238,9 +193,8 @@ def validate_verify_len_pin(self, value: Optional[int]) -> Optional[int]: def request_verify_len_pin(self, value: Optional[int]) -> Optional[int]: """Queue a pin (or ``None`` to clear) for the next decode step. - Validated here so a bad value is rejected at the call site rather than - on some later step, but NOT applied here: the step's allgather carries - it so all ranks adopt it together (see ``adopt_verify_len_pin``). + Validated here but NOT applied: the step's allgather carries it so all + ranks adopt it together (see ``adopt_verify_len_pin``). """ value = self.validate_verify_len_pin(value) self._pending_verify_len_pin = 0 if value is None else int(value) @@ -251,19 +205,18 @@ def pending_verify_len_pin(self) -> int: return int(self._pending_verify_len_pin) def adopt_verify_len_pin(self, wire_value: int) -> None: - """Apply the group's agreed pin. ``-1`` leaves the current one alone. + """Apply the group's agreed pin from the allgathered payload. - Called from the executor with a value every rank read out of the same - allgathered payload, so the pin takes effect on the same step - everywhere. 0 is the wire encoding for "clear the pin". + Every rank applies the same payload, so the pin takes effect on the + same step everywhere. ``-1`` leaves the current pin alone; ``0`` + clears it. """ wire_value = int(wire_value) if wire_value < 0: return - # Compare-and-clear: the RPC thread may queue a NEW request between - # the decode loop reading the payload and adopting it (they are an - # allgather apart). Unconditionally resetting here wiped that second - # request after its endpoint had already answered 200. + # Compare-and-clear: the RPC thread may queue a new request between + # the payload being read and adopted; an unconditional reset would + # drop it. if self._pending_verify_len_pin == wire_value: self._pending_verify_len_pin = -1 new_pin = None if wire_value == 0 else wire_value @@ -277,17 +230,9 @@ def adopt_verify_len_pin(self, wire_value: int) -> None: def _note_snapshot_stats(self, selected: torch.Tensor, rows: Optional[Sequence[int]]) -> None: - """Measure what the argmax is about to eat. Never changes behaviour. - - Two independent instruments: - - stamp-lag histogram: (staged_seq - row stamp) per gathered row. A - healthy relay shows one tight mode; mass at large lags is stale - replay masquerading as data, and the clamp bucket collects - never-drafted rows. - - neutral-row count: rows still carrying the neutral fill on every - position, i.e. "unknown" rows entering the argmax as certainties. - Together they say whether the per-step snapshot content can be - trusted; without them the planner's input was never measured at all. + """Record snapshot-quality stats: stamp-lag histogram (staleness per + gathered row) and neutral-row count ("unknown" rows entering the + argmax as certainties). Observability only; never changes behaviour. """ try: n = int(selected.shape[0]) @@ -318,17 +263,11 @@ def _gather_rows( """This step's confidence, one row per generation request, or None. ``rows`` is the per-request buffer row index (``DSparkWorker. - confidence_row_for``). Supplying it is what makes the lagged snapshot - correct: the batch is reshuffled between the step that wrote the - confidence and the step that reads it, so the request at position ``i`` - is routinely a different one. - - ``rows=None`` keeps the positional reading for callers that own the - ordering themselves (the unit tests stage a purpose-built buffer). That - path refuses to run on a snapshot with fewer rows than the batch rather - than silently returning a short answer -- a short answer becomes a - partially-assigned batch downstream, where half the requests get a - verify window and half do not. + confidence_row_for``); the batch is reshuffled across the one-iteration + lag, so positional reads (``rows=None``) are only correct for callers + that own the ordering. Declines (None) rather than returning fewer + rows than the batch, which would partially assign verify windows + downstream. """ snapshot = self._ready_snapshot() if snapshot is None: @@ -359,43 +298,25 @@ def decide_verify_lens( ) -> Optional[List[int]]: """Per-request verify lengths for a ragged step, or None to stay uniform. - The budget still comes from the cost-model argmax -- that is what knows - how many verified tokens the step can afford -- but instead of splitting - it evenly it is handed to the global survival top-k, so a request whose - draft is still alive at depth 5 can take positions from one whose draft - died at depth 1. - - Returns None whenever the uniform path would have fallen back (no - snapshot, flat cost table, empty batch): raggedness is a throughput - heuristic and an untrustworthy input must degrade to verifying - everything, not to a stale split. - - The batch-wide *maximum* is still reduced across ranks, because the - drafted-token buffer width and the per-request padding are sized from - it; the individual lengths stay local. - - The returned list always has exactly ``num_gen_requests`` entries, or is - None. A partial list would leave the tail of the batch without a verify - window, and the callers downstream disagree about what that means: the - input layout is built per request (so it goes ragged) while the spec - metadata sees a ``None`` and stays uniform. + The cost-model argmax sets the budget; the global survival top-k splits + it across requests. Fail-closed: any untrustworthy input (no snapshot, + flat cost table, empty batch) returns None and the step verifies + everything. The batch-wide maximum is reduced across ranks (it sizes + the drafted-token buffer and per-request padding); individual lengths + stay local. Returns exactly ``num_gen_requests`` entries or None, + never a partial list. """ - # Counted before any decline, so the fallback counters below have a - # denominator rather than being absolute counts against an unknown total. + # Counted before any decline so the fallback counters have a denominator. self.stats["decisions"] += 1 if num_gen_requests <= 0: - # Benign during ramp-up, but it has to be nameable: an unlabelled - # `return None` here is indistinguishable in the counters from a - # real decline, and this feature's failures are all silent. self.stats["fallback_no_gen_requests"] += 1 return None if self._forced_verify_len is not None: - # Ahead of the cost-table gate on purpose: the profiler runs - # without a table, so a pin placed after it would never fire. Sets - # `lens` rather than returning, so the pinned path still runs the - # cross-rank agreement and the length assert below. + # Must stay ahead of the cost-table gate (the profiler runs without + # a table) and must set `lens` rather than return, so the pinned + # path still runs the cross-rank agreement below. self.stats["forced_steps"] += 1 lens = [int(self._forced_verify_len)] * int(num_gen_requests) else: @@ -409,34 +330,21 @@ def decide_verify_lens( self._note_snapshot_stats(selected, rows) probs = self.apply_calibration(selected) survival = compute_survival(probs) - # Scored with the tau the step will actually collect. The rung is - # still chosen from the same ladder over the same cost grid -- the - # token totals, the buckets and the captured graphs are unchanged -- - # but budget_argmax_over_uniform_lens scores rung L as if every - # request took columns [min, L), while what runs is a top-k - # redistribution of the same B = n*(L - min) tokens. Top-k takes the - # B largest survivals, so the realised tau is >= the scored one at - # an identical cost, for every rung. The gap is not constant across - # rungs, so the argmax could sit on the wrong one; scoring what runs - # removes the question. + # Scored as the top-k redistribution that actually runs, on the + # same tier-aligned cost grid as the uniform argmax. budget = compute_verify_token_budget( survival=survival.numpy().astype(np.float64), num_gen_requests=int(num_gen_requests), cost_table=self.cost_table, min_verify_len=self.cfg.min_verify_len, max_verify_len=self.cfg.resolved_max_verify_len, - # Without this the answer is unrealisable: any non-tier total is - # rounded up to a captured bucket, which can push a budget - # chosen to sit below a cost riser back over it. + # Required: any non-tier total is rounded up to a captured + # bucket, pushing a budget chosen below a cost riser back over it. allowed_lens=self.tiers, ) - # The price the argmax just paid for this step, kept so the run can - # be reconciled against reality: hostStepTimeMS measures what the - # step actually cost, and a systematic gap between the two is the - # one number that catches a wrong table, a wrong lookup, and an - # engine/table config mismatch alike (same cell measured 23% apart - # across two engine configs). These keys ride along in self.stats, - # which the [final] summary prints wholesale -- no extra plumbing. + # Predicted step cost, recorded so runs can be reconciled against + # measured hostStepTimeMS; a systematic gap flags a wrong or + # mismatched cost table. floor_tokens = int(num_gen_requests) * (self.cfg.min_verify_len + 1) predicted = float(self.cost_table.step_times( np.asarray([floor_tokens + int(budget)]), @@ -447,11 +355,8 @@ def decide_verify_lens( self.stats["predicted_steps"] = ( self.stats.get("predicted_steps", 0) + 1) lens = schedule_verify_lens_topk(survival=survival, budget=budget, cfg=self.cfg).tolist() - # What the argmax bought and what top-k handed out, recorded - # SEPARATELY: the realized windows in verify_len_hist are read back - # after the bucket fit widened them, so without these two counters - # nothing says what this rank's own decision was -- and a fit-side - # defect is indistinguishable from a planner-side one. + # Local decision histograms, recorded before the bucket fit widens + # the realized windows (verify_len_hist is post-fit). rung = self.cfg.min_verify_len + int(budget) // max(int(num_gen_requests), 1) rhist = self.stats.setdefault("local_rung_hist", {}) rhist[rung] = rhist.get(rung, 0) + 1 @@ -459,14 +364,10 @@ def decide_verify_lens( for v in lens: lhist[int(v)] = lhist.get(int(v), 0) + 1 - # `reduce_across_ranks=False` means the caller is doing the cross-rank - # agreement itself. That is not just an optimization here: every early - # return above is rank-local (an empty batch, an unprofiled cost table, - # a confidence snapshot whose copy event has not landed yet -- that last - # one is timing-dependent), so a collective issued *after* them is issued - # by some ranks and not others. Deciding locally and reducing once, - # unconditionally, at a point every rank reaches is the only shape of - # this that cannot deadlock. + # The early returns above are rank-local (some timing-dependent), so no + # collective may run before this point; reduce once, unconditionally, + # at a point every rank reaches. reduce_across_ranks=False means the + # caller does that agreement itself. if reduce_across_ranks: all_rank_max = all_rank_max or self.all_rank_max if all_rank_max is not None: diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 3ddbd5abf1bb..342dd89bdd3c 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -132,8 +132,7 @@ def prepare_attn_metadata_for_draft_replay(attn_metadata, return None if not isinstance(attn_metadata, TrtllmAttentionMetadata): return None - draft_block_offsets = getattr(attn_metadata, 'draft_kv_cache_block_offsets', - None) + draft_block_offsets = attn_metadata.draft_kv_cache_block_offsets if draft_block_offsets is None: return None @@ -518,47 +517,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. ``runtime_draft_len`` above is one - # scalar for the whole iteration, which cannot express "this request - # verifies 5 positions and that one verifies 1". When confidence-scheduled - # verification runs in ragged mode these carry the per-request split: + # 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, so downstream code that only knows - # about ``runtime_draft_len`` keeps working unchanged. Built by + # Both stay None on every other path. Built by # ``_torch/speculative/dspark_ragged.RaggedVerifyLayout``. verify_lens: Optional[torch.Tensor] = None qo_indptr: Optional[torch.Tensor] = None - # Host-side ``sum(verify_lens)``. Kept as a Python int because the packed - # buffers have to be sliced before any kernel runs, and reading the total - # off ``verify_lens`` would sync. Equals the captured token bucket after - # ``RaggedVerifyLayout.fill_bucket``. + # 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 # ``cap-accept``: per-request ceilings on the accepted count, in the same # token units as ``verify_lens``, for a batch whose layout was NOT trimmed. - # - # Deliberately a separate field rather than reusing ``verify_lens``. That - # one carries two meanings at once -- "stop acceptance here" AND "the draft - # buffers are packed by these windows" (see ``_padded_gen_draft_tokens``, - # which slices ``draft_tokens[:total_verify_tokens - num_gens]``). Under - # cap-accept only the first holds: the target scored the full block and the - # buffers are the usual rectangle. Setting ``verify_lens`` to get the clamp - # would silently take the packed-buffer branch too, and unpack a rectangle - # as if it were ragged -- one request's drafts attributed to another, with - # no error anywhere. + # Deliberately separate from ``verify_lens``: that field also implies the + # packed buffer layout, and under cap-accept the buffers are the usual + # rectangle -- reusing it would silently unpack a rectangle as if ragged. accept_caps: Optional[torch.Tensor] = None - # [batch] positions the cap discarded, PER REQUEST -- the exact acceptance - # the schedule cost each one, and the number ``compact`` cannot produce - # because it never scores those positions. - # - # Per request rather than a running total, because the total cannot tell - # "every request gave up a little" from "a few requests were cut to - # pieces" -- and those call for opposite conclusions about the planner, - # while summing to exactly the same number. + # [batch] positions the cap discarded, per request -- the acceptance the + # schedule cost each one. accept_cap_trim: Optional[torch.Tensor] = None @property @@ -711,8 +692,7 @@ def write_padding_onehot_draft_probs(self, padding_slot_ids, draft_len): overwritten here. No-op unless rejection is enabled and slots exist. Static shapes -> CUDA-graph safe. """ - if (not padding_slot_ids - or not getattr(self, "use_rejection_sampling", False) + if (not padding_slot_ids or not self.use_rejection_sampling or self.draft_probs is None): return onehot_vocab = (self.draft_probs_last_dim if self.draft_probs_last_dim @@ -837,13 +817,10 @@ def _normalize_request_sampling_params( tk_val = _first_or_none(sampling_config.top_k) tp_val = _first_or_none(sampling_config.top_p) - # Context requests have no draft tokens yet. Generation requests - # contribute one row per verified position, which under ragged - # verification is per-request: these buffers are indexed by flat - # token position, so a fixed stride here would apply request r's - # temperature/top_k/top_p to request r+1's tokens. Note this whole - # block is skipped for all-greedy batches, so a greedy-only - # accuracy run (GSM8K) cannot exercise it. + # 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. num_tokens = get_request_tokens_per_gen_step( request, 1 + self.runtime_draft_len ) if request.state == LlmRequestState.GENERATION_IN_PROGRESS else 1 @@ -1062,22 +1039,11 @@ def apply_accept_caps(num_accepted_tokens: torch.Tensor, num_contexts: int, spec_metadata) -> None: """``cap-accept``: hold each request to its scheduled window, in place. - This has to run HERE, inside acceptance, and not on the host after the - sampler syncs. ``num_accepted_tokens`` is consumed on device later in the - same forward: the DSpark drafter reads it to pick the hidden state the next - block is conditioned on, to back-fill the accepted tokens into its rolling - KV window, and to advance its persistent decode position - (``_ctx_len += nacc``). Capping after that leaves the drafter advanced by - the uncapped count while only the capped prefix was committed -- the output - stays lossless, because verification is authoritative, but the draft state - drifts permanently and the mode stops being the faithful reference it - exists to be. SGLang caps at the same point, before anything downstream - reads the count. - - In-place on the caller's tensor and on persistent buffers, so this is safe - to record into a CUDA graph and never synchronizes. The per-request losses - ride back to the host on the sampler's existing copy, so measuring them - costs no sync of its own. + Must run inside acceptance, before the DSpark drafter consumes + ``num_accepted_tokens`` on device (hidden-state pick, rolling-KV + back-fill, ``_ctx_len`` advance); capping later leaves the draft state + permanently ahead of what was committed. In-place and sync-free, so safe + to record into a CUDA graph. """ caps = getattr(spec_metadata, "accept_caps", None) if caps is None: @@ -1600,7 +1566,7 @@ def _draft_logits_are_sharded(self, logits, spec_metadata): sharded and gather identical copies (overflowing the d2t table). """ mapping = self.mapping - if mapping is None or getattr(mapping, "tp_size", 1) <= 1: + if mapping is None or mapping.tp_size <= 1: return False # Under attention DP every rank is data-parallel -- it owns a distinct # set of requests -- so the draft logits must NOT be cross-rank gathered, @@ -1611,7 +1577,7 @@ def _draft_logits_are_sharded(self, logits, spec_metadata): # desyncing them into a hang / DeepEP launch failure. Only plain TP # (tp>1 without ADP), where all ranks share the same tokens sharded over # the vocab dim, needs the gather. - if getattr(mapping, "enable_attention_dp", False): + if mapping.enable_attention_dp: return False draft_full_vocab = (getattr(spec_metadata, "draft_vocab_size", 0) or getattr(spec_metadata, "vocab_size", 0) or 0) @@ -1709,11 +1675,9 @@ def _reshape_draft_tokens_for_accept(self, spec_metadata, num_gens, device): dtype=torch.int, device=device) if spec_metadata.is_ragged_verify: - # The buffer is packed by each request's own draft window, so it is - # not a [num_gens, runtime_draft_len] rectangle. Reshaping it would - # not even fail loudly -- fill_bucket pads the total to a captured - # bucket, which can happen to divide -- it would just attribute one - # request's drafts to another. + # 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) @@ -1773,17 +1737,8 @@ def _rejection_buffers_valid(self, draft_tokens, draft_len, stored_vocab, if draft_tokens.dim() != 2 or draft_tokens.shape[0] != num_gens: return False # logits must cover context rows (1 each) + however many gen rows the - # target was actually given. - # - # goal doc H7. Under ragged verification that is `total_verify_tokens`, - # not `num_gens * (draft_len + 1)` -- and the two differ by exactly the - # amount the scheduler trimmed, so requiring the uniform count made this - # guard fail whenever the feature did anything. Rejection sampling then - # fell back to strict acceptance silently, changing the sampled - # distribution for temperature > 0 while leaving the ragged branch below - # (which is complete: it scatters to the padded rectangle, fills the pad - # rows one-hot, and clamps the accepted count back to `verify_lens`) - # unreachable. Greedy GSM8K cannot see any of this. + # 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 spec_metadata.is_ragged_verify: # Fail closed on a half-built ragged layout: the branch below @@ -2070,7 +2025,7 @@ def greedy_sample_draft_with_tp_gather(self, ``[:batch]``, which belong to group rank 0. """ if (mapping_lm_head_tp is not None - and getattr(mapping_lm_head_tp, "tp_size", 1) > 1): + and mapping_lm_head_tp.tp_size > 1): from ..distributed.ops import allgather combined = self._get_local_max_and_combined(logits, mapping_lm_head_tp) @@ -2082,8 +2037,7 @@ def greedy_sample_draft_with_tp_gather(self, return self._get_draft_tokens_from_gathered(own_segment) mapping = self.mapping sharded = self._draft_logits_are_sharded(logits, spec_metadata) - if (sharded and mapping is not None - and getattr(mapping, "tp_size", 1) > 1 + if (sharded and mapping is not None and mapping.tp_size > 1 and not mapping.enable_attention_dp): from ..distributed.ops import allgather combined = self._get_local_max_and_combined(logits) @@ -2128,7 +2082,7 @@ def advanced_sample_draft_block(self, gen_logits: torch.Tensor, self._update_advance_draft_sampling_seed(gen_logits.device) flat_logits = gen_logits.reshape(num_gens * K, vocab) - if getattr(spec_metadata, "use_rejection_sampling", False): + if spec_metadata.use_rejection_sampling: flat_tokens, flat_probs = ( sampling_batch_spec_dec_one_model_for_rejection( flat_logits, @@ -2177,8 +2131,7 @@ def write_context_onehot_draft_probs(self, next iter reads this scalar). No-op unless rejection is enabled and the slot-indexed buffers exist. Static shapes -> CUDA-graph safe. """ - if (num_contexts <= 0 - or not getattr(spec_metadata, "use_rejection_sampling", False) + if (num_contexts <= 0 or not spec_metadata.use_rejection_sampling or spec_metadata.draft_probs is None): return ctx_slot_ids = spec_metadata.batch_slot_ids[:num_contexts] @@ -2347,8 +2300,7 @@ def draft_kv_cache_context(self, attn_metadata, draft_kv_cache_manager): return # Check if draft KV cache block offsets are allocated - draft_block_offsets = getattr(attn_metadata, - 'draft_kv_cache_block_offsets', None) + draft_block_offsets = attn_metadata.draft_kv_cache_block_offsets if draft_block_offsets is None: # Draft KV cache block offsets not allocated, skip switching yield From 59cdc060d5add36bef561cb1117dc1ddd7f4cf39 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 5 Aug 2026 09:31:37 -0700 Subject: [PATCH 099/177] [None][fix] Remove the debug harnesses that outlived their investigations Everything here existed to chase specific bugs that are now fixed; the observability counters and summaries stay untouched. - TLLM_DSPARK_ASSERT_LAYOUT and _assert_ragged_layout_consistent: the opt-in batch-walking cross-check built while hunting the ragged layout desyncs. The invariants it guarded are now pinned by unit tests (test_dspark_ragged, test_dspark_graph_key) and the window check in the graph-key derivation. - TLLM_DSPARK_ASSERT_ACTIVE, the mid-loop and shutdown assert gates, and DSparkRaggedStats.assert_ragged_active: the ragged-liveness test gate. Nothing consumes it any more; the shutdown summary (kept) carries the same counters for eyeballs. - TLLM_HANG_DETECTION_TIMEOUT: hang_detector.py reverted to upstream; the override existed so a frozen GPU could outlive the watchdog during coredump forensics. - tests/microbenchmarks/dspark_make_steep_sps_table.py: the synthetic steep-cost-table generator written to force trimming while reproducing the IMA. Two tests that tested the deleted gate itself are removed with it; the counter assertions they shared stay. DSpark suite: 226 passed on B300. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/pyexecutor/hang_detector.py | 42 +++- .../_torch/pyexecutor/model_engine.py | 85 +------ tensorrt_llm/_torch/pyexecutor/py_executor.py | 52 +--- .../speculative/dspark_observability.py | 77 +----- .../dspark_make_steep_sps_table.py | 222 ------------------ .../hw_agnostic/test_dspark_cap_accept.py | 13 - .../hw_agnostic/test_dspark_padded_ceiling.py | 5 +- 7 files changed, 36 insertions(+), 460 deletions(-) delete mode 100644 tests/microbenchmarks/dspark_make_steep_sps_table.py diff --git a/tensorrt_llm/_torch/pyexecutor/hang_detector.py b/tensorrt_llm/_torch/pyexecutor/hang_detector.py index e13768a9b64e..f6dde7c58959 100644 --- a/tensorrt_llm/_torch/pyexecutor/hang_detector.py +++ b/tensorrt_llm/_torch/pyexecutor/hang_detector.py @@ -91,13 +91,7 @@ class HangDetector: def __init__( self, timeout: Optional[int] = None, on_detected: Optional[Callable[[], None]] = None ): - # TLLM_HANG_DETECTION_TIMEOUT: diagnostic override. Debugging a device - # fault with CUDA_DEVICE_WAITS_ON_EXCEPTION needs the frozen GPU state - # to outlive this watchdog, whose hard-kill would otherwise destroy the - # very state the debugger is attached to. - if timeout is None: - timeout = int(os.environ.get("TLLM_HANG_DETECTION_TIMEOUT", "300")) - self.timeout = timeout + self.timeout = timeout if timeout is not None else 300 assert self.timeout > 0, "timeout must be greater than 0" self.on_detected = on_detected or (lambda: None) self.task = None @@ -106,6 +100,7 @@ def __init__( self.lock = threading.Lock() self.active = False self._detected = False + self._status_providers: list[Callable[[], str]] = [] def start(self): """Enable hang detection.""" @@ -119,13 +114,38 @@ def run_loop(): self.loop_thread = threading.Thread(target=run_loop, daemon=True, name="hang_detector_loop") self.loop_thread.start() - async def _detect_hang(self): + def register_status_provider(self, provider: Callable[[], str]) -> None: + """Register a nonblocking callable that returns status to dump on hang detection.""" + with self.lock: + self._status_providers.append(provider) + + async def _detect_hang(self) -> None: await asyncio.sleep(self.timeout) with self.lock: - self._detected = True - logger.error(f"Hang detected after {self.timeout} seconds.") + status_providers = tuple(self._status_providers) + + # All diagnostics are best-effort: nothing may prevent on_detected() + # (hard-kill propagation) from firing. + _best_effort_log_error(f"Hang detected after {self.timeout} seconds.") + for provider in status_providers: + try: + status = provider() + if status: + _best_effort_log_error(status) + except Exception as error: # noqa: BLE001 - isolate diagnostic providers + _best_effort_log_error( + f"HangDetector: status provider failed with {type(error).__name__}: {error}" + ) + try: print_all_stacks() - self.on_detected() + except Exception: # noqa: BLE001 - stack dump must not block hard kill + pass + + # Set _detected last so observers (and tests) see it only once + # diagnostics are done and on_detected is about to fire. + with self.lock: + self._detected = True + self.on_detected() def detected(self): """Return True if hang is detected.""" diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2bc99b84f031..6619922b7dba 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -259,12 +259,6 @@ def _filter_cuda_graph_seq_lens(cuda_graph_seq_lens: list[int], _DEEP_GEMM_PDL_CONFIGURED = False -# Ragged layout bugs are silent (rows get attributed to the wrong requests -# while every kernel-side divisibility check still passes); these assertions -# make them loud at the cost of walking the batch, so they are opt-in. -_DSPARK_ASSERT_LAYOUT = os.environ.get("TLLM_DSPARK_ASSERT_LAYOUT", - "0") in ("1", "true", "True") - def _configure_deep_gemm_pdl() -> None: global _DEEP_GEMM_PDL_CONFIGURED @@ -3882,79 +3876,6 @@ def _publish_gen_token_layout(self, attn_metadata, attn_metadata.ragged_verify_lens = self._ragged_token_lens( generation_requests) - def _assert_ragged_layout_consistent(self, spec_metadata, attn_metadata, - scheduled_requests) -> None: - """Assert the four descriptions of this step's generation-token count agree: - - sum(1 + py_verify_len) the requests, as fitted - num_tokens - num_ctx_tokens the flat input layout - total_verify_tokens what acceptance slices by - the ragged bucket the CUDA-graph key's token axis - - A mismatch is silent token misattribution, not a crash. Off by - default (walks the batch); enable with TLLM_DSPARK_ASSERT_LAYOUT=1. - """ - verify_lens = self._ragged_token_lens( - scheduled_requests.generation_requests) - if verify_lens is None: - # A uniform step must not carry a published ragged bucket: a later - # ragged step would key a graph the batch does not describe. - agreed_stale = self.cuda_graph_runner.agreed_ragged_bucket - assert agreed_stale is None, ( - f"uniform step carries a stale published ragged bucket " - f"{agreed_stale}; a failed fit exited without clearing it") - # Uniform step: assert the converse, that nothing ragged leaked - # through. A half-published layout is the dangerous state. - assert spec_metadata.total_verify_tokens is None, ( - "uniform batch left a stale total_verify_tokens on the spec " - f"metadata: {spec_metadata.total_verify_tokens}") - assert getattr(attn_metadata, "ragged_verify_lens", None) is None, ( - "uniform batch left stale ragged_verify_lens on the attention " - f"metadata: {attn_metadata.ragged_verify_lens}") - return - - from_requests = sum(verify_lens) - from_layout = attn_metadata.num_tokens - attn_metadata.num_ctx_tokens - from_spec = int(spec_metadata.total_verify_tokens) - agreed = self.cuda_graph_runner.agreed_ragged_bucket - - assert from_requests == from_layout, ( - f"ragged layout mismatch: requests carry {from_requests} generation " - f"tokens but the flat input layout has {from_layout} " - f"(attn_metadata.num_tokens={attn_metadata.num_tokens}, " - f"num_ctx_tokens={attn_metadata.num_ctx_tokens}). The token layout " - f"and the verify windows were built from different values.") - assert from_requests == from_spec, ( - f"ragged layout mismatch: requests carry {from_requests} generation " - f"tokens but spec_metadata.total_verify_tokens is {from_spec}; " - f"acceptance would slice the target logits at the wrong offsets.") - if agreed is not None: - # `from_requests` walks this batch's windows; `agreed` is what the - # graph key carries. A mismatch means the replayed graph is a - # different width than the batch, and nothing downstream says so. - assert from_requests == int(agreed), ( - f"ragged layout mismatch: the batch carries {from_requests} " - f"generation tokens but the fit agreed on {agreed}, which is " - f"what the CUDA-graph key's token axis says; the replayed " - f"graph was captured at a different width.") - - # Row axis: the token-major presentation can be malformed while every - # token count above agrees, so it needs its own check. - view = attn_metadata.token_major_gen_view() if hasattr( - attn_metadata, "token_major_gen_view") else None - if view is not None: - expected_rows = attn_metadata.num_contexts + from_requests - assert view.num_rows == expected_rows, ( - f"token-major presentation has {view.num_rows} rows but the " - f"batch has {attn_metadata.num_contexts} context rows plus " - f"{from_requests} generation tokens; the MLA generation ops " - f"would read a sequence count that does not describe this step.") - assert view.host_context_lengths.shape[0] == expected_rows, ( - f"token-major host_context_lengths has " - f"{view.host_context_lengths.shape[0]} entries but the " - f"presentation is {expected_rows} rows -- that tensor's shape " - f"IS what the ops read as their sequence count.") - def _pinned_host(self, key: str, values, dtype) -> torch.Tensor: """A persistent pinned staging buffer holding ``values``. @@ -5710,9 +5631,6 @@ def previous_seq_slots_device(): self._attach_ragged_verify_layout( spec_metadata, attn_metadata, scheduled_requests.generation_requests) - if _DSPARK_ASSERT_LAYOUT: - self._assert_ragged_layout_consistent( - spec_metadata, attn_metadata, scheduled_requests) if isinstance(spec_metadata, Eagle3SpecMetadata): spec_metadata.request_accepted_path = request_accepted_path # No-op for non 1-model @@ -6982,8 +6900,7 @@ def forward(self, can_run_graph = key is not None # A ragged step that misses its captured shape quietly runs eager; # count it. Only generation-only steps count: a batch with context - # requests has no generation graph to miss, and booking it as - # eager would make `assert_ragged_active` unsatisfiable. + # requests has no generation graph to miss. if padded_requests.can_run_cuda_graph: self._record_dspark_graph_use( replayed=can_run_graph, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index fe24886b903d..11995f0c821d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2399,42 +2399,11 @@ def _process_iter_stats( prev_device_step_time_ms=prev_device_step_time_ms, gpu_forward_time_ms=gpu_forward_time_ms) - def _maybe_assert_dspark_ragged_active(self, stats) -> None: - """Run the ragged gate mid-loop, once the step floor is reached. - - Raises mid-loop (a cleanup-time raise lands in a shutdown path the - test runner never sees); fires at most once. - """ - # Deliberately lazy: this class has several construction paths, and a - # missed __init__ attribute would silently disable the gate. - if getattr(self, "_dspark_asserted", False): - return - min_steps = self._dspark_assert_min_steps - # Count only multi-request steps: tiny ramp-up batches have trivially - # identical windows and would fail an innocent run. - if min_steps is None or stats.steps_multi_request < min_steps: - return - self._dspark_asserted = True - stats.assert_ragged_active(require_trim=True) - - @property - def _dspark_assert_min_steps(self) -> Optional[int]: - """Step floor from TLLM_DSPARK_ASSERT_ACTIVE, or None when disabled.""" - raw = os.environ.get("TLLM_DSPARK_ASSERT_ACTIVE", "") - if not raw or raw in ("0", "false", "False"): - return None - try: - return 32 if raw in ("1", "true", "True") else int(raw) - except ValueError: - return None - def _report_dspark_ragged_stats(self) -> None: - """Final DSpark scheduling summary, and optionally a hard gate on it. + """Final DSpark scheduling summary at shutdown. The counters live in the worker process and are unreachable from the - driver under multi-rank TP, so the summary is always logged at - shutdown, and with ``TLLM_DSPARK_ASSERT_ACTIVE`` the worker asserts on - its own counters (a worker exception does propagate to the driver). + driver under multi-rank TP, so the log is the shared channel. """ worker = getattr(self.model_engine, "_get_spec_worker", None) stats = getattr(worker(), "ragged_stats", None) if worker else None @@ -2446,26 +2415,10 @@ def _report_dspark_ragged_stats(self) -> None: if stats is None or not stats.steps_total: return stats.log_summary(prefix="DSpark ragged verify [final]") - # A MINIMUM STEP COUNT, not a boolean: throwaway executors (KV-cache - # estimation) run a few decode steps and must not trip the assert. - min_steps = self._dspark_assert_min_steps - if min_steps is None: - return - if stats.steps_total < min_steps: - logger.info( - f"DSpark ragged verify: skipping the active-path assertion, " - f"this executor ran {stats.steps_total} decode step(s) " - f"(< {min_steps}).") - return - # require_trim is deliberately on: a run that produced windows but - # never shortened one delivered the uniform baseline. - stats.assert_ragged_active(require_trim=True) def _executor_loop_cleanup(self): try: self._report_dspark_ragged_stats() - except AssertionError: - raise except Exception as exc: # noqa: BLE001 - never block shutdown on stats logger.warning(f"DSpark ragged stats reporting failed: {exc}") # Wake any waiters in await_responses BEFORE potentially-blocking @@ -3468,7 +3421,6 @@ def _dspark_confidence_draft_len(self, # boundary, so the log is the only shared channel. if stats.steps_total and stats.steps_total % 32 == 0: stats.log_summary(prefix="DSpark ragged verify [periodic]") - self._maybe_assert_dspark_ragged_active(stats) fitted = None delivered = None if cap_accept: diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index dc8950ae60f2..143300825977 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -15,9 +15,7 @@ """Verification-mode selection and per-step counters for DSpark ragged verify. Failures on this path are silent (output stays correct at baseline accuracy), -so activity is counted rather than inferred, and -:meth:`DSparkRaggedStats.assert_ragged_active` turns the counters into a test -gate. Modes: ``static`` verifies the full block; ``cap-accept`` computes and +so activity is counted rather than inferred. Modes: ``static`` verifies the full block; ``cap-accept`` computes and commits per-request windows but submits the full block, so its output must stay bit-identical to ``static``; ``compact`` submits only each request's window and needs a profiled :class:`~.dspark_planner.SpsCostTable` -- with a @@ -401,79 +399,6 @@ def log_summary(self, *, prefix: str = "DSpark ragged verify") -> None: summary = self.summary() logger.info(f"{prefix}: {summary}") - def assert_ragged_active(self, *, require_trim: bool = True) -> None: - """Fail loudly if the ragged path never actually ran. - - An accuracy pass alone proves nothing: the scheduler may have declined - every step. - - Args: - require_trim: also require that delivered tokens came in under the - no-trim ceiling. Leave on unless deliberately measuring a - workload where every draft survives. - - Raises: - AssertionError: with the full counter summary. - """ - summary = self.summary() - if self.steps_total == 0: - raise AssertionError( - f"DSpark ragged verify recorded no decode steps at all: {summary}" - ) - if self.steps_ragged == 0: - raise AssertionError( - f"DSpark ragged verify never produced differing per-request " - f"windows, so the ragged path was not exercised. The usual " - f"cause is a missing profiled SPS cost table, under which the " - f"planner's budget degenerates to verify-all: {summary}") - if self.distinct_verify_lens < 2: - raise AssertionError( - f"DSpark ragged verify handed out only one distinct window " - f"size, so the batch was uniform in all but name: {summary}") - # Zero planner decisions against non-zero windowed steps means the - # planner's stats are not wired to this object -- and its fallback - # counters are the only record of WHY a step declined to trim. - if self.planner_stats is not None: - decisions = int(self.planner_stats.get("decisions", 0)) - windowed = self.steps_ragged + self.steps_uniform_windows - if windowed > 0 and decisions == 0: - raise AssertionError( - f"DSpark scheduling produced windows on {windowed} steps " - f"but the planner recorded 0 decisions, so its fallback " - f"counters cannot be read -- and they are the only record " - f"of why a step declined to trim: {summary}") - # `cap-accept` submits the full block on purpose, so a zero trim ratio - # is expected there, not a silent no-op. - if (require_trim and self.mode.trims_submitted_tokens - and self.trim_ratio <= 0.0): - raise AssertionError( - f"DSpark ragged verify delivered as many tokens as a no-trim " - f"run, so nothing was saved: {summary}") - # Fail only on misses ragged is responsible for: context-filled batches - # fall out of replay on the uniform path too, and blaming those would - # make the check unsatisfiable under continuous batching. - # "key_not_captured"/"bs_not_supported" mean the batch was graphable - # and the bucket grid is wrong; "peer_shape_mismatch" means one rank - # chose a different ragged bucket and dragged the whole world eager. - blamed = {"key_not_captured", "bs_not_supported", "peer_shape_mismatch"} - attributable = { - reason: n - for reason, n in self.graph_miss_reasons.items() if reason in blamed - } - if attributable: - raise AssertionError( - f"DSpark ragged verify dropped out of CUDA graph replay for " - f"reasons attributable to ragged scheduling " - f"({attributable}); a ragged step without a captured shape " - f"costs far more than the tokens it trims: {summary}") - if self.graph_eager > 0 and not self.graph_miss_reasons: - raise AssertionError( - f"DSpark ragged verify dropped out of CUDA graph replay on " - f"{self.graph_eager} step(s) but recorded no reason for any of " - f"them, so the misses cannot be attributed. That is a broken " - f"probe, not a clean run: {summary}") - - def format_verify_len_histogram(lens: List[int]) -> str: """One-line ``len:count`` rendering, for per-step debug logging.""" hist = Counter(int(v) for v in lens) diff --git a/tests/microbenchmarks/dspark_make_steep_sps_table.py b/tests/microbenchmarks/dspark_make_steep_sps_table.py deleted file mode 100644 index b22260c5b90e..000000000000 --- a/tests/microbenchmarks/dspark_make_steep_sps_table.py +++ /dev/null @@ -1,222 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Emit an SPS cost table steep enough that the planner chooses to trim. - -WHAT THIS IS FOR, AND WHAT IT IS NOT FOR ----------------------------------------- -It exercises the ragged path end to end *through the planner*: with a steep -``theta(M)`` the confidence scheduler decides to trim, and its budget then -flows through the same confidence-ordered top-k that ships, which hands -different windows to requests with different survival. That is the only way to -see per-request verify lengths on a checkpoint whose acceptance is too high for -trimming to pay. - -**No throughput or latency number taken with this table means anything.** The -planner is optimising against a fictional cost curve, so it will trim where a -real deployment should not. Use it to answer "does the ragged path correctly -support different verify lengths per request", never "is ragged faster". The -output carries ``"SYNTHETIC": true`` so a run cannot quietly present it as a -measurement -- ``run_dspark_throughput.sh`` refuses to treat such a table as -evidence. - -WHY A TABLE RATHER THAN AN OVERRIDE ------------------------------------ -This replaces ``TLLM_DSPARK_FORCE_VERIFY_LENS``, which assigned windows by -batch position. Position is orthogonal to confidence, so that knob reproduced -the ragged *shape* while bypassing the policy entirely -- it could not tell a -correct assignment from an arbitrary one, and it had to sit in production code -ahead of the cost-table gate to be reachable. A table drives the real path. - -HOW STEEPNESS IS CHOSEN ------------------------ -The scheduler maximises expected accepted tokens per second, roughly - - (bs + sum of the top-K survival probabilities) / T(bs, K) - -With per-token acceptance around 0.9 each extra verify token adds about 0.9 to -the numerator, so trimming only wins once ``theta`` grows faster than that. The -default ``--exponent 2.0`` makes theta quadratic in the token count, which is -comfortably steeper than the measured curve (near-linear) and forces trimming -across the whole batch-size ladder. - - # derive from a real table, keeping its grid - python3 tests/microbenchmarks/dspark_make_steep_sps_table.py \ - --from sps_real_final.json --out steep.json - - # or synthesise a grid outright - python3 tests/microbenchmarks/dspark_make_steep_sps_table.py --out steep.json -""" - -import argparse -import json -import sys - - -def _default_grid(max_batch_size: int, width: int): - """Token totals a ragged batch can land on: bs * (verify_len + 1). - - The ladder always includes ``max_batch_size`` itself. A grid that stops - short of the deployment's widest step is worse than useless: theta clamps - at the last bucket, so the full window looks no more expensive than a - trimmed one and the planner keeps it -- a table that silently cannot - induce trimming, which is the failure this tool exists to prevent. - """ - ladder = sorted({b for b in (1, 4, 8, 16, 32, 64, 128, 256, 512) - if b <= max_batch_size} | {max_batch_size}) - tokens = sorted({b * w for b in ladder for w in range(2, width + 1)}) - return ladder, tokens - - - -def _preview_with_real_planner(*, path, batch_sizes, accept, max_verify_len): - """What budget_argmax_over_uniform_lens actually picks for this table. - - Returns True if any previewed batch size trims, False if none do, or None - when the planner cannot be imported (outside the container), so the caller - can distinguish "checked and fine" from "not checked". - """ - try: - import numpy as np - from tensorrt_llm._torch.speculative.dspark_planner import ( - SpsCostTable, budget_argmax_over_uniform_lens) - except Exception: - return None - - table = SpsCostTable.from_json_file(path) if hasattr( - SpsCostTable, "from_json_file") else None - if table is None: - with open(path, encoding="utf-8") as handle: - table = SpsCostTable(**{ - k: v for k, v in json.load(handle).items() - if not k.startswith("_") and k != "SYNTHETIC" - }) - - tiers = sorted({1, max(1, (max_verify_len + 1) // 2), max_verify_len}) - print() - print(f"planner's actual choice at acceptance {accept} " - f"(L = verify length, {max_verify_len} = no trim), tiers {tiers}:") - any_trim = False - for bs in [int(b) for b in str(batch_sizes).split(",") if str(b).strip()]: - # survival[r, k] = P(first k+1 drafted tokens all accepted) - survival = np.array( - [[accept**(k + 1) for k in range(max_verify_len)]] * bs, - dtype=np.float64) - chosen = budget_argmax_over_uniform_lens( - survival=survival, num_gen_requests=bs, cost_table=table, - allowed_lens=tiers, min_verify_len=1) - trims = int(chosen) < max_verify_len - any_trim = any_trim or trims - print(f" bs={bs:<5} -> L={chosen}" - f"{' TRIM' if trims else ' (full window)'}") - return any_trim - - -def main() -> int: - parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) - parser.add_argument("--out", required=True) - parser.add_argument("--from", dest="source", default=None, - help="real table to take the grid from; its timings " - "are replaced, only token_counts/batch_sizes are " - "reused so the shape matches the deployment") - parser.add_argument("--exponent", type=float, default=2.0, - help="theta(M) ~ M**exponent; >1 makes extra verify " - "tokens superlinearly expensive") - parser.add_argument("--peak-ms", type=float, default=400.0, - help="theta at the largest token count") - parser.add_argument("--fixed-overhead-ms", type=float, default=0.05, - help="the non-trimmable floor. MUST be small next to " - "theta in the operating range, or trimming cannot " - "pay however steep theta is -- see --preview-bs") - parser.add_argument("--preview-bs", default="8,32,128", - help="batch sizes to predict the planner's choice for, " - "so a table that cannot induce trimming is caught " - "here instead of 20 minutes into a run") - parser.add_argument("--preview-accept", type=float, default=0.9, - help="per-token acceptance to predict against; 0.9 is " - "what DeepSeek-V4-Pro-DSpark measures") - parser.add_argument("--max-batch-size", type=int, default=128) - parser.add_argument("--max-verify-len", type=int, default=5) - args = parser.parse_args() - - if args.exponent <= 1.0: - parser.error( - f"--exponent {args.exponent} is not steeper than linear, so the " - f"planner will keep verifying the full block and the ragged path " - f"will stay dark -- which is the exact failure this tool exists " - f"to avoid") - - needed = args.max_batch_size * (args.max_verify_len + 1) - if args.source: - with open(args.source, encoding="utf-8") as handle: - src = json.load(handle) - tokens = [int(t) for t in src["token_counts"]] - batch_sizes = [int(b) for b in src["batch_sizes"]] - if max(tokens) < needed: - parser.error( - f"--from {args.source} tops out at {max(tokens)} tokens but " - f"max_batch_size {args.max_batch_size} reaches {needed} at the " - f"full window. theta would clamp there, making the full window " - f"look as cheap as a trimmed one, so the planner would never " - f"trim. Drop --from to synthesise a grid that covers it.") - else: - batch_sizes, tokens = _default_grid(args.max_batch_size, - args.max_verify_len + 1) - - tokens = sorted({t for t in tokens if t > 0}) - peak = float(max(tokens)) - theta = [args.peak_ms * (t / peak)**args.exponent for t in tokens] - - table = { - "token_counts": tokens, - "step_time_ms": theta, - "fixed_overhead_ms": args.fixed_overhead_ms, - "batch_sizes": batch_sizes, - "batch_overhead_ms": [0.0] * len(batch_sizes), - "SYNTHETIC": True, - "_meta": { - "purpose": "force the planner to trim so the ragged path runs; " - "NOT a measurement of anything", - "exponent": args.exponent, - "derived_from": args.source, - }, - } - with open(args.out, "w", encoding="utf-8") as handle: - json.dump(table, handle, indent=2) - handle.write("\n") - - # Ask the REAL planner what it will do, rather than re-deriving its argmax. - # - # A hand-written approximation of the objective was wrong twice: it models - # the cost as fixed + theta(M), while budget_argmax_over_uniform_lens calls - # cost_table.step_time(total_verify_tokens(bs, length), num_gen_requests), - # which interpolates the table and adds the per-batch term. The - # approximation predicted "trim at every batch size" for a table under - # which the planner kept the full window through a 20-minute run. A preview - # that does not agree with the thing it previews is worse than none. - any_trim = _preview_with_real_planner( - path=args.out, batch_sizes=args.preview_bs, - accept=args.preview_accept, max_verify_len=args.max_verify_len) - if any_trim is None: # planner not importable here - print("preview skipped: run this inside the TensorRT-LLM container to " - "have the real planner check the table before you spend a run " - "on it", file=sys.stderr) - any_trim = True - print(f" bs={bs:<5} -> K={best_k}" - f"{' TRIM' if trims else ' (full window)'}") - if not any_trim: - print() - print("WARNING: this table induces no trimming at any previewed batch " - "size, so the ragged path will stay dark. Lower " - "--fixed-overhead-ms or raise --exponent.", file=sys.stderr) - - print() - print(f"wrote {args.out}: {len(tokens)} token buckets " - f"{tokens[0]}..{tokens[-1]}, theta {theta[0]:.3f}..{theta[-1]:.1f} ms " - f"(exponent {args.exponent}), SYNTHETIC=true") - print("Use only to check that per-request verify lengths work. Any " - "throughput taken with this table is meaningless.") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py index 909fce895ed4..67a93ab591b2 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py @@ -220,8 +220,6 @@ def test_cap_accept_must_not_be_credited_with_a_compute_saving(): # The schedule is still visible -- that is the whole point of the mode. assert stats.window_tokens == sum(1 + v for v in (5, 3, 2, 1)) assert stats.steps_ragged == 1 - # ...and the gate stays satisfiable in the mode it exists to validate. - stats.assert_ragged_active(require_trim=True) def test_compact_still_derives_delivered_from_the_bucket(): @@ -249,17 +247,6 @@ def test_cap_accept_computes_windows_but_does_not_trim_the_token_axis(): assert not RaggedVerifyMode.STATIC.trims_submitted_tokens -def test_the_active_gate_still_demands_a_saving_from_compact(): - """The regression the clause above must not introduce.""" - stats = _stats(mode=RaggedVerifyMode.COMPACT) - stats.record_step(num_gen_requests=4, - verify_lens=[5, 3, 2, 1], - bucket=4 * (1 + MAX_DRAFT_LEN)) - assert stats.trim_ratio == 0.0 - with pytest.raises(AssertionError, match="nothing was saved"): - stats.assert_ragged_active(require_trim=True) - - def test_a_slot_left_unwritten_cannot_inherit_the_previous_occupant_s_loss(): """The stale-buffer hazard, guarded at the one place it can arise. diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py index 94892660fbaf..73c8227f0421 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py @@ -10,10 +10,7 @@ ``num_gen_requests * (1 + block)`` over this rank's real, unpadded requests. Two different row bases in one ratio. -The gate reads that ratio: ``assert_ragged_active`` fires on ``trim_ratio <= 0`` -with the message "delivered as many tokens as a no-trim run, so nothing was -saved" -- the opposite of what a negative value means, so the defect also -misdiagnoses itself. +A wrong ratio silently misreports the feature's saving in every summary. """ import pytest From 1b135358bed8de6f1c1f841035431dd69318a8a1 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Thu, 6 Aug 2026 02:26:58 -0700 Subject: [PATCH 100/177] [None][feat] Profile the SPS table on the live ragged path (budget-frac sweep) Aligns the cost-table pipeline with the DSpark paper's deployment method (and SGLang's implementation of it): instead of pinning uniform verify lengths and measuring the uniform path, force the PLANNER'S BUDGET to a fraction of the trimmable maximum and let the real confidence top-k assign the windows -- the measured cells then run the production ragged geometry end to end (token-major presentation, bucket fit, captured graphs), removing both approximations the pinned sweep leaned on (token-major ~ batch-major, and small-block-full-width ~ big-block- trimmed). Production side: - TLLM_DSPARK_FORCE_BUDGET_FRAC (init env) plus a runtime setter that mirrors the verify-len pin end to end: request/pending/adopt through the step's own allgather (payload element 7), collective_rpc set_dspark_budget_frac allowlisted, rank-consistent by construction. - Guards, mutual and fail-fast: frac and the verify-len pin refuse each other (either would silently shadow the other); frac is refused on engines that do not trim (static AND cap-accept); the STS collector refuses to start under a trimming frac (label censoring); an env frac on a non-trimming engine fails at worker init instead of logging FORCED and doing nothing. - The executed (padded) generation-token total now rides iteration stats: captured on rank 0 at queue time, broadcast per rank row by the ADP fanout (group-uniform on graph steps by construction), emitted as numGenerationTokens -- the M axis the sweep aligns on. Profiler side: - --fracs mode: ONE engine build for the whole sweep (production tier ladder, so the bucket grid is real), per-frac runtime RPC, per-step (bs, M, T) samples keyed by the group's agreed executed shape with all-rank alignment; per-cell warmup discard; thin incidental buckets dropped loudly instead of aborting the fit. - StepSample/CellStat carry measured_tokens (persisted in the samples JSONL); bucket values are exact breakpoints, so no M binning is needed anywhere. - The stale Repeat.cu geometry guard is removed: the assert it fenced was another symptom of the graph key describing a different token width than the batch carried, fixed by the published-bucket key; verified gone by pinned L=3 runs under both real and synthetic acceptance (64+ ragged steps each). Adversarially reviewed (20-agent workflow): 14 confirmed findings fixed before first successful run, including three high-severity (stats field on an unreachable path, tier-ladder collapse flattening the bucket grid, STS guard bypass via the env entry). Unit suite: 226 passed on B300. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/pyexecutor/adp_iter_stats.py | 11 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 93 +++++- tensorrt_llm/_torch/speculative/dspark.py | 8 + tensorrt_llm/_torch/speculative/dspark_sts.py | 13 + .../_torch/speculative/dspark_verify.py | 111 +++++++ tensorrt_llm/executor/base_worker.py | 18 ++ tensorrt_llm/executor/proxy.py | 3 +- tests/microbenchmarks/dspark_sps_profiler.py | 293 +++++++++++++++--- 8 files changed, 507 insertions(+), 43 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py b/tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py index dff8d078eedb..f0f011da724e 100644 --- a/tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py +++ b/tensorrt_llm/_torch/pyexecutor/adp_iter_stats.py @@ -51,6 +51,11 @@ class ADPIterStatsRecord: host_step_time_ms: Optional[float] = None prev_device_step_time_ms: Optional[float] = None gpu_forward_time_ms: Optional[float] = None + #: Executed (padded) generation-token total, captured on rank 0 and + #: broadcast to every rank row like the timings above. Group-uniform by + #: construction on graph steps: every ADP rank pads and fills to the + #: group's agreed shape, so rank 0's value IS the group's. + num_generation_tokens: Optional[int] = None class ADPIterStatsBuffer: @@ -80,6 +85,7 @@ def __init__(self) -> None: self._rank0_host_step_time_ms: Dict[int, Optional[float]] = {} self._rank0_prev_device_step_time_ms: Dict[int, Optional[float]] = {} self._rank0_gpu_forward_time_ms: Dict[int, Optional[float]] = {} + self._rank0_num_generation_tokens: Dict[int, Optional[int]] = {} self._oldest_iter: Optional[int] = None @staticmethod @@ -108,6 +114,7 @@ def queue( host_step_time_ms: Optional[float] = None, prev_device_step_time_ms: Optional[float] = None, gpu_forward_time_ms: Optional[float] = None, + num_generation_tokens: Optional[int] = None, ) -> None: """Queue local stats; rank 0 also keeps objects needed for fanout.""" payload = self.make_payload(stats) @@ -129,6 +136,7 @@ def queue( self._rank0_host_step_time_ms[iter_id] = host_step_time_ms self._rank0_prev_device_step_time_ms[iter_id] = prev_device_step_time_ms self._rank0_gpu_forward_time_ms[iter_id] = gpu_forward_time_ms + self._rank0_num_generation_tokens[iter_id] = num_generation_tokens def next_payload(self) -> Optional[RankIterStatsPayload]: """Return the oldest pending stats payload to piggyback.""" @@ -163,6 +171,7 @@ def _discard(self, iter_id: int, *, recompute_oldest: bool = True) -> None: self._rank0_host_step_time_ms.pop(iter_id, None) self._rank0_prev_device_step_time_ms.pop(iter_id, None) self._rank0_gpu_forward_time_ms.pop(iter_id, None) + self._rank0_num_generation_tokens.pop(iter_id, None) if recompute_oldest and iter_id == self._oldest_iter: self._recompute_oldest_iter() @@ -294,6 +303,7 @@ def finalize( host_step_time_ms = self._rank0_host_step_time_ms.get(iter_stats_iter) prev_device_step_time_ms = self._rank0_prev_device_step_time_ms.get(iter_stats_iter) gpu_forward_time_ms = self._rank0_gpu_forward_time_ms.get(iter_stats_iter) + num_generation_tokens = self._rank0_num_generation_tokens.get(iter_stats_iter) for rank_state in sorted(matching_states, key=lambda s: s.rank): rank = rank_state.rank @@ -306,6 +316,7 @@ def finalize( host_step_time_ms=host_step_time_ms, prev_device_step_time_ms=prev_device_step_time_ms, gpu_forward_time_ms=gpu_forward_time_ms, + num_generation_tokens=num_generation_tokens, ) ) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 11995f0c821d..f56f64b7eb23 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2231,7 +2231,8 @@ def _append_iter_stats(self, attention_dp_rank: Optional[int] = None, host_step_time_ms: Optional[float] = None, prev_device_step_time_ms: Optional[float] = None, - gpu_forward_time_ms: Optional[float] = None): + gpu_forward_time_ms: Optional[float] = None, + num_generation_tokens: Optional[int] = None): """Append one iteration's finalized stats to the export buffer. The normal Attention-DP path fans out rank-local rows before calling @@ -2300,6 +2301,16 @@ def _append_iter_stats(self, local_dict["prevDeviceStepTimeMS"] = prev_device_step_time_ms if gpu_forward_time_ms is not None: local_dict["gpuForwardTimeMS"] = gpu_forward_time_ms + # The executed (padded) generation-token total -- the M axis of + # the SPS cost table. Same value the iteration log prints in + # `states`; carried per rank so ragged sweeps can align steps on + # the group's agreed shape. + num_gen_tokens = (num_generation_tokens + if num_generation_tokens is not None else + (self.model_engine.iter_states + or {}).get('num_generation_tokens')) + if num_gen_tokens is not None: + local_dict["numGenerationTokens"] = int(num_gen_tokens) local_dict["schedulerMode"] = scheduler_mode local_dict["rank"] = self.dist.tp_rank # Buffer for _flush_iter_stats_synced; in-place tp_allgather would @@ -2318,6 +2329,8 @@ def _append_iter_stats(self, # [5] prev_device_step_time_ms: Optional[float] # [6] scheduler_mode: "overlap" | "non_overlap" # [7] gpu_forward_time_ms: Optional[float] + # [8] num_generation_tokens: Optional[int] -- executed (padded) + # generation-token total, rank0's value broadcast per row with self.stats_lock: if (not _stats_buffer_is_unbounded(self.max_stats_len) and len(self.stats) > self.max_stats_len): @@ -2325,7 +2338,7 @@ def _append_iter_stats(self, self.stats.append( (stats, req_stats, kv_iter_stats, attention_dp_rank, host_step_time_ms, prev_device_step_time_ms, scheduler_mode, - gpu_forward_time_ms)) + gpu_forward_time_ms, num_generation_tokens)) def _process_iter_stats( self, @@ -2390,7 +2403,9 @@ def _process_iter_stats( is_rank0=self.dist.rank == 0, host_step_time_ms=host_step_time_ms, prev_device_step_time_ms=prev_device_step_time_ms, - gpu_forward_time_ms=gpu_forward_time_ms) + gpu_forward_time_ms=gpu_forward_time_ms, + num_generation_tokens=(self.model_engine.iter_states + or {}).get('num_generation_tokens')) else: self._append_iter_stats( stats, @@ -3252,8 +3267,65 @@ def set_dspark_verify_len_pin(self, verify_len: Optional[int]) -> Optional[int]: "below the full block censors the acceptance label for every " "sample collected after it. Pin the full block, clear the " "pin, or restart without the collector.") + if verify_len is not None and (planner.forced_budget_frac is not None + or planner.pending_budget_frac() >= 0): + raise RuntimeError( + "a budget fraction is in force (or queued) on this engine; a " + "pin would silently shadow it. Clear the fraction first.") return planner.request_verify_len_pin(verify_len) + def set_dspark_budget_frac(self, + frac: Optional[float]) -> Optional[float]: + """Queue a DSpark verify-budget fraction, or ``None`` to clear it. + + Replaces the cost-table argmax with ``frac`` of the maximum trimmable + budget while the confidence top-k keeps assigning the windows -- the + production ragged path at a controlled load point, for SPS cost-table + sweeps. Validated here, applied on the next decode step through the + step's cross-rank allgather. + + Raises: + RuntimeError: if this executor has no DSpark confidence planner, + runs static verification, or is collecting STS calibration + data (a trimmed budget censors the acceptance labels). + ValueError: if ``frac`` is outside ``(0.0, 1.0]``. + """ + worker_fn = getattr(self.model_engine, "_get_spec_worker", None) + worker = worker_fn() if worker_fn else None + planner = getattr(worker, "verify_planner", None) if worker else None + if planner is None: + raise RuntimeError( + "no DSpark confidence planner on this executor; the budget " + "fraction only applies to a DSpark engine with confidence " + "scheduling enabled") + mode = getattr(worker, "ragged_verify_mode", None) + trims = (mode.trims_submitted_tokens if mode is not None else + self.model_engine.spec_config.enable_ragged_verify) + if not trims: + # Static never computes windows; cap-accept computes them but + # submits the full block, so a frac there cannot move the load + # point -- it would only cap live acceptance, which is not what a + # cost sweep asks for. + raise RuntimeError( + "this engine does not trim submitted tokens (static or " + "cap-accept mode); the budget fraction only moves the load " + "point in compact mode, so setting it here would change " + "nothing or change acceptance instead") + if frac is not None and (planner.forced_verify_len is not None + or planner.pending_verify_len_pin() >= 0): + raise RuntimeError( + "a verify-length pin is in force (or queued) on this engine; " + "the pin takes precedence and the fraction would be adopted, " + "logged, and never applied. Clear the pin first.") + if (worker.sts_recorder is not None and frac is not None + and float(frac) < 1.0): + raise RuntimeError( + "STS calibration collection is active on this engine; a " + "trimmed budget censors the acceptance label for every sample " + "collected after it. Use 1.0, clear the fraction, or restart " + "without the collector.") + return planner.request_budget_frac(frac) + def _dspark_confidence_draft_len(self, scheduled_batch: ScheduledRequests) -> int: """Draft length for this iteration, from the DSpark confidence planner. @@ -3344,14 +3416,19 @@ def _dspark_confidence_draft_len(self, local_max_window, 1 if scheduled_batch.can_run_cuda_graph else 0, planner.pending_verify_len_pin(), + planner.pending_budget_frac(), ]) - # The pin rides the allgather so every rank adopts the same value - # on the same step (the endpoint lands on one rank). -1 means - # nobody queued anything. + # The pin and the budget fraction ride the allgather so every rank + # adopts the same value on the same step (the endpoint lands on + # one rank). -1 means nobody queued anything. agreed_pin = next( (int(payload[5]) for payload in payloads if len(payload) > 5 and int(payload[5]) >= 0), -1) planner.adopt_verify_len_pin(agreed_pin) + agreed_frac = next( + (int(payload[6]) for payload in payloads + if len(payload) > 6 and int(payload[6]) >= 0), -1) + planner.adopt_budget_frac(agreed_frac) if not all(int(payload[0]) for payload in payloads): # All-or-nothing across ranks: `all_rank_num_tokens` is frozen # into the captured graph, so a mixed step's mismatch would @@ -3372,6 +3449,7 @@ def _dspark_confidence_draft_len(self, else: # Single-rank world: the local queue IS the group's agreement. planner.adopt_verify_len_pin(planner.pending_verify_len_pin()) + planner.adopt_budget_frac(planner.pending_budget_frac()) # --- Phase 3: no more collectives ------------------------------------- # @@ -5344,7 +5422,8 @@ def _fetch_new_requests( host_step_time_ms=record.host_step_time_ms, prev_device_step_time_ms=record. prev_device_step_time_ms, - gpu_forward_time_ms=record.gpu_forward_time_ms) + gpu_forward_time_ms=record.gpu_forward_time_ms, + num_generation_tokens=record.num_generation_tokens) all_ranks_num_active_requests = [ s.num_active_requests for s in all_rank_states ] diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index 8116971abfa7..eca97fa8b688 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -490,6 +490,14 @@ def _build_verify_planner(self, draft_model, block_size: int, max_batch: int) -> if self.spec_config.enable_ragged_verify else RaggedVerifyMode.STATIC) self.ragged_verify_mode = read_ragged_verify_mode(default=configured) + if (self.verify_planner is not None + and self.verify_planner.forced_budget_frac is not None + and not self.ragged_verify_mode.trims_submitted_tokens): + raise ValueError( + "TLLM_DSPARK_FORCE_BUDGET_FRAC is set but ragged verify mode " + f"'{self.ragged_verify_mode.value}' does not trim submitted " + "tokens; the fraction would be logged as FORCED and silently " + "never applied. Unset it or run compact mode.") # STS collection is env-gated; it costs one device->host copy per # decode step. from .dspark_sts import make_recorder_from_env diff --git a/tensorrt_llm/_torch/speculative/dspark_sts.py b/tensorrt_llm/_torch/speculative/dspark_sts.py index c910a5cc0c2d..7267d61bd35f 100644 --- a/tensorrt_llm/_torch/speculative/dspark_sts.py +++ b/tensorrt_llm/_torch/speculative/dspark_sts.py @@ -293,6 +293,19 @@ def make_recorder_from_env(*, block_size: int, rank: int = 0, f"{STS_COLLECT_ENV} is set with ragged verify mode 'compact'. Padded " f"verify rows corrupt the per-position prefix label. Collect in " f"static mode (TLLM_DSPARK_RAGGED_VERIFY_MODE=static).") + + frac = os.environ.get("TLLM_DSPARK_FORCE_BUDGET_FRAC", "").strip() + if frac: + try: + frac_full = float(frac) >= 1.0 + except ValueError: + frac_full = False + if not frac_full: + raise ValueError( + f"{STS_COLLECT_ENV} is set while TLLM_DSPARK_FORCE_BUDGET_FRAC=" + f"{frac!r} trims the verify budget; a trimmed window censors " + f"the acceptance label for every collected sample. Collect " + f"with the fraction unset or at 1.0.") logger.warning( f"DSpark STS collection ON ({STS_COLLECT_ENV}={stem!r}). This adds one " f"device->host copy per decode step; it is a calibration mode, not a " diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index bde6001718bc..735054d6ad80 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -39,6 +39,18 @@ #: verify window at this length, bypassing the planner. Not a serving knob. FORCE_VERIFY_LEN_ENV = "TLLM_DSPARK_FORCE_VERIFY_LEN" +#: Profiling-only budget override (SPS profiling on the live ragged path): +#: replaces the cost-table argmax with ``frac`` of the maximum trimmable +#: budget, while windows still flow through the real confidence top-k. Unlike +#: the verify-length pin this keeps the executed shape production-faithful +#: (ragged windows, bucket fit); the profiler sweeps it to visit off-diagonal +#: (bs, M) cells. Not a serving knob. +FORCE_BUDGET_FRAC_ENV = "TLLM_DSPARK_FORCE_BUDGET_FRAC" + +#: Wire encoding for the frac's cross-rank agreement: fractions travel as +#: fixed-point ints so the allgather payload stays homogeneous ints. +_FRAC_WIRE_SCALE = 10_000 + __all__ = ["DSparkVerifyPlanner"] @@ -83,6 +95,11 @@ def __init__( # step's cross-rank allgather so every rank starts pinning on the # same step; applying it locally would diverge the shape gate. self._pending_verify_len_pin: int = -1 + # Budget-fraction override (same lifecycle as the pin: env at init, + # runtime requests adopted through the step's allgather). Wire values: + # -1 nothing queued, 0 clear, >0 frac * _FRAC_WIRE_SCALE. + self._forced_budget_frac = self._read_forced_budget_frac() + self._pending_budget_frac: int = -1 self._host_buffer: Optional[torch.Tensor] = None self._copy_event: Optional[torch.cuda.Event] = None @@ -190,6 +207,81 @@ def validate_verify_len_pin(self, value: Optional[int]) -> Optional[int]: f"and the measured times would not be step costs") return value + def _read_forced_budget_frac(self) -> Optional[float]: + raw = os.environ.get(FORCE_BUDGET_FRAC_ENV, "").strip() + if not raw: + return None + try: + value = float(raw) + except ValueError: + raise ValueError( + f"{FORCE_BUDGET_FRAC_ENV}={raw!r} is not a float") from None + value = self.validate_budget_frac(value) + logger.warning( + f"DSpark: verify budget FORCED to {value} of the trimmable " + f"maximum by {FORCE_BUDGET_FRAC_ENV}. Windows still follow " + f"confidence; this is a profiling mode, not a serving " + f"configuration.") + return value + + def validate_budget_frac(self, value: Optional[float]) -> Optional[float]: + """Check a budget fraction; None clears the override.""" + if value is None: + return None + value = float(value) + if not 0.0 < value <= 1.0: + raise ValueError(f"budget fraction {value} outside (0.0, 1.0]") + return value + + @property + def forced_budget_frac(self) -> Optional[float]: + """The budget fraction currently in force, or None.""" + return self._forced_budget_frac + + @property + def forced_verify_len(self) -> Optional[int]: + """The verify-length pin currently in force, or None.""" + return self._forced_verify_len + + def request_budget_frac(self, value: Optional[float]) -> Optional[float]: + """Queue a budget-fraction override (``None`` clears) for the next step. + + Validated here but NOT applied: the step's allgather carries it so all + ranks adopt it together (see :meth:`adopt_budget_frac`). Returns the + wire-quantized value that will actually be adopted, so the caller sees + what applies rather than what it asked for. + """ + value = self.validate_budget_frac(value) + wire = (0 if value is None else + max(1, int(round(value * _FRAC_WIRE_SCALE)))) + self._pending_budget_frac = wire + return None if value is None else min(1.0, wire / _FRAC_WIRE_SCALE) + + def pending_budget_frac(self) -> int: + """The queued request as a wire value: -1 none, 0 clear, >0 fixed-point.""" + return int(self._pending_budget_frac) + + def adopt_budget_frac(self, wire_value: int) -> None: + """Apply the group's agreed budget fraction from the allgathered payload.""" + wire_value = int(wire_value) + if wire_value < 0: + return + # Same non-atomic compare-and-clear as the pin above: an RPC-thread + # request landing between the compare and the store is dropped. Both + # knobs are single-operator profiling controls, so the race is + # tolerated for both rather than fixed for one. + if self._pending_budget_frac == wire_value: + self._pending_budget_frac = -1 + new_frac = (None if wire_value == 0 else + min(1.0, wire_value / _FRAC_WIRE_SCALE)) + if new_frac == self._forced_budget_frac: + return + self._forced_budget_frac = new_frac + logger.warning( + f"DSpark: verify budget fraction set to {new_frac} at runtime. " + f"This bypasses the cost table and is a profiling mode, not a " + f"serving configuration.") + def request_verify_len_pin(self, value: Optional[int]) -> Optional[int]: """Queue a pin (or ``None`` to clear) for the next decode step. @@ -319,6 +411,25 @@ def decide_verify_lens( # path still runs the cross-rank agreement below. self.stats["forced_steps"] += 1 lens = [int(self._forced_verify_len)] * int(num_gen_requests) + elif self._forced_budget_frac is not None: + # Also ahead of the cost-table gate (the frac sweep is how the + # table gets built). Unlike the pin, the windows stay real: the + # forced fraction only replaces the argmax'd budget, and the + # confidence top-k spends it exactly as in production. + selected = self._gather_rows(num_gen_requests=num_gen_requests, + rows=rows) + if selected is None: + return None + self._note_snapshot_stats(selected, rows) + survival = compute_survival(self.apply_calibration(selected)) + trimmable = (self.cfg.resolved_max_verify_len + - self.cfg.min_verify_len) + budget = int(round(self._forced_budget_frac + * int(num_gen_requests) * trimmable)) + self.stats["forced_budget_steps"] = ( + self.stats.get("forced_budget_steps", 0) + 1) + lens = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=self.cfg).tolist() else: if self.cost_table is None or self.cost_table.is_flat: self.stats["fallback_flat_cost"] += 1 diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 33aeab1ce4df..de0f774c8ded 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -886,6 +886,18 @@ def set_dspark_verify_len_pin(self, raise RuntimeError("engine is not initialized") return self.engine.set_dspark_verify_len_pin(verify_len) + def set_dspark_budget_frac(self, + frac: Optional[float]) -> Optional[float]: + """Set (or clear) the DSpark verify-budget fraction on this executor. + + Same delivery contract as :meth:`set_dspark_verify_len_pin`: a scalar + queued on rank 0 and handed to every rank by the decode loop's own + per-step allgather. + """ + if self.engine is None: + raise RuntimeError("engine is not initialized") + return self.engine.set_dspark_budget_frac(frac) + def sleep(self, sleep_tags: list[str]) -> None: """Release GPU virtual memory for the specified memory type tags. @@ -1014,6 +1026,7 @@ def _stats_serializer(stats) -> str: prev_device_step_time_ms = stats[5] if len(stats) > 5 else None scheduler_mode = stats[6] if len(stats) > 6 else None gpu_forward_time_ms = stats[7] if len(stats) > 7 else None + num_generation_tokens = stats[8] if len(stats) > 8 else None stats_dict = json.loads(iteration_stats.to_json_str()) # Always tag the row so Dynamo's adapter can read @@ -1056,6 +1069,11 @@ def _stats_serializer(stats) -> str: # CPU wall. Set per-record so consumers do not need server config. if scheduler_mode is not None: stats_dict["schedulerMode"] = scheduler_mode + # Executed (padded) generation-token total -- the M axis of the DSpark + # SPS cost table. Rank 0's value, broadcast per row by the ADP fanout; + # group-uniform on graph steps by construction. + if num_generation_tokens is not None: + stats_dict["numGenerationTokens"] = int(num_generation_tokens) # Convert back to JSON string return json.dumps(stats_dict) diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index b1080a4c3a31..f9cf259ea43a 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -64,7 +64,8 @@ # set_dspark_verify_len_pin() needs no control communicator at all: it queues # a scalar that the decode loop's own per-step allgather hands to every rank, # so rank-0 receiving it is sufficient and the loop is never paused. -_MULTI_RANK_ALLOWED_METHODS: frozenset[str] = frozenset({"sleep", "wakeup", "set_dspark_verify_len_pin"}) +_MULTI_RANK_ALLOWED_METHODS: frozenset[str] = frozenset( + {"sleep", "wakeup", "set_dspark_verify_len_pin", "set_dspark_budget_frac"}) def _check_collective_rpc_guard( diff --git a/tests/microbenchmarks/dspark_sps_profiler.py b/tests/microbenchmarks/dspark_sps_profiler.py index 43691fc441da..1a53c92ae642 100644 --- a/tests/microbenchmarks/dspark_sps_profiler.py +++ b/tests/microbenchmarks/dspark_sps_profiler.py @@ -268,27 +268,38 @@ class StepSample: verify_len: int step_time_ms: float iteration: int = -1 + #: Ragged (frac-sweep) steps: the group's executed token total, read from + #: the stats rows. None on pinned sweeps, where M is derived instead. + measured_tokens: Optional[int] = None @property def total_verify_tokens(self) -> int: - """``M`` for this step. Derived, never stored, so it cannot drift.""" + """``M`` for this step. Measured when the shape was ragged; derived + from the pin otherwise, so it cannot drift from the label.""" + if self.measured_tokens is not None: + return int(self.measured_tokens) return total_verify_tokens(self.batch_size, self.verify_len) def to_json(self) -> dict: - return { + payload = { "batch_size": int(self.batch_size), "verify_len": int(self.verify_len), "step_time_ms": float(self.step_time_ms), "iteration": int(self.iteration), } + if self.measured_tokens is not None: + payload["measured_tokens"] = int(self.measured_tokens) + return payload @classmethod def from_json(cls, raw: dict) -> "StepSample": + measured = raw.get("measured_tokens") return cls( batch_size=int(raw["batch_size"]), verify_len=int(raw["verify_len"]), step_time_ms=float(raw["step_time_ms"]), iteration=int(raw.get("iteration", -1)), + measured_tokens=None if measured is None else int(measured), ) @@ -302,9 +313,14 @@ class CellStat: num_samples: int p10_ms: float p90_ms: float + #: Ragged (frac-sweep) cells: the executed token total the samples agreed + #: on. None on pinned sweeps. + measured_tokens: Optional[int] = None @property def total_verify_tokens(self) -> int: + if self.measured_tokens is not None: + return int(self.measured_tokens) return total_verify_tokens(self.batch_size, self.verify_len) def to_json(self) -> dict: @@ -370,6 +386,53 @@ def aligned_steps_from_stats( return aligned +def ragged_steps_from_stats( + rows: Sequence[dict], + *, + expected_ranks: int = 1, + timing_key: str = HOST_STEP_TIME_KEY, +) -> List[Tuple[int, int, int, float]]: + """``(iteration, num_gen_requests, executed_tokens, step_time_ms)`` steps. + + The ragged (frac-sweep) sibling of :func:`aligned_steps_from_stats`: on top + of its alignment rules, every rank must also report the same + ``numGenerationTokens``. Under attention-DP the executed token total is the + group's agreed bucket -- every rank pads and fills to the same shape -- so + a step whose ranks disagree is a ramp/fallback step whose shape is not one + shape at all, and averaging it in would file its cost under a bucket it did + not run. + """ + by_iteration: Dict[int, List[dict]] = {} + for row in rows: + iteration = row.get("iter") + if iteration is None: + continue + by_iteration.setdefault(int(iteration), []).append(row) + + kept: List[Tuple[int, int, int, float]] = [] + for iteration in sorted(by_iteration): + group = by_iteration[iteration] + if len({int(r.get("attentionDpRank", 0)) for r in group}) != int(expected_ranks): + continue + batching = [row.get("inflightBatchingStats") or {} for row in group] + if any(int(stats.get("numContextRequests", 0) or 0) for stats in batching): + continue + gen_counts = {int(stats.get("numGenRequests", 0) or 0) for stats in batching} + if len(gen_counts) != 1 or (num_gen := gen_counts.pop()) <= 0: + continue + token_counts = {row.get("numGenerationTokens") for row in group} + if len(token_counts) != 1: + continue + tokens = token_counts.pop() + if tokens is None or int(tokens) <= 0: + continue + timing = group[0].get(timing_key) + if timing is None or float(timing) <= 0.0: + continue + kept.append((iteration, num_gen, int(tokens), float(timing))) + return kept + + def explain_alignment( rows: Sequence[dict], *, @@ -470,18 +533,27 @@ def summarize_cells( disconnect the batch-size ladder, which changes the fit rather than just widening its error bars. """ - grouped: Dict[Tuple[int, int], List[StepSample]] = {} + grouped: Dict[Tuple[int, int, Optional[int]], List[StepSample]] = {} for sample in samples: - grouped.setdefault((int(sample.batch_size), int(sample.verify_len)), []).append(sample) + key = (int(sample.batch_size), int(sample.verify_len), + None if sample.measured_tokens is None else int(sample.measured_tokens)) + grouped.setdefault(key, []).append(sample) requested = None if expected_cells is None else { (int(b), int(l)) for b, l in expected_cells} dropped: List[str] = [] cells: List[CellStat] = [] - for (batch_size, verify_len), cell_samples in sorted(grouped.items()): + for (batch_size, verify_len, measured_tokens), cell_samples in sorted( + grouped.items(), key=lambda kv: (kv[0][0], kv[0][1], kv[0][2] or -1)): ordered = sorted(cell_samples, key=lambda s: s.iteration) steady = [s.step_time_ms for s in ordered[int(warmup_steps) :]] if len(steady) < int(min_samples): + # Ragged sweeps visit whichever buckets the planner picked, so a + # thin incidental bucket is normal coverage noise, not a failed + # cell: drop it loudly rather than aborting the fit. + if measured_tokens is not None: + dropped.append(f"(bs={batch_size}, M={measured_tokens}, n={len(steady)})") + continue if requested is not None and (batch_size, verify_len) not in requested: dropped.append(f"(bs={batch_size}, L={verify_len}, n={len(steady)})") continue @@ -503,6 +575,7 @@ def summarize_cells( num_samples=int(values.size), p10_ms=float(np.percentile(values, 10)), p90_ms=float(np.percentile(values, 90)), + measured_tokens=measured_tokens, ) ) if dropped: @@ -1077,6 +1150,10 @@ class SweepConfig: seed_table_path: Optional[str] = None #: Restore the original sweep geometry; see the note in _build_llm. block_follows_verify_len: bool = False + #: Frac sweep (run_frac_sweep): budget fractions to visit on the live + #: ragged path. When set, verify_lens is pinned to [max_draft_len] by the + #: CLI and the per-cell M comes from measurement instead of the pin. + fracs: Optional[List[float]] = None attn_backend: str = "TRTLLM" disable_overlap_scheduler: bool = False #: On by default because the deployment being profiled runs with it on, and @@ -1133,26 +1210,26 @@ def max_tokens_for(self, verify_len: int, def validate(self) -> None: """Catch the geometry mistakes that would otherwise fail mid-sweep.""" + if self.fracs is not None: + bad = [f for f in self.fracs if not 0.0 < float(f) <= 1.0] + if bad: + raise SweepGeometryError( + f"--fracs values must be in (0, 1], got {sorted(bad)}") + if self.block_follows_verify_len: + raise SweepGeometryError( + "--fracs measures the live ragged path at the constant " + "block; --block-follows-verify-len contradicts it") if not self.batch_sizes or not self.verify_lens: raise SweepGeometryError("both --batch-sizes and --verify-lens must be non-empty") if min(self.verify_lens) < 1: raise SweepGeometryError("verify lengths must be >= 1 (the planner's floor)") - if (not self.block_follows_verify_len - and min(self.verify_lens) < int(self.max_draft_len)): - # Constant-block with a pinned window below the block currently - # dies on a device-side assert (Repeat.cu output_size vs - # sum(repeats)) at the first short cell -- AFTER minutes of engine - # build. Refusing here is this docstring's whole job; letting the - # default geometry walk into a known assert is not a validation. - raise SweepGeometryError( - f"constant-block geometry with verify lengths " - f"{sorted(self.verify_lens)} below the block " - f"{self.max_draft_len} runs the regime that currently trips " - f"the Repeat.cu output_size assert. Pass " - f"--block-follows-verify-len (each cell drafts exactly what " - f"it verifies), or sweep only L == {self.max_draft_len}, or " - f"use --base-url / --from-iter-log which measure a deployment " - f"instead of building one.") + # Constant-block with a pinned window below the block used to die on a + # device-side assert (Repeat.cu output_size vs sum(repeats)); a guard + # here refused that geometry. The assert was another symptom of the + # graph key describing a different token width than the batch carried, + # fixed by publishing the fit's bucket into the key; verified gone + # 2026-08-06 by pinned L=3 runs under both real and synthetic + # acceptance (64+ steps each, no assert). # A request that hits max_seq_len is evicted mid-cell, so the batch # drains and the steady window shrinks -- which surfaces much later as a # confusing "not enough steady samples". @@ -1222,12 +1299,14 @@ def _prepare_environment(config: SweepConfig) -> None: # only the axis the planner actually controls. os.environ[RAGGED_VERIFY_MODE_ENV] = RaggedVerifyMode.COMPACT.value if config.pin_acceptance: - from .interface import FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR + from tensorrt_llm._torch.speculative.interface import ( + FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR) os.environ[FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR] = repr(float(config.pin_acceptance)) -def _build_llm(config: SweepConfig, verify_len: int): +def _build_llm(config: SweepConfig, verify_len: int, *, + pin_verify_len: bool = True): """One engine per verify length; see the module docstring for why.""" from tensorrt_llm import LLM from tensorrt_llm.llmapi import CudaGraphConfig, DSparkDecodingConfig, KvCacheConfig @@ -1272,14 +1351,22 @@ def _build_llm(config: SweepConfig, verify_len: int): # their label does not describe. assert_swept_shape refuses such a sweep # downstream (verify_len_hist vs the nominal L), which is why it gates the # fit rather than merely warning. - os.environ[FORCE_VERIFY_LEN_ENV] = str(int(verify_len)) + # Frac sweeps leave the pin unset: their windows come from the live + # planner (that is the experiment), and the budget fraction travels by + # RPC + the step's own allgather, which reaches pre-spawned worlds the + # env mutation below cannot. + if pin_verify_len: + os.environ[FORCE_VERIFY_LEN_ENV] = str(int(verify_len)) + else: + os.environ.pop(FORCE_VERIFY_LEN_ENV, None) # Which knob moves M. # # constant block (default): the block stays at the checkpoint's trained # length for every cell and only the verify window moves, which is exactly - # what the deployment does. It is the right experiment -- but it makes the - # whole batch run with max(verify_len) < block_size, and the ragged packing - # currently asserts in that regime (Repeat.cu output_size vs sum(repeats)). + # what the deployment does. The all-short-window regime this creates used + # to trip the Repeat.cu output_size assert; that was the graph-key/batch + # width divergence, fixed with the published-bucket key and verified gone + # 2026-08-06. # # block follows verify_len: the original design. Never produces # max(lens) < block, so it runs today, at the price of shrinking the DRAFT @@ -1323,7 +1410,12 @@ def _build_llm(config: SweepConfig, verify_len: int): # every downstream path (bucket choice, padding reservation, graph # selection) runs exactly as it does in production. confidence_verify_len_tiers=( - [int(verify_len)] if config.enable_confidence_scheduling else None + # Frac sweeps need the PRODUCTION ladder: the token-bucket grid is + # derived from it, and a one-rung ladder would collapse every + # captured bucket to the full block -- the sweep would then only + # ever measure M = bs * (block + 1) regardless of the frac. + (None if not pin_verify_len else [int(verify_len)]) + if config.enable_confidence_scheduling else None ), ) llm_kwargs = dict( @@ -1360,8 +1452,15 @@ def _build_llm(config: SweepConfig, verify_len: int): return LLM(**llm_kwargs) -def _run_cell(llm, config: SweepConfig, *, batch_size: int, verify_len: int) -> List[StepSample]: - """Hold ``batch_size`` generating requests per rank and record every step.""" +def _run_cell(llm, config: SweepConfig, *, batch_size: int, verify_len: int, + measure_ragged: bool = False) -> List[StepSample]: + """Hold ``batch_size`` generating requests per rank and record every step. + + With ``measure_ragged`` the per-step token total is READ from the stats + (the group's agreed bucket) instead of derived from the pin, and samples + carry it in ``measured_tokens``; one cell then legitimately yields samples + at several M values. + """ from tensorrt_llm import SamplingParams rng = random.Random(config.seed + batch_size * 1000 + verify_len) @@ -1471,6 +1570,34 @@ def _drain() -> None: f"window may have closed early and the cell lost its tail", file=sys.stderr) + if measure_ragged: + ragged = ragged_steps_from_stats( + rows, expected_ranks=config.dp_size, timing_key=config.timing_key) + kept_ragged = [s for s in ragged if s[1] == int(batch_size)] + if not kept_ragged: + print( + f"[dspark-sps] bs={batch_size:>5} (ragged) kept nothing -- " + + explain_alignment(rows, expected_ranks=config.dp_size, + timing_key=config.timing_key) + + "; also requires numGenerationTokens on every rank row", + file=sys.stderr) + return [] + # Warmup is discarded per CELL by iteration order, here rather than in + # summarize_cells: one ragged cell yields several (bs, M) groups, and a + # per-group trim would re-charge the warmup against every bucket. + cutoff = sorted(s[0] for s in kept_ragged) + cutoff = cutoff[min(int(config.warmup_steps), len(cutoff) - 1)] + return [ + StepSample( + batch_size=int(batch_size), + verify_len=int(verify_len), + step_time_ms=step_time, + iteration=iteration, + measured_tokens=tokens, + ) + for iteration, num_gen, tokens, step_time in kept_ragged + if iteration >= cutoff + ] aligned = aligned_steps_from_stats( rows, expected_ranks=config.dp_size, timing_key=config.timing_key ) @@ -1536,6 +1663,70 @@ def run_sweep(config: SweepConfig) -> List[StepSample]: return samples +def run_frac_sweep(config: SweepConfig) -> List[StepSample]: + """Measure ``(batch_size, frac)`` cells on the LIVE ragged path. + + One engine build for the whole sweep: the block stays at the checkpoint's + trained length, no verify-length pin is set, and per cell the budget + fraction is switched at runtime (``set_dspark_budget_frac`` -- rank 0's + queue, delivered by the decode loop's own allgather). Windows come from + the real confidence top-k and the executed shape is the group's agreed + bucket, so the cells sample exactly the geometry production runs -- + including the token-major presentation and the bucket fit -- instead of + the pinned-uniform approximation. M is read per step from the stats + (``measured_tokens``), so one cell contributes samples at every bucket the + planner visited. + """ + config.validate() + _prepare_environment(config) + fracs = sorted({float(f) for f in (config.fracs or [])}) + if not fracs: + raise ValueError("run_frac_sweep requires a non-empty --fracs") + print( + f"[dspark-sps] building engine once: block_size={config.max_draft_len} " + f"(constant), live planner windows, frac sweep {fracs} ...", + file=sys.stderr, + ) + samples: List[StepSample] = [] + llm = _build_llm(config, int(config.max_draft_len), pin_verify_len=False) + + def _set_frac(value): + executor = llm._executor + rpc = getattr(executor, "collective_rpc", None) + if rpc is not None: + return rpc("set_dspark_budget_frac", args=(value,)) + # Single-process executor path: the worker IS the executor. + return executor.set_dspark_budget_frac(value) + + try: + for frac in fracs: + _set_frac(float(frac)) + for batch_size in sorted(config.batch_sizes): + started = time.time() + cell: List[StepSample] = [] + for _ in range(max(1, int(config.repeats))): + cell.extend( + _run_cell(llm, config, batch_size=batch_size, + verify_len=int(config.max_draft_len), + measure_ragged=True)) + buckets = sorted({s.measured_tokens for s in cell}) + print( + f"[dspark-sps] bs={batch_size:>5} frac={frac:.2f} " + f"aligned_steps={len(cell):>5} buckets={buckets[:8]}" + f"{' ...' if len(buckets) > 8 else ''} " + f"({time.time() - started:.1f}s)", + file=sys.stderr, + ) + samples.extend(cell) + finally: + try: + _set_frac(None) + except Exception: # noqa: BLE001 - shutdown path, engine may be gone + pass + llm.shutdown() + return samples + + # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- @@ -1668,6 +1859,14 @@ def _build_parser() -> argparse.ArgumentParser: help="Uniform verify lengths to sweep. Defaults to 1..--max-draft-len. " "Each value costs one engine build.", ) + sweep.add_argument( + "--fracs", + default=None, + help="Budget fractions in (0, 1] to sweep on the LIVE ragged path " + "(one engine build total; windows from the real confidence top-k, " + "per-step M measured from the group's agreed bucket). Mutually " + "exclusive with --verify-lens and --block-follows-verify-len.", + ) sweep.add_argument("--input-len", type=int, default=1024, help="Prompt length in tokens.") sweep.add_argument("--block-follows-verify-len", action="store_true", help="Shrink the drafted block with the verify window " @@ -2218,11 +2417,22 @@ def main(argv: Optional[Sequence[str]] = None) -> int: with open(args.extra_llm_api_options, encoding="utf-8") as handle: extra = yaml.safe_load(handle) or {} + fracs = _float_list(args.fracs) if args.fracs else None + if fracs: + if args.verify_lens: + raise SystemExit("--fracs and --verify-lens are mutually " + "exclusive; the frac sweep's M is measured, " + "not pinned") + # The frac sweep runs the constant block with live windows; the + # nominal verify_len only labels the samples and sizes the token + # budget, and max_draft_len satisfies every geometry guard. + verify_lens = [int(args.max_draft_len)] config = SweepConfig( model=args.model, speculative_model=args.speculative_model or args.model, batch_sizes=batch_sizes, verify_lens=verify_lens, + fracs=fracs, tp_size=args.tp_size, ep_size=args.ep_size, enable_attention_dp=args.enable_attention_dp, @@ -2247,9 +2457,15 @@ def main(argv: Optional[Sequence[str]] = None) -> int: seed=args.seed, extra_llm_api_options=extra, ) - samples = run_sweep(config) - expected_cells = frozenset( - (int(b), int(l)) for b in batch_sizes for l in verify_lens) + if fracs: + samples = run_frac_sweep(config) + # Which buckets a frac visits is the planner's choice, not a + # promise; thin incidental buckets are dropped by summarize_cells. + expected_cells = None + else: + samples = run_sweep(config) + expected_cells = frozenset( + (int(b), int(l)) for b in batch_sizes for l in verify_lens) _write_samples(samples_out, samples) print(f"[dspark-sps] wrote {len(samples)} raw samples to {samples_out}", file=sys.stderr) provenance = { @@ -2263,10 +2479,17 @@ def main(argv: Optional[Sequence[str]] = None) -> int: "acceptance_pinned": bool(config.pin_acceptance), "overlap_scheduler": not config.disable_overlap_scheduler, "confidence_scheduling": config.enable_confidence_scheduling, - "ragged_verify_mode": RaggedVerifyMode.STATIC.value, + # _prepare_environment pins COMPACT for both sweep flavours; the + # pinned sweep's uniform shape comes from the pin, not the mode. + "ragged_verify_mode": RaggedVerifyMode.COMPACT.value, + **({"fracs": [float(f) for f in fracs]} if fracs else {}), } - cells = summarize_cells(samples, warmup_steps=args.warmup_steps, + cells = summarize_cells(samples, + # Frac cells discard warmup inside _run_cell (per + # cell, by iteration); trimming again here would + # re-charge it against every (bs, M) bucket. + warmup_steps=(0 if args.fracs else args.warmup_steps), min_samples=args.min_samples, expected_cells=expected_cells) if not cells: From 08e1eea7cca11a8238b01e96955287e81bd9409f Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Thu, 6 Aug 2026 08:46:25 -0700 Subject: [PATCH 101/177] [None][feat] Device-side fill_bucket and row-map construction (capture-safe) The host round-robin in fill_bucket decomposes into full cycles (every row with headroom gains one) plus one partial cycle (the first spare headroom rows in index order gain one), so a statically-unrolled grant = headroom & (cumsum(headroom) <= spare) loop of at most max_verify_len - 1 cycles per phase reproduces it token-for-token with no host sync. The row maps prepare() stages from host lists are a pure function of the filled lens: correction o - v + 1 composed with a kv_lens gather is exactly refresh_ragged_row_kv_lens. Algorithms only; no integration yet. Parity is exact equality against the host implementations over the exhaustive small space and production-shaped randomized cases. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/speculative/dspark_ragged.py | 88 +++++++++++++ .../test_dspark_device_fill_parity.py | 118 ++++++++++++++++++ 2 files changed, 206 insertions(+) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py diff --git a/tensorrt_llm/_torch/speculative/dspark_ragged.py b/tensorrt_llm/_torch/speculative/dspark_ragged.py index 255fcfae6b04..64caa79ff7af 100644 --- a/tensorrt_llm/_torch/speculative/dspark_ragged.py +++ b/tensorrt_llm/_torch/speculative/dspark_ragged.py @@ -41,6 +41,8 @@ "ragged_gather_index_lists", "RaggedCaptureShape", "choose_ragged_capture_shape", + "fill_bucket_device", + "build_row_maps_device", ] @@ -358,6 +360,92 @@ def choose_ragged_capture_shape( f"together rather than splitting its graph keys") +def fill_bucket_device( + verify_lens: torch.Tensor, + *, + num_real: torch.Tensor, + graph_num_tokens: int, + max_verify_len: int, +) -> 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). + 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.ones(padded_bs, dtype=torch.int32, device=device), + ) + spare = graph_num_tokens - lens.sum() + for phase_mask in (is_real, ~is_real): + 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. + + The token-major row maps prepare() stages from host lists decompose into + a pure function of the filled ``verify_lens``: token ``t`` of a request + with window ``v`` at within-request offset ``o`` attends KV extent + ``kv_len - v + 1 + o``, i.e. correction ``o - v + 1``. Composing the + returned correction with a ``kv_lens`` gather is exactly + ``refresh_ragged_row_kv_lens``, so the caller reuses that op unchanged. + + ``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]]: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py new file mode 100644 index 000000000000..661903dc8e1e --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The device fill/row-map path must reproduce the host path token-for-token. + +``fill_bucket_device`` and ``build_row_maps_device`` are the capture-safe +replacements for ``RaggedVerifyLayout.fill_bucket`` and the host row-map +staging in prepare(). Any divergence in the per-row allocation or the +per-token KV correction is an attention-reads-wrong-KV bug that surfaces as +silent quality loss (or an IMA), never as an exception -- so the parity here +is exact equality against the host implementation over the whole feasible +space, not spot checks. +""" + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_ragged import ( + RaggedVerifyLayout, build_row_maps_device, fill_bucket_device) + + +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 from host lists: + token o of a request with window v gets correction o - v + 1, so the + gathered 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 From 116bb80f04c57c8f54749f63127d7ca3dbbe4638 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Thu, 6 Aug 2026 09:09:02 -0700 Subject: [PATCH 102/177] [None][feat] Assemble the device window-selection prologue select_windows_device mirrors the host chain (_gather_rows -> calibrate -> survival -> topk -> fill) with the lag removed: it reads the LIVE slot-indexed confidence buffer the previous draft scattered into, so the ranking is fresh while the capacity knobs (padded_bs, bucket, budget) stay host-decided from the lagged snapshot -- the paper's dual-timescale split. schedule_verify_lens_topk now honours a 0-d tensor budget without a host sync (the value stays device-resident through the rank < budget cut); the int path is unchanged. Stale-stamp rows degrade to the neutral row, the same fail-open the host planner applies. Pure algorithm, no integration yet. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../speculative/dspark_device_select.py | 159 ++++++++++++++++++ .../_torch/speculative/dspark_schedule.py | 13 +- .../test_dspark_device_fill_parity.py | 87 ++++++++++ 3 files changed, 256 insertions(+), 3 deletions(-) create mode 100644 tensorrt_llm/_torch/speculative/dspark_device_select.py diff --git a/tensorrt_llm/_torch/speculative/dspark_device_select.py b/tensorrt_llm/_torch/speculative/dspark_device_select.py new file mode 100644 index 000000000000..586e888ab69d --- /dev/null +++ b/tensorrt_llm/_torch/speculative/dspark_device_select.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Device-side verify-window selection: the in-graph prologue. + +The host planner decides windows from a confidence snapshot that is one +iteration old (the batch reshuffles across the lag, and the read itself must +not sync). This module ranks the block that is about to be verified by its +OWN confidence instead: everything from the slot-indexed gather to the +per-token row maps is a pure tensor function, so it can run at the head of +the captured graph, after the previous step's draft wrote the confidence +buffer and before this step's verify attention reads the layout. + +The split of responsibilities follows the paper's dual-timescale design: + +* the CAPACITY -- ``(padded_bs, bucket, budget)`` -- stays host-decided from + the lagged snapshot, because it is the CUDA-graph key and the attention-DP + agreement payload, both of which must exist before launch; +* the RANKING -- which requests win the budget -- is computed here from + fresh confidence, at zero staleness. + +Nothing in this module raises on data: a captured graph cannot branch. +Feasibility (the bucket bounds) is checkable from host-known scalars before +launch; per-row staleness degrades to the neutral row (verify everything for +that request), mirroring the host planner's fail-open semantics. +""" + +from dataclasses import dataclass +from typing import Callable, Optional, Tuple + +import torch + +from .dspark_ragged import (build_qo_indptr, build_row_maps_device, + fill_bucket_device) +from .dspark_schedule import (NEUTRAL_CONFIDENCE_LOGIT, DSparkScheduleConfig, + compute_survival, schedule_verify_lens_topk) + +__all__ = ["DeviceWindowResult", "select_windows_device"] + + +@dataclass +class DeviceWindowResult: + """Everything the ragged layout consumers need, all device-resident. + + Attributes: + verify_lens: ``[padded_bs]`` int32 token windows (bonus included), + summing to exactly ``graph_num_tokens``; pad rows carry their + fill. + qo_indptr: ``[padded_bs + 1]`` int32 exclusive prefix sum. + req_idx: ``[graph_num_tokens]`` int64 owning-row per packed token. + kv_correction: ``[graph_num_tokens]`` int32; composed with a + ``kv_lens`` gather it yields each token's KV extent + (``refresh_ragged_row_kv_lens``). + """ + + verify_lens: torch.Tensor + qo_indptr: torch.Tensor + req_idx: torch.Tensor + kv_correction: torch.Tensor + + +def select_windows_device( + *, + confidence_logits: torch.Tensor, + slot_idx: torch.Tensor, + num_real: torch.Tensor, + budget: torch.Tensor, + graph_num_tokens: int, + cfg: DSparkScheduleConfig, + apply_calibration: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, + stamp: Optional[torch.Tensor] = None, + expected_stamp: Optional[torch.Tensor] = None, +) -> DeviceWindowResult: + """Rank the batch by fresh confidence and pack it into the agreed bucket. + + Mirrors the host chain ``_gather_rows -> apply_calibration -> + compute_survival -> schedule_verify_lens_topk -> fill_bucket`` with the + lag removed: ``confidence_logits`` is the LIVE slot-indexed buffer the + previous draft pass scattered into, not a staged snapshot. + + Args: + confidence_logits: ``[num_slots, K]`` raw logits, slot-indexed. + slot_idx: ``[padded_bs]`` int64 buffer row per batch position; + entries at or beyond ``num_real`` are never read meaningfully + (point them at any valid row, e.g. 0). + num_real: 0-d integer tensor; rows past it are padding. + budget: 0-d integer tensor, verify tokens above the floor -- the + host-decided (lagged) capacity knob, written before launch. + graph_num_tokens: the captured token bucket (capture constant). + cfg: scheduling bounds; ``resolved_max_verify_len + 1`` is the + per-row token ceiling. + apply_calibration: logits -> probabilities; sigmoid when None. + stamp: optional ``[num_slots]`` last-writer stamps; with + ``expected_stamp`` (0-d), rows whose stamp mismatches fall back + to the neutral row (full survival: verify everything), the same + fail-open the host planner applies to unknown slots. + + Returns: + :class:`DeviceWindowResult`; every tensor is derived without a host + sync, so the call is capture-safe end to end. + """ + if confidence_logits.dim() != 2: + raise ValueError( + f"confidence_logits must be [num_slots, K], got " + f"{tuple(confidence_logits.shape)}") + device = confidence_logits.device + padded_bs = slot_idx.numel() + + selected = confidence_logits.index_select(0, slot_idx.to(torch.long)) + if stamp is not None: + if expected_stamp is None: + raise ValueError("stamp requires expected_stamp") + stale = stamp.index_select(0, slot_idx.to(torch.long)) != expected_stamp + selected = torch.where( + stale.unsqueeze(1), + torch.full_like(selected, NEUTRAL_CONFIDENCE_LOGIT), + selected, + ) + + calibrate = apply_calibration or torch.sigmoid + survival = compute_survival(calibrate(selected)) + # Pad rows must win no budget: zero survival sits below survival_eps, so + # the ranking never selects them and their windows stay at the floor + # until the fill tops them up. + is_real = torch.arange(padded_bs, device=device) < num_real + survival = torch.where(is_real.unsqueeze(1), survival, + torch.zeros_like(survival)) + + scheduled = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=cfg) + # The scheduler counts drafted positions; the token window adds the + # bonus/anchor, matching every host fill_bucket callsite. + token_lens = scheduled + 1 + max_token_len = int(cfg.resolved_max_verify_len) + 1 + filled = fill_bucket_device( + token_lens, + num_real=num_real, + graph_num_tokens=graph_num_tokens, + max_verify_len=max_token_len, + ) + req_idx, correction = build_row_maps_device( + filled, graph_num_tokens=graph_num_tokens) + return DeviceWindowResult( + verify_lens=filled, + qo_indptr=build_qo_indptr(filled), + req_idx=req_idx, + kv_correction=correction, + ) diff --git a/tensorrt_llm/_torch/speculative/dspark_schedule.py b/tensorrt_llm/_torch/speculative/dspark_schedule.py index b025f82c3382..643585b32e63 100644 --- a/tensorrt_llm/_torch/speculative/dspark_schedule.py +++ b/tensorrt_llm/_torch/speculative/dspark_schedule.py @@ -26,6 +26,7 @@ """ from dataclasses import dataclass +from typing import Union import torch @@ -106,7 +107,7 @@ def compute_survival(confidence: torch.Tensor) -> torch.Tensor: def schedule_verify_lens_topk( *, survival: torch.Tensor, - budget: int, + budget: Union[int, torch.Tensor], cfg: DSparkScheduleConfig, ) -> torch.Tensor: """Allocate ``budget`` verify tokens across the batch by global survival rank. @@ -115,6 +116,9 @@ def schedule_verify_lens_topk( survival: ``[bs, K]`` prefix-survival probabilities from :func:`compute_survival`. budget: total verify tokens to hand out *above* the per-request floor. + A 0-d integer tensor is honoured without a host sync (the value + stays device-resident through the ``rank < budget`` cut), which is + what lets the whole function run inside a captured graph. cfg: bounds; see :class:`DSparkScheduleConfig`. Returns: ``[bs]`` int32 verify lengths, each in @@ -137,8 +141,11 @@ def schedule_verify_lens_topk( floor = int(cfg.min_verify_len) schedulable = int(cfg.schedulable_per_request) verify_lens = torch.full((bs,), floor, dtype=torch.int32, device=device) - budget = int(budget) - if bs == 0 or schedulable <= 0 or budget <= 0: + if not isinstance(budget, torch.Tensor): + budget = int(budget) + if budget <= 0: + return verify_lens + if bs == 0 or schedulable <= 0: return verify_lens # Positions [0, floor) are already granted to every request; only those diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py index 661903dc8e1e..54bf4bf351fd 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py @@ -14,8 +14,12 @@ import pytest import torch +from tensorrt_llm._torch.speculative.dspark_device_select import \ + select_windows_device from tensorrt_llm._torch.speculative.dspark_ragged import ( RaggedVerifyLayout, build_row_maps_device, fill_bucket_device) +from tensorrt_llm._torch.speculative.dspark_schedule import ( + DSparkScheduleConfig, compute_survival, schedule_verify_lens_topk) def _host_fill(lens, padded_bs, bucket, max_verify_len): @@ -116,3 +120,86 @@ def test_row_maps_match_prepare_semantics(): ends = torch.cumsum(lens.to(torch.long), 0) - 1 assert torch.equal(extents[ends], kv_lens) assert int(extents.min()) >= 1 + + +def test_topk_tensor_budget_matches_int(): + """The 0-d tensor budget path (in-graph capacity knob) must allocate + identically to the host int path, including budget <= 0.""" + gen = torch.Generator().manual_seed(11) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + for _ in range(30): + bs = int(torch.randint(1, 65, (1,), generator=gen)) + survival = compute_survival( + torch.rand(bs, 5, generator=gen)) + for budget in (0, 1, bs, 2 * bs, bs * 4, bs * 10): + host = schedule_verify_lens_topk(survival=survival, + budget=budget, cfg=cfg) + dev = schedule_verify_lens_topk(survival=survival, + budget=torch.tensor(budget), + cfg=cfg) + assert torch.equal(host, dev) + + +def _host_select(rows, budget, bucket, padded_bs, cfg): + """The lag-free host reference: calibrate -> survival -> topk -> fill.""" + survival = compute_survival(torch.sigmoid(rows)) + scheduled = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=cfg) + token_lens = (scheduled + 1).tolist() + layout = RaggedVerifyLayout.from_verify_lens( + torch.tensor(token_lens, dtype=torch.int32), + graph_num_tokens=bucket, + total_verify_tokens=sum(token_lens)) + return layout.fill_bucket( + max_verify_len=cfg.resolved_max_verify_len + 1, padded_bs=padded_bs) + + +def test_select_windows_device_matches_host_chain(): + """The full prologue equals the host chain run on the same (fresh) rows, + across pad rows, indirection through slots, and stale-stamp fallback.""" + gen = torch.Generator().manual_seed(20260806) + K = 5 + cfg = DSparkScheduleConfig(block_size=K, min_verify_len=1) + max_token_len = cfg.resolved_max_verify_len + 1 + for _ in range(40): + n_real = int(torch.randint(1, 33, (1,), generator=gen)) + padded_bs = n_real + int(torch.randint(0, 5, (1,), generator=gen)) + num_slots = padded_bs + 8 + conf = torch.randn(num_slots, K, generator=gen) * 3 + # Requests sit at shuffled slots; pad rows point at slot 0. + slots = torch.randperm(num_slots, generator=gen)[:n_real] + slot_idx = torch.zeros(padded_bs, dtype=torch.long) + slot_idx[:n_real] = slots + # A third of the rows are stale: their gather must neutralize to + # "verify everything". + stamp = torch.full((num_slots,), 7, dtype=torch.int64) + stale_rows = torch.rand(n_real, generator=gen) < 0.33 + stamp[slots[stale_rows]] = 3 + + budget = int(torch.randint(0, n_real * K, (1,), generator=gen)) + lo = n_real * (cfg.min_verify_len + 1) + min( + budget, n_real * (K - cfg.min_verify_len)) + (padded_bs - n_real) + bucket = min(padded_bs * max_token_len, + lo + int(torch.randint(0, 8, (1,), generator=gen))) + + result = select_windows_device( + confidence_logits=conf, + slot_idx=slot_idx, + num_real=torch.tensor(n_real), + budget=torch.tensor(budget), + graph_num_tokens=bucket, + cfg=cfg, + stamp=stamp, + expected_stamp=torch.tensor(7), + ) + rows = conf[slots].clone() + rows[stale_rows] = 30.0 # NEUTRAL_CONFIDENCE_LOGIT + host = _host_select(rows, budget, bucket, padded_bs, cfg) + assert torch.equal(result.verify_lens, host.verify_lens) + assert torch.equal(result.qo_indptr, host.qo_indptr) + assert int(result.verify_lens.sum()) == bucket + # Row maps must describe exactly the filled lens. + want_req, want_corr = build_row_maps_device( + host.verify_lens, graph_num_tokens=bucket) + assert torch.equal(result.req_idx, want_req) + assert torch.equal(result.kv_correction, want_corr) From eac5f72ab95900e77dc82f030d7082654c90ae07 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Thu, 6 Aug 2026 17:45:23 -0700 Subject: [PATCH 103/177] [None][feat] Device-side window selection: fresh ranking, lag-2 capacity The host keeps deciding only the capacity -- (padded_bs, bucket) via the unchanged ADP agreement and fit, and the token budget, now argmax'd over a snapshot one staging OLDER (lag-2 by the block-relative count, matching the paper/SGLang: capacity tolerates staleness) -- and stamps a uniform SHAPE SPLIT of the budget through the normal fit, so every host copy width and account stays self-consistent. A pre-replay device prologue then re-ranks the same real/pad token totals with the verified block's OWN confidence (ranking lag zero, straight off the live slot-indexed buffer) and overwrites the layout's content in place: verify lens, qo_indptr, the per-token row maps, the packed input/position/draft tokens, and a per-request kv-length delta from the shape split to the true windows. No new sync anywhere: the prologue is device ops between input staging and replay, stream-ordered into the graph's reads. Host learns the true windows through the sampler's existing D2H batch (a slot-indexed verify_lens ride-along, same pattern as cap_trim_lens), which the rewind arithmetic now prefers; py_verify_len keeps the shape split as a presence sentinel for the graph key. The in-graph ragged expansion now reads repeats from a persistent staged buffer instead of a per-call pinned temporary whose host pointer capture used to bake in. STS collection refuses device mode (trimmed windows censor labels). Off by default; TLLM_DSPARK_DEVICE_WINDOWS=1. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa.py | 116 ++++++++++++- .../_torch/pyexecutor/model_engine.py | 150 +++++++++++++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 35 +++- tensorrt_llm/_torch/speculative/dspark.py | 38 +++++ .../speculative/dspark_device_select.py | 7 + .../_torch/speculative/dspark_ragged.py | 13 +- tensorrt_llm/_torch/speculative/dspark_sts.py | 7 + .../_torch/speculative/dspark_verify.py | 155 +++++++++++++++++- .../_torch/speculative/spec_sampler_base.py | 42 ++++- .../test_dspark_device_fill_parity.py | 106 ++++++++++++ 10 files changed, 644 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 4107f0aa2945..2f26ccbc627b 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -662,6 +662,12 @@ class DSAtrtllmAttentionMetadata(TrtllmAttentionMetadata): # `1 + max_draft_tokens` the strided expansions assume. Host-side because # reading repeats off a device tensor would sync. ragged_verify_lens: Optional[List[int]] = None + # Device-window mode (TLLM_DSPARK_DEVICE_WINDOWS): on ragged steps the + # host list above is only a SHAPE split summing to the bucket; the true + # per-request windows are chosen on device after prepare() and written + # through apply_device_ragged_layout. Host reads of window values must + # then be treated as bounds (see _prepare_impl). + device_windows_mode: bool = False # Query tokens per generation request on a *uniform* batch this step # (`1 + runtime_draft_len`). Zero means "not published"; expansions then # fall back to the static `1 + max_draft_tokens`. @@ -723,11 +729,19 @@ def expand_per_gen_token(self, if self.ragged_verify_lens is None: stride = self.gen_token_stride return values.repeat_interleave(stride, dim=dim), num_tokens - repeats_host = torch.tensor( - repeats, - dtype=torch.long, - pin_memory=(prefer_pinned() and values.device.type == 'cuda')) - repeats_dev = repeats_host.to(values.device, non_blocking=True) + # Device values take repeats from the PERSISTENT staged buffer, never + # a per-call pinned temporary: that call site runs inside the captured + # graph (via on_update_kv_lens), where a temporary would bake its host + # pointer into the captured H2D copy. The buffer is staged by + # prepare_for_spec_decode and overwritten in place by the + # device-window prologue. Host-tensor calls (prepare-time expansions + # of host kv lens / block tables) keep a plain host repeats tensor. + if values.device == self.gen_token_repeats_cuda.device: + repeats_dev = self.gen_token_repeats_cuda[:self.num_generations] + else: + repeats_dev = torch.tensor(repeats, + dtype=torch.long, + device=values.device) return (values.repeat_interleave(repeats_dev, dim=dim, output_size=num_tokens), num_tokens) @@ -847,6 +861,20 @@ def _prepare_impl(self): else: kv_lens = cached_token_lens + self.seq_lens_kv + if self.device_windows_mode and self.is_ragged_verify: + # Device-window steps: the host's per-request windows are only a + # SHAPE split; the true windows are chosen on device after + # prepare(). Every host-side read of kv_lens below is a bound or + # is rebuilt in-graph from kv_lens_cuda (on_update_kv_lens), so + # inflate the gen entries to the full drafted block -- a bound no + # device-chosen window can exceed. kv_lens_cuda keeps the exact + # shape-split values (staged by super().prepare() above); the + # prologue's per-request delta corrects it to the true windows. + nc, ns = self.num_contexts, self.num_seqs + kv_lens = kv_lens.clone() + kv_lens[nc:ns] += (1 + self.max_draft_tokens) - self.seq_lens_kv[nc:ns] + self.kv_lens[nc:ns] = kv_lens[nc:ns] + # For mla_rope_append_paged_kv_assign_q self.prepare_for_mla_rope_append(cached_token_lens, kv_lens) @@ -1226,6 +1254,23 @@ def create_buffers_for_indexer(self, capture_graph=False): ) self.host_req_idx_per_token = torch.empty_like( self.req_idx_per_token, device='cpu', pin_memory=prefer_pinned()) + # Per-gen-request query-token counts as a PERSISTENT device tensor. + # The in-graph ragged expansion (``on_update_kv_lens`` -> + # ``expand_per_gen_token``) must read repeats from a stable address: + # a per-call pinned temporary would bake the temporary's host pointer + # into the captured H2D. Staged in prepare_for_spec_decode; the + # device-window prologue overwrites it in place. + self.gen_token_repeats_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_sequences, ), + cache_name="gen_token_repeats_cuda", + dtype=torch.int64, + capture_graph=capture_graph, + ) + self.host_gen_token_repeats = torch.empty_like( + self.gen_token_repeats_cuda, + device='cpu', + pin_memory=prefer_pinned()) # Block table for topk_indices conversion (shared for context and generation) self.block_table = self.get_empty( self.cuda_graph_buffers, @@ -1800,6 +1845,16 @@ def prepare_for_spec_decode(self, kv_lens: torch.Tensor): ((self.max_draft_tokens == 2 or self.max_draft_tokens > 3) and get_sm_version() >= 100))) if self.use_expanded_buffers_for_mtp: + if self.is_ragged_verify: + # Stage the persistent device repeats the in-graph expansion + # reads (see expand_per_gen_token); must land before the + # expansions below and before any capture of this step. + n_gen = self.num_generations + self.host_gen_token_repeats[:n_gen].copy_( + torch.tensor(self.gen_token_repeat_list(), + dtype=torch.int64)) + self.gen_token_repeats_cuda[:n_gen].copy_( + self.host_gen_token_repeats[:n_gen], non_blocking=True) # Expand kv_lens_cuda (only generation): one row per query token, # repeating each request's kv_len as many times as it has tokens. gen_kv_lens = self.get_indexer_kv_lens( @@ -2083,6 +2138,57 @@ def ragged_row_kv_lens(self, num_tokens: int) -> Optional[torch.Tensor]: f"past what was populated") return self.row_kv_lens_cuda[:num_tokens] + def apply_device_ragged_layout(self, verify_lens: torch.Tensor, + req_idx: torch.Tensor, + kv_correction: torch.Tensor) -> None: + """Rebind this step's ragged layout to device-chosen windows. + + Pre-replay device ops only -- no host sync, stream-ordered before the + graph replay so every captured reader sees the new values. Overwrites + the same persistent buffers prepare() staged from the host shape + split; the per-row extents and every expanded/indexer artifact then + rebuild in-graph from them (``on_update_kv_lens``). + + Preconditions (the engine's prologue guarantees them): gen-only batch + (``num_contexts == 0``), ``verify_lens`` sums to the staged bucket + (``_ragged_num_rows``), and ``kv_lens_cuda`` has already been + delta-corrected from the shape split to these windows. + + Args: + verify_lens: ``[num_generations]`` int32 token windows. + req_idx: ``[bucket]`` int64 owning-row per packed token. + kv_correction: ``[bucket]`` int32, same ``1 - v .. 0`` convention + as the host staging (extent = kv_len + correction). + """ + nc = self.num_contexts + rows = self._ragged_num_rows + n_gen = self.num_generations + # Request-major: the true query lengths. Everything on_update_kv_lens + # derives per token (slot mappings, indptrs, dense top-k) follows. + self.seq_lens_cuda[nc:self.num_seqs] = verify_lens.to( + self.seq_lens_cuda.dtype) + self.gen_token_repeats_cuda[:n_gen] = verify_lens.to(torch.int64) + # Token-major maps, both the indexer's rows and the attention view's. + self.row_req_idx_cuda[:rows] = req_idx + self.row_kv_correction_cuda[:rows] = kv_correction + self.attn_row_req_idx_cuda[:rows] = req_idx + nc + self.attn_row_kv_correction_cuda[:rows] = kv_correction + self.req_idx_per_token[nc:nc + rows] = (req_idx + nc).to( + self.req_idx_per_token.dtype) + # Extents (row_kv_lens/attn_row_kv_lens) are rebuilt in-graph by the + # refresh functions; the block-table expansions are gathered per row + # here because block IDs do not move within a step and their refresh + # is prepare-time only. + self.refresh_token_major_block_table() + block_table_expanded = getattr(self, "block_table_expanded", None) + indexer_offsets = getattr(self, "indexer_k_cache_block_offsets", None) + if block_table_expanded is not None and indexer_offsets is not None: + width = min(indexer_offsets.shape[-1], + block_table_expanded.shape[-1]) + block_table_expanded[:rows, :width] = ( + indexer_offsets[nc:self.num_seqs, :width].index_select( + 0, req_idx).clamp_(min=0)) + def prepare_for_indexer_k_cache(self): # Build indexer_k_cache_block_offsets using pool block indices derived # from host_kv_cache_block_offsets (populated by super().prepare()). diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 6619922b7dba..ce95136fbf88 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3829,6 +3829,7 @@ def fit_ragged_verify_lens( # _get_padded_batch. self._dspark_last_total_tokens = int(bucket) self._dspark_last_padded_bs = int(padded_bs) + self._dspark_last_num_real = int(n_real) runner.agreed_ragged_bucket = int(bucket) runner.ragged_pad_verify_len = pad_len - 1 @@ -3837,6 +3838,130 @@ def fit_ragged_verify_lens( request.py_verify_len = int(tokens) - 1 return int(bucket) + def _apply_device_window_prologue(self, inputs, + new_tensors_device) -> bool: + """Re-rank this step's verify windows on device, with fresh confidence. + + Runs after ``_prepare_inputs`` and before the graph replay, all device + ops on the current stream -- stream order makes every write visible to + the replayed graph, and nothing here reads device data back to the + host. The host has already agreed ``(padded_bs, bucket)`` and staged a + SHAPE SPLIT of the (lagged) budget through the normal fit; this + prologue re-distributes the same real/pad token totals by the verified + block's OWN confidence (ranking lag zero) and overwrites the layout's + content: verify lens, qo_indptr, the per-token row maps, the packed + input/position/draft tokens, and the per-request kv-length delta from + the shape split to the true windows. + + Returns True when applied; False when a precondition fails (the step + then runs the shape split as-is, which is a valid window assignment). + """ + from ..speculative.dspark_device_select import select_windows_device + + runner = self.cuda_graph_runner + budget = getattr(self, "_dspark_device_budget", None) + self._dspark_device_budget = None + if budget is None or runner.agreed_ragged_bucket is None: + return False + if not getattr(self, "_dspark_prev_covers_batch", False): + # The input gathers below address new_tokens_device by + # previous_batch_indices in batch order; a request without a + # previous device tensor breaks that addressing. + return False + worker = self._get_spec_worker() + planner = getattr(worker, "verify_planner", None) + if planner is None or worker.staged_confidence_buffer() is None: + return False + if worker.batch_slot_view(1) is None: + return False + spec_metadata = inputs.get('spec_metadata') + attn_metadata = inputs.get('attn_metadata') + if spec_metadata is None or attn_metadata is None: + return False + if int(attn_metadata.num_contexts) != 0: + return False + + bucket = int(runner.agreed_ragged_bucket) + padded_bs = int(self._dspark_last_padded_bs) + n_real = int(self._dspark_last_num_real) + pad_len_tok = int(runner.ragged_pad_verify_len) + 1 + real_tokens = bucket - (padded_bs - n_real) * pad_len_tok + cfg = planner.cfg + # The scheduler cannot grant more than the real rows can absorb under + # the published split; the fill tops up any shortfall. + budget = max( + 0, + min(int(budget), real_tokens - n_real * (cfg.min_verify_len + 1))) + + # Snapshot the shape split BEFORE overwriting: past_seen per row is + # the staged position at each row's first token, and the kv delta + # needs the split the host baked into kv_lens_cuda. + lens_buf = self.ragged_verify_lens_cuda + qo_buf = self.ragged_qo_indptr_cuda + split_lens = lens_buf[:padded_bs].clone() + split_qo = qo_buf[:padded_bs + 1].to(torch.long) + past_seen = self.position_ids_cuda[split_qo[:-1]].clone() + + expected_stamp = worker.verified_draft_seq_cuda() + stamps = worker.confidence_stamp_buffer( + ) if expected_stamp is not None else None + result = select_windows_device( + confidence_logits=worker.staged_confidence_buffer(), + slot_idx=worker.batch_slot_view(padded_bs), + num_real=n_real, + budget=budget, + graph_num_tokens=bucket, + cfg=cfg, + apply_calibration=planner.apply_calibration, + stamp=stamps, + expected_stamp=expected_stamp, + pad_len=pad_len_tok, + ) + + lens_buf[:padded_bs].copy_(result.verify_lens) + qo_buf[:padded_bs + 1].copy_(result.qo_indptr) + # Every host-staged kv formula used the shape split consistently, so + # correcting by the difference lands on exactly the value the host + # path would have produced with these windows. + attn_metadata.kv_lens_cuda[:padded_bs] += result.verify_lens - split_lens + + req_idx = result.req_idx + device = req_idx.device + flat = torch.arange(bucket, device=device) + offset = flat - result.qo_indptr.to(torch.long)[req_idx] + prev_slots = self.previous_batch_indices_cuda[:n_real].to(torch.long) + slots_tok = prev_slots[req_idx.clamp(max=n_real - 1)] + + new_tokens_device = new_tensors_device.new_tokens + new_tokens_lens_device = new_tensors_device.new_tokens_lens + next_draft_tokens_device = new_tensors_device.next_draft_tokens + + self.input_ids_cuda[:bucket] = new_tokens_device.transpose( + 0, 1)[slots_tok, offset].flatten().to(self.input_ids_cuda.dtype) + self.position_ids_cuda[:bucket] = past_seen[req_idx] + offset.to( + past_seen.dtype) + # Overlap corrections gather by each token's OWNER; rebuild both the + # index and the per-token offset it feeds (mirrors the host staging at + # the previous_pos_indices block). + self.previous_pos_indices_cuda[:real_tokens] = slots_tok[:real_tokens].to( + self.previous_pos_indices_cuda.dtype) + self.previous_pos_id_offsets_cuda[:real_tokens].copy_( + new_tokens_lens_device[slots_tok[:real_tokens]]) + # Draft tokens pack compactly: token t at in-window offset o >= 1 goes + # to slot t - owner - 1 (one bonus token per preceding row). Anchor + # tokens (o == 0) are parked one past the real draft region, where any + # valid token id is discard-territory. + real_draft = real_tokens - n_real + mask = offset >= 1 + dst = torch.where(mask, flat - req_idx - 1, + torch.full_like(req_idx, real_draft)) + src = next_draft_tokens_device[slots_tok, (offset - 1).clamp_min(0)] + self.draft_tokens_cuda[dst] = src.to(self.draft_tokens_cuda.dtype) + + attn_metadata.apply_device_ragged_layout(result.verify_lens, req_idx, + result.kv_correction) + return True + @staticmethod def _ragged_token_lens(generation_requests) -> Optional[List[int]]: """Each generation request's token window, or None if the batch is uniform. @@ -3875,6 +4000,12 @@ def _publish_gen_token_layout(self, attn_metadata, if hasattr(attn_metadata, "ragged_verify_lens"): attn_metadata.ragged_verify_lens = self._ragged_token_lens( generation_requests) + if hasattr(attn_metadata, "device_windows_mode"): + # Tells the attention prepare that the host window VALUES are a + # shape split (bounds only); the true windows land on device + # through apply_device_ragged_layout after prepare. + from ..speculative.dspark_verify import device_windows_enabled + attn_metadata.device_windows_mode = device_windows_enabled() def _pinned_host(self, key: str, values, dtype) -> torch.Tensor: """A persistent pinned staging buffer holding ``values``. @@ -5087,6 +5218,15 @@ def append_cross_attention_state(request: LlmRequest, gen_request_seq_slots.append(request.py_seq_slot) previous_batch_len = len(previous_batch_indices) + # Device-window prologue precondition: it gathers this step's inputs + # by (slot, offset) through previous_batch_indices_cuda, which only + # covers requests that carried a previous device tensor and only in + # batch order when EVERY real generation request did. Recorded here, + # consumed by _apply_device_window_prologue. + self._dspark_prev_covers_batch = ( + previous_batch_len > 0 + and previous_batch_len == len(extend_requests) - + len(extend_dummy_requests)) def previous_seq_slots_device(): previous_batch_indices_host = torch.tensor( @@ -6934,6 +7074,16 @@ def forward(self, self._prepare_inputs_event = torch.cuda.Event() self._prepare_inputs_event.record() + # Device-window selection: re-rank the staged ragged layout by the + # verified block's own confidence, between input staging and the + # replay (stream-ordered, no host sync). Graph steps only: eager + # and capture steps run the host shape split, which is a valid + # window assignment. + if (can_run_graph and not self.is_warmup + and getattr(self, "_dspark_device_budget", None) + is not None): + self._apply_device_window_prologue(inputs, new_tensors_device) + with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): if not can_run_graph: # Fallback to eager execution if graph was not used diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index f56f64b7eb23..55a25dbbf86d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3382,13 +3382,29 @@ def _dspark_confidence_draft_len(self, # the single collective, and nothing after it may issue another one. ragged_lens = None fallback_reason = None + device_budget = None + # Cleared unconditionally: a budget left over from a previous step + # would re-rank a batch it was never computed for. + self.model_engine._dspark_device_budget = None if compute_windows: - ragged_lens = planner.decide_verify_lens( - num_gen_requests=num_gen_requests, - rows=confidence_rows, - reduce_across_ranks=False) - if ragged_lens is None: - fallback_reason = "planner_declined" + if planner.device_windows and mode.trims_submitted_tokens: + # Device-window mode: the host decides only the CAPACITY, from + # the lag-2 snapshot; the returned lens are a shape split used + # for the cross-rank agreement and the fit. The true windows + # are re-ranked on device just before the replay. + pair = planner.decide_verify_budget( + num_gen_requests=num_gen_requests, rows=confidence_rows) + if pair is None: + fallback_reason = "planner_declined" + else: + device_budget, ragged_lens = pair + else: + ragged_lens = planner.decide_verify_lens( + num_gen_requests=num_gen_requests, + rows=confidence_rows, + reduce_across_ranks=False) + if ragged_lens is None: + fallback_reason = "planner_declined" else: fallback_reason = "mode_static" # Nothing to compute: the fallback verifies the full drafted block. @@ -3482,6 +3498,13 @@ def _dspark_confidence_draft_len(self, ragged_active = bucket is not None if bucket is None: fallback_reason = "no_captured_shape" + elif device_budget is not None: + # The fit landed on a captured shape; hand the engine the + # capacity so the pre-replay prologue re-ranks it with + # fresh confidence. Published only here: every fallback + # leaves it None and the step runs the host windows. + self.model_engine._dspark_device_budget = int( + device_budget) else: # No rank can graph this step, so run the windows eagerly. Do # not fit a bucket: the fit sizes pad rows _get_padded_batch diff --git a/tensorrt_llm/_torch/speculative/dspark.py b/tensorrt_llm/_torch/speculative/dspark.py index eca97fa8b688..6d1ba60c3f7b 100644 --- a/tensorrt_llm/_torch/speculative/dspark.py +++ b/tensorrt_llm/_torch/speculative/dspark.py @@ -296,6 +296,30 @@ def verified_draft_seq(self) -> Optional[int]: """ return self._draft_seq_host - 1 if self._draft_seq_host > 1 else None + def verified_draft_seq_cuda(self) -> Optional[torch.Tensor]: + """:meth:`verified_draft_seq` as a device scalar, for the prologue. + + ``_draft_seq_cuda`` already holds the CURRENT step's sequence by + launch time (``bump_draft_seq`` fills it before the forward is + enqueued), so the pass that stamped the block being verified is one + less. Passing the buffer itself would make every fresh row read as + stale and silently degrade the whole batch to neutral. + """ + if self._draft_seq_cuda is None: + return None + return self._draft_seq_cuda - 1 + + def batch_slot_view(self, num_rows: int) -> Optional[torch.Tensor]: + """The batch-position -> confidence-row map staged this step. + + ``[num_rows]`` int64 on device; rows past the real batch point at the + scratch row (their survival is zeroed by the prologue's ``num_real`` + mask, so the content never matters). + """ + if self._batch_to_slot is None: + return None + return self._batch_to_slot[:num_rows] + @property def max_draft_len(self) -> int: return self.spec_config.max_draft_len @@ -846,4 +870,18 @@ def _forward_impl( cap_trim = spec_metadata.accept_cap_trim if cap_trim is not None and spec_metadata.accept_caps is not None: outputs["cap_trim_lens"] = cap_trim[:batch_size] + # Ragged steps: the TOKEN windows the step actually verified, straight + # off the persistent layout buffer. Rides the sampler's existing D2H + # batch; under device-window selection this is the ONLY place the host + # can learn the true windows (py_verify_len then holds only the host + # shape split), and the rewind arithmetic must use it. + # + # Gen-only steps only: the buffer is generation-indexed while every + # other output row is batch-indexed, so on a mixed step (eager ragged + # with finished contexts) the rows would misalign AND come up short. + # Mixed steps keep exact windows on py_verify_len, so the sampler's + # snapshot fallback stays correct there. + if (spec_metadata.verify_lens is not None + and spec_metadata.verify_lens.shape[0] >= batch_size): + outputs["verify_lens"] = spec_metadata.verify_lens[:batch_size] return outputs diff --git a/tensorrt_llm/_torch/speculative/dspark_device_select.py b/tensorrt_llm/_torch/speculative/dspark_device_select.py index 586e888ab69d..a2febd3ff914 100644 --- a/tensorrt_llm/_torch/speculative/dspark_device_select.py +++ b/tensorrt_llm/_torch/speculative/dspark_device_select.py @@ -81,6 +81,7 @@ def select_windows_device( apply_calibration: Optional[Callable[[torch.Tensor], torch.Tensor]] = None, stamp: Optional[torch.Tensor] = None, expected_stamp: Optional[torch.Tensor] = None, + pad_len: Optional[int] = None, ) -> DeviceWindowResult: """Rank the batch by fresh confidence and pack it into the agreed bucket. @@ -105,6 +106,11 @@ def select_windows_device( ``expected_stamp`` (0-d), rows whose stamp mismatches fall back to the neutral row (full survival: verify everything), the same fail-open the host planner applies to unknown slots. + pad_len: when given, pad rows carry EXACTLY this many tokens and all + fill slack goes to real rows -- the host fit's published split, + which its pre-launch copy widths depend on (see + :func:`fill_bucket_device`). The caller must clamp ``budget`` so + the real rows can absorb ``graph_num_tokens - n_pad * pad_len``. Returns: :class:`DeviceWindowResult`; every tensor is derived without a host @@ -148,6 +154,7 @@ def select_windows_device( num_real=num_real, graph_num_tokens=graph_num_tokens, max_verify_len=max_token_len, + pad_fill=pad_len, ) req_idx, correction = build_row_maps_device( filled, graph_num_tokens=graph_num_tokens) diff --git a/tensorrt_llm/_torch/speculative/dspark_ragged.py b/tensorrt_llm/_torch/speculative/dspark_ragged.py index 64caa79ff7af..746e9d8bd60e 100644 --- a/tensorrt_llm/_torch/speculative/dspark_ragged.py +++ b/tensorrt_llm/_torch/speculative/dspark_ragged.py @@ -366,6 +366,7 @@ def fill_bucket_device( num_real: torch.Tensor, graph_num_tokens: int, max_verify_len: int, + pad_fill: Optional[int] = None, ) -> torch.Tensor: """Capture-safe :meth:`RaggedVerifyLayout.fill_bucket`: same allocation, computed entirely on device. @@ -392,6 +393,12 @@ def fill_bucket_device( 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. """ @@ -404,10 +411,12 @@ def fill_bucket_device( lens = torch.where( is_real, verify_lens.to(torch.int32), - torch.ones(padded_bs, dtype=torch.int32, device=device), + 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() - for phase_mask in (is_real, ~is_real): + 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 diff --git a/tensorrt_llm/_torch/speculative/dspark_sts.py b/tensorrt_llm/_torch/speculative/dspark_sts.py index 7267d61bd35f..b0c67c785a97 100644 --- a/tensorrt_llm/_torch/speculative/dspark_sts.py +++ b/tensorrt_llm/_torch/speculative/dspark_sts.py @@ -294,6 +294,13 @@ def make_recorder_from_env(*, block_size: int, rank: int = 0, f"verify rows corrupt the per-position prefix label. Collect in " f"static mode (TLLM_DSPARK_RAGGED_VERIFY_MODE=static).") + if os.environ.get("TLLM_DSPARK_DEVICE_WINDOWS", "0") == "1": + raise ValueError( + f"{STS_COLLECT_ENV} is set with TLLM_DSPARK_DEVICE_WINDOWS=1. " + f"Device-selected windows trim per request on device, so every " + f"collected acceptance label is censored by a window the host " + f"never sees. Collect with device windows off.") + frac = os.environ.get("TLLM_DSPARK_FORCE_BUDGET_FRAC", "").strip() if frac: try: diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index 735054d6ad80..7054e268d64c 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -24,7 +24,7 @@ """ import os -from typing import Callable, List, Optional, Sequence +from typing import Callable, List, Optional, Sequence, Tuple import numpy as np import torch @@ -51,6 +51,20 @@ #: fixed-point ints so the allgather payload stays homogeneous ints. _FRAC_WIRE_SCALE = 10_000 +#: Device-side window selection (the paper's dual-timescale split): the host +#: decides only the CAPACITY -- (padded_bs, bucket, budget) -- from a snapshot +#: one step older than today's (lag-2 by the block-relative count, matching +#: SGLang), and a pre-replay device prologue re-ranks the batch with the +#: verified block's own fresh confidence (lag-0). Constant per process, read +#: once at init on every rank. +DEVICE_WINDOWS_ENV = "TLLM_DSPARK_DEVICE_WINDOWS" + + +def device_windows_enabled() -> bool: + """Whether device-side window selection is on (process-constant env).""" + return os.environ.get(DEVICE_WINDOWS_ENV, "0") == "1" + + __all__ = ["DSparkVerifyPlanner"] @@ -106,6 +120,24 @@ def __init__( self._host_stamps: Optional[torch.Tensor] = None self._staged_seq: Optional[int] = None self._snapshot_valid = False + + # Device-window mode: ranking happens on device with fresh confidence, + # so the host snapshot serves ONLY the capacity argmax -- and that one + # deliberately reads a snapshot one staging older (lag-2), per the + # paper's design: capacity tolerates staleness, ranking does not. + # Two pinned buffers alternate; the pair we are NOT writing this step + # is the older, guaranteed-landed one. + self.device_windows = os.environ.get(DEVICE_WINDOWS_ENV, "0") == "1" + if self.device_windows and self._forced_verify_len is not None: + raise ValueError( + f"{DEVICE_WINDOWS_ENV}=1 is incompatible with " + f"{FORCE_VERIFY_LEN_ENV}: the pin fixes host windows that " + f"device selection would silently override") + self._prev_buffer: Optional[torch.Tensor] = None + self._prev_event: Optional[torch.cuda.Event] = None + self._prev_stamps: Optional[torch.Tensor] = None + self._prev_seq: Optional[int] = None + self._prev_valid = False # Count every path that declines to trim; degradation here is silent. self.stats = { "decisions": 0, @@ -135,6 +167,19 @@ def stage_confidence(self, confidence_logits: torch.Tensor, """ if confidence_logits is None or confidence_logits.shape[0] == 0: return + if self.device_windows: + # Rotate before overwriting: what was current becomes the older + # (lag-2) snapshot the budget argmax reads; the old older buffer's + # storage is reused for this staging. Pure reference swaps. + self._host_buffer, self._prev_buffer = (self._prev_buffer, + self._host_buffer) + self._copy_event, self._prev_event = (self._prev_event, + self._copy_event) + self._host_stamps, self._prev_stamps = (self._prev_stamps, + self._host_stamps) + self._staged_seq, self._prev_seq = self._prev_seq, self._staged_seq + self._prev_valid = self._snapshot_valid + self._snapshot_valid = False if self._host_buffer is None or self._host_buffer.shape != confidence_logits.shape: self._host_buffer = torch.empty( confidence_logits.shape, @@ -321,23 +366,29 @@ def adopt_verify_len_pin(self, wire_value: int) -> None: f"configuration.") def _note_snapshot_stats(self, selected: torch.Tensor, - rows: Optional[Sequence[int]]) -> None: + rows: Optional[Sequence[int]], + stamps: Optional[torch.Tensor] = None, + staged_seq: Optional[int] = None) -> None: """Record snapshot-quality stats: stamp-lag histogram (staleness per gathered row) and neutral-row count ("unknown" rows entering the argmax as certainties). Observability only; never changes behaviour. + ``stamps``/``staged_seq`` override the current snapshot's pair when the + caller read the older (lag-2) one. """ try: + stamps = self._host_stamps if stamps is None else stamps + staged_seq = self._staged_seq if staged_seq is None else staged_seq n = int(selected.shape[0]) self.stats["snap_rows"] = self.stats.get("snap_rows", 0) + n neutral = int((selected >= NEUTRAL_CONFIDENCE_LOGIT - 1.0).all(dim=1).sum()) self.stats["snap_neutral_rows"] = ( self.stats.get("snap_neutral_rows", 0) + neutral) - if (self._host_stamps is not None and self._staged_seq is not None + if (stamps is not None and staged_seq is not None and rows is not None): idx = torch.as_tensor(list(rows), dtype=torch.long) - idx = idx.clamp_(0, self._host_stamps.shape[0] - 1) - lags = (int(self._staged_seq) - - self._host_stamps[idx].to(torch.int64)) + idx = idx.clamp_(0, stamps.shape[0] - 1) + lags = (int(staged_seq) + - stamps[idx].to(torch.int64)) hist = self.stats.setdefault("stamp_lag_hist", {}) for lag, cnt in zip(*[t.tolist() for t in lags.unique( return_counts=True)]): @@ -349,8 +400,22 @@ def _note_snapshot_stats(self, selected: torch.Tensor, if self.stats["snap_stats_errors"] == 1: logger.warning(f"DSpark snapshot stats failed once: {exc}") + def _ready_older_snapshot(self) -> Optional[torch.Tensor]: + """The lag-2 snapshot (device-window mode's capacity input), or None. + + By construction its copy was issued a full staging earlier, so its + event has practically always landed; the poll stays for the same + fail-closed reason as :meth:`_ready_snapshot`. + """ + if not self._prev_valid or self._prev_buffer is None: + return None + if self._prev_event is not None and not self._prev_event.query(): + return None + return self._prev_buffer + def _gather_rows( - self, *, num_gen_requests: int, rows: Optional[Sequence[int]] + self, *, num_gen_requests: int, rows: Optional[Sequence[int]], + snapshot: Optional[torch.Tensor] = None ) -> Optional[torch.Tensor]: """This step's confidence, one row per generation request, or None. @@ -359,9 +424,11 @@ def _gather_rows( lag, so positional reads (``rows=None``) are only correct for callers that own the ordering. Declines (None) rather than returning fewer rows than the batch, which would partially assign verify windows - downstream. + downstream. ``snapshot`` overrides the default (current) snapshot when + the caller already picked one (the lag-2 budget read). """ - snapshot = self._ready_snapshot() + if snapshot is None: + snapshot = self._ready_snapshot() if snapshot is None: self.stats["fallback_no_snapshot"] += 1 return None @@ -380,6 +447,76 @@ def _gather_rows( return None return selected + def decide_verify_budget( + self, + *, + num_gen_requests: int, + rows: Optional[Sequence[int]] = None, + ) -> Optional[Tuple[int, List[int]]]: + """Capacity only, for device-window mode: ``(budget, shape_lens)``. + + The budget (verify tokens above the per-request floor) comes from the + cost-table argmax over the OLDER staged snapshot -- lag-2 by the + block-relative count, the paper's prescription: capacity tolerates + staleness, ranking does not. ``shape_lens`` is a canonical uniform + spread of the budget used ONLY for the cross-rank shape agreement and + bucket fit (both read totals, never the split); the true per-request + windows are chosen on device from the verified block's own confidence. + Fail-closed like :meth:`decide_verify_lens`: None means verify the + full block. + """ + self.stats["decisions"] += 1 + if num_gen_requests <= 0: + self.stats["fallback_no_gen_requests"] += 1 + return None + n = int(num_gen_requests) + trimmable = self.cfg.resolved_max_verify_len - self.cfg.min_verify_len + if self._forced_budget_frac is not None: + # The frac sweep stays honest in device mode: it replaces only the + # argmax; the device top-k still spends the budget for real. + budget = int(round(self._forced_budget_frac * n * trimmable)) + self.stats["forced_budget_steps"] = ( + self.stats.get("forced_budget_steps", 0) + 1) + else: + if self.cost_table is None or self.cost_table.is_flat: + self.stats["fallback_flat_cost"] += 1 + return None + snapshot = self._ready_older_snapshot() + if snapshot is None: + self.stats["fallback_no_snapshot"] += 1 + return None + selected = self._gather_rows(num_gen_requests=n, rows=rows, + snapshot=snapshot) + if selected is None: + return None + self._note_snapshot_stats(selected, rows, + stamps=self._prev_stamps, + staged_seq=self._prev_seq) + survival = compute_survival(self.apply_calibration(selected)) + budget = compute_verify_token_budget( + survival=survival.numpy().astype(np.float64), + num_gen_requests=n, + cost_table=self.cost_table, + min_verify_len=self.cfg.min_verify_len, + max_verify_len=self.cfg.resolved_max_verify_len, + allowed_lens=self.tiers, + ) + floor_tokens = n * (self.cfg.min_verify_len + 1) + predicted = float(self.cost_table.step_times( + np.asarray([floor_tokens + int(budget)]), n)[0]) + self.last_predicted_step_ms = predicted + self.stats["predicted_ms_sum"] = ( + self.stats.get("predicted_ms_sum", 0.0) + predicted) + self.stats["predicted_steps"] = ( + self.stats.get("predicted_steps", 0) + 1) + budget = max(0, min(int(budget), n * trimmable)) + base, extra = divmod(budget, n) + shape_lens = [ + min(self.cfg.min_verify_len + base + (1 if i < extra else 0), + self.cfg.resolved_max_verify_len) for i in range(n) + ] + return budget, shape_lens + def decide_verify_lens( self, *, diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 6769ac5b5d6b..c2fbb9c8215e 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -49,6 +49,11 @@ class SampleStateTensorsSpec(SampleStateTensors): next_draft_tokens: torch.Tensor #: cap-accept only; None on every other path. cap_trim_lens: Optional[torch.Tensor] = None + #: ragged verification only; None on every other path. The TOKEN windows + #: the step actually verified (bonus included), slot-indexed. Under + #: device-window selection this D2H is the host's only source of the true + #: windows; zero marks "no window this step". + verify_lens: Optional[torch.Tensor] = None @dataclass(kw_only=True) @@ -112,6 +117,10 @@ class Store: #: the batch each step -- a slot left alone would hand the previous #: occupant's loss to whoever holds it now. cap_trim_lens: torch.Tensor + #: ragged verification: the TOKEN windows the step actually verified, + #: slot-indexed; zeroed for every batch slot on non-ragged steps (the + #: same stale-slot discipline as cap_trim_lens). + verify_lens: torch.Tensor def __init__(self, args: TorchSampler.Args, *, draft_len: int): """ @@ -139,6 +148,7 @@ def __init__(self, args: TorchSampler.Args, *, draft_len: int): next_draft_tokens=int_tensor((seq_slots, draft_tokens_size)), new_tokens_lens=int_tensor((seq_slots,)), cap_trim_lens=int_tensor((seq_slots,)), + verify_lens=int_tensor((seq_slots,)), ) def _get_max_tokens(self, args: TorchSampler.Args, draft_len: int) -> int: @@ -218,7 +228,8 @@ def _request_common_handling( @staticmethod def _verified_len(request: LlmRequest, runtime_draft_len: int, - verify_lens_snapshot: Optional[dict] = None) -> int: + verify_lens_snapshot: Optional[dict] = None, + ridden_verify_lens: Optional[list] = None) -> int: """How many draft positions THIS request was given to verify. Uniform scheduling gives every request the batch-wide @@ -234,7 +245,16 @@ def _verified_len(request: LlmRequest, runtime_draft_len: int, ``py_verify_len`` is set only by the ragged path, so this is exactly today's behavior for every other speculation mode. """ - # Prefer the snapshot taken when this step's tokens were sampled. The + # First preference: the windows the device layout actually verified, + # ridden back on this step's sampler D2H (slot-indexed TOKEN windows; + # zero = no window). Under device-window selection the host attribute + # holds only a shape split, so this is the only correct source; under + # host windows the two agree. + if ridden_verify_lens is not None: + token_window = ridden_verify_lens[request.py_seq_slot] + if token_window > 0: + return int(token_window) - 1 + # Second: the snapshot taken when this step's tokens were sampled. The # live attribute already belongs to the NEXT step by the time the # overlap scheduler rewinds this one. if verify_lens_snapshot is not None: @@ -263,6 +283,8 @@ def update_requests( new_tokens_lens_list = state.host.new_tokens_lens.tolist() cap_trim_list = (state.host.cap_trim_lens.tolist() if state.host.cap_trim_lens is not None else None) + ridden_verify_lens = (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) @@ -286,7 +308,8 @@ def update_requests( break req.py_num_accepted_draft_tokens = num_new_tokens - 1 verified_len = self._verified_len(req, runtime_draft_len, - state.verify_lens_snapshot) + state.verify_lens_snapshot, + ridden_verify_lens) req.py_rewind_len = verified_len - req.py_num_accepted_draft_tokens # Acceptance against the window the request was actually given. # Recorded here because this is the only place both numbers are @@ -405,6 +428,18 @@ def sample_async( o_cap_trim = o_cap_trim[num_skip:num_skip + num_sampling_requests] self.store.cap_trim_lens.index_copy_( 0, slots, o_cap_trim.to(self.store.cap_trim_lens.dtype)) + # Ragged verification: the true token windows, same zero-then-write + # slot discipline as cap_trim_lens. Under device-window selection the + # host py_verify_len holds only a shape split, so the rewind + # arithmetic in update_requests must read THIS. + o_verify_lens = outputs.get("verify_lens") + if o_verify_lens is None: + o_verify_lens = torch.zeros_like(o_new_tokens_lens) + else: + o_verify_lens = o_verify_lens[num_skip:num_skip + + num_sampling_requests] + self.store.verify_lens.index_copy_( + 0, slots, o_verify_lens.to(self.store.verify_lens.dtype)) # Create sample state with async D2H copy device_tensors = SampleStateTensorsSpec( @@ -420,6 +455,7 @@ def sample_async( # [max_num_sequences] int32 -- kilobytes next to new_tokens above, # and it rides the same event, so no extra synchronization. cap_trim_lens=self._copy_to_host(self.store.cap_trim_lens), + verify_lens=self._copy_to_host(self.store.verify_lens), ) sampler_event = self._record_sampler_event() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py index 54bf4bf351fd..f318d5f46300 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py @@ -122,6 +122,74 @@ def test_row_maps_match_prepare_semantics(): 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 -- + the host fit's copy widths depend on this exact split.""" + 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 + + +def test_select_windows_respects_published_split(): + """With pad_len the prologue must land on the host fit's real/pad token + split exactly: real tokens = bucket - n_pad * pad_len, pads = pad_len.""" + gen = torch.Generator().manual_seed(41) + K = 5 + cfg = DSparkScheduleConfig(block_size=K, min_verify_len=1) + max_tok = cfg.resolved_max_verify_len + 1 + for _ in range(30): + n_real = int(torch.randint(1, 33, (1,), generator=gen)) + n_pad = int(torch.randint(0, 5, (1,), generator=gen)) + padded_bs = n_real + n_pad + pad_len = int(torch.randint(1, max_tok + 1, (1,), generator=gen)) + real_target = int( + torch.randint(n_real * 2, n_real * max_tok + 1, (1,), + generator=gen)) + bucket = real_target + n_pad * pad_len + budget = max(0, min(int(torch.randint(0, n_real * K + 1, (1,), + generator=gen)), + real_target - n_real * 2)) + conf = torch.randn(padded_bs + 4, K, generator=gen) * 2 + slot_idx = torch.arange(padded_bs, dtype=torch.long) + res = select_windows_device( + confidence_logits=conf, + slot_idx=slot_idx, + num_real=torch.tensor(n_real), + budget=torch.tensor(budget), + graph_num_tokens=bucket, + cfg=cfg, + pad_len=pad_len, + ) + assert int(res.verify_lens.sum()) == bucket + assert int(res.verify_lens[:n_real].sum()) == real_target + assert (res.verify_lens[n_real:] == pad_len).all() + assert int(res.verify_lens[:n_real].min()) >= cfg.min_verify_len + 1 + assert int(res.verify_lens[:n_real].max()) <= max_tok + + def test_topk_tensor_budget_matches_int(): """The 0-d tensor budget path (in-graph capacity knob) must allocate identically to the host int path, including budget <= 0.""" @@ -140,6 +208,44 @@ def test_topk_tensor_budget_matches_int(): assert torch.equal(host, dev) +def test_planner_lag2_ring_and_budget_split(monkeypatch): + """Device-window mode: the budget reads the snapshot one staging OLDER + (lag-2), and the shape lens spread the budget uniformly, summing to + floor + budget.""" + monkeypatch.setenv("TLLM_DSPARK_DEVICE_WINDOWS", "1") + from tensorrt_llm._torch.speculative.dspark_verify import \ + DSparkVerifyPlanner + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + planner = DSparkVerifyPlanner(cfg=cfg) + assert planner.device_windows + + snap_a = torch.full((4, 5), 1.0) + snap_b = torch.full((4, 5), 2.0) + snap_c = torch.full((4, 5), 3.0) + planner.stage_confidence(snap_a) + assert planner._ready_older_snapshot() is None + planner.stage_confidence(snap_b) + older = planner._ready_older_snapshot() + assert older is not None and float(older[0, 0]) == 1.0 + planner.stage_confidence(snap_c) + older = planner._ready_older_snapshot() + assert float(older[0, 0]) == 2.0 + # The current snapshot is the freshest staging, untouched by the ring. + assert float(planner._ready_snapshot()[0, 0]) == 3.0 + + # Budget arithmetic through the frac override (snapshot-independent): + # shape lens sum to floor + budget and stay inside the bounds. + planner._forced_budget_frac = 0.5 + n = 7 + trimmable = cfg.resolved_max_verify_len - cfg.min_verify_len + budget, lens = planner.decide_verify_budget(num_gen_requests=n) + assert budget == int(round(0.5 * n * trimmable)) + assert len(lens) == n + assert sum(lens) == n * cfg.min_verify_len + budget + assert min(lens) >= cfg.min_verify_len + assert max(lens) <= cfg.resolved_max_verify_len + + def _host_select(rows, budget, bucket, padded_bs, cfg): """The lag-free host reference: calibrate -> survival -> topk -> fill.""" survival = compute_survival(torch.sigmoid(rows)) From 9d3858663eb6f6a19bc9afdd93e096ab224caf66 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Fri, 7 Aug 2026 04:08:46 -0700 Subject: [PATCH 104/177] [None][fix] cap-accept: never bake a capless capture into the graph apply_accept_caps reads spec_metadata.accept_caps as a Python attribute, so its None-check is evaluated at CAPTURE time: warmup batches carry no py_verify_cap, _attach_accept_caps published None, and the clamp was recorded as a no-op -- every replay silently skipped it. The positive control (window pinned to 1 at ~1.3 accepted) returned zero trims over 4679 requests, and unclamped accept_len exceeded the cap, which is how this surfaced; only eager steps ever clamped. In cap-accept mode, capless requests now take the full drafted block -- a numeric no-op clamp that keeps the ops and the persistent buffer address inside the captured graph. Other modes keep the drop-and-warn behavior for partial caps. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 46 +++++++++++++------ 1 file changed, 32 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index ce95136fbf88..8987d85cd0b8 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4062,21 +4062,39 @@ def _attach_accept_caps(self, spec_metadata, generation_requests) -> None: getattr(request, "py_verify_cap", None) for request in generation_requests ] + worker = self._get_spec_worker() + mode = getattr(worker, "ragged_verify_mode", None) if worker else None + cap_accept_mode = (mode is not None and mode.computes_windows + and not mode.trims_submitted_tokens) if not caps or any(cap is None for cap in caps): - # Partial caps are neither mode: drop them all, but warn, because - # the step then silently runs as `static`. - if caps and any(cap is not None for cap in caps): - if not getattr(self, "_warned_partial_accept_caps", False): - self._warned_partial_accept_caps = True - logger.warning( - "DSpark cap-accept: %d of %d generation requests carry " - "no verify cap, so caps were dropped for the whole " - "batch and this step ran as 'static'. The mode is " - "measuring nothing on steps like this one.", - sum(1 for cap in caps if cap is None), len(caps)) - spec_metadata.accept_caps = None - spec_metadata.accept_cap_trim = None - return + if cap_accept_mode and caps: + # cap-accept mode: NEVER publish None. `apply_accept_caps` + # reads `spec_metadata.accept_caps` as a Python attribute at + # CAPTURE time, so a None on the capture step (warmup batches + # carry no caps) bakes a no-op into the graph and every replay + # silently skips the clamp -- the zero-trim positive-control + # failure. Capless requests get the full block: numerically a + # no-op clamp, but the ops and buffer address stay captured. + caps = [ + self.spec_config.max_draft_len if cap is None else cap + for cap in caps + ] + else: + # Any other mode: partial caps are meaningless; drop them all, + # but warn, because the step then silently runs as `static`. + if caps and any(cap is not None for cap in caps): + if not getattr(self, "_warned_partial_accept_caps", False): + self._warned_partial_accept_caps = True + logger.warning( + "DSpark cap-accept: %d of %d generation requests " + "carry no verify cap, so caps were dropped for " + "the whole batch and this step ran as 'static'. " + "The mode is measuring nothing on steps like " + "this one.", + sum(1 for cap in caps if cap is None), len(caps)) + spec_metadata.accept_caps = None + spec_metadata.accept_cap_trim = None + return n = len(caps) # Persistent buffer: acceptance runs inside the captured graph, so a From 63dcc7d40c25a5b2ea0790aad666298c9e17e565 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Fri, 7 Aug 2026 04:45:59 -0700 Subject: [PATCH 105/177] [None][feat] Budget-override scheduling for cap-accept differentials decide_verify_lens gains a diagnostic budget_override that spends a caller-supplied budget through the normal current-snapshot top-k, and the executor wires it to cap-accept + device-window mode: the lag-2 argmax picks the budget, the block still executes in full, and the windows ride py_verify_cap as accounting caps. cap_trim_tokens then reports exactly how many target-accepted tokens that budget would have discarded, so two budgets can be compared at identical compute -- the instrument the first cap-accept differential ran on (and which caught the capless-capture bug). Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 18 ++++++++++++++++++ .../_torch/speculative/dspark_verify.py | 18 +++++++++++++++++- 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 55a25dbbf86d..28f6e1724bb3 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3398,6 +3398,24 @@ def _dspark_confidence_draft_len(self, fallback_reason = "planner_declined" else: device_budget, ragged_lens = pair + elif planner.device_windows: + # cap-accept + device mode: a pure measurement configuration. + # The lag-2 argmax picks the budget, the normal top-k spends + # it, and the block still executes IN FULL -- the windows ride + # py_verify_cap as accounting caps, so cap_trim_tokens reports + # exactly how many target-accepted tokens this budget would + # have discarded. Differencing against plain cap-accept + # (lag-1 budget) isolates the budget's true acceptance cost + # at identical compute. + pair = planner.decide_verify_budget( + num_gen_requests=num_gen_requests, rows=confidence_rows) + ragged_lens = (planner.decide_verify_lens( + num_gen_requests=num_gen_requests, + rows=confidence_rows, + reduce_across_ranks=False, + budget_override=pair[0]) if pair is not None else None) + if ragged_lens is None: + fallback_reason = "planner_declined" else: ragged_lens = planner.decide_verify_lens( num_gen_requests=num_gen_requests, diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index 7054e268d64c..635483c566a2 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -524,6 +524,7 @@ def decide_verify_lens( rows: Optional[Sequence[int]] = None, all_rank_max: Optional[Callable[[int], int]] = None, reduce_across_ranks: bool = True, + budget_override: Optional[int] = None, ) -> Optional[List[int]]: """Per-request verify lengths for a ragged step, or None to stay uniform. @@ -542,7 +543,22 @@ def decide_verify_lens( self.stats["fallback_no_gen_requests"] += 1 return None - if self._forced_verify_len is not None: + if budget_override is not None: + # Diagnostic: spend a caller-supplied budget (e.g. the lag-2 + # argmax's) through the normal CURRENT-snapshot top-k. Under + # cap-accept this measures a schedule's true acceptance cost -- + # execution stays full-block, the window is only an accounting + # cap -- so two budgets can be compared with identical compute. + selected = self._gather_rows(num_gen_requests=num_gen_requests, + rows=rows) + if selected is None: + return None + self._note_snapshot_stats(selected, rows) + survival = compute_survival(self.apply_calibration(selected)) + lens = schedule_verify_lens_topk(survival=survival, + budget=int(budget_override), + cfg=self.cfg).tolist() + elif self._forced_verify_len is not None: # Must stay ahead of the cost-table gate (the profiler runs without # a table) and must set `lens` rather than return, so the pinned # path still runs the cross-rank agreement below. From f5b1adef2b701f3e9ac674c05d14314b75e725f2 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Fri, 7 Aug 2026 04:46:09 -0700 Subject: [PATCH 106/177] [None][feat] Profiler: 128-step warmup default and --engine-tiers Per-cell warmup discard defaults to 128 iterations (empirically the existing 16 was sufficient in single-session sweeps -- the first cell's half-split medians agree to 0.1% -- but 128 defends colder setups at ~50s per sweep). --engine-tiers overrides the sweep engine's derived tier ladder: a denser ladder quantises the frac knob to more rungs, so one sweep measures more M columns per batch size; the dense-tier experiments use 2,3,4,5. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- tests/microbenchmarks/dspark_sps_profiler.py | 29 ++++++++++++++++++-- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/tests/microbenchmarks/dspark_sps_profiler.py b/tests/microbenchmarks/dspark_sps_profiler.py index 1a53c92ae642..01bc428335e4 100644 --- a/tests/microbenchmarks/dspark_sps_profiler.py +++ b/tests/microbenchmarks/dspark_sps_profiler.py @@ -1134,11 +1134,21 @@ class SweepConfig: enable_attention_dp: bool = False max_draft_len: int = 5 input_len: int = 1024 - warmup_steps: int = 16 + # 128, not a token gesture: autotuner cache misses ("using the fallback + # tactic") and allocator growth are still visible well past step 16 on a + # fresh engine, and a cell that starts on the cool side biases theta UP + # for exactly the small-M cells the argmax is most sensitive to. + warmup_steps: int = 128 measure_steps: int = 64 repeats: int = 1 min_samples: int = 16 max_seq_len: int = 4096 + # Frac sweeps only: override the engine's DERIVED tier ladder (cost-table + # shelf edges, typically [1,3,5]). A denser ladder makes the frac knob + # quantise to more rungs, so one sweep measures more M columns per bs -- + # the dense-tier experiments need the odd ratios {3,5} the default ladder + # cannot execute. + engine_tiers: Optional[List[int]] = None max_num_tokens: int = 8192 kv_cache_fraction: float = 0.5 #: Fitted STS temperatures. Not used to decide anything while the pin is on; @@ -1414,7 +1424,10 @@ def _build_llm(config: SweepConfig, verify_len: int, *, # derived from it, and a one-rung ladder would collapse every # captured bucket to the full block -- the sweep would then only # ever measure M = bs * (block + 1) regardless of the frac. - (None if not pin_verify_len else [int(verify_len)]) + # ``engine_tiers`` overrides the derived ladder for dense grids. + (([int(t) for t in config.engine_tiers] + if config.engine_tiers else None) + if not pin_verify_len else [int(verify_len)]) if config.enable_confidence_scheduling else None ), ) @@ -1867,6 +1880,14 @@ def _build_parser() -> argparse.ArgumentParser: "per-step M measured from the group's agreed bucket). Mutually " "exclusive with --verify-lens and --block-follows-verify-len.", ) + sweep.add_argument( + "--engine-tiers", + default=None, + help="Frac sweeps only: comma-separated tier ladder for the SWEEP " + "engine (e.g. 2,3,4,5), overriding the cost-table-derived default. " + "A denser ladder quantises the frac knob to more rungs, so one sweep " + "measures more M columns per batch size.", + ) sweep.add_argument("--input-len", type=int, default=1024, help="Prompt length in tokens.") sweep.add_argument("--block-follows-verify-len", action="store_true", help="Shrink the drafted block with the verify window " @@ -1886,7 +1907,7 @@ def _build_parser() -> argparse.ArgumentParser: help="Fitted STS temperatures. Required with confidence " "scheduling: ragged verification will not validate " "without either this or the cost table being built.") - sweep.add_argument("--warmup-steps", type=int, default=16) + sweep.add_argument("--warmup-steps", type=int, default=128) sweep.add_argument("--measure-steps", type=int, default=64) sweep.add_argument("--min-samples", type=int, default=16) sweep.add_argument("--seed", type=int, default=1234) @@ -2456,6 +2477,8 @@ def main(argv: Optional[Sequence[str]] = None) -> int: pin_acceptance=BONUS_TOKEN_ONLY_ACCEPTANCE if args.pin_acceptance else 0.0, seed=args.seed, extra_llm_api_options=extra, + engine_tiers=(_int_list(args.engine_tiers) + if args.engine_tiers else None), ) if fracs: samples = run_frac_sweep(config) From 19502c8c28ffbe6176fe7eb547d876367bc82d2f Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sat, 8 Aug 2026 06:36:56 -0700 Subject: [PATCH 107/177] [None][fix] Device-window kv_lens delta counts the window twice The host stages kv_lens as num_cached + seq_lens_kv, and BOTH terms bake the per-request token window: kv_lens = past + 2*S. The device-window prologue's rebase to the true windows w therefore needs 2*(w - S) on the kv side (paired with -(w - S) on previous_kv_lens_offsets); the previous single delta left every re-ranked request's kv_len off by (w - S), which shifted the indexer K-cache slot mapping and silently wrote K entries into the wrong cache slots. Verified by an in-process A/B tensor diff against a full host restage with identical windows; the pairing is locked by test_prologue_kv_pairing_matches_host_staging. Also ships the (env-gated, default-off) debug harness used to corner the remaining device-window corruption, kept temporarily so a second reviewer can reproduce the investigation: identity/steep window overrides, host-steep publishing, forced-eager/argmax paths, capture pass controls, the staging A/B differ, and ride/accept-bound logging. See docs/dspark_sts_program.md for the full evidence chain; all of it is slated for removal once the remaining bug is fixed. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 248 +++++++++++++++++- tensorrt_llm/_torch/speculative/interface.py | 8 +- .../_torch/speculative/spec_sampler_base.py | 23 ++ .../test_dspark_device_fill_parity.py | 44 ++++ 4 files changed, 310 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 8987d85cd0b8..8f41e0beb12e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2005,7 +2005,9 @@ def prepare_cross_batch(batch: ScheduledRequests, request._cached_tokens = cached_tokens request._cached_tokens_set = cached_tokens_set - def _run_capture_pass(force_non_greedy: bool, label: str) -> None: + def _run_capture_pass(force_non_greedy: bool, + label: str, + entries=None) -> None: spec_metadata = self.spec_metadata if force_non_greedy and spec_metadata is not None: spec_metadata._force_non_greedy_for_capture = True @@ -2015,11 +2017,17 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: # in this pass uses the non-greedy key; populate's override # below will keep it False on every subsequent iteration. spec_metadata.is_all_greedy_sample = False + elif spec_metadata is not None: + # Symmetric pre-flip for interleaved capture: a greedy capture + # right after an advanced one would otherwise key on the stale + # False before populate corrects it. + spec_metadata.is_all_greedy_sample = True _graph_mem_before = (torch.cuda.memory_reserved(), torch.cuda.memory_allocated()) _graph_count = 0 try: - for entry in graphs_to_capture: + for entry in (graphs_to_capture + if entries is None else entries): # Ragged verification adds a token-count axis, so entries # may be (bs, draft_len) or (bs, draft_len, token_bucket). bs, draft_len = entry[0], entry[1] @@ -2090,9 +2098,32 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: # the LAST shape's draft length; restore it after the passes, before # the generic token warmup builds requests at the full draft length. saved_runtime_draft_len = self.runtime_draft_len + needs_non_greedy_precheck = ( + self.spec_config is not None + and self.spec_config.spec_dec_mode.use_one_engine() + and os.environ.get("TLLM_DSPARK_SKIP_ADVANCED_CAPTURE") != "1") + if (needs_non_greedy_precheck and os.environ.get( + "TLLM_DSPARK_INTERLEAVED_CAPTURE") == "1"): + # Debug/fix candidate: capture both sampling variants of each + # shape back-to-back, so the two graphs bind under identical + # pool/workspace/persistent-buffer conditions instead of a whole + # sweep apart. + for _entry in graphs_to_capture: + _run_capture_pass(force_non_greedy=False, + label="greedy", + entries=[_entry]) + _run_capture_pass(force_non_greedy=True, + label="advanced sampling", + entries=[_entry]) + self.enable_spec_decode = self.is_spec_decode + self.runtime_draft_len = saved_runtime_draft_len + if self.spec_metadata is not None: + self.spec_metadata.is_all_greedy_sample = True + return # Pass 1: greedy fast-path (dummy requests carry no sampling params, # so is_all_greedy_sample is naturally True). - _run_capture_pass(force_non_greedy=False, label="greedy") + if os.environ.get("TLLM_DSPARK_SKIP_GREEDY_CAPTURE") != "1": + _run_capture_pass(force_non_greedy=False, label="greedy") # Pass 2: advanced sampling variant. Required because on-the-fly capture # is disabled outside warmup, so any inference batch that contains a # non-greedy request would otherwise fall back to eager. Only meaningful @@ -2101,7 +2132,8 @@ def _run_capture_pass(force_non_greedy: bool, label: str) -> None: # this variant. needs_non_greedy_capture = ( self.spec_config is not None - and self.spec_config.spec_dec_mode.use_one_engine()) + and self.spec_config.spec_dec_mode.use_one_engine() + and os.environ.get("TLLM_DSPARK_SKIP_ADVANCED_CAPTURE") != "1") if needs_non_greedy_capture: _run_capture_pass(force_non_greedy=True, label="advanced sampling") # Set the value back to the original value after cuda graph warmups are complete @@ -3833,8 +3865,23 @@ def fit_ragged_verify_lens( runner.agreed_ragged_bucket = int(bucket) runner.ragged_pad_verify_len = pad_len - 1 - for request, tokens in zip(generation_requests, - filled.verify_lens.tolist()): + published = filled.verify_lens.tolist() + if os.environ.get("TLLM_DSPARK_HOST_STEEP") == "1" and n_real > 1: + # Debug bisector: force truly non-uniform windows through the + # HOST staging (the reference implementation, no device + # prologue). The production host scheduler only ever publishes + # uniform windows (budget quantized to the tier ladder), so the + # ragged consumers' non-uniform handling is otherwise exercised + # only by the device path -- this isolates them. + lo = 1 + getattr(self.spec_config, "min_verify_len", 1) + hi = max_verify_len + pairs = (n_real // 2) * 2 + for i in range(0, pairs, 2): + give = min(published[i] - lo, hi - published[i + 1]) + if give > 0: + published[i] -= give + published[i + 1] += give + for request, tokens in zip(generation_requests, published): request.py_verify_len = int(tokens) - 1 return int(bucket) @@ -3901,6 +3948,9 @@ def _apply_device_window_prologue(self, inputs, split_lens = lens_buf[:padded_bs].clone() split_qo = qo_buf[:padded_bs + 1].to(torch.long) past_seen = self.position_ids_cuda[split_qo[:-1]].clone() + if os.environ.get("TLLM_DSPARK_AB_DUMP") == "1": + self._dbg_pre_kv = attn_metadata.kv_lens_cuda[:padded_bs].clone() + self._dbg_split = split_lens.clone() expected_stamp = worker.verified_draft_seq_cuda() stamps = worker.confidence_stamp_buffer( @@ -3917,13 +3967,54 @@ def _apply_device_window_prologue(self, inputs, expected_stamp=expected_stamp, pad_len=pad_len_tok, ) + debug_mode = os.environ.get("TLLM_DSPARK_DEVICE_IDENTITY", "") + if debug_mode in ("1", "steep"): + # Debug bisectors. "1": exercise every prologue write with the + # shape split itself (w := S) -- corruption surviving this is in + # the rewrite mechanics, not the window semantics. "steep": + # deterministic max-amplitude ragged windows (adjacent real rows + # trade tokens pairwise, conserving the total) -- forces w != S + # while staying independent of the confidence-gather chain. + from ..speculative.dspark_ragged import (build_qo_indptr, + build_row_maps_device) + id_lens = split_lens.clone() + if debug_mode == "steep" and n_real > 1: + lo = cfg.min_verify_len + 1 + hi = cfg.resolved_max_verify_len + 1 + pairs = (n_real // 2) * 2 + a = id_lens[0:pairs:2] + b = id_lens[1:pairs:2] + give = torch.minimum(a - lo, hi - b).clamp_min(0) + a -= give + b += give + id_req, id_corr = build_row_maps_device(id_lens, + graph_num_tokens=bucket) + result = type(result)(verify_lens=id_lens, + qo_indptr=build_qo_indptr(id_lens).to( + result.qo_indptr.dtype), + req_idx=id_req, + kv_correction=id_corr) lens_buf[:padded_bs].copy_(result.verify_lens) qo_buf[:padded_bs + 1].copy_(result.qo_indptr) - # Every host-staged kv formula used the shape split consistently, so - # correcting by the difference lands on exactly the value the host - # path would have produced with these windows. - attn_metadata.kv_lens_cuda[:padded_bs] += result.verify_lens - split_lens + # The host stages kv_lens as num_cached + seq_lens_kv, and BOTH terms + # bake the per-request token window (num_cached = past + tokens; + # seq_lens_kv = tokens), so kv_lens = past + 2*S -- the window counts + # TWICE. Moving to the true windows therefore needs 2*(w - S), not + # (w - S): the single-delta variant left every re-ranked request's + # kv_len off by (w - S), which shifted the indexer K-cache slot + # mapping (slot_mapping_fp8) and silently wrote K entries into the + # wrong cache slots. Established empirically by an A/B tensor diff + # against a full host restage with the same windows. + window_delta = result.verify_lens - split_lens + attn_metadata.kv_lens_cuda[:padded_bs] += 2 * window_delta + # The graph adds previous_kv_lens_offsets (staged as new_tokens_lens - + # shape_lens, per request) to kv_lens during replay; host-with-w + # stages new_tokens_lens - w there, so the offsets move by -(w - S). + # Combined: (past + 2S) + 2(w-S) + (new - S) - (w-S) = past + w + new, + # exactly the host-with-w in-graph sum. + self.previous_kv_lens_offsets_cuda[:padded_bs] -= window_delta.to( + self.previous_kv_lens_offsets_cuda.dtype) req_idx = result.req_idx device = req_idx.device @@ -3960,8 +4051,108 @@ def _apply_device_window_prologue(self, inputs, attn_metadata.apply_device_ragged_layout(result.verify_lens, req_idx, result.kv_correction) + if os.environ.get("TLLM_DSPARK_RIDE_LOG") == "1": + self._ride_log_ctr = getattr(self, "_ride_log_ctr", 0) + 1 + if self._ride_log_ctr % 64 == 0: + logger.info( + f"[ride-log] prologue #{self._ride_log_ctr} " + f"S={split_lens[:4].tolist()} w={result.verify_lens[:4].tolist()}") return True + def _dspark_ab_diff(self, padded_requests, kv_cache_manager, attn_metadata, + spec_metadata, new_tensors_device, + cache_indirection_buffer, num_accepted_tokens_device, + req_id_to_old_request, resource_manager, + can_run_graph) -> None: + """Debug: diff the prologue's rewritten staging against a host restage. + + Runs once, eager-only. Snapshot every tensor the forward reads after + the prologue rewrote the layout to the device windows w; then stamp w + onto each request's py_verify_len and rerun the full host staging (the + reference implementation, proven correct for non-uniform windows); + diff. Any mismatch names the artifact the prologue fails to mirror. + """ + torch.cuda.synchronize() + bucket = int(self.cuda_graph_runner.agreed_ragged_bucket or 0) + padded_bs = int(self._dspark_last_padded_bs) + gen_requests = padded_requests.generation_requests + n_real = int(self._dspark_last_num_real) + w_tokens = self.ragged_verify_lens_cuda[:padded_bs].tolist() + + def snap(): + fields = { + "input_ids": self.input_ids_cuda[:bucket], + "position_ids": self.position_ids_cuda[:bucket], + "draft_tokens": self.draft_tokens_cuda[:bucket], + "prev_pos_indices": self.previous_pos_indices_cuda[:bucket], + "prev_pos_offsets": self.previous_pos_id_offsets_cuda[:bucket], + "prev_kv_offsets": self.previous_kv_lens_offsets_cuda[:padded_bs], + "kv_lens": attn_metadata.kv_lens_cuda[:padded_bs], + "seq_lens_cuda": attn_metadata.seq_lens_cuda[:padded_bs], + "verify_lens": self.ragged_verify_lens_cuda[:padded_bs], + "qo_indptr": self.ragged_qo_indptr_cuda[:padded_bs + 1], + } + for name in ("gen_token_repeats_cuda", "row_req_idx_cuda", + "row_kv_correction_cuda", "attn_row_req_idx_cuda", + "attn_row_kv_correction_cuda", "req_idx_per_token", + "kv_lens_expanded_cuda", "block_table_expanded", + "slot_mapping_fp8", "gen_kv_indptr", + "gen_cached_token_indptr"): + t = getattr(attn_metadata, name, None) + if isinstance(t, torch.Tensor): + n = bucket if t.shape[0] >= bucket else t.shape[0] + fields[name] = t[:n] + for name in ("temperatures", "top_ks", "top_ps", "batch_slot_ids"): + t = getattr(spec_metadata, name, None) + if isinstance(t, torch.Tensor): + n = min(bucket, t.shape[0]) + fields["spec." + name] = t[:n] + return {k: v.detach().clone() for k, v in fields.items()} + + snap_a = snap() + prev_idx = getattr(self, "_dspark_ab_prev_idx", {}) + for request, tokens in zip(gen_requests, w_tokens): + request.py_verify_len = int(tokens) - 1 + if request.py_request_id in prev_idx: + request.py_batch_idx = prev_idx[request.py_request_id] + try: + self._prepare_inputs(padded_requests, kv_cache_manager, + attn_metadata, spec_metadata, + new_tensors_device, cache_indirection_buffer, + num_accepted_tokens_device, + req_id_to_old_request, resource_manager, + can_run_graph) + except Exception as exc: + logger.warning(f"[ab-diff] restage failed, disabling: {exc}") + self._dspark_ab_prints = 999 + return + torch.cuda.synchronize() + snap_b = snap() + self._dspark_ab_step = getattr(self, "_dspark_ab_step", 0) + 1 + # Expected-divergent names: staged from the shape split at prepare + # and rebuilt in-graph from the corrected kv_lens before any read. + benign = {"kv_lens_expanded_cuda", "slot_mapping_fp8"} + for name, a in snap_a.items(): + b = snap_b[name] + if name in benign or a.shape != b.shape: + continue + neq = (a != b) + cnt = int(neq.sum()) + if cnt: + idx = int(neq.reshape(-1).nonzero()[0]) + self._dspark_ab_prints = getattr(self, "_dspark_ab_prints", + 0) + 1 + logger.info( + f"[ab-diff] step#{self._dspark_ab_step} {name}: " + f"{cnt}/{a.numel()} MISMATCH first@{idx} " + f"prologue={a.reshape(-1)[idx].item()} " + f"host={b.reshape(-1)[idx].item()} " + f"(bs={padded_bs} n_real={n_real} bucket={bucket})") + if self._dspark_ab_step % 200 == 0: + logger.info(f"[ab-diff] step#{self._dspark_ab_step} checkpoint: " + f"{getattr(self, '_dspark_ab_prints', 0)} mismatch " + f"lines so far") + @staticmethod def _ragged_token_lens(generation_requests) -> Optional[List[int]]: """Each generation request's token window, or None if the batch is uniform. @@ -7055,6 +7246,13 @@ def forward(self, spec_resource_manager=spec_resource_manager, ) + if (key is not None and not self.is_warmup and os.environ.get( + "TLLM_DSPARK_FORCE_EAGER") == "1"): + # Debug bisector: capture normally (the ragged fit needs the + # captured bucket set) but veto every replay, so the exact + # same staged step runs eagerly. Splits capture-interaction + # bugs from layout-semantics bugs. + key = None can_run_graph = key is not None # A ragged step that misses its captured shape quietly runs eager; # count it. Only generation-only steps count: a batch with context @@ -7084,6 +7282,15 @@ def forward(self, padded_requests.generation_requests, ) + if os.environ.get("TLLM_DSPARK_AB_DUMP") == "1": + # The staging loop overwrites py_batch_idx (prev-batch index + # -> seq slot); the A/B restage must see the original or its + # reference staging reads the wrong new_tokens rows on any + # slot-churn step. + self._dspark_ab_prev_idx = { + r.py_request_id: r.py_batch_idx + for r in padded_requests.generation_requests + } inputs, gather_ids = self._prepare_inputs( padded_requests, kv_cache_manager, attn_metadata, spec_metadata, new_tensors_device, cache_indirection_buffer, @@ -7097,10 +7304,27 @@ def forward(self, # replay (stream-ordered, no host sync). Graph steps only: eager # and capture steps run the host shape split, which is a valid # window assignment. - if (can_run_graph and not self.is_warmup + if ((can_run_graph or os.environ.get( + "TLLM_DSPARK_DEVICE_EAGER_OK") == "1") + and not self.is_warmup and getattr(self, "_dspark_device_budget", None) is not None): - self._apply_device_window_prologue(inputs, new_tensors_device) + applied = self._apply_device_window_prologue( + inputs, new_tensors_device) + if applied and os.environ.get("TLLM_DSPARK_AB_DUMP") == "1": + self._dspark_ab_applied = getattr( + self, "_dspark_ab_applied", 0) + 1 + if (applied and os.environ.get("TLLM_DSPARK_AB_DUMP") == "1" + and getattr(self, "_dspark_ab_applied", 0) >= int( + os.environ.get("TLLM_DSPARK_AB_AFTER", "0")) + and getattr(self, "_dspark_ab_prints", 0) < 40): + self._dspark_ab_diff(padded_requests, kv_cache_manager, + attn_metadata, spec_metadata, + new_tensors_device, + cache_indirection_buffer, + num_accepted_tokens_device, + req_id_to_old_request, + resource_manager, can_run_graph) with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): if not can_run_graph: diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 342dd89bdd3c..0e5bd5ae309d 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -984,6 +984,11 @@ def populate_sampling_params_for_one_model( # read. Skip the heavier H->D copies. if self.is_all_greedy_sample: return + if os.environ.get("TLLM_DSPARK_SKIP_SAMPLING_H2D") == "1": + # Debug bisector (pair with TLLM_DSPARK_FORCE_ARGMAX, which stops + # anything reading these buffers): isolates whether the per-step + # non-greedy H2D staging itself is the corruption channel. + return # Phase 2: build per-token / per-request lists and copy to GPU. temperatures: list[float] = [] @@ -2348,7 +2353,8 @@ def _sample_tokens_for_batch( Returns: sampled_tokens: [num_tokens] - Sampled token ids """ - if not spec_metadata.is_all_greedy_sample: + if not spec_metadata.is_all_greedy_sample and os.environ.get( + "TLLM_DSPARK_FORCE_ARGMAX") != "1": # Use logits.shape[0] directly: for PARD under CUDA graph capture # runtime_draft_len may reflect the PARD-max while the captured # graph was built for a shorter draft_len, causing a shape mismatch diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index c2fbb9c8215e..bd1312a7d615 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -24,6 +24,8 @@ import torch +import os + from tensorrt_llm.logger import logger from ..pyexecutor.llm_request import LlmRequest, LlmRequestState @@ -294,6 +296,20 @@ def update_requests( if req.state == LlmRequestState.GENERATION_COMPLETE: continue num_new_tokens = new_tokens_lens_list[req.py_seq_slot] + if os.environ.get("TLLM_DSPARK_RIDE_LOG") == "1" and ( + ridden_verify_lens is not None): + _win = ridden_verify_lens[req.py_seq_slot] + if _win > 0 and num_new_tokens > _win: + logger.warning( + f"[accept-bound] slot={req.py_seq_slot} accepted+1=" + f"{num_new_tokens} > token window {_win} -- in-graph " + f"acceptance ran past this row's layout window") + elif _win == 0: + logger.warning( + f"[accept-bound] slot={req.py_seq_slot} ride window=0 " + f"on a live request -- rewind falling back to the " + f"S snapshot (py_verify_len=" + f"{getattr(req, 'py_verify_len', None)})") # cap-accept: `num_new_tokens` has ALREADY been capped, on device, # inside acceptance (`interface.apply_accept_caps`) -- it has to be, # because the drafter reads the same count later in that forward to @@ -440,6 +456,13 @@ def sample_async( num_sampling_requests] self.store.verify_lens.index_copy_( 0, slots, o_verify_lens.to(self.store.verify_lens.dtype)) + if os.environ.get("TLLM_DSPARK_RIDE_LOG") == "1": + self._ride_log_ctr = getattr(self, "_ride_log_ctr", 0) + 1 + if self._ride_log_ctr % 64 == 0: + logger.info( + f"[ride-log] sampler #{self._ride_log_ctr} " + f"lens_buf[:4]={o_verify_lens[:4].tolist()} " + f"slots[:4]={slots[:4].tolist()}") # Create sample state with async D2H copy device_tensors = SampleStateTensorsSpec( diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py index f318d5f46300..2a973ebd5e5c 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py @@ -309,3 +309,47 @@ def test_select_windows_device_matches_host_chain(): host.verify_lens, graph_num_tokens=bucket) assert torch.equal(result.req_idx, want_req) assert torch.equal(result.kv_correction, want_corr) + + +def test_prologue_kv_pairing_matches_host_staging(): + """The host stages ``kv_lens = num_cached + seq_lens_kv`` where BOTH terms + bake the per-request token window (num_cached = past + tokens, + seq_lens_kv = tokens): kv_lens = past + 2*S. The overlap graph then adds + ``previous_kv_lens_offsets`` (staged as ``new_tokens_lens - S``) during + replay. Rebasing the layout onto the true windows w therefore needs + kv_lens += 2*(w - S) and offsets -= (w - S); any other split leaves the + indexer's slot mapping (derived from kv_lens) pointing at the wrong + K-cache slots -- the silent KV corruption that inflated the device-window + throughput while degrading every completion. The 2x coefficient and the + reference values were established by an A/B tensor diff against a full + host restage with identical windows. + """ + gen = torch.Generator().manual_seed(20260807) + for _ in range(20): + padded_bs = int(torch.randint(2, 40, (1,), generator=gen)) + past_seen = torch.randint(10, 500, (padded_bs,), generator=gen) + new_tokens_lens = torch.randint(1, 7, (padded_bs,), generator=gen) + shape_lens = torch.randint(2, 7, (padded_bs,), generator=gen) + true_lens = torch.randint(2, 7, (padded_bs,), generator=gen) + + # Host staging under the shape split S: the window counts twice. + kv_lens = (past_seen + 2 * shape_lens).to(torch.int32) + offsets = (new_tokens_lens - shape_lens).to(torch.int32) + + # Prologue: 2x on the kv side, 1x on the offsets side. + window_delta = (true_lens - shape_lens).to(torch.int32) + kv_lens += 2 * window_delta + offsets -= window_delta + + # Pre-offset kv_lens must equal host-with-w staging: the indexer's + # slot mapping and the expanded per-token extents derive from it. + assert torch.equal(kv_lens, + (past_seen + 2 * true_lens).to(torch.int32)) + + # In-graph correction at replay. + effective = kv_lens + offsets + + # Reference: host staging directly with the true windows w. + want = (past_seen + 2 * true_lens).to(torch.int32) + ( + new_tokens_lens - true_lens).to(torch.int32) + assert torch.equal(effective, want) From 621c82c5315dfaba709e49d4b913d87bb78bab19 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sat, 8 Aug 2026 06:37:10 -0700 Subject: [PATCH 108/177] [None][doc] Device-window corruption investigation log (living doc) Full evidence chain for the two device-window bugs: the fixed kv_lens 2x pairing, the still-open steady-state corruption, the 20+ controlled discriminator experiments (with the single-probe-vs-census methodology correction), the reviewer briefing, the debug-switch inventory, and the minimal reproduction recipe. Working notes for joint review; not intended to merge -- it is removed together with the debug harness once the remaining bug is fixed. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 306 +++++++++++++++++++++++++++++++++++++ 1 file changed, 306 insertions(+) create mode 100644 docs/dspark_sts_program.md diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md new file mode 100644 index 000000000000..f8f371ffba11 --- /dev/null +++ b/docs/dspark_sts_program.md @@ -0,0 +1,306 @@ +# DSpark STS 校准实验程序(任务 #30) + +> 活文档:计划、进度、结果、分析都记在这里,随实验推进更新。 +> 分支 `dspark-p25-device-window`;负责段:2026-08-07 起。 + +## ⚠️ REVIEWER 导读(2026-08-08 07:00 快照,供另一 agent 协查) + +**现状一句话**:device 端窗口选择(#25,env `TLLM_DSPARK_DEVICE_WINDOWS=1` 门控,默认关)会静默腐蚀生成;两个 bug 已定位一个、修复一个;第二个已围至"序幕 × w≠S × 稳态多请求工况",正在用 at-scale 装配差分收网。**host 路径(PR 主线)经受住了所有压力测试,无恙。** + +**已定案**: +1. **Bug #1(已修,勿回退)**:`_apply_device_window_prologue` 对 kv_lens 的差量应为 `2×(w−S)`(host 装配公式 `kv_lens = num_cached + seq_lens_kv` 两项都含本步窗口),原单倍差量使 indexer K-cache 写槽(slot_mapping_fp8)错位。经进程内 A/B 张量差分逐位验证;单测 `test_prologue_kv_pairing_matches_host_staging` 锁定公式。 +2. **旧实验结论修正**:此前 device "+63.9%" 吞吐是腐蚀假象(垃圾 token 易接受 + 零 EOS 满批不排水);host 动态裁剪四轮复现 ≈ 0±3%(校准+修表后)。 + +**未定案(协查重点)**: +- Bug #2:序幕激活(w≠S)+ 多并发 → 逐请求 census 63-64/64 全腐蚀,**与温度/采样图变体无关**(单探针曾误导出变体理论,已被 census 推翻)。prepare 级 21 张量 A/B 在 ramp 期小批步上全等;稳态步未采样(at-scale A/B 进行中)。 +- 已排除(每项有受控实验):重写机制(恒等/hybrid)、共享 ragged 消费方(host+steep 净)、sampler ride(仪表直证)、接受越窗(不变量仪表 0 违规)、V2 rewind(窗口不变代数)、图内采样计算(force-argmax)、逐 token 采样参数 H2D(skip 实验)、双趟捕获分离(交错捕获)、eager(该模式本就不支持序幕,host 标量按 S 直喂 kernel)。 +- **建议协查视角**:① 序幕在 padded_bs>n_real / 大 bucket / 批组成变化步的数学(`_apply_device_window_prologue` 全部下标与边界);② 序幕写入与图内消费方的契约面(dsa.py `apply_device_ragged_layout`/`on_update_kv_lens` 链);③ 跨步状态链(store→next staging 的张量时序)。 +- 完整实验叙事(20+ 受控判别)见下方进度日志;**census 是唯一可信文本判据**(单请求探针会被幸存请求误导)。 + +**本分支的调试支架(全部 env 门控、默认关闭、定案后剥离,勿 review 为正式代码)**: +`TLLM_DSPARK_DEVICE_IDENTITY`(1/steep,序幕恒等/极端窗口)、`TLLM_DSPARK_HOST_STEEP`(host 发布非均匀窗口)、`TLLM_DSPARK_FORCE_EAGER`+`EAGER_OK`、`TLLM_DSPARK_AB_DUMP`+`AB_AFTER`(装配 A/B 差分)、`TLLM_DSPARK_RIDE_LOG`(ride/accept-bound 仪表)、`TLLM_DSPARK_FORCE_ARGMAX`、`TLLM_DSPARK_SKIP_SAMPLING_H2D`、`TLLM_DSPARK_SKIP_GREEDY_CAPTURE`/`SKIP_ADVANCED_CAPTURE`/`INTERLEAVED_CAPTURE`(捕获趟操控)。涉及文件:model_engine.py、spec_sampler_base.py、interface.py。 +**复现最小配方**:DEP8 serve + `p25_round_arm.yaml`(sts_v1+sps_frac_v2b+tiers[2,3,4,5])+ `DEVICE_WINDOWS=1 DEVICE_IDENTITY=steep` + 64 并发固定 prompt(`dspark-runs/percreq_payload.json`)→ census 全腐蚀;去掉 DEVICE_WINDOWS 或改 HOST_STEEP → 全净。 + +## 总目标 + +1. 采集并拟合 STS 校准(此前所有实验都是裸 sigmoid,未校准); +2. 在校准 + 修复表的条件下,验证细刀(tiers [2,3,4,5])下 host/device 对 notrim 的收益在 **bs 64/128 per-rank(agg DEP8)** 是否**跨节点多轮稳定复现**; +3. 若复现:**搞清楚 device 为什么高于 host**; +4. draft-len 维度:draft=5 的三模式对比,以及 "draft5+ragged vs 最优固定短 draft(3/4/5)+满块" 的零假设检验。draft 6 已砍(checkpoint 无第 6 块权重,用户确认不看)。 + +## 背景:为什么要做这一轮(初测结果及其瑕疵) + +2026-08-07 凌晨,细刀三臂初测(poetry temp0.7 out512,v2 表): + +| bs(per-rank) | notrim | host | device | +|---|---|---|---| +| 512(64) | 11269 | 11319 (+0.4%) | 14537 (+29.0%) | +| 1024(128) | 11300 | 12236 (+8.3%) | **18524 (+63.9%)** | + +调度统计:notrim trim 0% accept 1.162;host trim 33.7% accept 1.077;device trim 49.7% accept_stat 2.032(可疑)。 + +三个瑕疵(本轮逐一修正): +1. 无 STS 校准(裸 sigmoid); +2. v2 表拟合被 `compress_to_risers` 截断到 8 节点:θ(512) 实测 33.9ms 被插值低估 ~9%,384→480 的真实陡坡被抹平; +3. device 段与 notrim/host 段跑在不同节点(晚绑定切换),跨节点混杂未消。 + +同期对照(粗刀 [1,3,5],同负载):host −8.6%~+3.8%、device ≈0 —— **tier 粒度是收益解锁器**(#19 论点的实验证据)。 + +## 实验步骤 + +### Step 1:STS 采集与拟合 【✅ 完成 10:18】 +- **结果**:`sts_v1.json`,温度 = [2.512, 1.585, 1.122, 1.995, 1.413];ECE 前 [0.130, 0.145, 0.115, 0.083, 0.061] → 后 [0.074, 0.081, 0.064, 0.047, 0.030],**5/5 位置改善,约减半**; +- **解读**:全部温度 >1 —— 裸 sigmoid 在所有 draft 位置**过度自信**(位置 0 最重,T=2.51)。意味着此前未校准实验的存活概率被系统性高估 → 裁剪偏保守的方向性偏差;位置间非单调(位置 3 T≈2.0)属 per-position 自由拟合的正常形态。 +- static 模式满块 serve(采集守卫强制无审查标签:无成本表、无裁剪),DEP8; +- 负载:poetry + arena 混合,temp 0.7,bs 512 global,各 2 遍; +- 拟合:`dspark_fit_sts.py`(与 SGLang 同方法:41 点对数网格 / 左到右冻结 / cumprod-ECE / 15 bins)→ `sts_v1.json`; +- 与 SGLang 的差异仅在配对管道(overlap 边界 → draft_seq_ring 双轴配对);拟合结果可互换。 +- 注:SGLang blog 实验的 `speculative_dspark_confidence_sts_path=None` —— 对照系统同样未校准,我们此前的未校准实验与 blog 同口径。 + +### Step 2:成本表修复 【✅ 完成】 +- 同样本重拟合,`--max-breakpoints 16` → `sps_frac_v2b.json`(16 节点); +- 不需要重新采集(截断发生在拟合层,原始样本完整)。 + +### Step 3:复现性轮次(核心)【等 STS 就绪自动开跑】 +- **4 节点并行,每节点一轮独立实验**(不做同节点编排,跨节点复现本身就是判据——用户指示); +- **每轮 3 次冷启动 serve**(notrim → host → device)——每臂重启,杜绝 reuse/分配器/autotuner 遗留污染(用户指示); +- 每臂:**poetry 与 arena 各做一次完整网格**(用户指示):每数据集 bs {512, 1024} global(64/128 per-rank)× 5 遍 × output 768 —— 每臂测量 ~17 分钟,每轮 ~2 小时; +- 配置:`sts_v1.json` + `sps_frac_v2b.json` + tiers [2,3,4,5],DEP8,max_batch_size 128; +- 三臂环境变量:notrim=`TLLM_DSPARK_FORCE_BUDGET_FRAC=1.0`;host=无;device=`TLLM_DSPARK_DEVICE_WINDOWS=1`; +- **判决**:4 轮的 (host−notrim)/notrim 与 (device−notrim)/notrim,方向一致且轮间散布 ≪ 效应 → 复现成立。 + +### Step 4:根因分析(device 为何 > host)【门禁:Step 3 复现】 +- **H3**(优先):accept 统计审计——device accept_stat 2.03 ≈ 其平均窗口("每请求接满窗"特征)。三种解释与判别器: + - (a) **真机制**:device 按当步置信度排序,留下的正是会被接受的位置 → accept≈窗口本就是好调度器的形态(设计意图); + - (b) **统计口径缺陷**:ride 记账(token 窗 −1)与 drafted 口径错位 → 判别:算术闭环(门禁 3)——client 侧 tokens/steps 独立推出 1+accept,与调度器 ride 值对表;吞吐收益若与 steps 降幅成比例,则统计为真; + - (c) **接受判定损坏**(接了不该接的)→ 判别:**temp-0 逐 token 等价**——同 prompt 下 device 臂 vs static 满块臂输出必须逐字一致(投机解码无损性);不一致即判定损坏。GSM8K 已过但 bs 小,需在 bs128 复核; + - 判别顺序:先 (b)(轮次数据自带)→ 再 (c)(等价试验)→ 都排除则接受 (a),交由 H2 分离"预算更狠"与"排序更新鲜"。 +- **H2**(判决性):裁深匹配差分——host + FORCE_BUDGET_FRAC≈device 实际裁深(~0.5)vs device,拆开"预算更狠"与"排序更新鲜"; +- **H4**:槽位新鲜度——长输出(低流转)下差距是否收窄。 + +### Step 5:draft-len 维度 【门禁:Step 3 复现】 +- 5a:draft=5,三模式对比:ragged-compact / cap-accept / static 满块,@ bs 64/128; +- 5b:零假设检验——**ragged@draft5 vs 最优 static@draft{3,4,5}**:动态裁剪 ≈ 自适应短草稿,若 static-3 就能拿走大部分收益,调度机器的增量价值存疑; + - 前置:放宽 worker 的 `block_size == max_draft_len` 断言(截断用前 k 块)+ GSM8K@draft3 正确性验证; + - draft 6:已砍。 + +## 挂起项(本程序期间不动) +- #29 disagg 512/rank 解锁(gen 侧 V2+NIXL CUDA IMA 未解); +- PR 推送(6 commit 就绪,等用户侧 git 认证通道); +- #11/#19/#20 等。 + +--- + +# 进度日志 + +- **08-07 09:04** STS 采集首发:配置被拒(`enable_confidence_scheduling` 要求 `enable_ragged_verify=true`)。 +- **08-07 09:14** 修复(yaml ragged=true + static 环境变量)重发,735105 采集中。 +- **08-07 09:20** v2b 表重拟合完成:16 节点;θ(512)=33.9ms(v2 插值低估 9%);384→480 实测陡坡恢复(23.9→31.7ms)。 +- **08-07 09:25** 一次未经确认的抢跑被用户叫停,轮次与节点已回收;计划经 review 修订(每臂重启、reps=5/out=768、加 arena 对照、砍 draft6)后获准。 +- **08-07 09:30** 4 轮次节点排队(735828–31),派发器就绪;轮脚本带 STS+v2b 门禁,采集完成后自动开跑。 +- **08-07 09:40** 用户升级 arena 为完整实验(原为单格对照);轮脚本更新为 poetry+arena 双全网格,重派。 +- **08-07 09:16→09:42** STS 采集第 2 次发射死于配置连环套(24 分钟后才发现——盯梢失职,已换 120s 紧密监视 + serve 日志错误早检):`enable_confidence_scheduling` 要求 ragged=true,ragged=true 要求 cost table 或 sts path,而采集守卫拒绝 cost table。**解法:恒等 STS 垫门**(`sts_identity.json`,全 1.0 温度 = 裸 sigmoid,零偏置)满足校验;第 3 次发射 09:42 起。节点门禁同步放宽(≥1h,采集只需 ~45 分钟)。 +- **08-07 09:46** 第 3 次发射死于 **GPU 显存不足**:735105 是此前 disagg gen IMA 事故节点,`kill -9` 后 CUDA 上下文僵死(D-state),274GB/卡 无法回收——**IMA 事故节点即废节点**,只能整作业回收。教训:晚绑定节点选择两度肇事(交叉误杀 + 复用废节点),STS 战役改为**显式钉死节点**。735105 已 scancel。 +- **08-07 09:51** 第 4 次发射:钉死 735104(disagg ctx 节点,验证 0 进程 0 显存,3:58 余量),120s 紧监视 + serve 错误早检在岗。 +- **08-07 10:18** STS 采集+拟合完成(serve 就绪 10:16,4 段混合负载 ~2.5 分钟,168 rank 文件,拟合 exit=0);`sts_v1.json` 落盘,轮次门禁 10:18:58 放行,round1(735830)/round2(735831) notrim 臂开跑。round3/4 等 735828/829 出队后自动派发。 + +## 数据合理性门禁(每轮数据落地即检查;异常 → 当场 debug → 记录在下方 debug 日志) + +1. **调度自洽**:notrim 臂 trim 必须 ≈0;各臂 delivered/ceiling 与 trim_ratio 一致;fallback/graph_eager 计数不异常膨胀; +2. **接受合理性**:accept_len ≤ 平均窗口;**accept ≈ 满窗是饱和红旗**(现存疑点:初测 device accept_stat 2.03 ≈ 其平均窗,即 H3);accept 与 trim 的方向关系合理(裁得多不应接得更多); +3. **算术闭环**:吞吐 ↔ steps_total ↔ (1+accept)×并发/步时 三者量级互恰;completion_tokens 接近 bs×max_tokens(poetry EOS 折损 ~7% 内); +4. **统计卫生**:rep 间离群(>2σ)标记并查因;首 rep 爬坡污染;轮间方向翻转即触发 debug 而非平均掉; +5. **臂间守恒**:三臂 completion 总量应近似(同负载同上限),显著偏差 = 生成行为被改变的信号(优先查 rewind/接受判定,参考 08-07 抓到的 cap-accept capture bug 的教训)。 + +任何一条触发:暂停下结论 → 定位(优先顺序:统计口径 bug > 测量方法 bug > 真实机制)→ 结论与证据写入下方 debug 日志,修复后视影响面决定重跑范围。 + +# 实验结果 + +## Step 3 轮次数据(吞吐 tok/s,mean 去 rep0 爬坡,n=5) + +### notrim 基线(四轮全过门禁 1:trim 全 0.0;accept 1.78/1.73/1.77/1.74 轮间高度一致) + +| 轮(节点) | poetry bs512 | poetry bs1024 | arena bs512 | arena bs1024 | +|---|---|---|---|---| +| round1 (735830) | 11505 | 12520 | 16409 | 13761 | +| round2 (735831) | 11430 | 12128 | 15440 | 13751 | +| round3 (735828) | 11591 | 12136 | 15142 | 13626 | +| round4 (735829) | 11198 | 12167 | 16246 | 13755 | +| **轮间散布** | ±1.7% | ±1.6% | ±4.2% | ±0.5% | + +注:arena bs512 测量窗最短(~24s),散布最大;作对比分母时以同轮同臂配对为准,不跨轮混用。 + +### host 臂(完整,四轮全过门禁:trim 22.4-22.7%、accept 1.66-1.67 逐位一致) + +| 轮 | poetry bs512 | poetry bs1024 | arena bs512 | arena bs1024 | +|---|---|---|---|---| +| round1 | +3.2% | −0.1% | +1.0% | −2.4% | +| round2 | +0.7% | +1.9% | −1.0% | −0.1% | +| round3 | −0.7% | +1.3% | +1.6% | −1.4% | +| round4 | +5.7% | +2.2% | −3.7% | −5.6% | +| **均值** | **+2.2%** | **+1.3%** | **−0.5%** | **−2.4%** | + +### device 臂(数据作废——生成损坏,见 Debug 日志;吞吐 +50%~+55% 为假象) + +device 臂指纹(各轮一致):completion 恰好=bs×768(零 EOS)、trim ≈49.8%、accept ≈2.06 ≈ 平均窗口(饱和)。 + +# Debug 日志 + +## 【08-07 12:30】device 臂生成损坏实锤(门禁 5 触发 → 活体探针定罪) + +**触发**:round1 device poetry bs512 的 completion_tokens 连续多 rep **恰好 = 393,216 = 512×768(100.0% 上限,零 EOS)**,而同轮 notrim/host 均 ~96%(自然 EOS 折损)。无损投机解码不应改变文本分布——只有 device 臂 EOS 行为漂移。 + +**探针 1**(temp-0 事实问答,round1 device serve):`' Paris"\n '`,正常终止——排除全面乱码。 + +**探针 2(定罪)**(真 poetry prompt,temp 0.7,max_tokens 768,与批负载并行):`finish_reason: length`,**768 token 生成后全文只有 2 字符 `S 读越界/垃圾 KV,w ride 窗口` = 图内接受越过本行布局窗口的现行证据;`ride 窗口=0` = ride 缺失导致 rewind 回退 S 快照的现行证据。738957 已带此仪表重启,steep+graph,收割中。 + +## 【08-07 23:55】accept-bound 收割:零越窗违规;单请求净;转 bs 扫描找最小复现 + +- **零次"接受越窗"**:整段脏负载中,图内接受从未超过本行布局窗口——接受计数与布局一致,排除"越窗读写"类; +- **ride 零值波全部出现在混合/过渡步**,那里回退 S 恰好等于该步布局(无序幕),行为正确; +- **单请求(batch=1)探针完美干净**(987 字符连贯诗歌 temp0)——但注意 batch=1 时 steep 转移无操作(w≡S),判别力≈恒等; +- **推论收窄**:接受计数对、装配张量对、越窗不存在 → 腐蚀在"接受的 token 内容"或"跨请求 KV/状态污染",且需要多请求才触发; +- **过夜实验(738957)**:bs 扫描 {16,32,64,128,256,512} × temp0 × 每档文本探针——找腐蚀出现的最小批量,作为明天内核级挖掘的最小复现。 + +## 【08-08 00:55】当夜最终收口:残余 bug 锁定"advanced-sampling 图变体 × w≠S" + +- **bs 扫描(temp0)判决:2×修复后 greedy 路径全净**——bs 16/32/64/256/512 探针全部连贯(~4 chars/tok;128 的 2.4 是探针踩中 ramp 的噪声)。greedy 完整走 verify+draft+commit 管线 → **布局/KV/接受/提交机器全部无罪(修复生效)**。 +- **同一台 serve、同一份代码:temp0.7 负载 99.3% 撞顶(脏)**——残余腐蚀是**采样路径特异**的。temp0 replay greedy 图变体;temp0.7 replay **advanced-sampling 图变体**(独立捕获趟)——嫌疑收窄到该变体的捕获/绑定与 w≠S 的交互。 +- **ride 陈旧被证明良性**:ride=0+陈旧 py_verify_len 的波在两种温度下都大量出现(温度不相关),且 V2 管理器的两项 rewind(`py_rewind + max(reserved−runtime,0)`)总和 = `5−acc`,**与窗口无关**——rewind 正确性不依赖 ride。ride 只喂统计。 +- **明晨第一刀(廉价决定性)**:temp0.7 + w≠S + **禁用 advanced-sampling 捕获**(非 greedy 批将按注释回退 eager)。净 → bug 在 advanced 变体的捕获路径(对比两趟捕获的差异:`_force_non_greedy_for_capture`→populate→`_sample_tokens_for_batch` 非贪婪分支的捕获期绑定,重点:`spec_metadata.temperatures` 的对象身份/持久性、flashinfer `sampling_batch_spec_dec_one_model` 的 torch.compile 捕获);脏 → eager×0.7×w≠S 复现,范围另定。 +- 当夜战果:bug#1(kv_lens 2×)修复+验证;greedy 全尺寸验证干净;残余范围从"整个 device 路径"收窄到"一个图变体的捕获交互"。 + +## 【08-08 01:25】Bug #3 干净定罪:advanced-sampling 图变体 × device 窗口 + +- **禁 advanced 捕获 + temp0.7 = 干净**——但混杂(非贪婪回退 eager → 序幕不跑,等价 host 路径),仅作旁证; +- **去混杂终验(738957)**:`temperature=0.0001`(分布≈argmax;temp0 同负载 = 完美诗歌)→ **2 字符垃圾**。同 serve 同分布,唯一变量 = 路由到 advanced-sampling 图变体。**定罪成立**:greedy ragged 图 + 序幕 = 净;advanced ragged 图 + 序幕 = 毁;advanced ragged 图 + host 装配(host+steep@0.7)= 净。 +- **三元交互**:腐蚀 = advanced 图变体 × 序幕改写(w≠S)× 多请求。变体间的捕获差异只有 `_sample_tokens_for_batch` 非贪婪分支(flashinfer `sampling_batch_spec_dec_one_model`,torch.compiled)+ is_all_greedy 条件代码——明日在此范围内定位(候选:per-token 采样参数缓冲的捕获期对象绑定/装配路径、compiled kernel 的捕获期特化、seed/offset 语义;注意 host 装配可走通 → 嫌疑聚焦在"该分支读取了某个序幕没有重写、且按 S 装配的逐 token 缓冲——温度同质时值相等的推理可能有漏洞,如 dummy/probe 行的哨兵值在 w 布局下错位到真实行")。 +- **明日计划**:① 在 advanced 分支加 replay 后 D2H dump(sampled vs argmax @1e-4;temps/top_ks 实际值),一步定位;② 修复;③ graph steep 0.7 探针转净 → 真 device 吞吐网格 + GSM8K 大并发 → 剥支架、整理提交。 + +## 【08-08 01:30】变体内二分:捕获的算子无罪 → 嫌疑收缩到双趟捕获机制本身 + +- **force-argmax(捕获前置生效)+ advanced 变体 + 0.7 = 仍脏**(0.17 chars/tok)。此时 advanced 图的算子内容与 greedy 图完全一致(target argmax 已强制;draft 本就 argmax——`wants_advanced_draft_sampling` 需要 rejection,而 `use_rejection_sampling=False`);逐 token 采样参数缓冲同质中和。 +- **推论**:greedy 图净 / 计算等价的 advanced 图脏 ⇒ 差异在**双趟捕获机制**:图池共享/别名、第二趟捕获对共享持久状态的副作用、runner 对两变体的 key/元数据簿记——与序幕(w≠S)交互的面。 +- **明晨作战序列**:① 读 runner 的双变体捕获簿记(graph pool、shared_static_tensors、graph_metadata/outputs per key)找 w≠S 敏感的共享面;② 对照实验:只捕获 advanced 趟(跳过 greedy 趟)+ temp0.7 → 若净,则是"两趟共存"的交互;仍脏则是 advanced 趟自身的捕获环境(force_non_greedy 时的 populate/warmup 状态);③ 定位后修复 → 验证三件套 → 吞吐 + GSM8K 大并发 → 剥支架。 +- 注:所有调试开关(IDENTITY/FORCE_EAGER/EAGER_OK/AB_DUMP/RIDE_LOG/HOST_STEEP/SKIP_ADVANCED_CAPTURE/FORCE_ARGMAX/SKIP_GREEDY_CAPTURE)+ 仪表代码在修复定案后一并剥离;`kv_lens 2×` 修复与单测保留。 + +## 【08-08 02:00】当夜终判:双趟捕获共存被点名 + +- **只捕 advanced 趟(跳过 greedy 趟)+ steep + 0.7 = 大幅转好**(2.0 chars/tok,连贯语料延续;双趟时 0.01-0.17 垃圾)——**腐蚀 = 双趟捕获共存 × 序幕(w≠S)**。单独任一趟的图都(基本)健康;两趟共存时 temp0.7(pass-2 图)被毁。 +- 机制候选(明晨深挖):共享图池/attention 工作区在两趟重复捕获同形状时的指针/别名交互(代码内 dynamic-draft-len 的已知同类患:"resize 使先捕图的指针失效");pass-1 捕获对持久状态的遗留改变 pass-2 捕获期的 Python 条件求值;worker 持久状态(_ctx_len/_kv_windows/confidence)跨趟的捕获期录制差异。 +- 明晨序列:① adv-only 严格验证(完整网格 + completion 指纹,2.0 的余量需排除内容密度混杂);② runner 双趟簿记深读(shared pool、workspace resize、graph_metadata);③ 修复(候选:ragged bucket 仅在所需变体捕获 / 趟间池隔离 / max-shape 前置捕获同款 workaround);④ 验证三件套 → 吞吐 + GSM8K 大并发 → 剥支架。 +- morning 节点已排队:739289/739290(4h × 2)。 + +## 【08-08 02:30-04:45】参照指纹 + 逐请求普查:adv-only "转好" 是抽样错觉 + +- **干净参照线确立(host+steep@0.7 同网格)**:completion **96.6-96.7%**、accept 1.06、吞吐 9.9-12.1K——与 notrim 的 96% 一致。 +- **adv-only device 对照**:completion 99.0-99.8%、accept 2.14、吞吐 12.8-16.2K——**仍系统性偏移,腐蚀未消**。 +- **逐请求普查仪表(64 并发独立请求、逐个存文本)**:adv-only device @0.7 = **64/64 全腐蚀**(chars/tok 中位 ~0.03)。"单趟转好"“探针 3.67”均为抽样错觉(探针骑行时活跃请求少)。**双趟共存理论降级为噪声;回归主线:advanced 图变体(temp>0)× 序幕 = 全毁,greedy 图 = 净。** +- force-argmax 普查(判"eager 非贪婪主机侧代码 vs 图内 advanced 采样计算")因节点到期未完成;739900/739901 已排队,落地重跑。 +- **普查仪表沉淀**:`percreq_payload.json` + 64 并发 curl + chars/tok 分类——今后所有验证以此为准(单探针不再作为判据)。 + +# 分析 + +## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) + +1. **host 未复现**:初测 +8.3%(poetry 128/rank)在四轮中缩为 [0.0, +1.9, +1.3, +2.2]%;arena 侧 [−2.4, 0.0, −1.4, −5.6]%。host 整体 ≈ 0±3%,poetry 微正、arena 微负。host 调度行为本身跨节点完全可复现(trim 22.4-22.7%,accept 1.66-1.67,四轮逐位一致)——**是收益不存在,不是实验不稳**。 +2. **device "+63.9%" 判定为假象**:device 臂生成损坏(见 Debug 日志 12:30 条),吞吐虚高由重复特殊 token + 零 EOS + 满批不排水制造。本轮 device 吞吐(+50%~+55%)全部作废。**H3 判 (c):接受判定/停机记账损坏**,旧 accept_stat 2.03 即其指纹。 +3. 待办:device 路径根因定位与修复(736879/736880 debug 节点已排队);修复后 device 收益需重测才有结论。H2/H4 挂起,等 device 修复。 +4. 联动结论:#28(SGLang 原生动态在 B300 高 bs 亏损)+ 本轮 host≈0 → 细刀 tier + 校准 + 修表后,**host 动态裁剪在 B300 agg 上没有可复现的吞吐收益**;与 SGLang blog 的 +10-20%(H200/H100)相比,B300 的 θ 平坦区更宽,裁剪节省的 token 成本换不回接受长度损失。 From e2acaf24b06c004cf67d83e542581f61ebbd5a9e Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sat, 8 Aug 2026 07:03:19 -0700 Subject: [PATCH 109/177] [None][chore] Read-only invariant audit for the device-window prologue The restage-based staging differ crashes the executor under real load (request-state mutation breaks response bookkeeping and, on DEP8, rank divergence hangs collectives), so at-scale verification moves to read-only invariant checks: D2H a few derived values after the prologue and compare against the analytic host-with-w formulas (first invariant: kv_lens == past + 2*w per request). Env-gated debug-only (TLLM_DSPARK_INVARIANT_LOG); documented in the investigation log. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 6 ++++ .../_torch/pyexecutor/model_engine.py | 32 +++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index f8f371ffba11..4e7d48f4c53a 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -296,6 +296,12 @@ device 臂指纹(各轮一致):completion 恰好=bs×768(零 EOS)、trim ≈49.8 - force-argmax 普查(判"eager 非贪婪主机侧代码 vs 图内 advanced 采样计算")因节点到期未完成;739900/739901 已排队,落地重跑。 - **普查仪表沉淀**:`percreq_payload.json` + 64 并发 curl + chars/tok 分类——今后所有验证以此为准(单探针不再作为判据)。 +## 【08-08 06:45-07:45】重装配差分在规模下判死;转只读不变量审计 + +- **at-scale A/B(AB_AFTER=150)**:跨过阈值后开火即崩——executor 死于与猎手 v1 相同的响应管道错误(`'int' object is not iterable`)。**结论:重装配式差分只能在安静的 ramp 步上活着,规模下对请求态的突变不可控,该路线放弃**(其 ramp 步"21 张量全等"的判决仍有效,但只覆盖小批步)。 +- **新仪表(只读不变量审计,`TLLM_DSPARK_INVARIANT_LOG=1`)**:序幕末尾 D2H 少量派生量,对照 host-with-w 的解析公式验证不变量(首个:`kv_lens[r] == (max_beam_num_tokens−1) + 2×w_r`),零突变、零集合通信、零重装配——任意规模可常开。违规打印现场 (row, past, w, S, n_real, padded_bs, bucket);每 100 步打检查点(正向确认仪表存活)。 +- 协查提示:kv 不变量若稳态全绿,后续只读不变量按序追加:anchor 对位(`input_ids[qo[r]] == next_new_tokens[slot,0]`)、draft 对位、position 基。序幕输出面逐项排掉后,嫌疑收缩到图内消费方与跨步状态链。 + # 分析 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 8f41e0beb12e..ff2a9e8e5b1e 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -4057,6 +4057,34 @@ def _apply_device_window_prologue(self, inputs, logger.info( f"[ride-log] prologue #{self._ride_log_ctr} " f"S={split_lens[:4].tolist()} w={result.verify_lens[:4].tolist()}") + if (os.environ.get("TLLM_DSPARK_INVARIANT_LOG") == "1" + and getattr(self, "_dspark_inv_prints", 0) < 40): + # Read-only invariant audit against host-with-w formulas: no + # restage, no request mutation, no collectives -- safe at scale. + torch.cuda.synchronize() + gen_requests = [ + r for r in (self._dspark_inv_requests or []) + ] if getattr(self, "_dspark_inv_requests", None) else [] + w = result.verify_lens[:n_real].tolist() + qo = result.qo_indptr[:n_real + 1].tolist() + kv = attn_metadata.kv_lens_cuda[:n_real].tolist() + self._dspark_inv_step = getattr(self, "_dspark_inv_step", 0) + 1 + for i, req in enumerate(gen_requests[:n_real]): + past = int(req.max_beam_num_tokens) - 1 + want_kv = past + 2 * w[i] + if kv[i] != want_kv: + self._dspark_inv_prints = getattr( + self, "_dspark_inv_prints", 0) + 1 + logger.warning( + f"[invariant] step#{self._dspark_inv_step} row {i}: " + f"kv={kv[i]} want past+2w={want_kv} " + f"(past={past} w={w[i]} S={split_lens[i].item()} " + f"n_real={n_real} padded_bs={padded_bs} " + f"bucket={bucket})") + if self._dspark_inv_step % 100 == 0: + logger.info( + f"[invariant] step#{self._dspark_inv_step} checkpoint: " + f"{getattr(self, '_dspark_inv_prints', 0)} violations") return True def _dspark_ab_diff(self, padded_requests, kv_cache_manager, attn_metadata, @@ -7309,6 +7337,10 @@ def forward(self, and not self.is_warmup and getattr(self, "_dspark_device_budget", None) is not None): + if os.environ.get("TLLM_DSPARK_INVARIANT_LOG") == "1": + # Read-only reference for the prologue's invariant audit. + self._dspark_inv_requests = list( + padded_requests.generation_requests) applied = self._apply_device_window_prologue( inputs, new_tensors_device) if applied and os.environ.get("TLLM_DSPARK_AB_DUMP") == "1": From 57e9d56e8c4efa2fbe8ba6dd5e8e2e40ad418cc3 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:18:02 -0700 Subject: [PATCH 110/177] [None][fix] Device windows: compressor gen-token counts must follow the device buffer The compressor's per-request new-token counts (gen_new_tokens_per_seq_cuda) were staged as a pinned-H2D copy inside on_update_kv_lens, which is captured into the generation CUDA graph. Every replay therefore re-loaded the host shape split S from the pinned buffer, while the device-window prologue had already corrected kv_lens (and every other consumer) to the true windows w. The kernel then computed its state start position sp = kv_len - nn with a w-corrected kv_len but nn = S, and Phase 1 wrote the persistent paged KV/score recurrent state shifted by (w - S) per request -- silently and cumulatively corrupting compressed attention history for both the main compress layers and the indexer FP8 K-cache compressor on any step where a request's w != S (which requires >= 2 concurrent requests, since window re-ranking conserves the bucket total). Fix, two parts: - Source the captured copy from _seq_lens_cuda (D2D). It holds the same S at staging time and is rewritten to w by apply_device_ragged_layout before replay, so the captured copy node follows the device truth for both host and device windows. The pinned twin buffer is deleted. - Bake the next_n kernel template bound at the GLOBAL max window (1 + max_draft_tokens) instead of this step's batch max: ragged capture warms up with a uniform tier-value split, and the prologue may assign windows up to the config max on any bucket; the kernel treats the bound as guarded, so widening it is correctness-neutral for existing paths. Uniform and eager branches are untouched. Unit tests cover the D2D source switch, the global bound, context-row offsets, and the uniform branch. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../sparse/deepseek_v4/deepseek_v4.py | 41 +++++----- .../test_dspark_compressor_gen_counts.py | 78 +++++++++++++++++++ 2 files changed, 101 insertions(+), 18 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_compressor_gen_counts.py diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py index b7b81df73320..65350c249485 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py @@ -344,9 +344,6 @@ def __post_init__(self): cache_name="gen_new_tokens_per_seq_cuda", capture_graph=capture_graph, ) - self.gen_new_tokens_per_seq_host = torch.empty_like( - self.gen_new_tokens_per_seq_cuda, device="cpu", pin_memory=prefer_pinned() - ) # Set per step by prepare(); None means "uniform, use next_n". self.gen_new_tokens_per_seq: Optional[torch.Tensor] = None @@ -693,7 +690,7 @@ def _prepare_impl(self): # Overrides DSAtrtllmAttentionMetadata._prepare_impl, NOT prepare(): # the base prepare() wraps this whole body in the write-after-read # guard for the pinned staging rewritten below (cached_token_lens_cpu, - # cu_seq_lens, gen_new_tokens_per_seq_host, and everything the called + # cu_seq_lens, and everything the called # helpers stage). Overriding prepare() directly would silently opt this # class out of that guard. assert self.kv_cache_manager is not None @@ -754,8 +751,9 @@ def _prepare_impl(self): # Under ragged verification requests append different numbers of tokens, # so the integer division below is meaningless. The compressor kernel # still needs a compile-time loop bound, and its loops are already - # guarded, so the batch MAXIMUM is the correct scalar to hand it -- the - # exact per-request counts ride along in gen_new_tokens_per_seq. + # guarded, so the GLOBAL max window is the scalar to hand it (the + # device-window prologue may re-rank past this step's batch max) -- + # the exact per-request counts ride along in gen_new_tokens_per_seq. num_gen_tokens_per_seq = self._sync_gen_tokens_per_seq(num_gen_tokens) num_contexts = self.num_contexts num_generations = self.num_generations @@ -816,26 +814,33 @@ def _sync_gen_tokens_per_seq(self, num_gen_tokens: int) -> int: Uniform batches divide the flat generation-token count by the request count. Ragged batches have no such quotient, so the scalar becomes the - batch *maximum* -- which is exactly what the compressor kernel wants, - since ``NEXT_N`` is a template loop bound whose iterations are already - guarded -- and the true per-request counts travel beside it on the - device. + GLOBAL max window -- ``NEXT_N`` is a template loop bound whose + iterations are already guarded, so any upper bound is correct, and it + must be the global one because the device-window prologue may re-rank + rows up to the config max AFTER this staging runs (this step's batch + max would truncate Phase 1/3 for such rows). The true per-request + counts travel beside it on the device. + + The vector is a D2D copy from the seq-lens staging, NOT an H2D from a + pinned host list: this helper is captured inside ``on_update_kv_lens``, + and a captured H2D node would re-load the host's shape split S from + the pinned buffer on every replay -- silently clobbering the device + windows w and making the compressor place its persistent paged + KV/score state at positions shifted by (w - S). ``_seq_lens_cuda`` + holds the same S at staging time and is rewritten to w by + ``apply_device_ragged_layout`` before replay, so sourcing from it is + exact for both host and device windows. Single writer on purpose: the scalar and the vector describe the same thing, and every place that recomputed one without the other was a way for them to disagree. """ if self.is_ragged_verify and self.num_generations > 0: - gen_token_counts = self.gen_token_repeat_list() - self.num_gen_tokens_per_seq = max(gen_token_counts) + self.num_gen_tokens_per_seq = 1 + self.max_draft_tokens n = self.num_generations - self.gen_new_tokens_per_seq_host[:n].copy_( - torch.tensor(gen_token_counts, dtype=torch.int32) - ) + nc = self.num_contexts gen_new_tokens = self.gen_new_tokens_per_seq_cuda[:n] - gen_new_tokens.copy_( - self.gen_new_tokens_per_seq_host[:n], non_blocking=True - ) + gen_new_tokens.copy_(self._seq_lens_cuda[nc:nc + n]) self.gen_new_tokens_per_seq = gen_new_tokens else: self.num_gen_tokens_per_seq = ( diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_compressor_gen_counts.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_compressor_gen_counts.py new file mode 100644 index 000000000000..7bf647f680db --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_compressor_gen_counts.py @@ -0,0 +1,78 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The compressor's per-request new-token counts must follow the DEVICE buffer. + +``_sync_gen_tokens_per_seq`` runs inside the captured ``on_update_kv_lens``, +so its writes are replayed every step. The device-window prologue re-ranks +verify windows from the host shape split S to true windows w by rewriting +``_seq_lens_cuda`` (via ``apply_device_ragged_layout``) after host staging +but before replay. Two properties keep the compressor correct under that +rewrite, and both were violated by the original implementation: + +1. The per-request count vector must be a D2D copy sourced from + ``_seq_lens_cuda`` -- a captured H2D copy from a pinned host list re-loads + S on every replay, so the kernel computes its state start position + ``sp = kv_len - nn`` with a w-corrected ``kv_len`` but ``nn = S``, and + writes the persistent paged KV/score state shifted by (w - S). +2. The ``next_n`` template scalar must be the GLOBAL max window, not this + step's batch max: the prologue may assign any row up to the config max on + any bucket, and the kernel treats the scalar as a guarded upper bound. +""" + +import types + +import torch + +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import \ + DeepseekV4TrtllmAttentionMetadata + + +def _dummy(nc, staged_split, seq_lens_dev, max_draft_tokens=5, cap=16): + d = types.SimpleNamespace() + d.is_ragged_verify = True + d.num_generations = len(staged_split) + d.num_contexts = nc + d.max_draft_tokens = max_draft_tokens + # What the host shape split says (S); the fixed code must NOT read it. + d.ragged_verify_lens = list(staged_split) + d.gen_new_tokens_per_seq_cuda = torch.zeros(cap, dtype=torch.int) + d._seq_lens_cuda = torch.tensor(seq_lens_dev, dtype=torch.int) + return d + + +def _call(d, num_gen_tokens): + return DeepseekV4TrtllmAttentionMetadata._sync_gen_tokens_per_seq( + d, num_gen_tokens) + + +def test_vector_sources_from_seq_lens_not_host_split(): + # Simulate the prologue: host staged S, device rewrote seq lens to w. + S = [3, 3, 3] + w = [2, 4, 3] + d = _dummy(nc=0, staged_split=S, seq_lens_dev=w) + _call(d, sum(w)) + assert d.gen_new_tokens_per_seq.tolist() == w + + +def test_scalar_is_global_bound_not_batch_max(): + S = [2, 2] + d = _dummy(nc=0, staged_split=S, seq_lens_dev=S) + next_n = _call(d, sum(S)) + assert next_n == 1 + d.max_draft_tokens + assert d.num_gen_tokens_per_seq == 1 + d.max_draft_tokens + + +def test_gen_rows_slice_skips_context_rows(): + ctx_lens = [17, 9] + w = [4, 2] + d = _dummy(nc=2, staged_split=w, seq_lens_dev=ctx_lens + w) + _call(d, sum(w)) + assert d.gen_new_tokens_per_seq.tolist() == w + + +def test_uniform_branch_unchanged(): + d = _dummy(nc=0, staged_split=[3, 3], seq_lens_dev=[3, 3]) + d.is_ragged_verify = False + next_n = _call(d, 6) + assert next_n == 3 + assert d.gen_new_tokens_per_seq is None From 20686d3cde08ff7eec591c7eb687c172cf9df3c4 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:18:18 -0700 Subject: [PATCH 111/177] [None][doc] Root-cause log for the device-window corruption + rowdiv detector Living-doc update: the captured pinned-H2D reload of the compressor's per-request gen-token counts is the second (and with the census evidence, last) device-window corruption mechanism. Records the converging verdicts from the read-only invariant audits, the 64-way temp0 row-divergence detector (identical requests diverging from step 3 with degenerate token runs -- the fingerprint of per-request persistent-state damage), and the parallel static audit of every window-length-dependent buffer consumed during graph replay. Includes the alternative-fix evaluation for review. Instrumentation kept alongside: the rowdiv detector in the spec sampler (TLLM_DSPARK_ROWDIV_LOG), the extended batch-3 invariants in the prologue audit (row maps analytic recompute, seq_lens/repeats/position packing), and the profiler tier knobs. All debug switches are enumerated in the doc and will be stripped before the branch is finalized. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 42 ++++++++++++ .../_torch/pyexecutor/model_engine.py | 65 ++++++++++++++++--- .../_torch/speculative/spec_sampler_base.py | 30 +++++++++ tests/microbenchmarks/dspark_sps_profiler.py | 15 ++++- 4 files changed, 140 insertions(+), 12 deletions(-) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 4e7d48f4c53a..849fc3458405 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -302,6 +302,48 @@ device 臂指纹(各轮一致):completion 恰好=bs×768(零 EOS)、trim ≈49.8 - **新仪表(只读不变量审计,`TLLM_DSPARK_INVARIANT_LOG=1`)**:序幕末尾 D2H 少量派生量,对照 host-with-w 的解析公式验证不变量(首个:`kv_lens[r] == (max_beam_num_tokens−1) + 2×w_r`),零突变、零集合通信、零重装配——任意规模可常开。违规打印现场 (row, past, w, S, n_real, padded_bs, bucket);每 100 步打检查点(正向确认仪表存活)。 - 协查提示:kv 不变量若稳态全绿,后续只读不变量按序追加:anchor 对位(`input_ids[qo[r]] == next_new_tokens[slot,0]`)、draft 对位、position 基。序幕输出面逐项排掉后,嫌疑收缩到图内消费方与跨步状态链。 +## 【08-08 07:35-08:10】kv 不变量稳态全绿;扩展不变量 × 双线并行(steep 猎场 + identity 对照) + +- **kv 不变量:200+ 稳态步 0 违规**,检查点机制同时正向确认仪表存活——序幕 kv 输出面(修复后)在规模下逐请求正确。 +- **第二批只读不变量上线**(steep,739901):`kvoff == new_lens − w`、anchor 对位(`input_ids[qo[r]] == next_new_tokens[0,slot]`)、draft0 对位、`prev_pos_idx/off` 首尾 token 对位——覆盖序幕的全部 token 级输出。 +- **仪表对照组**(identity,739900):同一套不变量在 w=S 下应全绿——排除仪表自身假阳性;若对照组报违规则仪表公式有错。 +- 若 steep 全绿:第三批备选不变量 = row maps 解析重算(req_idx/kv_correction 是 w 的纯函数)、`seq_lens_cuda == w`、`gen_token_repeats == w`、pad 行窗口 == pad_len;全绿穿透后嫌疑正式移交图内消费方与跨步链。 + +## 【08-08 08:15-08:40】第二批全绿 + 对照全绿;batch-3 在跑;跨行分歧检测器上线 + +- **steep 第二批不变量:300+ 步 0 违规**(kvoff/anchor/draft0/prev_pos 首尾对位)——序幕 token 级输出面全部无罪;**identity 对照同样全绿**——仪表无假阳性,判定有效。 +- **batch-3(740252 加载中)**:row maps 解析重算、seq_lens==w、repeats==w、position 打包(行内位差)。 +- **新杀器(`TLLM_DSPARK_ROWDIV_LOG=1`,下个 serve 周期生效)**:census 的 64 个请求完全相同 → temp0 下所有行每步接受的 token 必须一致(窗口只改变验证量、不改变内容)。sampler 内纯 host 检查:行间 token 元组分歧即腐蚀现行,直接给出**首个分歧步 + 分歧行的窗口类别(w=2 vs w=6)**——回答"哪类窗口先烂",并把腐蚀定位从"输出面"推进到"图内效果"层。 + +## 【08-08 09:15-09:25】rowdiv serve 重发 ×2;守恒论证解释单请求永净;图内消费方静态围剿启动 + +- **rowdiv serve(740252)warmup 挂死**:日志冻结于 09:15:37,进程活但 20+ 分钟无进展 → 杀掉,双节点并行重发(740252→p25rowdiv2_serve.log,740251→p25rowdiv3_serve.log,同配置 steep+ROWDIV_LOG),first-ready-wins 收割流水线挂上(先 ready 者压 64 并发 temp0 census → grep `[rowdiv]`)。若二次 warmup 挂死 → 实锤 rowdiv 仪表自身与 warmup 步的交互问题。 +- **结构性论证(新)**:窗口守恒 sum(w)=sum(S) 意味着 **n_real=1 时数学上强制 w=S** —— 单请求下序幕是恒等变换。这解释了单请求探针为什么永净:不是运气,是结构。任何 w≠S 敏感的凶手**只可能**在多请求下现形,与全部观测(多并发 63-64/64 腐蚀、单请求净)精确吻合。 +- **静态围剿(并行 workflow)**:4 个读者分扫 ① host staging 面(_prepare_tp_inputs + spec_metadata.prepare)② attention+indexer 面(apply_device_ragged_layout 覆盖清单 vs MLA/索引器 replay 时全部读取)③ 图内 worker 面(fused verify+draft 的每个 gather/scatter 索引、confidence 装填索引、捕获期烘焙常量)④ 跨步链(host 簿记以 S 推进而 device 以 w 运行之处)。汇总产出"w 敏感 ∧ replay 消费 ∧ 序幕未重写"嫌疑排名 + 每个的最廉价判决实验。已知热点候选:`spec_metadata.gather_ids`(interface.py:478,drafting_loops 以它选 logits 行——若 DSpark 路径同类索引按 S 装填且未被序幕重写,则 w≠S 时采样/验证读错行,跨请求串染,恰好匹配全部现象)。 + +## 【08-08 10:15-10:50】BUG #2 根因定案:compressor 的逐请求新 token 数被捕获的 H2D 钉死在 S + +- **双线判决同时落地,互相咬合。** + - **rowdiv(动态,steep+temp0+64 相同请求)**:step#3-4 起行间 token 元组即分歧,同窗口类 (w=6) 行内容也互不相同,大批行迅速塌缩为 `(0,0,0)` 退化 token——腐蚀即时发生、按请求永久化,是"逐请求持久状态被写坏"的指纹。 + - **静态围剿(4+1 workflow,60 个缓冲逐一分类)**:唯一同时满足全部观测(多请求 only / identity 净 / temp0+force-argmax 也腐蚀 / 与捕获结构无关 / 序幕输出面逐项正确)的嫌疑就是 **`gen_new_tokens_per_seq_cuda`**。 +- **机制(全链路代码实证)**:host staging 在 `_sync_gen_tokens_per_seq`(deepseek_v4.py)把 host 形状split S 写入 pinned 缓冲并 H2D 到 `gen_new_tokens_per_seq_cuda`;该调用同时在捕获区 `on_update_kv_lens` 内执行,于是 **H2D memcpy 成为图内节点,每次 replay 从 pinned 缓冲重新加载当步的 S**——序幕改了也会被 replay 覆盖(序幕本来也没改它:`apply_device_ragged_layout` 写的是双胞胎 `seq_lens_cuda`/`gen_token_repeats_cuda`)。compressor kernel 算 `nn = new_tokens_per_seq[b]`、`sp = kv_len − nn`,而 `kv_len` 已被序幕修到 w 口径 → **`sp = past + (w−S)`**:Phase 1 把**持久分页 KV/score 递归状态**写错位(w>S 尾巴永远缺失;w 1 and drafts_l[qo[i] - i] != ndraft[s][0]: + _viol(i, f"draft0 {drafts_l[qo[i]-i]} want {ndraft[s][0]}") + for o in (0, w[i] - 1): + if ppi[qo[i] + o] != s: + _viol(i, f"prev_pos_idx[{qo[i]+o}]={ppi[qo[i]+o]} " + f"want slot {s}") + if ppo[qo[i] + o] != nlens[s]: + _viol(i, f"prev_pos_off[{qo[i]+o}]={ppo[qo[i]+o]} " + f"want nlens={nlens[s]}") if self._dspark_inv_step % 100 == 0: logger.info( f"[invariant] step#{self._dspark_inv_step} checkpoint: " diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index bd1312a7d615..e45496ebb8b5 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -291,6 +291,36 @@ def update_requests( beam_idx = DEFAULT_BEAM_IDX runtime_draft_len = getattr(state, "runtime_draft_len", self.draft_len) + if os.environ.get("TLLM_DSPARK_ROWDIV_LOG") == "1": + # Identical-request divergence detector: under a census load all + # requests share one prompt, so at temp 0 every row's accepted + # tokens must match its peers each step (windows change how MUCH + # is verified, never WHAT is accepted). The first divergent step + # names the corrupted rows and their window class. + sigs = {} + for _req in state.requests: + if _req.state == LlmRequestState.GENERATION_COMPLETE: + continue + _s = _req.py_seq_slot + _n = new_tokens_lens_list[_s] + _toks = tuple(new_tokens[i][_s][0] for i in range(_n)) + _w = (ridden_verify_lens[_s] + if ridden_verify_lens is not None else -1) + sigs.setdefault(_toks, []).append((_s, _w)) + self._rowdiv_step = getattr(self, "_rowdiv_step", 0) + 1 + if len(sigs) > 1 and getattr(self, "_rowdiv_prints", 0) < 30: + self._rowdiv_prints = getattr(self, "_rowdiv_prints", 0) + 1 + desc = "; ".join( + f"{len(v)}rows(w={sorted(set(x[1] for x in v))})" + f"toks={k[:3]}" for k, v in list(sigs.items())[:4]) + logger.warning( + f"[rowdiv] step#{self._rowdiv_step} {len(sigs)} distinct " + f"token-tuples across identical requests: {desc}") + if self._rowdiv_step % 100 == 0: + logger.info(f"[rowdiv] step#{self._rowdiv_step} checkpoint: " + f"{getattr(self, '_rowdiv_prints', 0)} " + f"divergent steps") + caps = state.verify_caps_snapshot for req in state.requests: if req.state == LlmRequestState.GENERATION_COMPLETE: diff --git a/tests/microbenchmarks/dspark_sps_profiler.py b/tests/microbenchmarks/dspark_sps_profiler.py index 01bc428335e4..934986f21143 100644 --- a/tests/microbenchmarks/dspark_sps_profiler.py +++ b/tests/microbenchmarks/dspark_sps_profiler.py @@ -1776,6 +1776,13 @@ def _build_parser() -> argparse.ArgumentParser: help="Sweep a RUNNING trtllm-serve instead of building an engine. The " "verify length is pinned over /dspark/verify_len_pin, so walking " "the ladder costs no engine rebuild.") + source.add_argument( + "--control-url", + help="Where the pin endpoint and /metrics live when they differ from " + "--base-url: under disaggregated serving the LOAD goes through the " + "proxy, but /dspark/verify_len_pin and the iteration stats belong to " + "the GENERATION server. Defaults to --base-url.", + ) source.add_argument("--server-settle-s", type=float, default=20.0, help="Seconds to let admission reach the plateau " "before a cell is measured (--base-url).") @@ -2161,6 +2168,7 @@ def _consume(line: str) -> None: def samples_from_server( *, base_url: str, + control_url: Optional[str] = None, batch_sizes: Sequence[int], verify_lens: Sequence[int], input_len: int, @@ -2197,6 +2205,7 @@ def samples_from_server( """ import requests + control_url = control_url or base_url rng = random.Random(seed) out: List[StepSample] = [] @@ -2228,7 +2237,7 @@ def _fire(body: dict) -> None: post_errors.append(f"{type(exc).__name__}: {exc}") def pin(verify_len: Optional[int]) -> None: - response = requests.post(f"{base_url}/dspark/verify_len_pin", + response = requests.post(f"{control_url}/dspark/verify_len_pin", json={"verify_len": verify_len}, timeout=60) if response.status_code != 200: raise RuntimeError( @@ -2236,7 +2245,7 @@ def pin(verify_len: Optional[int]) -> None: f"{response.status_code} {response.text}") def drain_metrics() -> List[dict]: - response = requests.get(f"{base_url}/metrics", timeout=60) + response = requests.get(f"{control_url}/metrics", timeout=60) response.raise_for_status() payload = response.json() return payload if isinstance(payload, list) else [payload] @@ -2391,6 +2400,8 @@ def main(argv: Optional[Sequence[str]] = None) -> int: elif args.base_url: samples = samples_from_server( base_url=args.base_url.rstrip("/"), + control_url=(args.control_url.rstrip("/") + if args.control_url else None), batch_sizes=batch_sizes, verify_lens=verify_lens, input_len=args.input_len, From 4cb6113fde6a97a7eca0a1632a364e4985c952a0 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sat, 8 Aug 2026 10:36:18 -0700 Subject: [PATCH 112/177] [None][doc] Device-window fix validated: both census verdicts clean 64-way temp0 census on the fixed build: steep mode (max-amplitude w != S every step) 64/64 clean with zero degenerate token runs (was 63-64/64 corrupt before the fix), and real confidence-driven device windows 64/64 clean with coherent text. Remaining rowdiv output under real windows is healthy pace-skew: lagging rows reproduce identical token tuples at later steps, i.e. same text at different progress. GSM8K accuracy at 1024 global concurrency (device vs no-trim reference) running next. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 849fc3458405..e1a8d05a34ef 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -335,6 +335,12 @@ device 臂指纹(各轮一致):completion 恰好=bs×768(零 EOS)、trim ≈49.8 - **验证中**:双节点重发修复后 serve——bia0003 steep(最狠压力)、bia0123 真实 device windows,各挂 64 并发 temp0 census + rowdiv 收割。判净标准:census 64/64 chars/tok≥1.5;steep 下 rowdiv 每步至多 2 个签名组(奇偶配对类)且无退化 token。 - 附:静态审计同时清点出两个**潜伏**(非本案)bug:① 图内逐 token 采样参数缓冲 (temperatures/top_ks/top_ps) 按 S 打包、非贪婪异参数批下 w≠S 会串染(temp0/force-argmax 不读,故与本案无关);② plain-DSA 的 block_table S 口径 −1 掩码对 device windows 不安全(V4 路径不可达)。修复后另行立项。 +### 【08-08 10:31】修复验证判决:双面 census 全净 + +- **steep 面(bia0003,最大幅度 w≠S 每步)**:64/64 干净(chars/tok 中位 4.07,min 3.41),rowdiv 输出中退化 token 串计数 **0**。修复前同配置:63-64/64 腐蚀、大批行塌缩 `(0,0,0)`。 +- **真实 device windows 面(bia0123)**:64/64 干净(中位 4.07),样本输出为连贯诗歌。rowdiv 行仍存在但形态为**健康节奏错峰**——真实窗口下各行推进速度不同,滞后行在后续步逐字复现相同元组(`(1956,7633,989)` step#8→9、`(13132,)` step#16→20),即同一文本、不同进度;这也说明 rowdiv 仪表的"锁步"前提只在 identity/同窗口类内成立,判净以 census 为准。 +- 结论:**bug #2 定案闭环——机制、观测、修复、验证互相咬合**。后续:completion 指纹对照(96.6% host+steep 参照线)、GSM8K 大并发精度(512/1024 global,device vs notrim)、诚实 device 吞吐网格。 + ### 备选修复方案评估(供评审;当前实施的是 A,标记为暂定) - **A(已实施,Python-only)**:捕获拷贝换 D2D 源(`_seq_lens_cuda`)+ NEXT_N 抬全局上界。最小爆炸半径,当场可在活挂载容器里验证;把双胞胎缓冲降级为 canonical 缓冲的衍生物。host 路径装填时刻数值逐位等价,uniform/eager 分支未触碰。 From f40804fb63e8827657166b3c5ff790f67f1836bb Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sat, 8 Aug 2026 11:05:27 -0700 Subject: [PATCH 113/177] [None][doc] GSM8K at 1024 global concurrency: device windows 96.323 vs no-trim 96.247 Both arms pass the 96.0 reference on the fixed build (DEP8, overlap scheduler, CUDA graphs, TRTLLM MoE); the +0.08pp delta is well inside GSM8K sampling noise. Accuracy is loss-free under device-window re-ranking at full concurrency. Three-arm throughput rounds (device/host/no-trim) running next. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index e1a8d05a34ef..83fdf05a7d2b 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -341,6 +341,11 @@ device 臂指纹(各轮一致):completion 恰好=bs×768(零 EOS)、trim ≈49.8 - **真实 device windows 面(bia0123)**:64/64 干净(中位 4.07),样本输出为连贯诗歌。rowdiv 行仍存在但形态为**健康节奏错峰**——真实窗口下各行推进速度不同,滞后行在后续步逐字复现相同元组(`(1956,7633,989)` step#8→9、`(13132,)` step#16→20),即同一文本、不同进度;这也说明 rowdiv 仪表的"锁步"前提只在 identity/同窗口类内成立,判净以 census 为准。 - 结论:**bug #2 定案闭环——机制、观测、修复、验证互相咬合**。后续:completion 指纹对照(96.6% host+steep 参照线)、GSM8K 大并发精度(512/1024 global,device vs notrim)、诚实 device 吞吐网格。 +### 【08-08 11:20】GSM8K 大并发精度判决:device 96.323 vs notrim 96.247 —— 无损 + +- 双臂并行(1024 global = 8 rank × bs128,DEP8 + overlap + CUDA graph,TRTLLM MoE):**notrim 参照**(`test_gsm8k_dep8_static_verify_overlap`)= **96.247**;**device windows 臂**(`test_gsm8k_dep8_ragged_verify` + `TLLM_DSPARK_DEVICE_WINDOWS=1` + v2b 表)= **96.323**。两臂均过 96.0 参考线;差值 +0.08pp 远小于 GSM8K 抽样噪声(σ≈0.5pp)——**修复后的 device 窗口在大并发下精度无损**(用户要求的验证项完成)。 +- 进行中:三臂吞吐复跑(rounds 方法论,poetry+arena × bs512/1024 global × 5 reps + warm rep)——device(bia0047)/ host(bia0048)/ notrim(bia0123),回答 #30 的原始问题:干净生成下 device 收益是否真实存在。 + ### 备选修复方案评估(供评审;当前实施的是 A,标记为暂定) - **A(已实施,Python-only)**:捕获拷贝换 D2D 源(`_seq_lens_cuda`)+ NEXT_N 抬全局上界。最小爆炸半径,当场可在活挂载容器里验证;把双胞胎缓冲降级为 canonical 缓冲的衍生物。host 路径装填时刻数值逐位等价,uniform/eager 分支未触碰。 From 0b9c9d4a44e81d5c9dce5e7048f7bfe41d467c68 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sat, 8 Aug 2026 11:58:25 -0700 Subject: [PATCH 114/177] [None][doc] Honest three-arm throughput closes #30: no device gain with clean generation Fixed-build rounds (node-local async client, medians of 5 steady reps): device sits at -6.6%..+0.3% vs host and both trim arms sit at ~0 +/- 4% vs no-trim across poetry/arena x bs512/1024 global. The original +64% device figure is conclusively a corruption artifact. Consistent with the earlier four-round host verdict (0 +/- 3%) and the SGLang-native dynamic experiment: no reproducible dynamic-trim throughput gain on B300 agg DEP8 under this calibrated config. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 83fdf05a7d2b..0392fac43712 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -346,6 +346,22 @@ device 臂指纹(各轮一致):completion 恰好=bs×768(零 EOS)、trim ≈49.8 - 双臂并行(1024 global = 8 rank × bs128,DEP8 + overlap + CUDA graph,TRTLLM MoE):**notrim 参照**(`test_gsm8k_dep8_static_verify_overlap`)= **96.247**;**device windows 臂**(`test_gsm8k_dep8_ragged_verify` + `TLLM_DSPARK_DEVICE_WINDOWS=1` + v2b 表)= **96.323**。两臂均过 96.0 参考线;差值 +0.08pp 远小于 GSM8K 抽样噪声(σ≈0.5pp)——**修复后的 device 窗口在大并发下精度无损**(用户要求的验证项完成)。 - 进行中:三臂吞吐复跑(rounds 方法论,poetry+arena × bs512/1024 global × 5 reps + warm rep)——device(bia0047)/ host(bia0048)/ notrim(bia0123),回答 #30 的原始问题:干净生成下 device 收益是否真实存在。 +### 【08-08 12:00】#30 终局:诚实三臂吞吐 —— device>host 不复现,原 +64% 判定为腐蚀假象(定案) + +修复后干净生成下的三臂对照(rounds 方法论,节点本地 asyncio 客户端,median tok/s,5 稳态 reps,warm rep 剔除;device=bia0047 / host=bia0048 / notrim=bia0123): + +| cell | notrim | host | device | host/notrim | device/host | +|---|---|---|---|---|---| +| poetry 512 | 12340 | 11840 | 11058 | −4.1% | −6.6% | +| poetry 1024 | 12912 | 12703 | 12385 | −1.6% | −2.5% | +| arena 512 | 16766 | 16221 | 16275 | −3.3% | +0.3% | +| arena 1024 | 13526 | 13636 | 13350 | +0.8% | −2.1% | + +- **主问题答案(#30 "若 device>host 复现,搞清楚为什么")**:干净生成下**不复现**——device 相对 host 在 −6.6%~+0.3% 区间(本轮未做节点对调,±3-5% 内视为噪声;poetry 512 的 −6.6% 或含真实序幕开销成分)。原 +64% 完全由腐蚀假象制造(重复特殊 token + 零 EOS + 满批不排水)。 +- host trim vs notrim 仍 ≈ 0±4%,与此前四轮判决(0±3%)一致。**B300 agg DEP8 上,该校准配置下 host/device 动态裁剪均无可复现吞吐收益**,与 #28(SGLang 原生动态在 B300 高 bs 亏损)闭环互证。 +- 量具备注:首版登录节点线程池客户端因线程耗尽产生 6× 失真(1998 tok/s),已废弃;节点本地 asyncio 客户端复现旧轮次量程(warm rep 11954 vs 旧轮 10.4-11.9k)。散在失败请求 <2%/格、三臂均摊,不影响中位数。 +- 至此 #30 全链闭环:STS 校准 → 稳定性轮次 → device>host 根因(= bug #1 kv_lens 双份 + bug #2 compressor 捕获 H2D)→ 修复 → census/GSM8K/吞吐三重验证。 + ### 备选修复方案评估(供评审;当前实施的是 A,标记为暂定) - **A(已实施,Python-only)**:捕获拷贝换 D2D 源(`_seq_lens_cuda`)+ NEXT_N 抬全局上界。最小爆炸半径,当场可在活挂载容器里验证;把双胞胎缓冲降级为 canonical 缓冲的衍生物。host 路径装填时刻数值逐位等价,uniform/eager 分支未触碰。 From fdfa9a45568fc6a90a196970133cf34a52bfe3f1 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sat, 8 Aug 2026 15:21:48 -0700 Subject: [PATCH 115/177] [None][doc] Quantitative decomposition of why dynamic trim shows no gain on B300 Per-phase counter deltas from the three fixed-build arms, composed with the decomposed cost table: (1) the untrimmable floor (fixed 29.5ms + batch overhead) is 47-60% of the step at bs64-128/rank and theta has wide flat shelves (tier 5->4 saves exactly 0ms at bs64); (2) on the low-accept workload the 32% trim buys 8.3% step time but costs 12% tokens/step (accept 1.53->1.23), netting the measured -4%; (3) on the high-accept workload the calibrated planner correctly declines to trim (5.8-7.6% delivered after bucket rounding) and the residual -3% is the scheduling machinery tax; (4) device re-ranking does raise accept (+0.036) but its prologue overhead eats the gain. Independently corroborated by the SGLang-native B300 experiment: +17.8% at bs1, -21% at bs128 with iter time rising as accept drops. The no-gain verdict is a B300 cost-structure property, not an implementation defect; the gain regime is low concurrency / steeper-theta hardware. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 0392fac43712..73d8ee73b11a 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -362,6 +362,34 @@ device 臂指纹(各轮一致):completion 恰好=bs×768(零 EOS)、trim ≈49.8 - 量具备注:首版登录节点线程池客户端因线程耗尽产生 6× 失真(1998 tok/s),已废弃;节点本地 asyncio 客户端复现旧轮次量程(warm rep 11954 vs 旧轮 10.4-11.9k)。散在失败请求 <2%/格、三臂均摊,不影响中位数。 - 至此 #30 全链闭环:STS 校准 → 稳定性轮次 → device>host 根因(= bug #1 kv_lens 双份 + bug #2 compressor 捕获 H2D)→ 修复 → census/GSM8K/吞吐三重验证。 +### 【08-08 15:45】"为什么没有收益"定量分解(三臂日志计数器差分 + 成本表 + SGLang 原生对照) + +**方法**:三臂 serve 日志的 periodic 计数器(全部为累积量)按 bench 阶段做差分,得到每格的逐阶段 accept_len / trim / bucket_hist(每步实际重放的捕获桶);与 v2b 成本表的分解语义(θ(M) 可裁项 + 29.5ms 固定 + α(bs) 批开销)合成。notrim 臂无统计(DSparkRaggedStats 只在 confidence-scheduling 分支创建——仪表盲区,已记录),其 accept 由吞吐比值模型反解并与四轮存档交叉验证。 + +**每阶段实测(rank-0 切片,稳态区)**: + +| 阶段 | 臂 | accept_len | 落地 trim | 请求 trim | 主导桶 | +|---|---|---|---|---|---| +| poetry512 | host | 1.225 | 31.8% | 34.4% | 256@bs64 (61%) | +| poetry512 | device | 1.261 | 32.1% | — | 256@bs64 (57%) | +| arena512 | host | 2.534 | **5.8%** | 13.5% | 384@bs64 (=满块) | +| arena512 | device | 2.585 | 6.3% | — | 384@bs64 | +| poetry1024 | host | 1.208 | 33.1% | 35.4% | 256@bs64 (53%) | +| arena1024 | host | 2.566 | 7.6% | 16.3% | 384/432/320 混 | + +**四个成因,量化到每一格:** + +1. **可裁蛋糕本来就小(B300 成本结构)**:bs64/rank 满块步 ≈ 60.3ms,其中固定 29.5 + α(64) 6.9 = **60% 不可裁**;bs128 也有 47%。θ(M) 还有宽平台(θ(320)=θ(384)=23.85、θ(256)=θ(288)):**tier 5→4 的第一档裁剪 kernel 时间省 0ms**。首个有实际节省的档要砍到 256 桶(省 5.0ms = 8.3%/步)。 +2. **accept 损失 > 步时节省(poetry,主项)**:host 裁 32% token 拿到 8.3% 步时节省,但 accept 1.53(反解的 notrim 值,poetry1024 反解 1.45 交叉一致;四轮混合存档 1.73-1.78 与 arena≈2.55 的混合构成也吻合)→ 1.23,tokens/step −12%。净:0.88 × 1.09 ≈ 0.96 = **实测的 −4.1%**。模型在 poetry1024 同样闭合(−1.6%)。 +3. **机制税(arena,唯一贡献项)**:arena 上校准后的 planner 正确地几乎不裁(请求 13.5-16.3%,**经桶取整落地只剩 5.8-7.6%**——fill_bucket 把取整余量还给真实验证),但仍付信心头 + planner 集合 + ragged fit 的常数税 → arena512 −3.3%(arena1024 +0.8% 噪声内;该阶段 host 臂另有 222 步 eager 逃逸)。 +4. **device 额外的序幕开销**:device 的 lag-0 重排**确实提升 accept**(1.225→1.261,窗口直方图向 3 集中——机制生效),但序幕的十几个小 kernel + D2H ride 的每步开销吃掉了它,净 −2.5~−6.6% vs host。 + +**独立复核(#28,SGLang 自家实现原生跑 B300)**:bs=1 时动态 **+17.8%**(iter 16.7→13.2ms),bs 8-64 持平,bs128 **−21±4%**(accept 4.65→4.08 而 iter 反升 28.1→33.2ms),bs256 −4~−8%。收益只存在于 θ 陡峭的低占用 regime;高并发下 B300 的 θ 平台宽到裁剪换不回 accept——**与实现无关,是硬件成本结构**。SGLang blog 的 +10-20% 在 H200/H100 上 θ 份额高、平台窄,不可迁移。 + +**次要成因(可改进,但理论上限也有限)**:① 表的已知偏差(L-sweep 连 draft pass 一起缩,θ 边际被高估 → 过度裁剪);② serve 用 tiers [2,3,4,5] 而表自荐 [1,2,5](密集档全躺在平台上,缺低 accept 时 argmax 选的深档);③ planner 按真实 bs 定价、执行按 padded_bs 取整(量化未定价);④ tier 边缘按 bs128 推导,不随 bs 迁移。按表内 shelf 重算,即便 tiers 换 [1,2,5],bs64 poetry 的理论净收益也只有 ~+3%。 + +**结论**:B300 agg 大并发下动态裁剪无收益不是 bug,是**成本结构性的**:不可裁地板占一半以上 + θ 宽平台 + 桶阶梯,使"省下的 token"折不成时间;而任何裁剪都真金白银付 accept。收益 regime 在低并发/延迟场景(SGLang 原生 bs1 +17.8% 亦证)与 θ 更陡的硬件上。 + ### 备选修复方案评估(供评审;当前实施的是 A,标记为暂定) - **A(已实施,Python-only)**:捕获拷贝换 D2D 源(`_seq_lens_cuda`)+ NEXT_N 抬全局上界。最小爆炸半径,当场可在活挂载容器里验证;把双胞胎缓冲降级为 canonical 缓冲的衍生物。host 路径装填时刻数值逐位等价,uniform/eager 分支未触碰。 From a26a319b9ebd195bed987614b24aab2d2e6b2917 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sat, 8 Aug 2026 23:05:26 -0700 Subject: [PATCH 116/177] =?UTF-8?q?[None][doc]=20Cross-stack=20check:=20SG?= =?UTF-8?q?Lang=20gains=20+2.6..+7.7%=20on=20our=20workload=20=E2=80=94=20?= =?UTF-8?q?no-gain=20verdict=20corrected?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SGLang (their #28 image, their optimal config, their profiler-collected table) driven with our exact rounds workload on the same B300 node type: poetry512 +7.7%, poetry1024 +2.6%, arena1024 +2.9%, arena512 -0.3% (compact vs no-trim). Mechanism data shows the same trade terms as ours (cap 6 -> 3.58, accept 2.79 -> 2.52 at poetry512) but a working exchange rate: their trimmed steps actually run faster. The earlier 'structural B300 no-gain' conclusion is narrowed to 'our implementation does not convert trimmed tokens into time'. Suspects ranked: graph-shape economics (their uniform per-step cap + no tier alignment vs our ragged windows + tier buckets, #19/#21), per-GPU load regime (their tp4/dp4 is 2x our DEP8), and trim-machinery tax. Follow-ups queued: DEP4@128/rank rerun, continuous-budget prototype, uniform-cap feasibility. Caveats: Flash model (Pro crashes their stack at graph capture), different per-rank concurrency. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 73d8ee73b11a..d8a8e10e6f0a 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -390,6 +390,26 @@ device 臂指纹(各轮一致):completion 恰好=bs×768(零 EOS)、trim ≈49.8 **结论**:B300 agg 大并发下动态裁剪无收益不是 bug,是**成本结构性的**:不可裁地板占一半以上 + θ 宽平台 + 桶阶梯,使"省下的 token"折不成时间;而任何裁剪都真金白银付 accept。收益 regime 在低并发/延迟场景(SGLang 原生 bs1 +17.8% 亦证)与 θ 更陡的硬件上。 +### 【08-08 23:00】跨栈对照(SGLang × 我们的负载):**不复现我们的无收益——结论修正** + +用户指令:让 SGLang 跑与我们完全相同的负载,表用他们自己的采集方法。设置:他们的 #28 镜像 + 他们的最优配置(Flash tp4/dp4、dsv4 后端、禁 autotune;Pro 在其镜像图捕获阶段崩溃 `Inplace update to inference tensor`,栈级兼容性 bug,已记录)+ 他们的 profiler 现采表(17 探针自检通过)+ **我们的 rounds 负载**(poetry/arena × 512/1024 global × warm+5 reps,同一客户端)。 + +| cell | notrim | compact(表)| delta | +|---|---|---|---| +| poetry 512 | 8744 | 9419 | **+7.7%** | +| poetry 1024 | 6293 | 6455 | +2.6% | +| arena 512 | 7959 | 7933 | −0.3% | +| arena 1024 | 6307 | 6493 | +2.9% | + +- **机制数据(poetry512,decode 日志逐步统计)**:notrim accept_len 2.79 @ cap 6.00;compact accept_len 2.52 @ cap 3.58。即:裁掉 ~40% verify token,付出 −9.7% tokens/step,**净 +7.7%**——他们把省下的 token 兑换成了时间。 +- **与我们的同工况对比**:我们在 128/rank(poetry1024)裁剪深度 (33%) 与 accept 代价 (~−10%) 几乎相同,但净 −1.6%。**两边的"贸易条款"一样,差在结算汇率**——我们的省 token 没有变成省时间。 +- **结论修正**:昨判"B300 结构性无收益"过强。正确表述:**在我们的实现里无收益;SGLang 在同硬件同负载上能兑现 +2.6~+7.7%**(高 accept 的 arena 512 例外 ≈ 0,方向与我们一致)。差距候选(按嫌疑排序): + 1. **图形状经济学**:他们 `align_verify_tokens_to_graph_tier=False` + 每步**均匀 cap**(标量 cap len)——裁剪直接改变每行 token 数,步进更小的图形状;我们是 per-request ragged 窗口 + tier 桶阶梯(#19/#21),裁剪常被桶取整吞掉。 + 2. **每 GPU 负载 regime**:他们 tp4/dp4 = 每 GPU 负载是我们 DEP8 的 2×,θ 份额更陡,兑换率天然更好。**建议补一个对照:我们的栈跑 4 卡 DEP4 @128/rank**,若收益出现则确认 regime 项,若仍无则全归实现项。 + 3. 机制税:他们不裁时 ≈ 0 损耗(arena −0.3%),我们 −3.3%。 +- Caveats:模型不同(Flash accept 2.79 vs Pro 1.45 于 poetry;Pro 在其栈不可跑),per-rank 并发不同(128/256 vs 我们 64/128);+2.6~+2.9% 的两格在 ±3% 噪声边缘,+7.7% 超出噪声带。 +- 后续实验队列:① 我们的栈 DEP4×128/rank 复测;② #19 连续预算(去 tier 量化)原型;③ 均匀 cap 模式(替代 ragged)可行性评估——若三者叠加能兑现 SGLang 级别收益,#25/#30 产品线复活。 + ### 备选修复方案评估(供评审;当前实施的是 A,标记为暂定) - **A(已实施,Python-only)**:捕获拷贝换 D2D 源(`_seq_lens_cuda`)+ NEXT_N 抬全局上界。最小爆炸半径,当场可在活挂载容器里验证;把双胞胎缓冲降级为 canonical 缓冲的衍生物。host 路径装填时刻数值逐位等价,uniform/eager 分支未触碰。 From f2b4a1dee2775f2f3a3075c2d21d657b1c6214e8 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 9 Aug 2026 00:11:54 -0700 Subject: [PATCH 117/177] [None][doc] DEP4 attempt infeasible (Pro OOM at tp4); regime question settled analytically Pro's 832GB weights leave no headroom for graph capture at tp4 even with kv 0.32 and a 5-entry capture set (three OOMs, ~1.5GB short). The regime discriminator does not need it: the per-GPU-load-matched cells already exist -- their poetry512 (128 reqs + ~128k MoE tokens per GPU) vs our poetry1024 (same per-GPU pressure) read +7.7% vs -1.6%. The gap is implementation-side; remaining confounds are EP width and model. The continuous-budget prototype (#19) plus uniform-cap evaluation are promoted to the primary fix track. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index d8a8e10e6f0a..084646c2881e 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -410,6 +410,11 @@ device 臂指纹(各轮一致):completion 恰好=bs×768(零 EOS)、trim ≈49.8 - Caveats:模型不同(Flash accept 2.79 vs Pro 1.45 于 poetry;Pro 在其栈不可跑),per-rank 并发不同(128/256 vs 我们 64/128);+2.6~+2.9% 的两格在 ±3% 噪声边缘,+7.7% 超出噪声带。 - 后续实验队列:① 我们的栈 DEP4×128/rank 复测;② #19 连续预算(去 tier 量化)原型;③ 均匀 cap 模式(替代 ragged)可行性评估——若三者叠加能兑现 SGLang 级别收益,#25/#30 产品线复活。 +**【08-09 00:25 追记】regime 判别的结论(实验 ① 改为解析判定)**: +- DEP4-Pro 实测尝试三次均 OOM(Pro 832GB 权重 tp4 = 208GB/卡,图捕获期差 1-2GB;KV 0.32 + 捕获集裁到 5 尺寸仍不够)——Pro 在 4 卡上不可行,与 SGLang 侧 Pro 崩溃互为印证:**跨栈同模型比较做不到,模型差异是本对照的固有 caveat**。 +- 但 regime 问题不需要 DEP4:**per-GPU 负载对齐的格子已经存在**。他们 poetry512(tp4/dp4:每 GPU 128 请求、~128k token MoE 负载)↔ 我们 poetry1024(DEP8:每 GPU 同样 128 请求、~128k token)。同等每 GPU 压力:**他们 +7.7% vs 我们 −1.6%**。→ **差距主体是实现侧**(图形状经济学 + 机制税),不是负载 regime;残余混杂仅剩 EP 宽度(EP8 vs EP4 的 a2a)与模型(Flash vs Pro 的 accept 分布)。 +- 优先级更新:#19 连续预算原型 + 均匀 cap 模式评估升为首位修复路线;DEP4 复测取消。 + ### 备选修复方案评估(供评审;当前实施的是 A,标记为暂定) - **A(已实施,Python-only)**:捕获拷贝换 D2D 源(`_seq_lens_cuda`)+ NEXT_N 抬全局上界。最小爆炸半径,当场可在活挂载容器里验证;把双胞胎缓冲降级为 canonical 缓冲的衍生物。host 路径装填时刻数值逐位等价,uniform/eager 分支未触碰。 From 9f12c06bbefd656bc4dd28a8fcd2b5b9365ea9a1 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 9 Aug 2026 00:55:58 -0700 Subject: [PATCH 118/177] [None][doc] Exchange-rate math closes: prime suspect is batch fragmentation, not trim economics Composing each stack's own cost table with its measured trim depth and accept cost: our v2b table predicts +7.7% at true 128-row steps (the same number SGLang measures), and -1.6% at 64-row steps (the number we measure). The padded_bs histogram shows the executor ran 64-row steps 53% of the time at 1024-global concurrency and 128-row steps only 3%, while SGLang held 126-170 running requests per rank under the same pressure. The trim mechanism, decisions, and cost table are fine -- the batch assembly slides every step into the floor-dominated regime where trimming cannot pay. New top task: root-cause the 64-row fragmentation (admission throttling, scheduler capacity, DP routing, or ragged-fit fallback); the fix criterion is the histogram mode moving to 128 and poetry1024 flipping positive. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 084646c2881e..5f573f964044 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -415,6 +415,21 @@ device 臂指纹(各轮一致):completion 恰好=bs×768(零 EOS)、trim ≈49.8 - 但 regime 问题不需要 DEP4:**per-GPU 负载对齐的格子已经存在**。他们 poetry512(tp4/dp4:每 GPU 128 请求、~128k token MoE 负载)↔ 我们 poetry1024(DEP8:每 GPU 同样 128 请求、~128k token)。同等每 GPU 压力:**他们 +7.7% vs 我们 −1.6%**。→ **差距主体是实现侧**(图形状经济学 + 机制税),不是负载 regime;残余混杂仅剩 EP 宽度(EP8 vs EP4 的 a2a)与模型(Flash vs Pro 的 accept 分布)。 - 优先级更新:#19 连续预算原型 + 均匀 cap 模式评估升为首位修复路线;DEP4 复测取消。 +### 【08-09 01:30】SGLang 深挖:兑换率数学闭合,第一嫌疑修正为**批碎片化** + +双方各自的成本表 × 各自实测的裁剪深度/accept 代价: + +| | 步时 | accept | 预测净 | 实测 | +|---|---|---|---|---| +| SGLang poetry512(128/rank 真实步)| 44.9→35.4ms(−21%)| −9.7% | +14.7% | +7.7% | +| 我们 @128 行/步(假想满行)| 92.2→77.2ms(−16%)| −9.8% | **+7.7%** | — | +| 我们 @64 行/步(实际形态)| 60.3→55.3ms(−8%)| −9.8% | **−1.6%** | **−1.6%** | + +- **我们自己的 v2b 表预测:真 128 行/步下裁剪净赚 +7.7%**(与 SGLang 实测同值)。但 poetry1024 的 padded_bs 直方图 = {64: 2364, 72: 443, **128: 仅 134**}/4448 步——**执行器在 1024-global 下多数步只装 64 行**;64 行 regime 地板占 60%,预测 −1.6% = 实测 −1.6%,逐位闭合。 +- **结论**:我们的裁剪机制、决策、成本表全都没大错——**是执行器的批装配把每一步滑进了裁剪必亏的 regime**。SGLang 同压力下每步稳定 126-170 行(decode 日志)。修好碎片化,poetry1024 按自家表应翻到 ~+7.7%,毋须改裁剪本身。 +- **新任务(第一优先)**:定位 1024-global 下 rank 步只有 ~64 行的机制——候选:serve 准入节流、调度器容量交互、DP 路由不均、ragged fit 大批回退。判据:修复后 padded_bs 直方图模式移到 128,poetry1024 转正。 +- 次级(降级但成立):64/rank 低压场景救不回(地板 60%,SGLang 也是靠 128/rank 赚);我们 θ 有平台他们没有(tp4 大形状);地板 29.5 vs 13ms(层数归一后仍 ~1.6×);机制税。SGLang 画像:高 accept 的 arena 他们也只打平(cap 4.01,−0.3%/+2.9%);超出 graph 覆盖(>192/rank)收益缩水到 +2.6%。 + ### 备选修复方案评估(供评审;当前实施的是 A,标记为暂定) - **A(已实施,Python-only)**:捕获拷贝换 D2D 源(`_seq_lens_cuda`)+ NEXT_N 抬全局上界。最小爆炸半径,当场可在活挂载容器里验证;把双胞胎缓冲降级为 canonical 缓冲的衍生物。host 路径装填时刻数值逐位等价,uniform/eager 分支未触碰。 From 4fe403b9b55079e3cd2a31f304ae32962ae05180 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 9 Aug 2026 01:31:17 -0700 Subject: [PATCH 119/177] [None][doc] Cap-accept verdict (#20): only 6.2%/16.1% of trimmed tokens would have been accepted Direct counterfactual measurement (cap-accept mode: same planner decisions, full-block verify, acceptance capped): poetry trims 2.08 positions per request per step of which 6.2% were alive (0.130 tokens, 9.8% of no-trim acceptance); arena trims 0.71 of which 16.1% were alive (4.3%). The decision layer is exonerated -- the confidence screen works and the planner correctly trims less on the high-accept workload. This closes the responsibility chain: decisions good, table economics sound at true 128-row steps, execution (batch fragmentation into 64-row floor-dominated steps) is the sole failure point. Also corrects the SGLang profiler characterization: it supports a 2D additive table via --fracs (the lineage of our #26 frac sweep); the canonical/blog flow and both cross-stack runs used the 1D static mode. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 5f573f964044..682437042641 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -430,6 +430,21 @@ device 臂指纹(各轮一致):completion 恰好=bs×768(零 EOS)、trim ≈49.8 - **新任务(第一优先)**:定位 1024-global 下 rank 步只有 ~64 行的机制——候选:serve 准入节流、调度器容量交互、DP 路由不均、ragged fit 大批回退。判据:修复后 padded_bs 直方图模式移到 128,poetry1024 转正。 - 次级(降级但成立):64/rank 低压场景救不回(地板 60%,SGLang 也是靠 128/rank 赚);我们 θ 有平台他们没有(tp4 大形状);地板 29.5 vs 13ms(层数归一后仍 ~1.6×);机制税。SGLang 画像:高 accept 的 arena 他们也只打平(cap 4.01,−0.3%/+2.9%);超出 graph 覆盖(>192/rank)收益缩水到 +2.6%。 +### 【08-09 01:35】cap-accept 判决(#20):被裁 token 只有 6.2%/16.1% 本会被接受——决策层无罪 + +cap-accept 模式(同一 planner 决策,提交完整块全量验证、接受端封顶,逐请求精确计数反事实)跑同负载: + +| | poetry | arena | +|---|---|---| +| 被裁位置数/请求/步 | 2.08 | 0.71 | +| 其中本会被接受 | **0.130(6.2%)** | **0.114(16.1%)** | +| 占无裁剪接受量 | 9.8% | 4.3% | +| trim_regret_rate | 22.7% | 34.3% | + +- **决策层定案无罪**:poetry 被裁位置 93.8% 反正会被拒,置信筛选有效;arena 上 planner 少裁(0.71 位)且被裁位置存活率高 2.6 倍——行为与风险结构一致。~10% 的 accept 代价与 compact 臂吞吐分解、SGLang 对照(9.7%)三方吻合。 +- **责任链闭合**:决策层 ✅(本节)→ 经济层 ✅(v2b 表:真 128 行/步预测 +7.7%)→ **执行层 ❌(批碎片化,64 行/步 → 地板 60% → −1.6%)**。修复目标唯一:批装配。 +- 更正记录:SGLang profiler 支持双模式(`--fracs` 给定 → 2D `SpsAdditiveCostTable`,离对角 frac 扫描,与我们 v2b 同构且为 #26 的方法来源;省略 → 对角 static → 1D `SpsCostTable`)。#28 与本次跨栈对照用的均为一维模式(manifest fracs=null)——他们拿一维粗表仍 +7.7%,进一步支持"表精度非胜负手"。 + ### 备选修复方案评估(供评审;当前实施的是 A,标记为暂定) - **A(已实施,Python-only)**:捕获拷贝换 D2D 源(`_seq_lens_cuda`)+ NEXT_N 抬全局上界。最小爆炸半径,当场可在活挂载容器里验证;把双胞胎缓冲降级为 canonical 缓冲的衍生物。host 路径装填时刻数值逐位等价,uniform/eager 分支未触碰。 From 327b9c0a7dd58cd0fc83f8d11954ba65fdfbf2c0 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 9 Aug 2026 07:35:14 -0700 Subject: [PATCH 120/177] [None][doc] Parity verdict: dynamic trim is a compensation mechanism for token-cost-dominated execution Same model (Flash-DSpark), same 4x B300, same per-rank load, same client, each stack's own profiler: TRT-LLM's no-trim baseline is 2.4-3.5x SGLang's. TRT-LLM's fresh table is judged inert by its own profiler (the goodput argmax never trims at any acceptance rate 0.3-0.95) because step time is a function of ROWS, not tokens (theta ~ 0, floor 93-95%: bs128 L=2 and L=5 both take 39.2ms); SGLang's step time is a function of TOKENS (13->45ms, theta ~70%) on the same GPU and model. Forcing SGLang's 40% trim depth on TRT-LLM loses 12-22% across all cells while the same depth gains SGLang +5.9%. Cap-accept ranking quality is comparable on both stacks (11-15% of trimmed positions alive). Their engineering envelope tops out at 192-row decode graphs (256 crashes cublas at capture) and Pro crashes their capture entirely. Product conclusion: on efficient execution the trim feature has no economic target; #31 (Pro DEP8 fragmentation) stays open since Pro's theta share is real, pending a Pro retest after the fix. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 682437042641..72a2a00d8072 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -445,6 +445,31 @@ cap-accept 模式(同一 planner 决策,提交完整块全量验证、接受端 - **责任链闭合**:决策层 ✅(本节)→ 经济层 ✅(v2b 表:真 128 行/步预测 +7.7%)→ **执行层 ❌(批碎片化,64 行/步 → 地板 60% → −1.6%)**。修复目标唯一:批装配。 - 更正记录:SGLang profiler 支持双模式(`--fracs` 给定 → 2D `SpsAdditiveCostTable`,离对角 frac 扫描,与我们 v2b 同构且为 #26 的方法来源;省略 → 对角 static → 1D `SpsCostTable`)。#28 与本次跨栈对照用的均为一维模式(manifest fracs=null)——他们拿一维粗表仍 +7.7%,进一步支持"表精度非胜负手"。 +### 【08-09 03:00-08:00】同台对照终判(parity):裁剪收益 = 慢执行的补偿;我们的栈无肉可裁 + +用户指令的实验:**同模型(Flash-DSpark)、同 4 卡 B300、同 per-rank 负载、同客户端,各自用各自的 profiler 采表**,对照 lbs128/256 下两栈行为。四节点七臂(双方 notrim/trim/cap-accept + 双表)。 + +**① 绝对吞吐(同一切条件,notrim)**: + +| cell | TRT-LLM | SGLang | 比值 | +|---|---|---|---| +| poetry 512 | 20,840 | 8,794 | 2.4× | +| poetry 1024 (lbs256) | 21,838 | 6,324 | **3.5×** | +| arena 512 | 26,279 | 7,941 | 3.3× | +| arena 1024 | 20,041 | 6,397 | 3.1× | + +**② 裁剪收益**:SGLang compact +5.9%(poetry512)/ +0.2~+0.8%(其余);TRT:**profiler 判自家新表 inert(argmax 在 bs32-256×接受率 0.3-0.95 全网格上零次选择裁剪)拒绝落盘**;强制按 SGLang 深度(frac 0.4)裁 → **−12~−22% 全线亏损**,表的判决被实证。 + +**③ 成本结构(双方 profiler 同台采集,机理定案)**:TRT 步时 ≈ 行数的函数——bs128 时 L=2 与 L=5 步时同为 39.2ms(token 翻倍零增长),fixed 20.1ms + α∝行数,**θ≈0,地板 93-95%**;SGLang 步时 ≈ token 的函数(13→45ms 随 M 涨 3.5×,θ 份额 ~70%)。**同一块 GPU、同一个模型,成本结构完全相反——动态裁剪是 token 成本主导执行的补偿机制,在行数成本主导的执行上无标的**。 + +**④ 决策/排序质量(cap-accept 对称测量)**:被裁位置存活率 TRT 15.2%(1.65 位/请求,裸 sigmoid,125 万请求)≈ SGLang 11.4-14.5%(差分估计)≈ Pro+STS 6.2-16.1%(昨日精确计数)——**双方筛子同级,排序不是胜负手**。 + +**⑤ 工程包线**:SGLang decode 图上限实测 192 行(256 行 cublas 捕获崩溃、profiler 自我拒绝),Pro 模型在其栈图捕获崩溃;其每 rank 在飞请求恒饱和 ~110-115(与需求无关)。TRT 256 行图一次捕获成功,notrim 吞吐 512→1024 global 仍在上升。 + +**⑥ 附带修正**:SGLang profiler 亦有二维 additive 表能力(`--fracs` 离对角扫描,#26 方法之源),canonical/blog 流程用一维 static 模式(两次对照均如此);其 STS 机制存在但 canonical 不启用(本对照两侧均未用)。 + +**产品线结论**:B300 + 高效执行(本仓)+ agg 场景下,**动态 verify 裁剪没有经济基础**——不是决策不准(④)、不是表不准(③已双向验证)、而是没有可回收的时间。#19(连续预算)在 θ≈0 下亦无标的。**保留价值**:①(2.4-3.5×)与 lbs256 扩展性本身;#31(Pro@DEP8 批碎片化)仍值得修——Pro 的 θ 份额(40-53%)非零,v2b 预测真 128 行步 +7.7% 未被 Flash 判决否定,须修完碎片化后用 Pro 复测定夺;#25 device 序幕的机制正确性已证(census/GSM8K),其激活条件同上。 + ### 备选修复方案评估(供评审;当前实施的是 A,标记为暂定) - **A(已实施,Python-only)**:捕获拷贝换 D2D 源(`_seq_lens_cuda`)+ NEXT_N 抬全局上界。最小爆炸半径,当场可在活挂载容器里验证;把双胞胎缓冲降级为 canonical 缓冲的衍生物。host 路径装填时刻数值逐位等价,uniform/eager 分支未触碰。 From 420cc4e3b16f823b179a890d022b18153ce728c0 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 9 Aug 2026 10:12:37 -0700 Subject: [PATCH 121/177] [None][doc] Iter-log ground truth overturns flat-theta: short-window bucket shapes are pathologically slow (#32) Five-phase isolation matrix on one serve (print_iter_log only, pin re-dial without reloads): pins verified via per-iteration token counts; enable_iter_perf_stats exonerated (0-2%); ragged machinery small (+1.5/+7.5ms at 128/256 rows); bucket churn ~0. The real finding: window-4 pinned steps run +12%/+34% SLOWER than full window-6 steps at fixed row count despite 33% fewer verify tokens, and window-3 is also slower -- theta is effectively negative on the serving path. The engine-mode profiler's flat rows came from synthetic long-context load masking the pathology. Cross-evidence: the GB300 Pro frac table's w=4 anomaly, our bs128 L=4 outlier, and SGLang's #28 bs128 iter-time rise under trim. Filed #32 (kernel tile/algorithm selection at non-full-width token shapes); with #31, these are the two keys gating the Pro cliff economics (w5->w3: -34% step time for -40% tokens on the GB300 table). Table collection methodology updated: collect from serving iter logs or real-shaped load, stats knob off. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 72a2a00d8072..049ddf956e0d 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -470,6 +470,25 @@ cap-accept 模式(同一 planner 决策,提交完整块全量验证、接受端 **产品线结论**:B300 + 高效执行(本仓)+ agg 场景下,**动态 verify 裁剪没有经济基础**——不是决策不准(④)、不是表不准(③已双向验证)、而是没有可回收的时间。#19(连续预算)在 θ≈0 下亦无标的。**保留价值**:①(2.4-3.5×)与 lbs256 扩展性本身;#31(Pro@DEP8 批碎片化)仍值得修——Pro 的 θ 份额(40-53%)非零,v2b 预测真 128 行步 +7.7% 未被 Flash 判决否定,须修完碎片化后用 Pro 复测定夺;#25 device 序幕的机制正确性已证(census/GSM8K),其激活条件同上。 +### 【08-09 08:50-10:15】iter-log 地面真值:θ 平坦结论被推翻,真凶 = 短窗桶 kernel 病灶(#32) + +用户质疑 profiler 可信度(stats 开销污染?pin 是否生效?)→ 重构为纯 `print_iter_log` 的五阶段隔离矩阵(Flash DEP4,逐迭代 `host_step_time`/`prev_device_step_time`/`num_generation_tokens` 地面真值,单 serve 上 pin 换档零重启): + +| 阶段(同负载)| tok/行 | ~128 行 device ms | ~256 行 device ms | +|---|---|---|---| +| U6 notrim(满块均匀图)| 6.30 | 46.4 | 72.3 | +| R5 pin 满窗(ragged 机制,同 token)| 6.30 | 47.9 | 79.8 | +| R5+`enable_iter_perf_stats` | 6.30 | 48.0 | 81.6 | +| **R3 pin 窗口4(token −33%)** | 4.20 | **53.7** | **107.0** | +| R2 pin 窗口3(token −50%)| 3.15 | 50.3 | 102.8 | +| FRAC 0.4(自然抖动)| 4.20 | 53.7 | 102.4 | + +**分解全部闭合**:① pin 真实生效(token 计数逐档精确变化)→ profiler 钉窗机制无罪;② `enable_iter_perf_stats` 税 0-2% → 污染假设基本排除(但采表默认关掉它,已列方法论改进);③ ragged 机制税小(+1.5ms@128 行 / +7.5ms@256 行);④ 桶抖动 ≈ 0(FRAC ≡ pin4);⑤ **真凶:短窗桶形状本身**——token 砍 1/3,步时反涨 **+12%/+34%**,且非单调(窗口 4 比 3 更糟)。**θ 在 serving 真值下不是 ~0,是负的。** + +- **跨证据链**:GB300 Pro frac 表的 w=4 异常(bs64: 62.7 > w5 的 60.8)、我方 Flash profiler 的 42.09 离群点(bs128 L=4)、SGLang #28 bs128 裁剪后 iter 28→33ms——同一病灶跨模型、跨节点、跨栈复现。嫌疑:attention/indexer/DeepGEMM 在非满宽 token 形状下的 tile/算法选型退化。已立 **#32**,下一步 nsys 对拍 pin5 vs pin3 单步 kernel 时间线点名退化 kernel。 +- **对既有结论的修订**:昨夜"θ≈0 → 裁剪无标的"修正为"**满窗形状高度优化、短窗形状病态劣化 → 当前裁剪不但无收益还倒贴**";profiler 引擎模式的平坦行是合成负载(input_len 1024)遮蔽了病灶,采表流程需改用真实负载形态或 serving 日志建表(`--from-iter-log`)。 +- **产品线含义**:#31(碎片化)+ **#32(短窗 kernel)** 是两把钥匙——#32 修复后 θ 恢复正斜率,Pro 的悬崖交换率(GB300 表:w5→w3 省 34% 步时/40% token)将**优于**当前一切估计;裁剪经济学的最终判决推迟到两者修复后的 Pro 复测。 + ### 备选修复方案评估(供评审;当前实施的是 A,标记为暂定) - **A(已实施,Python-only)**:捕获拷贝换 D2D 源(`_seq_lens_cuda`)+ NEXT_N 抬全局上界。最小爆炸半径,当场可在活挂载容器里验证;把双胞胎缓冲降级为 canonical 缓冲的衍生物。host 路径装填时刻数值逐位等价,uniform/eager 分支未触碰。 From 48a4f1d42abce181cec436f375f001d3a37e3a01 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 9 Aug 2026 21:33:08 -0700 Subject: [PATCH 122/177] [None][doc] nsys report-loss root cause: missing fork tracing; v6 flags Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 049ddf956e0d..ef820f157589 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -498,7 +498,13 @@ cap-accept 模式(同一 planner 决策,提交完整块全量验证、接受端 - **E(否)**:仅 device-windows 模式走 D2D。零收益模式分叉,制造未来漂移面。 - NEXT_N 抬全局上界同时排掉一个 host 侧潜伏雷:#19(取消 tier 阶梯量化)实施后 host 也可能在小 tier bucket 发超 tier 窗口,同样会截断 Phase 1。 -# 分析 +### #32 nsys 取证:四次报告丢失的根因与 v6 修正(08-09 00:20) + +四轮 nsys 采集(SIGINT 版 / TERM 版 / 手动收尾版 / v5 stop-shutdown 版)全部丢失报告,根因**不是收尾方式,而是采集从未开始**:v5 用 `--capture-range-end=stop-shutdown` 时应用理应在 iter 2050 被自动关停,实测 pin3 应用活过 iter 2110 → `cudaProfilerStart` 触发器从未被 nsys 看见。 + +- **代码侧验证(py_executor.py)**:触发门只有 `iter_counter ∈ profile_start_iters and not is_warmup`(L1739-1748),打印的 `iter = N` 与触发比较用的是**同一个计数器**(L1725),三条执行循环(pp/普通/overlap)都套着 `profile_step`——TRT-LLM 侧路径无问题,且触发时会打 `"Profiling started at iteration N."` 日志(此后可用该行判别触发与注入哪个失效)。 +- **真凶(官方配方对照,perf-analysis.md)**:`trtllm-serve` 的 worker 进程是 **fork-before-exec** 派生的,nsys 缺 `--trace-fork-before-exec=true` 时不会向 fork 出的子进程注入 → worker 里的 `cudaProfilerStart` 对 nsys 不可见,capture 永远不开始,报告永远不落盘。nsys 退出时抱怨的 "re-parented processes" 即其指纹。 +- **v6 修正**:`--trace-fork-before-exec=true` + `--cuda-graph-trace node`(后者让 CUDA graph 内的 kernel 逐个展开——我们全部计算都在图里,做 pin5 vs pin3 kernel 对拍必须要 node 粒度)。两轮已重发(pin5@743961 / pin3@743962)。 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 34ddf0440db13775fc3765b1c4a4e50d04c64713 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 9 Aug 2026 22:20:08 -0700 Subject: [PATCH 123/177] [None][doc] nsys v6 verdict (2026.2.1 injection failure) and v7 recipe (extracted 2025.5.1, range-end=stop, window steering) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index ef820f157589..6f7dd600683e 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -506,6 +506,10 @@ cap-accept 模式(同一 planner 决策,提交完整块全量验证、接受端 - **真凶(官方配方对照,perf-analysis.md)**:`trtllm-serve` 的 worker 进程是 **fork-before-exec** 派生的,nsys 缺 `--trace-fork-before-exec=true` 时不会向 fork 出的子进程注入 → worker 里的 `cudaProfilerStart` 对 nsys 不可见,capture 永远不开始,报告永远不落盘。nsys 退出时抱怨的 "re-parented processes" 即其指纹。 - **v6 修正**:`--trace-fork-before-exec=true` + `--cuda-graph-trace node`(后者让 CUDA graph 内的 kernel 逐个展开——我们全部计算都在图里,做 pin5 vs pin3 kernel 对拍必须要 node 粒度)。两轮已重发(pin5@743961 / pin3@743962)。 +**v6 判决(08-09 22:10)**:fork 追踪使触发器**开火了**(两轮 serve log 均出现 `"Profiling started at iteration 2000."`,这是 TRT-LLM 侧打的,证明 env 到位、门条件满足),但 nsys 退出时自供 `Processing events... → No reports were generated`——**容器内 nsys 2026.2.1 仍未注入到 worker 进程**,cudaProfilerStart 调了但没人接。此前"没看到 Profiling stopped"是预期行为(该日志只在 torch trace 开启时打,L1666-1671)。另收获:stop-shutdown 关不掉 re-parented worker,`--wait=primary` 让 nsys 干等,需手动 TERM app(留 nsys)才收尾。 + +**v7(执行中)——采用同事已验证配方**:peihengh 的 dsv4 nsys 战役(同集群同拓扑 B300 + trtllm-llmapi-launch + trtllm-serve)成功出过 `.nsys-rep`,其 start_worker.sh 注释直言 *"Lingjie's hang-avoidance: NSYS_BIN=extracted 2025.5.1"*——与用户"nsys 得用 2025.5.x"的提示互相印证:**容器自带 2026.2.1 有问题,须用 lustre 上提取版 2025.5.1.121**。v7 变更:① nsys 换 `peihengh/dsv4_pipe_clean_20260723/tools/nsys-2025.5.1.121`(容器内烟测可运行);② `--capture-range-end=stop`(cudaProfilerStop 时就地落盘,app 继续跑,绕开 shutdown 死结);③ 照抄其余旗标(`-t cuda,nvtx --cuda-graph-trace node --gpu-metrics-devices=none`);④ **窗口对准**:pin3 v6 的窗口正好砸进 burst 排水缝(iter 2000-2025 只有 1 个请求;pin5 幸运落在稳态 120 请求)——burst 跨度 ≈ max_tokens/accept 与批大小无关(pin5 ~470/pin3 ~515 iter,确定性强),v7 round 先放校准 burst 实测跨度,再用小 filler burst(64 req × 按缺口折算的 max_tokens)把 2000-2050 顶进整 burst 的稳态中段。 + ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) 1. **host 未复现**:初测 +8.3%(poetry 128/rank)在四轮中缩为 [0.0, +1.9, +1.3, +2.2]%;arena 侧 [−2.4, 0.0, −1.4, −5.6]%。host 整体 ≈ 0±3%,poetry 微正、arena 微负。host 调度行为本身跨节点完全可复现(trim 22.4-22.7%,accept 1.66-1.67,四轮逐位一致)——**是收益不存在,不是实验不稳**。 From ceae4e209fd58f85531c353123b2e366086e502b Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 9 Aug 2026 22:49:13 -0700 Subject: [PATCH 124/177] [None][doc] nsys v7 reports landed; preliminary kernel diff points at MoE grouped-GEMM tile ladder Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 6f7dd600683e..b89c46f64b55 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -510,6 +510,20 @@ cap-accept 模式(同一 planner 决策,提交完整块全量验证、接受端 **v7(执行中)——采用同事已验证配方**:peihengh 的 dsv4 nsys 战役(同集群同拓扑 B300 + trtllm-llmapi-launch + trtllm-serve)成功出过 `.nsys-rep`,其 start_worker.sh 注释直言 *"Lingjie's hang-avoidance: NSYS_BIN=extracted 2025.5.1"*——与用户"nsys 得用 2025.5.x"的提示互相印证:**容器自带 2026.2.1 有问题,须用 lustre 上提取版 2025.5.1.121**。v7 变更:① nsys 换 `peihengh/dsv4_pipe_clean_20260723/tools/nsys-2025.5.1.121`(容器内烟测可运行);② `--capture-range-end=stop`(cudaProfilerStop 时就地落盘,app 继续跑,绕开 shutdown 死结);③ 照抄其余旗标(`-t cuda,nvtx --cuda-graph-trace node --gpu-metrics-devices=none`);④ **窗口对准**:pin3 v6 的窗口正好砸进 burst 排水缝(iter 2000-2025 只有 1 个请求;pin5 幸运落在稳态 120 请求)——burst 跨度 ≈ max_tokens/accept 与批大小无关(pin5 ~470/pin3 ~515 iter,确定性强),v7 round 先放校准 burst 实测跨度,再用小 filler burst(64 req × 按缺口折算的 max_tokens)把 2000-2050 顶进整 burst 的稳态中段。 +### #32 nsys 取证 v7:报告到手 + 初步 kernel 对拍(08-09 22:50) + +**五连丢终结**:v7(提取版 nsys 2025.5.1.121 + `--capture-range-end=stop` + 窗口对准)两轮均落盘 27MB 报告(`nsys/pin5_rank0.nsys-rep` / `pin3_rank0.nsys-rep`),窗口 iter 2000-2050,graph 内 kernel node 粒度。对准偏差:两窗都落在 burst 的 ~65% 位置(排水段),pin5 占用 72→23 请求、pin3 57→6——pin3 没够到 padded_bs=128 档,已再发 v7b 两轮(对准收紧到 burst 20-35% 段)补清洁样本。 + +**混形状粗对拍(初步,待 v7b 清洁窗验证)**: +| 信号 | pin5(6 tok/req) | pin3(4 tok/req) | 读法 | +|---|---|---|---| +| MoE FC1 大 tile `t128x256x256`(swiGlu) | 17.0%,中位 178µs | 3.2%(几乎退场) | 大桶专属,效率最高 | +| MoE FC1 中 tile `t128x128x256u2` | 1.7% | **8.0%,中位 110µs** | 小桶主力:token 减半,GEMM 只省 38% | +| MoE FC1 小 tile `t128x64x256u2` | 0.1%,中位 69µs | **5.7%,中位 94µs(m3/m5=1.35)** | 同名 kernel 在 pin3 上下文里慢 35% | +| FMHA sparse / A2A / topK / 元素级 | — | 比值 0.6-0.8 | 正常随 token 缩放 | + +**初步指认**:短窗桶的病灶集中在 **MoE grouped GEMM 的 tile 阶梯**——token 桶变小 → 换小 tile → 每 token 效率近乎腰斩,GEMM 时间不随 token 线性下降,形成 #32 的"常数惩罚";与 GB300 Pro frac 表的平台+悬崖楼梯(MoE 更重 → 惩罚更大)跨模型同构。另见 pin3 独有的 FMHA 变体(1.0%)待查。**定论等 v7b 的 (128,768) vs (128,384) 同批次干净对拍。** + ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) 1. **host 未复现**:初测 +8.3%(poetry 128/rank)在四轮中缩为 [0.0, +1.9, +1.3, +2.2]%;arena 侧 [−2.4, 0.0, −1.4, −5.6]%。host 整体 ≈ 0±3%,poetry 微正、arena 微负。host 调度行为本身跨节点完全可复现(trim 22.4-22.7%,accept 1.66-1.67,四轮逐位一致)——**是收益不存在,不是实验不稳**。 From b356755a6c55b96402412b10c4d3d4509d62d754 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 9 Aug 2026 23:23:42 -0700 Subject: [PATCH 125/177] [None][doc] #32 verdict: short-window penalty is a host sync chain in _prepare_inputs, not kernel pathology Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index b89c46f64b55..ddd490c9931c 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -524,6 +524,27 @@ cap-accept 模式(同一 planner 决策,提交完整块全量验证、接受端 **初步指认**:短窗桶的病灶集中在 **MoE grouped GEMM 的 tile 阶梯**——token 桶变小 → 换小 tile → 每 token 效率近乎腰斩,GEMM 时间不随 token 线性下降,形成 #32 的"常数惩罚";与 GB300 Pro frac 表的平台+悬崖楼梯(MoE 更重 → 惩罚更大)跨模型同构。另见 pin3 独有的 FMHA 变体(1.0%)待查。**定论等 v7b 的 (128,768) vs (128,384) 同批次干净对拍。** +### #32 定论:短窗惩罚 = `_prepare_inputs` 的 host 同步链,不是 kernel 病态(08-09 23:40,v7b 清洁窗) + +v7b 两轮对准命中:pin5b(122±3 请求,桶 768)/ pin3b(122±2 请求,桶 512),窗口 iter 2000-2050 全程均质满批,**同 padded_bs=128 的干净对拍**。步时(窗内 `prev_device_step_time` 均值):pin5b **51.39ms**,pin3b **58.25ms**——token −33%,步时 **+13.3%**,在 nsys 在场下复现 iter-log 矩阵的短窗惩罚。 + +**三层分解(每 50-iter 窗口,rank0):** + +| 层 | pin5b | pin3b | 判读 | +|---|---|---|---| +| GPU kernel 总时间 | 2638.7ms | 2281.9ms(0.865) | kernel 变**快**了,不是 kernel 病态 | +| GPU 空转间隙(>merge)| 162ms(~1 个大隙/步)| **844ms(~7 个/步,单隙 6-8.7ms)** | 慢的全部来自空隙 | +| `_prepare_inputs` NVTX 相位 | 1253.7ms(25.1ms/步)| **1977.4ms(39.5ms/步)** | 惩罚源头相位 | +| 其中 cudaEventSynchronize | 100 次 / 0.9ms(免费)| **300 次 / 1078.2ms(21.6ms/步)** | **6 次/步阻塞 D2H,平均 3.6ms/次** | + +**机理**:短窗(w Date: Sun, 9 Aug 2026 23:56:01 -0700 Subject: [PATCH 126/177] [None][perf] Double-buffer pinned staging: short-window prepare no longer stalls the launch pipeline The _pinned_host WAR guard synchronized on the previous step's H2D from the same buffer, serializing step N+1's prepare against step N's stream position. On short-window (w 42.2 ms; pin5 control unchanged (48.6 ms). Step time now tracks GPU kernel time (0.869 vs 0.865 kernel ratio). Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/pyexecutor/model_engine.py | 33 +++++++++++++------ 1 file changed, 23 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 00b5c062fa23..ddbbdf070bf3 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -815,9 +815,14 @@ def __init__( self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) # Pinned staging buffers for async H2D copies on the ragged path; see - # _pinned_host / _pinned_host_record. + # _pinned_host / _pinned_host_record. Two slots per key: the WAR guard + # on a slot only waits for the H2D enqueued two steps ago, so in + # steady state the wait is free. A single slot would serialize step + # N+1's prepare against step N's stream position (measured 6 blocking + # event syncs / 21.6 ms per step on short-window batches). self._pinned_host_cache = {} self._pinned_host_events = {} + self._pinned_host_active = {} # Create Encoder CUDA graph config and runner. encoder_cuda_graph_runner_config = EncoderCUDAGraphRunnerConfig( @@ -4280,19 +4285,24 @@ def _pinned_host(self, key: str, values, dtype) -> torch.Tensor: grown monotonically. """ values = list(values) - cache = self._pinned_host_cache - # WAR guard: the buffer must not be rewritten while its previous - # non_blocking H2D is still queued. Callers record the key's event via - # _pinned_host_record after enqueuing; waiting here closes the window. - evt = self._pinned_host_events.get(key) + # WAR guard: a slot must not be rewritten while its previous + # non_blocking H2D is still queued. Callers record the active slot's + # event via _pinned_host_record after enqueuing; two slots alternate, + # so this wait targets the copy from two steps ago and is free in + # steady state (a single slot would stall prepare behind the previous + # step's stream position). + slot = 1 - self._pinned_host_active.get(key, 1) + evt = self._pinned_host_events.get((key, slot)) if evt is not None: evt.synchronize() - buf = cache.get(key) + bufs = self._pinned_host_cache.setdefault(key, [None, None]) + buf = bufs[slot] if buf is None or buf.numel() < len(values) or buf.dtype != dtype: buf = torch.empty(max(len(values), 1), dtype=dtype, pin_memory=prefer_pinned()) - cache[key] = buf + bufs[slot] = buf + self._pinned_host_active[key] = slot view = buf[:len(values)] if values: view.copy_(torch.tensor(values, dtype=dtype)) @@ -4308,9 +4318,12 @@ def _pinned_host_record(self, *keys: str) -> None: if torch.cuda.is_current_stream_capturing(): return for key in keys: - evt = self._pinned_host_events.get(key) + slot = self._pinned_host_active.get(key) + if slot is None: + continue + evt = self._pinned_host_events.get((key, slot)) if evt is None: - evt = self._pinned_host_events[key] = torch.cuda.Event() + evt = self._pinned_host_events[(key, slot)] = torch.cuda.Event() evt.record() def _attach_accept_caps(self, spec_metadata, generation_requests) -> None: From 24bd23ca57a1ebd7f3371a301191e55a83815ced Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:07:10 -0700 Subject: [PATCH 127/177] [None][doc] #32 fix validated: pin3 53.7->42.2ms, theta positive again Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index ddd490c9931c..deec8bc01088 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -545,6 +545,22 @@ v7b 两轮对准命中:pin5b(122±3 请求,桶 768)/ pin3b(122±2 请求,桶 512 - **裁剪经济学翻案路径**:修掉 prep 同步链(方向:元数据形状改由 host 端窗口计划推导,或直接用 padded 均匀形状——文档注释明言 reserved slots 会在 postprocess 被 mask,按桶静态化本就是容忍的设计),短窗步时应回到 ≤ 满窗水平,届时 Pro@DEP8 复测才是 trim 产品线的真判决。#11(ragged 路径无额外同步验证)就此有了答案:**有,6 次/步,已点名**。 - 跨栈注记:GB300 Pro 的 w4 异常应为同一 prep 路径(同代码);SGLang 的 28→33ms 是另一栈,不并案。 +### #32 修复落地并验证:双缓冲 ping-pong 暂存(08-10 00:15) + +**同步点最终点名**:6 次/步 = 基线 2 次 + `_pinned_host` 的 4 个短窗 key(`accept_caps` / `ragged_verify_lens` / `gather_rows` / `gather_cols`,model_engine.py)。其 WAR 守卫在复用暂存缓冲前 `evt.synchronize()` 等**上一步**从同一缓冲发出的 H2D 被流消费——overlap 模式下即把第 N+1 步的 prepare 串行到第 N 步的流位置上。 + +**修复(b87ea28bf3,已提交)**:每 key 双槽 ping-pong——守卫改等**两步前**同槽的拷贝(稳态必已完成,等待≈0);WAR 保证原样保留(槽自身上次 H2D 完成后才重写),与 confidence relay 两深环同一设计语言,23 行。 + +**验证(122±1 请求/rank,DEP4 Flash,无 nsys,窗口 iter 2000-2050 全稳态)**: + +| 臂 | 修复前 | 修复后 | 判读 | +|---|---|---|---| +| pin5(满窗,对照)| 47.9ms(iter-log 矩阵)| **48.55 ± 0.54ms** | 无回归 ✓ | +| pin3(4 tok/req,桶 512)| 53.7ms | **42.20 ± 0.76ms(−21%)** | **短窗首次快过满窗** | +| 步时比 42.20/48.55 | — | **0.869** | 与 nsys kernel 总时间比 0.865 重合:步时回归纯 GPU 决定 | + +**含义:θ 由负转正**——w5→w3 现在省 13% 步时(此前反贴 12%)。此前所有"裁剪无收益/负收益"的吞吐结论(强裁臂 −12~−22%、host 动态 ≈0)都建立在这条被串行化拖垮的短窗路径上,需要在修复后重测;Pro@DEP8 复测(#31 修复后)仍是产品线最终判决。单测:hw_agnostic 套件 DSpark 相关 296/296 通过(`test_user_provided.py` 的 33 个失败为 llama 权重环境性失败,与本补丁无关;测试需 `--mpi=pmix` 起容器,否则 mpi4py 初始化即中止)。FRAC0.4 混合窗验证轮在跑(gather 路径覆盖)。 + ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) 1. **host 未复现**:初测 +8.3%(poetry 128/rank)在四轮中缩为 [0.0, +1.9, +1.3, +2.2]%;arena 侧 [−2.4, 0.0, −1.4, −5.6]%。host 整体 ≈ 0±3%,poetry 微正、arena 微负。host 调度行为本身跨节点完全可复现(trim 22.4-22.7%,accept 1.66-1.67,四轮逐位一致)——**是收益不存在,不是实验不稳**。 From 81367a42baa63fab94a7c3e94e9bc7ff9b0f7545 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 00:21:59 -0700 Subject: [PATCH 128/177] [None][doc] #32 closed: FRAC0.4 mixed-window arm confirms the fix (42.23ms) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index deec8bc01088..6dc08186e3d9 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -559,7 +559,9 @@ v7b 两轮对准命中:pin5b(122±3 请求,桶 768)/ pin3b(122±2 请求,桶 512 | pin3(4 tok/req,桶 512)| 53.7ms | **42.20 ± 0.76ms(−21%)** | **短窗首次快过满窗** | | 步时比 42.20/48.55 | — | **0.869** | 与 nsys kernel 总时间比 0.865 重合:步时回归纯 GPU 决定 | -**含义:θ 由负转正**——w5→w3 现在省 13% 步时(此前反贴 12%)。此前所有"裁剪无收益/负收益"的吞吐结论(强裁臂 −12~−22%、host 动态 ≈0)都建立在这条被串行化拖垮的短窗路径上,需要在修复后重测;Pro@DEP8 复测(#31 修复后)仍是产品线最终判决。单测:hw_agnostic 套件 DSpark 相关 296/296 通过(`test_user_provided.py` 的 33 个失败为 llama 权重环境性失败,与本补丁无关;测试需 `--mpi=pmix` 起容器,否则 mpi4py 初始化即中止)。FRAC0.4 混合窗验证轮在跑(gather 路径覆盖)。 +**含义:θ 由负转正**——w5→w3 现在省 13% 步时(此前反贴 12%)。此前所有"裁剪无收益/负收益"的吞吐结论(强裁臂 −12~−22%、host 动态 ≈0)都建立在这条被串行化拖垮的短窗路径上,需要在修复后重测;Pro@DEP8 复测(#31 修复后)仍是产品线最终判决。单测:hw_agnostic 套件 DSpark 相关 296/296 通过(`test_user_provided.py` 的 33 个失败为 llama 权重环境性失败,与本补丁无关;测试需 `--mpi=pmix` 起容器,否则 mpi4py 初始化即中止)。 + +**FRAC0.4 混合窗验证(00:25)**:**42.23 ± 0.78ms @ 121 请求**——与 pin3 uniform(42.20)逐位一致,真 ragged 混合窗(gather_rows/gather_cols 全覆盖)同样 −21%。修复前该臂 53.7ms。**三臂验证闭合,#32 关闭。** 下一步:修复后的三臂吞吐重测(scheduled vs off vs forced,Flash DEP4)——此前"裁剪无收益"的结论全部失效待重验。 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 5dfca9b1499f0da520d0ac3d82edde8b3e4e32ae Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 01:59:49 -0700 Subject: [PATCH 129/177] [None][doc] Post-fix retest first batch: forced trim flips to +5.3/+7.6% on poetry Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 6dc08186e3d9..da1ee2066b48 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -563,6 +563,19 @@ v7b 两轮对准命中:pin5b(122±3 请求,桶 768)/ pin3b(122±2 请求,桶 512 **FRAC0.4 混合窗验证(00:25)**:**42.23 ± 0.78ms @ 121 请求**——与 pin3 uniform(42.20)逐位一致,真 ragged 混合窗(gather_rows/gather_cols 全覆盖)同样 −21%。修复前该臂 53.7ms。**三臂验证闭合,#32 关闭。** 下一步:修复后的三臂吞吐重测(scheduled vs off vs forced,Flash DEP4)——此前"裁剪无收益"的结论全部失效待重验。 +### #33 修复后吞吐重测(08-10 02:00,首批;两节点臂序对调,poetry cell 全齐 n=5) + +强裁臂(FORCE_BUDGET_FRAC=0.4)vs 满窗基线,Flash DEP4,修复 b87ea28bf3 之后: + +| 载荷 | 全局 bs | frac04 中位 | notrim 中位 | Δ | 修复前同臂 | +|---|---|---|---|---|---| +| poetry | 512 | 21,369 | 20,292 | **+5.3%** | −12~−22% | +| poetry | 1024 | 23,540 | 21,876 | **+7.6%** | 同上 | +| arena | 512 | 25,854 | 25,863 | −0.0% | — | +| arena | 1024 | 19,704 | 19,848 | −0.7%(部分 cell 在跑) | — | + +**判读**:①裁剪的步时节省首次真实转化为吞吐收益(poetry +5~8%,与 SGLang 原生动态在 B300 的 +5.9% 同量级);②两节点交叉复现(poetry 各 cell 双节点中位差 <2%);③arena ≈ 0 提示机械 40% 裁剪在多样负载上把接受损失顶回去了——**这正是置信度调度(按请求裁)该赢过机械裁的场景**;④scheduled 臂待重采表:旧表(v2b/parity)嵌入了修复前的序列化惩罚,会误导 planner 拒绝裁剪,须在修复后代码上重采(server-sweep pin 阶梯 + --from-iter-log)。 + ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) 1. **host 未复现**:初测 +8.3%(poetry 128/rank)在四轮中缩为 [0.0, +1.9, +1.3, +2.2]%;arena 侧 [−2.4, 0.0, −1.4, −5.6]%。host 整体 ≈ 0±3%,poetry 微正、arena 微负。host 调度行为本身跨节点完全可复现(trim 22.4-22.7%,accept 1.66-1.67,四轮逐位一致)——**是收益不存在,不是实验不稳**。 From 8be4662e3a89e3edd1a613d32a8388edd9fdfce5 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 02:02:42 -0700 Subject: [PATCH 130/177] [None][doc] Retest final: arena bs1024 +1.5% (n=10); table re-collection started Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index da1ee2066b48..f3fa7bb80554 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -572,7 +572,7 @@ v7b 两轮对准命中:pin5b(122±3 请求,桶 768)/ pin3b(122±2 请求,桶 512 | poetry | 512 | 21,369 | 20,292 | **+5.3%** | −12~−22% | | poetry | 1024 | 23,540 | 21,876 | **+7.6%** | 同上 | | arena | 512 | 25,854 | 25,863 | −0.0% | — | -| arena | 1024 | 19,704 | 19,848 | −0.7%(部分 cell 在跑) | — | +| arena | 1024 | 20,197 | 19,903 | +1.5%(终值,n=10) | — | **判读**:①裁剪的步时节省首次真实转化为吞吐收益(poetry +5~8%,与 SGLang 原生动态在 B300 的 +5.9% 同量级);②两节点交叉复现(poetry 各 cell 双节点中位差 <2%);③arena ≈ 0 提示机械 40% 裁剪在多样负载上把接受损失顶回去了——**这正是置信度调度(按请求裁)该赢过机械裁的场景**;④scheduled 臂待重采表:旧表(v2b/parity)嵌入了修复前的序列化惩罚,会误导 planner 拒绝裁剪,须在修复后代码上重采(server-sweep pin 阶梯 + --from-iter-log)。 From e6723980bbff008237af06bc42525686d9a395ab Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 02:34:24 -0700 Subject: [PATCH 131/177] [None][doc] Post-fix cost table: theta monotone with real mass, matches measured pins Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index f3fa7bb80554..72844f72ac6f 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -576,6 +576,12 @@ v7b 两轮对准命中:pin5b(122±3 请求,桶 768)/ pin3b(122±2 请求,桶 512 **判读**:①裁剪的步时节省首次真实转化为吞吐收益(poetry +5~8%,与 SGLang 原生动态在 B300 的 +5.9% 同量级);②两节点交叉复现(poetry 各 cell 双节点中位差 <2%);③arena ≈ 0 提示机械 40% 裁剪在多样负载上把接受损失顶回去了——**这正是置信度调度(按请求裁)该赢过机械裁的场景**;④scheduled 臂待重采表:旧表(v2b/parity)嵌入了修复前的序列化惩罚,会误导 planner 拒绝裁剪,须在修复后代码上重采(server-sweep pin 阶梯 + --from-iter-log)。 +### 修复后成本表(08-10 02:45,--from-iter-log,今晚全部 serving 日志) + +`postfix_flash_table.json`(bs 32-128;bs 1/8 排水 cell 因网格连通性剔除):**θ(M) 单调且有真实质量**——θ(96)=3.1 → θ(512)=19.6 → θ(768)=25.9ms;固定 18.4ms + α(32/64/128)=0/1.3/3.9ms。验算与实测严丝合缝:bs128 满窗预测 48.2(实测 48.55)、bs128 L3 预测 41.9(实测 42.20)。**bs128 步时的 54% 是 token 比例项**(修复前 parity 表 θ 全程 0.9-3.1ms"惰性"、serving 真值为负)——置信度调度首次拥有真实杠杆。方法论注:旧表(v2b/parity/steep)全部作废,嵌入了 #32 的序列化惩罚。 + +scheduled 臂(新表 + tiers 2-5)已在两节点开跑,基线 notrim(今晚同节点 n=10)可直接配对。 + ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) 1. **host 未复现**:初测 +8.3%(poetry 128/rank)在四轮中缩为 [0.0, +1.9, +1.3, +2.2]%;arena 侧 [−2.4, 0.0, −1.4, −5.6]%。host 整体 ≈ 0±3%,poetry 微正、arena 微负。host 调度行为本身跨节点完全可复现(trim 22.4-22.7%,accept 1.66-1.67,四轮逐位一致)——**是收益不存在,不是实验不稳**。 From 7333a130f8f109cac4832ef3640541ae1b819516 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 03:17:05 -0700 Subject: [PATCH 132/177] [None][doc] #33 final: scheduled trim positive on all four cells (+6.0/+5.7/+1.0/+1.6%) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 72844f72ac6f..2b516e00d24e 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -582,6 +582,23 @@ v7b 两轮对准命中:pin5b(122±3 请求,桶 768)/ pin3b(122±2 请求,桶 512 scheduled 臂(新表 + tiers 2-5)已在两节点开跑,基线 notrim(今晚同节点 n=10)可直接配对。 +### #33 终局:三臂全齐,动态裁剪在我方栈上首次全面转正(08-10 03:20) + +Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10(双节点臂序对调): + +| 载荷 | 全局 bs | notrim | frac04 | **scheduled** | frac Δ | **sched Δ** | +|---|---|---|---|---|---|---| +| poetry | 512 | 20,292 | 21,369 | 21,502 | +5.3% | **+6.0%** | +| poetry | 1024 | 21,876 | 23,540 | 23,114 | +7.6% | **+5.7%** | +| arena | 512 | 25,863 | 25,854 | 26,125 | −0.0% | **+1.0%** | +| arena | 1024 | 19,903 | 20,197 | 20,222 | +1.5% | **+1.6%** | + +**判决**:①scheduled 四个 cell 全部为正,无一倒贴——修复前(host 动态 ≈0±3%、强裁 −12~−22%)的"无收益"结论整体翻案,根因就是 #32 的 WAR 守卫串行化;②在机械裁剪打平的 arena 上,置信度调度仍 +1.0~+1.6%——"按请求该裁才裁"的选择性兑现;③poetry +5.7~+6.0% 与 SGLang 博客量级(H200 +10-20%、其 B300 原生复现 +5.9%)对齐;④master arc "定位为什么没有收益"就此闭环:收益一直存在,被我们自己 4 个暂存 key 的同步链吃掉了。 + +**余留**:Pro@DEP8 复测(依赖 #31 碎片化修复)= 产品线最终判决;#19(连续预算)在新表 θ 单调后价值上升;GSM8K 精度回归(新表只改窗口选择不改接受判定,风险低,例行补测)。 + +## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) + ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) 1. **host 未复现**:初测 +8.3%(poetry 128/rank)在四轮中缩为 [0.0, +1.9, +1.3, +2.2]%;arena 侧 [−2.4, 0.0, −1.4, −5.6]%。host 整体 ≈ 0±3%,poetry 微正、arena 微负。host 调度行为本身跨节点完全可复现(trim 22.4-22.7%,accept 1.66-1.67,四轮逐位一致)——**是收益不存在,不是实验不稳**。 From 98e385a8c3b8ace08b90029b534caff3821a1216 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 05:35:20 -0700 Subject: [PATCH 133/177] [None][doc] #31 repro + elimination chain + block-granularity token-budget hypothesis Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 2b516e00d24e..68e0e6fbde12 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -597,7 +597,14 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( **余留**:Pro@DEP8 复测(依赖 #31 碎片化修复)= 产品线最终判决;#19(连续预算)在新表 θ 单调后价值上升;GSM8K 精度回归(新表只改窗口选择不改接受判定,风险低,例行补测)。 -## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) +### #31 碎片化:复现 + 排除链 + 工作假设(08-10 05:50,进行中) + +**复现(Pro@DEP8 单 8-GPU 节点,poetry 1024 全局,notrim + iter log)**:rank0 每步调度请求数众数 **60-62**(上尾 75,偶发 65 个 iter@128),远低于公平份额 128;累计分配 639/5120 ≈ 完全均衡 → **路由公平,并发准入被卡 ~60/rank**。一个 1024-burst 要 ~4480 iter 才消化完(等效并发 ~60)。 + +**排除链**:① KV 容量——`free_gpu_memory_fraction` 0.5→0.85 后最终配额 35.4→60.1GiB(定容管线无恙;注意 v2 有两条 quota 日志,第一条是估算期临时缓存,勿误读),直方图不变;KV 需求(~830 token/req × 8438 B/token)远低于配额,排除。② 客户端连接池——nsys_burst2 的 `TCPConnector(limit=600)` 曾是干扰源(600/8=75 恰是上尾!),提到 1300 后众数仍 60-62,排除(**但注意:此前所有用该客户端打 >600 全局并发的实验都被客户端限流污染,幸而 p25fix_bench 是 1100,bs1024 吞吐数据勉强幸免**)。③ ADP per-rank cap(nvbug-6133201 派生上限)——serve log 无 capping 警告,4096/6=682 不咬,排除。 + +**工作假设(判别中)**:调度器步 token 预算按 **KV block 粒度**给 gen 请求计费——Pro block=64 → 4096/64 = **64 请求/步**(众数 60-62 恰在其下,且与最初"64 行步"观察逐位吻合);Flash block=32 → 4096/32 = 128(DEP4 从未撞上,v7b 实测 122 通行)。判别:Pro `max_num_tokens` 4096→8192,若调度跳 ~128 即定案;修复方向 = gen 请求按实际步 token(6)计费,或部署侧提高 max_num_tokens。 + ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 10b6aad308121acfb961d6a4f4306156a1b65f99 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 06:33:30 -0700 Subject: [PATCH 134/177] [None][doc] #31 experiment void notice: zombie serves on 2-node allocation contaminated discriminators Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 68e0e6fbde12..85facf187611 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -605,6 +605,8 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( **工作假设(判别中)**:调度器步 token 预算按 **KV block 粒度**给 gen 请求计费——Pro block=64 → 4096/64 = **64 请求/步**(众数 60-62 恰在其下,且与最初"64 行步"观察逐位吻合);Flash block=32 → 4096/32 = 128(DEP4 从未撞上,v7b 实测 122 通行)。判别:Pro `max_num_tokens` 4096→8192,若调度跳 ~128 即定案;修复方向 = gen 请求按实际步 token(6)计费,或部署侧提高 max_num_tokens。 +**⚠ 实验作废公告(08-10 06:35)**:上述 KV-fraction 排除与 max_num_tokens=8192 判别**全部作废**——744499 是 2 节点作业,而各轮的 serve/客户端/pkill srun(-N1)随机落节点,导致**僵尸 serve 跨轮累积**(清点:bia0038 上 12 个、bia0042 上 2 个 trtllm 进程),burst 常打到旧 serve(带统计客户端坐实:1024 全 OK 而新 serve 日志 0 流量)。可信的仅有首轮复现(60-62 众数,fresh serve)。已全节点清理 + 所有 srun 钉死 `-w bia0038`,判别链从头重跑。教训入册:**多节点分配上做单节点实验,每个 srun 必须显式 -w**。 + ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 27733a7495dc6c37b7d39d711130c1e6695091e5 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 08:31:29 -0700 Subject: [PATCH 135/177] [None][doc] #31 clean elimination chain: scheduler picks ~60 of 128 active; v1-manager/python-scheduler bisection running Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 85facf187611..cfacf7fed61e 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -607,6 +607,12 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( **⚠ 实验作废公告(08-10 06:35)**:上述 KV-fraction 排除与 max_num_tokens=8192 判别**全部作废**——744499 是 2 节点作业,而各轮的 serve/客户端/pkill srun(-N1)随机落节点,导致**僵尸 serve 跨轮累积**(清点:bia0038 上 12 个、bia0042 上 2 个 trtllm 进程),burst 常打到旧 serve(带统计客户端坐实:1024 全 OK 而新 serve 日志 0 流量)。可信的仅有首轮复现(60-62 众数,fresh serve)。已全节点清理 + 所有 srun 钉死 `-w bia0038`,判别链从头重跑。教训入册:**多节点分配上做单节点实验,每个 srun 必须显式 -w**。 +**干净判决链(08-10 08:30,单节点钉死后全部重跑)**: +- 复现坐实:1024 全局真实在飞(带统计客户端 1024 全 OK 且与服务端 200 计数一致),rank0 持有满份额 128 活跃请求(`currank_total_requests = 128/1024`),ctx 在前 3 个 iter 内全部完成(准入一次到位),此后**每步只调度 60-63 个**,kv_cache_util 仅 0.014。 +- 干净排除:① KV 配额(35.4 vs 60.1GiB,直方图不变);② max_num_tokens(4096 vs 8192,不变)→ block 粒度计费假设也随之排除;③ 客户端;④ 准入速率。 +- 新观测:步时 ~290ms @ 62 行(缺健康基线,eager 嫌疑未证);SWA 固定成本 8.1GiB = max_batch_size(128) × 63.6MB/请求,定容公式按满批预留(sizing 侧无恙)。 +- 进行中的二分:v1 KV 管理器(`use_kv_cache_manager_v2: false`)vs python 调度器(`use_python_scheduler: true`)各跑一轮——若 v1 恢复 128 → 病灶在 v2 管理器的调度集成;若 python 调度器恢复 → 病灶在 C++ 容量调度器与 v2 的适配。 + ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 042700795cefeb0a1879a0fd2a1b3f5aecb81843 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 09:36:18 -0700 Subject: [PATCH 136/177] [None][doc] #31 bisection lands: cap = max_batch_size/2; maxbs=256 restores full batches and 2.7x step time Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index cfacf7fed61e..dd41439a6dd2 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -613,6 +613,18 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( - 新观测:步时 ~290ms @ 62 行(缺健康基线,eager 嫌疑未证);SWA 固定成本 8.1GiB = max_batch_size(128) × 63.6MB/请求,定容公式按满批预留(sizing 侧无恙)。 - 进行中的二分:v1 KV 管理器(`use_kv_cache_manager_v2: false`)vs python 调度器(`use_python_scheduler: true`)各跑一轮——若 v1 恢复 128 → 病灶在 v2 管理器的调度集成;若 python 调度器恢复 → 病灶在 C++ 容量调度器与 v2 的适配。 +**二分收网(08-10 09:40)——上限 ≈ max_batch_size/2,ماxbs=256 即刻解锁**: +| 判别 | 结果 | 结论 | +|---|---|---| +| v1 KV 管理器 | 仍 ~60-63 | v2 管理器免罪 | +| python 调度器 | 仍 ~60-63 | C++/python 调度器共同层 | +| 关 overlap | 仍 ~62-65 | overlap 免罪 | +| **max_batch_size 128→256** | **调度 115-128(满份额),步时 290→108ms(2.7×)** | **上限 ≈ maxbs/2 坐实** | + +- **部署级解法(立即可用)**:Pro@DEP8 设 `max_batch_size=256`(或 2× 目标并发/rank)→ 128 行满批 + 步时 2.7× 改善;Pro 复测(trim 产品线最终判决)就此解锁。 +- **代码级根因(待定位)**:两个调度器共享的容量/槽记账把每个 DSpark 请求按 ~2 份计(嫌疑:spec 序列槽为 draft 复占、或 GNE 容量对 spec 请求的双倍预留);`get_max_num_sequences`=pp×bs=128 本身无恙。290ms 步时 = 减半批的伴生病(62 行时慢 2.7×,机理与槽争抢/eager 待nsys确认)。 +- 观测注记:cap 尾部到 67-75(非硬 64)——完成中请求可能释放半槽,支持"槽复占"方向。 + ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 40bd8f6806ad4308219cb53b36b0a53079193d62 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 10:55:24 -0700 Subject: [PATCH 137/177] [None][doc] Post-fix Pro table: monotone theta, w4 anomaly and staircase gone; 56% of step trimmable Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index dd41439a6dd2..f51c015207e9 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -625,6 +625,8 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( - **代码级根因(待定位)**:两个调度器共享的容量/槽记账把每个 DSpark 请求按 ~2 份计(嫌疑:spec 序列槽为 draft 复占、或 GNE 容量对 spec 请求的双倍预留);`get_max_num_sequences`=pp×bs=128 本身无恙。290ms 步时 = 减半批的伴生病(62 行时慢 2.7×,机理与槽争抢/eager 待nsys确认)。 - 观测注记:cap 尾部到 67-75(非硬 64)——完成中请求可能释放半槽,支持"槽复占"方向。 +**修复后 Pro 表(08-10 10:55,maxbs=256 满批 + #32 修复,--from-iter-log,4 pin × 3 burst 双节点)**:`postfix_pro_table.json`,bs128 实测 L2/3/4/5 = **80.8/90.3/102.7/110.3ms**(n=47-81/cell)——**完美单调,w4 异常消失,GB300"平台+悬崖"楼梯消失**(确证均为 #32 prep 串行化 + #31 半批的伪影);固定 39.3ms + α(128)=9.1ms,θ(768)=61.9ms = **步时 56% 可裁**(Flash 为 54%,但 Pro 绝对杠杆 2.4×)。w5→w2 省 26.7% 步时。下一步:Pro@DEP8 三臂终局(scheduled-新表 vs notrim,maxbs=256),节点已排。 + ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From ad169e92e5a261f26d8bf67b17f7c82be6acb926 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 13:47:46 -0700 Subject: [PATCH 138/177] [None][doc] FINAL VERDICT: Pro@DEP8 confidence scheduling +5.9~+24.7% across all cells Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index f51c015207e9..6fc19a0c42d6 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -597,6 +597,21 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( **余留**:Pro@DEP8 复测(依赖 #31 碎片化修复)= 产品线最终判决;#19(连续预算)在新表 θ 单调后价值上升;GSM8K 精度回归(新表只改窗口选择不改接受判定,风险低,例行补测)。 +### 产品线终局判决:Pro@DEP8 置信度调度全面转正(08-10 13:50) + +Flash 三臂翻案后,Pro@DEP8 终局(maxbs=256 解锁满批 + `postfix_pro_table.json` 新表,scheduled vs notrim,双节点臂序对调,每 cell n=10): + +| 载荷 | 全局 bs | notrim | **scheduled** | Δ | +|---|---|---|---|---| +| poetry | 512 | 12,351 | 13,568 | **+9.9%** | +| poetry | 1024 | 15,586 | 19,431 | **+24.7%** | +| arena | 512 | 15,362 | 16,400 | **+6.8%** | +| arena | 1024 | 20,887 | 22,124 | **+5.9%** | + +**判决**:①四 cell 全正,最难的 poetry@1024 达 **+24.7%**,超过 SGLang blog 量级(H200 +10-20%);②master arc"为什么没有收益"终章——收益从未缺席,被我们自己的三个问题吃掉:#32 prep 同步链(每步 ~14ms)、#31 半批上限(cap=maxbs/2,兼致 2.7× 步时惩罚)、以及嵌入这两个伪影的旧成本表;③修复组合 = b87ea28bf3(双缓冲)+ maxbs=2× 目标并发(部署参数)+ 修复后重采表。 +**方法论注记**:745396 首次 S 臂遇 nvlink peer 显存错误(warmup 期),同配置两节点随后均干净复跑 → 判硬件抖动,非确定性 bug;若复发再立案。 +**余留**:#31 代码级根因(spec 请求 2× 槽记账,已有部署级规避)、GSM8K 精度例行回归、#19 连续预算(θ 单调后价值上升)、反馈 SGLang(2026.2.1 nsys、192 行图上限)。 + ### #31 碎片化:复现 + 排除链 + 工作假设(08-10 05:50,进行中) **复现(Pro@DEP8 单 8-GPU 节点,poetry 1024 全局,notrim + iter log)**:rank0 每步调度请求数众数 **60-62**(上尾 75,偶发 65 个 iter@128),远低于公平份额 128;累计分配 639/5120 ≈ 完全均衡 → **路由公平,并发准入被卡 ~60/rank**。一个 1024-burst 要 ~4480 iter 才消化完(等效并发 ~60)。 From 722852c6ae36752fa1af27b1ff87a9bfc4b9819c Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 19:57:45 -0700 Subject: [PATCH 139/177] [None][doc] #34 breakthrough: cap is not spec-specific; V2 scheduler deadlock at ~61 gen allocs, quota-insensitive Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 6fc19a0c42d6..9db7320a4794 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -597,6 +597,13 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( **余留**:Pro@DEP8 复测(依赖 #31 碎片化修复)= 产品线最终判决;#19(连续预算)在新表 θ 单调后价值上升;GSM8K 精度回归(新表只改窗口选择不改接受判定,风险低,例行补测)。 +### #34 破题:上限不是 spec 专属,是 V2 调度器/KV-v2 的按请求分配在 ~maxbs/2 处耗尽(08-10 20:10) + +**no-spec 对照的意外收获**:同配置去掉 DSpark(maxbs=128,poetry 1024 burst)——服务器**直接崩**:`scheduler/scheduler_v2.py:439 RuntimeError: V2 scheduler deadlock: 61 generation request(s) active but none could be scheduled or evicted. KV cache pool is likely exhausted…`(1018/1024 请求收到 400)。三个信息: +1. **上限与投机解码无关**(no-spec 也卡 61)——"spec 槽 2× 记账"假设推翻;有 spec 时优雅降级到 ~60,无 spec 时死锁崩溃(本身是个产品级 bug:普通高并发 serving 会炸)。 +2. **`scheduler_v2.py` 是 use_python_scheduler 两档共走的路径**——解释当晚二分"python/C++ 调度器都一样"。 +3. 调度门 = `kv_cache_manager.try_allocate_generation(req)` → `kv_cache.resize(...)` 在第 ~62 个 gen 请求失败,而 kv_cache_util 仅 0.014、字节配额 35→60GiB 不敏感——**耗尽的是某个按 maxbs 定容、每 gen 请求消耗 ~2 份的结构**(嫌疑:SWA/压缩器窗口池的 scratch-reuse 定容与运行时逐请求分配不一致;首波满批因 ctx 阶段不走 gen 分配)。深挖 agent 与 no-spec@maxbs256 象限实验进行中。 + ### 产品线终局判决:Pro@DEP8 置信度调度全面转正(08-10 13:50) Flash 三臂翻案后,Pro@DEP8 终局(maxbs=256 解锁满批 + `postfix_pro_table.json` 新表,scheduled vs notrim,双节点臂序对调,每 cell n=10): From efb283d0a78d17598f252df9b4fd76ec3d8b0789 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:15:10 -0700 Subject: [PATCH 140/177] [None][doc] #34 root cause: windowed pool sized 1 block/request, runtime locks 2 (window 133 vs block 128); cap=0.51*maxbs arithmetic closed Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 9db7320a4794..bdc5894b15cf 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -604,6 +604,12 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( 2. **`scheduler_v2.py` 是 use_python_scheduler 两档共走的路径**——解释当晚二分"python/C++ 调度器都一样"。 3. 调度门 = `kv_cache_manager.try_allocate_generation(req)` → `kv_cache.resize(...)` 在第 ~62 个 gen 请求失败,而 kv_cache_util 仅 0.014、字节配额 35→60GiB 不敏感——**耗尽的是某个按 maxbs 定容、每 gen 请求消耗 ~2 份的结构**(嫌疑:SWA/压缩器窗口池的 scratch-reuse 定容与运行时逐请求分配不一致;首波满批因 ctx 阶段不走 gen 分配)。深挖 agent 与 no-spec@maxbs256 象限实验进行中。 +### #34 定案:窗口池槽数按"每请求 1 块"定容,运行时每请求锁 2 块(08-10 20:50) + +**机制(算术闭合)**:DSv4 的窗口池组(SWA+COMPRESSOR_KV+COMPRESSOR_SCORE,窗口=128+max_draft_len=133 token)是固定槽分配器,槽数被钉在下限 `min_slots≈1.05×maxbs`——①比例定容用 `typical_seq_len=max_seq_len`(serve 推断 1M)把配额喂给全注意力池,窗口池饿到下限(kv_cache_manager_v2.py:1775-1779);②下限按 history=0 的"刚起步 decode"建模 → 每请求 1 块(kv_cache_manager_v2.py:1803-1812 → _storage_manager.py:909-1001)。运行时窗口 133>块 128 → 每 gen 请求锁 2(偶 3)块(_life_cycle_registry.py:38-51)。**容量 = 1.05·maxbs/2.07 ≈ 0.51·maxbs**:maxbs=128 预测 63-65/实测 61-63;maxbs=256 预测 131-135/实测 131-136。配额不敏感(下限无配额项)、kv_util 0.014 是 token 口径(窗口池实为 100% 锁死)、首波满批(短 prompt 1 块,h>128 才要第二块)、spec 降级 vs no-spec 死锁(接受方差去相位 vs +1/步全队锁相集体撞边界)——全部观测闭合。字节估算器还漏了 +draft 窗口放大(cache_manager.py:1316-1326),但字节不是绑定资源。 + +**修复**:A(首选)min-slots 按稳态 decode(2-3 块/请求)建模;B 估算器补 draft(须配 A);C `typical_seq_len` 不默认 max_seq_len / 部署侧显式设 max_seq_len 解钉(判别实验在跑)。上游三报:下限建模、估算器漏项、no-spec 高并发死锁崩溃(报错信息误导)。附:python-scheduler 旗标在 v2 下是 no-op(_util.py:2618),DSv4 强制 v2 管理器——当晚两条二分腿均为无效对照,已修订。 + ### 产品线终局判决:Pro@DEP8 置信度调度全面转正(08-10 13:50) Flash 三臂翻案后,Pro@DEP8 终局(maxbs=256 解锁满批 + `postfix_pro_table.json` 新表,scheduled vs notrim,双节点臂序对调,每 cell n=10): From 6e65c54cfcb3de3fb48ce0b2962d03d9ed49b4fe Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:22:56 -0700 Subject: [PATCH 141/177] [None][test] DSpark hw-agnostic unit tests: planner parity, cost interpolation, iter-log sampling, guards Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../dspark_make_steep_sps_table.py | 222 +++++++ .../test_dspark_budget_parity_sglang.py | 232 +++++++ .../test_dspark_cost_interpolation.py | 61 ++ .../test_dspark_iter_log_samples.py | 159 +++++ .../test_dspark_reconcile_guards.py | 99 +++ .../test_dspark_snapshot_instruments.py | 110 ++++ .../hw_agnostic/test_dspark_sps_profiler.py | 614 ++++++++++++++++++ .../test_dspark_sts_collection_guards.py | 135 ++++ .../test_dspark_sts_ring_pairing.py | 168 +++++ .../test_dspark_tier_aligned_budget.py | 96 +++ 10 files changed, 1896 insertions(+) create mode 100644 tests/microbenchmarks/dspark_make_steep_sps_table.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py diff --git a/tests/microbenchmarks/dspark_make_steep_sps_table.py b/tests/microbenchmarks/dspark_make_steep_sps_table.py new file mode 100644 index 000000000000..b22260c5b90e --- /dev/null +++ b/tests/microbenchmarks/dspark_make_steep_sps_table.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Emit an SPS cost table steep enough that the planner chooses to trim. + +WHAT THIS IS FOR, AND WHAT IT IS NOT FOR +---------------------------------------- +It exercises the ragged path end to end *through the planner*: with a steep +``theta(M)`` the confidence scheduler decides to trim, and its budget then +flows through the same confidence-ordered top-k that ships, which hands +different windows to requests with different survival. That is the only way to +see per-request verify lengths on a checkpoint whose acceptance is too high for +trimming to pay. + +**No throughput or latency number taken with this table means anything.** The +planner is optimising against a fictional cost curve, so it will trim where a +real deployment should not. Use it to answer "does the ragged path correctly +support different verify lengths per request", never "is ragged faster". The +output carries ``"SYNTHETIC": true`` so a run cannot quietly present it as a +measurement -- ``run_dspark_throughput.sh`` refuses to treat such a table as +evidence. + +WHY A TABLE RATHER THAN AN OVERRIDE +----------------------------------- +This replaces ``TLLM_DSPARK_FORCE_VERIFY_LENS``, which assigned windows by +batch position. Position is orthogonal to confidence, so that knob reproduced +the ragged *shape* while bypassing the policy entirely -- it could not tell a +correct assignment from an arbitrary one, and it had to sit in production code +ahead of the cost-table gate to be reachable. A table drives the real path. + +HOW STEEPNESS IS CHOSEN +----------------------- +The scheduler maximises expected accepted tokens per second, roughly + + (bs + sum of the top-K survival probabilities) / T(bs, K) + +With per-token acceptance around 0.9 each extra verify token adds about 0.9 to +the numerator, so trimming only wins once ``theta`` grows faster than that. The +default ``--exponent 2.0`` makes theta quadratic in the token count, which is +comfortably steeper than the measured curve (near-linear) and forces trimming +across the whole batch-size ladder. + + # derive from a real table, keeping its grid + python3 tests/microbenchmarks/dspark_make_steep_sps_table.py \ + --from sps_real_final.json --out steep.json + + # or synthesise a grid outright + python3 tests/microbenchmarks/dspark_make_steep_sps_table.py --out steep.json +""" + +import argparse +import json +import sys + + +def _default_grid(max_batch_size: int, width: int): + """Token totals a ragged batch can land on: bs * (verify_len + 1). + + The ladder always includes ``max_batch_size`` itself. A grid that stops + short of the deployment's widest step is worse than useless: theta clamps + at the last bucket, so the full window looks no more expensive than a + trimmed one and the planner keeps it -- a table that silently cannot + induce trimming, which is the failure this tool exists to prevent. + """ + ladder = sorted({b for b in (1, 4, 8, 16, 32, 64, 128, 256, 512) + if b <= max_batch_size} | {max_batch_size}) + tokens = sorted({b * w for b in ladder for w in range(2, width + 1)}) + return ladder, tokens + + + +def _preview_with_real_planner(*, path, batch_sizes, accept, max_verify_len): + """What budget_argmax_over_uniform_lens actually picks for this table. + + Returns True if any previewed batch size trims, False if none do, or None + when the planner cannot be imported (outside the container), so the caller + can distinguish "checked and fine" from "not checked". + """ + try: + import numpy as np + from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, budget_argmax_over_uniform_lens) + except Exception: + return None + + table = SpsCostTable.from_json_file(path) if hasattr( + SpsCostTable, "from_json_file") else None + if table is None: + with open(path, encoding="utf-8") as handle: + table = SpsCostTable(**{ + k: v for k, v in json.load(handle).items() + if not k.startswith("_") and k != "SYNTHETIC" + }) + + tiers = sorted({1, max(1, (max_verify_len + 1) // 2), max_verify_len}) + print() + print(f"planner's actual choice at acceptance {accept} " + f"(L = verify length, {max_verify_len} = no trim), tiers {tiers}:") + any_trim = False + for bs in [int(b) for b in str(batch_sizes).split(",") if str(b).strip()]: + # survival[r, k] = P(first k+1 drafted tokens all accepted) + survival = np.array( + [[accept**(k + 1) for k in range(max_verify_len)]] * bs, + dtype=np.float64) + chosen = budget_argmax_over_uniform_lens( + survival=survival, num_gen_requests=bs, cost_table=table, + allowed_lens=tiers, min_verify_len=1) + trims = int(chosen) < max_verify_len + any_trim = any_trim or trims + print(f" bs={bs:<5} -> L={chosen}" + f"{' TRIM' if trims else ' (full window)'}") + return any_trim + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + parser.add_argument("--out", required=True) + parser.add_argument("--from", dest="source", default=None, + help="real table to take the grid from; its timings " + "are replaced, only token_counts/batch_sizes are " + "reused so the shape matches the deployment") + parser.add_argument("--exponent", type=float, default=2.0, + help="theta(M) ~ M**exponent; >1 makes extra verify " + "tokens superlinearly expensive") + parser.add_argument("--peak-ms", type=float, default=400.0, + help="theta at the largest token count") + parser.add_argument("--fixed-overhead-ms", type=float, default=0.05, + help="the non-trimmable floor. MUST be small next to " + "theta in the operating range, or trimming cannot " + "pay however steep theta is -- see --preview-bs") + parser.add_argument("--preview-bs", default="8,32,128", + help="batch sizes to predict the planner's choice for, " + "so a table that cannot induce trimming is caught " + "here instead of 20 minutes into a run") + parser.add_argument("--preview-accept", type=float, default=0.9, + help="per-token acceptance to predict against; 0.9 is " + "what DeepSeek-V4-Pro-DSpark measures") + parser.add_argument("--max-batch-size", type=int, default=128) + parser.add_argument("--max-verify-len", type=int, default=5) + args = parser.parse_args() + + if args.exponent <= 1.0: + parser.error( + f"--exponent {args.exponent} is not steeper than linear, so the " + f"planner will keep verifying the full block and the ragged path " + f"will stay dark -- which is the exact failure this tool exists " + f"to avoid") + + needed = args.max_batch_size * (args.max_verify_len + 1) + if args.source: + with open(args.source, encoding="utf-8") as handle: + src = json.load(handle) + tokens = [int(t) for t in src["token_counts"]] + batch_sizes = [int(b) for b in src["batch_sizes"]] + if max(tokens) < needed: + parser.error( + f"--from {args.source} tops out at {max(tokens)} tokens but " + f"max_batch_size {args.max_batch_size} reaches {needed} at the " + f"full window. theta would clamp there, making the full window " + f"look as cheap as a trimmed one, so the planner would never " + f"trim. Drop --from to synthesise a grid that covers it.") + else: + batch_sizes, tokens = _default_grid(args.max_batch_size, + args.max_verify_len + 1) + + tokens = sorted({t for t in tokens if t > 0}) + peak = float(max(tokens)) + theta = [args.peak_ms * (t / peak)**args.exponent for t in tokens] + + table = { + "token_counts": tokens, + "step_time_ms": theta, + "fixed_overhead_ms": args.fixed_overhead_ms, + "batch_sizes": batch_sizes, + "batch_overhead_ms": [0.0] * len(batch_sizes), + "SYNTHETIC": True, + "_meta": { + "purpose": "force the planner to trim so the ragged path runs; " + "NOT a measurement of anything", + "exponent": args.exponent, + "derived_from": args.source, + }, + } + with open(args.out, "w", encoding="utf-8") as handle: + json.dump(table, handle, indent=2) + handle.write("\n") + + # Ask the REAL planner what it will do, rather than re-deriving its argmax. + # + # A hand-written approximation of the objective was wrong twice: it models + # the cost as fixed + theta(M), while budget_argmax_over_uniform_lens calls + # cost_table.step_time(total_verify_tokens(bs, length), num_gen_requests), + # which interpolates the table and adds the per-batch term. The + # approximation predicted "trim at every batch size" for a table under + # which the planner kept the full window through a 20-minute run. A preview + # that does not agree with the thing it previews is worse than none. + any_trim = _preview_with_real_planner( + path=args.out, batch_sizes=args.preview_bs, + accept=args.preview_accept, max_verify_len=args.max_verify_len) + if any_trim is None: # planner not importable here + print("preview skipped: run this inside the TensorRT-LLM container to " + "have the real planner check the table before you spend a run " + "on it", file=sys.stderr) + any_trim = True + print(f" bs={bs:<5} -> K={best_k}" + f"{' TRIM' if trims else ' (full window)'}") + if not any_trim: + print() + print("WARNING: this table induces no trimming at any previewed batch " + "size, so the ragged path will stay dark. Lower " + "--fixed-overhead-ms or raise --exponent.", file=sys.stderr) + + print() + print(f"wrote {args.out}: {len(tokens)} token buckets " + f"{tokens[0]}..{tokens[-1]}, theta {theta[0]:.3f}..{theta[-1]:.1f} ms " + f"(exponent {args.exponent}), SYNTHETIC=true") + print("Use only to check that per-request verify lengths work. Any " + "throughput taken with this table is meaningless.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py new file mode 100644 index 000000000000..ac989a423428 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py @@ -0,0 +1,232 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Decision-level parity with the SGLang budget planner this was ported from. + +A throughput gain alone does not validate the port: the wrong economics can +still move the needle in the right direction. This pins the DECISION function +against a line-for-line reference copy of SGLang's +``compute_verify_token_budget`` + ``_additive_step_time_tensor`` +(python/sglang/srt/speculative/dspark_components/dspark_planner.py), so a +divergence in either port shows up as an argmax mismatch here rather than as +an unexplained fleet-level regression. + +The two planners parameterize the same optimization differently: + +- SGLang admits EVERY position (including position 0) as a budget candidate + and its floor is one bonus token per request: total M = bs + n. +- TensorRT-LLM grants positions below ``min_verify_len`` for free (their + survival rides in the tau base) and only positions at or beyond it are + schedulable: total M = bs*(min_verify_len+1) + m. + +For ``min_verify_len=1`` these describe the same curve shifted by ``bs``: +survival is non-increasing along a row, so the top-(bs+m) SGLang candidates +are exactly the bs position-0 entries plus the top-m deeper entries, giving +``tau_sgl(bs+m) == tau_trt(m)`` and ``M_sgl(bs+m) == M_trt(m)``. The argmaxes +must therefore agree up to that shift whenever SGLang admits at least the +full position-0 set. Known, deliberate divergences are pinned separately: +SGLang filters candidates by ``survival_eps`` inside the budget function +(TensorRT-LLM applies eps at allocation time instead), and SGLang can choose +sub-floor budgets (windows of zero draft positions) that TensorRT-LLM's +``min_verify_len=1`` floor excludes. +""" + +import numpy as np +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, compute_verify_token_budget) + +# -------------------------------------------------------------------------- +# Reference implementation, copied from SGLang dspark_components/dspark_planner +# (additive-table path) with only cosmetic renames. Do not "fix" this copy: +# its job is to be what SGLang runs. +# -------------------------------------------------------------------------- + + +def _sgl_interp_clamped(xs, ys, x: float) -> float: + xs = torch.tensor(xs, dtype=torch.float64) + ys = torch.tensor(ys, dtype=torch.float64) + x_t = torch.tensor(float(x), dtype=torch.float64).clamp_(xs[0], xs[-1]) + hi = torch.bucketize(x_t, xs, right=True).clamp_(1, xs.numel() - 1) + lo = hi - 1 + span = (xs[hi] - xs[lo]).clamp_(min=1e-9) + frac = (x_t - xs[lo]) / span + return float(ys[lo] + frac * (ys[hi] - ys[lo])) + + +def _sgl_additive_step_time(table, num_requests: int, num_budgets: int): + floor = table["bias_seconds"] + _sgl_interp_clamped( + table["bs_probes"], table["alpha_seconds"], float(num_requests)) + m_probes = torch.tensor(table["m_probes"], dtype=torch.float64) + theta_vals = torch.tensor(table["theta_seconds"], dtype=torch.float64) + m = (num_requests + torch.arange(num_budgets, dtype=torch.float64)).clamp_( + min=float(table["m_probes"][0]), max=float(table["m_probes"][-1])) + hi = torch.bucketize(m, m_probes, right=True).clamp_(1, m_probes.numel() - 1) + lo = hi - 1 + span = (m_probes[hi] - m_probes[lo]).clamp_(min=1e-9) + frac = (m - m_probes[lo]) / span + theta_at_m = theta_vals[lo] + frac * (theta_vals[hi] - theta_vals[lo]) + return floor + theta_at_m + + +def _sgl_compute_verify_token_budget(*, history_survival_probs, table, + max_verify_len, survival_eps): + num_requests = history_survival_probs.shape[0] + candidates = history_survival_probs[:, :max_verify_len].flatten() + candidates = candidates[candidates >= survival_eps].to(torch.float64) + candidates_sorted = torch.sort(candidates, descending=True).values + prefix_sum = torch.cumsum(candidates_sorted, dim=0) + tau_star = num_requests + torch.cat( + [torch.zeros(1, dtype=torch.float64), prefix_sum]) + step_time = _sgl_additive_step_time(table, int(num_requests), + int(tau_star.numel())) + theta = tau_star / step_time + return int(torch.argmax(theta)) + + +# -------------------------------------------------------------------------- +# One certified-table shape shared by both sides. TensorRT-LLM stores ms, +# SGLang stores seconds; the argmax is scale-invariant but the tables are +# built with the honest factor anyway. +# -------------------------------------------------------------------------- + +M_PROBES = (64, 96, 128, 192, 384, 512, 768, 1536) +THETA_MS = (9.68, 11.08, 11.95, 13.95, 22.09, 23.85, 35.61, 105.64) +BIAS_MS = 25.244443 +BS_PROBES = (32, 64, 128, 256) +ALPHA_MS = (0.0, 2.177477, 11.462212, 19.342661) +BLOCK = 5 + +TRT_TABLE = SpsCostTable(token_counts=M_PROBES, step_time_ms=THETA_MS, + fixed_overhead_ms=BIAS_MS, batch_sizes=BS_PROBES, + batch_overhead_ms=ALPHA_MS) +SGL_TABLE = { + "m_probes": M_PROBES, + "theta_seconds": tuple(v / 1e3 for v in THETA_MS), + "bias_seconds": BIAS_MS / 1e3, + "bs_probes": BS_PROBES, + "alpha_seconds": tuple(v / 1e3 for v in ALPHA_MS), +} + + +def _inversion_free_survival(rng, bs): + """Survivals where every position-0 beats every deeper candidate. + + The exact-parity theorem needs this: SGLang admits position 0 as a budget + candidate while TensorRT-LLM grants it in the floor, so the two describe + the same curve (shifted by ``bs``) only when the top-(bs+m) SGLang set is + exactly "all of position 0, plus the top-m deeper entries". A weak row's + position-0 falling below a strong row's position-1 breaks that -- see + test_starving_weak_rows_is_a_known_gap for what SGLang does with it. + """ + conf0 = rng.uniform(0.85, 0.99, size=(bs, 1)) + deeper = rng.uniform(0.25, 0.80, size=(bs, BLOCK - 1)) + return np.cumprod(np.concatenate([conf0, deeper], axis=1), axis=1) + + +@pytest.mark.parametrize("bs", [1, 3, 8, 64, 128, 252, 256]) +def test_budget_argmax_matches_sglang_up_to_the_floor_shift(bs): + """Same survival, same table, eps off, no inversions: shift-exact parity.""" + rng = np.random.default_rng(17 + bs) + for _ in range(5): + surv = _inversion_free_survival(rng, bs) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.0) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + if sgl_n >= bs: + assert trt_m == sgl_n - bs, ( + f"bs={bs}: SGLang admits {sgl_n} candidates " + f"(= floor {bs} + {sgl_n - bs}), TRT-LLM budget {trt_m}") + else: + # SGLang went below the one-draft floor; the closest budget + # TensorRT-LLM can express under min_verify_len=1 is zero. + assert trt_m == 0 + + +def test_live_pooled_survivals_park_on_the_breakpoint_together(): + """The campaign's operating point, both consumers, both argmax modes. + + With interpolated segments the continuous optimum parks where the theta + slope jumps: buying position-2 survivals (0.502/token) beats the + 512..768 slope (0.046 ms/token) but not the 768..1536 one (0.091), so + both planners stop at exactly M = 768 -- TensorRT-LLM at budget + 768 - 2*252 = 264, SGLang at 768 - 252 = 516. Tier alignment then rounds + down to the rung-2 budget the executor can capture. + """ + bs = 252 + surv = np.tile([0.771, 0.693, 0.502, 0.360, 0.315], (bs, 1)) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.0) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + assert trt_m == 768 - 2 * bs + assert sgl_n - bs == trt_m + tiered = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1, allowed_lens=[1, 2, 5]) + assert tiered == bs * 1 # rung-2: one scheduled position past the floor + + +def test_starving_weak_rows_is_a_known_gap(): + """SGLang can verify NOTHING for a hopeless request; TRT-LLM cannot. + + SGLang's budget admits position-0 candidates individually, so a row whose + draft is dead on arrival can be starved to a zero-length window and its + would-be tokens spent on a strong row's deep positions. TensorRT-LLM's + ``min_verify_len=1`` floor buys position 0 for every request + unconditionally. On a batch that mixes both kinds, SGLang's achievable + Theta is therefore strictly better -- pinned here so a future + min_verify_len=0 port is a deliberate decision with a test to flip. + """ + # Half the batch is strong all the way down, half is dead on arrival. + # The batch must be big enough that its token range sits on a genuinely + # rising part of the theta curve: below the first breakpoint the clamp + # makes cost flat, every candidate is free, and both planners trivially + # buy everything at the same theta -- the two-row version of this test + # sat in that regime and proved nothing. + bs = 64 + strong = np.tile([0.99, 0.98, 0.97, 0.96, 0.95], (bs // 2, 1)) + dead = np.tile([0.01, 0.008, 0.006, 0.004, 0.002], (bs // 2, 1)) + surv = np.concatenate([strong, dead], axis=0) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.0) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + + def theta_at(tau, tokens): + return tau / float(TRT_TABLE.step_times(np.asarray([tokens]), bs)[0]) + + cand = np.sort(surv.reshape(-1))[::-1] + sgl_theta = theta_at(bs + cand[:sgl_n].sum(), bs + sgl_n) + deeper = np.sort(surv[:, 1:].reshape(-1))[::-1] + trt_theta = theta_at(bs + surv[:, 0].sum() + deeper[:trt_m].sum(), + 2 * bs + trt_m) + assert sgl_theta > trt_theta + + +def test_eps_filtering_is_a_known_divergence(): + """SGLang drops sub-eps candidates inside the budget; TRT-LLM does not. + + TensorRT-LLM applies survival_eps at allocation time + (schedule_verify_lens_topk's eligibility mask), so its BUDGET can count a + candidate the allocator then refuses. This pins the divergence so a + future unification is a deliberate change, not an accident: with a heavy + sub-eps tail the SGLang budget must not exceed TRT-LLM's shifted one. + """ + bs = 16 + surv = np.tile([0.9, 0.4, 0.008, 0.004, 0.002], (bs, 1)) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.01) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + assert sgl_n - bs <= trt_m diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py new file mode 100644 index 000000000000..acbf6eba383f --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py @@ -0,0 +1,61 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Cost lookup must interpolate between breakpoints, not floor onto them. + +The floor lookup this replaces silently disabled the whole feature: with +measured points at 768 and 1536, every total in 769..1535 was billed at the +768 price, so a full block at bs=252 (1512 tokens, true cost ~148 ms) was +priced 72.3 ms. The planner's cost ratio T(6bs)/T(3bs) collapsed from a true +~2.0 to 1.19 -- below the GSM8K survival ratio tau5/tau2 ~ 1.48 -- so the +argmax bought the full block on ~95% of decisions, the bucket landed on +padded_bs*6 on every step of every run, and trimming never happened. The only +exception ever observed was bs=256 exactly, where the tier totals land on +breakpoints and floor happens to be exact. + +The consumer this was ported from (SGLang ``_additive_step_time_tensor``) +interpolates both theta(M) and alpha(bs); this port must match. +""" + +import pytest + +from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable + +# The certified GB300 table's shape, abbreviated to the two breakpoints whose +# gap hid the bug. +TABLE = SpsCostTable(token_counts=(512, 768, 1536), + step_time_ms=(23.85, 35.61, 105.64), + fixed_overhead_ms=25.244, + batch_sizes=(128, 256), + batch_overhead_ms=(11.462, 19.343)) + + +def test_theta_interpolates_between_breakpoints(): + """Exact on breakpoints, linear between them. + + The regression case: 1512 tokens must cost ~the 1536 price, not the 768 + one. + """ + assert TABLE.step_time(768, 256) == pytest.approx(25.244 + 19.343 + 35.61) + assert TABLE.step_time(1536, 256) == pytest.approx(25.244 + 19.343 + 105.64) + got = TABLE.step_time(1512, 252) + theta = 35.61 + (105.64 - 35.61) * (1512 - 768) / (1536 - 768) + alpha = 11.462 + (19.343 - 11.462) * (252 - 128) / (256 - 128) + assert got == pytest.approx(25.244 + alpha + theta) + # The floor price it used to return -- and must never return again. + assert got > 130.0 + + +def test_alpha_interpolates_and_clamps(): + assert TABLE.batch_overhead(128) == pytest.approx(11.462) + assert TABLE.batch_overhead(192) == pytest.approx((11.462 + 19.343) / 2) + assert TABLE.batch_overhead(64) == pytest.approx(11.462) # clamp low + assert TABLE.batch_overhead(512) == pytest.approx(19.343) # clamp high + + +def test_theta_clamps_outside_the_measured_range(): + """Theta clamps to the end values; a table without a batch axis adds no alpha.""" + plain = SpsCostTable(token_counts=(512, 768, 1536), + step_time_ms=(23.85, 35.61, 105.64), + fixed_overhead_ms=25.244) + assert plain.step_time(100, 0) == pytest.approx(25.244 + 23.85) + assert plain.step_time(4096, 0) == pytest.approx(25.244 + 105.64) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py new file mode 100644 index 000000000000..e2b9f049807b --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py @@ -0,0 +1,159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Reading cost-table samples out of production iteration logs. + +The in-process sweep measures an engine the profiler built itself, and when +that engine's configuration drifts from the deployment the table describes a +machine nobody is running -- a swept (bs=256, M=1536) cell read 150.2 ms +against 118.8 ms measured live at the same shape, with nothing in the +fingerprint to record the difference. ``print_iter_log`` already emits the row +count, the token total and the host step time for every step of the real +deployment, so a table can be fitted from traffic that actually ran. + +What must hold: the token total identifies the verify length exactly (no +clustering, no guessing), steps that cannot be filed honestly are dropped +rather than mislabelled, and the first step of each executor instance -- whose +host time contains the wait for the first request, measured at 62 s live -- is +never treated as a measurement. +""" + +import importlib.util +import pathlib +import sys +import tempfile + +_SPEC = importlib.util.spec_from_file_location( + "dspark_sps_profiler", + pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / + "dspark_sps_profiler.py") +_PROFILER = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _PROFILER # @dataclass needs the module registered +_SPEC.loader.exec_module(_PROFILER) + +samples_from_iter_log = _PROFILER.samples_from_iter_log +resolve_padded_shape = _PROFILER.resolve_padded_shape + +# A deployment's captured ladder, abbreviated. Padding rounds a step's rows up +# to the smallest entry at or above them, which is what turns a padded token +# total into exactly one verify length. +LADDER = [1, 2, 4, 8, 16, 32, 64, 128, 192, 256] + + +def _line(iteration, bs, step_ms, gen_tokens, rank="0"): + return (f"{rank}: [08/04/2026-19:08:13] [TRT-LLM] [I] [_torch][RANK 0] " + f"iter = {iteration}, global_rank = 0, rank = 0, " + f"num_scheduled_requests = {bs}, kv_cache_util = 0.125, " + f"currank_total_requests = 1/8, host_step_time = {step_ms}ms, " + f"prev_device_step_time = {step_ms}ms, " + f"timestamp = 2026-08-04 19:08:13, " + f"states = {{'num_ctx_requests': 0, 'num_ctx_tokens': 0, " + f"'num_generation_tokens': {gen_tokens}, " + f"'cached_kv_tokens': 63821}}") + + +def _write(lines, name="gen.log"): + # tempfile rather than pytest's tmp_path: these tests also run under the + # fixture-less in-container runner used for the GPU images. + path = pathlib.Path(tempfile.mkdtemp()) / name + path.write_text("\n".join(lines) + "\n") + return str(path) + + +def test_verify_length_comes_from_the_padded_token_total(): + """The label is read out of the log, never inferred from step times. + + Steps the ladder cannot explain are skipped, not mislabelled: filing a + step under a shape it did not run is the mislabelling that once shipped + eight fictional cells with near-zero residuals. + """ + log = _write([ + _line(1, 1, 62098.0, 6), # executor's first step: dropped + _line(2, 244, 79.2, 768), # 244 rows padded to 256 -> rung-2 + _line(3, 242, 117.1, 1536), # padded to 256 -> full block + _line(4, 128, 61.0, 384), # exactly 128 rows -> rung-2 + _line(5, 256, 300.0, 2048), # L = 7, past the block: skipped + _line(6, 300, 400.0, 1800), # more rows than the ladder holds: skipped + ]) + samples = samples_from_iter_log([log], max_draft_len=5, + padded_batch_sizes=LADDER) + got = sorted((s.batch_size, s.verify_len, s.step_time_ms) + for s in samples) + assert got == [(128, 2, 61.0), (256, 2, 79.2), (256, 5, 117.1)] + # The table is indexed by the padded token total, which each sample exposes. + assert sorted(s.total_verify_tokens for s in samples) == [384, 768, 1536] + + +def test_the_padded_width_is_what_gets_recorded(): + """A 244-row step cost what 256 rows cost; the table is indexed by that. + + The ladder is also what removes the ambiguity: 768 tokens over 128 rows + is L=5 at 128, L=3 at 192, or L=2 at 256 -- without the ladder all three + are admissible and the step is dropped; with it, padding rounds 128 up to + 128 and the answer is L=5 (the first version of this collector kept 312 + of 5682 production steps for lack of this). And a step whose token total + the padded width cannot explain (_get_padded_batch can bail) is dropped: + nothing in the log says padding declined, so that is the only honest move. + """ + assert resolve_padded_shape(num_rows=244, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) == (256, 2) + assert resolve_padded_shape(num_rows=128, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256) is None + assert resolve_padded_shape(num_rows=128, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) == (128, 5) + assert resolve_padded_shape(num_rows=200, num_generation_tokens=600, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) is None + + +def test_the_first_step_of_each_executor_is_dropped(): + """A build spins up more than one executor, and each numbers from 1. + + The first step of each carries the wait for its first request inside + host_step_time (62 s on a live server), which would otherwise enter the + fit as a cell measurement. + """ + log = _write([ + _line(1, 1, 200.0, 6), # KV-estimation executor + _line(2, 256, 79.0, 768), + _line(1, 1, 62098.0, 6), # real executor starts over + _line(2, 256, 79.4, 768), + ]) + times = [s.step_time_ms for s in samples_from_iter_log([log], max_draft_len=5)] + assert times == [79.0, 79.4] + + +def test_only_the_requested_rank_is_read(): + """Every rank logs; counting all of them would multiply every cell. + + Outlier steps are bounded out on the same pass. + """ + log = _write([ + _line(2, 256, 79.2, 768, rank="0"), + _line(2, 256, 79.9, 768, rank="3"), # another rank's copy: dropped + _line(3, 256, 5000.0, 768), # outlier past max_step_ms + ]) + times = [s.step_time_ms for s in + samples_from_iter_log([log], max_draft_len=5, max_step_ms=1000.0)] + assert times == [79.2] + + +def test_several_logs_are_pooled(): + """One arm gives one rung; the table needs the ladder, so pool the runs.""" + a = _write([_line(1, 1, 100.0, 6), _line(2, 256, 117.1, 1536)], + name="a.log") + c = _write([_line(1, 1, 100.0, 6), _line(2, 256, 79.2, 768)], + name="c.log") + got = sorted(s.verify_len for s in + samples_from_iter_log([a, c], max_draft_len=5)) + assert got == [2, 5] + + +def test_a_log_without_the_states_dict_yields_nothing(): + """enable_iter_perf_stats off: fail loudly upstream, not with a bad table.""" + path = pathlib.Path(tempfile.mkdtemp()) / "bare.log" + path.write_text( + "0: [TRT-LLM] iter = 2, num_scheduled_requests = 256, " + "host_step_time = 79.2ms\n") + assert samples_from_iter_log([str(path)], max_draft_len=5) == [] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py new file mode 100644 index 000000000000..78d882c0ec0f --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The two reconciliation guards: engine fingerprint and price prediction. + +Both exist because this pipeline's characteristic failure is a PREDICTION +recorded as a MEASUREMENT with nothing positioned to contradict it: a cost +table measured on the wrong MoE kernel almost planned MegaMoE steps, the same +nominal cell measured 23% apart across two engine configs, and a planner +argmax'd against a table whose labels described shapes that never ran. The +fingerprint refuses the wrong table at load; the recorded prediction gives +every run a number that reality (hostStepTimeMS) can contradict. +""" + +import numpy as np +import pytest + +from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, check_table_fingerprint) +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig +from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + +LIVE = {"tp": 8, "ep": 8, "attention_dp": True, "block": 5, + "max_batch_size": 256} + + +def _payload(engine): + return {"token_counts": [512, 1536], "step_time_ms": [60.0, 150.0], + "_meta": {"engine": engine} if engine is not None else {}} + + +def test_mismatched_engine_is_refused(): + """The 23%-apart case: same cell, different max_batch_size.""" + wrong = dict(LIVE, max_batch_size=64) + with pytest.raises(ValueError, match="different engine configuration"): + check_table_fingerprint(payload=_payload(wrong), live=dict(LIVE)) + + +def test_fingerprints_that_must_load(): + """Every non-contradicted fingerprint loads; only a contradicted fact refuses. + + Matching facts load silently; string facts must not fail on spelling case + (MEGAMOE vs MegaMoe); facts the consumer cannot see (image hash, geometry) + are informational and logged, not refused -- refusing on unverifiable keys + would make every fingerprint addition a breaking change; and an old table + with no fingerprint warns rather than breaking existing deployments. + """ + check_table_fingerprint(payload=_payload(dict(LIVE)), live=dict(LIVE)) + check_table_fingerprint( + payload=_payload(dict(LIVE, moe_backend="megamoe_cutedsl")), + live=dict(LIVE, moe_backend="MEGAMOE_CUTEDSL")) + check_table_fingerprint( + payload=_payload(dict(LIVE, image="faf2c60935", + geometry="constant_block")), + live=dict(LIVE)) + check_table_fingerprint(payload=_payload(None), live=dict(LIVE)) + + +def test_planner_records_the_price_it_paid(): + """Every budget decision must leave behind the step time it assumed. + + Without this the run has no number reality can contradict: the argmax + consults the table and the table's opinion evaporates. With it, the + [final] planner block carries predicted_ms_sum / predicted_steps, and a + systematic gap against measured hostStepTimeMS is the one signal that + catches a wrong table, a wrong lookup, and an engine mismatch alike. + """ + table = SpsCostTable(token_counts=(0, 512, 768, 1536), + step_time_ms=(5.0, 68.4, 80.2, 150.5), + fixed_overhead_ms=1.0) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + planner = DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=[1, 2, 5]) + + bs = 8 + rng = np.random.default_rng(11) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() + + import torch + # Only the snapshot is faked; calibration stays the constructor's default + # sigmoid (setting it to None post-construction would bypass the + # `apply_calibration or torch.sigmoid` fallback and crash the decide path). + planner._gather_rows = lambda **_: torch.tensor(survival, + dtype=torch.float32) + + lens = planner.decide_verify_lens(num_gen_requests=bs, + reduce_across_ranks=False) + assert lens is not None + + steps = planner.stats.get("predicted_steps", 0) + total = planner.stats.get("predicted_ms_sum", 0.0) + assert steps == 1 and total > 0.0 + + # The recorded price must be the table's own answer for the tokens the + # decision implies -- anything else and the reconciliation compares two + # unrelated numbers. + budget = sum(lens) - bs * cfg.min_verify_len + tokens = bs * (cfg.min_verify_len + 1) + budget + expected = float(table.step_times(np.asarray([tokens]), bs)[0]) + assert planner.last_predicted_step_ms == pytest.approx(expected) + assert total == pytest.approx(expected) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py new file mode 100644 index 000000000000..85517f5c4c30 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The two snapshot instruments: stamp-lag histogram and neutral-row count. + +They exist because the planner's inputs have never been measured: the argmax +eats a slot-indexed snapshot staged one step earlier, and nothing recorded how +old each gathered row was (relay staleness) or how many rows were still the ++30 neutral fill (unknowns entering as certainties). The live anomaly they +instrument: 445/476 full-batch steps chose the full block while calibrated +pooled survivals favour rung-2 by 36% -- so the per-step snapshot content must +differ from the pooled shards, and these counters say in which way. + +Both are measurement-only. A failure inside them must never kill a step. +""" + +import numpy as np +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig +from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + + +def _planner(): + table = SpsCostTable(token_counts=(0, 512, 768, 1536), + step_time_ms=(5.0, 68.4, 80.2, 150.5), + fixed_overhead_ms=1.0) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + return DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=[1, 2, 5]) + + +def test_neutral_rows_are_counted_and_real_rows_are_not(): + planner = _planner() + selected = torch.zeros((8, 5)) + selected[3] = 30.0 # the neutral fill: sigmoid ~ 1.0 at any temperature + planner._note_snapshot_stats(selected, rows=None) + assert planner.stats["snap_rows"] == 8 + assert planner.stats["snap_neutral_rows"] == 1 + + +def test_lag_histogram_reads_staged_seq_minus_stamp(): + planner = _planner() + planner._host_stamps = torch.tensor([7, 6, 5, 0], dtype=torch.int32) + planner._staged_seq = 7 + selected = torch.zeros((4, 5)) + planner._note_snapshot_stats(selected, rows=[0, 1, 2, 3]) + hist = planner.stats["stamp_lag_hist"] + # Stamp 0 is "never drafted": lag 7 in truth, and it must land visibly at + # the histogram's far end rather than vanish -- these are exactly the rows + # the neutral counter should agree about. + assert hist == {0: 1, 1: 1, 2: 1, 7: 1} + + +def test_large_and_negative_lags_clamp_instead_of_exploding_the_dict(): + planner = _planner() + planner._host_stamps = torch.tensor([0, 500], dtype=torch.int32) + planner._staged_seq = 400 + planner._note_snapshot_stats(torch.zeros((2, 5)), rows=[0, 1]) + assert planner.stats["stamp_lag_hist"] == {8: 1, -2: 1} + + +def test_counts_accumulate_across_steps(): + planner = _planner() + planner._host_stamps = torch.tensor([3, 3], dtype=torch.int32) + planner._staged_seq = 4 + for _ in range(3): + planner._note_snapshot_stats(torch.full((2, 5), 30.0), rows=[0, 1]) + assert planner.stats["snap_rows"] == 6 + assert planner.stats["snap_neutral_rows"] == 6 + assert planner.stats["stamp_lag_hist"] == {1: 6} + + +def test_instrument_failure_is_counted_not_raised(): + planner = _planner() + planner._host_stamps = torch.tensor([1], dtype=torch.int32) + planner._staged_seq = 1 + # Rows that cannot index the stamp buffer must clamp, and anything worse + # must be swallowed into the error counter: instruments never kill steps. + planner._note_snapshot_stats(torch.zeros((1, 5)), rows=[10_000]) + assert planner.stats.get("snap_stats_errors", 0) == 0 + planner._note_snapshot_stats(object(), rows=None) # type: ignore[arg-type] + assert planner.stats["snap_stats_errors"] == 1 + + +def test_decide_path_feeds_the_instruments(): + """The counters must fill from the real decide path, not only direct calls.""" + planner = _planner() + bs = 8 + rng = np.random.default_rng(3) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() + planner._gather_rows = lambda **_: torch.tensor(survival, + dtype=torch.float32) + lens = planner.decide_verify_lens(num_gen_requests=bs, + reduce_across_ranks=False) + assert lens is not None + assert planner.stats["snap_rows"] == bs + assert planner.stats["snap_neutral_rows"] == 0 + # The local-decision instruments must fill on the same path: one rung per + # decision, one window per request, and the two must agree -- the windows + # are the top-k realisation of exactly that rung's budget. + rung_hist = planner.stats["local_rung_hist"] + assert sum(rung_hist.values()) == 1 + (rung,) = rung_hist + assert rung in (1, 2, 5) + len_hist = planner.stats["local_len_hist"] + assert sum(len_hist.values()) == bs + budget = bs * (rung - planner.cfg.min_verify_len) + extra = sum((k - planner.cfg.min_verify_len) * v for k, v in len_hist.items()) + assert extra <= budget diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py new file mode 100644 index 000000000000..39633aa600b5 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py @@ -0,0 +1,614 @@ +# 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. +"""Unit tests for the DSpark SPS cost-table profiler (CPU only, no model load). + +Everything here exercises the pure half of the profiler: turning ``get_stats`` +rows into aligned samples, summarizing cells, separating ``alpha(bs)`` from +``theta(M)``, and assembling / refusing the emitted JSON. The sweep driver needs +a GPU and a checkpoint and is deliberately not covered. + +The load-bearing case is :func:`test_round_trip_reproduces_measured_cells`: it +takes a synthetic machine with a known cost curve, runs it through the whole +fit-and-emit pipeline, and reloads the JSON exactly the way +``dspark.py::_build_verify_planner`` does. If the profiler and the loader ever +disagree about units, keys or the meaning of a token count, that test is where +it shows up. +""" + +import importlib.util +import json +import pathlib +import sys + +import numpy as np +import pytest + +from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable, total_verify_tokens +# The profiler is a development tool under tests/microbenchmarks, not part of +# the shipped package: nothing on the serving path imports it, and keeping a +# 1.8k-line sweep driver inside tensorrt_llm/ would ship it to every user. +# Load it by path so this test follows the file rather than an import root. +_SPEC = importlib.util.spec_from_file_location( + "dspark_sps_profiler", + pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / + "dspark_sps_profiler.py") +_PROFILER = importlib.util.module_from_spec(_SPEC) +# Registered BEFORE exec: @dataclass resolves annotations through +# sys.modules[cls.__module__].__dict__, so a module executed without being +# registered raises "'NoneType' object has no attribute '__dict__'" on its +# first dataclass. +sys.modules[_SPEC.name] = _PROFILER +_SPEC.loader.exec_module(_PROFILER) + +globals().update({name: getattr(_PROFILER, name) for name in ('CellStat', 'FlatCostTableError', 'InertCostTableError', 'InsufficientSamplesError', 'StepSample', 'SweepConfig', 'SweepGeometryError', 'aligned_steps_from_stats', 'build_cost_table_payload', 'check_table_is_informative', 'compress_to_risers', 'fit_additive_cost_model', 'load_cost_table', 'profitability_probe', 'running_max', 'summarize_cells')}) + +#: The exact set of keys ``dspark.py::_build_verify_planner`` reads. Anything +#: else in the file is ignored by the loader, so the profiler must not rely on +#: it -- and must not omit any of these. +LOADER_KEYS = { + "token_counts", + "step_time_ms", + "fixed_overhead_ms", + "batch_sizes", + "batch_overhead_ms", +} + + +def _stats_row(*, iteration, num_gen, rank=0, num_ctx=0, host_ms=10.0, gpu_ms=8.0): + """One /metrics-shaped iteration record.""" + row = { + "iter": iteration, + "attentionDpRank": rank, + "inflightBatchingStats": { + "numContextRequests": num_ctx, + "numGenRequests": num_gen, + }, + } + if host_ms is not None: + row["hostStepTimeMS"] = host_ms + if gpu_ms is not None: + row["gpuForwardTimeMS"] = gpu_ms + return row + + +# -------------------------------------------------------------------------- +# aligned step extraction +# -------------------------------------------------------------------------- + + +def test_aligned_steps_drops_prefill_iterations(): + """A step with context requests is a different shape, not a cheap decode.""" + steady = [_stats_row(iteration=i, num_gen=8, host_ms=10.0 + i) for i in range(3)] + assert aligned_steps_from_stats(steady) == [(0, 8, 10.0), (1, 8, 11.0), (2, 8, 12.0)] + rows = [ + _stats_row(iteration=0, num_gen=0, num_ctx=8, host_ms=90.0), + _stats_row(iteration=1, num_gen=8, num_ctx=2, host_ms=40.0), + _stats_row(iteration=2, num_gen=8, host_ms=10.0), + ] + assert aligned_steps_from_stats(rows) == [(2, 8, 10.0)] + + +def test_aligned_steps_requires_every_dp_rank(): + """Multi-rank alignment: complete, agreeing, and timed by rank 0. + + A missing rank means the iteration was not observed batch-wide; ranks at + different batch sizes are not one (bs, M) point; and ADP fans out + rank-local counters with rank-0's clock, so timing must never be averaged. + """ + rows = [ + _stats_row(iteration=0, num_gen=8, rank=0), + _stats_row(iteration=0, num_gen=8, rank=1), + _stats_row(iteration=1, num_gen=8, rank=0), + ] + assert aligned_steps_from_stats(rows, expected_ranks=2) == [(0, 8, 10.0)] + + disagreeing = [ + _stats_row(iteration=0, num_gen=8, rank=0), + _stats_row(iteration=0, num_gen=7, rank=1), + ] + assert aligned_steps_from_stats(disagreeing, expected_ranks=2) == [] + + rank0_not_first = [ + _stats_row(iteration=0, num_gen=8, rank=1, host_ms=99.0), + _stats_row(iteration=0, num_gen=8, rank=0, host_ms=10.0), + ] + assert aligned_steps_from_stats(rank0_not_first, expected_ranks=2) == [(0, 8, 10.0)] + + +def test_aligned_steps_skips_missing_timing_field(): + """gpuForwardTimeMS is read without a sync and is absent on some steps.""" + rows = [ + _stats_row(iteration=0, num_gen=8, gpu_ms=None), + _stats_row(iteration=1, num_gen=8, gpu_ms=7.5), + ] + assert aligned_steps_from_stats(rows, timing_key="gpuForwardTimeMS") == [(1, 8, 7.5)] + + +# -------------------------------------------------------------------------- +# warmup + median +# -------------------------------------------------------------------------- + + +def _samples(batch_size, verify_len, times): + return [ + StepSample(batch_size=batch_size, verify_len=verify_len, step_time_ms=t, iteration=i) + for i, t in enumerate(times) + ] + + +def test_summarize_discards_warmup(): + warm = [500.0] * 4 + steady = [10.0] * 8 + cells = summarize_cells(_samples(8, 3, warm + steady), warmup_steps=4, min_samples=8) + assert len(cells) == 1 + assert cells[0].step_time_ms == pytest.approx(10.0) + assert cells[0].num_samples == 8 + + +def test_summarize_reports_median_not_mean(): + """The right tail is real and must not move the point.""" + times = [10.0, 10.0, 10.0, 10.0, 10.0, 1000.0] + cells = summarize_cells(_samples(8, 3, times), warmup_steps=0, min_samples=1) + assert cells[0].step_time_ms == pytest.approx(10.0) + assert cells[0].step_time_ms != pytest.approx(float(np.mean(times))) + # ...but the tail is still reported, so a pathological cell is visible. + assert cells[0].p90_ms > cells[0].step_time_ms + + +def test_summarize_drops_thin_cells_outside_the_requested_grid(): + """A deployment log carries ramp/drain occupancy nobody asked to measure. + + Job 2585129: a five-rung server sweep with 984-step cells at every + requested (bs, L) died at the very end on (bs=2, L=1) with one sample -- + the admission ramp of the first cell's client, filed as a cell and then + held to the requested-grid bar. Incidental cells must be dropped when + thin, kept when thick, and a thin REQUESTED cell must still refuse -- + with or without a requested grid, since a silently dropped cell can + disconnect the grid, not just widen error bars. + """ + grid = [(8, 3)] + requested = _samples(8, 3, [10.0] * 8) + ramp = _samples(2, 3, [30.0]) # thin, incidental: drop + thick_extra = _samples(4, 3, [20.0] * 8) # thick, incidental: keep + cells = summarize_cells(requested + ramp + thick_extra, warmup_steps=0, + min_samples=4, expected_cells=grid) + assert [(c.batch_size, c.verify_len) for c in cells] == [(4, 3), (8, 3)] + with pytest.raises(InsufficientSamplesError, match="steady samples"): + summarize_cells(_samples(8, 3, [10.0]) + thick_extra, warmup_steps=0, + min_samples=4, expected_cells=grid) + # No grid given means every observed cell was requested: same refusal. + with pytest.raises(InsufficientSamplesError, match="steady samples"): + summarize_cells(_samples(8, 3, [10.0] * 5), warmup_steps=4, min_samples=8) + + +def test_summarize_orders_by_iteration_before_cutting_warmup(): + """Stats arrive per queue, not necessarily in iteration order.""" + out_of_order = [ + StepSample(batch_size=8, verify_len=1, step_time_ms=t, iteration=i) + for i, t in sorted(enumerate([500.0, 500.0, 10.0, 10.0]), key=lambda p: -p[0]) + ] + cells = summarize_cells(out_of_order, warmup_steps=2, min_samples=2) + assert cells[0].step_time_ms == pytest.approx(10.0) + + +def test_step_sample_json_round_trip(): + sample = StepSample(batch_size=16, verify_len=3, step_time_ms=12.5, iteration=7) + assert StepSample.from_json(json.loads(json.dumps(sample.to_json()))) == sample + assert sample.total_verify_tokens == total_verify_tokens(16, 3) == 64 + + +# -------------------------------------------------------------------------- +# additive fit +# -------------------------------------------------------------------------- + + +def _synthetic_cells(batch_sizes, verify_lens, *, alpha, theta, noise=None): + """Cells generated by an exactly additive machine ``alpha(bs) + theta(M)``.""" + cells = [] + for batch_size in batch_sizes: + for verify_len in verify_lens: + tokens = total_verify_tokens(batch_size, verify_len) + value = alpha(batch_size) + theta(tokens) + if noise is not None: + value += noise(batch_size, verify_len) + cells.append( + CellStat( + batch_size=batch_size, + verify_len=verify_len, + step_time_ms=value, + num_samples=32, + p10_ms=value, + p90_ms=value, + ) + ) + return cells + + +def _staircase(edges_and_costs): + def theta(tokens): + value = edges_and_costs[0][1] + for edge, cost in edges_and_costs: + if tokens >= edge: + value = cost + return value + + return theta + + +BATCH_SIZES = [8, 16, 32, 64] +VERIFY_LENS = [1, 2, 3, 4, 5] +ALPHA = {8: 20.0, 16: 22.0, 32: 26.0, 64: 34.0} +THETA = _staircase([(0, 1.0), (64, 2.0), (128, 4.0), (256, 8.0)]) + + +def test_fit_recovers_relative_structure_exactly(): + """The fit is only identified up to one shared constant -- pin that, then + every cell must be reproduced to the last bit.""" + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) + fit = fit_additive_cost_model(cells) + assert fit.warnings == () + + for cell in cells: + predicted = fit.intercept_ms[cell.batch_size] + fit.theta_ms[cell.total_verify_tokens] + assert predicted == pytest.approx(cell.step_time_ms, abs=1e-6) + assert fit.max_rel_residual == pytest.approx(0.0, abs=1e-9) + + # The batch axis keeps its measured *differences*, which is what alpha means. + for smaller, larger in zip(BATCH_SIZES, BATCH_SIZES[1:]): + measured = ALPHA[larger] - ALPHA[smaller] + fitted = fit.intercept_ms[larger] - fit.intercept_ms[smaller] + assert fitted == pytest.approx(measured, abs=1e-6) + + +def test_fit_surfaces_a_bad_cell_rather_than_hiding_it(): + """A cell the additive model cannot explain is evidence, not noise. + + Dropping it would produce a table that fits beautifully and predicts badly, + and would delete the one signal that says ``theta`` is not a function of + ``M`` alone. + """ + + def noise(batch_size, verify_len): + return 500.0 if (batch_size, verify_len) == (32, 3) else 0.0 + + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA, noise=noise) + fit = fit_additive_cost_model(cells) + + assert fit.max_abs_residual_ms > 100.0 + assert fit.max_rel_residual > 1.0 + report = " ".join(fit.warnings) + assert "additive model misses" in report + # The least-squares solve spreads the bad cell over everything sharing its + # batch size or token count, so the report names a neighbourhood -- the bad + # cell is in it, but so are innocent ones, which is why it is worded that way. + assert "bs=32, L=3" in report + + +def test_fit_clamps_a_dip_upward(): + """A non-monotone theta would offer the planner a cheaper *longer* verify.""" + dipped = _staircase([(0, 1.0), (64, 5.0), (128, 3.0), (256, 8.0)]) + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=dipped) + fit = fit_additive_cost_model(cells) + + tokens = sorted(fit.theta_ms) + values = [fit.theta_ms[m] for m in tokens] + assert values == running_max(values) + assert any("monoton" in w for w in fit.warnings) + + +def test_fit_refuses_a_disconnected_sweep(): + """Two islands each carry their own unknown offset; splicing them lies.""" + cells = _synthetic_cells([8], [1, 2], alpha=ALPHA.get, theta=THETA) + cells += _synthetic_cells([100], [1, 2], alpha=lambda _: 40.0, theta=THETA) + with pytest.raises(SweepGeometryError, match="disconnected"): + fit_additive_cost_model(cells) + + +# -------------------------------------------------------------------------- +# staircase shaping +# -------------------------------------------------------------------------- + + +def test_running_max_is_monotone(): + assert running_max([1.0, 3.0, 2.0, 2.5, 9.0]) == [1.0, 3.0, 3.0, 3.0, 9.0] + + +def test_compress_keeps_risers_and_closes_their_shelves(): + """A shelf must survive as a breakpoint PAIR under the interp consumer. + + Keeping only [16, 64] would turn the measured flat 16..48 into one long + rising segment: mid-shelf totals get billed part of the riser -- the + over-trim mirror of the floor-lookup bug. The shelf's last point (48) + rides along so interpolation stays flat where the hardware was flat. + """ + tokens = [16, 32, 48, 64, 80] + times = [1.0, 1.002, 1.004, 2.0, 2.001] + kept_tokens, kept_times = compress_to_risers( + tokens, times, min_riser_ms=0.02, max_breakpoints=8 + ) + assert kept_tokens == [16, 48, 64] + assert kept_times == [1.0, 1.004, 2.0] + + +def test_truncation_drops_the_point_cheapest_to_interpolate(): + """Over budget, remove the interior point interpolation best reconstructs. + + Here 64 goes: its neighbors predict it within 0.27ms, while removing 32 + would ramp the near-flat 16..32 stretch into the 64 riser and misprice it + by 0.9ms. The previous rule (drop the smallest incoming jump first) chose + 32 -- by construction it always chose the shelf-side point, re-billing + mid-shelf totals upward, the exact over-charge the pair encoding of + shelves exists to prevent. + """ + tokens = [16, 32, 64, 128] + times = [1.0, 1.1, 4.0, 9.0] + kept_tokens, _ = compress_to_risers(tokens, times, min_riser_ms=0.0, max_breakpoints=3) + assert kept_tokens == [16, 32, 128] + + +def test_compress_never_drops_the_base_shelf(): + """Index 0 is what every below-range query clamps onto.""" + tokens = [16, 32, 64] + times = [1.0, 100.0, 101.0] + kept_tokens, _ = compress_to_risers(tokens, times, min_riser_ms=0.0, max_breakpoints=2) + assert kept_tokens[0] == 16 + + +# -------------------------------------------------------------------------- +# payload assembly and JSON round-trip through SpsCostTable +# -------------------------------------------------------------------------- + + +def _payload(): + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=THETA) + return build_cost_table_payload(cells, riser_tolerance=0.01, max_breakpoints=8), cells + + +def test_round_trip_reproduces_measured_cells(): + """The whole point: what the planner reads back must be what was measured. + + The payload must carry exactly the keys ``dspark.py::_build_verify_planner`` + reads (everything else under ``_meta``, where the loader ignores it), + survive a JSON file round trip into a real ``SpsCostTable``, and price + every measured cell back to within one shelf width -- the emitted curve is + a floor staircase by construction, so a cell between two risers is + deliberately priced at the lower shelf. + """ + payload, cells = _payload() + assert LOADER_KEYS.issubset(payload) + assert set(payload) - LOADER_KEYS == {"_meta"} + + table = load_cost_table(json.loads(json.dumps(payload))) + assert isinstance(table, SpsCostTable) + assert not table.is_flat + # The loader builds SpsCostTable directly, so the profiler's own validation + # must be the real constructor rather than a copy of its rules. + assert list(table.token_counts) == payload["token_counts"] + assert len(table.batch_sizes) == len(table.batch_overhead_ms) + + for cell in cells: + predicted = table.step_time(cell.total_verify_tokens, cell.batch_size) + assert predicted == pytest.approx(cell.step_time_ms, abs=1e-4) + + +def test_emitted_terms_are_not_double_counted(): + """step_time_ms is theta only; the floor lives in the two overhead fields.""" + payload, cells = _payload() + table = load_cost_table(payload) + smallest = min(BATCH_SIZES) + floor = payload["fixed_overhead_ms"] + table.batch_overhead(smallest) + # The floor is the M-independent part, i.e. the intercept of the measured + # row at the smallest batch size -- not the whole step time. + row = sorted((c for c in cells if c.batch_size == smallest), key=lambda c: c.verify_len) + assert floor < row[0].step_time_ms + assert floor == pytest.approx(ALPHA[smallest], abs=0.5) + + +def test_batch_overhead_grows_with_batch_size(): + payload, cells = _payload() + overheads = payload["batch_overhead_ms"] + assert overheads == sorted(overheads) + assert payload["batch_sizes"] == sorted(BATCH_SIZES) + assert payload["fixed_overhead_ms"] >= 0.0 + # Provenance rides along under _meta: the encoding marker and the cells. + assert payload["_meta"]["encoding"] == "decomposed" + assert len(payload["_meta"]["cells"]) == len(cells) + + +# -------------------------------------------------------------------------- +# refusals +# -------------------------------------------------------------------------- + + +def test_refuses_a_flat_curve(): + """A machine whose step cost ignores M must not produce a table at all.""" + cells = _synthetic_cells(BATCH_SIZES, VERIFY_LENS, alpha=ALPHA.get, theta=lambda _: 3.0) + payload = build_cost_table_payload(cells) + assert load_cost_table(payload).is_flat + with pytest.raises(FlatCostTableError, match="is_flat"): + check_table_is_informative(payload, batch_sizes=BATCH_SIZES, tiers=[1, 3, 5]) + + +def test_refuses_a_curve_that_only_moves_by_noise(): + """is_flat is exact float equality, so it cannot catch this one.""" + cells = _synthetic_cells( + BATCH_SIZES, + VERIFY_LENS, + alpha=ALPHA.get, + theta=lambda tokens: 3.0 + 1e-4 * tokens, + ) + payload = build_cost_table_payload(cells, riser_tolerance=0.0, max_breakpoints=8) + assert not load_cost_table(payload).is_flat + with pytest.raises(FlatCostTableError, match="noise floor"): + check_table_is_informative( + payload, batch_sizes=BATCH_SIZES, tiers=[1, 3, 5], min_step_time_spread=0.02 + ) + + +# A real but shallow riser: at bs=32 tier 3 (M=128) sits below it and tier 5 +# (M=192) above, so the curve is genuinely non-flat and clears the noise floor +# (2.5%), yet the step only gets 2.5% more expensive while the acceptance yield +# given up between those tiers is ~36% at p=0.9. The argmax therefore never +# leaves max_tier -- non-flat and inert at the same time. +# +# The shelf is ENCODED as a breakpoint pair: under the interpolating consumer +# a bare [0, 160] would ramp from 2.0 to 2.5 across the whole range, dropping +# the ladder spread below the noise floor and turning this fixture into the +# flat-noise case instead of the inert one it exists to pin. +INERT_PAYLOAD = { + "token_counts": [0, 159, 160], + "step_time_ms": [2.0, 2.0, 2.5], + "fixed_overhead_ms": 18.0, + "batch_sizes": [32], + "batch_overhead_ms": [0.0], +} + + +def test_refuses_an_inert_table(): + """Non-flat, but no tier pair is worth trading yield for. + + This is the failure mode that does *not* trip ``fallback_flat_cost``, so + nothing downstream would report it. ``allow_inert=True`` is the explicit + escape hatch -- and it must still say that nothing trims. + """ + kwargs = dict(batch_sizes=[32], tiers=[1, 3, 5], acceptance_rates=[0.9, 0.95]) + with pytest.raises(InertCostTableError, match="inert"): + check_table_is_informative(INERT_PAYLOAD, **kwargs) + diagnostics = check_table_is_informative(INERT_PAYLOAD, allow_inert=True, **kwargs) + assert diagnostics["trims_somewhere"] is False + + +def test_informative_table_passes_and_reports_where_trimming_wins(): + payload, _ = _payload() + diagnostics = check_table_is_informative( + payload, + batch_sizes=BATCH_SIZES, + tiers=[1, 3, 5], + acceptance_rates=[0.3, 0.6, 0.9], + ) + assert diagnostics["trims_somewhere"] is True + assert diagnostics["best_step_time_spread"] > 0.02 + # Even a good table cannot rescue a one-rung ladder: nothing to choose. + with pytest.raises(InertCostTableError, match="single entry"): + check_table_is_informative(payload, batch_sizes=BATCH_SIZES, tiers=[5]) + + +# -------------------------------------------------------------------------- +# sweep geometry guards (pure; the sweep itself needs a GPU) +# -------------------------------------------------------------------------- + + +def _sweep(**overrides) -> SweepConfig: + kwargs = dict( + model="/models/x", + speculative_model="/models/x", + batch_sizes=[8, 16], + verify_lens=[1, 3, 5], + input_len=1024, + warmup_steps=8, + measure_steps=32, + max_seq_len=4096, + max_num_tokens=8192, + ) + kwargs.update(overrides) + return SweepConfig(**kwargs) + + +def test_token_budget_is_the_step_count_when_acceptance_is_pinned(): + """Pinned, a step commits exactly one token, so the two are the same + number; unpinned it can commit verify_len + 1, so the budget must scale.""" + pinned = _sweep() + assert pinned.max_tokens_for(5) == pinned.warmup_steps + pinned.measure_steps + 16 + loose = _sweep(pin_acceptance=0.0) + assert loose.max_tokens_for(5) == pinned.max_tokens_for(5) * 6 + + +def test_validate_rejects_a_sequence_that_would_be_evicted(): + """An evicted request drains the batch mid-cell and shrinks the shape.""" + with pytest.raises(SweepGeometryError, match="max-seq-len"): + _sweep(input_len=4064, block_follows_verify_len=True).validate() + + +def test_validate_rejects_a_decode_step_wider_than_max_num_tokens(): + """A split decode step never has the shape its cell is filed under.""" + with pytest.raises(SweepGeometryError, match="widest decode step"): + _sweep(batch_sizes=[8, 4096], max_num_tokens=1024, max_seq_len=8192, + block_follows_verify_len=True).validate() + + +def test_validate_accepts_a_sane_sweep(): + """Sane includes verifying below a constant drafted block: the Repeat.cu + output_size assert that geometry used to trip was an IMA-family symptom, + fixed by the published-bucket graph key and verified gone 2026-08-06.""" + _sweep(block_follows_verify_len=True).validate() + _sweep(verify_lens=[5], max_draft_len=5).validate() + _sweep().validate() + + +def test_dp_size_only_counts_ranks_under_attention_dp(): + """Non-ADP publishes one stats row per iteration however wide TP is.""" + assert _sweep(tp_size=8).dp_size == 1 + assert _sweep(tp_size=8, enable_attention_dp=True).dp_size == 8 + + +# -------------------------------------------------------------------------- +# profitability probe (runs the real planner argmax) +# -------------------------------------------------------------------------- + + +def test_probe_trims_at_low_acceptance_and_not_at_high(): + """Trimming is a trade: it only wins when the drafts are unlikely to survive.""" + payload = { + # A riser between bs*(3+1)=128 and bs*(5+1)=192 at bs=32, so tiers 3 and + # 5 sit on different shelves and the argmax has something to choose: + # T(3) = 3 ms, T(5) = 4 ms, a 33% step for a yield gain that is worth it + # only when the drafts survive. + "token_counts": [0, 160], + "step_time_ms": [2.0, 3.0], + "fixed_overhead_ms": 0.0, + "batch_sizes": [32], + "batch_overhead_ms": [1.0], + } + probes = profitability_probe( + payload, batch_size=32, tiers=[1, 3, 5], acceptance_rates=[0.05, 0.99] + ) + by_rate = {p["acceptance_rate"]: p for p in probes} + assert by_rate[0.05]["chosen_verify_len"] < 5 + assert by_rate[0.05]["trims"] is True + assert by_rate[0.99]["chosen_verify_len"] == 5 + + +def test_probe_uses_the_bonus_token_convention(): + """A tier's cost must be looked up at bs*(L+1), never bs*L. + + The two differ by a whole batch, which is easily a shelf's width; this pins + the profiler to the same convention as ``total_verify_tokens``. + """ + payload = { + # Priced so that tier 5 (M = 8*6 = 48) is expensive and tier 3 + # (M = 8*4 = 32) is cheap. Under the wrong bs*L convention tier 5 would + # be M = 40 and land on the cheap shelf instead. The shelf is a + # breakpoint pair so the interpolating consumer keeps it flat. + "token_counts": [0, 43, 44], + "step_time_ms": [1.0, 1.0, 50.0], + "fixed_overhead_ms": 0.0, + "batch_sizes": [8], + "batch_overhead_ms": [0.0], + } + table = load_cost_table(payload) + assert table.step_time(total_verify_tokens(8, 5), 8) == pytest.approx(50.0) + assert table.step_time(total_verify_tokens(8, 3), 8) == pytest.approx(1.0) + probes = profitability_probe(payload, batch_size=8, tiers=[1, 3, 5], acceptance_rates=[0.9]) + assert probes[0]["chosen_verify_len"] == 3 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py new file mode 100644 index 000000000000..1324ad0c7684 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py @@ -0,0 +1,135 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A collection run must refuse the regimes that fit the scheduler, not the head. + +Both regimes below produce shards that look perfectly well-formed. The run +completes, the row counts are plausible, the fitter converges, and the emitted +temperatures are wrong -- there is no downstream check that can tell. That is +why these are construction-time failures rather than warnings. + +Measured on this branch: collecting while the planner trimmed gave +``ece_before = [.244 .319 .361 .389 .394]``; pinning the window to the full +block, same model and same fitter, gave ``[.157 .208 .270 .321 .344]``. The +difference is entirely the censored label. + +SGLang guards the first regime next to its own collector +(``dspark_planner.py:100-108``) and names the second in the sibling +``ConfidenceMetricsProbe`` -- "padded verify rows corrupt the per-position +prefix label" (``dspark_observability.py:688-697``). +""" + +import contextlib +import os +import tempfile + +import pytest + +from tensorrt_llm._torch.speculative.dspark_sts import (STS_COLLECT_ENV, + make_recorder_from_env) + +_PIN = "TLLM_DSPARK_FORCE_VERIFY_LEN" + + +@contextlib.contextmanager +def collecting(pin=None): + """Set the collection env for the block, restore it after. + + A context manager rather than a pytest fixture: these run inside the + container, which ships no pytest, under a minimal shim that does not + implement fixtures or tmp_path. + """ + saved = {k: os.environ.get(k) for k in (STS_COLLECT_ENV, _PIN)} + try: + with tempfile.TemporaryDirectory() as tmp: + os.environ[STS_COLLECT_ENV] = os.path.join(tmp, "collect") + if pin is None: + os.environ.pop(_PIN, None) + else: + os.environ[_PIN] = str(pin) + yield + finally: + for k, v in saved.items(): + os.environ.pop(k, None) if v is None else os.environ.__setitem__(k, v) + + +def test_collection_off_builds_nothing(): + """The serving default must be untouched by any of this.""" + os.environ.pop(STS_COLLECT_ENV, None) + assert make_recorder_from_env(block_size=5) is None + assert make_recorder_from_env(block_size=5, has_cost_table=True, + ragged_mode="compact") is None + + +def test_a_loaded_cost_table_without_a_pin_is_refused(): + """A table means the planner trims, and trimming censors the label.""" + with collecting(): + with pytest.raises(ValueError, match="censors the acceptance label"): + make_recorder_from_env(block_size=5, has_cost_table=True) + + +def test_a_pinned_window_makes_the_table_harmless(): + """The uncensored regimes must all build a recorder. + + Pinning to the full block makes the label uncensored by construction; + without a table the planner cannot trim, so nothing censors anything; + static mode contributes no padding rows. + """ + with collecting(pin=5): + rec = make_recorder_from_env(block_size=5, has_cost_table=True) + assert rec is not None and rec.block_size == 5 + with collecting(): + assert make_recorder_from_env(block_size=5, + has_cost_table=False) is not None + assert make_recorder_from_env(block_size=5, + ragged_mode="static") is not None + + +def test_compact_mode_is_refused(): + """Padding rows contribute prefix labels that measure nothing. + + Refused even with the pin set: the pin fixes the *window*, while this is + about rows that belong to no real request at all. + """ + with collecting(pin=5): + with pytest.raises(ValueError, match="Padded verify rows"): + make_recorder_from_env(block_size=5, ragged_mode="compact") + + +class _Head: + pass + + +class _Bare: + """DSparkDraftModel layout: stages under .mtp_layers.""" + + def __init__(self, head): + stage = type("Stage", (), {})() + if head is not None: + stage.confidence_head = head + self.mtp_layers = [object(), object(), stage] + + +class _Wrapper: + """DSparkForCausalLM layout: bare model under .dspark_model. + + This is the layout the chained getattr missed: the wrapper has no + mtp_layers of its own, so `getattr(model, "mtp_layers", [None])[-1]` + resolved to None and silently disabled calibration for every + confidence-scheduled run of the first campaign. + """ + + def __init__(self, head): + self.dspark_model = _Bare(head) + + +def test_head_resolves_across_both_layouts_and_absence_is_none(): + from tensorrt_llm._torch.speculative.dspark_sts import resolve_confidence_head + h = _Head() + assert resolve_confidence_head(_Bare(h)) is h + # The wrapper layout is the one the chained getattr missed: resolving to + # None here silently disabled calibration for every confidence-scheduled + # run of the first campaign. + assert resolve_confidence_head(_Wrapper(h)) is h + assert resolve_confidence_head(_Bare(None)) is None + assert resolve_confidence_head(_Wrapper(None)) is None + assert resolve_confidence_head(object()) is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py new file mode 100644 index 000000000000..23f2467ce68c --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py @@ -0,0 +1,168 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""STS pairing must join on draft-pass identity, never on arrival order. + +The recorder's label (accepted count) and its feature (confidence logits) are +produced by different steps: the block verified at step ``i`` was drafted at +step ``i-1``, and under the overlap scheduler the sampler consumes the label +while pass ``i+1`` has already overwritten the shared confidence buffer. The +single mutable stash this ring replaced had NO execution order under which +the pair was correct -- measured on job 2562577, 82.2% of rows paired +confidence(t+1) with accepted(t), correlation 0.14 -- and the shards looked +perfectly healthy. + +These tests pin the join contract: a pair is appended only when the ring +snapshot's per-row stamp equals the pass the label verifies, everything else +is a *counted* decline, and the shard carries provenance the fitter can +refuse. No CUDA: the recorder itself is device-agnostic (CPU tensors skip the +pinned-copy event), which is exactly what lets this run in pre-merge CI. +""" + +import importlib.util +import os +import pathlib +import sys +import tempfile + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_sts import DSparkStsRecorder + +BLOCK = 5 +ROWS = 8 + + +def _recorder(tmp, flush_every=1000): + return DSparkStsRecorder(path_stem=os.path.join(tmp, "shard"), + block_size=BLOCK, rank=0, + flush_every=flush_every) + + +def _stage(rec, seq, written_rows, fill): + """Stage one pass's snapshot: `written_rows` stamped `seq`, rest stale. + + Mirrors the producer: the in-graph scatter stamps only the rows the pass + actually drafted; every other row keeps whatever stamp it had (0 for + never-written, matching the worker's zero-initialized stamp buffer). + """ + logits = torch.full((ROWS, BLOCK), float(fill)) + stamps = torch.zeros(ROWS, dtype=torch.int32) + stamps[list(written_rows)] = seq + rec.stage_snapshot(device_logits=logits, device_stamps=stamps, + staged_seq=seq) + + +def test_pairs_by_pass_identity_not_arrival_order(): + """A label for pass 1 gets pass-1 logits even after pass 2 was staged.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + _stage(rec, seq=1, written_rows=[2], fill=1.0) + _stage(rec, seq=2, written_rows=[2], fill=2.0) # overwrite arrives first + rec.record(row=2, accepted=3, target_seq=1) + rec.record(row=2, accepted=1, target_seq=2) + assert rec.stats["recorded"] == 2 + assert rec._logits[0][0].item() == 1.0 # pass 1 content, not latest + assert rec._logits[1][0].item() == 2.0 + + +def test_row_not_written_for_that_pass_is_declined(): + """A row the pass never drafted must not contribute a fabricated pair.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + _stage(rec, seq=1, written_rows=[2], fill=1.0) + rec.record(row=3, accepted=2, target_seq=1) # row 3 stamp is 0, not 1 + assert rec.stats["recorded"] == 0 + assert rec.stats["stale_stamp"] == 1 + + +def test_ring_eviction_declines_instead_of_mispairing(): + """A label older than the ring depth is dropped, never matched to the + + evicting pass's content: the slot's stamps say who actually wrote it. + """ + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + for seq in range(1, 2 + DSparkStsRecorder._RING_DEPTH): + _stage(rec, seq=seq, written_rows=[2], fill=float(seq)) + rec.record(row=2, accepted=3, target_seq=1) # slot now holds seq 5 + assert rec.stats["recorded"] == 0 + assert rec.stats["stale_stamp"] == 1 + + +def test_every_giving_up_path_is_counted(): + """Silent declines are how the mispaired fit shipped; none may be silent.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + rec.record(row=2, accepted=1, target_seq=None) # no pass id + rec.record(row=2, accepted=1, target_seq=7) # ring slot never staged + _stage(rec, seq=1, written_rows=[2], fill=1.0) + rec.record(row=None, accepted=1, target_seq=1) # unresolvable request + rec.record(row=ROWS + 3, accepted=1, target_seq=1) + assert rec.stats == {"recorded": 0, "no_snapshot": 2, + "stale_stamp": 0, "row_out_of_range": 1, + "no_row": 1, "snapshots_staged": 1} + + +def test_staging_without_stamps_is_a_noop(): + """No stamps means no way to verify identity -- refuse to guess.""" + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + rec.stage_snapshot(device_logits=torch.zeros(ROWS, BLOCK), + device_stamps=None, staged_seq=1) + rec.stage_snapshot(device_logits=torch.zeros(ROWS, BLOCK), + device_stamps=torch.zeros(ROWS, dtype=torch.int32), + staged_seq=None) + assert rec.stats["snapshots_staged"] == 0 + rec.record(row=2, accepted=1, target_seq=1) + assert rec.stats["no_snapshot"] == 1 + + +def test_shard_carries_provenance_and_prefix_semantics(): + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + _stage(rec, seq=1, written_rows=[2], fill=1.5) + rec.record(row=2, accepted=3, target_seq=1) + rec.flush() + blob = torch.load(os.path.join(tmp, "shard.r0.0.pt"), + map_location="cpu") + assert blob["meta"]["pairing"] == "draft_seq_ring" + assert blob["meta"]["stats"]["recorded"] == 1 + assert blob["logits"].shape == (1, BLOCK) + # accepted=3 -> positions 0..2 were ALL accepted: the prefix event + # whose probability survival[k] estimates. + assert blob["prefix_mask"].tolist() == [[1.0, 1.0, 1.0, 0.0, 0.0]] + + +# --------------------------------------------------------------------------- +# The fitter must refuse shards from the pre-ring recorders: their pairs are +# mislabeled in ways the tensors themselves cannot reveal. +# --------------------------------------------------------------------------- + +_SPEC = importlib.util.spec_from_file_location( + "dspark_fit_sts", + pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / + "dspark_fit_sts.py") +_FITTER = importlib.util.module_from_spec(_SPEC) +sys.modules[_SPEC.name] = _FITTER +_SPEC.loader.exec_module(_FITTER) + + +def test_fitter_refuses_shards_without_ring_provenance(): + with tempfile.TemporaryDirectory() as tmp: + torch.save({"logits": torch.zeros(4, BLOCK), + "prefix_mask": torch.zeros(4, BLOCK)}, + os.path.join(tmp, "old.r0.0.pt")) + with pytest.raises(SystemExit, match="draft_seq_ring"): + _FITTER.load_shards(os.path.join(tmp, "*.pt")) + + +def test_fitter_accepts_ring_shards(): + with tempfile.TemporaryDirectory() as tmp: + rec = _recorder(tmp) + _stage(rec, seq=1, written_rows=[2, 3], fill=0.5) + rec.record(row=2, accepted=2, target_seq=1) + rec.record(row=3, accepted=5, target_seq=1) + rec.flush() + logits, mask = _FITTER.load_shards(os.path.join(tmp, "*.pt")) + assert logits.shape == (2, BLOCK) and mask.shape == (2, BLOCK) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py new file mode 100644 index 000000000000..2160e1d0ba49 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py @@ -0,0 +1,96 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The budget must be scored with the tau the step actually collects. + +``decide_verify_lens`` picks a rung, converts it to ``B = n*(L - min)`` tokens, +and hands B to ``schedule_verify_lens_topk``, which spends it on the B largest +survivals in the batch. But the rung was chosen by +``budget_argmax_over_uniform_lens``, which scores rung L as if every request took +columns ``[min, L)``. Top-k is by construction the largest sum over all size-B +subsets, so the realised tau is >= the scored one at an identical token cost -- +and the gap is not the same at every rung, so the argmax can land on the wrong +one. + +The fix is not to let the planner choose an arbitrary budget: the executor can +only run ``bs*(t+1)`` tokens for a captured tier, and any other total rounds UP, +which would push a budget deliberately chosen below a cost riser back over it. +So the ladder stays the grid and only the numerator changes. +""" + +import numpy as np +import pytest + +from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, budget_argmax_over_uniform_lens, compute_verify_token_budget) + + +def _table() -> SpsCostTable: + """A staircase with genuine risers, so trimming is worth something. + + Breakpoints are spaced so that neither tier is trivially optimal: a flat + table degenerates to verify-all and would make every assertion here vacuous. + """ + token_counts = tuple(range(0, 400, 16)) + step_time_ms = tuple(4.0 + 0.6 * (tok // 96) + 0.004 * tok + for tok in token_counts) + return SpsCostTable(token_counts=token_counts, step_time_ms=step_time_ms, + fixed_overhead_ms=1.0) + + +@pytest.mark.parametrize("tiers", [[1, 2, 5], [1, 3, 5], [1, 5]]) +def test_restricted_answer_is_always_realisable(tiers): + """Every returned budget corresponds to a rung the executor captured.""" + rng = np.random.default_rng(20260803) + table = _table() + for _ in range(200): + bs = int(rng.integers(2, 17)) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] + n = compute_verify_token_budget( + survival=survival, num_gen_requests=bs, cost_table=table, + min_verify_len=1, allowed_lens=tiers) + assert n % bs == 0, ( + f"budget {n} for {bs} requests is not n*(t-min) for any tier") + assert (n // bs) + 1 in tiers, ( + f"budget {n} implies tier {(n // bs) + 1}, not in {tiers}") + + +def test_restricted_never_scores_worse_than_the_uniform_choice(): + """Same grid, better numerator: the chosen rung's realised theta wins. + + Both routes evaluate the same discrete cost points, so the restricted form + cannot be worse under the tau that the step actually collects. This is the + property that makes the swap safe -- if it ever failed, the change would be + trading a modelling error for a real one. At bs=1 the uniform and top-k + allocations are the same set, so there the two scorers must agree on the + rung outright: any disagreement is a plain bug, not the intended + improvement. + """ + rng = np.random.default_rng(31337) + table = _table() + tiers = [1, 2, 5] + for _ in range(300): + bs = int(rng.integers(1, 17)) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] + + def realised_theta(tier: int) -> float: + budget = bs * (tier - 1) + cand = np.sort(survival[:, 1:].reshape(-1))[::-1] + tau = float(bs) + float(survival[:, :1].sum()) + float( + cand[:budget].sum()) + tokens = np.array([bs * (tier + 1)]) + return tau / float(table.step_times(tokens, bs)[0]) + + n = compute_verify_token_budget( + survival=survival, num_gen_requests=bs, cost_table=table, + min_verify_len=1, allowed_lens=tiers) + chosen = (n // bs) + 1 + uniform = budget_argmax_over_uniform_lens( + survival=survival, num_gen_requests=bs, cost_table=table, + allowed_lens=tiers, min_verify_len=1) + assert realised_theta(chosen) >= realised_theta(uniform) - 1e-12, ( + f"restricted picked tier {chosen} whose realised theta is below " + f"the uniform scorer's tier {uniform}") + if bs == 1: + assert chosen == uniform, ( + "at bs=1 the uniform and top-k allocations are the same set; " + "the two scorers must pick the same rung") From 493996f6f127f358c4c4b6f12985781626583429 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:23:26 -0700 Subject: [PATCH 142/177] [None][test] DSpark verify-len pin, planner trim, and forced-acceptance tests Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../hw_agnostic/test_dspark_verify_len_pin.py | 97 +++++ .../speculative/test_dspark_planner_trims.py | 75 ++++ .../speculative/test_force_accepted_tokens.py | 363 ++++++++++++++++++ 3 files changed, 535 insertions(+) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py create mode 100644 tests/unittest/_torch/speculative/test_dspark_planner_trims.py create mode 100644 tests/unittest/_torch/speculative/test_force_accepted_tokens.py diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py new file mode 100644 index 000000000000..04b8d5c0dce6 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The runtime verify-length pin: queue, agree, adopt. + +A cost-table sweep has to hold every request at one verify length while it +measures a cell, or the cell's label describes a shape that never ran (this +branch shipped a sweep where eight of twelve cells were mislabelled and the +fit dutifully concluded the curve was flat). The environment variable that +does this today is read once at planner construction, so walking the ladder +means rebuilding the engine per length -- and on a pre-spawned world it +cannot reach the ranks that are already running. + +The runtime pin fixes both, but it must never let one rank pin while its +peers do not: the ranks would derive different token totals, and the +attention-DP shape gate would drop the whole group out of graph replay for as +long as they disagreed. So the pin is queued locally, carried by the step's +existing allgather, and adopted from that agreed payload -- which is what +these tests pin down. +""" + +import pytest + +from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable +from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig +from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + + +def _planner(tiers=(1, 2, 5)): + table = SpsCostTable(token_counts=(0, 512, 768, 1536), + step_time_ms=(5.0, 68.4, 80.2, 150.5), + fixed_overhead_ms=1.0) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + return DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=list(tiers)) + + +def test_adopting_the_agreed_value_applies_it(): + """Queuing must not change this rank's behaviour before the group agrees; + adopting the agreed value applies it and consumes the pending request.""" + planner = _planner() + assert planner.request_verify_len_pin(2) == 2 + # Queued, not yet in force. + assert planner._forced_verify_len is None + assert planner.pending_verify_len_pin() == 2 + planner.adopt_verify_len_pin(2) + assert planner._forced_verify_len == 2 + # Consumed: the next step must not re-broadcast a pin nobody asked for. + assert planner.pending_verify_len_pin() == -1 + + +def test_a_rank_that_queued_nothing_still_adopts_the_group_value(): + """The endpoint lands on ONE rank; the others learn it from the payload.""" + peer = _planner() + assert peer.pending_verify_len_pin() == -1 + peer.adopt_verify_len_pin(5) + assert peer._forced_verify_len == 5 + + +def test_zero_clears_the_pin(): + """Wire protocol: -1 means nobody asked (leave the pin alone); clearing + needs its own value, 0, because None cannot travel in an int payload.""" + planner = _planner() + planner.adopt_verify_len_pin(2) + planner.adopt_verify_len_pin(-1) + assert planner._forced_verify_len == 2 + assert planner.request_verify_len_pin(None) is None + assert planner.pending_verify_len_pin() == 0 + planner.adopt_verify_len_pin(0) + assert planner._forced_verify_len is None + + +def test_an_uncaptured_length_is_refused_at_the_call_site(): + """Rejected when requested, not on some later step, and never queued.""" + planner = _planner(tiers=(1, 2, 5)) + with pytest.raises(ValueError, match="captured tier ladder"): + planner.request_verify_len_pin(3) + with pytest.raises(ValueError, match="outside"): + planner.request_verify_len_pin(9) + assert planner.pending_verify_len_pin() == -1 + assert planner._forced_verify_len is None + + +def test_pinned_steps_hand_every_request_the_same_window(): + """The whole point: the shape a sweep cell can honestly label.""" + import numpy as np + import torch + + planner = _planner() + bs = 8 + rng = np.random.default_rng(5) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() + planner._gather_rows = lambda **_: torch.tensor(survival, + dtype=torch.float32) + planner.adopt_verify_len_pin(2) + lens = planner.decide_verify_lens(num_gen_requests=bs, + reduce_across_ranks=False) + assert lens == [2] * bs + assert planner.stats["forced_steps"] == 1 diff --git a/tests/unittest/_torch/speculative/test_dspark_planner_trims.py b/tests/unittest/_torch/speculative/test_dspark_planner_trims.py new file mode 100644 index 000000000000..bcb12eecb37c --- /dev/null +++ b/tests/unittest/_torch/speculative/test_dspark_planner_trims.py @@ -0,0 +1,75 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Make the planner itself produce per-request windows, with a crafted table. + +This replaces TLLM_DSPARK_FORCE_VERIFY_LENS. That knob assigned windows by +*batch position*, which is orthogonal to confidence: it reproduced the ragged +*shape* but never the ragged *policy*, so it could not tell a correct +confidence-driven assignment from an arbitrary one. It also had to sit in +production code, ahead of the cost-table gate, to be reachable at all. + +A cost table does the same job through the real path. The planner trims when +the marginal cost of extra verify tokens outweighs their expected acceptance, +so a table whose theta(M) rises steeply makes it trim at any acceptance rate -- +and the budget then flows through the same confidence-ordered top-k, which +hands different windows to requests with different survival. That exercises +what actually ships. + +Why this matters beyond tidiness: the KV leak fixed in this branch (KV reserved +for the full drafted block but only the verified window rewound) is reachable +only when verified_len < draft_len. On this checkpoint the planner declines to +trim at every batch size, because draft acceptance is ~90% and the fitted table +says verifying the full block wins -- so without a way to make it trim, that +code path ships untested. +""" + +import torch + +from tensorrt_llm._torch.speculative.dspark_schedule import ( + DSparkScheduleConfig, schedule_verify_lens_topk) + + +def _Cfg(min_verify_len=1, max_verify_len=5): + """The real config, so the test cannot drift from what ships.""" + return DSparkScheduleConfig(block_size=max_verify_len, + min_verify_len=min_verify_len, + max_verify_len=max_verify_len) + + +def test_topk_hands_longer_windows_to_more_confident_requests(): + """The property the position-rotating knob could never check. + + Request 0 is confident throughout, request 2 collapses immediately. A + correct scheduler spends the budget where survival is highest. + """ + # survival[r, k] = P(the first k+1 drafted tokens all get accepted) + survival = torch.tensor([ + [0.99, 0.98, 0.97, 0.96, 0.95], # confident + [0.90, 0.70, 0.50, 0.30, 0.10], # middling + [0.20, 0.04, 0.01, 0.00, 0.00], # collapses + ]) + lens = schedule_verify_lens_topk(survival=survival, budget=6, + cfg=_Cfg()).tolist() + + assert len(lens) == 3 + assert lens[0] > lens[2], ( + f"confident request got {lens[0]}, collapsing request got {lens[2]}; " + f"the budget must follow survival, not batch position") + assert lens[0] >= lens[1] >= lens[2] + # Every request keeps at least the floor, and the budget is respected. + assert min(lens) >= 1 + assert sum(lens) - len(lens) * 1 <= 6 + + +def test_a_full_budget_degenerates_to_the_uniform_full_window(): + """The no-trim case must stay exactly uniform. + + This is what the planner correctly chooses on this checkpoint, and it is + the case where ragged must cost nothing. + """ + num_reqs, max_len = 4, 5 + survival = torch.full((num_reqs, max_len), 0.99) + budget = num_reqs * (max_len - 1) + lens = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=_Cfg()).tolist() + assert lens == [max_len] * num_reqs diff --git a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py new file mode 100644 index 000000000000..144de6a6c174 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py @@ -0,0 +1,363 @@ +# 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. +"""Unit tests for the synthetic forced acceptance rate in 1-model spec dec. + +These tests exercise ``SpecWorkerBase._apply_force_accepted_tokens``, which +honors the ``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` environment +variable. The implementation must be: + + - Faithful to fractional rates: integer part always accepts, fractional + part acts as the per-iteration probability of accepting one extra draft + token. + - Tensor-parallel deterministic: identical seed/counter state across + ranks so all ranks agree on per-request accepted counts (otherwise + downstream collectives expecting matching shapes deadlock). + - CUDA-graph compatible: pure-tensor ops with statically-known shapes, + captureable and replayable. +""" + +import os +from typing import Optional +from unittest.mock import patch + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata +from tensorrt_llm._torch.speculative.interface import ( + FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR, + SpecWorkerBase, + get_force_num_accepted_tokens, + get_force_num_accepted_tokens_float, +) + + +class _StubSpecWorker(SpecWorkerBase): + """Concrete ``SpecWorkerBase`` that stubs out the abstract API. + + Used purely to drive ``_apply_force_accepted_tokens`` in isolation. + """ + + @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: + worker = _StubSpecWorker() + if value is not None: + worker.force_num_accepted_tokens = value + return worker + + +def _make_input(batch_size: int, device: str = "cuda") -> torch.Tensor: + # Initial counts: 1 (target token, no draft accepts yet). The forced + # override only ever rewrites entries from ``num_contexts:`` onward. + return torch.ones(batch_size, dtype=torch.int32, device=device) + + +def _require_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required for this test.") + + +# ---------------- env-var parsing helpers ----------------------------------- + + +@pytest.mark.parametrize( + "env_value, expected_float, expected_int", + [ + ("0", 0.0, 0), + ("3", 3.0, 3), + ("2.6", 2.6, 0), + ("0.5", 0.5, 0), + ("not-a-number", 0.0, 0), + ], +) +def test_get_force_num_accepted_tokens_float(env_value, expected_float, expected_int): + """The float helper parses fractional rates; the int helper (used by the + 2-model path) must keep its original behavior on the same values.""" + with patch.dict(os.environ, {FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR: env_value}): + assert get_force_num_accepted_tokens_float() == pytest.approx(expected_float) + assert get_force_num_accepted_tokens() == expected_int + + +# ---------------- _apply_force_accepted_tokens semantics -------------------- + + +def test_zero_value_is_noop(): + _require_cuda() + worker = _make_worker(0.0) + base = _make_input(batch_size=4) + before = base.clone() + out = worker._apply_force_accepted_tokens(base, num_contexts=0, runtime_draft_len=4) + assert torch.equal(out, before) + # No RNG state should have been touched in the early-exit path. + assert worker._force_accept_rng_pool is None + assert worker._force_accept_rng_counter is None + + +@pytest.mark.parametrize( + "value, runtime_draft_len", + [ + (1.0, 4), + (2.0, 4), + (4.0, 4), + # 10 draft tokens requested but only 2 available -> capped at 3. + (10.0, 2), + # Fractional but no room: max_total = 3, base_total = min(10, 3) = 3, + # so ``base_total < max_total`` is False -> all 3, no RNG draw. + (9.5, 2), + ], +) +def test_integer_value_matches_legacy_behavior(value, runtime_draft_len): + _require_cuda() + worker = _make_worker(value) + out = worker._apply_force_accepted_tokens( + _make_input(batch_size=8), num_contexts=0, runtime_draft_len=runtime_draft_len + ) + expected = min(int(value) + 1, runtime_draft_len + 1) + assert torch.all(out == expected) + # The deterministic (non-random) path must not allocate the RNG pool. + assert worker._force_accept_rng_pool is None + + +def test_fractional_distribution_matches_target_probability(): + """Fractional draws emit only base/base+1 and average to the target rate.""" + _require_cuda() + worker = _make_worker(2.6) + target_frac = 0.6 + n_iters = 200 + batch = 64 + seen = set() + extra_count = 0 + total = 0 + for _ in range(n_iters): + out = worker._apply_force_accepted_tokens( + _make_input(batch_size=batch), num_contexts=0, runtime_draft_len=4 + ) + seen.update(out.unique().tolist()) + extra_count += int((out == 4).sum().item()) + total += batch + # Either 2 draft + target = 3, or 3 draft + target = 4 -- never anything else. + assert seen == {3, 4} + measured = extra_count / total + # Pool-based RNG is deterministic; the empirical mean over ~13k draws + # is tight against 0.6. + assert abs(measured - target_frac) < 0.03, ( + f"measured fraction {measured:.4f} differs from target {target_frac}" + ) + + +def test_num_contexts_offset_is_respected(): + """Context (prefill) rows must be left untouched by the override.""" + _require_cuda() + worker = _make_worker(2.6) + batch_size = 8 + num_contexts = 3 + base = _make_input(batch_size=batch_size) + sentinel = torch.iinfo(base.dtype).max + base[:num_contexts] = sentinel + out = worker._apply_force_accepted_tokens(base, num_contexts=num_contexts, runtime_draft_len=4) + assert torch.all(out[:num_contexts] == sentinel) + assert torch.all((out[num_contexts:] == 3) | (out[num_contexts:] == 4)) + + +# ---------------- TP determinism -------------------------------------------- + + +def test_tp_determinism_across_independent_workers(): + """ + Two independently-instantiated workers (simulating two TP ranks) must + produce bit-identical accepted-token counts when fed the same call + sequence. This is the property whose absence used to hang TP>1. + """ + _require_cuda() + rank0 = _make_worker(2.6) + rank1 = _make_worker(2.6) + # Use a non-power-of-two batch to make accidental shape coincidences + # less likely to mask divergence. + for _ in range(32): + out0 = rank0._apply_force_accepted_tokens( + _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 + ) + out1 = rank1._apply_force_accepted_tokens( + _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 + ) + assert torch.equal(out0, out1) + + +def test_tp_determinism_survives_default_generator_drift(): + """ + Even if the default CUDA generator state diverges between ranks (e.g. + because some other component consumed random numbers), the synthetic + AR draws must still match. This is the regression guard for the + original bug, where ``torch.rand`` against the default generator gave + each rank a different sequence. + """ + _require_cuda() + rank0 = _make_worker(2.6) + rank1 = _make_worker(2.6) + + def _advance_default_generator(seed: int): + # Seed the default CUDA generator on the same device differently + # for each "rank" and consume some randoms, simulating drift. + torch.cuda.manual_seed(seed) + _ = torch.rand(128, device="cuda") + + _advance_default_generator(seed=11) + out0 = rank0._apply_force_accepted_tokens( + _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 + ) + _advance_default_generator(seed=999) + out1 = rank1._apply_force_accepted_tokens( + _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 + ) + assert torch.equal(out0, out1) + + +# ---------------- CUDA graph capture/replay --------------------------------- + + +def test_cuda_graph_capture_and_replay_match_eager(): + """ + Capturing the override into a CUDA graph and replaying it must produce + the exact same per-replay output sequence as running it eagerly. + """ + _require_cuda() + + n_iters = 8 + batch_size = 16 + runtime_draft_len = 4 + value = 2.6 + + # 1) Eager reference. + eager_worker = _make_worker(value) + eager_outputs = [] + for _ in range(n_iters): + out = eager_worker._apply_force_accepted_tokens( + _make_input(batch_size=batch_size), num_contexts=0, runtime_draft_len=runtime_draft_len + ) + eager_outputs.append(out.clone()) + + # 2) Captured-graph run, aligned to the eager reference. + graph_worker = _make_worker(value) + static_input = _make_input(batch_size=batch_size) + # Eager warmup: forces lazy RNG state allocation OUTSIDE capture. + graph_worker._apply_force_accepted_tokens( + static_input, num_contexts=0, runtime_draft_len=runtime_draft_len + ) + # Realign the device-side counter so the captured graph's first replay + # advances 0 → 1, matching the eager loop's first iteration. + graph_worker._force_accept_rng_counter.zero_() + + static_input.fill_(1) + static_output = torch.empty_like(static_input) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_worker._apply_force_accepted_tokens( + static_input, num_contexts=0, runtime_draft_len=runtime_draft_len + ) + # ``_apply_force_accepted_tokens`` mutates ``static_input`` in place; + # mirror it into ``static_output`` so we can snapshot per replay. + static_output.copy_(static_input) + + graph_outputs = [] + for _ in range(n_iters): + # Reset input each replay so the override semantics match eager + # (which also starts from all-1s every iteration). + static_input.fill_(1) + graph.replay() + torch.cuda.synchronize() + graph_outputs.append(static_output.clone()) + + for i, (eager, graphed) in enumerate(zip(eager_outputs, graph_outputs)): + assert torch.equal(eager, graphed), ( + f"Iteration {i}: eager={eager.tolist()} vs graphed={graphed.tolist()}" + ) + + +def test_cuda_graph_replay_advances_rng_state(): + """ + Across replays, the captured counter must advance, producing different + (but deterministic) draws. A constant output across replays would mean + the counter increment was not captured. + """ + _require_cuda() + + worker = _make_worker(0.5) + batch_size = 64 + runtime_draft_len = 4 + + static_input = _make_input(batch_size=batch_size) + worker._apply_force_accepted_tokens( + static_input, num_contexts=0, runtime_draft_len=runtime_draft_len + ) + worker._force_accept_rng_counter.zero_() + + static_input.fill_(1) + static_output = torch.empty_like(static_input) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + worker._apply_force_accepted_tokens( + static_input, num_contexts=0, runtime_draft_len=runtime_draft_len + ) + static_output.copy_(static_input) + + snapshots = [] + for _ in range(8): + static_input.fill_(1) + graph.replay() + torch.cuda.synchronize() + snapshots.append(static_output.clone()) + + # At least two replays must differ; otherwise the captured counter is + # not advancing (the original bug-class). + assert any(not torch.equal(snapshots[0], s) for s in snapshots[1:]), ( + "Captured graph produced identical outputs on every replay — " + "RNG counter is not advancing inside the captured graph." + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) + + +class _FailingSpecWorker(_StubSpecWorker): + """Stub whose forward saves spec-dec metadata state and then fails.""" + + def _forward_impl(self, *args: object, **kwargs: object) -> None: + attn_metadata = kwargs["attn_metadata"] + attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") + raise RuntimeError("simulated draft failure") + + +def test_forward_restores_spec_dec_state_on_failure() -> None: + """A failure between prepare_for_spec_dec and restore must not leak saved + attn-metadata state; SpecWorkerBase.forward restores it in its cleanup + (https://nvbugs/6442074).""" + _require_cuda() + attn_metadata = AttentionMetadata(max_num_requests=2, max_num_tokens=16) + attn_metadata.seq_lens = torch.ones(2, dtype=torch.int32) + worker = _FailingSpecWorker() + for _ in range(2): + # The second iteration would trip the pairing assert inside + # prepare_for_spec_dec if the first failure had leaked saved state. + with pytest.raises(RuntimeError, match="simulated draft failure"): + worker(attn_metadata=attn_metadata, spec_metadata=None) + assert not attn_metadata.has_spec_dec_saved_state From 25f279e7709470d7dfd120a24f2f48475a2a1d0a Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:23:45 -0700 Subject: [PATCH 143/177] [None][test] Ragged verify config validation test Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../test_dspark_ragged_requires_table.py | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) create mode 100644 tests/unittest/llmapi/test_dspark_ragged_requires_table.py diff --git a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py new file mode 100644 index 000000000000..898699b69e91 --- /dev/null +++ b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Ragged verification must refuse to start without a profiled cost table. + +Without one the planner's budget degenerates to verify-all -- correctly, since +a flat cost model makes every extra verify token look free -- so the ragged +path runs and changes nothing. As a runtime warning that produced 27-minute +runs which "succeeded" having done nothing, detectable only from a counter +(``steps_ragged: 0``) that nobody reads unless they already suspect the +problem. Config errors belong at construction. + +There is no exception: a crafted (non-flat) cost table is how a test makes the +planner trim, which exercises the real confidence-driven path rather than +bypassing it. See tests/unittest/_torch/speculative/test_dspark_planner_trims.py. +""" + +import pytest + +from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig + + +def _cfg(**kwargs): + base = dict(max_draft_len=5, + speculative_model="/nonexistent/model", + enable_confidence_scheduling=True, + enable_ragged_verify=True) + base.update(kwargs) + return DSparkDecodingConfig(**base) + + +def test_ragged_without_a_cost_table_is_rejected(): + with pytest.raises(ValueError, match="requires a profiled cost table"): + _cfg() + # The same gate must open once a table is named -- a gate that rejects + # everything would pass the raise above. + assert _cfg(confidence_sps_table_path="/tmp/table.json").enable_ragged_verify + + +def test_confidence_scheduling_without_ragged_is_rejected(): + """There is no uniform tier ladder to fall back to. + + Scheduling used to have a middle state: pick one verify length for the + whole batch from the captured tiers. It could not act on per-request + confidence -- which is the entire point of a confidence head -- and it + degenerated to the full block whenever acceptance was high, i.e. on every + workload measured on this checkpoint. It is gone, so this combination now + names a path that does not exist and must not silently mean something else. + + The two remaining states are: schedule per request (ragged), or verify the + whole drafted block (confidence scheduling off). + """ + with pytest.raises(ValueError, match="enable_ragged_verify"): + DSparkDecodingConfig(max_draft_len=5, + speculative_model="/nonexistent/model", + enable_confidence_scheduling=True, + enable_ragged_verify=False) From 8874ddf6678e3f7b6218c13bc290764ddc7bbf78 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 20:43:40 -0700 Subject: [PATCH 144/177] [None][doc] #34 option C validated: explicit max_seq_len unpins the windowed pool (128 full batches, zero deadlock) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index bdc5894b15cf..db49a7ef6803 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -608,6 +608,18 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( **机制(算术闭合)**:DSv4 的窗口池组(SWA+COMPRESSOR_KV+COMPRESSOR_SCORE,窗口=128+max_draft_len=133 token)是固定槽分配器,槽数被钉在下限 `min_slots≈1.05×maxbs`——①比例定容用 `typical_seq_len=max_seq_len`(serve 推断 1M)把配额喂给全注意力池,窗口池饿到下限(kv_cache_manager_v2.py:1775-1779);②下限按 history=0 的"刚起步 decode"建模 → 每请求 1 块(kv_cache_manager_v2.py:1803-1812 → _storage_manager.py:909-1001)。运行时窗口 133>块 128 → 每 gen 请求锁 2(偶 3)块(_life_cycle_registry.py:38-51)。**容量 = 1.05·maxbs/2.07 ≈ 0.51·maxbs**:maxbs=128 预测 63-65/实测 61-63;maxbs=256 预测 131-135/实测 131-136。配额不敏感(下限无配额项)、kv_util 0.014 是 token 口径(窗口池实为 100% 锁死)、首波满批(短 prompt 1 块,h>128 才要第二块)、spec 降级 vs no-spec 死锁(接受方差去相位 vs +1/步全队锁相集体撞边界)——全部观测闭合。字节估算器还漏了 +draft 窗口放大(cache_manager.py:1316-1326),但字节不是绑定资源。 +**C 方案实验判决(08-10 20:45,判别象限全齐)**: + +| 象限(Pro@DEP8,poetry 1024 burst)| 每步调度 | 结局 | +|---|---|---| +| spec, maxbs=128 | 60-63 | 优雅降级 | +| spec, maxbs=256 | 115-128 | 健康 | +| no-spec, maxbs=128 | 61 | **死锁崩溃** | +| no-spec, maxbs=256 | 128×254 iter | 健康 | +| **no-spec, maxbs=128 + `max_seq_len: 8192`** | **128×230 iter** | **健康,零死锁** | + +`max_seq_len` 显式设为真实值即解钉(推断值 1M 经 typical_seq_len 把窗口池饿到下限)——**C 方案(部署侧)验证通过,按用户决定采用**。部署规则:DSv4 serving 必须显式设 `max_seq_len`(或 `avg_seq_len`),且 `max_batch_size ≥ 目标并发/rank`(设了 max_seq_len 后无需 2×)。v2 定容的代码级修复(min-slots 稳态建模 / typical_seq_len 默认值 / 池钉死告警)属上游 KV-v2 关注面,与本 PR 分离,另开 PR 跟进。 + **修复**:A(首选)min-slots 按稳态 decode(2-3 块/请求)建模;B 估算器补 draft(须配 A);C `typical_seq_len` 不默认 max_seq_len / 部署侧显式设 max_seq_len 解钉(判别实验在跑)。上游三报:下限建模、估算器漏项、no-spec 高并发死锁崩溃(报错信息误导)。附:python-scheduler 旗标在 v2 下是 no-op(_util.py:2618),DSv4 强制 v2 管理器——当晚两条二分腿均为无效对照,已修订。 ### 产品线终局判决:Pro@DEP8 置信度调度全面转正(08-10 13:50) From d7f3c402a8579182ce6b887149ab69f9504f2651 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:05:30 -0700 Subject: [PATCH 145/177] [None][doc] #34 experiment loop closed: 2x2 quadrants + explicit max_seq_len unpins the windowed pool (option C validated) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index db49a7ef6803..c3db1972bc4a 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -622,6 +622,18 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( **修复**:A(首选)min-slots 按稳态 decode(2-3 块/请求)建模;B 估算器补 draft(须配 A);C `typical_seq_len` 不默认 max_seq_len / 部署侧显式设 max_seq_len 解钉(判别实验在跑)。上游三报:下限建模、估算器漏项、no-spec 高并发死锁崩溃(报错信息误导)。附:python-scheduler 旗标在 v2 下是 no-op(_util.py:2618),DSv4 强制 v2 管理器——当晚两条二分腿均为无效对照,已修订。 +### #34 实验环闭合:2×2 象限 + C 方案判决(08-10 21:05) + +| 象限(Pro,poetry 1024 burst) | 结果 | +|---|---| +| spec, maxbs=128 | 稳态 ~60-63(优雅降级) | +| spec, maxbs=256 | 满批 115-128 | +| **no-spec, maxbs=128** | **61 处死锁崩溃**(scheduler_v2.py:439) | +| **no-spec, maxbs=256** | **满批 128,零死锁**(254 iter) | +| **no-spec, maxbs=128 + `max_seq_len: 8192`** | **满批 127-128,零死锁(254 iter)——C 方案判决命中** | + +结论:上限=窗口池按 min_slots 下限定容(比例定容被推断的 max_seq_len=1M 饿死)+ 运行时每请求锁 2 块;**显式设 max_seq_len 即解钉**(用户已拍板 C 方向)。部署规约:DSv4 serving 必须显式设 `max_seq_len`(或 `avg_seq_len`),否则窗口池被钉死在 ~maxbs/2 并发,no-spec 高并发直接死锁。代码侧 C 修复(typical_seq_len 不默认 max_seq_len)+ 上游三报(下限建模、估算器漏 draft、死锁而非降级)待落。 + ### 产品线终局判决:Pro@DEP8 置信度调度全面转正(08-10 13:50) Flash 三臂翻案后,Pro@DEP8 终局(maxbs=256 解锁满批 + `postfix_pro_table.json` 新表,scheduled vs notrim,双节点臂序对调,每 cell n=10): From 313662604a532c9038e9580a063fcc9b1a2ad155 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:09:15 -0700 Subject: [PATCH 146/177] [None][chore] Strip debug scaffolding and dev-history comments from the DSpark surface Removes the investigation-era bisectors and instrumentation (AB dump harness, invariant audit, ride/rowdiv logs, capture-order and eager bisectors, corruption-hunt sampling switches), dead helpers, and war-story comments across 15 files. Verified by a two-phase sweep (inventory + adversarial reference check); no runtime semantics change. hw_agnostic suite: 337 passed, same as before the sweep. Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/attention_backend/trtllm.py | 14 +- tensorrt_llm/_torch/models/dspark/draft.py | 4 +- tensorrt_llm/_torch/models/modeling_dspark.py | 16 +- .../_torch/pyexecutor/model_engine.py | 277 +----------------- .../speculative/dspark_observability.py | 7 +- tensorrt_llm/_torch/speculative/dspark_sts.py | 25 -- tensorrt_llm/_torch/speculative/interface.py | 8 +- .../_torch/speculative/spec_sampler_base.py | 56 +--- tensorrt_llm/evaluate/lm_eval.py | 22 +- tensorrt_llm/llmapi/llm_args.py | 6 +- .../dspark_make_steep_sps_table.py | 2 - tests/microbenchmarks/dspark_sps_profiler.py | 34 +-- .../test_deepseek_v4_sparse_mla.py | 76 +---- .../test_dspark_sts_collection_guards.py | 3 - .../thop/parallel/test_compressor_ragged.py | 11 +- 15 files changed, 32 insertions(+), 529 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 0438fbd85c6f..7e41d33a111c 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -565,14 +565,8 @@ def presented_token_major(self, view): # batch, not from here. # # The attention workspace is sized from the same number, so it - # grows with the row ceiling. An earlier version of this comment - # quoted 68 MiB -> 604 MiB per rank as the cost; that figure came - # from a single ragged log with no uniform control, and the one - # memory claim that WAS controlled -- CUDA graph pool size -- came - # out at parity (8.746 vs 8.740 GiB at max_batch_size 128), so the - # uncontrolled figure should not be relied on. The op resizes - # itself, so whatever the true delta is, it costs GPU memory that - # would otherwise hold KV cache rather than correctness. + # grows with the row ceiling; that memory would otherwise hold + # KV cache. self.max_num_requests = view.max_num_rows # Every runtime view has to agree on batch_size with # kv_lens_cuda_runtime, and the op reads request_types over all of @@ -1839,7 +1833,7 @@ def forward( if not self.fmha_libs: self.create_fmha_libs() - # goal doc K2. On a ragged generation step the attention op is handed + # On a ragged generation step the attention op is handed # one row per query token so its `num_tokens % num_seqs == 0` check # passes with seq_len == 1. Wrapping the dispatch rather than the # library keeps every fmha implementation, and its tests, untouched. @@ -2112,7 +2106,7 @@ def mla_rope_generation( metadata.helix_position_offsets, metadata.helix_is_inactive_rank ] - # goal doc K1. Same presentation as the attention dispatch above: the + # Same presentation as the attention dispatch above: the # rope kernel derives batch_idx = tok / seq_len and the KV write offset # from a single scalar seq_len, which only describes a batch where every # request contributes the same number of tokens. diff --git a/tensorrt_llm/_torch/models/dspark/draft.py b/tensorrt_llm/_torch/models/dspark/draft.py index 8aa3df8016b5..a70ccc503155 100644 --- a/tensorrt_llm/_torch/models/dspark/draft.py +++ b/tensorrt_llm/_torch/models/dspark/draft.py @@ -87,8 +87,8 @@ def dspark_propose( """ # ``draft_logits`` are the per-position distributions the draft token is drawn # from (markov-corrected when a head is present, else the raw base logits). - # Surfaced under ``return_logits`` for the §7.9 probabilistic-acceptance - # (1-TV) measurement; the normal path ignores them. + # Surfaced under ``return_logits`` to feed the rejection-sampling verifier; + # the normal path ignores them. draft_logits = base_logits if markov_head is not None: draft_tokens, corrected = markov_head.sample_block_tokens( diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 60761ee6003a..cdd3097a3dec 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -146,15 +146,6 @@ def remap_dspark_draft_keys(weights: Dict, num_stages: int) -> Dict: return out -# The checkpoint stores -# ``mtp.{s}.attn.wo_a`` as fp8_e4m3 + a UE8M0 128x128 block scale (verified), and -# the reference (`inference/model.py`, ``self.wo_a`` is a bf16 ColumnParallelLinear -# loaded from the fp8 ckpt) uses the DEQUANTIZED bf16 ``wo_a`` (== ``wo_a_fp8 * -# scale`` ~ absmean 0.065). The bf16 captured-context path historically skipped -# this dequant (raw fp8-cast-to-bf16, ~993x too large); the correct behavior is to -# dequantize ``wo_a`` (cos 1.0 vs ``wo_a_fp8 * scale``). Always dequantize now. - - class DSparkBlock(DeepseekV4DecoderLayer): """One DSpark draft stage = a DeepSeek-V4 decoder block + DSpark extras. @@ -833,10 +824,7 @@ def _forward_stage( # distinct set of requests) and needs no cross-rank reduction; but under # plain tensor parallelism (attention_dp off, tp_size > 1) the expert-sharded # MoE output must be all-reduced across ranks -- exactly what the target MoE - # does (enable_allreduce = not (POST_MOE_FUSION or tp_size == 1)). Previously - # this was hard-coded to False, which dropped the reduction on the non-ADP - # path, corrupting the draft block proposals and roughly halving DSpark - # acceptance length whenever attention_dp was disabled. + # does (enable_allreduce = not (POST_MOE_FUSION or tp_size == 1)). moe_enable_allreduce = ( not self.model_config.mapping.enable_attention_dp and self.model_config.mapping.tp_size > 1 @@ -1036,7 +1024,7 @@ def forward_head( ``block_hidden`` is the last stage's mHC residual ``[*, block, hc_mult, hidden]``. Returns (draft_tokens [*, block], confidence [*, block] or None); with ``return_logits`` - also returns the per-position draft logits [*, block, vocab] (§7.9 1-TV). + also returns the per-position draft logits [*, block, vocab]. """ last = self.mtp_layers[-1] h = last.hc_head(block_hidden) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index ddbbdf070bf3..ff72f570c571 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2027,8 +2027,6 @@ def _run_capture_pass(force_non_greedy: bool, # right after an advanced one would otherwise key on the stale # False before populate corrects it. spec_metadata.is_all_greedy_sample = True - _graph_mem_before = (torch.cuda.memory_reserved(), - torch.cuda.memory_allocated()) _graph_count = 0 try: for entry in (graphs_to_capture @@ -2103,32 +2101,9 @@ def _run_capture_pass(force_non_greedy: bool, # the LAST shape's draft length; restore it after the passes, before # the generic token warmup builds requests at the full draft length. saved_runtime_draft_len = self.runtime_draft_len - needs_non_greedy_precheck = ( - self.spec_config is not None - and self.spec_config.spec_dec_mode.use_one_engine() - and os.environ.get("TLLM_DSPARK_SKIP_ADVANCED_CAPTURE") != "1") - if (needs_non_greedy_precheck and os.environ.get( - "TLLM_DSPARK_INTERLEAVED_CAPTURE") == "1"): - # Debug/fix candidate: capture both sampling variants of each - # shape back-to-back, so the two graphs bind under identical - # pool/workspace/persistent-buffer conditions instead of a whole - # sweep apart. - for _entry in graphs_to_capture: - _run_capture_pass(force_non_greedy=False, - label="greedy", - entries=[_entry]) - _run_capture_pass(force_non_greedy=True, - label="advanced sampling", - entries=[_entry]) - self.enable_spec_decode = self.is_spec_decode - self.runtime_draft_len = saved_runtime_draft_len - if self.spec_metadata is not None: - self.spec_metadata.is_all_greedy_sample = True - return # Pass 1: greedy fast-path (dummy requests carry no sampling params, # so is_all_greedy_sample is naturally True). - if os.environ.get("TLLM_DSPARK_SKIP_GREEDY_CAPTURE") != "1": - _run_capture_pass(force_non_greedy=False, label="greedy") + _run_capture_pass(force_non_greedy=False, label="greedy") # Pass 2: advanced sampling variant. Required because on-the-fly capture # is disabled outside warmup, so any inference batch that contains a # non-greedy request would otherwise fall back to eager. Only meaningful @@ -2137,8 +2112,7 @@ def _run_capture_pass(force_non_greedy: bool, # this variant. needs_non_greedy_capture = ( self.spec_config is not None - and self.spec_config.spec_dec_mode.use_one_engine() - and os.environ.get("TLLM_DSPARK_SKIP_ADVANCED_CAPTURE") != "1") + and self.spec_config.spec_dec_mode.use_one_engine()) if needs_non_greedy_capture: _run_capture_pass(force_non_greedy=True, label="advanced sampling") # Set the value back to the original value after cuda graph warmups are complete @@ -3861,10 +3835,9 @@ def fit_ragged_verify_lens( # Published only now, past the last way this fit can fail: a stale # published bucket on a fallback step is the state behind the ragged - # IMA. `_dspark_last_total_tokens`/`_dspark_last_padded_bs` feed the + # IMA. `_dspark_last_padded_bs` feeds the # stats; `ragged_pad_verify_len` is stamped on the shared dummy by # _get_padded_batch. - self._dspark_last_total_tokens = int(bucket) self._dspark_last_padded_bs = int(padded_bs) self._dspark_last_num_real = int(n_real) runner.agreed_ragged_bucket = int(bucket) @@ -3953,9 +3926,6 @@ def _apply_device_window_prologue(self, inputs, split_lens = lens_buf[:padded_bs].clone() split_qo = qo_buf[:padded_bs + 1].to(torch.long) past_seen = self.position_ids_cuda[split_qo[:-1]].clone() - if os.environ.get("TLLM_DSPARK_AB_DUMP") == "1": - self._dbg_pre_kv = attn_metadata.kv_lens_cuda[:padded_bs].clone() - self._dbg_split = split_lens.clone() expected_stamp = worker.verified_draft_seq_cuda() stamps = worker.confidence_stamp_buffer( @@ -3972,33 +3942,6 @@ def _apply_device_window_prologue(self, inputs, expected_stamp=expected_stamp, pad_len=pad_len_tok, ) - debug_mode = os.environ.get("TLLM_DSPARK_DEVICE_IDENTITY", "") - if debug_mode in ("1", "steep"): - # Debug bisectors. "1": exercise every prologue write with the - # shape split itself (w := S) -- corruption surviving this is in - # the rewrite mechanics, not the window semantics. "steep": - # deterministic max-amplitude ragged windows (adjacent real rows - # trade tokens pairwise, conserving the total) -- forces w != S - # while staying independent of the confidence-gather chain. - from ..speculative.dspark_ragged import (build_qo_indptr, - build_row_maps_device) - id_lens = split_lens.clone() - if debug_mode == "steep" and n_real > 1: - lo = cfg.min_verify_len + 1 - hi = cfg.resolved_max_verify_len + 1 - pairs = (n_real // 2) * 2 - a = id_lens[0:pairs:2] - b = id_lens[1:pairs:2] - give = torch.minimum(a - lo, hi - b).clamp_min(0) - a -= give - b += give - id_req, id_corr = build_row_maps_device(id_lens, - graph_num_tokens=bucket) - result = type(result)(verify_lens=id_lens, - qo_indptr=build_qo_indptr(id_lens).to( - result.qo_indptr.dtype), - req_idx=id_req, - kv_correction=id_corr) lens_buf[:padded_bs].copy_(result.verify_lens) qo_buf[:padded_bs + 1].copy_(result.qo_indptr) @@ -4056,181 +3999,8 @@ def _apply_device_window_prologue(self, inputs, attn_metadata.apply_device_ragged_layout(result.verify_lens, req_idx, result.kv_correction) - if os.environ.get("TLLM_DSPARK_RIDE_LOG") == "1": - self._ride_log_ctr = getattr(self, "_ride_log_ctr", 0) + 1 - if self._ride_log_ctr % 64 == 0: - logger.info( - f"[ride-log] prologue #{self._ride_log_ctr} " - f"S={split_lens[:4].tolist()} w={result.verify_lens[:4].tolist()}") - if (os.environ.get("TLLM_DSPARK_INVARIANT_LOG") == "1" - and getattr(self, "_dspark_inv_prints", 0) < 40): - # Read-only invariant audit against host-with-w formulas: no - # restage, no request mutation, no collectives -- safe at scale. - torch.cuda.synchronize() - gen_requests = [ - r for r in (self._dspark_inv_requests or []) - ] if getattr(self, "_dspark_inv_requests", None) else [] - w = result.verify_lens[:n_real].tolist() - qo = result.qo_indptr[:n_real + 1].tolist() - kv = attn_metadata.kv_lens_cuda[:n_real].tolist() - kvoff = self.previous_kv_lens_offsets_cuda[:n_real].tolist() - slots_l = prev_slots.tolist() - ids = self.input_ids_cuda[:bucket].tolist() - drafts_l = self.draft_tokens_cuda[:bucket].tolist() - ppi = self.previous_pos_indices_cuda[:bucket].tolist() - ppo = self.previous_pos_id_offsets_cuda[:bucket].tolist() - anchors = new_tokens_device[0].squeeze(-1).tolist() - nlens = new_tokens_lens_device.tolist() - ndraft = next_draft_tokens_device.tolist() - self._dspark_inv_step = getattr(self, "_dspark_inv_step", 0) + 1 - - def _viol(i, msg): - self._dspark_inv_prints = getattr(self, "_dspark_inv_prints", - 0) + 1 - logger.warning( - f"[invariant] step#{self._dspark_inv_step} row {i}: {msg} " - f"(w={w[i]} S={split_lens[i].item()} n_real={n_real} " - f"padded_bs={padded_bs} bucket={bucket})") - - # Batch 3: layout-derived tensors recomputed analytically from w. - from ..speculative.dspark_ragged import build_row_maps_device - want_req, want_corr = build_row_maps_device( - result.verify_lens[:padded_bs], graph_num_tokens=bucket) - if not torch.equal(attn_metadata.row_req_idx_cuda[:bucket], - want_req): - _viol(-1, "row_req_idx mismatch vs analytic recompute") - if not torch.equal(attn_metadata.row_kv_correction_cuda[:bucket], - want_corr): - _viol(-1, "row_kv_correction mismatch vs analytic recompute") - seq_l = attn_metadata.seq_lens_cuda[:n_real].tolist() - reps_l = attn_metadata.gen_token_repeats_cuda[:n_real].tolist() - pos = self.position_ids_cuda[:bucket].tolist() - for i, req in enumerate(gen_requests[:n_real]): - past = int(req.max_beam_num_tokens) - 1 - s = slots_l[i] - if seq_l[i] != w[i]: - _viol(i, f"seq_lens={seq_l[i]} want w={w[i]}") - if reps_l[i] != w[i]: - _viol(i, f"repeats={reps_l[i]} want w={w[i]}") - if pos[qo[i] + w[i] - 1] - pos[qo[i]] != w[i] - 1: - _viol(i, f"position packing: pos[last]-pos[first]=" - f"{pos[qo[i]+w[i]-1]-pos[qo[i]]} want {w[i]-1}") - if kv[i] != past + 2 * w[i]: - _viol(i, f"kv={kv[i]} want past+2w={past + 2 * w[i]}") - if kvoff[i] != nlens[s] - w[i]: - _viol(i, f"kvoff={kvoff[i]} want nlens-w={nlens[s]-w[i]}") - if ids[qo[i]] != anchors[s]: - _viol(i, f"anchor ids[{qo[i]}]={ids[qo[i]]} " - f"want new_tokens[0,{s}]={anchors[s]}") - if w[i] > 1 and drafts_l[qo[i] - i] != ndraft[s][0]: - _viol(i, f"draft0 {drafts_l[qo[i]-i]} want {ndraft[s][0]}") - for o in (0, w[i] - 1): - if ppi[qo[i] + o] != s: - _viol(i, f"prev_pos_idx[{qo[i]+o}]={ppi[qo[i]+o]} " - f"want slot {s}") - if ppo[qo[i] + o] != nlens[s]: - _viol(i, f"prev_pos_off[{qo[i]+o}]={ppo[qo[i]+o]} " - f"want nlens={nlens[s]}") - if self._dspark_inv_step % 100 == 0: - logger.info( - f"[invariant] step#{self._dspark_inv_step} checkpoint: " - f"{getattr(self, '_dspark_inv_prints', 0)} violations") return True - def _dspark_ab_diff(self, padded_requests, kv_cache_manager, attn_metadata, - spec_metadata, new_tensors_device, - cache_indirection_buffer, num_accepted_tokens_device, - req_id_to_old_request, resource_manager, - can_run_graph) -> None: - """Debug: diff the prologue's rewritten staging against a host restage. - - Runs once, eager-only. Snapshot every tensor the forward reads after - the prologue rewrote the layout to the device windows w; then stamp w - onto each request's py_verify_len and rerun the full host staging (the - reference implementation, proven correct for non-uniform windows); - diff. Any mismatch names the artifact the prologue fails to mirror. - """ - torch.cuda.synchronize() - bucket = int(self.cuda_graph_runner.agreed_ragged_bucket or 0) - padded_bs = int(self._dspark_last_padded_bs) - gen_requests = padded_requests.generation_requests - n_real = int(self._dspark_last_num_real) - w_tokens = self.ragged_verify_lens_cuda[:padded_bs].tolist() - - def snap(): - fields = { - "input_ids": self.input_ids_cuda[:bucket], - "position_ids": self.position_ids_cuda[:bucket], - "draft_tokens": self.draft_tokens_cuda[:bucket], - "prev_pos_indices": self.previous_pos_indices_cuda[:bucket], - "prev_pos_offsets": self.previous_pos_id_offsets_cuda[:bucket], - "prev_kv_offsets": self.previous_kv_lens_offsets_cuda[:padded_bs], - "kv_lens": attn_metadata.kv_lens_cuda[:padded_bs], - "seq_lens_cuda": attn_metadata.seq_lens_cuda[:padded_bs], - "verify_lens": self.ragged_verify_lens_cuda[:padded_bs], - "qo_indptr": self.ragged_qo_indptr_cuda[:padded_bs + 1], - } - for name in ("gen_token_repeats_cuda", "row_req_idx_cuda", - "row_kv_correction_cuda", "attn_row_req_idx_cuda", - "attn_row_kv_correction_cuda", "req_idx_per_token", - "kv_lens_expanded_cuda", "block_table_expanded", - "slot_mapping_fp8", "gen_kv_indptr", - "gen_cached_token_indptr"): - t = getattr(attn_metadata, name, None) - if isinstance(t, torch.Tensor): - n = bucket if t.shape[0] >= bucket else t.shape[0] - fields[name] = t[:n] - for name in ("temperatures", "top_ks", "top_ps", "batch_slot_ids"): - t = getattr(spec_metadata, name, None) - if isinstance(t, torch.Tensor): - n = min(bucket, t.shape[0]) - fields["spec." + name] = t[:n] - return {k: v.detach().clone() for k, v in fields.items()} - - snap_a = snap() - prev_idx = getattr(self, "_dspark_ab_prev_idx", {}) - for request, tokens in zip(gen_requests, w_tokens): - request.py_verify_len = int(tokens) - 1 - if request.py_request_id in prev_idx: - request.py_batch_idx = prev_idx[request.py_request_id] - try: - self._prepare_inputs(padded_requests, kv_cache_manager, - attn_metadata, spec_metadata, - new_tensors_device, cache_indirection_buffer, - num_accepted_tokens_device, - req_id_to_old_request, resource_manager, - can_run_graph) - except Exception as exc: - logger.warning(f"[ab-diff] restage failed, disabling: {exc}") - self._dspark_ab_prints = 999 - return - torch.cuda.synchronize() - snap_b = snap() - self._dspark_ab_step = getattr(self, "_dspark_ab_step", 0) + 1 - # Expected-divergent names: staged from the shape split at prepare - # and rebuilt in-graph from the corrected kv_lens before any read. - benign = {"kv_lens_expanded_cuda", "slot_mapping_fp8"} - for name, a in snap_a.items(): - b = snap_b[name] - if name in benign or a.shape != b.shape: - continue - neq = (a != b) - cnt = int(neq.sum()) - if cnt: - idx = int(neq.reshape(-1).nonzero()[0]) - self._dspark_ab_prints = getattr(self, "_dspark_ab_prints", - 0) + 1 - logger.info( - f"[ab-diff] step#{self._dspark_ab_step} {name}: " - f"{cnt}/{a.numel()} MISMATCH first@{idx} " - f"prologue={a.reshape(-1)[idx].item()} " - f"host={b.reshape(-1)[idx].item()} " - f"(bs={padded_bs} n_real={n_real} bucket={bucket})") - if self._dspark_ab_step % 200 == 0: - logger.info(f"[ab-diff] step#{self._dspark_ab_step} checkpoint: " - f"{getattr(self, '_dspark_ab_prints', 0)} mismatch " - f"lines so far") - @staticmethod def _ragged_token_lens(generation_requests) -> Optional[List[int]]: """Each generation request's token window, or None if the batch is uniform. @@ -7332,13 +7102,6 @@ def forward(self, spec_resource_manager=spec_resource_manager, ) - if (key is not None and not self.is_warmup and os.environ.get( - "TLLM_DSPARK_FORCE_EAGER") == "1"): - # Debug bisector: capture normally (the ragged fit needs the - # captured bucket set) but veto every replay, so the exact - # same staged step runs eagerly. Splits capture-interaction - # bugs from layout-semantics bugs. - key = None can_run_graph = key is not None # A ragged step that misses its captured shape quietly runs eager; # count it. Only generation-only steps count: a batch with context @@ -7368,15 +7131,6 @@ def forward(self, padded_requests.generation_requests, ) - if os.environ.get("TLLM_DSPARK_AB_DUMP") == "1": - # The staging loop overwrites py_batch_idx (prev-batch index - # -> seq slot); the A/B restage must see the original or its - # reference staging reads the wrong new_tokens rows on any - # slot-churn step. - self._dspark_ab_prev_idx = { - r.py_request_id: r.py_batch_idx - for r in padded_requests.generation_requests - } inputs, gather_ids = self._prepare_inputs( padded_requests, kv_cache_manager, attn_metadata, spec_metadata, new_tensors_device, cache_indirection_buffer, @@ -7390,31 +7144,10 @@ def forward(self, # replay (stream-ordered, no host sync). Graph steps only: eager # and capture steps run the host shape split, which is a valid # window assignment. - if ((can_run_graph or os.environ.get( - "TLLM_DSPARK_DEVICE_EAGER_OK") == "1") - and not self.is_warmup + if (can_run_graph and not self.is_warmup and getattr(self, "_dspark_device_budget", None) is not None): - if os.environ.get("TLLM_DSPARK_INVARIANT_LOG") == "1": - # Read-only reference for the prologue's invariant audit. - self._dspark_inv_requests = list( - padded_requests.generation_requests) - applied = self._apply_device_window_prologue( - inputs, new_tensors_device) - if applied and os.environ.get("TLLM_DSPARK_AB_DUMP") == "1": - self._dspark_ab_applied = getattr( - self, "_dspark_ab_applied", 0) + 1 - if (applied and os.environ.get("TLLM_DSPARK_AB_DUMP") == "1" - and getattr(self, "_dspark_ab_applied", 0) >= int( - os.environ.get("TLLM_DSPARK_AB_AFTER", "0")) - and getattr(self, "_dspark_ab_prints", 0) < 40): - self._dspark_ab_diff(padded_requests, kv_cache_manager, - attn_metadata, spec_metadata, - new_tensors_device, - cache_indirection_buffer, - num_accepted_tokens_device, - req_id_to_old_request, - resource_manager, can_run_graph) + self._apply_device_window_prologue(inputs, new_tensors_device) with with_shared_pool(self.cuda_graph_runner.get_graph_pool()): if not can_run_graph: diff --git a/tensorrt_llm/_torch/speculative/dspark_observability.py b/tensorrt_llm/_torch/speculative/dspark_observability.py index 143300825977..4fcce61c6016 100644 --- a/tensorrt_llm/_torch/speculative/dspark_observability.py +++ b/tensorrt_llm/_torch/speculative/dspark_observability.py @@ -25,7 +25,7 @@ import os from collections import Counter from enum import Enum -from typing import Dict, List, Optional, Sequence, Tuple +from typing import Dict, Optional, Sequence, Tuple from ...logger import logger @@ -398,8 +398,3 @@ def summary(self) -> Dict[str, object]: def log_summary(self, *, prefix: str = "DSpark ragged verify") -> None: summary = self.summary() logger.info(f"{prefix}: {summary}") - -def format_verify_len_histogram(lens: List[int]) -> str: - """One-line ``len:count`` rendering, for per-step debug logging.""" - hist = Counter(int(v) for v in lens) - return " ".join(f"{k}:{v}" for k, v in sorted(hist.items())) diff --git a/tensorrt_llm/_torch/speculative/dspark_sts.py b/tensorrt_llm/_torch/speculative/dspark_sts.py index b0c67c785a97..40f3d50afc2b 100644 --- a/tensorrt_llm/_torch/speculative/dspark_sts.py +++ b/tensorrt_llm/_torch/speculative/dspark_sts.py @@ -24,7 +24,6 @@ import json import os -from dataclasses import dataclass, field from pathlib import Path from typing import List, Optional @@ -35,7 +34,6 @@ __all__ = [ "STS_COLLECT_ENV", - "DSparkStsCalibration", "DSparkStsRecorder", "load_sts_temperatures_from_path", "make_recorder_from_env", @@ -49,29 +47,6 @@ _DEFAULT_FLUSH_EVERY = 64 -@dataclass -class DSparkStsCalibration: - """A fitted per-position temperature vector, plus how well it did.""" - - temperatures: List[float] - ece_before: List[float] = field(default_factory=list) - ece_after: List[float] = field(default_factory=list) - dataset: str = "" - num_samples: int = 0 - - def to_json(self) -> dict: - # Both key spellings on purpose: this repo reads ``sts_temperatures``, - # SGLang uses ``temperatures``; tables stay interchangeable. - return { - "sts_temperatures": list(self.temperatures), - "temperatures": list(self.temperatures), - "ece_before": list(self.ece_before), - "ece_after": list(self.ece_after), - "dataset": self.dataset, - "num_samples": int(self.num_samples), - } - - def load_sts_temperatures_from_path(path: str) -> List[float]: """Read a temperature vector, accepting either spelling of the key. diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 0e5bd5ae309d..342dd89bdd3c 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -984,11 +984,6 @@ def populate_sampling_params_for_one_model( # read. Skip the heavier H->D copies. if self.is_all_greedy_sample: return - if os.environ.get("TLLM_DSPARK_SKIP_SAMPLING_H2D") == "1": - # Debug bisector (pair with TLLM_DSPARK_FORCE_ARGMAX, which stops - # anything reading these buffers): isolates whether the per-step - # non-greedy H2D staging itself is the corruption channel. - return # Phase 2: build per-token / per-request lists and copy to GPU. temperatures: list[float] = [] @@ -2353,8 +2348,7 @@ def _sample_tokens_for_batch( Returns: sampled_tokens: [num_tokens] - Sampled token ids """ - if not spec_metadata.is_all_greedy_sample and os.environ.get( - "TLLM_DSPARK_FORCE_ARGMAX") != "1": + if not spec_metadata.is_all_greedy_sample: # Use logits.shape[0] directly: for PARD under CUDA graph capture # runtime_draft_len may reflect the PARD-max while the captured # graph was built for a shorter draft_len, causing a shape mismatch diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index e45496ebb8b5..95996ca16f62 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -24,8 +24,6 @@ import torch -import os - from tensorrt_llm.logger import logger from ..pyexecutor.llm_request import LlmRequest, LlmRequestState @@ -291,55 +289,11 @@ def update_requests( beam_idx = DEFAULT_BEAM_IDX runtime_draft_len = getattr(state, "runtime_draft_len", self.draft_len) - if os.environ.get("TLLM_DSPARK_ROWDIV_LOG") == "1": - # Identical-request divergence detector: under a census load all - # requests share one prompt, so at temp 0 every row's accepted - # tokens must match its peers each step (windows change how MUCH - # is verified, never WHAT is accepted). The first divergent step - # names the corrupted rows and their window class. - sigs = {} - for _req in state.requests: - if _req.state == LlmRequestState.GENERATION_COMPLETE: - continue - _s = _req.py_seq_slot - _n = new_tokens_lens_list[_s] - _toks = tuple(new_tokens[i][_s][0] for i in range(_n)) - _w = (ridden_verify_lens[_s] - if ridden_verify_lens is not None else -1) - sigs.setdefault(_toks, []).append((_s, _w)) - self._rowdiv_step = getattr(self, "_rowdiv_step", 0) + 1 - if len(sigs) > 1 and getattr(self, "_rowdiv_prints", 0) < 30: - self._rowdiv_prints = getattr(self, "_rowdiv_prints", 0) + 1 - desc = "; ".join( - f"{len(v)}rows(w={sorted(set(x[1] for x in v))})" - f"toks={k[:3]}" for k, v in list(sigs.items())[:4]) - logger.warning( - f"[rowdiv] step#{self._rowdiv_step} {len(sigs)} distinct " - f"token-tuples across identical requests: {desc}") - if self._rowdiv_step % 100 == 0: - logger.info(f"[rowdiv] step#{self._rowdiv_step} checkpoint: " - f"{getattr(self, '_rowdiv_prints', 0)} " - f"divergent steps") - caps = state.verify_caps_snapshot for req in state.requests: if req.state == LlmRequestState.GENERATION_COMPLETE: continue num_new_tokens = new_tokens_lens_list[req.py_seq_slot] - if os.environ.get("TLLM_DSPARK_RIDE_LOG") == "1" and ( - ridden_verify_lens is not None): - _win = ridden_verify_lens[req.py_seq_slot] - if _win > 0 and num_new_tokens > _win: - logger.warning( - f"[accept-bound] slot={req.py_seq_slot} accepted+1=" - f"{num_new_tokens} > token window {_win} -- in-graph " - f"acceptance ran past this row's layout window") - elif _win == 0: - logger.warning( - f"[accept-bound] slot={req.py_seq_slot} ride window=0 " - f"on a live request -- rewind falling back to the " - f"S snapshot (py_verify_len=" - f"{getattr(req, 'py_verify_len', None)})") # cap-accept: `num_new_tokens` has ALREADY been capped, on device, # inside acceptance (`interface.apply_accept_caps`) -- it has to be, # because the drafter reads the same count later in that forward to @@ -465,8 +419,7 @@ def sample_async( # cap-accept: absent on every other path, and written as ZEROS then # rather than skipped. The buffer is persistent and slot-indexed, so a # slot left untouched would keep whatever the previous occupant lost - # and hand it to the current one -- the stale-buffer bug this branch - # has already paid for twice. + # and hand it to the current one. o_cap_trim = outputs.get("cap_trim_lens") if o_cap_trim is None: o_cap_trim = torch.zeros_like(o_new_tokens_lens) @@ -486,13 +439,6 @@ def sample_async( num_sampling_requests] self.store.verify_lens.index_copy_( 0, slots, o_verify_lens.to(self.store.verify_lens.dtype)) - if os.environ.get("TLLM_DSPARK_RIDE_LOG") == "1": - self._ride_log_ctr = getattr(self, "_ride_log_ctr", 0) + 1 - if self._ride_log_ctr % 64 == 0: - logger.info( - f"[ride-log] sampler #{self._ride_log_ctr} " - f"lens_buf[:4]={o_verify_lens[:4].tolist()} " - f"slots[:4]={slots[:4].tolist()}") # Create sample state with async D2H copy device_tensors = SampleStateTensorsSpec( diff --git a/tensorrt_llm/evaluate/lm_eval.py b/tensorrt_llm/evaluate/lm_eval.py index f0729fdc4538..adcc76a2945a 100644 --- a/tensorrt_llm/evaluate/lm_eval.py +++ b/tensorrt_llm/evaluate/lm_eval.py @@ -56,9 +56,7 @@ def _is_metric_key(key: str) -> bool: lm-eval keys every metric as ``","`` -- ``"exact_match, strict-match"``, ``"acc,none"``. Entries without a comma are bookkeeping: - ``"alias"`` is the task name, ``"samples"`` is the example count. Both had - been treated as scores, which scaled gsm8k's 1319 samples to 131900 and - averaged it into the reported accuracy. + ``"alias"`` is the task name, ``"samples"`` is the example count. """ return "," in key @@ -647,13 +645,7 @@ def evaluate(self, system_instruction=self.system_prompt, log_samples=self.log_samples) - # Normalize scores to range 0~100. - # - # Only real metrics. lm-eval keys a metric as "," (for - # example "exact_match,strict-match"); everything without a comma is - # bookkeeping -- "alias" holds the task name, "samples" holds the - # example count. Scaling those too turned gsm8k's 1319 samples into - # 131900, which then dominated the average below. + # Normalize real metric scores to range 0~100 (see _is_metric_key). scores = results["results"][self.task_name] for metric in scores.keys(): if _is_metric_key(metric) and isinstance(scores[metric], @@ -673,14 +665,8 @@ def evaluate(self, f"lm-eval {self.task_name} {scores_filter} accuracy: {result_acc:.2f}" ) else: - # Average the metrics only. The score dict also carries - # bookkeeping: "alias" is the task name (a string, which made - # numpy fail with "resolved dtypes are not compatible with - # add.reduce"), and "samples" is the example count -- for gsm8k - # that is 1319, and averaging it in produced 44030.98 against a - # reference of 96, which *passed* the accuracy assertion. A wrong - # number that clears the threshold is worse than the crash it - # replaced, so both are excluded structurally rather than by name. + # Average real metrics only; bookkeeping keys ("alias", + # "samples") are excluded structurally (see _is_metric_key). numeric = [ acc for m, acc in scores.items() if _is_metric_key(m) and "_stderr" not in m diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 126ca66122b7..45ecf39443a0 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2893,11 +2893,7 @@ def validate_ragged_verify(self): and not self.confidence_sts_path): # Without a profiled cost table the planner's budget degenerates to # verify-all -- correctly, since a flat model makes every extra - # token look free -- so the ragged path silently never runs. That - # used to be a runtime warning, which means a 27-minute run - # "succeeds" having done nothing and the only evidence is a counter - # nobody reads. Fail at construction instead. - # + # token look free -- so the ragged path silently never runs. raise ValueError( "enable_ragged_verify=True requires a profiled cost table: pass " "confidence_sps_table_path (produced by " diff --git a/tests/microbenchmarks/dspark_make_steep_sps_table.py b/tests/microbenchmarks/dspark_make_steep_sps_table.py index b22260c5b90e..db8a4ae57075 100644 --- a/tests/microbenchmarks/dspark_make_steep_sps_table.py +++ b/tests/microbenchmarks/dspark_make_steep_sps_table.py @@ -201,8 +201,6 @@ def main() -> int: "have the real planner check the table before you spend a run " "on it", file=sys.stderr) any_trim = True - print(f" bs={bs:<5} -> K={best_k}" - f"{' TRIM' if trims else ' (full window)'}") if not any_trim: print() print("WARNING: this table induces no trimming at any previewed batch " diff --git a/tests/microbenchmarks/dspark_sps_profiler.py b/tests/microbenchmarks/dspark_sps_profiler.py index 934986f21143..c44a9a2b41f1 100644 --- a/tests/microbenchmarks/dspark_sps_profiler.py +++ b/tests/microbenchmarks/dspark_sps_profiler.py @@ -1233,13 +1233,6 @@ def validate(self) -> None: raise SweepGeometryError("both --batch-sizes and --verify-lens must be non-empty") if min(self.verify_lens) < 1: raise SweepGeometryError("verify lengths must be >= 1 (the planner's floor)") - # Constant-block with a pinned window below the block used to die on a - # device-side assert (Repeat.cu output_size vs sum(repeats)); a guard - # here refused that geometry. The assert was another symptom of the - # graph key describing a different token width than the batch carried, - # fixed by publishing the fit's bucket into the key; verified gone - # 2026-08-06 by pinned L=3 runs under both real and synthetic - # acceptance (64+ steps each, no assert). # A request that hits max_seq_len is evicted mid-cell, so the batch # drains and the steady window shrinks -- which surfaces much later as a # confusing "not enough steady samples". @@ -1338,20 +1331,6 @@ def _build_llm(config: SweepConfig, verify_len: int, *, # pinned -- a rank-divergent shape under attention-DP. With one verify # length per sweep process the pin always equals the ladder rung below, so # it can never be rejected either. - # - # (An earlier revision popped it and relied on the one-rung ladder alone, - # citing a "pin skips the padding reservation" theory for the Repeat.cu - # assert. That theory was wrong -- the assert reproduced with no pin at - # all; the actual cause was the CUDA-graph key describing a - # different token width than the batch carried. The key publishes the - # fit's bucket (so attention-DP ranks agree by construction) and the - # layout assert cross-checks it against the batch walk; a divergence now - # has to show up there rather than inside a replayed graph. And the one-rung ladder was - # never one rung anyway: the llm_args validator appends max_draft_len, so - # under a constant block [L] is really [L, block]. With a seed table - # present the planner argmax'd its way to the block on 93-95% of steps - # while the samples file recorded M = bs*(L+1) as if L had run -- twelve - # cells, eight mislabelled, and the fit dutifully concluded q ~ 0.) from tensorrt_llm._torch.speculative.dspark_verify import FORCE_VERIFY_LEN_ENV # LAUNCH-MODE CAVEAT: this mutation reaches ranks 1..N-1 only when the @@ -1373,18 +1352,14 @@ def _build_llm(config: SweepConfig, verify_len: int, *, # # constant block (default): the block stays at the checkpoint's trained # length for every cell and only the verify window moves, which is exactly - # what the deployment does. The all-short-window regime this creates used - # to trip the Repeat.cu output_size assert; that was the graph-key/batch - # width divergence, fixed with the published-bucket key and verified gone - # 2026-08-06. + # what the deployment does. # # block follows verify_len: the original design. Never produces # max(lens) < block, so it runs today, at the price of shrinking the DRAFT # pass along with the window -- theta(M) then absorbs a per-L draft term. # That term is recoverable rather than fatal: with several batch sizes the # same M is reached at different L (M=192 from bs=32/L=5 and bs=64/L=2), and - # those shared points separate a token cost from an L cost. See - # tmp/dspark_build/verdict_theta.py. + # those shared points separate a token cost from an L cost. block = int(verify_len) if config.block_follows_verify_len else int(config.max_draft_len) spec_config = DSparkDecodingConfig( max_draft_len=block, @@ -1509,8 +1484,7 @@ def _run_cell(llm, config: SweepConfig, *, batch_size: int, verify_len: int, # and the cell reports "kept nothing" while the plateau it was measuring # scrolled out of the buffer unseen (jobs 2561986, 2563525). # - # The hazard this replaces is real but narrower than the old comment - # claimed: `IterationResult` latches done when its queue runs dry, so an + # `IterationResult` latches done when its queue runs dry, so an # empty read ends collection. Polling only while requests are still # outstanding keeps the queue fed; the count of empty reads is reported so a # cell that lost its tail this way is visible rather than merely short. @@ -2263,7 +2237,7 @@ def drain_metrics() -> List[dict]: f"({exc}); skipping the rung", file=sys.stderr) continue for batch_size in sorted(batch_sizes): - # H2: the load must outlive the measurement window, or the + # The load must outlive the measurement window, or the # cell measures ramp-down rows filed under smaller batch # sizes. Budget the output by the window at a conservative # 30 ms/step lower bound. diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py index 45b55e394567..4c40e44e51cf 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py @@ -18,7 +18,6 @@ """ import math -import os from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple @@ -27,7 +26,6 @@ from utils.util import skip_pre_blackwell from tensorrt_llm._torch.attention_backend.interface import ( - AttentionForwardArgs, AttentionInputType, MLAParams, PositionalEmbeddingParams, @@ -40,7 +38,6 @@ ) from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import ( DeepseekV4TrtllmAttentionMetadata, - get_token_bytes, ) from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest @@ -676,7 +673,7 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): torch.manual_seed(42) context_lengths = [1] - local_num_heads = int(os.environ.get("DSV4_REPRO_LOCAL_NUM_HEADS", "16")) + local_num_heads = 16 layer_idx = 0 ratio = scenario.compress_ratios[layer_idx] assert ratio == 1 @@ -727,8 +724,7 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): ) layer.update_quant_config(None) attn_sink = torch.randn(local_num_heads, dtype=torch.float32, device=device).mul_(0.5) - if not os.environ.get("DSV4_REPRO_NO_SINK"): - layer.attn_sink = torch.nn.Parameter(attn_sink, requires_grad=False) + layer.attn_sink = torch.nn.Parameter(attn_sink, requires_grad=False) attn_metadata = DeepseekV4TrtllmAttentionMetadata( seq_lens=torch.tensor(context_lengths, dtype=torch.int), @@ -759,75 +755,7 @@ def test_deepseek_v4_sparse_mla_single_token_tp4_local_heads_repro(): fused_q = torch.cat([q_nope, q_pe], dim=-1).view(1, local_num_heads * head_dim) latent_cache = torch.cat([compressed_kv, k_pe], dim=-1) - simple_swa_pool = None - if os.environ.get("DSV4_REPRO_SIMPLE_POOL"): - simple_swa_pool = torch.empty( - (cache_manager.tokens_per_block, head_dim), dtype=dtype, device=device - ) - simple_swa_pool.zero_() - simple_swa_pool_ptr = simple_swa_pool.data_ptr() - cache_manager.kv_cache_pool_pointers[0, 0] = simple_swa_pool_ptr - attn_metadata.host_kv_cache_pool_pointers[0, 0] = simple_swa_pool_ptr - attn_metadata.sparse_mla_base_ptrs[1] = simple_swa_pool_ptr - attn_metadata.swa_buffer_ptrs[layer_idx] = simple_swa_pool_ptr - attn_metadata.block_tables[(1, DeepseekV4AttentionType.SWA)][0].fill_(-1) - attn_metadata.block_tables[(1, DeepseekV4AttentionType.SWA)][0, 0] = 0 - torch.cuda.synchronize() - - if os.environ.get("DSV4_REPRO_PRINT_PARAMS"): - sparse_attn_indices, _ = layer.sparse_attn_predict( - fused_q, - None, - attn_metadata, - AttentionForwardArgs( - attention_input_type=AttentionInputType.context_only, - latent_cache=latent_cache, - q_pe=q_pe, - ), - ) - torch.cuda.synchronize() - print("DSV4_REPRO local_num_heads", local_num_heads) - print( - "DSV4_REPRO token_stride", - get_token_bytes( - head_dim, - sparse_config.index_head_dim, - ratio, - DeepseekV4AttentionType.SWA, - False, - ), - ) - print("DSV4_REPRO tokens_per_block", cache_manager.tokens_per_block) - print("DSV4_REPRO max_blocks_per_seq", cache_manager.max_blocks_per_seq) - print("DSV4_REPRO max_seq_len", cache_manager.max_seq_len) - print("DSV4_REPRO metadata_max_seq_len", attn_metadata.max_seq_len) - print("DSV4_REPRO kv_lens_runtime", attn_metadata.kv_lens_runtime[:1].cpu().tolist()) - print("DSV4_REPRO host_total_kv_lens", attn_metadata.host_total_kv_lens.cpu().tolist()) - print( - "DSV4_REPRO prompt_lens_runtime", - attn_metadata.prompt_lens_cpu_runtime[:1].cpu().tolist(), - ) - print("DSV4_REPRO swa_pool_base_ptr", attn_metadata.sparse_mla_base_ptrs[1]) - print("DSV4_REPRO swa_buffer_ptr", attn_metadata.swa_buffer_ptrs[layer_idx]) - print( - "DSV4_REPRO block_table_swa", - attn_metadata.block_tables[(1, DeepseekV4AttentionType.SWA)][:1, :4].cpu().tolist(), - ) - print("DSV4_REPRO sparse_attn_indices", sparse_attn_indices.cpu().tolist()) - print("DSV4_REPRO sparse_attn_indices_dtype", sparse_attn_indices.dtype) - print( - "DSV4_REPRO sparse_mla_topk_lens", - attn_metadata.sparse_mla_topk_lens[ratio][:1].cpu().tolist(), - ) - print( - "DSV4_REPRO sparse_mla_topk_lens_dtype", attn_metadata.sparse_mla_topk_lens[ratio].dtype - ) - softmax_stats_tensor = None - if os.environ.get("DSV4_REPRO_SOFTMAX_STATS"): - softmax_stats_tensor = torch.empty( - (1, local_num_heads, 2), dtype=torch.float32, device=device - ) result = layer.forward( fused_q.clone(), diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py index 1324ad0c7684..b394d8644034 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py @@ -126,9 +126,6 @@ def test_head_resolves_across_both_layouts_and_absence_is_none(): from tensorrt_llm._torch.speculative.dspark_sts import resolve_confidence_head h = _Head() assert resolve_confidence_head(_Bare(h)) is h - # The wrapper layout is the one the chained getattr missed: resolving to - # None here silently disabled calibration for every confidence-scheduled - # run of the first campaign. assert resolve_confidence_head(_Wrapper(h)) is h assert resolve_confidence_head(_Bare(None)) is None assert resolve_confidence_head(_Wrapper(None)) is None diff --git a/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py b/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py index cfd9ea48b587..9387f1732de6 100644 --- a/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py +++ b/tests/unittest/_torch/thop/parallel/test_compressor_ragged.py @@ -14,11 +14,10 @@ given to verify, so ``next_n`` was demoted to a compile-time *upper bound* and the true per-request counts now arrive in an optional ``new_tokens_per_seq``. -The property that matters is the one the goal doc calls out as missing -(§6.3): a ragged call whose ``new_tokens_per_seq`` is filled with the uniform -value must be **bitwise identical** to passing ``None``. That is what makes the -optional argument safe to add -- it says the ragged code path did not perturb -the uniform one, which every non-DSpark model still uses. +The property that matters: a ragged call whose ``new_tokens_per_seq`` is filled +with the uniform value must be **bitwise identical** to passing ``None``. That +is what makes the optional argument safe to add -- it says the ragged code path +did not perturb the uniform one, which every non-DSpark model still uses. The second test covers the case the first cannot: genuinely differing counts. It checks the ragged call against per-request uniform calls, which is the only @@ -189,7 +188,7 @@ def run(self, next_n: int, ) def test_uniform_new_tokens_matches_none(kv_lens: List[int], next_n: int, dtype: torch.dtype): - """goal doc §6.3: filling ``new_tokens_per_seq`` uniformly == passing None. + """Filling ``new_tokens_per_seq`` uniformly must equal passing None. This is the differential that licenses the optional argument. If it fails, the ragged plumbing changed the uniform path -- which every non-DSpark model From 52fc0e70a85843204216c9e1d92ce7e747889dc2 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:21:26 -0700 Subject: [PATCH 147/177] [None][chore] Cleanup remainder: HOST_STEEP bisector, dead predicates, dev-history comment trims Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/attention_backend/sparse/dsa.py | 13 ++++------- tensorrt_llm/_torch/models/modeling_dspark.py | 2 +- .../_torch/pyexecutor/model_engine.py | 23 ++----------------- .../_torch/speculative/dspark_ragged.py | 7 +----- .../kv_cache_manager_v2/_core/_kv_cache.py | 4 +--- .../defs/accuracy/test_llm_api_pytorch.py | 6 ----- .../dspark_make_steep_sps_table.py | 7 ++---- tests/microbenchmarks/dspark_sps_profiler.py | 6 ++--- .../hw_agnostic/test_dspark_padded_ceiling.py | 13 +++++------ .../hw_agnostic/test_dspark_ragged.py | 7 ++---- .../hw_agnostic/test_dspark_sps_profiler.py | 4 +--- .../test_dspark_sts_collection_guards.py | 7 +----- .../test_dspark_sts_ring_pairing.py | 4 +--- .../test_dspark_ragged_requires_table.py | 3 ++- 14 files changed, 28 insertions(+), 78 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py index 2f26ccbc627b..c5540472285e 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/dsa.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/dsa.py @@ -2180,14 +2180,11 @@ def apply_device_ragged_layout(self, verify_lens: torch.Tensor, # here because block IDs do not move within a step and their refresh # is prepare-time only. self.refresh_token_major_block_table() - block_table_expanded = getattr(self, "block_table_expanded", None) - indexer_offsets = getattr(self, "indexer_k_cache_block_offsets", None) - if block_table_expanded is not None and indexer_offsets is not None: - width = min(indexer_offsets.shape[-1], - block_table_expanded.shape[-1]) - block_table_expanded[:rows, :width] = ( - indexer_offsets[nc:self.num_seqs, :width].index_select( - 0, req_idx).clamp_(min=0)) + width = min(self.indexer_k_cache_block_offsets.shape[-1], + self.block_table_expanded.shape[-1]) + self.block_table_expanded[:rows, :width] = ( + self.indexer_k_cache_block_offsets[nc:self.num_seqs, :width] + .index_select(0, req_idx).clamp_(min=0)) def prepare_for_indexer_k_cache(self): # Build indexer_k_cache_block_offsets using pool block indices derived diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index cdd3097a3dec..20beb8aebfc6 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -432,7 +432,7 @@ def cache_attn_weights_from_checkpoint(self, ckpt_dir: str, weight_map: Dict[str """Populate ``_dspark_attn`` by reading the ``mtp.{s}.attn.*`` tensors from the checkpoint shards on disk, then dequantizing via :meth:`_cache_attn_weights`. - TODO(step 3): source these from the loaded ``MLA`` modules instead, once the + TODO: source these from the loaded ``MLA`` modules instead, once the fused/interleaved fp8 scale layout is decoded, to drop the checkpoint I/O. """ from safetensors import safe_open diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index ff72f570c571..87553c8e6bba 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -814,12 +814,8 @@ def __init__( ) self.cuda_graph_runner = CUDAGraphRunner(cuda_graph_runner_config) - # Pinned staging buffers for async H2D copies on the ragged path; see - # _pinned_host / _pinned_host_record. Two slots per key: the WAR guard - # on a slot only waits for the H2D enqueued two steps ago, so in - # steady state the wait is free. A single slot would serialize step - # N+1's prepare against step N's stream position (measured 6 blocking - # event syncs / 21.6 ms per step on short-window batches). + # Pinned staging buffers for async H2D copies on the ragged path; + # see _pinned_host for the two-slot WAR-guard rationale. self._pinned_host_cache = {} self._pinned_host_events = {} self._pinned_host_active = {} @@ -3844,21 +3840,6 @@ def fit_ragged_verify_lens( runner.ragged_pad_verify_len = pad_len - 1 published = filled.verify_lens.tolist() - if os.environ.get("TLLM_DSPARK_HOST_STEEP") == "1" and n_real > 1: - # Debug bisector: force truly non-uniform windows through the - # HOST staging (the reference implementation, no device - # prologue). The production host scheduler only ever publishes - # uniform windows (budget quantized to the tier ladder), so the - # ragged consumers' non-uniform handling is otherwise exercised - # only by the device path -- this isolates them. - lo = 1 + getattr(self.spec_config, "min_verify_len", 1) - hi = max_verify_len - pairs = (n_real // 2) * 2 - for i in range(0, pairs, 2): - give = min(published[i] - lo, hi - published[i + 1]) - if give > 0: - published[i] -= give - published[i + 1] += give for request, tokens in zip(generation_requests, published): request.py_verify_len = int(tokens) - 1 return int(bucket) diff --git a/tensorrt_llm/_torch/speculative/dspark_ragged.py b/tensorrt_llm/_torch/speculative/dspark_ragged.py index 746e9d8bd60e..8031b5d0e1ef 100644 --- a/tensorrt_llm/_torch/speculative/dspark_ragged.py +++ b/tensorrt_llm/_torch/speculative/dspark_ragged.py @@ -50,7 +50,7 @@ def round_up_to_bucket(total: int, buckets: Sequence[int]) -> int: """Smallest captured bucket that fits ``total``. Raises rather than clamping; callers must reject an oversized batch before - selecting a layout (see :func:`exceeds_captured_buckets`). + selecting a layout. """ if not buckets: raise ValueError("round_up_to_bucket requires a non-empty bucket list") @@ -62,11 +62,6 @@ def round_up_to_bucket(total: int, buckets: Sequence[int]) -> int: return buckets[bisect.bisect_left(buckets, total)] -def exceeds_captured_buckets(total: int, buckets: Sequence[int]) -> bool: - """Whether ``total`` has no captured bucket (check before building a layout).""" - return bool(buckets) and total > buckets[-1] - - def build_qo_indptr(verify_lens: torch.Tensor) -> torch.Tensor: """``[bs]`` lengths -> ``[bs + 1]`` exclusive prefix sum, on the same device. diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index f445db5f122f..9be5fb87e0da 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -870,9 +870,7 @@ def resize(self, capacity: int | None, history_length: int | None = None) -> boo # and is therefore max_blocks_per_seq long, so this # means the sequence wants more blocks than # max_seq_len allows -- i.e. it is consuming KV - # faster than it emits tokens. The bare message gave - # no way to tell that from a genuinely oversized - # request, which cost several 20-minute runs. + # faster than it emits tokens. raise ValueError( f"User-provided base page indices is too short: " f"buffer holds {len(indices)} blocks " diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 57575cad3271..52fd89a314d8 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4070,12 +4070,6 @@ def _run_gsm8k(self, llm): def test_gsm8k_dep8_static_verify_overlap(self, moe_backend): """DSpark verifying the full block, with the overlap scheduler ON. - This was a confidence-scheduling test until the uniform tier ladder was - removed. That path chose one verify length for the whole batch and, on - this checkpoint, always chose the full block -- so it never differed - from what this test does now, and the combination it configured - (scheduling on, ragged off) is rejected outright by llm_args today. - Kept, rather than deleted, because it is the STATIC reference the other two modes are differenced against: same prompts, same seed, verify everything. cap-accept must match it token for token (the schedule only diff --git a/tests/microbenchmarks/dspark_make_steep_sps_table.py b/tests/microbenchmarks/dspark_make_steep_sps_table.py index db8a4ae57075..077a78390306 100644 --- a/tests/microbenchmarks/dspark_make_steep_sps_table.py +++ b/tests/microbenchmarks/dspark_make_steep_sps_table.py @@ -21,11 +21,8 @@ WHY A TABLE RATHER THAN AN OVERRIDE ----------------------------------- -This replaces ``TLLM_DSPARK_FORCE_VERIFY_LENS``, which assigned windows by -batch position. Position is orthogonal to confidence, so that knob reproduced -the ragged *shape* while bypassing the policy entirely -- it could not tell a -correct assignment from an arbitrary one, and it had to sit in production code -ahead of the cost-table gate to be reachable. A table drives the real path. +A table drives the real confidence-scheduled path; it is the only sanctioned +mechanism for forcing window assignments. HOW STEEPNESS IS CHOSEN ----------------------- diff --git a/tests/microbenchmarks/dspark_sps_profiler.py b/tests/microbenchmarks/dspark_sps_profiler.py index c44a9a2b41f1..4500b666b4b3 100644 --- a/tests/microbenchmarks/dspark_sps_profiler.py +++ b/tests/microbenchmarks/dspark_sps_profiler.py @@ -1200,7 +1200,7 @@ def max_tokens_for(self, verify_len: int, can ever yield a sample. A cell is only counted while all ``dp_size`` ranks hold exactly ``batch_size`` generating requests, so the request admitted first must still be alive when the last one is admitted. - Measured on job 2561986 (bs=256, input_len=1024): the scheduler admits + Measured at bs=256, input_len=1024: the scheduler admits roughly one context request per iteration per rank, so the ramp is ~``batch_size`` iterations, while the budget was 85 -- the first requests died ~170 iterations before the last arrived, concurrency @@ -1482,7 +1482,7 @@ def _run_cell(llm, config: SweepConfig, *, batch_size: int, verify_len: int, # bs>=128 the whole retained window therefore lands inside the drain: every # step it contains has 0..11 generating requests, none has `batch_size`, # and the cell reports "kept nothing" while the plateau it was measuring - # scrolled out of the buffer unseen (jobs 2561986, 2563525). + # scrolled out of the buffer unseen. # # `IterationResult` latches done when its queue runs dry, so an # empty read ends collection. Polling only while requests are still @@ -2407,7 +2407,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int: # is enforceable. Under attention DP the router spreads the load and # each rank's occupancy is emergent -- 1024 requests over 8 ranks # oscillates around 128/rank and pads to whatever ladder entry the - # drift lands on (job 2585368: requested 128, measured at 192; only + # drift lands on (observed: requested 128, measured at 192; only # the cap-saturated 256 landed exactly). The pin guarantees L; bs is # whatever the deployment did. Require nothing, keep what is thick, # and let the fit's grid-connectivity check catch a real hole. diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py index 73c8227f0421..5e03d515d855 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py @@ -2,13 +2,12 @@ # SPDX-License-Identifier: Apache-2.0 """``trim_ratio`` must not report a saving the run did not make -- or a loss. -Four runs reported a NEGATIVE ``trim_ratio`` (jobs 2560992 -0.001, 2561519 --0.0438, 2563581 -0.0082), which reads as "the ragged path delivered more tokens -than a no-trim run". It did not. ``delivered_tokens`` books the fitted bucket, -which is ``padded_bs * (tier + 1)`` where ``padded_bs`` is the CUDA-graph-rounded -maximum row count *across attention-DP ranks*; ``ceiling_tokens`` booked -``num_gen_requests * (1 + block)`` over this rank's real, unpadded requests. Two -different row bases in one ratio. +Real runs reported a NEGATIVE ``trim_ratio``, which reads as "the ragged path +delivered more tokens than a no-trim run". It did not. ``delivered_tokens`` +books the fitted bucket, which is ``padded_bs * (tier + 1)`` where ``padded_bs`` +is the CUDA-graph-rounded maximum row count *across attention-DP ranks*; +``ceiling_tokens`` booked ``num_gen_requests * (1 + block)`` over this rank's +real, unpadded requests. Two different row bases in one ratio. A wrong ratio silently misreports the feature's saving in every summary. """ diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py index 6b18e41b943d..3c9f5c7fd1ec 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py @@ -22,7 +22,6 @@ build_qo_indptr, choose_ragged_capture_shape, count_accepted_ragged, - exceeds_captured_buckets, fill_padded_rows_onehot, round_up_to_bucket, row_ids_from_lens, @@ -82,8 +81,6 @@ def test_round_up_raises_past_the_largest_bucket(): """Clamping would silently drop the step out of graph replay into eager.""" with pytest.raises(ValueError, match="exceeds the largest captured bucket"): round_up_to_bucket(65, [8, 16, 32, 64]) - assert exceeds_captured_buckets(65, [8, 16, 32, 64]) - assert not exceeds_captured_buckets(64, [8, 16, 32, 64]) def test_layout_derives_the_bucket_from_the_host_total(): @@ -425,7 +422,7 @@ def test_shape_requires_bs_buckets(): # shared dummy), so decomposability is a divisibility question -- sharpest # under a full-block pin where real rows have zero slack -- and a rank that # failed it declined ALONE while its peers published the bucket. The ADP -# shape gate then took the whole group eager: 70/962 steps on job 2586075. +# shape gate then took the whole group eager. # --------------------------------------------------------------------------- MAXLEN_BUCKETS = [16, 24, 32, 40, 48] # padded_bs 8, tiers 1..5 -> 8*(t+1) @@ -443,7 +440,7 @@ def _shape6(num_real, total, peers): def test_pinned_ranks_agree_on_a_bucket_every_rank_can_decompose(): - """The job-2586075 divergence, miniaturised. + """The production ADP divergence, miniaturised. Both ranks are pinned to the full block (floor == capacity). needed = 8 + 25 = 33 rounds to 40, which rank (5, 30) cannot decompose (3 pad rows, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py index 39633aa600b5..07e12fac11e4 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py @@ -549,9 +549,7 @@ def test_validate_rejects_a_decode_step_wider_than_max_num_tokens(): def test_validate_accepts_a_sane_sweep(): - """Sane includes verifying below a constant drafted block: the Repeat.cu - output_size assert that geometry used to trip was an IMA-family symptom, - fixed by the published-bucket graph key and verified gone 2026-08-06.""" + """Sane includes verifying below a constant drafted block.""" _sweep(block_follows_verify_len=True).validate() _sweep(verify_lens=[5], max_draft_len=5).validate() _sweep().validate() diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py index b394d8644034..1dab83acaaa1 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py @@ -32,12 +32,7 @@ @contextlib.contextmanager def collecting(pin=None): - """Set the collection env for the block, restore it after. - - A context manager rather than a pytest fixture: these run inside the - container, which ships no pytest, under a minimal shim that does not - implement fixtures or tmp_path. - """ + """Set the collection env for the block, restore it after.""" saved = {k: os.environ.get(k) for k in (STS_COLLECT_ENV, _PIN)} try: with tempfile.TemporaryDirectory() as tmp: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py index 23f2467ce68c..d75d95a74495 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py @@ -7,9 +7,7 @@ step ``i-1``, and under the overlap scheduler the sampler consumes the label while pass ``i+1`` has already overwritten the shared confidence buffer. The single mutable stash this ring replaced had NO execution order under which -the pair was correct -- measured on job 2562577, 82.2% of rows paired -confidence(t+1) with accepted(t), correlation 0.14 -- and the shards looked -perfectly healthy. +the pair was correct -- and the shards looked perfectly healthy. These tests pin the join contract: a pair is appended only when the ring snapshot's per-row stamp equals the pass the label verifies, everything else diff --git a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py index 898699b69e91..7e7112e86643 100644 --- a/tests/unittest/llmapi/test_dspark_ragged_requires_table.py +++ b/tests/unittest/llmapi/test_dspark_ragged_requires_table.py @@ -11,7 +11,8 @@ There is no exception: a crafted (non-flat) cost table is how a test makes the planner trim, which exercises the real confidence-driven path rather than -bypassing it. See tests/unittest/_torch/speculative/test_dspark_planner_trims.py. +bypassing it. See +tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py. """ import pytest From 11bf02a7e12d5163c58d18d2b8329f03fe5d59c9 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:48:41 -0700 Subject: [PATCH 148/177] [None][chore] Trim branch-relative narration from planner-trims test docstring Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../_torch/speculative/test_dspark_planner_trims.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tests/unittest/_torch/speculative/test_dspark_planner_trims.py b/tests/unittest/_torch/speculative/test_dspark_planner_trims.py index bcb12eecb37c..1262c8ff368f 100644 --- a/tests/unittest/_torch/speculative/test_dspark_planner_trims.py +++ b/tests/unittest/_torch/speculative/test_dspark_planner_trims.py @@ -15,12 +15,11 @@ hands different windows to requests with different survival. That exercises what actually ships. -Why this matters beyond tidiness: the KV leak fixed in this branch (KV reserved -for the full drafted block but only the verified window rewound) is reachable -only when verified_len < draft_len. On this checkpoint the planner declines to -trim at every batch size, because draft acceptance is ~90% and the fitted table -says verifying the full block wins -- so without a way to make it trim, that -code path ships untested. +Why this matters beyond tidiness: the KV rewind path (KV reserved for the full +drafted block but only the verified window rewound) is reachable only when +verified_len < draft_len, and a well-fitted table can decline to trim at every +batch size -- so without a table that forces trimming, that code path ships +untested. """ import torch From c94388439ef722b15541ed9aaea578a43f4d34f5 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:50:50 -0700 Subject: [PATCH 149/177] [None][doc] GSM8K regression passes on the ragged path with the post-fix table Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index c3db1972bc4a..c9226e35b19e 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -634,6 +634,10 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( 结论:上限=窗口池按 min_slots 下限定容(比例定容被推断的 max_seq_len=1M 饿死)+ 运行时每请求锁 2 块;**显式设 max_seq_len 即解钉**(用户已拍板 C 方向)。部署规约:DSv4 serving 必须显式设 `max_seq_len`(或 `avg_seq_len`),否则窗口池被钉死在 ~maxbs/2 并发,no-spec 高并发直接死锁。代码侧 C 修复(typical_seq_len 不默认 max_seq_len)+ 上游三报(下限建模、估算器漏 draft、死锁而非降级)待落。 +### GSM8K 精度回归:新表真实裁剪下通过(08-10 21:40) + +`test_gsm8k_dep8_ragged_verify[moe_backend=TRTLLM]`(DEP8 + overlap + CUDA graph + ragged verify + 置信度调度),`TLLM_DSPARK_SPS_TABLE=postfix_pro_table.json` 注入修复后真实表(planner 主动裁剪而非退化均匀),**1 passed(27 分钟)——精度 ≥ 记录参考值,断言通过**。终局判决的精度一环闭合。注:MEGAMOE_DEEPGEMM 变体在该节点因 DeepGEMM JIT "NVCC compilation failed" 无法初始化(nvcc 在位、/tmp 空,原因未深究;TRTLLM 变体即吞吐判决所用后端,更具代表性)。 + ### 产品线终局判决:Pro@DEP8 置信度调度全面转正(08-10 13:50) Flash 三臂翻案后,Pro@DEP8 终局(maxbs=256 解锁满批 + `postfix_pro_table.json` 新表,scheduled vs notrim,双节点臂序对调,每 cell n=10): From 27191b69551adbd65ab6af9e13a487ac83e8cc50 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Mon, 10 Aug 2026 22:37:37 -0700 Subject: [PATCH 150/177] [None][chore] Consolidate DSpark unit tests and simplify defensive access Merge 23 small test files into 7 subsystem modules (planner/budget, ragged/buckets, model, profiler, STS, accept-stats, sampler), dropping two verbatim duplicate tests and trimming docstrings to the invariant statements; 221 tests preserved (was 223). Replace 14 provably-safe defensive getattr calls with direct attribute access (init sites cited in review), keeping the legitimate lazy-attribute probes. Remove the TLLM_DSPARK_DISABLE_OVERLAP diagnostic from the GSM8K test. Regression: 378 passed, failures unchanged (environmental llama file). Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/_torch/models/dspark/draft.py | 2 +- tensorrt_llm/_torch/models/modeling_dspark.py | 14 +- .../_torch/pyexecutor/model_engine.py | 4 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 7 +- tensorrt_llm/_torch/speculative/interface.py | 2 +- .../defs/accuracy/test_llm_api_pytorch.py | 7 +- ..._accept.py => test_dspark_accept_stats.py} | 214 ++-- .../hw_agnostic/test_dspark_attention.py | 419 ------- .../test_dspark_budget_parity_sglang.py | 232 ---- .../test_dspark_compressor_gen_counts.py | 78 -- .../test_dspark_confidence_schedule.py | 509 -------- .../test_dspark_cost_interpolation.py | 61 - .../hw_agnostic/test_dspark_cuda_graph.py | 445 ------- .../test_dspark_device_fill_parity.py | 355 ------ .../hw_agnostic/test_dspark_draft.py | 118 -- .../hw_agnostic/test_dspark_heads.py | 158 --- .../test_dspark_iter_log_samples.py | 159 --- .../hw_agnostic/test_dspark_model.py | 1036 +++++++++++++++++ .../hw_agnostic/test_dspark_overlap_layout.py | 112 -- .../hw_agnostic/test_dspark_padded_ceiling.py | 70 -- .../hw_agnostic/test_dspark_planner.py | 903 ++++++++++++++ .../hw_agnostic/test_dspark_ragged.py | 658 ++++++++++- .../test_dspark_reconcile_guards.py | 99 -- .../test_dspark_snapshot_instruments.py | 110 -- .../hw_agnostic/test_dspark_sps_profiler.py | 222 ++-- ...sts_ring_pairing.py => test_dspark_sts.py} | 131 ++- .../test_dspark_sts_collection_guards.py | 127 -- .../test_dspark_tier_aligned_budget.py | 96 -- .../hw_agnostic/test_dspark_verify_len_pin.py | 97 -- .../speculative/test_dspark_graph_key.py | 146 --- .../speculative/test_dspark_planner_trims.py | 74 -- .../speculative/test_force_accepted_tokens.py | 363 ------ .../speculative/test_group_all_greedy_sync.py | 83 -- .../test_rejection_buffers_guard.py | 270 ----- .../_torch/speculative/test_spec_sampler.py | 650 +++++++++++ 35 files changed, 3553 insertions(+), 4478 deletions(-) rename tests/unittest/_torch/speculative/hw_agnostic/{test_dspark_cap_accept.py => test_dspark_accept_stats.py} (64%) delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_compressor_gen_counts.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_planner.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py rename tests/unittest/_torch/speculative/hw_agnostic/{test_dspark_sts_ring_pairing.py => test_dspark_sts.py} (58%) delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py delete mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py delete mode 100644 tests/unittest/_torch/speculative/test_dspark_graph_key.py delete mode 100644 tests/unittest/_torch/speculative/test_dspark_planner_trims.py delete mode 100644 tests/unittest/_torch/speculative/test_force_accepted_tokens.py delete mode 100644 tests/unittest/_torch/speculative/test_group_all_greedy_sync.py delete mode 100644 tests/unittest/_torch/speculative/test_rejection_buffers_guard.py create mode 100644 tests/unittest/_torch/speculative/test_spec_sampler.py diff --git a/tensorrt_llm/_torch/models/dspark/draft.py b/tensorrt_llm/_torch/models/dspark/draft.py index a70ccc503155..ec02b64aeb83 100644 --- a/tensorrt_llm/_torch/models/dspark/draft.py +++ b/tensorrt_llm/_torch/models/dspark/draft.py @@ -111,7 +111,7 @@ def dspark_propose( prev_ids = torch.cat([bonus_token_ids.unsqueeze(1), draft_tokens[:, :-1]], dim=1) prev_emb = ( markov_head.get_prev_embeddings(prev_ids) - if (markov_head is not None and getattr(confidence_head, "with_markov", False)) + if (markov_head is not None and confidence_head.with_markov) else None ) confidence = ( diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 20beb8aebfc6..b46ef07aeffa 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -174,7 +174,7 @@ def __init__( disable_post_moe_fusion=True, ) config = model_config.pretrained_config - spec_cfg = getattr(model_config, "spec_config", None) + spec_cfg = model_config.spec_config self.stage_id = int(stage_id) self.num_stages = int(num_stages) # mask_token_id is a user override on the speculative_config; None means @@ -266,7 +266,7 @@ def __init__( # ``mtp.{0..n-1}.*`` weight namespace. Resolve it from (in priority): # an explicit override, the spec config's ``num_draft_layers``, a # pretrained-config ``n_mtp_layers``, else fall back to nextn. - spec_cfg = getattr(model_config, "spec_config", None) + spec_cfg = model_config.spec_config self.num_stages = int( num_stages if num_stages is not None @@ -521,8 +521,8 @@ def _draft_normalized_quant_config(model_config): when nothing changes (e.g. the MXFP4 checkpoint, whose global algo is not ``MIXED_PRECISION``) so the draft config is left byte-identical. """ - qc = getattr(model_config, "quant_config", None) - if qc is None or getattr(qc, "quant_algo", None) != QuantAlgo.MIXED_PRECISION: + qc = model_config.quant_config + if qc is None or qc.quant_algo != QuantAlgo.MIXED_PRECISION: return None # Reuse the target normalizer on a throwaway shallow copy so we only # extract the resolved global quant_config; the shared ``model_config`` @@ -530,7 +530,7 @@ def _draft_normalized_quant_config(model_config): probe = copy.copy(model_config) object.__setattr__(probe, "_frozen", False) normalized = _normalize_deepseek_v4_nvfp4_mixed_precision_config(probe) - resolved = getattr(normalized, "quant_config", qc) + resolved = normalized.quant_config return resolved if resolved is not qc else None @staticmethod @@ -541,7 +541,7 @@ def _draft_sparse_config(model_config, base: int, num_stages: int): sparse config must expose the draft layers' per-layer ratios at indices ``[0, num_stages)``. """ - sa = getattr(model_config, "sparse_attention_config", None) + sa = model_config.sparse_attention_config compress_ratios = getattr(sa, "compress_ratios", None) if sa is not None else None if not compress_ratios or len(compress_ratios) < base + num_stages: return None @@ -564,7 +564,7 @@ def _draft_quant_config_dict(model_config, base: int, num_stages: int): physically MXFP4 (identical to the main MoE layers), so copy a representative main MoE layer's experts quant onto the draft layer keys. """ - qcd = getattr(model_config, "quant_config_dict", None) + qcd = model_config.quant_config_dict if not qcd: return None src = next( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 87553c8e6bba..fd3ca223d424 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3865,7 +3865,7 @@ def _apply_device_window_prologue(self, inputs, from ..speculative.dspark_device_select import select_windows_device runner = self.cuda_graph_runner - budget = getattr(self, "_dspark_device_budget", None) + budget = self._dspark_device_budget self._dspark_device_budget = None if budget is None or runner.agreed_ragged_bucket is None: return False @@ -3875,7 +3875,7 @@ def _apply_device_window_prologue(self, inputs, # previous device tensor breaks that addressing. return False worker = self._get_spec_worker() - planner = getattr(worker, "verify_planner", None) + planner = worker.verify_planner if planner is None or worker.staged_confidence_buffer() is None: return False if worker.batch_slot_view(1) is None: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 28f6e1724bb3..2d345cbedd4b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3246,7 +3246,7 @@ def set_dspark_verify_len_pin(self, verify_len: Optional[int]) -> Optional[int]: "no DSpark confidence planner on this executor; the verify-len " "pin only applies to a DSpark engine with confidence " "scheduling enabled") - mode = getattr(worker, "ragged_verify_mode", None) + mode = worker.ragged_verify_mode computes = (mode.computes_windows if mode is not None else self.model_engine.spec_config.enable_ragged_verify) if not computes: @@ -3298,7 +3298,7 @@ def set_dspark_budget_frac(self, "no DSpark confidence planner on this executor; the budget " "fraction only applies to a DSpark engine with confidence " "scheduling enabled") - mode = getattr(worker, "ragged_verify_mode", None) + mode = worker.ragged_verify_mode trims = (mode.trims_submitted_tokens if mode is not None else self.model_engine.spec_config.enable_ragged_verify) if not trims: @@ -3562,8 +3562,7 @@ def _dspark_confidence_draft_len(self, # Only when the fit delivered: otherwise this # books the LAST fit's row count on steps that # never entered the fit. - padded_bs=(getattr(self.model_engine, - "_dspark_last_padded_bs", None) + padded_bs=(self.model_engine._dspark_last_padded_bs if bucket is not None else None), fallback=fallback_reason, delivered=delivered) diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 342dd89bdd3c..0fc2168ad54f 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1045,7 +1045,7 @@ def apply_accept_caps(num_accepted_tokens: torch.Tensor, num_contexts: int, permanently ahead of what was committed. In-place and sync-free, so safe to record into a CUDA graph. """ - caps = getattr(spec_metadata, "accept_caps", None) + caps = spec_metadata.accept_caps if caps is None: return gen = num_accepted_tokens[num_contexts:] diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 52fd89a314d8..966c6f3e1082 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4174,12 +4174,7 @@ def test_gsm8k_dep8_ragged_verify(self, moe_backend): kv_cache_config=kv_cache_config, cuda_graph_config=CudaGraphConfig(enable_padding=True), enable_chunked_prefill=False, - # Diagnostic override only. The overlap scheduler is the - # configuration this test exists to cover, so it stays on by - # default; turning it off isolates whether a failure needs the - # previous step's kernels to still be in flight. - disable_overlap_scheduler=bool( - os.environ.get("TLLM_DSPARK_DISABLE_OVERLAP")), + disable_overlap_scheduler=False, custom_tokenizer="deepseek_v4", speculative_config=spec_config) as llm: score, acc_params = self._run_gsm8k(llm) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_accept_stats.py similarity index 64% rename from tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py rename to tests/unittest/_torch/speculative/hw_agnostic/test_dspark_accept_stats.py index 67a93ab591b2..080294abb8a4 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cap_accept.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_accept_stats.py @@ -1,24 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""``cap-accept``: run the schedule's policy without the layout that pays for it. +"""DSpark accept-caps and ragged-verify observability (hardware-agnostic, CPU). -The mode exists to split one question into two. ``compact`` changes both what -the scheduler commits *and* how the batch is laid out, so when its output moves, -nothing says which half did it. ``cap-accept`` applies the same per-request -windows to commitment while submitting the full block, giving two independent -assertions instead of one entangled one: - - cap-accept output != static output -> the scheduling policy is wrong - compact output != cap-accept -> the layout compaction is wrong - -It also produces a number ``compact`` structurally cannot: ``cap_trim_tokens``, -the positions the target accepted and the window then discarded. Under -``compact`` those positions are never scored, so the acceptance cost of trimming -can only be bounded (``trim_regret_rate``), never measured. - -The price is that it saves no compute at all, which is why it is a diagnostic -mode and never a serving one -- and why the counters must not report it as -though it had saved any. +Covers `apply_accept_caps` clamping semantics, the cap-accept diagnostic mode's +counters (`cap_trim_tokens`, trim shape), `trim_ratio` row-base accounting under +CUDA-graph padding, and trim-regret acceptance stats. """ import os @@ -34,10 +20,15 @@ MAX_DRAFT_LEN = 5 -def _stats(mode=RaggedVerifyMode.CAP_ACCEPT): +def _cap_stats(mode=RaggedVerifyMode.CAP_ACCEPT): return DSparkRaggedStats(mode=mode, max_draft_len=MAX_DRAFT_LEN) +def _compact_stats(max_draft_len=5): + return DSparkRaggedStats(mode=RaggedVerifyMode.COMPACT, + max_draft_len=max_draft_len) + + class _Meta: """Just the two fields `apply_accept_caps` reads.""" @@ -49,16 +40,7 @@ def __init__(self, caps=None, track_trim=True, rows=8): if track_trim else None) -# --- the cap itself --------------------------------------------------------- -# -# WHERE this runs is the whole point, and it is not testable from here: it has -# to be inside acceptance, because the DSpark drafter reads the same -# `num_accepted_tokens` later in the SAME forward to advance its rolling KV -# window and its persistent decode position (`_ctx_len += nacc`, dspark.py). -# Capping on the host afterwards leaves the drafter advanced by the uncapped -# count while only the capped prefix was committed -- lossless output, but the -# draft state drifts and the mode stops being a faithful reference. -# +# --- the cap itself ---------------------------------------------------------- # Units: `num_accepted_tokens` is `accepted drafts + 1`, so a window of w # drafted positions is a cap of `w + 1`, matching `verify_lens`. @@ -71,13 +53,8 @@ def test_no_caps_is_a_no_op(): def test_counts_beyond_their_windows_are_clamped_and_the_loss_accumulated(): - """Clamp plus per-step (not cumulative) loss recording. - - The host accumulates across steps, so the buffer must hold what THIS step - lost: a buffer that added would double-count every request still in the - batch next step. Repeating the same step must therefore leave the same - losses, not multiples of them. - """ + """The trim buffer holds this step's loss, not a running sum: repeating the + same step leaves the same losses.""" meta = _Meta(caps=[3, 5, 6]) for _ in range(3): counts = torch.tensor([6, 5, 2], dtype=torch.int32) @@ -125,11 +102,11 @@ def test_nothing_is_lost_or_invented(accepted, window): assert committed <= window + 1 -# --- what the counters must say --------------------------------------------- +# --- what the counters must say ---------------------------------------------- def test_cap_trim_is_recorded_per_request(): - stats = _stats() + stats = _cap_stats() # Given 5, allowed 2 -> 3 accepted positions thrown away. stats.record_acceptance(accepted=2, window=2, cap_trim=3) # Died on its own inside a wide window: the schedule cost it nothing. @@ -154,17 +131,13 @@ def test_cap_trim_is_recorded_per_request(): def test_the_same_total_loss_is_distinguished_by_its_shape(): - """The reason this is per request at all. - - Both runs lose 20 tokens over 10 requests. The first spreads it; the - second takes it all out of one request. A total cannot tell them apart, - and they mean opposite things about the planner. - """ - spread = _stats() + """Equal totals with different per-request shapes must stay + distinguishable: spread loss vs. loss concentrated in one request.""" + spread = _cap_stats() for _ in range(10): spread.record_acceptance(accepted=1, window=1, cap_trim=2) - concentrated = _stats() + concentrated = _cap_stats() concentrated.record_acceptance(accepted=1, window=1, cap_trim=20) for _ in range(9): concentrated.record_acceptance(accepted=3, window=5, cap_trim=0) @@ -184,7 +157,7 @@ def test_the_same_total_loss_is_distinguished_by_its_shape(): def test_the_loss_is_absent_rather_than_zero_when_not_measured(): """`compact` cannot produce this number, and must not appear to.""" - stats = _stats(mode=RaggedVerifyMode.COMPACT) + stats = _cap_stats(mode=RaggedVerifyMode.COMPACT) stats.record_acceptance(accepted=3, window=3) assert stats.cap_trim_tokens == 0 assert stats.cap_trim_concentration == 0.0 @@ -194,19 +167,9 @@ def test_the_loss_is_absent_rather_than_zero_when_not_measured(): def test_cap_accept_must_not_be_credited_with_a_compute_saving(): - """The regression this guards: cap-accept booking the windows as delivered. - - It submits the full block by construction. Deriving delivered tokens from - the windows would report a trim ratio for a run that trimmed no compute at - all -- and that ratio is the headline number for whether the feature works. - - For the same reason the active gate must not demand a saving here: - `require_trim` asks "did delivered tokens come in under the no-trim - ceiling", cap-accept answers no by construction, and keying the gate on - the ratio alone would fail every cap-accept run -- the reference run - against which the others are judged. - """ - stats = _stats() + """cap-accept submits the full block, so delivered tokens must not be + derived from the windows and `trim_ratio` must stay 0.""" + stats = _cap_stats() num_reqs = 4 full_block = num_reqs * (1 + MAX_DRAFT_LEN) stats.record_step(num_gen_requests=num_reqs, @@ -224,7 +187,7 @@ def test_cap_accept_must_not_be_credited_with_a_compute_saving(): def test_compact_still_derives_delivered_from_the_bucket(): """The default path is unchanged by the new argument.""" - stats = _stats(mode=RaggedVerifyMode.COMPACT) + stats = _cap_stats(mode=RaggedVerifyMode.COMPACT) stats.record_step(num_gen_requests=4, verify_lens=[5, 3, 2, 1], bucket=16) @@ -232,7 +195,7 @@ def test_compact_still_derives_delivered_from_the_bucket(): assert stats.trim_ratio > 0.0 -# --- mode plumbing ---------------------------------------------------------- +# --- mode plumbing ----------------------------------------------------------- def test_cap_accept_computes_windows_but_does_not_trim_the_token_axis(): @@ -248,13 +211,8 @@ def test_cap_accept_computes_windows_but_does_not_trim_the_token_axis(): def test_a_slot_left_unwritten_cannot_inherit_the_previous_occupant_s_loss(): - """The stale-buffer hazard, guarded at the one place it can arise. - - cap_trim_lens is persistent and slot-indexed. A step that produces no - caps -- static, or any other speculation mode -- must still write ZEROS - into this batch's slots, or a slot recycled to a new request would report - whatever its previous occupant lost. - """ + """cap_trim_lens is persistent and slot-indexed: a step producing no caps + must still zero this batch's slots or recycled slots report stale loss.""" import inspect from tensorrt_llm._torch.speculative.spec_sampler_base import SpecSamplerBase @@ -266,16 +224,8 @@ def test_a_slot_left_unwritten_cannot_inherit_the_previous_occupant_s_loss(): def test_the_cuda_graph_padding_dummy_carries_a_cap(): - """The regression that made the whole mode a silent no-op. - - `pad_batch` appends a shared dummy request to reach a captured batch size, - and `_attach_accept_caps` drops the batch's caps if ANY request lacks one. - The dummy has no cap of its own, so before this it disabled cap-accept -- - on padded steps only, so the mode measured nothing on some steps and - worked on others, with nothing in the output to say which. - - Mirrors the treatment `py_verify_len` already got for the ragged path. - """ + """The CUDA-graph padding dummy must carry a verify cap, or cap-accept + silently degrades to static on every padded step.""" import inspect from tensorrt_llm._torch.pyexecutor import cuda_graph_runner @@ -292,20 +242,8 @@ def test_the_cuda_graph_padding_dummy_carries_a_cap(): def test_the_capture_set_follows_the_mode_not_the_config_flag(): - """The bug the first end-to-end run found, in CUDA-graph capture. - - cap-accept is configured with enable_ragged_verify=True -- the config is - byte-identical to compact, only the environment differs. Anything that - SHAPES the batch must therefore ask whether the token axis actually - shrinks, not whether ragged was requested. The capture set asked the flag, - so it built synthetic batches at ragged token totals (bs*(t+1)) for a mode - whose real steps always submit the full block, and the capture pass died - inside MoE on `unmatched tensor shape` -- before any request had run, with - nothing pointing at the scheduler. - - Unit tests could not have caught it: it needs a real capture pass. This - pins the predicate instead. - """ + """Batch shaping must ask whether the token axis actually shrinks, not + whether ragged was requested: cap-accept always submits the full block.""" from types import SimpleNamespace from tensorrt_llm._torch.speculative.dspark_observability import ( @@ -333,15 +271,8 @@ def test_the_capture_set_follows_the_mode_not_the_config_flag(): def test_accept_caps_is_not_the_ragged_flag(): - """The distinction the whole fix rests on. - - `is_ragged_verify` means two things at once -- cap acceptance AND the draft - buffers are packed by window (`_padded_gen_draft_tokens` slices - `draft_tokens[:total_verify_tokens - num_gens]`). cap-accept wants only the - first: its buffers are the ordinary rectangle. If `accept_caps` ever starts - implying ragged, a rectangle gets unpacked as if it were packed and one - request's drafts are attributed to another, silently. - """ + """`accept_caps` must not imply ragged buffer packing: cap-accept's draft + buffers are the ordinary rectangle.""" from tensorrt_llm._torch.speculative.interface import SpecMetadata fields = SpecMetadata.__dataclass_fields__ @@ -351,3 +282,82 @@ def test_accept_caps_is_not_the_ragged_flag(): meta = _Meta(caps=[2, 3]) # The mock stands in for a metadata object carrying caps but no windows. assert getattr(meta, "verify_lens", None) is None + + +# --- trim_ratio row-base under CUDA-graph padding ---------------------------- + + +def test_padding_rows_do_not_make_the_trim_look_negative(): + """A top-tier step on padded rows saved nothing -- and lost nothing.""" + stats = _compact_stats(max_draft_len=5) + # 3 real requests, graph padded to 4 rows, every window at the full block. + # delivered = 4 * 6 = 24. Against a ceiling of 3 * 6 = 18 that is -0.333. + stats.record_step(num_gen_requests=3, verify_lens=[5, 5, 5], + bucket=4 * 6, padded_bs=4) + summary = stats.summary() + assert summary["trim_ratio"] == pytest.approx(0.0), ( + f"a step that verified the full block on every row trimmed nothing, so " + f"the ratio is 0 -- got {summary['trim_ratio']}, which came from " + f"comparing {summary['delivered_tokens']} padded-row tokens against a " + f"ceiling counted over {3} real requests") + # The padded row count itself must land in the summary: the executor once + # passed padded_bs=None unconditionally, and an always-empty padded_bs_hist + # is indistinguishable from "padding never happened". + assert summary["padded_bs_hist"] == {4: 1} + + +def test_a_real_trim_still_reads_as_a_saving(): + """The fix must not flatten genuine trimming to zero.""" + stats = _compact_stats(max_draft_len=5) + # 4 rows, no padding, every window trimmed to 2 -> bucket 4*3 = 12 against a + # ceiling of 4*6 = 24. + stats.record_step(num_gen_requests=4, verify_lens=[2, 2, 2, 2], + bucket=4 * 3, padded_bs=4) + assert stats.summary()["trim_ratio"] == pytest.approx(0.5) + + +def test_paths_without_a_bucket_keep_the_local_row_base(): + """Only the bucket path rebases to the padded row count; widening the + no-window and cap-accept paths would invent a saving out of padding.""" + no_window = _compact_stats(max_draft_len=5) + no_window.record_step(num_gen_requests=3, verify_lens=None, padded_bs=8) + assert no_window.summary()["trim_ratio"] == pytest.approx(0.0) + + explicit = _compact_stats(max_draft_len=5) + # cap-accept: windows are computed but the full block is submitted anyway. + explicit.record_step(num_gen_requests=3, verify_lens=[2, 3, 1], + bucket=8 * 6, padded_bs=8, delivered=3 * 6) + assert explicit.summary()["trim_ratio"] == pytest.approx(0.0) + + +# --- trim regret ------------------------------------------------------------- + + +def test_trim_regret_counts_only_drafts_alive_at_the_cut(): + """Regret separates "trimming was free, those drafts would have died + anyway" from "trimming discarded acceptance".""" + stats = _compact_stats() + # Given the full block and died early: not trimmed, no regret. + stats.record_acceptance(accepted=2, window=5) + # Trimmed to 3 and died at 1: the cut cost nothing. + stats.record_acceptance(accepted=1, window=3) + # Trimmed to 3 and accepted all 3: still alive at the cut -> regret. + stats.record_acceptance(accepted=3, window=3) + + assert stats.requests_scored == 3 + assert stats.requests_trimmed == 2 + assert stats.trimmed_hit_ceiling == 1 + assert stats.trim_regret_rate == 0.5 + assert stats.accept_len == pytest.approx(2.0) + summary = stats.summary() + for key in ("accept_len", "requests_scored", "requests_trimmed", + "trim_regret_rate"): + assert key in summary, f"{key} missing from the summary" + + # Nothing trimmed: the regret rate is 0/0-safe and stays 0. + untrimmed = _compact_stats() + for accepted in (0, 3, 5): + untrimmed.record_acceptance(accepted=accepted, window=5) + assert untrimmed.requests_trimmed == 0 + assert untrimmed.trim_regret_rate == 0.0 + assert untrimmed.summary()["accept_len"] == pytest.approx(8 / 3, abs=1e-4) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py deleted file mode 100644 index b8cce346cd75..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_attention.py +++ /dev/null @@ -1,419 +0,0 @@ -# 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. -"""Unit tests for the DSpark captured-context attention primitives (CPU). - -The two primitives are validated against fully independent computations: -``dspark_sparse_attn`` vs (a) a per-element loop reference of the kernel formula -and (b) ``torch.nn.functional.scaled_dot_product_attention`` for the no-sink, -all-valid case; ``get_dspark_topk_idxs`` vs the reference index formula. -""" - -import types -from unittest.mock import Mock - -import pytest -import torch -import torch.nn.functional as F - -import tensorrt_llm._torch.models.modeling_dspark as modeling_dspark -from tensorrt_llm._torch.models.dspark.attention import ( - apply_dspark_rotary, - dspark_attention_forward, - dspark_sparse_attn, - get_dspark_topk_idxs, - precompute_dspark_freqs_cis, -) -from tensorrt_llm._torch.models.modeling_dspark import DSparkDraftModel - - -def test_rope_table_is_cached_once_per_device(): - model = types.SimpleNamespace( - _attn_params={"rope_head_dim": 16}, - _freqs_cap=64, - _rope_theta=10000.0, - _freqs_table_cache={}, - ) - - first = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) - second = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) - - assert first.data_ptr() == second.data_ptr() - assert len(model._freqs_table_cache) == 1 - positions = torch.tensor([1, 17, 63]) - expected = precompute_dspark_freqs_cis(16, 64, rope_theta=10000.0) - torch.testing.assert_close(first[positions], expected[positions]) - - -def test_dspark_block_uses_stage_id_as_attention_layer_idx(monkeypatch): - captured = {} - - def fake_decoder_layer_init( - self, - model_config, - layer_idx, - aux_stream_dict, - attention_layer_idx=None, - mapping_with_cp=None, - disable_post_moe_fusion=False, - ): - torch.nn.Module.__init__(self) - self.model_config = model_config - self.config = model_config.pretrained_config - self.layer_idx = layer_idx - captured.update( - layer_idx=layer_idx, - attention_layer_idx=attention_layer_idx, - aux_stream_dict=aux_stream_dict, - mapping_with_cp=mapping_with_cp, - disable_post_moe_fusion=disable_post_moe_fusion, - ) - - monkeypatch.setattr( - modeling_dspark.DeepseekV4DecoderLayer, - "__init__", - fake_decoder_layer_init, - ) - model_config = types.SimpleNamespace( - pretrained_config=types.SimpleNamespace(vocab_size=128, hc_mult=2), - spec_config=None, - ) - - block = modeling_dspark.DSparkBlock( - model_config, - layer_idx=10, - aux_stream_dict={}, - stage_id=1, - num_stages=3, - num_capture_layers=0, - ) - - assert block.layer_idx == captured["layer_idx"] == 10 - assert captured["attention_layer_idx"] == block.stage_id == 1 - assert captured["disable_post_moe_fusion"] is True - - -@pytest.mark.parametrize("enable_fused_hc", [True, False]) -def test_forward_stage_honors_enable_fused_hc(monkeypatch, enable_fused_hc): - """The draft stage must use the inherited fused-HC rollback setting.""" - torch.manual_seed(71) - num_requests, block_size, hc_mult, hidden_size = 1, 2, 2, 3 - h = torch.randn(num_requests, block_size, hc_mult, hidden_size) - attention_input = torch.randn(num_requests, block_size, hidden_size) - attention_output = torch.randn_like(attention_input) - mid_residual = torch.randn_like(h) - attention_post_mix = torch.randn(num_requests, block_size, hc_mult, 1) - attention_comb_mix = torch.randn(num_requests, block_size, hc_mult, hc_mult) - ffn_post_mix = torch.randn_like(attention_post_mix) - ffn_comb_mix = torch.randn_like(attention_comb_mix) - raw_ffn_input = torch.randn_like(attention_input) - normed_ffn_input = torch.randn_like(attention_input) - moe_output = torch.randn(num_requests * block_size, hidden_size) - final_h = torch.randn_like(h) - events = [] - - def record(name, result): - def call(*args, **kwargs): - events.append(name) - return result - - return call - - monkeypatch.setattr( - modeling_dspark, - "dspark_attention_forward", - Mock(return_value=attention_output), - ) - - hc_attn = types.SimpleNamespace( - pre_mapping=Mock(return_value=(attention_post_mix, attention_comb_mix, attention_input)), - post_mapping=Mock(side_effect=record("attention_post", mid_residual)), - ) - hc_ffn = types.SimpleNamespace( - fused_hc=Mock( - side_effect=record( - "fused", - (mid_residual, ffn_post_mix, ffn_comb_mix, normed_ffn_input), - ) - ), - pre_mapping=Mock( - side_effect=record("ffn_pre", (ffn_post_mix, ffn_comb_mix, raw_ffn_input)) - ), - post_mapping=Mock(side_effect=record("ffn_post", final_h)), - ) - post_attention_layernorm = Mock(side_effect=record("ffn_norm", normed_ffn_input)) - post_attention_layernorm.weight = torch.ones(hidden_size) - post_attention_layernorm.variance_epsilon = 1e-6 - stage = types.SimpleNamespace( - enable_fused_hc=enable_fused_hc, - hc_attn=hc_attn, - hc_ffn=hc_ffn, - input_layernorm=Mock(side_effect=lambda tensor: tensor), - post_attention_layernorm=post_attention_layernorm, - mlp=Mock(return_value=moe_output), - _dspark_attn={}, - ) - model = types.SimpleNamespace( - use_real_mla=False, - _attn_params={"window_size": 2, "head_dim": 1}, - model_config=types.SimpleNamespace( - mapping=types.SimpleNamespace(enable_attention_dp=False, tp_size=8) - ), - ) - - actual = DSparkDraftModel._forward_stage( - model, - stage, - h, - torch.randn(num_requests, hidden_size), - 1, - torch.empty(0), - torch.zeros(num_requests, block_size, dtype=torch.long), - ) - - assert actual is final_h - torch.testing.assert_close( - stage.mlp.call_args.args[0], - normed_ffn_input.reshape(num_requests * block_size, hidden_size), - ) - # Non-attention-DP multi-GPU (enable_attention_dp=False, tp_size>1): the draft - # MoE must all-reduce its TP-sharded output, mirroring the target MoE. The - # attention-DP and single-GPU paths keep it disabled. - assert stage.mlp.call_args.kwargs["final_all_reduce_params"].enable_allreduce is True - if enable_fused_hc: - assert events == ["fused", "ffn_post"] - hc_ffn.fused_hc.assert_called_once() - hc_attn.post_mapping.assert_not_called() - hc_ffn.pre_mapping.assert_not_called() - post_attention_layernorm.assert_not_called() - fused_kwargs = hc_ffn.fused_hc.call_args.kwargs - assert fused_kwargs["norm_weight"] is post_attention_layernorm.weight - assert fused_kwargs["norm_eps"] == post_attention_layernorm.variance_epsilon - else: - assert events == ["attention_post", "ffn_pre", "ffn_norm", "ffn_post"] - hc_ffn.fused_hc.assert_not_called() - hc_attn.post_mapping.assert_called_once() - hc_ffn.pre_mapping.assert_called_once_with(mid_residual) - post_attention_layernorm.assert_called_once_with(raw_ffn_input) - - -def _ref_precompute_freqs_cis(dim, seqlen, base): - """DeepSpec precompute_freqs_cis with original_seq_len == 0 (no YaRN).""" - freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) - t = torch.arange(seqlen) - freqs = torch.outer(t, freqs) - return torch.polar(torch.ones_like(freqs), freqs) - - -def _ref_apply_rotary_emb(x, freqs_cis, inverse=False): - """DeepSpec apply_rotary_emb (returns a fresh tensor instead of in-place).""" - xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) - if inverse: - freqs_cis = freqs_cis.conj() - if xc.ndim == 3: - fc = freqs_cis.view(1, xc.size(1), xc.size(-1)) - else: - fc = freqs_cis.view(1, xc.size(1), 1, xc.size(-1)) - return torch.view_as_real(xc * fc).flatten(-2).to(x.dtype) - - -def _loop_reference(q, kv, attn_sink, topk_idxs, scale): - """Obvious, slow per-(b,m,h) reference of the exact kernel math.""" - b, m, h, d = q.shape - out = torch.zeros(b, m, h, d, dtype=torch.float32) - qf, kvf, sink = q.float(), kv.float(), attn_sink.float() - for bi in range(b): - for mi in range(m): - idxs = topk_idxs[bi, mi].tolist() - for hi in range(h): - scores, vecs = [], [] - for j in idxs: - if j < 0: - continue - k = kvf[bi, j] - scores.append(torch.dot(qf[bi, mi, hi], k) * scale) - vecs.append(k) - if not scores: - continue - s = torch.stack(scores) - smax = s.max() - p = torch.exp(s - smax) - denom = p.sum() + torch.exp(sink[hi] - smax) - num = (p.unsqueeze(-1) * torch.stack(vecs)).sum(0) - out[bi, mi, hi] = num / denom - return out - - -def test_sparse_attn_matches_loop_reference(): - torch.manual_seed(0) - b, m, h, d, n, topk = 2, 5, 3, 16, 40, 12 - q = torch.randn(b, m, h, d) - kv = torch.randn(b, n, d) - attn_sink = torch.randn(h) - idx = torch.stack( - [torch.stack([torch.randperm(n)[:topk] for _ in range(m)]) for _ in range(b)] - ).int() - scale = d**-0.5 - got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() - ref = _loop_reference(q, kv, attn_sink, idx, scale) - torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) - - -def test_sparse_attn_no_sink_matches_sdpa(): - """With a -inf sink (no sink mass) and all-valid contiguous indices, the - primitive must equal standard scaled-dot-product attention over the gathered - KV — an independent implementation.""" - torch.manual_seed(0) - b, m, h, d, topk = 2, 4, 5, 16, 9 - q = torch.randn(b, m, h, d) - kv = torch.randn(b, topk, d) # n == topk, attend to all - attn_sink = torch.full((h,), float("-inf")) - idx = torch.arange(topk).view(1, 1, -1).expand(b, m, topk).int() - scale = d**-0.5 - got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() - - # SDPA: q [b,h,m,d], k/v [b,h,topk,d] (broadcast the shared KV over heads). - qh = q.permute(0, 2, 1, 3) - kvh = kv.unsqueeze(1).expand(b, h, topk, d) - ref = F.scaled_dot_product_attention(qh, kvh, kvh, scale=scale).permute(0, 2, 1, 3) - torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) - - -def test_sparse_attn_masked_indices_excluded(): - """An index of -1 must be excluded exactly (equiv. to dropping that column).""" - torch.manual_seed(0) - b, m, h, d = 1, 1, 2, 16 - q = torch.randn(b, m, h, d) - kv = torch.randn(b, 5, d) - sink = torch.full((h,), float("-inf")) - scale = d**-0.5 - full = torch.tensor([[[0, 1, 2, 3]]]).int() - masked = torch.tensor([[[0, 1, 2, -1]]]).int() - drop3 = torch.tensor([[[0, 1, 2]]]).int() - got_masked = dspark_sparse_attn(q, kv, sink, masked, scale) - got_drop = dspark_sparse_attn(q, kv, sink, drop3, scale) - torch.testing.assert_close(got_masked, got_drop, rtol=1e-5, atol=1e-5) - # And masking genuinely changes the result vs attending to position 3. - got_full = dspark_sparse_attn(q, kv, sink, full, scale) - assert not torch.allclose(got_full, got_masked, rtol=1e-3, atol=1e-3) - - -@pytest.mark.parametrize("start_pos,window,block", [(3, 128, 5), (10, 4, 5)]) -def test_get_dspark_topk_idxs_matches_reference(start_pos, window, block): - bsz = 3 - got = get_dspark_topk_idxs(window, bsz, block, start_pos) - # Reference formula (DeepSpec get_dspark_topk_idxs). - ctx = torch.arange(min(window, start_pos + 1)) - blk = window + torch.arange(block) - ref_row = torch.cat([ctx, blk]).int() - assert got.shape == (bsz, block, ref_row.numel()) - for bi in range(bsz): - for mi in range(block): - torch.testing.assert_close(got[bi, mi], ref_row) - - -@pytest.mark.parametrize("rope_head_dim,seqlen", [(64, 16), (64, 1), (128, 8)]) -def test_precompute_freqs_cis_matches_reference(rope_head_dim, seqlen): - got = precompute_dspark_freqs_cis(rope_head_dim, seqlen, rope_theta=10000.0) - ref = _ref_precompute_freqs_cis(rope_head_dim, seqlen, 10000.0) - torch.testing.assert_close(got, ref) - - -@pytest.mark.parametrize("ndim", [3, 4]) -def test_apply_rotary_matches_reference(ndim): - torch.manual_seed(0) - b, s, h, rd = 2, 5, 4, 64 - x = torch.randn(b, s, h, rd) if ndim == 4 else torch.randn(b, s, rd) - fc = precompute_dspark_freqs_cis(rd, s) - got = apply_dspark_rotary(x, fc) - ref = _ref_apply_rotary_emb(x, fc) - torch.testing.assert_close(got, ref) - # De-rotation (inverse) must undo the forward rotation (property test). - roundtrip = apply_dspark_rotary(got, fc, inverse=True) - torch.testing.assert_close(roundtrip, x, rtol=1e-5, atol=1e-5) - - -def _make_attn_inputs(seed=0): - """Small synthetic DSpark attention inputs/weights (CPU bf16).""" - torch.manual_seed(seed) - dim, n_heads, head_dim, rd = 12, 4, 8, 4 - q_lora, o_lora, n_groups = 6, 5, 2 - window, block, start_pos = 8, 3, 5 - b = 2 - g = dict( - dim=dim, - n_heads=n_heads, - head_dim=head_dim, - rope_head_dim=rd, - q_lora=q_lora, - o_lora=o_lora, - n_groups=n_groups, - window=window, - block=block, - start_pos=start_pos, - b=b, - eps=1e-6, - softmax_scale=head_dim**-0.5, - ) - bf = torch.bfloat16 - g["x"] = torch.randn(b, block, dim, dtype=bf) - g["main_x"] = torch.randn(b, 1, dim, dtype=bf) - g["kv_cache0"] = torch.randn(b, window, head_dim, dtype=bf) - g["wq_a"] = torch.randn(q_lora, dim, dtype=bf) * 0.1 - g["wq_b"] = torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1 - g["wkv"] = torch.randn(head_dim, dim, dtype=bf) * 0.1 - g["wo_a"] = torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1 - g["wo_b"] = torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1 - g["q_norm"] = torch.ones(q_lora) - g["kv_norm"] = torch.ones(head_dim) - g["attn_sink"] = torch.randn(n_heads) - g["freqs"] = precompute_dspark_freqs_cis(rd, start_pos + 1 + block + 2) - return g - - -def _run(g): - return dspark_attention_forward( - g["x"], - g["main_x"], - g["start_pos"], - g["kv_cache0"], - wq_a=g["wq_a"], - q_norm_w=g["q_norm"], - wq_b=g["wq_b"], - wkv=g["wkv"], - kv_norm_w=g["kv_norm"], - wo_a=g["wo_a"], - wo_b=g["wo_b"], - attn_sink=g["attn_sink"], - n_heads=g["n_heads"], - head_dim=g["head_dim"], - rope_head_dim=g["rope_head_dim"], - n_groups=g["n_groups"], - o_lora_rank=g["o_lora"], - window_size=g["window"], - eps=g["eps"], - softmax_scale=g["softmax_scale"], - freqs_cis=g["freqs"], - ) - - -def test_attention_forward_shape_and_determinism(): - g = _make_attn_inputs() - before = g["kv_cache0"].clone() - o = _run(g) - assert tuple(o.shape) == (g["b"], g["block"], g["dim"]) - assert torch.isfinite(o.float()).all() - torch.testing.assert_close(o, _run(g)) # deterministic - # The rolling window write must be functional (cache cloned, not mutated). - torch.testing.assert_close(g["kv_cache0"], before) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py deleted file mode 100644 index ac989a423428..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_budget_parity_sglang.py +++ /dev/null @@ -1,232 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Decision-level parity with the SGLang budget planner this was ported from. - -A throughput gain alone does not validate the port: the wrong economics can -still move the needle in the right direction. This pins the DECISION function -against a line-for-line reference copy of SGLang's -``compute_verify_token_budget`` + ``_additive_step_time_tensor`` -(python/sglang/srt/speculative/dspark_components/dspark_planner.py), so a -divergence in either port shows up as an argmax mismatch here rather than as -an unexplained fleet-level regression. - -The two planners parameterize the same optimization differently: - -- SGLang admits EVERY position (including position 0) as a budget candidate - and its floor is one bonus token per request: total M = bs + n. -- TensorRT-LLM grants positions below ``min_verify_len`` for free (their - survival rides in the tau base) and only positions at or beyond it are - schedulable: total M = bs*(min_verify_len+1) + m. - -For ``min_verify_len=1`` these describe the same curve shifted by ``bs``: -survival is non-increasing along a row, so the top-(bs+m) SGLang candidates -are exactly the bs position-0 entries plus the top-m deeper entries, giving -``tau_sgl(bs+m) == tau_trt(m)`` and ``M_sgl(bs+m) == M_trt(m)``. The argmaxes -must therefore agree up to that shift whenever SGLang admits at least the -full position-0 set. Known, deliberate divergences are pinned separately: -SGLang filters candidates by ``survival_eps`` inside the budget function -(TensorRT-LLM applies eps at allocation time instead), and SGLang can choose -sub-floor budgets (windows of zero draft positions) that TensorRT-LLM's -``min_verify_len=1`` floor excludes. -""" - -import numpy as np -import pytest -import torch - -from tensorrt_llm._torch.speculative.dspark_planner import ( - SpsCostTable, compute_verify_token_budget) - -# -------------------------------------------------------------------------- -# Reference implementation, copied from SGLang dspark_components/dspark_planner -# (additive-table path) with only cosmetic renames. Do not "fix" this copy: -# its job is to be what SGLang runs. -# -------------------------------------------------------------------------- - - -def _sgl_interp_clamped(xs, ys, x: float) -> float: - xs = torch.tensor(xs, dtype=torch.float64) - ys = torch.tensor(ys, dtype=torch.float64) - x_t = torch.tensor(float(x), dtype=torch.float64).clamp_(xs[0], xs[-1]) - hi = torch.bucketize(x_t, xs, right=True).clamp_(1, xs.numel() - 1) - lo = hi - 1 - span = (xs[hi] - xs[lo]).clamp_(min=1e-9) - frac = (x_t - xs[lo]) / span - return float(ys[lo] + frac * (ys[hi] - ys[lo])) - - -def _sgl_additive_step_time(table, num_requests: int, num_budgets: int): - floor = table["bias_seconds"] + _sgl_interp_clamped( - table["bs_probes"], table["alpha_seconds"], float(num_requests)) - m_probes = torch.tensor(table["m_probes"], dtype=torch.float64) - theta_vals = torch.tensor(table["theta_seconds"], dtype=torch.float64) - m = (num_requests + torch.arange(num_budgets, dtype=torch.float64)).clamp_( - min=float(table["m_probes"][0]), max=float(table["m_probes"][-1])) - hi = torch.bucketize(m, m_probes, right=True).clamp_(1, m_probes.numel() - 1) - lo = hi - 1 - span = (m_probes[hi] - m_probes[lo]).clamp_(min=1e-9) - frac = (m - m_probes[lo]) / span - theta_at_m = theta_vals[lo] + frac * (theta_vals[hi] - theta_vals[lo]) - return floor + theta_at_m - - -def _sgl_compute_verify_token_budget(*, history_survival_probs, table, - max_verify_len, survival_eps): - num_requests = history_survival_probs.shape[0] - candidates = history_survival_probs[:, :max_verify_len].flatten() - candidates = candidates[candidates >= survival_eps].to(torch.float64) - candidates_sorted = torch.sort(candidates, descending=True).values - prefix_sum = torch.cumsum(candidates_sorted, dim=0) - tau_star = num_requests + torch.cat( - [torch.zeros(1, dtype=torch.float64), prefix_sum]) - step_time = _sgl_additive_step_time(table, int(num_requests), - int(tau_star.numel())) - theta = tau_star / step_time - return int(torch.argmax(theta)) - - -# -------------------------------------------------------------------------- -# One certified-table shape shared by both sides. TensorRT-LLM stores ms, -# SGLang stores seconds; the argmax is scale-invariant but the tables are -# built with the honest factor anyway. -# -------------------------------------------------------------------------- - -M_PROBES = (64, 96, 128, 192, 384, 512, 768, 1536) -THETA_MS = (9.68, 11.08, 11.95, 13.95, 22.09, 23.85, 35.61, 105.64) -BIAS_MS = 25.244443 -BS_PROBES = (32, 64, 128, 256) -ALPHA_MS = (0.0, 2.177477, 11.462212, 19.342661) -BLOCK = 5 - -TRT_TABLE = SpsCostTable(token_counts=M_PROBES, step_time_ms=THETA_MS, - fixed_overhead_ms=BIAS_MS, batch_sizes=BS_PROBES, - batch_overhead_ms=ALPHA_MS) -SGL_TABLE = { - "m_probes": M_PROBES, - "theta_seconds": tuple(v / 1e3 for v in THETA_MS), - "bias_seconds": BIAS_MS / 1e3, - "bs_probes": BS_PROBES, - "alpha_seconds": tuple(v / 1e3 for v in ALPHA_MS), -} - - -def _inversion_free_survival(rng, bs): - """Survivals where every position-0 beats every deeper candidate. - - The exact-parity theorem needs this: SGLang admits position 0 as a budget - candidate while TensorRT-LLM grants it in the floor, so the two describe - the same curve (shifted by ``bs``) only when the top-(bs+m) SGLang set is - exactly "all of position 0, plus the top-m deeper entries". A weak row's - position-0 falling below a strong row's position-1 breaks that -- see - test_starving_weak_rows_is_a_known_gap for what SGLang does with it. - """ - conf0 = rng.uniform(0.85, 0.99, size=(bs, 1)) - deeper = rng.uniform(0.25, 0.80, size=(bs, BLOCK - 1)) - return np.cumprod(np.concatenate([conf0, deeper], axis=1), axis=1) - - -@pytest.mark.parametrize("bs", [1, 3, 8, 64, 128, 252, 256]) -def test_budget_argmax_matches_sglang_up_to_the_floor_shift(bs): - """Same survival, same table, eps off, no inversions: shift-exact parity.""" - rng = np.random.default_rng(17 + bs) - for _ in range(5): - surv = _inversion_free_survival(rng, bs) - sgl_n = _sgl_compute_verify_token_budget( - history_survival_probs=torch.tensor(surv, dtype=torch.float32), - table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.0) - trt_m = compute_verify_token_budget( - survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, - min_verify_len=1) - if sgl_n >= bs: - assert trt_m == sgl_n - bs, ( - f"bs={bs}: SGLang admits {sgl_n} candidates " - f"(= floor {bs} + {sgl_n - bs}), TRT-LLM budget {trt_m}") - else: - # SGLang went below the one-draft floor; the closest budget - # TensorRT-LLM can express under min_verify_len=1 is zero. - assert trt_m == 0 - - -def test_live_pooled_survivals_park_on_the_breakpoint_together(): - """The campaign's operating point, both consumers, both argmax modes. - - With interpolated segments the continuous optimum parks where the theta - slope jumps: buying position-2 survivals (0.502/token) beats the - 512..768 slope (0.046 ms/token) but not the 768..1536 one (0.091), so - both planners stop at exactly M = 768 -- TensorRT-LLM at budget - 768 - 2*252 = 264, SGLang at 768 - 252 = 516. Tier alignment then rounds - down to the rung-2 budget the executor can capture. - """ - bs = 252 - surv = np.tile([0.771, 0.693, 0.502, 0.360, 0.315], (bs, 1)) - sgl_n = _sgl_compute_verify_token_budget( - history_survival_probs=torch.tensor(surv, dtype=torch.float32), - table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.0) - trt_m = compute_verify_token_budget( - survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, - min_verify_len=1) - assert trt_m == 768 - 2 * bs - assert sgl_n - bs == trt_m - tiered = compute_verify_token_budget( - survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, - min_verify_len=1, allowed_lens=[1, 2, 5]) - assert tiered == bs * 1 # rung-2: one scheduled position past the floor - - -def test_starving_weak_rows_is_a_known_gap(): - """SGLang can verify NOTHING for a hopeless request; TRT-LLM cannot. - - SGLang's budget admits position-0 candidates individually, so a row whose - draft is dead on arrival can be starved to a zero-length window and its - would-be tokens spent on a strong row's deep positions. TensorRT-LLM's - ``min_verify_len=1`` floor buys position 0 for every request - unconditionally. On a batch that mixes both kinds, SGLang's achievable - Theta is therefore strictly better -- pinned here so a future - min_verify_len=0 port is a deliberate decision with a test to flip. - """ - # Half the batch is strong all the way down, half is dead on arrival. - # The batch must be big enough that its token range sits on a genuinely - # rising part of the theta curve: below the first breakpoint the clamp - # makes cost flat, every candidate is free, and both planners trivially - # buy everything at the same theta -- the two-row version of this test - # sat in that regime and proved nothing. - bs = 64 - strong = np.tile([0.99, 0.98, 0.97, 0.96, 0.95], (bs // 2, 1)) - dead = np.tile([0.01, 0.008, 0.006, 0.004, 0.002], (bs // 2, 1)) - surv = np.concatenate([strong, dead], axis=0) - sgl_n = _sgl_compute_verify_token_budget( - history_survival_probs=torch.tensor(surv, dtype=torch.float32), - table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.0) - trt_m = compute_verify_token_budget( - survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, - min_verify_len=1) - - def theta_at(tau, tokens): - return tau / float(TRT_TABLE.step_times(np.asarray([tokens]), bs)[0]) - - cand = np.sort(surv.reshape(-1))[::-1] - sgl_theta = theta_at(bs + cand[:sgl_n].sum(), bs + sgl_n) - deeper = np.sort(surv[:, 1:].reshape(-1))[::-1] - trt_theta = theta_at(bs + surv[:, 0].sum() + deeper[:trt_m].sum(), - 2 * bs + trt_m) - assert sgl_theta > trt_theta - - -def test_eps_filtering_is_a_known_divergence(): - """SGLang drops sub-eps candidates inside the budget; TRT-LLM does not. - - TensorRT-LLM applies survival_eps at allocation time - (schedule_verify_lens_topk's eligibility mask), so its BUDGET can count a - candidate the allocator then refuses. This pins the divergence so a - future unification is a deliberate change, not an accident: with a heavy - sub-eps tail the SGLang budget must not exceed TRT-LLM's shifted one. - """ - bs = 16 - surv = np.tile([0.9, 0.4, 0.008, 0.004, 0.002], (bs, 1)) - sgl_n = _sgl_compute_verify_token_budget( - history_survival_probs=torch.tensor(surv, dtype=torch.float32), - table=SGL_TABLE, max_verify_len=BLOCK, survival_eps=0.01) - trt_m = compute_verify_token_budget( - survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, - min_verify_len=1) - assert sgl_n - bs <= trt_m diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_compressor_gen_counts.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_compressor_gen_counts.py deleted file mode 100644 index 7bf647f680db..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_compressor_gen_counts.py +++ /dev/null @@ -1,78 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The compressor's per-request new-token counts must follow the DEVICE buffer. - -``_sync_gen_tokens_per_seq`` runs inside the captured ``on_update_kv_lens``, -so its writes are replayed every step. The device-window prologue re-ranks -verify windows from the host shape split S to true windows w by rewriting -``_seq_lens_cuda`` (via ``apply_device_ragged_layout``) after host staging -but before replay. Two properties keep the compressor correct under that -rewrite, and both were violated by the original implementation: - -1. The per-request count vector must be a D2D copy sourced from - ``_seq_lens_cuda`` -- a captured H2D copy from a pinned host list re-loads - S on every replay, so the kernel computes its state start position - ``sp = kv_len - nn`` with a w-corrected ``kv_len`` but ``nn = S``, and - writes the persistent paged KV/score state shifted by (w - S). -2. The ``next_n`` template scalar must be the GLOBAL max window, not this - step's batch max: the prologue may assign any row up to the config max on - any bucket, and the kernel treats the scalar as a guarded upper bound. -""" - -import types - -import torch - -from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import \ - DeepseekV4TrtllmAttentionMetadata - - -def _dummy(nc, staged_split, seq_lens_dev, max_draft_tokens=5, cap=16): - d = types.SimpleNamespace() - d.is_ragged_verify = True - d.num_generations = len(staged_split) - d.num_contexts = nc - d.max_draft_tokens = max_draft_tokens - # What the host shape split says (S); the fixed code must NOT read it. - d.ragged_verify_lens = list(staged_split) - d.gen_new_tokens_per_seq_cuda = torch.zeros(cap, dtype=torch.int) - d._seq_lens_cuda = torch.tensor(seq_lens_dev, dtype=torch.int) - return d - - -def _call(d, num_gen_tokens): - return DeepseekV4TrtllmAttentionMetadata._sync_gen_tokens_per_seq( - d, num_gen_tokens) - - -def test_vector_sources_from_seq_lens_not_host_split(): - # Simulate the prologue: host staged S, device rewrote seq lens to w. - S = [3, 3, 3] - w = [2, 4, 3] - d = _dummy(nc=0, staged_split=S, seq_lens_dev=w) - _call(d, sum(w)) - assert d.gen_new_tokens_per_seq.tolist() == w - - -def test_scalar_is_global_bound_not_batch_max(): - S = [2, 2] - d = _dummy(nc=0, staged_split=S, seq_lens_dev=S) - next_n = _call(d, sum(S)) - assert next_n == 1 + d.max_draft_tokens - assert d.num_gen_tokens_per_seq == 1 + d.max_draft_tokens - - -def test_gen_rows_slice_skips_context_rows(): - ctx_lens = [17, 9] - w = [4, 2] - d = _dummy(nc=2, staged_split=w, seq_lens_dev=ctx_lens + w) - _call(d, sum(w)) - assert d.gen_new_tokens_per_seq.tolist() == w - - -def test_uniform_branch_unchanged(): - d = _dummy(nc=0, staged_split=[3, 3], seq_lens_dev=[3, 3]) - d.is_ragged_verify = False - next_n = _call(d, 6) - assert next_n == 3 - assert d.gen_new_tokens_per_seq is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py deleted file mode 100644 index ffad376de6e4..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.py +++ /dev/null @@ -1,509 +0,0 @@ -# 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. -"""Unit tests for DSpark confidence-scheduled verification (hardware-agnostic, CPU). - -These cover the parts where a silent error would be expensive and invisible at -runtime: the survival semantics, the budget argmax over a staircase cost curve, -the prefix property of the allocation, and cross-rank determinism. -""" - -import itertools - -import numpy as np -import pytest -import torch - -from tensorrt_llm._torch.speculative.dspark_planner import ( - SpsCostTable, - budget_argmax_over_uniform_lens, - compute_verify_token_budget, - derive_verify_len_tiers, -) -from tensorrt_llm._torch.speculative.dspark_schedule import ( - DSparkScheduleConfig, - compute_survival, - schedule_verify_lens_topk, -) - -BLOCK = 7 - - -def _cfg(**kw) -> DSparkScheduleConfig: - return DSparkScheduleConfig(block_size=BLOCK, **kw) - - -# -------------------------------------------------------------------------- -# survival -# -------------------------------------------------------------------------- - - -def test_survival_is_cumulative_product(): - conf = torch.tensor([[0.9, 0.8, 0.5, 1.0, 1.0, 1.0, 1.0]]) - surv = compute_survival(conf) - assert torch.allclose(surv[0, :3], torch.tensor([0.9, 0.72, 0.36]), atol=1e-6) - - -# -------------------------------------------------------------------------- -# budget argmax -# -------------------------------------------------------------------------- - - -def _brute_force_budget(survival, num_gen, table, min_len=1, max_len=None): - """Independent O(N) reimplementation used as the oracle for the argmax. - - Spells out the objective: expected emitted tokens per millisecond, where the - expectation counts one bonus token per request, the always-verified floor - positions, and the admitted candidates. The cost is looked up on the tokens - actually submitted to the target -- ``bs * (min_len + 1)`` at the floor, - because every request also submits its bonus position. - """ - bs, blk = survival.shape - cap = min(max_len or blk, blk) - cand = np.sort(survival[:, min_len:cap].reshape(-1))[::-1] - base = num_gen + float(survival[:, :min_len].sum()) - floor_tokens = bs * (min_len + 1) - best_n, best = 0, -np.inf - for n in range(cand.size + 1): - theta = (base + float(cand[:n].sum())) / table.step_time(floor_tokens + n, num_gen) - if theta > best: - best, best_n = theta, n - return best_n - - -def test_budget_matches_brute_force_on_staircase(): - """Theta is not unimodal on a staircase; a greedy first-descent scan is wrong.""" - rng = np.random.default_rng(7) - # A cost curve with genuine shelves and risers. - table = SpsCostTable( - token_counts=(0, 64, 128, 192, 256, 320), - step_time_ms=(4.0, 4.05, 6.5, 6.55, 9.0, 9.05), - ) - for trial in range(25): - bs = int(rng.integers(1, 40)) - conf = rng.uniform(0.55, 0.99, size=(bs, BLOCK)) - surv = np.cumprod(conf, axis=1) - got = compute_verify_token_budget(survival=surv, num_gen_requests=bs, cost_table=table) - want = _brute_force_budget(surv, bs, table) - assert got == want, f"trial {trial}: bs={bs} got={got} want={want}" - - -def test_flat_cost_table_degenerates_to_verify_all(): - """A flat cost model makes every token free -> spend the whole budget.""" - table = SpsCostTable.flat(2.0) - assert table.is_flat - surv = np.cumprod(np.full((6, BLOCK), 0.9), axis=1) - budget = compute_verify_token_budget(survival=surv, num_gen_requests=6, cost_table=table) - assert budget == 6 * (BLOCK - 1) - # An empty batch has no candidates to buy, whatever the table says. - assert ( - compute_verify_token_budget( - survival=np.zeros((0, BLOCK)), num_gen_requests=0, cost_table=table - ) - == 0 - ) - - -def test_cost_table_interpolates_between_breakpoints(): - """The consumer contract is clamped linear interpolation, not a floor. - - The floor variant this replaces billed every total below the next - breakpoint at the previous breakpoint's price. On the sparse live table - (points at 768 and 1536) that priced a 1512-token full block at the 768 - price -- the planner's cost ratio collapsed below the survival ratio and - it bought the full block on ~95% of decisions. Flatness between points is - a claim about the hardware; if a table wants a shelf, it must MEASURE the - shelf (two breakpoints with equal values), not have the consumer assume it. - """ - table = SpsCostTable(token_counts=(0, 100), step_time_ms=(1.0, 5.0)) - assert table.step_time(99) > table.step_time(1) - assert table.step_time(50) == pytest.approx(3.0) - # A measured shelf stays flat under interpolation. - shelf = SpsCostTable(token_counts=(0, 99, 100), step_time_ms=(1.0, 1.0, 5.0)) - assert shelf.step_time(1) == shelf.step_time(99) - # Above the last breakpoint we clamp rather than extrapolate. - assert table.step_time(10_000) == table.step_time(100) - - -def test_cost_table_rejects_malformed_input(): - with pytest.raises(ValueError, match="strictly increasing"): - SpsCostTable(token_counts=(0, 0), step_time_ms=(1.0, 2.0)) - with pytest.raises(ValueError, match="same length"): - SpsCostTable(token_counts=(0, 1), step_time_ms=(1.0,)) - with pytest.raises(ValueError, match="positive"): - SpsCostTable(token_counts=(0,), step_time_ms=(0.0,)) - - -# -------------------------------------------------------------------------- -# per-request allocation -# -------------------------------------------------------------------------- - - -def test_allocation_is_a_prefix_and_respects_bounds(): - torch.manual_seed(3) - cfg = _cfg(min_verify_len=1) - surv = compute_survival(torch.rand(12, BLOCK)) - lens = schedule_verify_lens_topk(survival=surv, budget=20, cfg=cfg) - assert lens.dtype == torch.int32 - assert torch.all(lens >= cfg.min_verify_len) - assert torch.all(lens <= cfg.resolved_max_verify_len) - # Budget is spent above the floor and never exceeded. - assert int((lens - cfg.min_verify_len).sum()) <= 20 - # Zero budget degenerates to the floor for everyone. - cfg_floor = _cfg(min_verify_len=2) - lens = schedule_verify_lens_topk( - survival=compute_survival(torch.rand(5, BLOCK)), budget=0, cfg=cfg_floor) - assert torch.equal(lens, torch.full((5,), 2, dtype=torch.int32)) - # A huge budget saturates at the cap. - cfg_cap = _cfg(min_verify_len=1, max_verify_len=4) - lens = schedule_verify_lens_topk( - survival=compute_survival(torch.full((5, BLOCK), 0.99)), budget=10**6, cfg=cfg_cap) - assert torch.equal(lens, torch.full((5,), 4, dtype=torch.int32)) - # An empty batch allocates nothing. - assert schedule_verify_lens_topk( - survival=torch.zeros(0, BLOCK), budget=5, cfg=_cfg()).numel() == 0 - - -def test_survival_eps_excludes_hopeless_positions(): - """Even with unlimited budget, sub-eps candidates are never admitted.""" - cfg = _cfg(min_verify_len=1, survival_eps=1e-3) - conf = torch.tensor([[1.0, 1.0, 1e-9, 1.0, 1.0, 1.0, 1.0]]) - lens = schedule_verify_lens_topk(survival=compute_survival(conf), budget=10**6, cfg=cfg) - # positions 0,1 survive; from position 2 on, survival collapses below eps. - assert int(lens[0]) == 2 - - -def test_allocation_is_deterministic_under_tied_confidences(): - """Two ranks with identical input must produce identical verify lengths.""" - cfg = _cfg(min_verify_len=1) - surv = compute_survival(torch.full((9, BLOCK), 0.9)) - first = schedule_verify_lens_topk(survival=surv, budget=13, cfg=cfg) - for _ in range(8): - assert torch.equal(schedule_verify_lens_topk(survival=surv, budget=13, cfg=cfg), first) - - -def test_ties_break_toward_earlier_positions(): - """A tie must not hand a later position to one request over an earlier one. - - Every request has identical confidence, so all candidates at a given - position tie. The allocation must stay balanced (a prefix front), not give - one request a deep suffix while another gets nothing. - """ - cfg = _cfg(min_verify_len=1) - surv = compute_survival(torch.full((4, BLOCK), 1.0)) - lens = schedule_verify_lens_topk(survival=surv, budget=4, cfg=cfg) - assert int(lens.max()) - int(lens.min()) <= 1 - - -def test_schedule_config_rejects_bad_bounds(): - with pytest.raises(ValueError, match="min_verify_len must be >= 1"): - DSparkScheduleConfig(block_size=BLOCK, min_verify_len=0) - with pytest.raises(ValueError, match="block_size must be >= 1"): - DSparkScheduleConfig(block_size=0) - with pytest.raises(ValueError, match="< min_verify_len"): - DSparkScheduleConfig(block_size=BLOCK, min_verify_len=5, max_verify_len=3) - with pytest.raises(ValueError, match=r"\[bs, K\]"): - compute_survival(torch.rand(BLOCK)) - - -# -------------------------------------------------------------------------- -# budget conservation: allocation must never exceed what the planner sized -# -------------------------------------------------------------------------- - - -@pytest.mark.parametrize("bs,min_len", list(itertools.product([1, 3, 17], [1, 2]))) -def test_planner_and_allocator_agree_on_the_budget(bs, min_len): - rng = np.random.default_rng(bs * 100 + min_len) - table = SpsCostTable(token_counts=(0, 16, 32, 64, 128), step_time_ms=(3.0, 3.1, 5.0, 5.1, 9.0)) - conf = rng.uniform(0.4, 0.99, size=(bs, BLOCK)) - surv_np = np.cumprod(conf, axis=1) - budget = compute_verify_token_budget( - survival=surv_np, num_gen_requests=bs, cost_table=table, min_verify_len=min_len - ) - lens = schedule_verify_lens_topk( - survival=torch.from_numpy(surv_np).float(), - budget=budget, - cfg=_cfg(min_verify_len=min_len), - ) - assert int((lens - min_len).sum()) <= budget - assert torch.all(lens >= min_len) - - -# -------------------------------------------------------------------------- -# uniform-K projection (what TRT-LLM's graph key can actually express today) -# -------------------------------------------------------------------------- - - -def test_discrete_argmax_beats_round_then_snap_across_a_riser(): - """Optimizing over runnable lengths is not the same as rounding the budget. - - The continuous optimum sits just past a cost riser; snapping down lands on a - length whose Theta is worse than a length the discrete search finds. - """ - table = SpsCostTable(token_counts=(0, 8, 24, 40), step_time_ms=(2.0, 2.05, 2.1, 20.0)) - surv = np.cumprod(np.full((8, BLOCK), 0.97), axis=1) - chosen = budget_argmax_over_uniform_lens( - survival=surv, - num_gen_requests=8, - cost_table=table, - allowed_lens=[1, 3, 5, 7], - ) - # 8*(7+1)=64 tokens lands on the 20ms riser; the search must avoid it. - assert chosen < 7 - - def theta(length): - # Mirrors the implementation exactly, including the bonus token every - # request submits alongside its drafts. - tau = 8 + float(surv[:, :length].sum()) - return tau / table.step_time(8 * (length + 1), 8) - - assert theta(chosen) == max(theta(v) for v in (1, 3, 5, 7)) - # With nothing runnable, fall back to the floor. - assert ( - budget_argmax_over_uniform_lens( - survival=surv, num_gen_requests=8, cost_table=table, allowed_lens=[] - ) - == 1 - ) - - -def test_derived_tiers_lose_nothing_versus_continuous_k(): - """The shelf-right-edge property: derived tiers contain the true optimum. - - Within a cost shelf the step time is constant while tau strictly increases, - so Theta rises monotonically across the shelf and the optimum can only sit - at a right edge. A tier set built from the right edges must therefore match - an exhaustive search over *every* length. - """ - rng = np.random.default_rng(11) - # The shelves are ENCODED, breakpoint pairs with equal values: under the - # interpolating consumer, flatness between points is only real when the - # table measured it. (Assuming it -- the old floor lookup -- is what - # priced a 1512-token step at the 768-token price on the live table and - # disabled trimming outright; the theorem below is conditional on genuine - # shelves and this encoding is what "genuine" means now.) - table = SpsCostTable( - token_counts=(0, 11, 12, 29, 30, 47, 48, 89, 90, 159, 160), - step_time_ms=(2.0, 2.0, 2.4, 2.4, 3.9, 3.9, 4.0, 4.0, 7.5, 7.5, 12.0), - ) - for bs in (3, 8, 16, 31): - surv = np.cumprod(rng.uniform(0.5, 0.995, size=(bs, BLOCK)), axis=1) - tiers = derive_verify_len_tiers( - cost_table=table, num_requests=bs, block_size=BLOCK, max_tiers=99 - ) - - def theta(length): - return (bs + float(surv[:, :length].sum())) / table.step_time(bs * (length + 1), bs) - - best_any = max(theta(v) for v in range(1, BLOCK + 1)) - best_tier = max(theta(v) for v in tiers) - assert best_tier == pytest.approx(best_any), f"bs={bs} tiers={tiers}" - - -def test_derived_tiers_respect_the_capture_budget(): - table = SpsCostTable( - token_counts=(0, 5, 9, 14, 20, 27, 35), step_time_ms=(1.0, 2, 3, 4, 5, 6, 7) - ) - tiers = derive_verify_len_tiers(cost_table=table, num_requests=2, block_size=BLOCK, max_tiers=3) - assert len(tiers) <= 3 - # Endpoints are always kept: the floor is always runnable and the full block - # is always a right edge (the last shelf is unbounded). - assert tiers[0] == 1 and tiers[-1] == BLOCK - - -def test_derived_tiers_are_a_function_of_batch_size(): - """Tiers move with bs because total tokens = bs * (length + 1). - - With a riser at 40 tokens: at bs=8, length 3 costs 8*4=32 tokens (still on - the shelf) while length 4 costs 8*5=40 (past it), so 3 is the right edge; at - bs=32 even length 1 is already past the riser, so the shelf yields nothing - and only the endpoints remain. A tier set derived once and reused across - batch sizes would be wrong for one of them. - """ - table = SpsCostTable(token_counts=(0, 40), step_time_ms=(1.0, 5.0)) - assert derive_verify_len_tiers( - cost_table=table, num_requests=8, block_size=BLOCK, max_tiers=99 - ) == [1, 3, BLOCK] - assert derive_verify_len_tiers( - cost_table=table, num_requests=32, block_size=BLOCK, max_tiers=99 - ) == [1, BLOCK] - - -# -------------------------------------------------------------------------- -# host-side planner: relay, fallbacks, cross-rank agreement -# -------------------------------------------------------------------------- - - -def _planner(**kw): - from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner - - kw.setdefault("cfg", _cfg(min_verify_len=1)) - kw.setdefault( - "cost_table", - SpsCostTable(token_counts=(0, 8, 24, 48), step_time_ms=(2.0, 2.1, 4.0, 9.0)), - ) - kw.setdefault("tiers", [1, 3, BLOCK]) - return DSparkVerifyPlanner(**kw) - - -def test_planner_reads_confidence_by_row_not_by_batch_position(): - """The snapshot is slot-indexed; ``rows`` is what re-associates it. - - The buffer is written by slot and read one iteration later, by which point - joins and departures have reshuffled the batch. Position 0 of this step's - batch is routinely a different request than the one scored in row 0. - """ - p = _planner() - # Row 3 is a confident request, row 0 a hopeless one. - buf = torch.full((5, BLOCK), -8.0) - buf[3] = 8.0 - p._host_buffer, p._copy_event, p._snapshot_valid = buf, None, True - - # This step schedules the confident request first, then the hopeless one. - lens = p.decide_verify_lens(num_gen_requests=2, rows=[3, 0]) - assert lens is not None and lens[0] > lens[1], f"row mapping ignored: {lens}" - - -def test_planner_rejects_a_row_list_that_does_not_match_the_batch(): - p = _planner() - p._host_buffer, p._copy_event, p._snapshot_valid = torch.rand(8, BLOCK), None, True - assert p.decide_verify_lens(num_gen_requests=4, rows=[0, 1]) is None - assert p.stats["fallback_short_snapshot"] == 1 - - -def _capture_set(batch_sizes, tiers, max_draft_len=BLOCK): - """Invoke PyTorchModelEngine._get_graphs_to_capture with a stub engine.""" - import types - - from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine - - spec_config = types.SimpleNamespace( - enable_confidence_scheduling=True, - verify_len_tiers=list(tiers), - draft_len_schedule=None, - max_concurrency=None, - is_linear_tree=True, - spec_dec_mode=types.SimpleNamespace( - support_dynamic_draft_len=lambda: False, use_one_engine=lambda: False - ), - ) - engine = types.SimpleNamespace( - is_draft_model=False, - spec_config=spec_config, - max_draft_len=max_draft_len, - max_total_draft_tokens=max_draft_len, - original_max_draft_len=max_draft_len, - _dynamic_draft_len_mapping=None, - _cuda_graph_batch_sizes=list(batch_sizes), - ) - return PyTorchModelEngine._get_graphs_to_capture(engine, list(batch_sizes), None) - - -def test_capture_pins_the_draft_length_to_the_top_tier(): - """One graph per batch size, not a cross product with the tier ladder. - - The cross product was right while the uniform tier path existed and chose - a different draft_len per step. That path is gone and nothing varies - draft_len at runtime now: cap-accept pins it to planner.max_tier and - compact moves the ladder to the token axis. Capturing the lower tiers - anyway spent KV cache -- captured graphs come out of the same pool -- on - graphs that could never be selected. Measured on a real run: 57 graphs - where 19 were reachable. - """ - bss, tiers = [1, 8, 64], [1, 3, BLOCK] - graphs = _capture_set(bss, tiers) - assert set(graphs) == {(bs, BLOCK) for bs in bss} - assert len(graphs) == len(bss), ( - f"expected one graph per batch size, got {len(graphs)}: {graphs}") - - -def _stats(max_draft_len=5): - from tensorrt_llm._torch.speculative.dspark_observability import ( - DSparkRaggedStats, RaggedVerifyMode) - return DSparkRaggedStats(mode=RaggedVerifyMode.COMPACT, - max_draft_len=max_draft_len) - - -def test_trim_regret_counts_only_drafts_alive_at_the_cut(): - """A trimmed request that accepted its whole window lost something. - - This is the quantity a delivered-only acceptance metric cannot recover: it - separates "trimming was free, those drafts would have died anyway" from - "trimming bought throughput by discarding acceptance". - """ - stats = _stats() - # Given the full block and died early: not trimmed, no regret. - stats.record_acceptance(accepted=2, window=5) - # Trimmed to 3 and died at 1: the cut cost nothing. - stats.record_acceptance(accepted=1, window=3) - # Trimmed to 3 and accepted all 3: still alive at the cut -> regret. - stats.record_acceptance(accepted=3, window=3) - - assert stats.requests_scored == 3 - assert stats.requests_trimmed == 2 - assert stats.trimmed_hit_ceiling == 1 - assert stats.trim_regret_rate == 0.5 - assert stats.accept_len == pytest.approx(2.0) - summary = stats.summary() - for key in ("accept_len", "requests_scored", "requests_trimmed", - "trim_regret_rate"): - assert key in summary, f"{key} missing from the summary" - - # Nothing trimmed: the regret rate is 0/0-safe and stays 0. - untrimmed = _stats() - for accepted in (0, 3, 5): - untrimmed.record_acceptance(accepted=accepted, window=5) - assert untrimmed.requests_trimmed == 0 - assert untrimmed.trim_regret_rate == 0.0 - assert untrimmed.summary()["accept_len"] == pytest.approx(8 / 3, abs=1e-4) - - -# -------------------------------------------------------------------------- -# top-k allocation follows confidence (moved from test_dspark_planner_trims.py) -# -------------------------------------------------------------------------- - - -def _block5_cfg() -> DSparkScheduleConfig: - return DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) - - -def test_topk_hands_longer_windows_to_more_confident_requests(): - """The budget must follow survival, not batch position.""" - # survival[r, k] = P(the first k+1 drafted tokens all get accepted) - survival = torch.tensor([ - [0.99, 0.98, 0.97, 0.96, 0.95], # confident - [0.90, 0.70, 0.50, 0.30, 0.10], # middling - [0.20, 0.04, 0.01, 0.00, 0.00], # collapses - ]) - lens = schedule_verify_lens_topk(survival=survival, budget=6, - cfg=_block5_cfg()).tolist() - - assert len(lens) == 3 - assert lens[0] > lens[2], ( - f"confident request got {lens[0]}, collapsing request got {lens[2]}; " - f"the budget must follow survival, not batch position") - assert lens[0] >= lens[1] >= lens[2] - # Every request keeps at least the floor, and the budget is respected. - assert min(lens) >= 1 - assert sum(lens) - len(lens) * 1 <= 6 - - -def test_a_full_budget_degenerates_to_the_uniform_full_window(): - """The no-trim case must stay exactly uniform: ragged must cost nothing.""" - num_reqs, max_len = 4, 5 - survival = torch.full((num_reqs, max_len), 0.99) - budget = num_reqs * (max_len - 1) - lens = schedule_verify_lens_topk(survival=survival, budget=budget, - cfg=_block5_cfg()).tolist() - assert lens == [max_len] * num_reqs diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py deleted file mode 100644 index acbf6eba383f..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cost_interpolation.py +++ /dev/null @@ -1,61 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Cost lookup must interpolate between breakpoints, not floor onto them. - -The floor lookup this replaces silently disabled the whole feature: with -measured points at 768 and 1536, every total in 769..1535 was billed at the -768 price, so a full block at bs=252 (1512 tokens, true cost ~148 ms) was -priced 72.3 ms. The planner's cost ratio T(6bs)/T(3bs) collapsed from a true -~2.0 to 1.19 -- below the GSM8K survival ratio tau5/tau2 ~ 1.48 -- so the -argmax bought the full block on ~95% of decisions, the bucket landed on -padded_bs*6 on every step of every run, and trimming never happened. The only -exception ever observed was bs=256 exactly, where the tier totals land on -breakpoints and floor happens to be exact. - -The consumer this was ported from (SGLang ``_additive_step_time_tensor``) -interpolates both theta(M) and alpha(bs); this port must match. -""" - -import pytest - -from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable - -# The certified GB300 table's shape, abbreviated to the two breakpoints whose -# gap hid the bug. -TABLE = SpsCostTable(token_counts=(512, 768, 1536), - step_time_ms=(23.85, 35.61, 105.64), - fixed_overhead_ms=25.244, - batch_sizes=(128, 256), - batch_overhead_ms=(11.462, 19.343)) - - -def test_theta_interpolates_between_breakpoints(): - """Exact on breakpoints, linear between them. - - The regression case: 1512 tokens must cost ~the 1536 price, not the 768 - one. - """ - assert TABLE.step_time(768, 256) == pytest.approx(25.244 + 19.343 + 35.61) - assert TABLE.step_time(1536, 256) == pytest.approx(25.244 + 19.343 + 105.64) - got = TABLE.step_time(1512, 252) - theta = 35.61 + (105.64 - 35.61) * (1512 - 768) / (1536 - 768) - alpha = 11.462 + (19.343 - 11.462) * (252 - 128) / (256 - 128) - assert got == pytest.approx(25.244 + alpha + theta) - # The floor price it used to return -- and must never return again. - assert got > 130.0 - - -def test_alpha_interpolates_and_clamps(): - assert TABLE.batch_overhead(128) == pytest.approx(11.462) - assert TABLE.batch_overhead(192) == pytest.approx((11.462 + 19.343) / 2) - assert TABLE.batch_overhead(64) == pytest.approx(11.462) # clamp low - assert TABLE.batch_overhead(512) == pytest.approx(19.343) # clamp high - - -def test_theta_clamps_outside_the_measured_range(): - """Theta clamps to the end values; a table without a batch axis adds no alpha.""" - plain = SpsCostTable(token_counts=(512, 768, 1536), - step_time_ms=(23.85, 35.61, 105.64), - fixed_overhead_ms=25.244) - assert plain.step_time(100, 0) == pytest.approx(25.244 + 23.85) - assert plain.step_time(4096, 0) == pytest.approx(25.244 + 105.64) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py deleted file mode 100644 index 9b1a27bc24ff..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_cuda_graph.py +++ /dev/null @@ -1,445 +0,0 @@ -# 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. -"""CUDA-graph-safety unit tests for the DSpark batched draft-attention path. - -The load-bearing invariant for the DSpark batched draft path (the default, used -whenever ``cuda_graph_config`` is set, since the one-engine drafter is captured in -the target's graph) is that the batched, sync-free primitives are **numerically -identical, per request**, to the validated scalar path — only the host-int -``start_pos`` and the per-request window indexing are tensorized. These tests assert -that equivalence on CPU (so they run in pre-merge CI without a GPU), plus a -GPU-gated capture+replay smoke test that proves the batched attention is actually -graph-capturable. -""" - -import pytest -import torch - -from tensorrt_llm._torch.models.dspark.attention import ( - apply_dspark_rotary, - apply_dspark_rotary_batched, - dspark_attention_forward, - dspark_attention_forward_batched, - get_dspark_topk_idxs, - get_dspark_topk_idxs_batched, - precompute_dspark_freqs_cis, -) - - -def _make_batched_inputs(seed=0, start_positions=(1, 3, 20)): - """Per-request DSpark attention inputs/weights (shared weights, distinct pos). - - Mirrors ``test_dspark_attention._make_attn_inputs`` but builds ``G`` requests - each with its own ``start_pos`` (small => partial context, large => full - rolling window) and its own pre-seeded window, so the batched-vs-scalar - comparison exercises the per-request RoPE gather + windowed context read. - """ - torch.manual_seed(seed) - dim, n_heads, head_dim, rd = 12, 4, 8, 4 - q_lora, o_lora, n_groups = 6, 5, 2 - window, block = 8, 3 - G = len(start_positions) - bf = torch.bfloat16 - # A single fixed RoPE table covering every request's positions (both paths - # index/gather the same values, so freqs are identical across paths). - maxlen = max(start_positions) + 1 + block + 4 - g = dict( - dim=dim, - n_heads=n_heads, - head_dim=head_dim, - rope_head_dim=rd, - q_lora=q_lora, - o_lora=o_lora, - n_groups=n_groups, - window=window, - block=block, - G=G, - start_positions=list(start_positions), - eps=1e-6, - softmax_scale=head_dim**-0.5, - x=torch.randn(G, block, dim, dtype=bf), - main_x=torch.randn(G, 1, dim, dtype=bf), - # Distinct, non-zero seeded window per request to exercise context reads. - kv_cache=torch.randn(G, window, head_dim, dtype=bf), - wq_a=torch.randn(q_lora, dim, dtype=bf) * 0.1, - wq_b=torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1, - wkv=torch.randn(head_dim, dim, dtype=bf) * 0.1, - wo_a=torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1, - wo_b=torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1, - q_norm=torch.ones(q_lora), - kv_norm=torch.ones(head_dim), - attn_sink=torch.randn(n_heads), - freqs=precompute_dspark_freqs_cis(rd, maxlen), - ) - return g - - -def _attn_kwargs(g): - return dict( - wq_a=g["wq_a"], - q_norm_w=g["q_norm"], - wq_b=g["wq_b"], - wkv=g["wkv"], - kv_norm_w=g["kv_norm"], - wo_a=g["wo_a"], - wo_b=g["wo_b"], - attn_sink=g["attn_sink"], - n_heads=g["n_heads"], - head_dim=g["head_dim"], - rope_head_dim=g["rope_head_dim"], - n_groups=g["n_groups"], - o_lora_rank=g["o_lora"], - window_size=g["window"], - eps=g["eps"], - softmax_scale=g["softmax_scale"], - freqs_cis=g["freqs"], - ) - - -def _scalar_reference(g): - """Run the validated scalar attention once per request and stack the outputs.""" - outs = [] - for i, sp in enumerate(g["start_positions"]): - out_i = dspark_attention_forward( - g["x"][i : i + 1], - g["main_x"][i : i + 1], - int(sp), - g["kv_cache"][i : i + 1].clone(), - **_attn_kwargs(g), - ) - outs.append(out_i) - return torch.cat(outs, dim=0) - - -def _batched(g, persist=False): - G = g["G"] - start_pos = torch.tensor(g["start_positions"], dtype=torch.long) - slots = torch.arange(G, dtype=torch.long) - return dspark_attention_forward_batched( - g["x"], - g["main_x"], - start_pos, - g["kv_cache"].clone(), - slots, - persist=persist, - **_attn_kwargs(g), - ) - - -@pytest.mark.parametrize("seed", [0, 1, 2]) -def test_batched_attention_matches_scalar_per_request(seed): - """Batched attention == per-request scalar attention at distinct start_pos. - - This is the invariant that lets the batched path replace the per-request loop - under CUDA graphs without changing draft quality / greedy parity. - """ - g = _make_batched_inputs(seed=seed) - ref = _scalar_reference(g) - got = _batched(g) - assert tuple(got.shape) == (g["G"], g["block"], g["dim"]) - torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2) - - -def test_batched_attention_persist_writes_through_window(): - """persist gates the window write: False mutates nothing, True writes - main_kv into exactly the start_pos%window row of each request.""" - g = _make_batched_inputs(seed=3) - G, win = g["G"], g["window"] - start_pos = torch.tensor(g["start_positions"], dtype=torch.long) - slots = torch.arange(G, dtype=torch.long) - cache = g["kv_cache"].clone() - before = cache.clone() - - dspark_attention_forward_batched( - g["x"], g["main_x"], start_pos, cache, slots, persist=False, **_attn_kwargs(g) - ) - # persist=False must not mutate the caller's window (functional). - torch.testing.assert_close(cache, before) - - dspark_attention_forward_batched( - g["x"], g["main_x"], start_pos, cache, slots, persist=True, **_attn_kwargs(g) - ) - # Exactly the start_pos%win row of each request changed. - for i, sp in enumerate(g["start_positions"]): - changed = (cache[i] != before[i]).any(dim=-1) - expected = torch.zeros(win, dtype=torch.bool) - expected[sp % win] = True - assert torch.equal(changed, expected), f"req {i}: wrong window row written" - - -@pytest.mark.parametrize("start_positions", [(1, 3, 20), (5, 5, 5), (2, 7, 200)]) -def test_batched_topk_matches_scalar(start_positions): - """Fixed-size masked batched topk == scalar topk per request (valid slots).""" - window, block = 8, 3 - start_pos = torch.tensor(start_positions, dtype=torch.long) - batched = get_dspark_topk_idxs_batched(window, block, start_pos) - assert tuple(batched.shape) == (len(start_positions), block, window + block) - for i, sp in enumerate(start_positions): - scalar = get_dspark_topk_idxs(window, 1, block, int(sp))[0] # [block, topk_i] - # The batched row drops the -1-masked context slots to recover the scalar - # (variable-width) index list; both must then be identical. - for m in range(block): - valid = batched[i, m][batched[i, m] >= 0] - torch.testing.assert_close(valid, scalar[m].to(valid.dtype)) - - -@pytest.mark.parametrize("ndim", [3, 4]) -def test_batched_rotary_matches_scalar_per_row(ndim): - """apply_dspark_rotary_batched (per-row freqs) == scalar applied row by row.""" - torch.manual_seed(0) - G, s, h, rd = 3, 4, 2, 8 - x = torch.randn(G, s, h, rd) if ndim == 4 else torch.randn(G, s, rd) - table = precompute_dspark_freqs_cis(rd, 64) - # Per-row absolute start positions -> per-row freq windows. - starts = [1, 9, 30] - per_row = torch.stack([table[sp : sp + s] for sp in starts], dim=0) # [G, s, rd/2] - got = apply_dspark_rotary_batched(x, per_row) - for i, sp in enumerate(starts): - ref_i = apply_dspark_rotary(x[i : i + 1], table[sp : sp + s]) - torch.testing.assert_close(got[i : i + 1], ref_i) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") -def test_batched_attention_cuda_graph_capture_replay(): - """The batched attention captures + replays and matches eager output. - - Proves the path is free of capture-illegal ops (host syncs, dynamic shapes). - """ - g = _make_batched_inputs(seed=0) - dev = "cuda" - G = g["G"] - start_pos = torch.tensor(g["start_positions"], dtype=torch.long, device=dev) - slots = torch.arange(G, dtype=torch.long, device=dev) - # Static input tensors the graph reads/writes. - x = g["x"].to(dev) - main_x = g["main_x"].to(dev) - cache = g["kv_cache"].to(dev) - kw = {k: (v.to(dev) if torch.is_tensor(v) else v) for k, v in _attn_kwargs(g).items()} - - def run(persist): - return dspark_attention_forward_batched( - x, main_x, start_pos, cache, slots, persist=persist, **kw - ) - - eager = run(persist=False) - - # Warmup (PyTorch CUDA-graph semantics) then capture on a non-persist call so - # the comparison isn't perturbed by the window write-through. - s = torch.cuda.Stream() - s.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(s): - for _ in range(3): - run(persist=False) - torch.cuda.current_stream().wait_stream(s) - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - out = run(persist=False) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(out, eager, rtol=2e-2, atol=2e-2) - - -# -------------------------------------------------------------------------- -# Confidence-scheduled verification: the scoring path must be graph-capturable. -# -# DSpark is a one-engine drafter, so its worker forward runs INSIDE the target's -# CUDA graph. Anything the confidence head adds to that forward has to capture. -# The head itself is a fixed-shape Linear, but the surrounding code is where -# capture breaks: the previous implementation followed it with a `.item()` and a -# `torch.nonzero`, both of which are illegal under capture. -# -------------------------------------------------------------------------- - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") -@pytest.mark.parametrize("with_markov", [False, True]) -def test_confidence_head_is_cuda_graph_capturable(with_markov): - from tensorrt_llm._torch.models.dspark.heads import DSparkConfidenceHead - - B, BLK, HID, RANK = 4, 5, 32, 16 - head = ( - DSparkConfidenceHead( - hidden_size=HID, markov_rank=RANK, with_markov=with_markov, block_size=BLK - ) - .cuda() - .eval() - ) - hid = torch.randn(B, BLK, HID, device="cuda") - prev = torch.randn(B, BLK, RANK, device="cuda") if with_markov else None - out = torch.empty(B, BLK, device="cuda", dtype=torch.float32) - - def run(): - return head.apply_sts(head(hid, prev_embeddings=prev)) - - # Warm up on a side stream, as CUDA graph capture requires. - s = torch.cuda.Stream() - s.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(s): - for _ in range(3): - run() - torch.cuda.current_stream().wait_stream(s) - - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): - out.copy_(run()) - - eager = run() - g.replay() - torch.cuda.synchronize() - assert torch.allclose(out, eager, atol=1e-5), "replay diverged from eager" - assert torch.all((out >= 0.0) & (out <= 1.0)), "calibrated confidence must be a probability" - - # New inputs must flow through on replay -- a graph that captured constants - # would keep returning the old scores and silently freeze the scheduler. - hid.copy_(torch.randn_like(hid)) - fresh = run() - g.replay() - torch.cuda.synchronize() - assert torch.allclose(out, fresh, atol=1e-5), "replay ignored updated inputs" - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") -def test_sts_table_update_is_visible_to_a_captured_graph(): - """STS must be updated in place; rebinding the buffer would be invisible. - - ``nn.Module.__setattr__`` on a registered buffer replaces the tensor object, - changing its data_ptr. A graph captured before that keeps reading the OLD - storage, so calibration would silently not apply. - """ - from tensorrt_llm._torch.models.dspark.heads import DSparkConfidenceHead - - B, BLK, HID = 2, 5, 32 - head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK).cuda().eval() - hid = torch.randn(B, BLK, HID, device="cuda") - out = torch.empty(B, BLK, device="cuda", dtype=torch.float32) - - s = torch.cuda.Stream() - s.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(s): - for _ in range(3): - head.apply_sts(head(hid)) - torch.cuda.current_stream().wait_stream(s) - - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): - out.copy_(head.apply_sts(head(hid))) - - g.replay() - torch.cuda.synchronize() - before = out.clone() - - storage_before = head.sts_temperatures.data_ptr() - head.load_sts_temperatures(torch.full((BLK,), 4.0)) - assert head.sts_temperatures.data_ptr() == storage_before, ( - "load_sts_temperatures must write in place; rebinding hides it from the graph" - ) - - g.replay() - torch.cuda.synchronize() - assert not torch.allclose(out, before), ( - "the captured graph did not see the new STS temperatures" - ) - assert torch.allclose(out, torch.sigmoid(head(hid) / 4.0), atol=1e-5) - - -def test_apply_sts_accepts_host_logits_from_a_device_head(): - """The planner calibrates on the host; the head lives on the device. - - Confidence is staged to pinned CPU memory and the verify planner calls - ``apply_sts`` on those rows, while ``sts_temperatures`` is a buffer that - moved to CUDA with the draft model. Without the transfer this is a device - mismatch -- and it is invisible until a run supplies a profiled cost table, - because the flat-table fallback returns before calibration is ever reached. - """ - from tensorrt_llm._torch.models.dspark.heads import DSparkConfidenceHead - - BLK = 5 - head = DSparkConfidenceHead(hidden_size=8, block_size=BLK) - if torch.cuda.is_available(): - head = head.cuda() - head.load_sts_temperatures(torch.full((BLK,), 2.0)) - - host_logits = torch.zeros(3, BLK) # CPU, as the planner stages them - probs = head.apply_sts(host_logits) - assert probs.device.type == "cpu" - assert torch.allclose(probs, torch.full((3, BLK), 0.5), atol=1e-6) - - -# -------------------------------------------------------------------------- -# Ragged verification runs inside the same captured graph: DSpark's acceptance -# is part of the target's forward. The scatter that unpacks a packed ragged -# batch into the rectangle acceptance wants must therefore capture -- which it -# only does if ``repeat_interleave`` is told its output size up front. Without -# that, torch reads the cumulative sum back to the host to size the result, -# which is illegal under capture. -# -------------------------------------------------------------------------- - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") -def test_ragged_scatter_is_cuda_graph_capturable(): - """scatter_ragged_to_padded and count_accepted_ragged capture and replay. - - Both run inside the target's captured graph; the scatter is the op that - needs ``repeat_interleave`` told its output size up front, and the accept - count must never gain a host sync either. - """ - from tensorrt_llm._torch.speculative.dspark_ragged import ( - build_qo_indptr, - count_accepted_ragged, - scatter_ragged_to_padded, - ) - - lens = torch.tensor([3, 1, 2], dtype=torch.int32, device="cuda") - indptr = build_qo_indptr(lens) - flat = torch.tensor([10, 11, 12, 20, 30, 31], dtype=torch.int32, device="cuda") - draft = torch.tensor([[1, 2, 3, 9], [4, 9, 9, 9], [5, 6, 9, 9]], device="cuda") - target = torch.tensor([[1, 2, 7, 9], [4, 9, 9, 9], [5, 6, 9, 9]], device="cuda") - out = torch.empty(3, 4, dtype=torch.int32, device="cuda") - accepted = torch.empty(3, dtype=torch.int64, device="cuda") - - def run(): - padded = scatter_ragged_to_padded( - flat, verify_lens=lens, qo_indptr=indptr, max_len=4, pad_value=-1 - ) - counts = count_accepted_ragged( - draft_tokens=draft, target_tokens=target, verify_lens=lens - ) - return padded, counts - - s = torch.cuda.Stream() - s.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(s): - for _ in range(3): - run() - torch.cuda.current_stream().wait_stream(s) - - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): - padded, counts = run() - out.copy_(padded) - accepted.copy_(counts) - - g.replay() - torch.cuda.synchronize() - assert out.tolist() == [[10, 11, 12, -1], [20, -1, -1, -1], [30, 31, -1, -1]] - # req 0 matches 2 then diverges; req 1 matches its single position; req 2 - # matches both. Column 3 is padding for every row and must never count. - assert accepted.tolist() == [2, 1, 2] - - # New packed contents must flow through on replay. - flat.copy_(torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32, device="cuda")) - g.replay() - torch.cuda.synchronize() - assert out.tolist() == [[1, 2, 3, -1], [4, -1, -1, -1], [5, 6, -1, -1]] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py deleted file mode 100644 index 2a973ebd5e5c..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_device_fill_parity.py +++ /dev/null @@ -1,355 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The device fill/row-map path must reproduce the host path token-for-token. - -``fill_bucket_device`` and ``build_row_maps_device`` are the capture-safe -replacements for ``RaggedVerifyLayout.fill_bucket`` and the host row-map -staging in prepare(). Any divergence in the per-row allocation or the -per-token KV correction is an attention-reads-wrong-KV bug that surfaces as -silent quality loss (or an IMA), never as an exception -- so the parity here -is exact equality against the host implementation over the whole feasible -space, not spot checks. -""" - -import pytest -import torch - -from tensorrt_llm._torch.speculative.dspark_device_select import \ - select_windows_device -from tensorrt_llm._torch.speculative.dspark_ragged import ( - RaggedVerifyLayout, build_row_maps_device, fill_bucket_device) -from tensorrt_llm._torch.speculative.dspark_schedule import ( - DSparkScheduleConfig, compute_survival, schedule_verify_lens_topk) - - -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 from host lists: - token o of a request with window v gets correction o - v + 1, so the - gathered 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 -- - the host fit's copy widths depend on this exact split.""" - 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 - - -def test_select_windows_respects_published_split(): - """With pad_len the prologue must land on the host fit's real/pad token - split exactly: real tokens = bucket - n_pad * pad_len, pads = pad_len.""" - gen = torch.Generator().manual_seed(41) - K = 5 - cfg = DSparkScheduleConfig(block_size=K, min_verify_len=1) - max_tok = cfg.resolved_max_verify_len + 1 - for _ in range(30): - n_real = int(torch.randint(1, 33, (1,), generator=gen)) - n_pad = int(torch.randint(0, 5, (1,), generator=gen)) - padded_bs = n_real + n_pad - pad_len = int(torch.randint(1, max_tok + 1, (1,), generator=gen)) - real_target = int( - torch.randint(n_real * 2, n_real * max_tok + 1, (1,), - generator=gen)) - bucket = real_target + n_pad * pad_len - budget = max(0, min(int(torch.randint(0, n_real * K + 1, (1,), - generator=gen)), - real_target - n_real * 2)) - conf = torch.randn(padded_bs + 4, K, generator=gen) * 2 - slot_idx = torch.arange(padded_bs, dtype=torch.long) - res = select_windows_device( - confidence_logits=conf, - slot_idx=slot_idx, - num_real=torch.tensor(n_real), - budget=torch.tensor(budget), - graph_num_tokens=bucket, - cfg=cfg, - pad_len=pad_len, - ) - assert int(res.verify_lens.sum()) == bucket - assert int(res.verify_lens[:n_real].sum()) == real_target - assert (res.verify_lens[n_real:] == pad_len).all() - assert int(res.verify_lens[:n_real].min()) >= cfg.min_verify_len + 1 - assert int(res.verify_lens[:n_real].max()) <= max_tok - - -def test_topk_tensor_budget_matches_int(): - """The 0-d tensor budget path (in-graph capacity knob) must allocate - identically to the host int path, including budget <= 0.""" - gen = torch.Generator().manual_seed(11) - cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) - for _ in range(30): - bs = int(torch.randint(1, 65, (1,), generator=gen)) - survival = compute_survival( - torch.rand(bs, 5, generator=gen)) - for budget in (0, 1, bs, 2 * bs, bs * 4, bs * 10): - host = schedule_verify_lens_topk(survival=survival, - budget=budget, cfg=cfg) - dev = schedule_verify_lens_topk(survival=survival, - budget=torch.tensor(budget), - cfg=cfg) - assert torch.equal(host, dev) - - -def test_planner_lag2_ring_and_budget_split(monkeypatch): - """Device-window mode: the budget reads the snapshot one staging OLDER - (lag-2), and the shape lens spread the budget uniformly, summing to - floor + budget.""" - monkeypatch.setenv("TLLM_DSPARK_DEVICE_WINDOWS", "1") - from tensorrt_llm._torch.speculative.dspark_verify import \ - DSparkVerifyPlanner - cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) - planner = DSparkVerifyPlanner(cfg=cfg) - assert planner.device_windows - - snap_a = torch.full((4, 5), 1.0) - snap_b = torch.full((4, 5), 2.0) - snap_c = torch.full((4, 5), 3.0) - planner.stage_confidence(snap_a) - assert planner._ready_older_snapshot() is None - planner.stage_confidence(snap_b) - older = planner._ready_older_snapshot() - assert older is not None and float(older[0, 0]) == 1.0 - planner.stage_confidence(snap_c) - older = planner._ready_older_snapshot() - assert float(older[0, 0]) == 2.0 - # The current snapshot is the freshest staging, untouched by the ring. - assert float(planner._ready_snapshot()[0, 0]) == 3.0 - - # Budget arithmetic through the frac override (snapshot-independent): - # shape lens sum to floor + budget and stay inside the bounds. - planner._forced_budget_frac = 0.5 - n = 7 - trimmable = cfg.resolved_max_verify_len - cfg.min_verify_len - budget, lens = planner.decide_verify_budget(num_gen_requests=n) - assert budget == int(round(0.5 * n * trimmable)) - assert len(lens) == n - assert sum(lens) == n * cfg.min_verify_len + budget - assert min(lens) >= cfg.min_verify_len - assert max(lens) <= cfg.resolved_max_verify_len - - -def _host_select(rows, budget, bucket, padded_bs, cfg): - """The lag-free host reference: calibrate -> survival -> topk -> fill.""" - survival = compute_survival(torch.sigmoid(rows)) - scheduled = schedule_verify_lens_topk(survival=survival, budget=budget, - cfg=cfg) - token_lens = (scheduled + 1).tolist() - layout = RaggedVerifyLayout.from_verify_lens( - torch.tensor(token_lens, dtype=torch.int32), - graph_num_tokens=bucket, - total_verify_tokens=sum(token_lens)) - return layout.fill_bucket( - max_verify_len=cfg.resolved_max_verify_len + 1, padded_bs=padded_bs) - - -def test_select_windows_device_matches_host_chain(): - """The full prologue equals the host chain run on the same (fresh) rows, - across pad rows, indirection through slots, and stale-stamp fallback.""" - gen = torch.Generator().manual_seed(20260806) - K = 5 - cfg = DSparkScheduleConfig(block_size=K, min_verify_len=1) - max_token_len = cfg.resolved_max_verify_len + 1 - for _ in range(40): - n_real = int(torch.randint(1, 33, (1,), generator=gen)) - padded_bs = n_real + int(torch.randint(0, 5, (1,), generator=gen)) - num_slots = padded_bs + 8 - conf = torch.randn(num_slots, K, generator=gen) * 3 - # Requests sit at shuffled slots; pad rows point at slot 0. - slots = torch.randperm(num_slots, generator=gen)[:n_real] - slot_idx = torch.zeros(padded_bs, dtype=torch.long) - slot_idx[:n_real] = slots - # A third of the rows are stale: their gather must neutralize to - # "verify everything". - stamp = torch.full((num_slots,), 7, dtype=torch.int64) - stale_rows = torch.rand(n_real, generator=gen) < 0.33 - stamp[slots[stale_rows]] = 3 - - budget = int(torch.randint(0, n_real * K, (1,), generator=gen)) - lo = n_real * (cfg.min_verify_len + 1) + min( - budget, n_real * (K - cfg.min_verify_len)) + (padded_bs - n_real) - bucket = min(padded_bs * max_token_len, - lo + int(torch.randint(0, 8, (1,), generator=gen))) - - result = select_windows_device( - confidence_logits=conf, - slot_idx=slot_idx, - num_real=torch.tensor(n_real), - budget=torch.tensor(budget), - graph_num_tokens=bucket, - cfg=cfg, - stamp=stamp, - expected_stamp=torch.tensor(7), - ) - rows = conf[slots].clone() - rows[stale_rows] = 30.0 # NEUTRAL_CONFIDENCE_LOGIT - host = _host_select(rows, budget, bucket, padded_bs, cfg) - assert torch.equal(result.verify_lens, host.verify_lens) - assert torch.equal(result.qo_indptr, host.qo_indptr) - assert int(result.verify_lens.sum()) == bucket - # Row maps must describe exactly the filled lens. - want_req, want_corr = build_row_maps_device( - host.verify_lens, graph_num_tokens=bucket) - assert torch.equal(result.req_idx, want_req) - assert torch.equal(result.kv_correction, want_corr) - - -def test_prologue_kv_pairing_matches_host_staging(): - """The host stages ``kv_lens = num_cached + seq_lens_kv`` where BOTH terms - bake the per-request token window (num_cached = past + tokens, - seq_lens_kv = tokens): kv_lens = past + 2*S. The overlap graph then adds - ``previous_kv_lens_offsets`` (staged as ``new_tokens_lens - S``) during - replay. Rebasing the layout onto the true windows w therefore needs - kv_lens += 2*(w - S) and offsets -= (w - S); any other split leaves the - indexer's slot mapping (derived from kv_lens) pointing at the wrong - K-cache slots -- the silent KV corruption that inflated the device-window - throughput while degrading every completion. The 2x coefficient and the - reference values were established by an A/B tensor diff against a full - host restage with identical windows. - """ - gen = torch.Generator().manual_seed(20260807) - for _ in range(20): - padded_bs = int(torch.randint(2, 40, (1,), generator=gen)) - past_seen = torch.randint(10, 500, (padded_bs,), generator=gen) - new_tokens_lens = torch.randint(1, 7, (padded_bs,), generator=gen) - shape_lens = torch.randint(2, 7, (padded_bs,), generator=gen) - true_lens = torch.randint(2, 7, (padded_bs,), generator=gen) - - # Host staging under the shape split S: the window counts twice. - kv_lens = (past_seen + 2 * shape_lens).to(torch.int32) - offsets = (new_tokens_lens - shape_lens).to(torch.int32) - - # Prologue: 2x on the kv side, 1x on the offsets side. - window_delta = (true_lens - shape_lens).to(torch.int32) - kv_lens += 2 * window_delta - offsets -= window_delta - - # Pre-offset kv_lens must equal host-with-w staging: the indexer's - # slot mapping and the expanded per-token extents derive from it. - assert torch.equal(kv_lens, - (past_seen + 2 * true_lens).to(torch.int32)) - - # In-graph correction at replay. - effective = kv_lens + offsets - - # Reference: host staging directly with the true windows w. - want = (past_seen + 2 * true_lens).to(torch.int32) + ( - new_tokens_lens - true_lens).to(torch.int32) - assert torch.equal(effective, want) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py deleted file mode 100644 index d4270eb836ee..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_draft.py +++ /dev/null @@ -1,118 +0,0 @@ -# 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. -"""Unit tests for the DSpark draft I/O proposal stage.""" - -import torch - -from tensorrt_llm._torch.models.dspark.draft import build_draft_input_ids, dspark_propose -from tensorrt_llm._torch.models.dspark.heads import DSparkConfidenceHead, build_markov_head - -VOCAB, HID, RANK, B, BLK = 257, 32, 16, 2, 5 -NOISE_ID = 199 - - -def test_build_draft_input_ids(): - bonus = torch.tensor([7, 9]) - ids = build_draft_input_ids(bonus, block_size=BLK, noise_token_id=NOISE_ID) - assert ids.shape == (B, BLK) - assert torch.equal(ids[:, 0], bonus) - assert torch.all(ids[:, 1:] == NOISE_ID) - - -def test_dspark_propose_full_block_no_confidence(): - torch.manual_seed(0) - markov = build_markov_head( - markov_head_type="rnn", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID - ).eval() - base = torch.randn(B, BLK, VOCAB) - bonus = torch.randint(0, VOCAB, (B,)) - hid = torch.randn(B, BLK, HID) - with torch.no_grad(): - tokens, confidence = dspark_propose( - base, - bonus_token_ids=bonus, - block_hidden=hid, - markov_head=markov, - confidence_head=None, - block_size=BLK, - ) - assert tokens.shape == (B, BLK) - assert confidence is None - # Tokens match the markov head's own greedy block sampling. - ref_tokens, _ = markov.sample_block_tokens( - base, first_prev_token_ids=bonus, hidden_states=hid, temperature=0.0 - ) - assert torch.equal(tokens, ref_tokens) - - -def _propose_with_confidence(conf, markov, *, return_confidence): - base = torch.randn(1, BLK, VOCAB) - bonus = torch.randint(0, VOCAB, (1,)) - hid = torch.ones(1, BLK, HID) - with torch.no_grad(): - return dspark_propose( - base, - bonus_token_ids=bonus, - block_hidden=hid, - markov_head=markov, - confidence_head=conf, - block_size=BLK, - return_confidence=return_confidence, - ) - - -def test_dspark_propose_scores_without_shortening_the_block(): - """The block is always proposed in full; confidence only scores it, and - only when asked (return_confidence is the run-constant opt-in).""" - torch.manual_seed(1) - markov = build_markov_head( - markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID - ).eval() - conf = DSparkConfidenceHead(hidden_size=HID, block_size=BLK).eval() - # The confidence proj is bias-free, so drive the logit via a constant weight - # against a constant hidden: logit = weight_val * HID per position. - with torch.no_grad(): - conf.proj.weight.fill_(-5.0 / HID) # sigmoid ~ 0.0067, i.e. hopeless - - tokens, confidence = _propose_with_confidence(conf, markov, return_confidence=True) - assert tokens.shape == (1, BLK) - assert confidence.shape == (1, BLK) - # Low confidence must NOT shorten the proposal -- that decision belongs to - # the verification scheduler, not the drafter. - assert torch.all(confidence < 0.0) - assert torch.all(conf.apply_sts(confidence) < 0.5) - - # Opt-in: with the flag off the head is not consulted at all. - _, confidence = _propose_with_confidence(conf, markov, return_confidence=False) - assert confidence is None - - -def test_dspark_propose_is_free_of_host_syncs(): - """No ``.item()``/``nonzero`` on this path: it runs inside the target's graph.""" - import inspect - import io - import tokenize - - from tensorrt_llm._torch.models.dspark import draft as draft_mod - - src = inspect.getsource(draft_mod.dspark_propose) - # Strip comments and string literals so the check reads the code, not the - # prose describing it (the implementation comments name these very calls). - code = "".join( - tok.string if tok.type not in (tokenize.COMMENT, tokenize.STRING) else " " - for tok in tokenize.generate_tokens(io.StringIO(src).readline) - ) - for banned in (".item(", "nonzero", "range("): - assert banned not in code, f"dspark_propose must stay capture-safe: found {banned!r}" diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py deleted file mode 100644 index c8de13d02ff9..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_heads.py +++ /dev/null @@ -1,158 +0,0 @@ -# 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. -"""Unit tests for the DSpark draft-network heads (hardware-agnostic, CPU).""" - -import pytest -import torch - -from tensorrt_llm._torch.models.dspark.heads import ( - DSparkConfidenceHead, - RNNHead, - VanillaMarkov, - build_markov_head, -) - -VOCAB, RANK, HID, B, BLK = 257, 16, 32, 3, 5 - - -@pytest.mark.parametrize("head_type", ["vanilla", "gated", "rnn"]) -def test_markov_block_sampling_shapes_and_determinism(head_type): - torch.manual_seed(0) - head = build_markov_head( - markov_head_type=head_type, vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID - ).eval() - base = torch.randn(B, BLK, VOCAB) - first = torch.randint(0, VOCAB, (B,)) - hid = torch.randn(B, BLK, HID) - with torch.no_grad(): - tok, logits = head.sample_block_tokens( - base, first_prev_token_ids=first, hidden_states=hid, temperature=0.0 - ) - tok2, _ = head.sample_block_tokens( - base, first_prev_token_ids=first, hidden_states=hid, temperature=0.0 - ) - assert tok.shape == (B, BLK) - assert logits.shape == (B, BLK, VOCAB) - # Greedy is deterministic. - assert torch.equal(tok, tok2) - # Each sampled token is the argmax of its (bias-corrected) step logits. - assert torch.equal(tok, logits.argmax(dim=-1)) - - -def test_markov_bias_is_additive_low_rank(): - # bias = W2(W1[token]); the corrected first-step logits == base + bias. - torch.manual_seed(1) - head = VanillaMarkov(vocab_size=VOCAB, markov_rank=RANK).eval() - base = torch.randn(B, BLK, VOCAB) - first = torch.randint(0, VOCAB, (B,)) - with torch.no_grad(): - _, corrected = head.sample_block_tokens( - base, first_prev_token_ids=first, hidden_states=None, temperature=0.0 - ) - expected0 = base[:, 0] + head.markov_w2(head.markov_w1(first)) - assert torch.allclose(corrected[:, 0], expected0, atol=1e-5) - - -def test_rnn_state_carries_across_positions(): - torch.manual_seed(2) - head = RNNHead(vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID).eval() - initial_state = torch.zeros(1, RANK) - prev_embedding = head.get_prev_embeddings(torch.zeros(1, dtype=torch.long)) - prefix_hidden = torch.randn(1, HID) - current_hidden = torch.randn(1, HID) - - with torch.no_grad(): - state_a, _ = head._rnn_step(initial_state, prev_embedding, prefix_hidden) - state_b, _ = head._rnn_step(initial_state, prev_embedding, -prefix_hidden) - _, bias_a = head._rnn_step(state_a, prev_embedding, current_hidden) - _, bias_b = head._rnn_step(state_b, prev_embedding, current_hidden) - - assert not torch.allclose(state_a, state_b) - assert not torch.allclose(bias_a, bias_b) - - -def test_build_markov_head_rank_zero_returns_none(): - assert ( - build_markov_head( - markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=0, hidden_size=HID - ) - is None - ) - - -def test_confidence_head_emits_raw_logits(): - head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) - conf = head(torch.randn(B, BLK, HID)) - assert conf.shape == (B, BLK) - assert conf.dtype == torch.float32 - - -def test_apply_sts_uses_per_position_temperatures(): - head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) - # Default (all-ones) calibration is a plain sigmoid, bounded in [0, 1]. - raw = torch.tensor([[2.0, 0.0, -2.0, 1.0, -1.0]]) - assert torch.allclose(head.apply_sts(raw), torch.sigmoid(raw), atol=1e-6) - assert torch.all((head.apply_sts(raw) >= 0.0) & (head.apply_sts(raw) <= 1.0)) - - temps = torch.tensor([1.0, 2.0, 4.0, 1.0, 1.0]) - head.load_sts_temperatures(temps) - raw = torch.full((1, BLK), 4.0) - got = head.apply_sts(raw) - # A larger temperature pulls the probability toward 0.5. - assert got[0, 0] > got[0, 1] > got[0, 2] - assert torch.allclose(got, torch.sigmoid(raw / temps), atol=1e-6) - - -def test_load_sts_temperatures_updates_in_place(): - """Rebinding the buffer would be invisible to an already-captured graph; - and it is calibration, not a checkpoint weight, so it must stay out of - state_dict and refuse malformed tables.""" - head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) - assert "sts_temperatures" not in head.state_dict() - with pytest.raises(ValueError, match="one per block position"): - head.load_sts_temperatures(torch.ones(BLK + 1)) - with pytest.raises(ValueError, match="strictly positive"): - head.load_sts_temperatures(torch.zeros(BLK)) - before = head.sts_temperatures - head.load_sts_temperatures(torch.full((BLK,), 2.0)) - assert head.sts_temperatures is before - assert torch.allclose(before, torch.full((BLK,), 2.0)) - - -def test_confidence_head_load_weights_rejects_a_dropped_bias(): - """A checkpoint bias silently ignored would shift every confidence score.""" - head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) - good = {"proj.weight": torch.randn(1, HID)} - head.load_weights([good]) - assert torch.allclose(head.proj.weight, good["proj.weight"]) - - with pytest.raises(ValueError, match="bias=True"): - head.load_weights([{**good, "proj.bias": torch.zeros(1)}]) - - # Constructed with bias=True, the same checkpoint bias is accepted. - biased = DSparkConfidenceHead(hidden_size=HID, block_size=BLK, bias=True) - w = {"proj.weight": torch.randn(1, HID), "proj.bias": torch.randn(1)} - biased.load_weights([w]) - assert torch.allclose(biased.proj.bias, w["proj.bias"]) - - -def test_confidence_head_with_markov_concat_dim(): - head = DSparkConfidenceHead(hidden_size=HID, markov_rank=RANK, with_markov=True) - hid = torch.randn(B, BLK, HID) - prev_emb = torch.randn(B, BLK, RANK) - out = head(hid, prev_embeddings=prev_emb) - assert out.shape == (B, BLK) - with pytest.raises(AssertionError): - head(hid) # with_markov requires prev_embeddings diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py deleted file mode 100644 index e2b9f049807b..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_iter_log_samples.py +++ /dev/null @@ -1,159 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Reading cost-table samples out of production iteration logs. - -The in-process sweep measures an engine the profiler built itself, and when -that engine's configuration drifts from the deployment the table describes a -machine nobody is running -- a swept (bs=256, M=1536) cell read 150.2 ms -against 118.8 ms measured live at the same shape, with nothing in the -fingerprint to record the difference. ``print_iter_log`` already emits the row -count, the token total and the host step time for every step of the real -deployment, so a table can be fitted from traffic that actually ran. - -What must hold: the token total identifies the verify length exactly (no -clustering, no guessing), steps that cannot be filed honestly are dropped -rather than mislabelled, and the first step of each executor instance -- whose -host time contains the wait for the first request, measured at 62 s live -- is -never treated as a measurement. -""" - -import importlib.util -import pathlib -import sys -import tempfile - -_SPEC = importlib.util.spec_from_file_location( - "dspark_sps_profiler", - pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / - "dspark_sps_profiler.py") -_PROFILER = importlib.util.module_from_spec(_SPEC) -sys.modules[_SPEC.name] = _PROFILER # @dataclass needs the module registered -_SPEC.loader.exec_module(_PROFILER) - -samples_from_iter_log = _PROFILER.samples_from_iter_log -resolve_padded_shape = _PROFILER.resolve_padded_shape - -# A deployment's captured ladder, abbreviated. Padding rounds a step's rows up -# to the smallest entry at or above them, which is what turns a padded token -# total into exactly one verify length. -LADDER = [1, 2, 4, 8, 16, 32, 64, 128, 192, 256] - - -def _line(iteration, bs, step_ms, gen_tokens, rank="0"): - return (f"{rank}: [08/04/2026-19:08:13] [TRT-LLM] [I] [_torch][RANK 0] " - f"iter = {iteration}, global_rank = 0, rank = 0, " - f"num_scheduled_requests = {bs}, kv_cache_util = 0.125, " - f"currank_total_requests = 1/8, host_step_time = {step_ms}ms, " - f"prev_device_step_time = {step_ms}ms, " - f"timestamp = 2026-08-04 19:08:13, " - f"states = {{'num_ctx_requests': 0, 'num_ctx_tokens': 0, " - f"'num_generation_tokens': {gen_tokens}, " - f"'cached_kv_tokens': 63821}}") - - -def _write(lines, name="gen.log"): - # tempfile rather than pytest's tmp_path: these tests also run under the - # fixture-less in-container runner used for the GPU images. - path = pathlib.Path(tempfile.mkdtemp()) / name - path.write_text("\n".join(lines) + "\n") - return str(path) - - -def test_verify_length_comes_from_the_padded_token_total(): - """The label is read out of the log, never inferred from step times. - - Steps the ladder cannot explain are skipped, not mislabelled: filing a - step under a shape it did not run is the mislabelling that once shipped - eight fictional cells with near-zero residuals. - """ - log = _write([ - _line(1, 1, 62098.0, 6), # executor's first step: dropped - _line(2, 244, 79.2, 768), # 244 rows padded to 256 -> rung-2 - _line(3, 242, 117.1, 1536), # padded to 256 -> full block - _line(4, 128, 61.0, 384), # exactly 128 rows -> rung-2 - _line(5, 256, 300.0, 2048), # L = 7, past the block: skipped - _line(6, 300, 400.0, 1800), # more rows than the ladder holds: skipped - ]) - samples = samples_from_iter_log([log], max_draft_len=5, - padded_batch_sizes=LADDER) - got = sorted((s.batch_size, s.verify_len, s.step_time_ms) - for s in samples) - assert got == [(128, 2, 61.0), (256, 2, 79.2), (256, 5, 117.1)] - # The table is indexed by the padded token total, which each sample exposes. - assert sorted(s.total_verify_tokens for s in samples) == [384, 768, 1536] - - -def test_the_padded_width_is_what_gets_recorded(): - """A 244-row step cost what 256 rows cost; the table is indexed by that. - - The ladder is also what removes the ambiguity: 768 tokens over 128 rows - is L=5 at 128, L=3 at 192, or L=2 at 256 -- without the ladder all three - are admissible and the step is dropped; with it, padding rounds 128 up to - 128 and the answer is L=5 (the first version of this collector kept 312 - of 5682 production steps for lack of this). And a step whose token total - the padded width cannot explain (_get_padded_batch can bail) is dropped: - nothing in the log says padding declined, so that is the only honest move. - """ - assert resolve_padded_shape(num_rows=244, num_generation_tokens=768, - max_draft_len=5, max_batch_size=256, - padded_batch_sizes=LADDER) == (256, 2) - assert resolve_padded_shape(num_rows=128, num_generation_tokens=768, - max_draft_len=5, max_batch_size=256) is None - assert resolve_padded_shape(num_rows=128, num_generation_tokens=768, - max_draft_len=5, max_batch_size=256, - padded_batch_sizes=LADDER) == (128, 5) - assert resolve_padded_shape(num_rows=200, num_generation_tokens=600, - max_draft_len=5, max_batch_size=256, - padded_batch_sizes=LADDER) is None - - -def test_the_first_step_of_each_executor_is_dropped(): - """A build spins up more than one executor, and each numbers from 1. - - The first step of each carries the wait for its first request inside - host_step_time (62 s on a live server), which would otherwise enter the - fit as a cell measurement. - """ - log = _write([ - _line(1, 1, 200.0, 6), # KV-estimation executor - _line(2, 256, 79.0, 768), - _line(1, 1, 62098.0, 6), # real executor starts over - _line(2, 256, 79.4, 768), - ]) - times = [s.step_time_ms for s in samples_from_iter_log([log], max_draft_len=5)] - assert times == [79.0, 79.4] - - -def test_only_the_requested_rank_is_read(): - """Every rank logs; counting all of them would multiply every cell. - - Outlier steps are bounded out on the same pass. - """ - log = _write([ - _line(2, 256, 79.2, 768, rank="0"), - _line(2, 256, 79.9, 768, rank="3"), # another rank's copy: dropped - _line(3, 256, 5000.0, 768), # outlier past max_step_ms - ]) - times = [s.step_time_ms for s in - samples_from_iter_log([log], max_draft_len=5, max_step_ms=1000.0)] - assert times == [79.2] - - -def test_several_logs_are_pooled(): - """One arm gives one rung; the table needs the ladder, so pool the runs.""" - a = _write([_line(1, 1, 100.0, 6), _line(2, 256, 117.1, 1536)], - name="a.log") - c = _write([_line(1, 1, 100.0, 6), _line(2, 256, 79.2, 768)], - name="c.log") - got = sorted(s.verify_len for s in - samples_from_iter_log([a, c], max_draft_len=5)) - assert got == [2, 5] - - -def test_a_log_without_the_states_dict_yields_nothing(): - """enable_iter_perf_stats off: fail loudly upstream, not with a bad table.""" - path = pathlib.Path(tempfile.mkdtemp()) / "bare.log" - path.write_text( - "0: [TRT-LLM] iter = 2, num_scheduled_requests = 256, " - "host_step_time = 79.2ms\n") - assert samples_from_iter_log([str(path)], max_draft_len=5) == [] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model.py new file mode 100644 index 000000000000..7b87d5375a2e --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_model.py @@ -0,0 +1,1036 @@ +# 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. +"""Unit tests for the DSpark draft model: captured-context attention primitives, +the batched CUDA-graph-safe attention path, the draft I/O proposal stage, and +the draft-network heads (markov + confidence). CPU except per-test GPU gates. +""" + +import types +from unittest.mock import Mock + +import pytest +import torch +import torch.nn.functional as F + +import tensorrt_llm._torch.models.modeling_dspark as modeling_dspark +from tensorrt_llm._torch.models.dspark.attention import ( + apply_dspark_rotary, + apply_dspark_rotary_batched, + dspark_attention_forward, + dspark_attention_forward_batched, + dspark_sparse_attn, + get_dspark_topk_idxs, + get_dspark_topk_idxs_batched, + precompute_dspark_freqs_cis, +) +from tensorrt_llm._torch.models.dspark.draft import build_draft_input_ids, dspark_propose +from tensorrt_llm._torch.models.dspark.heads import ( + DSparkConfidenceHead, + RNNHead, + VanillaMarkov, + build_markov_head, +) +from tensorrt_llm._torch.models.modeling_dspark import DSparkDraftModel + +VOCAB, RANK, HID, BLK = 257, 16, 32, 5 +DRAFT_B, HEADS_B = 2, 3 +NOISE_ID = 199 + +# --------------------------------------------------------------------------- +# Captured-context attention primitives (CPU) +# --------------------------------------------------------------------------- + + +def test_rope_table_is_cached_once_per_device(): + model = types.SimpleNamespace( + _attn_params={"rope_head_dim": 16}, + _freqs_cap=64, + _rope_theta=10000.0, + _freqs_table_cache={}, + ) + + first = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) + second = DSparkDraftModel._dspark_freqs_table(model, torch.device("cpu")) + + assert first.data_ptr() == second.data_ptr() + assert len(model._freqs_table_cache) == 1 + positions = torch.tensor([1, 17, 63]) + expected = precompute_dspark_freqs_cis(16, 64, rope_theta=10000.0) + torch.testing.assert_close(first[positions], expected[positions]) + + +def test_dspark_block_uses_stage_id_as_attention_layer_idx(monkeypatch): + captured = {} + + def fake_decoder_layer_init( + self, + model_config, + layer_idx, + aux_stream_dict, + attention_layer_idx=None, + mapping_with_cp=None, + disable_post_moe_fusion=False, + ): + torch.nn.Module.__init__(self) + self.model_config = model_config + self.config = model_config.pretrained_config + self.layer_idx = layer_idx + captured.update( + layer_idx=layer_idx, + attention_layer_idx=attention_layer_idx, + aux_stream_dict=aux_stream_dict, + mapping_with_cp=mapping_with_cp, + disable_post_moe_fusion=disable_post_moe_fusion, + ) + + monkeypatch.setattr( + modeling_dspark.DeepseekV4DecoderLayer, + "__init__", + fake_decoder_layer_init, + ) + model_config = types.SimpleNamespace( + pretrained_config=types.SimpleNamespace(vocab_size=128, hc_mult=2), + spec_config=None, + ) + + block = modeling_dspark.DSparkBlock( + model_config, + layer_idx=10, + aux_stream_dict={}, + stage_id=1, + num_stages=3, + num_capture_layers=0, + ) + + assert block.layer_idx == captured["layer_idx"] == 10 + assert captured["attention_layer_idx"] == block.stage_id == 1 + assert captured["disable_post_moe_fusion"] is True + + +@pytest.mark.parametrize("enable_fused_hc", [True, False]) +def test_forward_stage_honors_enable_fused_hc(monkeypatch, enable_fused_hc): + """The draft stage must use the inherited fused-HC rollback setting.""" + torch.manual_seed(71) + num_requests, block_size, hc_mult, hidden_size = 1, 2, 2, 3 + h = torch.randn(num_requests, block_size, hc_mult, hidden_size) + attention_input = torch.randn(num_requests, block_size, hidden_size) + attention_output = torch.randn_like(attention_input) + mid_residual = torch.randn_like(h) + attention_post_mix = torch.randn(num_requests, block_size, hc_mult, 1) + attention_comb_mix = torch.randn(num_requests, block_size, hc_mult, hc_mult) + ffn_post_mix = torch.randn_like(attention_post_mix) + ffn_comb_mix = torch.randn_like(attention_comb_mix) + raw_ffn_input = torch.randn_like(attention_input) + normed_ffn_input = torch.randn_like(attention_input) + moe_output = torch.randn(num_requests * block_size, hidden_size) + final_h = torch.randn_like(h) + events = [] + + def record(name, result): + def call(*args, **kwargs): + events.append(name) + return result + + return call + + monkeypatch.setattr( + modeling_dspark, + "dspark_attention_forward", + Mock(return_value=attention_output), + ) + + hc_attn = types.SimpleNamespace( + pre_mapping=Mock(return_value=(attention_post_mix, attention_comb_mix, attention_input)), + post_mapping=Mock(side_effect=record("attention_post", mid_residual)), + ) + hc_ffn = types.SimpleNamespace( + fused_hc=Mock( + side_effect=record( + "fused", + (mid_residual, ffn_post_mix, ffn_comb_mix, normed_ffn_input), + ) + ), + pre_mapping=Mock( + side_effect=record("ffn_pre", (ffn_post_mix, ffn_comb_mix, raw_ffn_input)) + ), + post_mapping=Mock(side_effect=record("ffn_post", final_h)), + ) + post_attention_layernorm = Mock(side_effect=record("ffn_norm", normed_ffn_input)) + post_attention_layernorm.weight = torch.ones(hidden_size) + post_attention_layernorm.variance_epsilon = 1e-6 + stage = types.SimpleNamespace( + enable_fused_hc=enable_fused_hc, + hc_attn=hc_attn, + hc_ffn=hc_ffn, + input_layernorm=Mock(side_effect=lambda tensor: tensor), + post_attention_layernorm=post_attention_layernorm, + mlp=Mock(return_value=moe_output), + _dspark_attn={}, + ) + model = types.SimpleNamespace( + use_real_mla=False, + _attn_params={"window_size": 2, "head_dim": 1}, + model_config=types.SimpleNamespace( + mapping=types.SimpleNamespace(enable_attention_dp=False, tp_size=8) + ), + ) + + actual = DSparkDraftModel._forward_stage( + model, + stage, + h, + torch.randn(num_requests, hidden_size), + 1, + torch.empty(0), + torch.zeros(num_requests, block_size, dtype=torch.long), + ) + + assert actual is final_h + torch.testing.assert_close( + stage.mlp.call_args.args[0], + normed_ffn_input.reshape(num_requests * block_size, hidden_size), + ) + # Non-attention-DP multi-GPU (enable_attention_dp=False, tp_size>1): the draft + # MoE must all-reduce its TP-sharded output, mirroring the target MoE. The + # attention-DP and single-GPU paths keep it disabled. + assert stage.mlp.call_args.kwargs["final_all_reduce_params"].enable_allreduce is True + if enable_fused_hc: + assert events == ["fused", "ffn_post"] + hc_ffn.fused_hc.assert_called_once() + hc_attn.post_mapping.assert_not_called() + hc_ffn.pre_mapping.assert_not_called() + post_attention_layernorm.assert_not_called() + fused_kwargs = hc_ffn.fused_hc.call_args.kwargs + assert fused_kwargs["norm_weight"] is post_attention_layernorm.weight + assert fused_kwargs["norm_eps"] == post_attention_layernorm.variance_epsilon + else: + assert events == ["attention_post", "ffn_pre", "ffn_norm", "ffn_post"] + hc_ffn.fused_hc.assert_not_called() + hc_attn.post_mapping.assert_called_once() + hc_ffn.pre_mapping.assert_called_once_with(mid_residual) + post_attention_layernorm.assert_called_once_with(raw_ffn_input) + + +def _ref_precompute_freqs_cis(dim, seqlen, base): + """DeepSpec precompute_freqs_cis with original_seq_len == 0 (no YaRN).""" + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + t = torch.arange(seqlen) + freqs = torch.outer(t, freqs) + return torch.polar(torch.ones_like(freqs), freqs) + + +def _ref_apply_rotary_emb(x, freqs_cis, inverse=False): + """DeepSpec apply_rotary_emb (returns a fresh tensor instead of in-place).""" + xc = torch.view_as_complex(x.float().unflatten(-1, (-1, 2))) + if inverse: + freqs_cis = freqs_cis.conj() + if xc.ndim == 3: + fc = freqs_cis.view(1, xc.size(1), xc.size(-1)) + else: + fc = freqs_cis.view(1, xc.size(1), 1, xc.size(-1)) + return torch.view_as_real(xc * fc).flatten(-2).to(x.dtype) + + +def _loop_reference(q, kv, attn_sink, topk_idxs, scale): + """Obvious, slow per-(b,m,h) reference of the exact kernel math.""" + b, m, h, d = q.shape + out = torch.zeros(b, m, h, d, dtype=torch.float32) + qf, kvf, sink = q.float(), kv.float(), attn_sink.float() + for bi in range(b): + for mi in range(m): + idxs = topk_idxs[bi, mi].tolist() + for hi in range(h): + scores, vecs = [], [] + for j in idxs: + if j < 0: + continue + k = kvf[bi, j] + scores.append(torch.dot(qf[bi, mi, hi], k) * scale) + vecs.append(k) + if not scores: + continue + s = torch.stack(scores) + smax = s.max() + p = torch.exp(s - smax) + denom = p.sum() + torch.exp(sink[hi] - smax) + num = (p.unsqueeze(-1) * torch.stack(vecs)).sum(0) + out[bi, mi, hi] = num / denom + return out + + +def test_sparse_attn_matches_loop_reference(): + torch.manual_seed(0) + b, m, h, d, n, topk = 2, 5, 3, 16, 40, 12 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, n, d) + attn_sink = torch.randn(h) + idx = torch.stack( + [torch.stack([torch.randperm(n)[:topk] for _ in range(m)]) for _ in range(b)] + ).int() + scale = d**-0.5 + got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() + ref = _loop_reference(q, kv, attn_sink, idx, scale) + torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) + + +def test_sparse_attn_no_sink_matches_sdpa(): + """With a -inf sink and all-valid contiguous indices, the primitive must + equal standard scaled-dot-product attention over the gathered KV.""" + torch.manual_seed(0) + b, m, h, d, topk = 2, 4, 5, 16, 9 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, topk, d) # n == topk, attend to all + attn_sink = torch.full((h,), float("-inf")) + idx = torch.arange(topk).view(1, 1, -1).expand(b, m, topk).int() + scale = d**-0.5 + got = dspark_sparse_attn(q, kv, attn_sink, idx, scale).float() + + # SDPA: q [b,h,m,d], k/v [b,h,topk,d] (broadcast the shared KV over heads). + qh = q.permute(0, 2, 1, 3) + kvh = kv.unsqueeze(1).expand(b, h, topk, d) + ref = F.scaled_dot_product_attention(qh, kvh, kvh, scale=scale).permute(0, 2, 1, 3) + torch.testing.assert_close(got, ref, rtol=1e-4, atol=1e-4) + + +def test_sparse_attn_masked_indices_excluded(): + """An index of -1 must be excluded exactly (equiv. to dropping that column).""" + torch.manual_seed(0) + b, m, h, d = 1, 1, 2, 16 + q = torch.randn(b, m, h, d) + kv = torch.randn(b, 5, d) + sink = torch.full((h,), float("-inf")) + scale = d**-0.5 + full = torch.tensor([[[0, 1, 2, 3]]]).int() + masked = torch.tensor([[[0, 1, 2, -1]]]).int() + drop3 = torch.tensor([[[0, 1, 2]]]).int() + got_masked = dspark_sparse_attn(q, kv, sink, masked, scale) + got_drop = dspark_sparse_attn(q, kv, sink, drop3, scale) + torch.testing.assert_close(got_masked, got_drop, rtol=1e-5, atol=1e-5) + # And masking genuinely changes the result vs attending to position 3. + got_full = dspark_sparse_attn(q, kv, sink, full, scale) + assert not torch.allclose(got_full, got_masked, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("start_pos,window,block", [(3, 128, 5), (10, 4, 5)]) +def test_get_dspark_topk_idxs_matches_reference(start_pos, window, block): + bsz = 3 + got = get_dspark_topk_idxs(window, bsz, block, start_pos) + # Reference formula (DeepSpec get_dspark_topk_idxs). + ctx = torch.arange(min(window, start_pos + 1)) + blk = window + torch.arange(block) + ref_row = torch.cat([ctx, blk]).int() + assert got.shape == (bsz, block, ref_row.numel()) + for bi in range(bsz): + for mi in range(block): + torch.testing.assert_close(got[bi, mi], ref_row) + + +@pytest.mark.parametrize("rope_head_dim,seqlen", [(64, 16), (64, 1), (128, 8)]) +def test_precompute_freqs_cis_matches_reference(rope_head_dim, seqlen): + got = precompute_dspark_freqs_cis(rope_head_dim, seqlen, rope_theta=10000.0) + ref = _ref_precompute_freqs_cis(rope_head_dim, seqlen, 10000.0) + torch.testing.assert_close(got, ref) + + +@pytest.mark.parametrize("ndim", [3, 4]) +def test_apply_rotary_matches_reference(ndim): + torch.manual_seed(0) + b, s, h, rd = 2, 5, 4, 64 + x = torch.randn(b, s, h, rd) if ndim == 4 else torch.randn(b, s, rd) + fc = precompute_dspark_freqs_cis(rd, s) + got = apply_dspark_rotary(x, fc) + ref = _ref_apply_rotary_emb(x, fc) + torch.testing.assert_close(got, ref) + # De-rotation (inverse) must undo the forward rotation (property test). + roundtrip = apply_dspark_rotary(got, fc, inverse=True) + torch.testing.assert_close(roundtrip, x, rtol=1e-5, atol=1e-5) + + +def _make_attn_inputs(seed=0): + """Small synthetic DSpark attention inputs/weights (CPU bf16).""" + torch.manual_seed(seed) + dim, n_heads, head_dim, rd = 12, 4, 8, 4 + q_lora, o_lora, n_groups = 6, 5, 2 + window, block, start_pos = 8, 3, 5 + b = 2 + g = dict( + dim=dim, + n_heads=n_heads, + head_dim=head_dim, + rope_head_dim=rd, + q_lora=q_lora, + o_lora=o_lora, + n_groups=n_groups, + window=window, + block=block, + start_pos=start_pos, + b=b, + eps=1e-6, + softmax_scale=head_dim**-0.5, + ) + bf = torch.bfloat16 + g["x"] = torch.randn(b, block, dim, dtype=bf) + g["main_x"] = torch.randn(b, 1, dim, dtype=bf) + g["kv_cache0"] = torch.randn(b, window, head_dim, dtype=bf) + g["wq_a"] = torch.randn(q_lora, dim, dtype=bf) * 0.1 + g["wq_b"] = torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1 + g["wkv"] = torch.randn(head_dim, dim, dtype=bf) * 0.1 + g["wo_a"] = torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1 + g["wo_b"] = torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1 + g["q_norm"] = torch.ones(q_lora) + g["kv_norm"] = torch.ones(head_dim) + g["attn_sink"] = torch.randn(n_heads) + g["freqs"] = precompute_dspark_freqs_cis(rd, start_pos + 1 + block + 2) + return g + + +def _run(g): + return dspark_attention_forward( + g["x"], + g["main_x"], + g["start_pos"], + g["kv_cache0"], + wq_a=g["wq_a"], + q_norm_w=g["q_norm"], + wq_b=g["wq_b"], + wkv=g["wkv"], + kv_norm_w=g["kv_norm"], + wo_a=g["wo_a"], + wo_b=g["wo_b"], + attn_sink=g["attn_sink"], + n_heads=g["n_heads"], + head_dim=g["head_dim"], + rope_head_dim=g["rope_head_dim"], + n_groups=g["n_groups"], + o_lora_rank=g["o_lora"], + window_size=g["window"], + eps=g["eps"], + softmax_scale=g["softmax_scale"], + freqs_cis=g["freqs"], + ) + + +def test_attention_forward_shape_and_determinism(): + g = _make_attn_inputs() + before = g["kv_cache0"].clone() + o = _run(g) + assert tuple(o.shape) == (g["b"], g["block"], g["dim"]) + assert torch.isfinite(o.float()).all() + torch.testing.assert_close(o, _run(g)) # deterministic + # The rolling window write must be functional (cache cloned, not mutated). + torch.testing.assert_close(g["kv_cache0"], before) + + +# --------------------------------------------------------------------------- +# Batched draft-attention path (CUDA-graph safety) +# --------------------------------------------------------------------------- + + +def _make_batched_inputs(seed=0, start_positions=(1, 3, 20)): + """Per-request DSpark attention inputs/weights (shared weights, distinct + start_pos and pre-seeded windows) for batched-vs-scalar comparison.""" + torch.manual_seed(seed) + dim, n_heads, head_dim, rd = 12, 4, 8, 4 + q_lora, o_lora, n_groups = 6, 5, 2 + window, block = 8, 3 + G = len(start_positions) + bf = torch.bfloat16 + # A single fixed RoPE table covering every request's positions (both paths + # index/gather the same values, so freqs are identical across paths). + maxlen = max(start_positions) + 1 + block + 4 + g = dict( + dim=dim, + n_heads=n_heads, + head_dim=head_dim, + rope_head_dim=rd, + q_lora=q_lora, + o_lora=o_lora, + n_groups=n_groups, + window=window, + block=block, + G=G, + start_positions=list(start_positions), + eps=1e-6, + softmax_scale=head_dim**-0.5, + x=torch.randn(G, block, dim, dtype=bf), + main_x=torch.randn(G, 1, dim, dtype=bf), + # Distinct, non-zero seeded window per request to exercise context reads. + kv_cache=torch.randn(G, window, head_dim, dtype=bf), + wq_a=torch.randn(q_lora, dim, dtype=bf) * 0.1, + wq_b=torch.randn(n_heads * head_dim, q_lora, dtype=bf) * 0.1, + wkv=torch.randn(head_dim, dim, dtype=bf) * 0.1, + wo_a=torch.randn(n_groups * o_lora, n_heads * head_dim // n_groups, dtype=bf) * 0.1, + wo_b=torch.randn(dim, n_groups * o_lora, dtype=bf) * 0.1, + q_norm=torch.ones(q_lora), + kv_norm=torch.ones(head_dim), + attn_sink=torch.randn(n_heads), + freqs=precompute_dspark_freqs_cis(rd, maxlen), + ) + return g + + +def _attn_kwargs(g): + return dict( + wq_a=g["wq_a"], + q_norm_w=g["q_norm"], + wq_b=g["wq_b"], + wkv=g["wkv"], + kv_norm_w=g["kv_norm"], + wo_a=g["wo_a"], + wo_b=g["wo_b"], + attn_sink=g["attn_sink"], + n_heads=g["n_heads"], + head_dim=g["head_dim"], + rope_head_dim=g["rope_head_dim"], + n_groups=g["n_groups"], + o_lora_rank=g["o_lora"], + window_size=g["window"], + eps=g["eps"], + softmax_scale=g["softmax_scale"], + freqs_cis=g["freqs"], + ) + + +def _scalar_reference(g): + """Run the validated scalar attention once per request and stack the outputs.""" + outs = [] + for i, sp in enumerate(g["start_positions"]): + out_i = dspark_attention_forward( + g["x"][i : i + 1], + g["main_x"][i : i + 1], + int(sp), + g["kv_cache"][i : i + 1].clone(), + **_attn_kwargs(g), + ) + outs.append(out_i) + return torch.cat(outs, dim=0) + + +def _batched(g, persist=False): + G = g["G"] + start_pos = torch.tensor(g["start_positions"], dtype=torch.long) + slots = torch.arange(G, dtype=torch.long) + return dspark_attention_forward_batched( + g["x"], + g["main_x"], + start_pos, + g["kv_cache"].clone(), + slots, + persist=persist, + **_attn_kwargs(g), + ) + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_batched_attention_matches_scalar_per_request(seed): + """Batched attention == per-request scalar attention at distinct start_pos.""" + g = _make_batched_inputs(seed=seed) + ref = _scalar_reference(g) + got = _batched(g) + assert tuple(got.shape) == (g["G"], g["block"], g["dim"]) + torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2) + + +def test_batched_attention_persist_writes_through_window(): + """persist gates the window write: False mutates nothing, True writes + main_kv into exactly the start_pos%window row of each request.""" + g = _make_batched_inputs(seed=3) + G, win = g["G"], g["window"] + start_pos = torch.tensor(g["start_positions"], dtype=torch.long) + slots = torch.arange(G, dtype=torch.long) + cache = g["kv_cache"].clone() + before = cache.clone() + + dspark_attention_forward_batched( + g["x"], g["main_x"], start_pos, cache, slots, persist=False, **_attn_kwargs(g) + ) + # persist=False must not mutate the caller's window (functional). + torch.testing.assert_close(cache, before) + + dspark_attention_forward_batched( + g["x"], g["main_x"], start_pos, cache, slots, persist=True, **_attn_kwargs(g) + ) + # Exactly the start_pos%win row of each request changed. + for i, sp in enumerate(g["start_positions"]): + changed = (cache[i] != before[i]).any(dim=-1) + expected = torch.zeros(win, dtype=torch.bool) + expected[sp % win] = True + assert torch.equal(changed, expected), f"req {i}: wrong window row written" + + +@pytest.mark.parametrize("start_positions", [(1, 3, 20), (5, 5, 5), (2, 7, 200)]) +def test_batched_topk_matches_scalar(start_positions): + """Fixed-size masked batched topk == scalar topk per request (valid slots).""" + window, block = 8, 3 + start_pos = torch.tensor(start_positions, dtype=torch.long) + batched = get_dspark_topk_idxs_batched(window, block, start_pos) + assert tuple(batched.shape) == (len(start_positions), block, window + block) + for i, sp in enumerate(start_positions): + scalar = get_dspark_topk_idxs(window, 1, block, int(sp))[0] # [block, topk_i] + # The batched row drops the -1-masked context slots to recover the scalar + # (variable-width) index list; both must then be identical. + for m in range(block): + valid = batched[i, m][batched[i, m] >= 0] + torch.testing.assert_close(valid, scalar[m].to(valid.dtype)) + + +@pytest.mark.parametrize("ndim", [3, 4]) +def test_batched_rotary_matches_scalar_per_row(ndim): + """apply_dspark_rotary_batched (per-row freqs) == scalar applied row by row.""" + torch.manual_seed(0) + G, s, h, rd = 3, 4, 2, 8 + x = torch.randn(G, s, h, rd) if ndim == 4 else torch.randn(G, s, rd) + table = precompute_dspark_freqs_cis(rd, 64) + # Per-row absolute start positions -> per-row freq windows. + starts = [1, 9, 30] + per_row = torch.stack([table[sp : sp + s] for sp in starts], dim=0) # [G, s, rd/2] + got = apply_dspark_rotary_batched(x, per_row) + for i, sp in enumerate(starts): + ref_i = apply_dspark_rotary(x[i : i + 1], table[sp : sp + s]) + torch.testing.assert_close(got[i : i + 1], ref_i) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +def test_batched_attention_cuda_graph_capture_replay(): + """The batched attention captures + replays and matches eager output, + proving the path is free of capture-illegal ops.""" + g = _make_batched_inputs(seed=0) + dev = "cuda" + G = g["G"] + start_pos = torch.tensor(g["start_positions"], dtype=torch.long, device=dev) + slots = torch.arange(G, dtype=torch.long, device=dev) + # Static input tensors the graph reads/writes. + x = g["x"].to(dev) + main_x = g["main_x"].to(dev) + cache = g["kv_cache"].to(dev) + kw = {k: (v.to(dev) if torch.is_tensor(v) else v) for k, v in _attn_kwargs(g).items()} + + def run(persist): + return dspark_attention_forward_batched( + x, main_x, start_pos, cache, slots, persist=persist, **kw + ) + + eager = run(persist=False) + + # Warmup (PyTorch CUDA-graph semantics) then capture on a non-persist call so + # the comparison isn't perturbed by the window write-through. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run(persist=False) + torch.cuda.current_stream().wait_stream(s) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = run(persist=False) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(out, eager, rtol=2e-2, atol=2e-2) + + +# Confidence-scheduled verification: the scoring path runs inside the target's +# captured graph, so it must itself be graph-capturable. + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +@pytest.mark.parametrize("with_markov", [False, True]) +def test_confidence_head_is_cuda_graph_capturable(with_markov): + B, BLK, HID, RANK = 4, 5, 32, 16 + head = ( + DSparkConfidenceHead( + hidden_size=HID, markov_rank=RANK, with_markov=with_markov, block_size=BLK + ) + .cuda() + .eval() + ) + hid = torch.randn(B, BLK, HID, device="cuda") + prev = torch.randn(B, BLK, RANK, device="cuda") if with_markov else None + out = torch.empty(B, BLK, device="cuda", dtype=torch.float32) + + def run(): + return head.apply_sts(head(hid, prev_embeddings=prev)) + + # Warm up on a side stream, as CUDA graph capture requires. + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out.copy_(run()) + + eager = run() + g.replay() + torch.cuda.synchronize() + assert torch.allclose(out, eager, atol=1e-5), "replay diverged from eager" + assert torch.all((out >= 0.0) & (out <= 1.0)), "calibrated confidence must be a probability" + + # New inputs must flow through on replay -- a graph that captured constants + # would keep returning the old scores and silently freeze the scheduler. + hid.copy_(torch.randn_like(hid)) + fresh = run() + g.replay() + torch.cuda.synchronize() + assert torch.allclose(out, fresh, atol=1e-5), "replay ignored updated inputs" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +def test_sts_table_update_is_visible_to_a_captured_graph(): + """STS must be updated in place; rebinding the buffer changes its data_ptr + and a graph captured earlier would keep reading the old storage.""" + B, BLK, HID = 2, 5, 32 + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK).cuda().eval() + hid = torch.randn(B, BLK, HID, device="cuda") + out = torch.empty(B, BLK, device="cuda", dtype=torch.float32) + + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + head.apply_sts(head(hid)) + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + out.copy_(head.apply_sts(head(hid))) + + g.replay() + torch.cuda.synchronize() + before = out.clone() + + storage_before = head.sts_temperatures.data_ptr() + head.load_sts_temperatures(torch.full((BLK,), 4.0)) + assert head.sts_temperatures.data_ptr() == storage_before, ( + "load_sts_temperatures must write in place; rebinding hides it from the graph" + ) + + g.replay() + torch.cuda.synchronize() + assert not torch.allclose(out, before), ( + "the captured graph did not see the new STS temperatures" + ) + assert torch.allclose(out, torch.sigmoid(head(hid) / 4.0), atol=1e-5) + + +def test_apply_sts_accepts_host_logits_from_a_device_head(): + """apply_sts must handle host logits while sts_temperatures lives on the + device: the planner calibrates on CPU-staged rows from a CUDA head.""" + BLK = 5 + head = DSparkConfidenceHead(hidden_size=8, block_size=BLK) + if torch.cuda.is_available(): + head = head.cuda() + head.load_sts_temperatures(torch.full((BLK,), 2.0)) + + host_logits = torch.zeros(3, BLK) # CPU, as the planner stages them + probs = head.apply_sts(host_logits) + assert probs.device.type == "cpu" + assert torch.allclose(probs, torch.full((3, BLK), 0.5), atol=1e-6) + + +# Ragged verification runs inside the same captured graph, so its scatter and +# accept-count primitives must also capture without host syncs. + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graph capture needs a GPU") +def test_ragged_scatter_is_cuda_graph_capturable(): + """scatter_ragged_to_padded and count_accepted_ragged capture and replay + inside the target's graph without host syncs.""" + from tensorrt_llm._torch.speculative.dspark_ragged import ( + build_qo_indptr, + count_accepted_ragged, + scatter_ragged_to_padded, + ) + + lens = torch.tensor([3, 1, 2], dtype=torch.int32, device="cuda") + indptr = build_qo_indptr(lens) + flat = torch.tensor([10, 11, 12, 20, 30, 31], dtype=torch.int32, device="cuda") + draft = torch.tensor([[1, 2, 3, 9], [4, 9, 9, 9], [5, 6, 9, 9]], device="cuda") + target = torch.tensor([[1, 2, 7, 9], [4, 9, 9, 9], [5, 6, 9, 9]], device="cuda") + out = torch.empty(3, 4, dtype=torch.int32, device="cuda") + accepted = torch.empty(3, dtype=torch.int64, device="cuda") + + def run(): + padded = scatter_ragged_to_padded( + flat, verify_lens=lens, qo_indptr=indptr, max_len=4, pad_value=-1 + ) + counts = count_accepted_ragged( + draft_tokens=draft, target_tokens=target, verify_lens=lens + ) + return padded, counts + + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + for _ in range(3): + run() + torch.cuda.current_stream().wait_stream(s) + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + padded, counts = run() + out.copy_(padded) + accepted.copy_(counts) + + g.replay() + torch.cuda.synchronize() + assert out.tolist() == [[10, 11, 12, -1], [20, -1, -1, -1], [30, 31, -1, -1]] + # req 0 matches 2 then diverges; req 1 matches its single position; req 2 + # matches both. Column 3 is padding for every row and must never count. + assert accepted.tolist() == [2, 1, 2] + + # New packed contents must flow through on replay. + flat.copy_(torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32, device="cuda")) + g.replay() + torch.cuda.synchronize() + assert out.tolist() == [[1, 2, 3, -1], [4, -1, -1, -1], [5, 6, -1, -1]] + + +# --------------------------------------------------------------------------- +# Draft I/O proposal stage +# --------------------------------------------------------------------------- + + +def test_build_draft_input_ids(): + bonus = torch.tensor([7, 9]) + ids = build_draft_input_ids(bonus, block_size=BLK, noise_token_id=NOISE_ID) + assert ids.shape == (DRAFT_B, BLK) + assert torch.equal(ids[:, 0], bonus) + assert torch.all(ids[:, 1:] == NOISE_ID) + + +def test_dspark_propose_full_block_no_confidence(): + torch.manual_seed(0) + markov = build_markov_head( + markov_head_type="rnn", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID + ).eval() + base = torch.randn(DRAFT_B, BLK, VOCAB) + bonus = torch.randint(0, VOCAB, (DRAFT_B,)) + hid = torch.randn(DRAFT_B, BLK, HID) + with torch.no_grad(): + tokens, confidence = dspark_propose( + base, + bonus_token_ids=bonus, + block_hidden=hid, + markov_head=markov, + confidence_head=None, + block_size=BLK, + ) + assert tokens.shape == (DRAFT_B, BLK) + assert confidence is None + # Tokens match the markov head's own greedy block sampling. + ref_tokens, _ = markov.sample_block_tokens( + base, first_prev_token_ids=bonus, hidden_states=hid, temperature=0.0 + ) + assert torch.equal(tokens, ref_tokens) + + +def _propose_with_confidence(conf, markov, *, return_confidence): + base = torch.randn(1, BLK, VOCAB) + bonus = torch.randint(0, VOCAB, (1,)) + hid = torch.ones(1, BLK, HID) + with torch.no_grad(): + return dspark_propose( + base, + bonus_token_ids=bonus, + block_hidden=hid, + markov_head=markov, + confidence_head=conf, + block_size=BLK, + return_confidence=return_confidence, + ) + + +def test_dspark_propose_scores_without_shortening_the_block(): + """The block is always proposed in full; confidence only scores it, and + only when asked (return_confidence is the run-constant opt-in).""" + torch.manual_seed(1) + markov = build_markov_head( + markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID + ).eval() + conf = DSparkConfidenceHead(hidden_size=HID, block_size=BLK).eval() + # The confidence proj is bias-free, so drive the logit via a constant weight + # against a constant hidden: logit = weight_val * HID per position. + with torch.no_grad(): + conf.proj.weight.fill_(-5.0 / HID) # sigmoid ~ 0.0067, i.e. hopeless + + tokens, confidence = _propose_with_confidence(conf, markov, return_confidence=True) + assert tokens.shape == (1, BLK) + assert confidence.shape == (1, BLK) + # Low confidence must NOT shorten the proposal -- that decision belongs to + # the verification scheduler, not the drafter. + assert torch.all(confidence < 0.0) + assert torch.all(conf.apply_sts(confidence) < 0.5) + + # Opt-in: with the flag off the head is not consulted at all. + _, confidence = _propose_with_confidence(conf, markov, return_confidence=False) + assert confidence is None + + +def test_dspark_propose_is_free_of_host_syncs(): + """No ``.item()``/``nonzero`` on this path: it runs inside the target's graph.""" + import inspect + import io + import tokenize + + from tensorrt_llm._torch.models.dspark import draft as draft_mod + + src = inspect.getsource(draft_mod.dspark_propose) + # Strip comments and string literals so the check reads the code, not the + # prose describing it (the implementation comments name these very calls). + code = "".join( + tok.string if tok.type not in (tokenize.COMMENT, tokenize.STRING) else " " + for tok in tokenize.generate_tokens(io.StringIO(src).readline) + ) + for banned in (".item(", "nonzero", "range("): + assert banned not in code, f"dspark_propose must stay capture-safe: found {banned!r}" + + +# --------------------------------------------------------------------------- +# Draft-network heads (markov + confidence, CPU) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("head_type", ["vanilla", "gated", "rnn"]) +def test_markov_block_sampling_shapes_and_determinism(head_type): + torch.manual_seed(0) + head = build_markov_head( + markov_head_type=head_type, vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID + ).eval() + base = torch.randn(HEADS_B, BLK, VOCAB) + first = torch.randint(0, VOCAB, (HEADS_B,)) + hid = torch.randn(HEADS_B, BLK, HID) + with torch.no_grad(): + tok, logits = head.sample_block_tokens( + base, first_prev_token_ids=first, hidden_states=hid, temperature=0.0 + ) + tok2, _ = head.sample_block_tokens( + base, first_prev_token_ids=first, hidden_states=hid, temperature=0.0 + ) + assert tok.shape == (HEADS_B, BLK) + assert logits.shape == (HEADS_B, BLK, VOCAB) + # Greedy is deterministic. + assert torch.equal(tok, tok2) + # Each sampled token is the argmax of its (bias-corrected) step logits. + assert torch.equal(tok, logits.argmax(dim=-1)) + + +def test_markov_bias_is_additive_low_rank(): + # bias = W2(W1[token]); the corrected first-step logits == base + bias. + torch.manual_seed(1) + head = VanillaMarkov(vocab_size=VOCAB, markov_rank=RANK).eval() + base = torch.randn(HEADS_B, BLK, VOCAB) + first = torch.randint(0, VOCAB, (HEADS_B,)) + with torch.no_grad(): + _, corrected = head.sample_block_tokens( + base, first_prev_token_ids=first, hidden_states=None, temperature=0.0 + ) + expected0 = base[:, 0] + head.markov_w2(head.markov_w1(first)) + assert torch.allclose(corrected[:, 0], expected0, atol=1e-5) + + +def test_rnn_state_carries_across_positions(): + torch.manual_seed(2) + head = RNNHead(vocab_size=VOCAB, markov_rank=RANK, hidden_size=HID).eval() + initial_state = torch.zeros(1, RANK) + prev_embedding = head.get_prev_embeddings(torch.zeros(1, dtype=torch.long)) + prefix_hidden = torch.randn(1, HID) + current_hidden = torch.randn(1, HID) + + with torch.no_grad(): + state_a, _ = head._rnn_step(initial_state, prev_embedding, prefix_hidden) + state_b, _ = head._rnn_step(initial_state, prev_embedding, -prefix_hidden) + _, bias_a = head._rnn_step(state_a, prev_embedding, current_hidden) + _, bias_b = head._rnn_step(state_b, prev_embedding, current_hidden) + + assert not torch.allclose(state_a, state_b) + assert not torch.allclose(bias_a, bias_b) + + +def test_build_markov_head_rank_zero_returns_none(): + assert ( + build_markov_head( + markov_head_type="vanilla", vocab_size=VOCAB, markov_rank=0, hidden_size=HID + ) + is None + ) + + +def test_confidence_head_emits_raw_logits(): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + conf = head(torch.randn(HEADS_B, BLK, HID)) + assert conf.shape == (HEADS_B, BLK) + assert conf.dtype == torch.float32 + + +def test_apply_sts_uses_per_position_temperatures(): + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + # Default (all-ones) calibration is a plain sigmoid, bounded in [0, 1]. + raw = torch.tensor([[2.0, 0.0, -2.0, 1.0, -1.0]]) + assert torch.allclose(head.apply_sts(raw), torch.sigmoid(raw), atol=1e-6) + assert torch.all((head.apply_sts(raw) >= 0.0) & (head.apply_sts(raw) <= 1.0)) + + temps = torch.tensor([1.0, 2.0, 4.0, 1.0, 1.0]) + head.load_sts_temperatures(temps) + raw = torch.full((1, BLK), 4.0) + got = head.apply_sts(raw) + # A larger temperature pulls the probability toward 0.5. + assert got[0, 0] > got[0, 1] > got[0, 2] + assert torch.allclose(got, torch.sigmoid(raw / temps), atol=1e-6) + + +def test_load_sts_temperatures_updates_in_place(): + """The STS table is calibration, not a checkpoint weight: it must update + in place, stay out of state_dict, and refuse malformed tables.""" + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + assert "sts_temperatures" not in head.state_dict() + with pytest.raises(ValueError, match="one per block position"): + head.load_sts_temperatures(torch.ones(BLK + 1)) + with pytest.raises(ValueError, match="strictly positive"): + head.load_sts_temperatures(torch.zeros(BLK)) + before = head.sts_temperatures + head.load_sts_temperatures(torch.full((BLK,), 2.0)) + assert head.sts_temperatures is before + assert torch.allclose(before, torch.full((BLK,), 2.0)) + + +def test_confidence_head_load_weights_rejects_a_dropped_bias(): + """A checkpoint bias silently ignored would shift every confidence score.""" + head = DSparkConfidenceHead(hidden_size=HID, block_size=BLK) + good = {"proj.weight": torch.randn(1, HID)} + head.load_weights([good]) + assert torch.allclose(head.proj.weight, good["proj.weight"]) + + with pytest.raises(ValueError, match="bias=True"): + head.load_weights([{**good, "proj.bias": torch.zeros(1)}]) + + # Constructed with bias=True, the same checkpoint bias is accepted. + biased = DSparkConfidenceHead(hidden_size=HID, block_size=BLK, bias=True) + w = {"proj.weight": torch.randn(1, HID), "proj.bias": torch.randn(1)} + biased.load_weights([w]) + assert torch.allclose(biased.proj.bias, w["proj.bias"]) + + +def test_confidence_head_with_markov_concat_dim(): + head = DSparkConfidenceHead(hidden_size=HID, markov_rank=RANK, with_markov=True) + hid = torch.randn(HEADS_B, BLK, HID) + prev_emb = torch.randn(HEADS_B, BLK, RANK) + out = head(hid, prev_embeddings=prev_emb) + assert out.shape == (HEADS_B, BLK) + with pytest.raises(AssertionError): + head(hid) # with_markov requires prev_embeddings diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py deleted file mode 100644 index 9619cb9eeaa0..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.py +++ /dev/null @@ -1,112 +0,0 @@ -# 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. -"""Overlap-scheduler input layout under ragged (per-request) verification. - -The overlap path builds this iteration's inputs from the *previous* iteration's -device tensors, and does it with a fixed stride: every request is assumed to -contribute ``runtime_tokens_per_gen_step`` tokens. The flat gathers out of -``new_tokens`` / ``next_draft_tokens`` and the ``position_ids`` / -``gather_ids`` / ``seq_lens`` the host appends break silently when that -assumption is dropped: a wrong count shifts every following request, so RoPE -phases and attention windows go wrong without any shape error. - -These tests pin that arithmetic, and -- most importantly -- that a -uniform batch reproduces the strided layout exactly, since that is the path -every non-DSpark model keeps taking. -""" - -import pytest -import torch - -from tensorrt_llm._torch.pyexecutor.llm_request import get_request_tokens_per_gen_step -from tensorrt_llm._torch.speculative.dspark_ragged import ragged_gather_index_lists - - -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 - - -# -------------------------------------------------------------------------- -# get_request_tokens_per_gen_step -# -------------------------------------------------------------------------- - - -_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 - - -# -------------------------------------------------------------------------- -# ragged_gather_index_lists -# -------------------------------------------------------------------------- - - -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_dspark_padded_ceiling.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py deleted file mode 100644 index 5e03d515d855..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_padded_ceiling.py +++ /dev/null @@ -1,70 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""``trim_ratio`` must not report a saving the run did not make -- or a loss. - -Real runs reported a NEGATIVE ``trim_ratio``, which reads as "the ragged path -delivered more tokens than a no-trim run". It did not. ``delivered_tokens`` -books the fitted bucket, which is ``padded_bs * (tier + 1)`` where ``padded_bs`` -is the CUDA-graph-rounded maximum row count *across attention-DP ranks*; -``ceiling_tokens`` booked ``num_gen_requests * (1 + block)`` over this rank's -real, unpadded requests. Two different row bases in one ratio. - -A wrong ratio silently misreports the feature's saving in every summary. -""" - -import pytest - -from tensorrt_llm._torch.speculative.dspark_observability import ( - DSparkRaggedStats, RaggedVerifyMode) - - -def _stats(block_size: int = 5) -> DSparkRaggedStats: - return DSparkRaggedStats(mode=RaggedVerifyMode.COMPACT, - max_draft_len=block_size) - - -def test_padding_rows_do_not_make_the_trim_look_negative(): - """A top-tier step on padded rows saved nothing -- and lost nothing.""" - stats = _stats(block_size=5) - # 3 real requests, graph padded to 4 rows, every window at the full block. - # delivered = 4 * 6 = 24. Against a ceiling of 3 * 6 = 18 that is -0.333. - stats.record_step(num_gen_requests=3, verify_lens=[5, 5, 5], - bucket=4 * 6, padded_bs=4) - summary = stats.summary() - assert summary["trim_ratio"] == pytest.approx(0.0), ( - f"a step that verified the full block on every row trimmed nothing, so " - f"the ratio is 0 -- got {summary['trim_ratio']}, which came from " - f"comparing {summary['delivered_tokens']} padded-row tokens against a " - f"ceiling counted over {3} real requests") - # The padded row count itself must land in the summary: the executor once - # passed padded_bs=None unconditionally, and an always-empty padded_bs_hist - # is indistinguishable from "padding never happened". - assert summary["padded_bs_hist"] == {4: 1} - - -def test_a_real_trim_still_reads_as_a_saving(): - """The fix must not flatten genuine trimming to zero.""" - stats = _stats(block_size=5) - # 4 rows, no padding, every window trimmed to 2 -> bucket 4*3 = 12 against a - # ceiling of 4*6 = 24. - stats.record_step(num_gen_requests=4, verify_lens=[2, 2, 2, 2], - bucket=4 * 3, padded_bs=4) - assert stats.summary()["trim_ratio"] == pytest.approx(0.5) - - -def test_paths_without_a_bucket_keep_the_local_row_base(): - """Only the bucket path rebases; widening the others understates the trim. - - The no-window path and the cap-accept path each book ``delivered`` over the - real rows, so pulling their ceiling up to the padded count would invent a - saving out of padding that those paths never submitted. - """ - no_window = _stats(block_size=5) - no_window.record_step(num_gen_requests=3, verify_lens=None, padded_bs=8) - assert no_window.summary()["trim_ratio"] == pytest.approx(0.0) - - explicit = _stats(block_size=5) - # cap-accept: windows are computed but the full block is submitted anyway. - explicit.record_step(num_gen_requests=3, verify_lens=[2, 3, 1], - bucket=8 * 6, padded_bs=8, delivered=3 * 6) - assert explicit.summary()["trim_ratio"] == pytest.approx(0.0) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_planner.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_planner.py new file mode 100644 index 000000000000..f59f7bad6fb8 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_planner.py @@ -0,0 +1,903 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""DSpark planner/budget/cost-table unit tests (hardware-agnostic, CPU). + +Covers survival, the budget argmax, cost-table interpolation, top-k +allocation, tier derivation, SGLang decision parity, the verify-length pin, +snapshot instruments, and reconciliation guards. +""" + +import itertools + +import numpy as np +import pytest +import torch + +from tensorrt_llm._torch.speculative.dspark_planner import ( + SpsCostTable, budget_argmax_over_uniform_lens, check_table_fingerprint, + compute_verify_token_budget, derive_verify_len_tiers) +from tensorrt_llm._torch.speculative.dspark_schedule import ( + DSparkScheduleConfig, compute_survival, schedule_verify_lens_topk) +from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner + +BLOCK = 7 + + +def _cfg(**kw) -> DSparkScheduleConfig: + return DSparkScheduleConfig(block_size=BLOCK, **kw) + + +# -------------------------------------------------------------------------- +# survival +# -------------------------------------------------------------------------- + + +def test_survival_is_cumulative_product(): + conf = torch.tensor([[0.9, 0.8, 0.5, 1.0, 1.0, 1.0, 1.0]]) + surv = compute_survival(conf) + assert torch.allclose(surv[0, :3], torch.tensor([0.9, 0.72, 0.36]), atol=1e-6) + + +# -------------------------------------------------------------------------- +# budget argmax +# -------------------------------------------------------------------------- + + +def _brute_force_budget(survival, num_gen, table, min_len=1, max_len=None): + """Independent O(N) reimplementation used as the oracle for the argmax: + expected emitted tokens per millisecond over submitted-token cost.""" + bs, blk = survival.shape + cap = min(max_len or blk, blk) + cand = np.sort(survival[:, min_len:cap].reshape(-1))[::-1] + base = num_gen + float(survival[:, :min_len].sum()) + floor_tokens = bs * (min_len + 1) + best_n, best = 0, -np.inf + for n in range(cand.size + 1): + theta = (base + float(cand[:n].sum())) / table.step_time(floor_tokens + n, num_gen) + if theta > best: + best, best_n = theta, n + return best_n + + +def test_budget_matches_brute_force_on_staircase(): + """Theta is not unimodal on a staircase; a greedy first-descent scan is wrong.""" + rng = np.random.default_rng(7) + # A cost curve with genuine shelves and risers. + table = SpsCostTable( + token_counts=(0, 64, 128, 192, 256, 320), + step_time_ms=(4.0, 4.05, 6.5, 6.55, 9.0, 9.05), + ) + for trial in range(25): + bs = int(rng.integers(1, 40)) + conf = rng.uniform(0.55, 0.99, size=(bs, BLOCK)) + surv = np.cumprod(conf, axis=1) + got = compute_verify_token_budget(survival=surv, num_gen_requests=bs, cost_table=table) + want = _brute_force_budget(surv, bs, table) + assert got == want, f"trial {trial}: bs={bs} got={got} want={want}" + + +def test_flat_cost_table_degenerates_to_verify_all(): + """A flat cost model makes every token free -> spend the whole budget.""" + table = SpsCostTable.flat(2.0) + assert table.is_flat + surv = np.cumprod(np.full((6, BLOCK), 0.9), axis=1) + budget = compute_verify_token_budget(survival=surv, num_gen_requests=6, cost_table=table) + assert budget == 6 * (BLOCK - 1) + # An empty batch has no candidates to buy, whatever the table says. + assert ( + compute_verify_token_budget( + survival=np.zeros((0, BLOCK)), num_gen_requests=0, cost_table=table + ) + == 0 + ) + + +def test_cost_table_interpolates_between_breakpoints(): + """The consumer contract is clamped linear interpolation, not a floor; + a shelf is only flat if the table measured it as two equal breakpoints.""" + table = SpsCostTable(token_counts=(0, 100), step_time_ms=(1.0, 5.0)) + assert table.step_time(99) > table.step_time(1) + assert table.step_time(50) == pytest.approx(3.0) + # A measured shelf stays flat under interpolation. + shelf = SpsCostTable(token_counts=(0, 99, 100), step_time_ms=(1.0, 1.0, 5.0)) + assert shelf.step_time(1) == shelf.step_time(99) + # Above the last breakpoint we clamp rather than extrapolate. + assert table.step_time(10_000) == table.step_time(100) + + +def test_cost_table_rejects_malformed_input(): + with pytest.raises(ValueError, match="strictly increasing"): + SpsCostTable(token_counts=(0, 0), step_time_ms=(1.0, 2.0)) + with pytest.raises(ValueError, match="same length"): + SpsCostTable(token_counts=(0, 1), step_time_ms=(1.0,)) + with pytest.raises(ValueError, match="positive"): + SpsCostTable(token_counts=(0,), step_time_ms=(0.0,)) + + +# -------------------------------------------------------------------------- +# per-request allocation +# -------------------------------------------------------------------------- + + +def test_allocation_is_a_prefix_and_respects_bounds(): + torch.manual_seed(3) + cfg = _cfg(min_verify_len=1) + surv = compute_survival(torch.rand(12, BLOCK)) + lens = schedule_verify_lens_topk(survival=surv, budget=20, cfg=cfg) + assert lens.dtype == torch.int32 + assert torch.all(lens >= cfg.min_verify_len) + assert torch.all(lens <= cfg.resolved_max_verify_len) + # Budget is spent above the floor and never exceeded. + assert int((lens - cfg.min_verify_len).sum()) <= 20 + # Zero budget degenerates to the floor for everyone. + cfg_floor = _cfg(min_verify_len=2) + lens = schedule_verify_lens_topk( + survival=compute_survival(torch.rand(5, BLOCK)), budget=0, cfg=cfg_floor) + assert torch.equal(lens, torch.full((5,), 2, dtype=torch.int32)) + # A huge budget saturates at the cap. + cfg_cap = _cfg(min_verify_len=1, max_verify_len=4) + lens = schedule_verify_lens_topk( + survival=compute_survival(torch.full((5, BLOCK), 0.99)), budget=10**6, cfg=cfg_cap) + assert torch.equal(lens, torch.full((5,), 4, dtype=torch.int32)) + # An empty batch allocates nothing. + assert schedule_verify_lens_topk( + survival=torch.zeros(0, BLOCK), budget=5, cfg=_cfg()).numel() == 0 + + +def test_survival_eps_excludes_hopeless_positions(): + """Even with unlimited budget, sub-eps candidates are never admitted.""" + cfg = _cfg(min_verify_len=1, survival_eps=1e-3) + conf = torch.tensor([[1.0, 1.0, 1e-9, 1.0, 1.0, 1.0, 1.0]]) + lens = schedule_verify_lens_topk(survival=compute_survival(conf), budget=10**6, cfg=cfg) + # positions 0,1 survive; from position 2 on, survival collapses below eps. + assert int(lens[0]) == 2 + + +def test_allocation_is_deterministic_under_tied_confidences(): + """Two ranks with identical input must produce identical verify lengths.""" + cfg = _cfg(min_verify_len=1) + surv = compute_survival(torch.full((9, BLOCK), 0.9)) + first = schedule_verify_lens_topk(survival=surv, budget=13, cfg=cfg) + for _ in range(8): + assert torch.equal(schedule_verify_lens_topk(survival=surv, budget=13, cfg=cfg), first) + + +def test_ties_break_toward_earlier_positions(): + """All candidates tie, so the allocation must stay a balanced prefix front, + not hand one request a deep suffix while another gets nothing.""" + cfg = _cfg(min_verify_len=1) + surv = compute_survival(torch.full((4, BLOCK), 1.0)) + lens = schedule_verify_lens_topk(survival=surv, budget=4, cfg=cfg) + assert int(lens.max()) - int(lens.min()) <= 1 + + +def test_schedule_config_rejects_bad_bounds(): + with pytest.raises(ValueError, match="min_verify_len must be >= 1"): + DSparkScheduleConfig(block_size=BLOCK, min_verify_len=0) + with pytest.raises(ValueError, match="block_size must be >= 1"): + DSparkScheduleConfig(block_size=0) + with pytest.raises(ValueError, match="< min_verify_len"): + DSparkScheduleConfig(block_size=BLOCK, min_verify_len=5, max_verify_len=3) + with pytest.raises(ValueError, match=r"\[bs, K\]"): + compute_survival(torch.rand(BLOCK)) + + +# -------------------------------------------------------------------------- +# budget conservation: allocation must never exceed what the planner sized +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("bs,min_len", list(itertools.product([1, 3, 17], [1, 2]))) +def test_planner_and_allocator_agree_on_the_budget(bs, min_len): + rng = np.random.default_rng(bs * 100 + min_len) + table = SpsCostTable(token_counts=(0, 16, 32, 64, 128), step_time_ms=(3.0, 3.1, 5.0, 5.1, 9.0)) + conf = rng.uniform(0.4, 0.99, size=(bs, BLOCK)) + surv_np = np.cumprod(conf, axis=1) + budget = compute_verify_token_budget( + survival=surv_np, num_gen_requests=bs, cost_table=table, min_verify_len=min_len + ) + lens = schedule_verify_lens_topk( + survival=torch.from_numpy(surv_np).float(), + budget=budget, + cfg=_cfg(min_verify_len=min_len), + ) + assert int((lens - min_len).sum()) <= budget + assert torch.all(lens >= min_len) + + +# -------------------------------------------------------------------------- +# uniform-K projection (what TRT-LLM's graph key can actually express today) +# -------------------------------------------------------------------------- + + +def test_discrete_argmax_beats_round_then_snap_across_a_riser(): + """Optimizing over runnable lengths is not the same as rounding the budget: + snapping down past a cost riser lands on a worse Theta than the search.""" + table = SpsCostTable(token_counts=(0, 8, 24, 40), step_time_ms=(2.0, 2.05, 2.1, 20.0)) + surv = np.cumprod(np.full((8, BLOCK), 0.97), axis=1) + chosen = budget_argmax_over_uniform_lens( + survival=surv, + num_gen_requests=8, + cost_table=table, + allowed_lens=[1, 3, 5, 7], + ) + # 8*(7+1)=64 tokens lands on the 20ms riser; the search must avoid it. + assert chosen < 7 + + def theta(length): + # Mirrors the implementation exactly, including the bonus token every + # request submits alongside its drafts. + tau = 8 + float(surv[:, :length].sum()) + return tau / table.step_time(8 * (length + 1), 8) + + assert theta(chosen) == max(theta(v) for v in (1, 3, 5, 7)) + # With nothing runnable, fall back to the floor. + assert ( + budget_argmax_over_uniform_lens( + survival=surv, num_gen_requests=8, cost_table=table, allowed_lens=[] + ) + == 1 + ) + + +def test_derived_tiers_lose_nothing_versus_continuous_k(): + """The shelf-right-edge property: within a shelf Theta rises monotonically, + so tiers built from right edges must match an exhaustive length search.""" + rng = np.random.default_rng(11) + # Shelves are encoded as breakpoint pairs with equal values: under the + # interpolating consumer, flatness between points is only real when the + # table measured it. + table = SpsCostTable( + token_counts=(0, 11, 12, 29, 30, 47, 48, 89, 90, 159, 160), + step_time_ms=(2.0, 2.0, 2.4, 2.4, 3.9, 3.9, 4.0, 4.0, 7.5, 7.5, 12.0), + ) + for bs in (3, 8, 16, 31): + surv = np.cumprod(rng.uniform(0.5, 0.995, size=(bs, BLOCK)), axis=1) + tiers = derive_verify_len_tiers( + cost_table=table, num_requests=bs, block_size=BLOCK, max_tiers=99 + ) + + def theta(length): + return (bs + float(surv[:, :length].sum())) / table.step_time(bs * (length + 1), bs) + + best_any = max(theta(v) for v in range(1, BLOCK + 1)) + best_tier = max(theta(v) for v in tiers) + assert best_tier == pytest.approx(best_any), f"bs={bs} tiers={tiers}" + + +def test_derived_tiers_respect_the_capture_budget(): + table = SpsCostTable( + token_counts=(0, 5, 9, 14, 20, 27, 35), step_time_ms=(1.0, 2, 3, 4, 5, 6, 7) + ) + tiers = derive_verify_len_tiers(cost_table=table, num_requests=2, block_size=BLOCK, max_tiers=3) + assert len(tiers) <= 3 + # Endpoints are always kept: the floor is always runnable and the full block + # is always a right edge (the last shelf is unbounded). + assert tiers[0] == 1 and tiers[-1] == BLOCK + + +def test_derived_tiers_are_a_function_of_batch_size(): + """Tiers move with bs because total tokens = bs * (length + 1); a tier set + derived once and reused across batch sizes would be wrong for one of them.""" + table = SpsCostTable(token_counts=(0, 40), step_time_ms=(1.0, 5.0)) + assert derive_verify_len_tiers( + cost_table=table, num_requests=8, block_size=BLOCK, max_tiers=99 + ) == [1, 3, BLOCK] + assert derive_verify_len_tiers( + cost_table=table, num_requests=32, block_size=BLOCK, max_tiers=99 + ) == [1, BLOCK] + + +# -------------------------------------------------------------------------- +# host-side planner: relay, fallbacks, cross-rank agreement +# -------------------------------------------------------------------------- + + +def _planner(**kw): + kw.setdefault("cfg", _cfg(min_verify_len=1)) + kw.setdefault( + "cost_table", + SpsCostTable(token_counts=(0, 8, 24, 48), step_time_ms=(2.0, 2.1, 4.0, 9.0)), + ) + kw.setdefault("tiers", [1, 3, BLOCK]) + return DSparkVerifyPlanner(**kw) + + +def test_planner_reads_confidence_by_row_not_by_batch_position(): + """The snapshot is slot-indexed and read one iteration later; ``rows`` is + what re-associates it after joins and departures reshuffle the batch.""" + p = _planner() + # Row 3 is a confident request, row 0 a hopeless one. + buf = torch.full((5, BLOCK), -8.0) + buf[3] = 8.0 + p._host_buffer, p._copy_event, p._snapshot_valid = buf, None, True + + # This step schedules the confident request first, then the hopeless one. + lens = p.decide_verify_lens(num_gen_requests=2, rows=[3, 0]) + assert lens is not None and lens[0] > lens[1], f"row mapping ignored: {lens}" + + +def test_planner_rejects_a_row_list_that_does_not_match_the_batch(): + p = _planner() + p._host_buffer, p._copy_event, p._snapshot_valid = torch.rand(8, BLOCK), None, True + assert p.decide_verify_lens(num_gen_requests=4, rows=[0, 1]) is None + assert p.stats["fallback_short_snapshot"] == 1 + + +def _capture_set(batch_sizes, tiers, max_draft_len=BLOCK): + """Invoke PyTorchModelEngine._get_graphs_to_capture with a stub engine.""" + import types + + from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + + spec_config = types.SimpleNamespace( + enable_confidence_scheduling=True, + verify_len_tiers=list(tiers), + draft_len_schedule=None, + max_concurrency=None, + is_linear_tree=True, + spec_dec_mode=types.SimpleNamespace( + support_dynamic_draft_len=lambda: False, use_one_engine=lambda: False + ), + ) + engine = types.SimpleNamespace( + is_draft_model=False, + spec_config=spec_config, + max_draft_len=max_draft_len, + max_total_draft_tokens=max_draft_len, + original_max_draft_len=max_draft_len, + _dynamic_draft_len_mapping=None, + _cuda_graph_batch_sizes=list(batch_sizes), + ) + return PyTorchModelEngine._get_graphs_to_capture(engine, list(batch_sizes), None) + + +def test_capture_pins_the_draft_length_to_the_top_tier(): + """One graph per batch size, not a cross product with the tier ladder: + nothing varies draft_len at runtime, so lower-tier graphs are unreachable.""" + bss, tiers = [1, 8, 64], [1, 3, BLOCK] + graphs = _capture_set(bss, tiers) + assert set(graphs) == {(bs, BLOCK) for bs in bss} + assert len(graphs) == len(bss), ( + f"expected one graph per batch size, got {len(graphs)}: {graphs}") + + +# -------------------------------------------------------------------------- +# top-k allocation follows confidence +# -------------------------------------------------------------------------- + + +def _block5_cfg() -> DSparkScheduleConfig: + return DSparkScheduleConfig(block_size=5, min_verify_len=1, max_verify_len=5) + + +def test_topk_hands_longer_windows_to_more_confident_requests(): + """The budget must follow survival, not batch position.""" + # survival[r, k] = P(the first k+1 drafted tokens all get accepted) + survival = torch.tensor([ + [0.99, 0.98, 0.97, 0.96, 0.95], # confident + [0.90, 0.70, 0.50, 0.30, 0.10], # middling + [0.20, 0.04, 0.01, 0.00, 0.00], # collapses + ]) + lens = schedule_verify_lens_topk(survival=survival, budget=6, + cfg=_block5_cfg()).tolist() + + assert len(lens) == 3 + assert lens[0] > lens[2], ( + f"confident request got {lens[0]}, collapsing request got {lens[2]}; " + f"the budget must follow survival, not batch position") + assert lens[0] >= lens[1] >= lens[2] + # Every request keeps at least the floor, and the budget is respected. + assert min(lens) >= 1 + assert sum(lens) - len(lens) * 1 <= 6 + + +def test_a_full_budget_degenerates_to_the_uniform_full_window(): + """The no-trim case must stay exactly uniform: ragged must cost nothing.""" + num_reqs, max_len = 4, 5 + survival = torch.full((num_reqs, max_len), 0.99) + budget = num_reqs * (max_len - 1) + lens = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=_block5_cfg()).tolist() + assert lens == [max_len] * num_reqs + + +# -------------------------------------------------------------------------- +# cost-table interpolation: clamped linear on both the token and batch axes +# -------------------------------------------------------------------------- + +# The certified GB300 table's shape, abbreviated to two breakpoints. +TABLE = SpsCostTable(token_counts=(512, 768, 1536), + step_time_ms=(23.85, 35.61, 105.64), + fixed_overhead_ms=25.244, + batch_sizes=(128, 256), + batch_overhead_ms=(11.462, 19.343)) + + +def test_theta_interpolates_between_breakpoints(): + """Exact on breakpoints, linear between them: 1512 tokens must cost ~the + 1536 price, not the 768 one a floor lookup would return.""" + assert TABLE.step_time(768, 256) == pytest.approx(25.244 + 19.343 + 35.61) + assert TABLE.step_time(1536, 256) == pytest.approx(25.244 + 19.343 + 105.64) + got = TABLE.step_time(1512, 252) + theta = 35.61 + (105.64 - 35.61) * (1512 - 768) / (1536 - 768) + alpha = 11.462 + (19.343 - 11.462) * (252 - 128) / (256 - 128) + assert got == pytest.approx(25.244 + alpha + theta) + # The floor price it used to return -- and must never return again. + assert got > 130.0 + + +def test_alpha_interpolates_and_clamps(): + assert TABLE.batch_overhead(128) == pytest.approx(11.462) + assert TABLE.batch_overhead(192) == pytest.approx((11.462 + 19.343) / 2) + assert TABLE.batch_overhead(64) == pytest.approx(11.462) # clamp low + assert TABLE.batch_overhead(512) == pytest.approx(19.343) # clamp high + + +def test_theta_clamps_outside_the_measured_range(): + """Theta clamps to the end values; a table without a batch axis adds no alpha.""" + plain = SpsCostTable(token_counts=(512, 768, 1536), + step_time_ms=(23.85, 35.61, 105.64), + fixed_overhead_ms=25.244) + assert plain.step_time(100, 0) == pytest.approx(25.244 + 23.85) + assert plain.step_time(4096, 0) == pytest.approx(25.244 + 105.64) + + +# -------------------------------------------------------------------------- +# decision-level parity with the SGLang budget planner this was ported from +# -------------------------------------------------------------------------- + +# Reference implementation, copied from SGLang dspark_components/dspark_planner +# (additive-table path) with only cosmetic renames. Do not "fix" this copy: +# its job is to be what SGLang runs. + + +def _sgl_interp_clamped(xs, ys, x: float) -> float: + xs = torch.tensor(xs, dtype=torch.float64) + ys = torch.tensor(ys, dtype=torch.float64) + x_t = torch.tensor(float(x), dtype=torch.float64).clamp_(xs[0], xs[-1]) + hi = torch.bucketize(x_t, xs, right=True).clamp_(1, xs.numel() - 1) + lo = hi - 1 + span = (xs[hi] - xs[lo]).clamp_(min=1e-9) + frac = (x_t - xs[lo]) / span + return float(ys[lo] + frac * (ys[hi] - ys[lo])) + + +def _sgl_additive_step_time(table, num_requests: int, num_budgets: int): + floor = table["bias_seconds"] + _sgl_interp_clamped( + table["bs_probes"], table["alpha_seconds"], float(num_requests)) + m_probes = torch.tensor(table["m_probes"], dtype=torch.float64) + theta_vals = torch.tensor(table["theta_seconds"], dtype=torch.float64) + m = (num_requests + torch.arange(num_budgets, dtype=torch.float64)).clamp_( + min=float(table["m_probes"][0]), max=float(table["m_probes"][-1])) + hi = torch.bucketize(m, m_probes, right=True).clamp_(1, m_probes.numel() - 1) + lo = hi - 1 + span = (m_probes[hi] - m_probes[lo]).clamp_(min=1e-9) + frac = (m - m_probes[lo]) / span + theta_at_m = theta_vals[lo] + frac * (theta_vals[hi] - theta_vals[lo]) + return floor + theta_at_m + + +def _sgl_compute_verify_token_budget(*, history_survival_probs, table, + max_verify_len, survival_eps): + num_requests = history_survival_probs.shape[0] + candidates = history_survival_probs[:, :max_verify_len].flatten() + candidates = candidates[candidates >= survival_eps].to(torch.float64) + candidates_sorted = torch.sort(candidates, descending=True).values + prefix_sum = torch.cumsum(candidates_sorted, dim=0) + tau_star = num_requests + torch.cat( + [torch.zeros(1, dtype=torch.float64), prefix_sum]) + step_time = _sgl_additive_step_time(table, int(num_requests), + int(tau_star.numel())) + theta = tau_star / step_time + return int(torch.argmax(theta)) + + +# One certified-table shape shared by both sides. TensorRT-LLM stores ms, +# SGLang stores seconds; the argmax is scale-invariant but the tables are +# built with the honest factor anyway. + +M_PROBES = (64, 96, 128, 192, 384, 512, 768, 1536) +THETA_MS = (9.68, 11.08, 11.95, 13.95, 22.09, 23.85, 35.61, 105.64) +BIAS_MS = 25.244443 +BS_PROBES = (32, 64, 128, 256) +ALPHA_MS = (0.0, 2.177477, 11.462212, 19.342661) +SGL_BLOCK = 5 + +TRT_TABLE = SpsCostTable(token_counts=M_PROBES, step_time_ms=THETA_MS, + fixed_overhead_ms=BIAS_MS, batch_sizes=BS_PROBES, + batch_overhead_ms=ALPHA_MS) +SGL_TABLE = { + "m_probes": M_PROBES, + "theta_seconds": tuple(v / 1e3 for v in THETA_MS), + "bias_seconds": BIAS_MS / 1e3, + "bs_probes": BS_PROBES, + "alpha_seconds": tuple(v / 1e3 for v in ALPHA_MS), +} + + +def _inversion_free_survival(rng, bs): + """Survivals where every position-0 beats every deeper candidate; the + shift-exact parity theorem holds only under this condition.""" + conf0 = rng.uniform(0.85, 0.99, size=(bs, 1)) + deeper = rng.uniform(0.25, 0.80, size=(bs, SGL_BLOCK - 1)) + return np.cumprod(np.concatenate([conf0, deeper], axis=1), axis=1) + + +@pytest.mark.parametrize("bs", [1, 3, 8, 64, 128, 252, 256]) +def test_budget_argmax_matches_sglang_up_to_the_floor_shift(bs): + """Same survival, same table, eps off, no inversions: shift-exact parity.""" + rng = np.random.default_rng(17 + bs) + for _ in range(5): + surv = _inversion_free_survival(rng, bs) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=SGL_BLOCK, survival_eps=0.0) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + if sgl_n >= bs: + assert trt_m == sgl_n - bs, ( + f"bs={bs}: SGLang admits {sgl_n} candidates " + f"(= floor {bs} + {sgl_n - bs}), TRT-LLM budget {trt_m}") + else: + # SGLang went below the one-draft floor; the closest budget + # TensorRT-LLM can express under min_verify_len=1 is zero. + assert trt_m == 0 + + +def test_live_pooled_survivals_park_on_the_breakpoint_together(): + """The campaign's operating point: both planners stop at exactly M = 768, + and tier alignment then rounds down to the capturable rung-2 budget.""" + bs = 252 + surv = np.tile([0.771, 0.693, 0.502, 0.360, 0.315], (bs, 1)) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=SGL_BLOCK, survival_eps=0.0) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + assert trt_m == 768 - 2 * bs + assert sgl_n - bs == trt_m + tiered = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1, allowed_lens=[1, 2, 5]) + assert tiered == bs * 1 # rung-2: one scheduled position past the floor + + +def test_starving_weak_rows_is_a_known_gap(): + """SGLang can verify NOTHING for a hopeless request while TRT-LLM's + min_verify_len=1 floor cannot, so SGLang's achievable Theta is better.""" + # Half the batch is strong all the way down, half is dead on arrival. + # The batch must be big enough that its token range sits on a genuinely + # rising part of the theta curve, or every candidate is free and the + # comparison proves nothing. + bs = 64 + strong = np.tile([0.99, 0.98, 0.97, 0.96, 0.95], (bs // 2, 1)) + dead = np.tile([0.01, 0.008, 0.006, 0.004, 0.002], (bs // 2, 1)) + surv = np.concatenate([strong, dead], axis=0) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=SGL_BLOCK, survival_eps=0.0) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + + def theta_at(tau, tokens): + return tau / float(TRT_TABLE.step_times(np.asarray([tokens]), bs)[0]) + + cand = np.sort(surv.reshape(-1))[::-1] + sgl_theta = theta_at(bs + cand[:sgl_n].sum(), bs + sgl_n) + deeper = np.sort(surv[:, 1:].reshape(-1))[::-1] + trt_theta = theta_at(bs + surv[:, 0].sum() + deeper[:trt_m].sum(), + 2 * bs + trt_m) + assert sgl_theta > trt_theta + + +def test_eps_filtering_is_a_known_divergence(): + """SGLang drops sub-eps candidates inside the budget; TRT-LLM applies eps + at allocation time, so its budget can count a candidate the allocator refuses.""" + bs = 16 + surv = np.tile([0.9, 0.4, 0.008, 0.004, 0.002], (bs, 1)) + sgl_n = _sgl_compute_verify_token_budget( + history_survival_probs=torch.tensor(surv, dtype=torch.float32), + table=SGL_TABLE, max_verify_len=SGL_BLOCK, survival_eps=0.01) + trt_m = compute_verify_token_budget( + survival=surv, num_gen_requests=bs, cost_table=TRT_TABLE, + min_verify_len=1) + assert sgl_n - bs <= trt_m + + +# -------------------------------------------------------------------------- +# tier-aligned budgets: score with the tau the step actually collects +# -------------------------------------------------------------------------- + + +def _table() -> SpsCostTable: + """A staircase with genuine risers, so trimming is worth something.""" + token_counts = tuple(range(0, 400, 16)) + step_time_ms = tuple(4.0 + 0.6 * (tok // 96) + 0.004 * tok + for tok in token_counts) + return SpsCostTable(token_counts=token_counts, step_time_ms=step_time_ms, + fixed_overhead_ms=1.0) + + +@pytest.mark.parametrize("tiers", [[1, 2, 5], [1, 3, 5], [1, 5]]) +def test_restricted_answer_is_always_realisable(tiers): + """Every returned budget corresponds to a rung the executor captured.""" + rng = np.random.default_rng(20260803) + table = _table() + for _ in range(200): + bs = int(rng.integers(2, 17)) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] + n = compute_verify_token_budget( + survival=survival, num_gen_requests=bs, cost_table=table, + min_verify_len=1, allowed_lens=tiers) + assert n % bs == 0, ( + f"budget {n} for {bs} requests is not n*(t-min) for any tier") + assert (n // bs) + 1 in tiers, ( + f"budget {n} implies tier {(n // bs) + 1}, not in {tiers}") + + +def test_restricted_never_scores_worse_than_the_uniform_choice(): + """Same grid, better numerator: the chosen rung's realised theta wins; + at bs=1 the two scorers must pick the same rung outright.""" + rng = np.random.default_rng(31337) + table = _table() + tiers = [1, 2, 5] + for _ in range(300): + bs = int(rng.integers(1, 17)) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] + + def realised_theta(tier: int) -> float: + budget = bs * (tier - 1) + cand = np.sort(survival[:, 1:].reshape(-1))[::-1] + tau = float(bs) + float(survival[:, :1].sum()) + float( + cand[:budget].sum()) + tokens = np.array([bs * (tier + 1)]) + return tau / float(table.step_times(tokens, bs)[0]) + + n = compute_verify_token_budget( + survival=survival, num_gen_requests=bs, cost_table=table, + min_verify_len=1, allowed_lens=tiers) + chosen = (n // bs) + 1 + uniform = budget_argmax_over_uniform_lens( + survival=survival, num_gen_requests=bs, cost_table=table, + allowed_lens=tiers, min_verify_len=1) + assert realised_theta(chosen) >= realised_theta(uniform) - 1e-12, ( + f"restricted picked tier {chosen} whose realised theta is below " + f"the uniform scorer's tier {uniform}") + if bs == 1: + assert chosen == uniform, ( + "at bs=1 the uniform and top-k allocations are the same set; " + "the two scorers must pick the same rung") + + +# -------------------------------------------------------------------------- +# runtime verify-length pin: queue, agree, adopt +# -------------------------------------------------------------------------- + + +def _pin_planner(tiers=(1, 2, 5)): + table = SpsCostTable(token_counts=(0, 512, 768, 1536), + step_time_ms=(5.0, 68.4, 80.2, 150.5), + fixed_overhead_ms=1.0) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + return DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=list(tiers)) + + +def test_adopting_the_agreed_value_applies_it(): + """Queuing must not change this rank's behaviour before the group agrees; + adopting the agreed value applies it and consumes the pending request.""" + planner = _pin_planner() + assert planner.request_verify_len_pin(2) == 2 + # Queued, not yet in force. + assert planner._forced_verify_len is None + assert planner.pending_verify_len_pin() == 2 + planner.adopt_verify_len_pin(2) + assert planner._forced_verify_len == 2 + # Consumed: the next step must not re-broadcast a pin nobody asked for. + assert planner.pending_verify_len_pin() == -1 + + +def test_a_rank_that_queued_nothing_still_adopts_the_group_value(): + """The endpoint lands on ONE rank; the others learn it from the payload.""" + peer = _pin_planner() + assert peer.pending_verify_len_pin() == -1 + peer.adopt_verify_len_pin(5) + assert peer._forced_verify_len == 5 + + +def test_zero_clears_the_pin(): + """Wire protocol: -1 means nobody asked (leave the pin alone); clearing + needs its own value, 0, because None cannot travel in an int payload.""" + planner = _pin_planner() + planner.adopt_verify_len_pin(2) + planner.adopt_verify_len_pin(-1) + assert planner._forced_verify_len == 2 + assert planner.request_verify_len_pin(None) is None + assert planner.pending_verify_len_pin() == 0 + planner.adopt_verify_len_pin(0) + assert planner._forced_verify_len is None + + +def test_an_uncaptured_length_is_refused_at_the_call_site(): + """Rejected when requested, not on some later step, and never queued.""" + planner = _pin_planner(tiers=(1, 2, 5)) + with pytest.raises(ValueError, match="captured tier ladder"): + planner.request_verify_len_pin(3) + with pytest.raises(ValueError, match="outside"): + planner.request_verify_len_pin(9) + assert planner.pending_verify_len_pin() == -1 + assert planner._forced_verify_len is None + + +def test_pinned_steps_hand_every_request_the_same_window(): + """The whole point: the shape a sweep cell can honestly label.""" + planner = _pin_planner() + bs = 8 + rng = np.random.default_rng(5) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() + planner._gather_rows = lambda **_: torch.tensor(survival, + dtype=torch.float32) + planner.adopt_verify_len_pin(2) + lens = planner.decide_verify_lens(num_gen_requests=bs, + reduce_across_ranks=False) + assert lens == [2] * bs + assert planner.stats["forced_steps"] == 1 + + +# -------------------------------------------------------------------------- +# snapshot instruments: stamp-lag histogram and neutral-row count +# -------------------------------------------------------------------------- + + +def test_neutral_rows_are_counted_and_real_rows_are_not(): + planner = _pin_planner() + selected = torch.zeros((8, 5)) + selected[3] = 30.0 # the neutral fill: sigmoid ~ 1.0 at any temperature + planner._note_snapshot_stats(selected, rows=None) + assert planner.stats["snap_rows"] == 8 + assert planner.stats["snap_neutral_rows"] == 1 + + +def test_lag_histogram_reads_staged_seq_minus_stamp(): + planner = _pin_planner() + planner._host_stamps = torch.tensor([7, 6, 5, 0], dtype=torch.int32) + planner._staged_seq = 7 + selected = torch.zeros((4, 5)) + planner._note_snapshot_stats(selected, rows=[0, 1, 2, 3]) + hist = planner.stats["stamp_lag_hist"] + # Stamp 0 is "never drafted": lag 7 in truth, and it must land visibly at + # the histogram's far end rather than vanish -- these are exactly the rows + # the neutral counter should agree about. + assert hist == {0: 1, 1: 1, 2: 1, 7: 1} + + +def test_large_and_negative_lags_clamp_instead_of_exploding_the_dict(): + planner = _pin_planner() + planner._host_stamps = torch.tensor([0, 500], dtype=torch.int32) + planner._staged_seq = 400 + planner._note_snapshot_stats(torch.zeros((2, 5)), rows=[0, 1]) + assert planner.stats["stamp_lag_hist"] == {8: 1, -2: 1} + + +def test_counts_accumulate_across_steps(): + planner = _pin_planner() + planner._host_stamps = torch.tensor([3, 3], dtype=torch.int32) + planner._staged_seq = 4 + for _ in range(3): + planner._note_snapshot_stats(torch.full((2, 5), 30.0), rows=[0, 1]) + assert planner.stats["snap_rows"] == 6 + assert planner.stats["snap_neutral_rows"] == 6 + assert planner.stats["stamp_lag_hist"] == {1: 6} + + +def test_instrument_failure_is_counted_not_raised(): + planner = _pin_planner() + planner._host_stamps = torch.tensor([1], dtype=torch.int32) + planner._staged_seq = 1 + # Rows that cannot index the stamp buffer must clamp, and anything worse + # must be swallowed into the error counter: instruments never kill steps. + planner._note_snapshot_stats(torch.zeros((1, 5)), rows=[10_000]) + assert planner.stats.get("snap_stats_errors", 0) == 0 + planner._note_snapshot_stats(object(), rows=None) # type: ignore[arg-type] + assert planner.stats["snap_stats_errors"] == 1 + + +def test_decide_path_feeds_the_instruments(): + """The counters must fill from the real decide path, not only direct calls.""" + planner = _pin_planner() + bs = 8 + rng = np.random.default_rng(3) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() + planner._gather_rows = lambda **_: torch.tensor(survival, + dtype=torch.float32) + lens = planner.decide_verify_lens(num_gen_requests=bs, + reduce_across_ranks=False) + assert lens is not None + assert planner.stats["snap_rows"] == bs + assert planner.stats["snap_neutral_rows"] == 0 + # The local-decision instruments must fill on the same path: one rung per + # decision, one window per request, and the two must agree -- the windows + # are the top-k realisation of exactly that rung's budget. + rung_hist = planner.stats["local_rung_hist"] + assert sum(rung_hist.values()) == 1 + (rung,) = rung_hist + assert rung in (1, 2, 5) + len_hist = planner.stats["local_len_hist"] + assert sum(len_hist.values()) == bs + budget = bs * (rung - planner.cfg.min_verify_len) + extra = sum((k - planner.cfg.min_verify_len) * v for k, v in len_hist.items()) + assert extra <= budget + + +# -------------------------------------------------------------------------- +# reconciliation guards: engine fingerprint and price prediction +# -------------------------------------------------------------------------- + +LIVE = {"tp": 8, "ep": 8, "attention_dp": True, "block": 5, + "max_batch_size": 256} + + +def _payload(engine): + return {"token_counts": [512, 1536], "step_time_ms": [60.0, 150.0], + "_meta": {"engine": engine} if engine is not None else {}} + + +def test_mismatched_engine_is_refused(): + """Same cell, different max_batch_size: the wrong table must not load.""" + wrong = dict(LIVE, max_batch_size=64) + with pytest.raises(ValueError, match="different engine configuration"): + check_table_fingerprint(payload=_payload(wrong), live=dict(LIVE)) + + +def test_fingerprints_that_must_load(): + """Every non-contradicted fingerprint loads; only a contradicted fact + refuses (case-insensitive strings, unverifiable keys logged, no-meta warns).""" + check_table_fingerprint(payload=_payload(dict(LIVE)), live=dict(LIVE)) + check_table_fingerprint( + payload=_payload(dict(LIVE, moe_backend="megamoe_cutedsl")), + live=dict(LIVE, moe_backend="MEGAMOE_CUTEDSL")) + check_table_fingerprint( + payload=_payload(dict(LIVE, image="faf2c60935", + geometry="constant_block")), + live=dict(LIVE)) + check_table_fingerprint(payload=_payload(None), live=dict(LIVE)) + + +def test_planner_records_the_price_it_paid(): + """Every budget decision must leave behind the step time it assumed, so a + gap against measured hostStepTimeMS can contradict a wrong table.""" + table = SpsCostTable(token_counts=(0, 512, 768, 1536), + step_time_ms=(5.0, 68.4, 80.2, 150.5), + fixed_overhead_ms=1.0) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + planner = DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=[1, 2, 5]) + + bs = 8 + rng = np.random.default_rng(11) + survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() + + # Only the snapshot is faked; calibration stays the constructor's default + # sigmoid (setting it to None post-construction would bypass the + # `apply_calibration or torch.sigmoid` fallback and crash the decide path). + planner._gather_rows = lambda **_: torch.tensor(survival, + dtype=torch.float32) + + lens = planner.decide_verify_lens(num_gen_requests=bs, + reduce_across_ranks=False) + assert lens is not None + + steps = planner.stats.get("predicted_steps", 0) + total = planner.stats.get("predicted_ms_sum", 0.0) + assert steps == 1 and total > 0.0 + + # The recorded price must be the table's own answer for the tokens the + # decision implies -- anything else and the reconciliation compares two + # unrelated numbers. + budget = sum(lens) - bs * cfg.min_verify_len + tokens = bs * (cfg.min_verify_len + 1) + budget + expected = float(table.step_times(np.asarray([tokens]), bs)[0]) + assert planner.last_predicted_step_ms == pytest.approx(expected) + assert total == pytest.approx(expected) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py index 3c9f5c7fd1ec..15f4885cd679 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.py @@ -12,17 +12,29 @@ # 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. -"""Unit tests for the ragged (per-request) verification layout.""" +"""DSpark ragged (per-request) verification subsystem: packed layout and +bucket selection, host/device fill parity, overlap-scheduler gather layout, +compressor per-request new-token counts, and the ragged CUDA-graph key.""" + +import types import pytest import torch +from tensorrt_llm._torch.attention_backend.sparse.deepseek_v4.deepseek_v4 import \ + DeepseekV4TrtllmAttentionMetadata +from tensorrt_llm._torch.pyexecutor.llm_request import get_request_tokens_per_gen_step +from tensorrt_llm._torch.speculative.dspark_device_select import \ + select_windows_device from tensorrt_llm._torch.speculative.dspark_ragged import ( RaggedVerifyLayout, build_qo_indptr, + build_row_maps_device, choose_ragged_capture_shape, count_accepted_ragged, + fill_bucket_device, fill_padded_rows_onehot, + ragged_gather_index_lists, round_up_to_bucket, row_ids_from_lens, scatter_ragged_to_padded, @@ -219,9 +231,8 @@ def test_validate_rejects_a_stale_indptr(): def test_ragged_beats_uniform_at_equal_budget(): - """Pins the reason ragged exists: heterogeneous confidence must yield - heterogeneous lengths that pack into a valid layout and out-earn the best - uniform split at the same token spend.""" + """Heterogeneous confidence must yield heterogeneous lengths that pack into + a valid layout and out-earn the best uniform split at the same token spend.""" conf = torch.tensor( [ [0.98, 0.97, 0.96, 0.95, 0.94, 0.93, 0.92], # deep, worth verifying @@ -268,13 +279,8 @@ def test_row_ids_map_each_packed_token_to_its_request(): def test_row_ids_requires_the_host_total(): - """``total`` is what keeps the expansion sync-free, so it cannot be optional. - - Without ``output_size`` torch reads the cumulative sum back to the host to - size its output. DSpark's acceptance runs inside the target's captured CUDA - graph, where that sync is illegal -- so the signature forces every caller to - supply a total it already knows. - """ + """``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) @@ -414,15 +420,9 @@ def test_shape_requires_bs_buckets(): # --------------------------------------------------------------------------- -# group-consistent bucket choice (max_verify_len given) -# -# The regression these pin: the chooser used to pick the smallest bucket >= -# needed and leave "can MY rank decompose it into real + pad rows" to a LOCAL -# check in fit_ragged_verify_lens. Pad rows all share one window (single -# shared dummy), so decomposability is a divisibility question -- sharpest -# under a full-block pin where real rows have zero slack -- and a rank that -# failed it declined ALONE while its peers published the bucket. The ADP -# shape gate then took the whole group eager. +# group-consistent bucket choice (max_verify_len given): pad rows share one +# window, so decomposability is a divisibility question every rank must answer +# identically -- a lone decline used to take the whole ADP group eager. # --------------------------------------------------------------------------- MAXLEN_BUCKETS = [16, 24, 32, 40, 48] # padded_bs 8, tiers 1..5 -> 8*(t+1) @@ -440,14 +440,8 @@ def _shape6(num_real, total, peers): def test_pinned_ranks_agree_on_a_bucket_every_rank_can_decompose(): - """The production ADP divergence, miniaturised. - - Both ranks are pinned to the full block (floor == capacity). needed = - 8 + 25 = 33 rounds to 40, which rank (5, 30) cannot decompose (3 pad rows, - (40-30)/3 is not an integer window) while rank (4, 24) can (4 rows at - window 4 exactly) -- the old code split the group here. The chooser must - now walk past 40 and land BOTH ranks on 48. - """ + """Full-block-pinned ranks must land together on a bucket every rank can + decompose (48 here), not split the group at the first fitting bucket (40).""" peers = [(5, 30), (4, 24)] shapes = {_shape6(num_real=r, total=t, peers=peers) for r, t in peers} assert len(shapes) == 1 @@ -455,9 +449,8 @@ def test_pinned_ranks_agree_on_a_bucket_every_rank_can_decompose(): def test_slacky_batch_keeps_the_smallest_bucket(): - """Feasibility filtering must not grow the bucket when real rows have - slack to absorb the remainder -- the ordinary trimmed batch keeps the - exact answer the ungated chooser gave.""" + """Feasibility filtering must not grow the bucket when real rows have slack + to absorb the remainder -- the trimmed batch keeps the ungated answer.""" peers = [(3, 9), (6, 30), (5, 11)] got = _shape6(num_real=3, total=9, peers=peers) ungated = choose_ragged_capture_shape( @@ -470,14 +463,7 @@ def test_slacky_batch_keeps_the_smallest_bucket(): def test_group_declines_together_when_nothing_decomposes(): """When no captured bucket works for every rank, every rank must see the - same ValueError -- a group-consistent decline replays the uniform graph, - where a lone decline used to split the graph keys. - - A full-block-pinned rank can always decompose the FULL bucket (pad rows - at the max window), so total infeasibility requires a grid whose top tier - is missing -- which capture produces for real when max_num_tokens - truncates the ladder at the largest batch size. - """ + same ValueError -- a group-consistent decline replays the uniform graph.""" truncated = [16, 24, 32, 40] # top tier 48 not captured peers = [(5, 30), (4, 24)] # pinned: 40 fails (5,30), 32 fails it too for r, t in peers: @@ -555,12 +541,7 @@ def test_warmup_split_reaches_any_captured_bucket(): def _pad_len(*, bucket, n_real, n_pad, max_verify_len, floor=None): """Mirror of the pad-row window ModelEngine.fit_ragged_verify_lens derives. - - Pad rows take the one-token minimum so real requests keep the slack, and - only grow when the real rows physically cannot absorb the bucket. Returns - None when no integer window leaves the real rows a reachable target -- the - scheduler then declines and the step stays uniform. - """ + Returns None when no integer window leaves the real rows a reachable target.""" real_capacity = n_real * max_verify_len floor = n_real if floor is None else floor if n_pad == 0: @@ -604,15 +585,8 @@ def test_pad_rows_grow_when_the_real_rows_cannot_absorb_the_bucket(bucket, expec def test_pad_length_never_yields_an_infeasible_real_target(): - """Exhaustive sweep of every shape the scheduler can actually produce. - - A wrong pad window is the worst failure mode here: the batch's token total - silently misses every captured bucket and the step drops out of graph - replay into eager. So sweep the real bucket grid (one bucket per tier at - each captured batch size), pick the bucket the way the scheduler does, and - require that the derived split either reconstructs the bucket exactly or - declines outright. - """ + """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): @@ -651,13 +625,575 @@ def test_pad_length_never_yields_an_infeasible_real_target(): def test_extra_peer_stat_fields_are_ignored_here(): - """The payload grows; this function's contract does not. - - peer_stats gained a third element -- the group's answer to whether row - padding will take the cross-rank maximum at all -- which the caller uses - before reaching this function. A two-element unpack here turned that into - a ValueError on the first ragged step. - """ + """peer_stats may grow extra elements consumed by the caller; a rigid + two-element unpack here would raise on the first ragged step.""" two = _shape(num_real=3, total=12, peers=[[3, 12], [4, 16]]) three = _shape(num_real=3, total=12, peers=[[3, 12, 1], [4, 16, 1]]) assert (three.padded_bs, three.bucket) == (two.padded_bs, two.bucket) + + +# --------------------------------------------------------------------------- +# 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 + + +def test_select_windows_respects_published_split(): + """With pad_len the prologue must land on the host fit's real/pad token + split exactly: real tokens = bucket - n_pad * pad_len, pads = pad_len.""" + gen = torch.Generator().manual_seed(41) + K = 5 + cfg = DSparkScheduleConfig(block_size=K, min_verify_len=1) + max_tok = cfg.resolved_max_verify_len + 1 + for _ in range(30): + n_real = int(torch.randint(1, 33, (1,), generator=gen)) + n_pad = int(torch.randint(0, 5, (1,), generator=gen)) + padded_bs = n_real + n_pad + pad_len = int(torch.randint(1, max_tok + 1, (1,), generator=gen)) + real_target = int( + torch.randint(n_real * 2, n_real * max_tok + 1, (1,), + generator=gen)) + bucket = real_target + n_pad * pad_len + budget = max(0, min(int(torch.randint(0, n_real * K + 1, (1,), + generator=gen)), + real_target - n_real * 2)) + conf = torch.randn(padded_bs + 4, K, generator=gen) * 2 + slot_idx = torch.arange(padded_bs, dtype=torch.long) + res = select_windows_device( + confidence_logits=conf, + slot_idx=slot_idx, + num_real=torch.tensor(n_real), + budget=torch.tensor(budget), + graph_num_tokens=bucket, + cfg=cfg, + pad_len=pad_len, + ) + assert int(res.verify_lens.sum()) == bucket + assert int(res.verify_lens[:n_real].sum()) == real_target + assert (res.verify_lens[n_real:] == pad_len).all() + assert int(res.verify_lens[:n_real].min()) >= cfg.min_verify_len + 1 + assert int(res.verify_lens[:n_real].max()) <= max_tok + + +def test_topk_tensor_budget_matches_int(): + """The 0-d tensor budget path (in-graph capacity knob) must allocate + identically to the host int path, including budget <= 0.""" + gen = torch.Generator().manual_seed(11) + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + for _ in range(30): + bs = int(torch.randint(1, 65, (1,), generator=gen)) + survival = compute_survival( + torch.rand(bs, 5, generator=gen)) + for budget in (0, 1, bs, 2 * bs, bs * 4, bs * 10): + host = schedule_verify_lens_topk(survival=survival, + budget=budget, cfg=cfg) + dev = schedule_verify_lens_topk(survival=survival, + budget=torch.tensor(budget), + cfg=cfg) + assert torch.equal(host, dev) + + +def test_planner_lag2_ring_and_budget_split(monkeypatch): + """Device-window mode: the budget reads the snapshot one staging older + (lag-2), and the shape lens spread the budget uniformly to floor + budget.""" + monkeypatch.setenv("TLLM_DSPARK_DEVICE_WINDOWS", "1") + from tensorrt_llm._torch.speculative.dspark_verify import \ + DSparkVerifyPlanner + cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) + planner = DSparkVerifyPlanner(cfg=cfg) + assert planner.device_windows + + snap_a = torch.full((4, 5), 1.0) + snap_b = torch.full((4, 5), 2.0) + snap_c = torch.full((4, 5), 3.0) + planner.stage_confidence(snap_a) + assert planner._ready_older_snapshot() is None + planner.stage_confidence(snap_b) + older = planner._ready_older_snapshot() + assert older is not None and float(older[0, 0]) == 1.0 + planner.stage_confidence(snap_c) + older = planner._ready_older_snapshot() + assert float(older[0, 0]) == 2.0 + # The current snapshot is the freshest staging, untouched by the ring. + assert float(planner._ready_snapshot()[0, 0]) == 3.0 + + # Budget arithmetic through the frac override (snapshot-independent): + # shape lens sum to floor + budget and stay inside the bounds. + planner._forced_budget_frac = 0.5 + n = 7 + trimmable = cfg.resolved_max_verify_len - cfg.min_verify_len + budget, lens = planner.decide_verify_budget(num_gen_requests=n) + assert budget == int(round(0.5 * n * trimmable)) + assert len(lens) == n + assert sum(lens) == n * cfg.min_verify_len + budget + assert min(lens) >= cfg.min_verify_len + assert max(lens) <= cfg.resolved_max_verify_len + + +def _host_select(rows, budget, bucket, padded_bs, cfg): + """The lag-free host reference: calibrate -> survival -> topk -> fill.""" + survival = compute_survival(torch.sigmoid(rows)) + scheduled = schedule_verify_lens_topk(survival=survival, budget=budget, + cfg=cfg) + token_lens = (scheduled + 1).tolist() + layout = RaggedVerifyLayout.from_verify_lens( + torch.tensor(token_lens, dtype=torch.int32), + graph_num_tokens=bucket, + total_verify_tokens=sum(token_lens)) + return layout.fill_bucket( + max_verify_len=cfg.resolved_max_verify_len + 1, padded_bs=padded_bs) + + +def test_select_windows_device_matches_host_chain(): + """The full prologue equals the host chain run on the same (fresh) rows, + across pad rows, indirection through slots, and stale-stamp fallback.""" + gen = torch.Generator().manual_seed(20260806) + K = 5 + cfg = DSparkScheduleConfig(block_size=K, min_verify_len=1) + max_token_len = cfg.resolved_max_verify_len + 1 + for _ in range(40): + n_real = int(torch.randint(1, 33, (1,), generator=gen)) + padded_bs = n_real + int(torch.randint(0, 5, (1,), generator=gen)) + num_slots = padded_bs + 8 + conf = torch.randn(num_slots, K, generator=gen) * 3 + # Requests sit at shuffled slots; pad rows point at slot 0. + slots = torch.randperm(num_slots, generator=gen)[:n_real] + slot_idx = torch.zeros(padded_bs, dtype=torch.long) + slot_idx[:n_real] = slots + # A third of the rows are stale: their gather must neutralize to + # "verify everything". + stamp = torch.full((num_slots,), 7, dtype=torch.int64) + stale_rows = torch.rand(n_real, generator=gen) < 0.33 + stamp[slots[stale_rows]] = 3 + + budget = int(torch.randint(0, n_real * K, (1,), generator=gen)) + lo = n_real * (cfg.min_verify_len + 1) + min( + budget, n_real * (K - cfg.min_verify_len)) + (padded_bs - n_real) + bucket = min(padded_bs * max_token_len, + lo + int(torch.randint(0, 8, (1,), generator=gen))) + + result = select_windows_device( + confidence_logits=conf, + slot_idx=slot_idx, + num_real=torch.tensor(n_real), + budget=torch.tensor(budget), + graph_num_tokens=bucket, + cfg=cfg, + stamp=stamp, + expected_stamp=torch.tensor(7), + ) + rows = conf[slots].clone() + rows[stale_rows] = 30.0 # NEUTRAL_CONFIDENCE_LOGIT + host = _host_select(rows, budget, bucket, padded_bs, cfg) + assert torch.equal(result.verify_lens, host.verify_lens) + assert torch.equal(result.qo_indptr, host.qo_indptr) + assert int(result.verify_lens.sum()) == bucket + # Row maps must describe exactly the filled lens. + want_req, want_corr = build_row_maps_device( + host.verify_lens, graph_num_tokens=bucket) + assert torch.equal(result.req_idx, want_req) + assert torch.equal(result.kv_correction, want_corr) + + +def test_prologue_kv_pairing_matches_host_staging(): + """Rebasing the layout onto the true windows w needs kv_lens += 2*(w - S) + and offsets -= (w - S); any other split corrupts the indexer's slot mapping.""" + gen = torch.Generator().manual_seed(20260807) + for _ in range(20): + padded_bs = int(torch.randint(2, 40, (1,), generator=gen)) + past_seen = torch.randint(10, 500, (padded_bs,), generator=gen) + new_tokens_lens = torch.randint(1, 7, (padded_bs,), generator=gen) + shape_lens = torch.randint(2, 7, (padded_bs,), generator=gen) + true_lens = torch.randint(2, 7, (padded_bs,), generator=gen) + + # Host staging under the shape split S: the window counts twice. + kv_lens = (past_seen + 2 * shape_lens).to(torch.int32) + offsets = (new_tokens_lens - shape_lens).to(torch.int32) + + # Prologue: 2x on the kv side, 1x on the offsets side. + window_delta = (true_lens - shape_lens).to(torch.int32) + kv_lens += 2 * window_delta + offsets -= window_delta + + # Pre-offset kv_lens must equal host-with-w staging: the indexer's + # slot mapping and the expanded per-token extents derive from it. + assert torch.equal(kv_lens, + (past_seen + 2 * true_lens).to(torch.int32)) + + # In-graph correction at replay. + effective = kv_lens + offsets + + # Reference: host staging directly with the true windows w. + want = (past_seen + 2 * true_lens).to(torch.int32) + ( + new_tokens_lens - true_lens).to(torch.int32) + assert torch.equal(effective, want) + + +# -------------------------------------------------------------------------- +# overlap-scheduler input layout: the strided gather assumption and its +# ragged replacement, with the uniform batch reproducing the strided layout. +# -------------------------------------------------------------------------- + + +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]) + + +# -------------------------------------------------------------------------- +# compressor gen counts: _sync_gen_tokens_per_seq must source its per-request +# vector from the DEVICE seq-lens buffer (D2D, not a captured H2D of the host +# split) and use the global max window as the next_n scalar bound. +# -------------------------------------------------------------------------- + + +def _dummy(nc, staged_split, seq_lens_dev, max_draft_tokens=5, cap=16): + d = types.SimpleNamespace() + d.is_ragged_verify = True + d.num_generations = len(staged_split) + d.num_contexts = nc + d.max_draft_tokens = max_draft_tokens + # What the host shape split says (S); the fixed code must NOT read it. + d.ragged_verify_lens = list(staged_split) + d.gen_new_tokens_per_seq_cuda = torch.zeros(cap, dtype=torch.int) + d._seq_lens_cuda = torch.tensor(seq_lens_dev, dtype=torch.int) + return d + + +def _call(d, num_gen_tokens): + return DeepseekV4TrtllmAttentionMetadata._sync_gen_tokens_per_seq( + d, num_gen_tokens) + + +def test_vector_sources_from_seq_lens_not_host_split(): + # Simulate the prologue: host staged S, device rewrote seq lens to w. + S = [3, 3, 3] + w = [2, 4, 3] + d = _dummy(nc=0, staged_split=S, seq_lens_dev=w) + _call(d, sum(w)) + assert d.gen_new_tokens_per_seq.tolist() == w + + +def test_scalar_is_global_bound_not_batch_max(): + S = [2, 2] + d = _dummy(nc=0, staged_split=S, seq_lens_dev=S) + next_n = _call(d, sum(S)) + assert next_n == 1 + d.max_draft_tokens + assert d.num_gen_tokens_per_seq == 1 + d.max_draft_tokens + + +def test_gen_rows_slice_skips_context_rows(): + ctx_lens = [17, 9] + w = [4, 2] + d = _dummy(nc=2, staged_split=w, seq_lens_dev=ctx_lens + w) + _call(d, sum(w)) + assert d.gen_new_tokens_per_seq.tolist() == w + + +def test_uniform_branch_unchanged(): + d = _dummy(nc=0, staged_split=[3, 3], seq_lens_dev=[3, 3]) + d.is_ragged_verify = False + next_n = _call(d, 6) + assert next_n == 3 + assert d.gen_new_tokens_per_seq is None + + +# -------------------------------------------------------------------------- +# ragged CUDA-graph key: the key's token axis must describe the batch the +# graph will actually replay over (fit, publication, and the padding ladder). +# -------------------------------------------------------------------------- + + +class _Runner: + """Just the fields will_pad_to reads.""" + + def __init__(self, supported, cfg_batch_size): + self.supported_batch_sizes = supported + self.config = type("C", (), {"batch_size": cfg_batch_size})() + + +def test_will_pad_to_rejects_row_counts_the_ladder_cannot_reach(): + """The ragged fit must not derive a bucket grid from row counts the + captured padding ladder cannot reach.""" + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ + CUDAGraphRunner + + runner = _Runner([1, 2, 4, 8, 128, 192, 256], cfg_batch_size=256) + will = CUDAGraphRunner.will_pad_to + + assert will(runner, 256, 193), "193 -> 256 is reachable and should be taken" + assert not will(runner, 193, 193), "193 is not a captured size" + assert will(runner, 192, 192), "already at a captured size" + # Padding past the configured batch size is refused by _get_padded_batch, + # so the fit must not assume it. + assert not will(_Runner([256, 512], cfg_batch_size=256), 512, 300) + assert not will(runner, 0, 10) + assert not will(runner, 256, 0) + + +def test_graph_key_uses_the_agreed_bucket_not_a_fresh_sum(): + """Every attention-DP rank must key on the same fitted bucket -- never a + fresh re-sum of the padded batch, which can diverge across ranks.""" + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ + CUDAGraphRunner + + class _R: + spec_config = type("S", (), {"enable_ragged_verify": True})() + agreed_ragged_bucket = None + + def _batch(*verify_lens): + return types.SimpleNamespace(generation_requests=[ + types.SimpleNamespace(py_verify_len=v) for v in verify_lens + ]) + + r = _R() + get = CUDAGraphRunner._ragged_verify_bucket + windowed = _batch(3, 2, 1) + + # Nothing fitted yet -> not ragged, key unchanged from the uniform one. + assert get(r, windowed) is None + + # A fitted bucket is returned verbatim: the batch's own token total is + # never re-summed here, which is what makes every rank agree. + r.agreed_ragged_bucket = 449 + assert get(r, windowed) == 449 + # ... including when that total plainly differs from the fitted value. + assert get(r, _batch(1, 1)) == 449 + + # But a request without a window means the token-major staging chain never + # ran this step, so a ragged key would replay the PREVIOUS ragged step's + # row maps -- the ragged IMA. Decline, whatever the fit published. + assert get(r, _batch(3, None, 1)) is None + + # Cleared between steps so a stale bucket cannot key into the wrong graph. + r.agreed_ragged_bucket = None + assert get(r, windowed) is None + + # Config gate still wins: a non-ragged engine keeps the original key shape. + r.agreed_ragged_bucket = 512 + r.spec_config = type("S", (), {"enable_ragged_verify": False})() + assert get(r, windowed) is None + + +def test_capture_publishes_the_bucket_it_shaped_the_batch_to(): + """Whoever shapes a warmup batch must publish its bucket, or the graph is + keyed as if the batch were uniform and matches no runtime step.""" + from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine + + runner = types.SimpleNamespace(agreed_ragged_bucket=None) + requests = [types.SimpleNamespace(py_verify_len=None) for _ in range(4)] + batch = types.SimpleNamespace(generation_requests=requests) + engine = types.SimpleNamespace(cuda_graph_runner=runner) + + # bs=4, draft_len=5 -> the t=5 bucket is 4 * 6 = 24 tokens. + PyTorchModelEngine._set_warmup_ragged_windows(engine, batch, 24, 5) + + assert sum(1 + r.py_verify_len for r in requests) == 24, ( + "the warmup batch must hit the bucket exactly") + assert runner.agreed_ragged_bucket == 24, ( + "capture shaped the batch to 24 tokens but did not publish it, so the " + "graph would be keyed as if the batch were uniform") + + # A narrower bucket at the same batch size, to prove it is per-entry and + # not set once: capture walks every tier for every batch size. + PyTorchModelEngine._set_warmup_ragged_windows(engine, batch, 8, 5) + assert sum(1 + r.py_verify_len for r in requests) == 8 + assert runner.agreed_ragged_bucket == 8 + + # An empty batch publishes nothing rather than a stale or zero bucket. + runner.agreed_ragged_bucket = None + PyTorchModelEngine._set_warmup_ragged_windows( + engine, types.SimpleNamespace(generation_requests=[]), 24, 5) + assert runner.agreed_ragged_bucket is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py deleted file mode 100644 index 78d882c0ec0f..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_reconcile_guards.py +++ /dev/null @@ -1,99 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The two reconciliation guards: engine fingerprint and price prediction. - -Both exist because this pipeline's characteristic failure is a PREDICTION -recorded as a MEASUREMENT with nothing positioned to contradict it: a cost -table measured on the wrong MoE kernel almost planned MegaMoE steps, the same -nominal cell measured 23% apart across two engine configs, and a planner -argmax'd against a table whose labels described shapes that never ran. The -fingerprint refuses the wrong table at load; the recorded prediction gives -every run a number that reality (hostStepTimeMS) can contradict. -""" - -import numpy as np -import pytest - -from tensorrt_llm._torch.speculative.dspark_planner import ( - SpsCostTable, check_table_fingerprint) -from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig -from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner - -LIVE = {"tp": 8, "ep": 8, "attention_dp": True, "block": 5, - "max_batch_size": 256} - - -def _payload(engine): - return {"token_counts": [512, 1536], "step_time_ms": [60.0, 150.0], - "_meta": {"engine": engine} if engine is not None else {}} - - -def test_mismatched_engine_is_refused(): - """The 23%-apart case: same cell, different max_batch_size.""" - wrong = dict(LIVE, max_batch_size=64) - with pytest.raises(ValueError, match="different engine configuration"): - check_table_fingerprint(payload=_payload(wrong), live=dict(LIVE)) - - -def test_fingerprints_that_must_load(): - """Every non-contradicted fingerprint loads; only a contradicted fact refuses. - - Matching facts load silently; string facts must not fail on spelling case - (MEGAMOE vs MegaMoe); facts the consumer cannot see (image hash, geometry) - are informational and logged, not refused -- refusing on unverifiable keys - would make every fingerprint addition a breaking change; and an old table - with no fingerprint warns rather than breaking existing deployments. - """ - check_table_fingerprint(payload=_payload(dict(LIVE)), live=dict(LIVE)) - check_table_fingerprint( - payload=_payload(dict(LIVE, moe_backend="megamoe_cutedsl")), - live=dict(LIVE, moe_backend="MEGAMOE_CUTEDSL")) - check_table_fingerprint( - payload=_payload(dict(LIVE, image="faf2c60935", - geometry="constant_block")), - live=dict(LIVE)) - check_table_fingerprint(payload=_payload(None), live=dict(LIVE)) - - -def test_planner_records_the_price_it_paid(): - """Every budget decision must leave behind the step time it assumed. - - Without this the run has no number reality can contradict: the argmax - consults the table and the table's opinion evaporates. With it, the - [final] planner block carries predicted_ms_sum / predicted_steps, and a - systematic gap against measured hostStepTimeMS is the one signal that - catches a wrong table, a wrong lookup, and an engine mismatch alike. - """ - table = SpsCostTable(token_counts=(0, 512, 768, 1536), - step_time_ms=(5.0, 68.4, 80.2, 150.5), - fixed_overhead_ms=1.0) - cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) - planner = DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=[1, 2, 5]) - - bs = 8 - rng = np.random.default_rng(11) - survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() - - import torch - # Only the snapshot is faked; calibration stays the constructor's default - # sigmoid (setting it to None post-construction would bypass the - # `apply_calibration or torch.sigmoid` fallback and crash the decide path). - planner._gather_rows = lambda **_: torch.tensor(survival, - dtype=torch.float32) - - lens = planner.decide_verify_lens(num_gen_requests=bs, - reduce_across_ranks=False) - assert lens is not None - - steps = planner.stats.get("predicted_steps", 0) - total = planner.stats.get("predicted_ms_sum", 0.0) - assert steps == 1 and total > 0.0 - - # The recorded price must be the table's own answer for the tokens the - # decision implies -- anything else and the reconciliation compares two - # unrelated numbers. - budget = sum(lens) - bs * cfg.min_verify_len - tokens = bs * (cfg.min_verify_len + 1) + budget - expected = float(table.step_times(np.asarray([tokens]), bs)[0]) - assert planner.last_predicted_step_ms == pytest.approx(expected) - assert total == pytest.approx(expected) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py deleted file mode 100644 index 85517f5c4c30..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_snapshot_instruments.py +++ /dev/null @@ -1,110 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The two snapshot instruments: stamp-lag histogram and neutral-row count. - -They exist because the planner's inputs have never been measured: the argmax -eats a slot-indexed snapshot staged one step earlier, and nothing recorded how -old each gathered row was (relay staleness) or how many rows were still the -+30 neutral fill (unknowns entering as certainties). The live anomaly they -instrument: 445/476 full-batch steps chose the full block while calibrated -pooled survivals favour rung-2 by 36% -- so the per-step snapshot content must -differ from the pooled shards, and these counters say in which way. - -Both are measurement-only. A failure inside them must never kill a step. -""" - -import numpy as np -import pytest -import torch - -from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable -from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig -from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner - - -def _planner(): - table = SpsCostTable(token_counts=(0, 512, 768, 1536), - step_time_ms=(5.0, 68.4, 80.2, 150.5), - fixed_overhead_ms=1.0) - cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) - return DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=[1, 2, 5]) - - -def test_neutral_rows_are_counted_and_real_rows_are_not(): - planner = _planner() - selected = torch.zeros((8, 5)) - selected[3] = 30.0 # the neutral fill: sigmoid ~ 1.0 at any temperature - planner._note_snapshot_stats(selected, rows=None) - assert planner.stats["snap_rows"] == 8 - assert planner.stats["snap_neutral_rows"] == 1 - - -def test_lag_histogram_reads_staged_seq_minus_stamp(): - planner = _planner() - planner._host_stamps = torch.tensor([7, 6, 5, 0], dtype=torch.int32) - planner._staged_seq = 7 - selected = torch.zeros((4, 5)) - planner._note_snapshot_stats(selected, rows=[0, 1, 2, 3]) - hist = planner.stats["stamp_lag_hist"] - # Stamp 0 is "never drafted": lag 7 in truth, and it must land visibly at - # the histogram's far end rather than vanish -- these are exactly the rows - # the neutral counter should agree about. - assert hist == {0: 1, 1: 1, 2: 1, 7: 1} - - -def test_large_and_negative_lags_clamp_instead_of_exploding_the_dict(): - planner = _planner() - planner._host_stamps = torch.tensor([0, 500], dtype=torch.int32) - planner._staged_seq = 400 - planner._note_snapshot_stats(torch.zeros((2, 5)), rows=[0, 1]) - assert planner.stats["stamp_lag_hist"] == {8: 1, -2: 1} - - -def test_counts_accumulate_across_steps(): - planner = _planner() - planner._host_stamps = torch.tensor([3, 3], dtype=torch.int32) - planner._staged_seq = 4 - for _ in range(3): - planner._note_snapshot_stats(torch.full((2, 5), 30.0), rows=[0, 1]) - assert planner.stats["snap_rows"] == 6 - assert planner.stats["snap_neutral_rows"] == 6 - assert planner.stats["stamp_lag_hist"] == {1: 6} - - -def test_instrument_failure_is_counted_not_raised(): - planner = _planner() - planner._host_stamps = torch.tensor([1], dtype=torch.int32) - planner._staged_seq = 1 - # Rows that cannot index the stamp buffer must clamp, and anything worse - # must be swallowed into the error counter: instruments never kill steps. - planner._note_snapshot_stats(torch.zeros((1, 5)), rows=[10_000]) - assert planner.stats.get("snap_stats_errors", 0) == 0 - planner._note_snapshot_stats(object(), rows=None) # type: ignore[arg-type] - assert planner.stats["snap_stats_errors"] == 1 - - -def test_decide_path_feeds_the_instruments(): - """The counters must fill from the real decide path, not only direct calls.""" - planner = _planner() - bs = 8 - rng = np.random.default_rng(3) - survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() - planner._gather_rows = lambda **_: torch.tensor(survival, - dtype=torch.float32) - lens = planner.decide_verify_lens(num_gen_requests=bs, - reduce_across_ranks=False) - assert lens is not None - assert planner.stats["snap_rows"] == bs - assert planner.stats["snap_neutral_rows"] == 0 - # The local-decision instruments must fill on the same path: one rung per - # decision, one window per request, and the two must agree -- the windows - # are the top-k realisation of exactly that rung's budget. - rung_hist = planner.stats["local_rung_hist"] - assert sum(rung_hist.values()) == 1 - (rung,) = rung_hist - assert rung in (1, 2, 5) - len_hist = planner.stats["local_len_hist"] - assert sum(len_hist.values()) == bs - budget = bs * (rung - planner.cfg.min_verify_len) - extra = sum((k - planner.cfg.min_verify_len) * v for k, v in len_hist.items()) - assert extra <= budget diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py index 07e12fac11e4..aa223abcb9b9 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sps_profiler.py @@ -14,45 +14,35 @@ # limitations under the License. """Unit tests for the DSpark SPS cost-table profiler (CPU only, no model load). -Everything here exercises the pure half of the profiler: turning ``get_stats`` -rows into aligned samples, summarizing cells, separating ``alpha(bs)`` from -``theta(M)``, and assembling / refusing the emitted JSON. The sweep driver needs -a GPU and a checkpoint and is deliberately not covered. - -The load-bearing case is :func:`test_round_trip_reproduces_measured_cells`: it -takes a synthetic machine with a known cost curve, runs it through the whole -fit-and-emit pipeline, and reloads the JSON exactly the way -``dspark.py::_build_verify_planner`` does. If the profiler and the loader ever -disagree about units, keys or the meaning of a token count, that test is where -it shows up. +Covers stats-row alignment, cell summarizing, the additive alpha/theta fit, +staircase compression, payload assembly/refusals, sweep geometry guards, the +profitability probe, and reading samples out of production iteration logs. """ import importlib.util import json import pathlib import sys +import tempfile import numpy as np import pytest from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable, total_verify_tokens + # The profiler is a development tool under tests/microbenchmarks, not part of -# the shipped package: nothing on the serving path imports it, and keeping a -# 1.8k-line sweep driver inside tensorrt_llm/ would ship it to every user. -# Load it by path so this test follows the file rather than an import root. +# the shipped package; load it by path so this test follows the file. The +# module must be registered in sys.modules BEFORE exec so @dataclass can +# resolve annotations through sys.modules[cls.__module__].__dict__. _SPEC = importlib.util.spec_from_file_location( "dspark_sps_profiler", pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / "dspark_sps_profiler.py") _PROFILER = importlib.util.module_from_spec(_SPEC) -# Registered BEFORE exec: @dataclass resolves annotations through -# sys.modules[cls.__module__].__dict__, so a module executed without being -# registered raises "'NoneType' object has no attribute '__dict__'" on its -# first dataclass. sys.modules[_SPEC.name] = _PROFILER _SPEC.loader.exec_module(_PROFILER) -globals().update({name: getattr(_PROFILER, name) for name in ('CellStat', 'FlatCostTableError', 'InertCostTableError', 'InsufficientSamplesError', 'StepSample', 'SweepConfig', 'SweepGeometryError', 'aligned_steps_from_stats', 'build_cost_table_payload', 'check_table_is_informative', 'compress_to_risers', 'fit_additive_cost_model', 'load_cost_table', 'profitability_probe', 'running_max', 'summarize_cells')}) +globals().update({name: getattr(_PROFILER, name) for name in ('CellStat', 'FlatCostTableError', 'InertCostTableError', 'InsufficientSamplesError', 'StepSample', 'SweepConfig', 'SweepGeometryError', 'aligned_steps_from_stats', 'build_cost_table_payload', 'check_table_is_informative', 'compress_to_risers', 'fit_additive_cost_model', 'load_cost_table', 'profitability_probe', 'resolve_padded_shape', 'running_max', 'samples_from_iter_log', 'summarize_cells')}) #: The exact set of keys ``dspark.py::_build_verify_planner`` reads. Anything #: else in the file is ignored by the loader, so the profiler must not rely on @@ -101,12 +91,8 @@ def test_aligned_steps_drops_prefill_iterations(): def test_aligned_steps_requires_every_dp_rank(): - """Multi-rank alignment: complete, agreeing, and timed by rank 0. - - A missing rank means the iteration was not observed batch-wide; ranks at - different batch sizes are not one (bs, M) point; and ADP fans out - rank-local counters with rank-0's clock, so timing must never be averaged. - """ + """Multi-rank alignment: iterations must be complete, agree on batch size, + and be timed by rank 0's clock, never averaged.""" rows = [ _stats_row(iteration=0, num_gen=8, rank=0), _stats_row(iteration=0, num_gen=8, rank=1), @@ -168,16 +154,8 @@ def test_summarize_reports_median_not_mean(): def test_summarize_drops_thin_cells_outside_the_requested_grid(): - """A deployment log carries ramp/drain occupancy nobody asked to measure. - - Job 2585129: a five-rung server sweep with 984-step cells at every - requested (bs, L) died at the very end on (bs=2, L=1) with one sample -- - the admission ramp of the first cell's client, filed as a cell and then - held to the requested-grid bar. Incidental cells must be dropped when - thin, kept when thick, and a thin REQUESTED cell must still refuse -- - with or without a requested grid, since a silently dropped cell can - disconnect the grid, not just widen error bars. - """ + """Thin incidental cells are dropped, thick ones kept, and a thin + REQUESTED cell must still refuse -- with or without a requested grid.""" grid = [(8, 3)] requested = _samples(8, 3, [10.0] * 8) ramp = _samples(2, 3, [30.0]) # thin, incidental: drop @@ -273,12 +251,8 @@ def test_fit_recovers_relative_structure_exactly(): def test_fit_surfaces_a_bad_cell_rather_than_hiding_it(): - """A cell the additive model cannot explain is evidence, not noise. - - Dropping it would produce a table that fits beautifully and predicts badly, - and would delete the one signal that says ``theta`` is not a function of - ``M`` alone. - """ + """A cell the additive model cannot explain is surfaced in the warnings, + not dropped as noise.""" def noise(batch_size, verify_len): return 500.0 if (batch_size, verify_len) == (32, 3) else 0.0 @@ -326,13 +300,8 @@ def test_running_max_is_monotone(): def test_compress_keeps_risers_and_closes_their_shelves(): - """A shelf must survive as a breakpoint PAIR under the interp consumer. - - Keeping only [16, 64] would turn the measured flat 16..48 into one long - rising segment: mid-shelf totals get billed part of the riser -- the - over-trim mirror of the floor-lookup bug. The shelf's last point (48) - rides along so interpolation stays flat where the hardware was flat. - """ + """A shelf survives as a breakpoint PAIR so the interpolating consumer + stays flat where the hardware was flat.""" tokens = [16, 32, 48, 64, 80] times = [1.0, 1.002, 1.004, 2.0, 2.001] kept_tokens, kept_times = compress_to_risers( @@ -343,15 +312,8 @@ def test_compress_keeps_risers_and_closes_their_shelves(): def test_truncation_drops_the_point_cheapest_to_interpolate(): - """Over budget, remove the interior point interpolation best reconstructs. - - Here 64 goes: its neighbors predict it within 0.27ms, while removing 32 - would ramp the near-flat 16..32 stretch into the 64 riser and misprice it - by 0.9ms. The previous rule (drop the smallest incoming jump first) chose - 32 -- by construction it always chose the shelf-side point, re-billing - mid-shelf totals upward, the exact over-charge the pair encoding of - shelves exists to prevent. - """ + """Over budget, remove the interior point interpolation best reconstructs + -- never the shelf-side point, which would re-bill mid-shelf totals.""" tokens = [16, 32, 64, 128] times = [1.0, 1.1, 4.0, 9.0] kept_tokens, _ = compress_to_risers(tokens, times, min_riser_ms=0.0, max_breakpoints=3) @@ -377,15 +339,8 @@ def _payload(): def test_round_trip_reproduces_measured_cells(): - """The whole point: what the planner reads back must be what was measured. - - The payload must carry exactly the keys ``dspark.py::_build_verify_planner`` - reads (everything else under ``_meta``, where the loader ignores it), - survive a JSON file round trip into a real ``SpsCostTable``, and price - every measured cell back to within one shelf width -- the emitted curve is - a floor staircase by construction, so a cell between two risers is - deliberately priced at the lower shelf. - """ + """The payload carries exactly the loader's keys, survives a JSON round + trip into a real SpsCostTable, and prices every cell within one shelf.""" payload, cells = _payload() assert LOADER_KEYS.issubset(payload) assert set(payload) - LOADER_KEYS == {"_meta"} @@ -461,12 +416,8 @@ def test_refuses_a_curve_that_only_moves_by_noise(): # (M=192) above, so the curve is genuinely non-flat and clears the noise floor # (2.5%), yet the step only gets 2.5% more expensive while the acceptance yield # given up between those tiers is ~36% at p=0.9. The argmax therefore never -# leaves max_tier -- non-flat and inert at the same time. -# -# The shelf is ENCODED as a breakpoint pair: under the interpolating consumer -# a bare [0, 160] would ramp from 2.0 to 2.5 across the whole range, dropping -# the ladder spread below the noise floor and turning this fixture into the -# flat-noise case instead of the inert one it exists to pin. +# leaves max_tier -- non-flat and inert at the same time. The shelf is encoded +# as a breakpoint pair so the interpolating consumer keeps it flat. INERT_PAYLOAD = { "token_counts": [0, 159, 160], "step_time_ms": [2.0, 2.0, 2.5], @@ -477,12 +428,8 @@ def test_refuses_a_curve_that_only_moves_by_noise(): def test_refuses_an_inert_table(): - """Non-flat, but no tier pair is worth trading yield for. - - This is the failure mode that does *not* trip ``fallback_flat_cost``, so - nothing downstream would report it. ``allow_inert=True`` is the explicit - escape hatch -- and it must still say that nothing trims. - """ + """Non-flat, but no tier pair is worth trading yield for; allow_inert=True + is the escape hatch and must still say that nothing trims.""" kwargs = dict(batch_sizes=[32], tiers=[1, 3, 5], acceptance_rates=[0.9, 0.95]) with pytest.raises(InertCostTableError, match="inert"): check_table_is_informative(INERT_PAYLOAD, **kwargs) @@ -589,11 +536,8 @@ def test_probe_trims_at_low_acceptance_and_not_at_high(): def test_probe_uses_the_bonus_token_convention(): - """A tier's cost must be looked up at bs*(L+1), never bs*L. - - The two differ by a whole batch, which is easily a shelf's width; this pins - the profiler to the same convention as ``total_verify_tokens``. - """ + """A tier's cost must be looked up at bs*(L+1), never bs*L -- the same + convention as ``total_verify_tokens``.""" payload = { # Priced so that tier 5 (M = 8*6 = 48) is expensive and tier 3 # (M = 8*4 = 32) is cheap. Under the wrong bs*L convention tier 5 would @@ -610,3 +554,115 @@ def test_probe_uses_the_bonus_token_convention(): assert table.step_time(total_verify_tokens(8, 3), 8) == pytest.approx(1.0) probes = profitability_probe(payload, batch_size=8, tiers=[1, 3, 5], acceptance_rates=[0.9]) assert probes[0]["chosen_verify_len"] == 3 + + +# -------------------------------------------------------------------------- +# reading cost-table samples out of production iteration logs +# -------------------------------------------------------------------------- + +# A deployment's captured ladder, abbreviated. Padding rounds a step's rows up +# to the smallest entry at or above them, which is what turns a padded token +# total into exactly one verify length. +LADDER = [1, 2, 4, 8, 16, 32, 64, 128, 192, 256] + + +def _line(iteration, bs, step_ms, gen_tokens, rank="0"): + return (f"{rank}: [08/04/2026-19:08:13] [TRT-LLM] [I] [_torch][RANK 0] " + f"iter = {iteration}, global_rank = 0, rank = 0, " + f"num_scheduled_requests = {bs}, kv_cache_util = 0.125, " + f"currank_total_requests = 1/8, host_step_time = {step_ms}ms, " + f"prev_device_step_time = {step_ms}ms, " + f"timestamp = 2026-08-04 19:08:13, " + f"states = {{'num_ctx_requests': 0, 'num_ctx_tokens': 0, " + f"'num_generation_tokens': {gen_tokens}, " + f"'cached_kv_tokens': 63821}}") + + +def _write(lines, name="gen.log"): + # tempfile rather than pytest's tmp_path: these tests also run under the + # fixture-less in-container runner used for the GPU images. + path = pathlib.Path(tempfile.mkdtemp()) / name + path.write_text("\n".join(lines) + "\n") + return str(path) + + +def test_verify_length_comes_from_the_padded_token_total(): + """The label is read out of the log, never inferred from step times; steps + the ladder cannot explain are skipped, not mislabelled.""" + log = _write([ + _line(1, 1, 62098.0, 6), # executor's first step: dropped + _line(2, 244, 79.2, 768), # 244 rows padded to 256 -> rung-2 + _line(3, 242, 117.1, 1536), # padded to 256 -> full block + _line(4, 128, 61.0, 384), # exactly 128 rows -> rung-2 + _line(5, 256, 300.0, 2048), # L = 7, past the block: skipped + _line(6, 300, 400.0, 1800), # more rows than the ladder holds: skipped + ]) + samples = samples_from_iter_log([log], max_draft_len=5, + padded_batch_sizes=LADDER) + got = sorted((s.batch_size, s.verify_len, s.step_time_ms) + for s in samples) + assert got == [(128, 2, 61.0), (256, 2, 79.2), (256, 5, 117.1)] + # The table is indexed by the padded token total, which each sample exposes. + assert sorted(s.total_verify_tokens for s in samples) == [384, 768, 1536] + + +def test_the_padded_width_is_what_gets_recorded(): + """A padded step is filed under the padded width, the ladder disambiguates + the verify length, and an unexplainable token total drops the step.""" + assert resolve_padded_shape(num_rows=244, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) == (256, 2) + assert resolve_padded_shape(num_rows=128, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256) is None + assert resolve_padded_shape(num_rows=128, num_generation_tokens=768, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) == (128, 5) + assert resolve_padded_shape(num_rows=200, num_generation_tokens=600, + max_draft_len=5, max_batch_size=256, + padded_batch_sizes=LADDER) is None + + +def test_the_first_step_of_each_executor_is_dropped(): + """Each executor numbers from iter 1 and its first step's host time + contains the wait for the first request, so it is never a measurement.""" + log = _write([ + _line(1, 1, 200.0, 6), # KV-estimation executor + _line(2, 256, 79.0, 768), + _line(1, 1, 62098.0, 6), # real executor starts over + _line(2, 256, 79.4, 768), + ]) + times = [s.step_time_ms for s in samples_from_iter_log([log], max_draft_len=5)] + assert times == [79.0, 79.4] + + +def test_only_the_requested_rank_is_read(): + """Every rank logs; counting all of them would multiply every cell. + Outlier steps are bounded out on the same pass.""" + log = _write([ + _line(2, 256, 79.2, 768, rank="0"), + _line(2, 256, 79.9, 768, rank="3"), # another rank's copy: dropped + _line(3, 256, 5000.0, 768), # outlier past max_step_ms + ]) + times = [s.step_time_ms for s in + samples_from_iter_log([log], max_draft_len=5, max_step_ms=1000.0)] + assert times == [79.2] + + +def test_several_logs_are_pooled(): + """One arm gives one rung; the table needs the ladder, so pool the runs.""" + a = _write([_line(1, 1, 100.0, 6), _line(2, 256, 117.1, 1536)], + name="a.log") + c = _write([_line(1, 1, 100.0, 6), _line(2, 256, 79.2, 768)], + name="c.log") + got = sorted(s.verify_len for s in + samples_from_iter_log([a, c], max_draft_len=5)) + assert got == [2, 5] + + +def test_a_log_without_the_states_dict_yields_nothing(): + """enable_iter_perf_stats off: fail loudly upstream, not with a bad table.""" + path = pathlib.Path(tempfile.mkdtemp()) / "bare.log" + path.write_text( + "0: [TRT-LLM] iter = 2, num_scheduled_requests = 256, " + "host_step_time = 79.2ms\n") + assert samples_from_iter_log([str(path)], max_draft_len=5) == [] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py similarity index 58% rename from tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py rename to tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py index d75d95a74495..f9644d183e3d 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_ring_pairing.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts.py @@ -1,21 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""STS pairing must join on draft-pass identity, never on arrival order. - -The recorder's label (accepted count) and its feature (confidence logits) are -produced by different steps: the block verified at step ``i`` was drafted at -step ``i-1``, and under the overlap scheduler the sampler consumes the label -while pass ``i+1`` has already overwritten the shared confidence buffer. The -single mutable stash this ring replaced had NO execution order under which -the pair was correct -- and the shards looked perfectly healthy. - -These tests pin the join contract: a pair is appended only when the ring -snapshot's per-row stamp equals the pass the label verifies, everything else -is a *counted* decline, and the shard carries provenance the fitter can -refuse. No CUDA: the recorder itself is device-agnostic (CPU tensors skip the -pinned-copy event), which is exactly what lets this run in pre-merge CI. +"""DSpark STS calibration subsystem: recorder ring pairing, shard provenance, +fitter acceptance, collection-regime guards, and confidence-head resolution. +Runs on CPU only (the recorder is device-agnostic), so it fits pre-merge CI. """ +import contextlib import importlib.util import os import pathlib @@ -25,11 +15,15 @@ import pytest import torch -from tensorrt_llm._torch.speculative.dspark_sts import DSparkStsRecorder +from tensorrt_llm._torch.speculative.dspark_sts import (DSparkStsRecorder, + STS_COLLECT_ENV, + make_recorder_from_env) BLOCK = 5 ROWS = 8 +_PIN = "TLLM_DSPARK_FORCE_VERIFY_LEN" + def _recorder(tmp, flush_every=1000): return DSparkStsRecorder(path_stem=os.path.join(tmp, "shard"), @@ -38,12 +32,7 @@ def _recorder(tmp, flush_every=1000): def _stage(rec, seq, written_rows, fill): - """Stage one pass's snapshot: `written_rows` stamped `seq`, rest stale. - - Mirrors the producer: the in-graph scatter stamps only the rows the pass - actually drafted; every other row keeps whatever stamp it had (0 for - never-written, matching the worker's zero-initialized stamp buffer). - """ + """Stage one pass's snapshot: `written_rows` stamped `seq`, rest stale.""" logits = torch.full((ROWS, BLOCK), float(fill)) stamps = torch.zeros(ROWS, dtype=torch.int32) stamps[list(written_rows)] = seq @@ -76,9 +65,7 @@ def test_row_not_written_for_that_pass_is_declined(): def test_ring_eviction_declines_instead_of_mispairing(): """A label older than the ring depth is dropped, never matched to the - - evicting pass's content: the slot's stamps say who actually wrote it. - """ + evicting pass's content.""" with tempfile.TemporaryDirectory() as tmp: rec = _recorder(tmp) for seq in range(1, 2 + DSparkStsRecorder._RING_DEPTH): @@ -89,7 +76,7 @@ def test_ring_eviction_declines_instead_of_mispairing(): def test_every_giving_up_path_is_counted(): - """Silent declines are how the mispaired fit shipped; none may be silent.""" + """No decline path may be silent; every one increments a named counter.""" with tempfile.TemporaryDirectory() as tmp: rec = _recorder(tmp) rec.record(row=2, accepted=1, target_seq=None) # no pass id @@ -117,6 +104,7 @@ def test_staging_without_stamps_is_a_noop(): def test_shard_carries_provenance_and_prefix_semantics(): + """Flushed shards carry ring provenance, stats, and prefix-mask labels.""" with tempfile.TemporaryDirectory() as tmp: rec = _recorder(tmp) _stage(rec, seq=1, written_rows=[2], fill=1.5) @@ -132,11 +120,6 @@ def test_shard_carries_provenance_and_prefix_semantics(): assert blob["prefix_mask"].tolist() == [[1.0, 1.0, 1.0, 0.0, 0.0]] -# --------------------------------------------------------------------------- -# The fitter must refuse shards from the pre-ring recorders: their pairs are -# mislabeled in ways the tensors themselves cannot reveal. -# --------------------------------------------------------------------------- - _SPEC = importlib.util.spec_from_file_location( "dspark_fit_sts", pathlib.Path(__file__).resolve().parents[4] / "microbenchmarks" / @@ -147,6 +130,7 @@ def test_shard_carries_provenance_and_prefix_semantics(): def test_fitter_refuses_shards_without_ring_provenance(): + """The fitter must refuse shards lacking draft_seq_ring provenance.""" with tempfile.TemporaryDirectory() as tmp: torch.save({"logits": torch.zeros(4, BLOCK), "prefix_mask": torch.zeros(4, BLOCK)}, @@ -156,6 +140,7 @@ def test_fitter_refuses_shards_without_ring_provenance(): def test_fitter_accepts_ring_shards(): + """Shards produced by the ring recorder load cleanly in the fitter.""" with tempfile.TemporaryDirectory() as tmp: rec = _recorder(tmp) _stage(rec, seq=1, written_rows=[2, 3], fill=0.5) @@ -164,3 +149,89 @@ def test_fitter_accepts_ring_shards(): rec.flush() logits, mask = _FITTER.load_shards(os.path.join(tmp, "*.pt")) assert logits.shape == (2, BLOCK) and mask.shape == (2, BLOCK) + + +@contextlib.contextmanager +def collecting(pin=None): + """Set the collection env for the block, restore it after.""" + saved = {k: os.environ.get(k) for k in (STS_COLLECT_ENV, _PIN)} + try: + with tempfile.TemporaryDirectory() as tmp: + os.environ[STS_COLLECT_ENV] = os.path.join(tmp, "collect") + if pin is None: + os.environ.pop(_PIN, None) + else: + os.environ[_PIN] = str(pin) + yield + finally: + for k, v in saved.items(): + os.environ.pop(k, None) if v is None else os.environ.__setitem__(k, v) + + +def test_collection_off_builds_nothing(): + """The serving default must be untouched by any of this.""" + os.environ.pop(STS_COLLECT_ENV, None) + assert make_recorder_from_env(block_size=5) is None + assert make_recorder_from_env(block_size=5, has_cost_table=True, + ragged_mode="compact") is None + + +def test_a_loaded_cost_table_without_a_pin_is_refused(): + """A table means the planner trims, and trimming censors the label.""" + with collecting(): + with pytest.raises(ValueError, match="censors the acceptance label"): + make_recorder_from_env(block_size=5, has_cost_table=True) + + +def test_a_pinned_window_makes_the_table_harmless(): + """Every uncensored regime (pinned window, no table, static mode) must + build a recorder.""" + with collecting(pin=5): + rec = make_recorder_from_env(block_size=5, has_cost_table=True) + assert rec is not None and rec.block_size == 5 + with collecting(): + assert make_recorder_from_env(block_size=5, + has_cost_table=False) is not None + assert make_recorder_from_env(block_size=5, + ragged_mode="static") is not None + + +def test_compact_mode_is_refused(): + """Padding rows contribute prefix labels that measure nothing; refused + even with the pin set.""" + with collecting(pin=5): + with pytest.raises(ValueError, match="Padded verify rows"): + make_recorder_from_env(block_size=5, ragged_mode="compact") + + +class _Head: + pass + + +class _Bare: + """DSparkDraftModel layout: stages under .mtp_layers.""" + + def __init__(self, head): + stage = type("Stage", (), {})() + if head is not None: + stage.confidence_head = head + self.mtp_layers = [object(), object(), stage] + + +class _Wrapper: + """DSparkForCausalLM layout: bare model under .dspark_model.""" + + def __init__(self, head): + self.dspark_model = _Bare(head) + + +def test_head_resolves_across_both_layouts_and_absence_is_none(): + """resolve_confidence_head finds the head under either model layout and + returns None when absent.""" + from tensorrt_llm._torch.speculative.dspark_sts import resolve_confidence_head + h = _Head() + assert resolve_confidence_head(_Bare(h)) is h + assert resolve_confidence_head(_Wrapper(h)) is h + assert resolve_confidence_head(_Bare(None)) is None + assert resolve_confidence_head(_Wrapper(None)) is None + assert resolve_confidence_head(object()) is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py deleted file mode 100644 index 1dab83acaaa1..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_sts_collection_guards.py +++ /dev/null @@ -1,127 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""A collection run must refuse the regimes that fit the scheduler, not the head. - -Both regimes below produce shards that look perfectly well-formed. The run -completes, the row counts are plausible, the fitter converges, and the emitted -temperatures are wrong -- there is no downstream check that can tell. That is -why these are construction-time failures rather than warnings. - -Measured on this branch: collecting while the planner trimmed gave -``ece_before = [.244 .319 .361 .389 .394]``; pinning the window to the full -block, same model and same fitter, gave ``[.157 .208 .270 .321 .344]``. The -difference is entirely the censored label. - -SGLang guards the first regime next to its own collector -(``dspark_planner.py:100-108``) and names the second in the sibling -``ConfidenceMetricsProbe`` -- "padded verify rows corrupt the per-position -prefix label" (``dspark_observability.py:688-697``). -""" - -import contextlib -import os -import tempfile - -import pytest - -from tensorrt_llm._torch.speculative.dspark_sts import (STS_COLLECT_ENV, - make_recorder_from_env) - -_PIN = "TLLM_DSPARK_FORCE_VERIFY_LEN" - - -@contextlib.contextmanager -def collecting(pin=None): - """Set the collection env for the block, restore it after.""" - saved = {k: os.environ.get(k) for k in (STS_COLLECT_ENV, _PIN)} - try: - with tempfile.TemporaryDirectory() as tmp: - os.environ[STS_COLLECT_ENV] = os.path.join(tmp, "collect") - if pin is None: - os.environ.pop(_PIN, None) - else: - os.environ[_PIN] = str(pin) - yield - finally: - for k, v in saved.items(): - os.environ.pop(k, None) if v is None else os.environ.__setitem__(k, v) - - -def test_collection_off_builds_nothing(): - """The serving default must be untouched by any of this.""" - os.environ.pop(STS_COLLECT_ENV, None) - assert make_recorder_from_env(block_size=5) is None - assert make_recorder_from_env(block_size=5, has_cost_table=True, - ragged_mode="compact") is None - - -def test_a_loaded_cost_table_without_a_pin_is_refused(): - """A table means the planner trims, and trimming censors the label.""" - with collecting(): - with pytest.raises(ValueError, match="censors the acceptance label"): - make_recorder_from_env(block_size=5, has_cost_table=True) - - -def test_a_pinned_window_makes_the_table_harmless(): - """The uncensored regimes must all build a recorder. - - Pinning to the full block makes the label uncensored by construction; - without a table the planner cannot trim, so nothing censors anything; - static mode contributes no padding rows. - """ - with collecting(pin=5): - rec = make_recorder_from_env(block_size=5, has_cost_table=True) - assert rec is not None and rec.block_size == 5 - with collecting(): - assert make_recorder_from_env(block_size=5, - has_cost_table=False) is not None - assert make_recorder_from_env(block_size=5, - ragged_mode="static") is not None - - -def test_compact_mode_is_refused(): - """Padding rows contribute prefix labels that measure nothing. - - Refused even with the pin set: the pin fixes the *window*, while this is - about rows that belong to no real request at all. - """ - with collecting(pin=5): - with pytest.raises(ValueError, match="Padded verify rows"): - make_recorder_from_env(block_size=5, ragged_mode="compact") - - -class _Head: - pass - - -class _Bare: - """DSparkDraftModel layout: stages under .mtp_layers.""" - - def __init__(self, head): - stage = type("Stage", (), {})() - if head is not None: - stage.confidence_head = head - self.mtp_layers = [object(), object(), stage] - - -class _Wrapper: - """DSparkForCausalLM layout: bare model under .dspark_model. - - This is the layout the chained getattr missed: the wrapper has no - mtp_layers of its own, so `getattr(model, "mtp_layers", [None])[-1]` - resolved to None and silently disabled calibration for every - confidence-scheduled run of the first campaign. - """ - - def __init__(self, head): - self.dspark_model = _Bare(head) - - -def test_head_resolves_across_both_layouts_and_absence_is_none(): - from tensorrt_llm._torch.speculative.dspark_sts import resolve_confidence_head - h = _Head() - assert resolve_confidence_head(_Bare(h)) is h - assert resolve_confidence_head(_Wrapper(h)) is h - assert resolve_confidence_head(_Bare(None)) is None - assert resolve_confidence_head(_Wrapper(None)) is None - assert resolve_confidence_head(object()) is None diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py deleted file mode 100644 index 2160e1d0ba49..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_tier_aligned_budget.py +++ /dev/null @@ -1,96 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The budget must be scored with the tau the step actually collects. - -``decide_verify_lens`` picks a rung, converts it to ``B = n*(L - min)`` tokens, -and hands B to ``schedule_verify_lens_topk``, which spends it on the B largest -survivals in the batch. But the rung was chosen by -``budget_argmax_over_uniform_lens``, which scores rung L as if every request took -columns ``[min, L)``. Top-k is by construction the largest sum over all size-B -subsets, so the realised tau is >= the scored one at an identical token cost -- -and the gap is not the same at every rung, so the argmax can land on the wrong -one. - -The fix is not to let the planner choose an arbitrary budget: the executor can -only run ``bs*(t+1)`` tokens for a captured tier, and any other total rounds UP, -which would push a budget deliberately chosen below a cost riser back over it. -So the ladder stays the grid and only the numerator changes. -""" - -import numpy as np -import pytest - -from tensorrt_llm._torch.speculative.dspark_planner import ( - SpsCostTable, budget_argmax_over_uniform_lens, compute_verify_token_budget) - - -def _table() -> SpsCostTable: - """A staircase with genuine risers, so trimming is worth something. - - Breakpoints are spaced so that neither tier is trivially optimal: a flat - table degenerates to verify-all and would make every assertion here vacuous. - """ - token_counts = tuple(range(0, 400, 16)) - step_time_ms = tuple(4.0 + 0.6 * (tok // 96) + 0.004 * tok - for tok in token_counts) - return SpsCostTable(token_counts=token_counts, step_time_ms=step_time_ms, - fixed_overhead_ms=1.0) - - -@pytest.mark.parametrize("tiers", [[1, 2, 5], [1, 3, 5], [1, 5]]) -def test_restricted_answer_is_always_realisable(tiers): - """Every returned budget corresponds to a rung the executor captured.""" - rng = np.random.default_rng(20260803) - table = _table() - for _ in range(200): - bs = int(rng.integers(2, 17)) - survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] - n = compute_verify_token_budget( - survival=survival, num_gen_requests=bs, cost_table=table, - min_verify_len=1, allowed_lens=tiers) - assert n % bs == 0, ( - f"budget {n} for {bs} requests is not n*(t-min) for any tier") - assert (n // bs) + 1 in tiers, ( - f"budget {n} implies tier {(n // bs) + 1}, not in {tiers}") - - -def test_restricted_never_scores_worse_than_the_uniform_choice(): - """Same grid, better numerator: the chosen rung's realised theta wins. - - Both routes evaluate the same discrete cost points, so the restricted form - cannot be worse under the tau that the step actually collects. This is the - property that makes the swap safe -- if it ever failed, the change would be - trading a modelling error for a real one. At bs=1 the uniform and top-k - allocations are the same set, so there the two scorers must agree on the - rung outright: any disagreement is a plain bug, not the intended - improvement. - """ - rng = np.random.default_rng(31337) - table = _table() - tiers = [1, 2, 5] - for _ in range(300): - bs = int(rng.integers(1, 17)) - survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1] - - def realised_theta(tier: int) -> float: - budget = bs * (tier - 1) - cand = np.sort(survival[:, 1:].reshape(-1))[::-1] - tau = float(bs) + float(survival[:, :1].sum()) + float( - cand[:budget].sum()) - tokens = np.array([bs * (tier + 1)]) - return tau / float(table.step_times(tokens, bs)[0]) - - n = compute_verify_token_budget( - survival=survival, num_gen_requests=bs, cost_table=table, - min_verify_len=1, allowed_lens=tiers) - chosen = (n // bs) + 1 - uniform = budget_argmax_over_uniform_lens( - survival=survival, num_gen_requests=bs, cost_table=table, - allowed_lens=tiers, min_verify_len=1) - assert realised_theta(chosen) >= realised_theta(uniform) - 1e-12, ( - f"restricted picked tier {chosen} whose realised theta is below " - f"the uniform scorer's tier {uniform}") - if bs == 1: - assert chosen == uniform, ( - "at bs=1 the uniform and top-k allocations are the same set; " - "the two scorers must pick the same rung") diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py deleted file mode 100644 index 04b8d5c0dce6..000000000000 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_verify_len_pin.py +++ /dev/null @@ -1,97 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The runtime verify-length pin: queue, agree, adopt. - -A cost-table sweep has to hold every request at one verify length while it -measures a cell, or the cell's label describes a shape that never ran (this -branch shipped a sweep where eight of twelve cells were mislabelled and the -fit dutifully concluded the curve was flat). The environment variable that -does this today is read once at planner construction, so walking the ladder -means rebuilding the engine per length -- and on a pre-spawned world it -cannot reach the ranks that are already running. - -The runtime pin fixes both, but it must never let one rank pin while its -peers do not: the ranks would derive different token totals, and the -attention-DP shape gate would drop the whole group out of graph replay for as -long as they disagreed. So the pin is queued locally, carried by the step's -existing allgather, and adopted from that agreed payload -- which is what -these tests pin down. -""" - -import pytest - -from tensorrt_llm._torch.speculative.dspark_planner import SpsCostTable -from tensorrt_llm._torch.speculative.dspark_schedule import DSparkScheduleConfig -from tensorrt_llm._torch.speculative.dspark_verify import DSparkVerifyPlanner - - -def _planner(tiers=(1, 2, 5)): - table = SpsCostTable(token_counts=(0, 512, 768, 1536), - step_time_ms=(5.0, 68.4, 80.2, 150.5), - fixed_overhead_ms=1.0) - cfg = DSparkScheduleConfig(block_size=5, min_verify_len=1) - return DSparkVerifyPlanner(cfg=cfg, cost_table=table, tiers=list(tiers)) - - -def test_adopting_the_agreed_value_applies_it(): - """Queuing must not change this rank's behaviour before the group agrees; - adopting the agreed value applies it and consumes the pending request.""" - planner = _planner() - assert planner.request_verify_len_pin(2) == 2 - # Queued, not yet in force. - assert planner._forced_verify_len is None - assert planner.pending_verify_len_pin() == 2 - planner.adopt_verify_len_pin(2) - assert planner._forced_verify_len == 2 - # Consumed: the next step must not re-broadcast a pin nobody asked for. - assert planner.pending_verify_len_pin() == -1 - - -def test_a_rank_that_queued_nothing_still_adopts_the_group_value(): - """The endpoint lands on ONE rank; the others learn it from the payload.""" - peer = _planner() - assert peer.pending_verify_len_pin() == -1 - peer.adopt_verify_len_pin(5) - assert peer._forced_verify_len == 5 - - -def test_zero_clears_the_pin(): - """Wire protocol: -1 means nobody asked (leave the pin alone); clearing - needs its own value, 0, because None cannot travel in an int payload.""" - planner = _planner() - planner.adopt_verify_len_pin(2) - planner.adopt_verify_len_pin(-1) - assert planner._forced_verify_len == 2 - assert planner.request_verify_len_pin(None) is None - assert planner.pending_verify_len_pin() == 0 - planner.adopt_verify_len_pin(0) - assert planner._forced_verify_len is None - - -def test_an_uncaptured_length_is_refused_at_the_call_site(): - """Rejected when requested, not on some later step, and never queued.""" - planner = _planner(tiers=(1, 2, 5)) - with pytest.raises(ValueError, match="captured tier ladder"): - planner.request_verify_len_pin(3) - with pytest.raises(ValueError, match="outside"): - planner.request_verify_len_pin(9) - assert planner.pending_verify_len_pin() == -1 - assert planner._forced_verify_len is None - - -def test_pinned_steps_hand_every_request_the_same_window(): - """The whole point: the shape a sweep cell can honestly label.""" - import numpy as np - import torch - - planner = _planner() - bs = 8 - rng = np.random.default_rng(5) - survival = np.sort(rng.random((bs, 5)), axis=1)[:, ::-1].copy() - planner._gather_rows = lambda **_: torch.tensor(survival, - dtype=torch.float32) - planner.adopt_verify_len_pin(2) - lens = planner.decide_verify_lens(num_gen_requests=bs, - reduce_across_ranks=False) - assert lens == [2] * bs - assert planner.stats["forced_steps"] == 1 diff --git a/tests/unittest/_torch/speculative/test_dspark_graph_key.py b/tests/unittest/_torch/speculative/test_dspark_graph_key.py deleted file mode 100644 index df9ba57412c2..000000000000 --- a/tests/unittest/_torch/speculative/test_dspark_graph_key.py +++ /dev/null @@ -1,146 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""The ragged CUDA-graph key: fit, publication, and the padding ladder. - -Three regression guards for one invariant -- the graph key's token axis must -describe the batch the graph will actually replay over. Each test pins a bug -that shipped: a fit against unreachable pad rows, a key derived by re-summing -the batch instead of reading the fitted value (the ragged-IMA family), and a -capture pass that shaped a batch without publishing its bucket. -""" - - -class _Runner: - """Just the fields will_pad_to reads.""" - - def __init__(self, supported, cfg_batch_size): - self.supported_batch_sizes = supported - self.config = type("C", (), {"batch_size": cfg_batch_size})() - - -def test_will_pad_to_rejects_row_counts_the_ladder_cannot_reach(): - """The ragged fit must not derive a grid from rows that never appear. - - A run fitted its bucket grid to padded_bs=256 while the batch stayed at 193 - -- the captured ladder held 192 and 256 but not 193 -- so the filled total - 449 matched no captured bucket and every rank lost a graph replay on 16 of - 318 steps. Only reachable with real trimming: without it the total is - padded_bs * (max_verify_len + 1), a captured bucket by construction. - """ - from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ - CUDAGraphRunner - - runner = _Runner([1, 2, 4, 8, 128, 192, 256], cfg_batch_size=256) - will = CUDAGraphRunner.will_pad_to - - assert will(runner, 256, 193), "193 -> 256 is reachable and should be taken" - assert not will(runner, 193, 193), "193 is not a captured size" - assert will(runner, 192, 192), "already at a captured size" - # Padding past the configured batch size is refused by _get_padded_batch, - # so the fit must not assume it. - assert not will(_Runner([256, 512], cfg_batch_size=256), 512, 300) - assert not will(runner, 0, 10) - assert not will(runner, 256, 0) - - -def test_graph_key_uses_the_agreed_bucket_not_a_fresh_sum(): - """Every attention-DP rank must key on the same bucket, by construction. - - The key used to re-sum 1 + py_verify_len over generation_requests, which - walks the batch *after* padding is appended -- a second, independent - derivation of a value all ranks must match exactly. When the two - derivations diverged the shape gate dropped every rank out of replay - (peer_shape_mismatch: 2/265 steps without trimming, 16/318 with it, since - trimming widens the per-rank window spread). Reading the fitted value - instead makes agreement structural: it comes from allgathered peer stats - through rules every rank runs identically. - """ - from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ - CUDAGraphRunner - - import types - - class _R: - spec_config = type("S", (), {"enable_ragged_verify": True})() - agreed_ragged_bucket = None - - def _batch(*verify_lens): - return types.SimpleNamespace(generation_requests=[ - types.SimpleNamespace(py_verify_len=v) for v in verify_lens - ]) - - r = _R() - get = CUDAGraphRunner._ragged_verify_bucket - windowed = _batch(3, 2, 1) - - # Nothing fitted yet -> not ragged, key unchanged from the uniform one. - assert get(r, windowed) is None - - # A fitted bucket is returned verbatim: the batch's own token total is - # never re-summed here, which is what makes every rank agree. - r.agreed_ragged_bucket = 449 - assert get(r, windowed) == 449 - # ... including when that total plainly differs from the fitted value. - assert get(r, _batch(1, 1)) == 449 - - # But a request without a window means the token-major staging chain never - # ran this step, so a ragged key would replay the PREVIOUS ragged step's - # row maps -- the ragged IMA. Decline, whatever the fit published. - assert get(r, _batch(3, None, 1)) is None - - # Cleared between steps so a stale bucket cannot key into the wrong graph. - r.agreed_ragged_bucket = None - assert get(r, windowed) is None - - # Config gate still wins: a non-ragged engine keeps the original key shape. - r.agreed_ragged_bucket = 512 - r.spec_config = type("S", (), {"enable_ragged_verify": False})() - assert get(r, windowed) is None - - -def test_capture_publishes_the_bucket_it_shaped_the_batch_to(): - """The other half of the test above, and the regression it let through. - - Once the key reads a fitted value instead of re-summing the batch, whoever - shapes a batch has to publish the shape. Real steps do both in - fit_ragged_verify_lens. CUDA-graph capture shapes its warmup batch in - _set_warmup_ragged_windows and never called the fit, so the bucket stayed - None while the batch really held `verify_bucket` tokens. - - Both consequences are silent in different ways. The capture-time forward - ran with metadata describing a uniform batch over ragged tensors and died - hundreds of frames later in the MoE on `unmatched tensor shape`. Had it - survived, all three of a batch size's buckets would have been stored under - one key with no token axis -- colliding with each other, and matching no - runtime step, so every ragged step would have fallen back to eager without - an error anywhere. - """ - import types - - from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine - - runner = types.SimpleNamespace(agreed_ragged_bucket=None) - requests = [types.SimpleNamespace(py_verify_len=None) for _ in range(4)] - batch = types.SimpleNamespace(generation_requests=requests) - engine = types.SimpleNamespace(cuda_graph_runner=runner) - - # bs=4, draft_len=5 -> the t=5 bucket is 4 * 6 = 24 tokens. - PyTorchModelEngine._set_warmup_ragged_windows(engine, batch, 24, 5) - - assert sum(1 + r.py_verify_len for r in requests) == 24, ( - "the warmup batch must hit the bucket exactly") - assert runner.agreed_ragged_bucket == 24, ( - "capture shaped the batch to 24 tokens but did not publish it, so the " - "graph would be keyed as if the batch were uniform") - - # A narrower bucket at the same batch size, to prove it is per-entry and - # not set once: capture walks every tier for every batch size. - PyTorchModelEngine._set_warmup_ragged_windows(engine, batch, 8, 5) - assert sum(1 + r.py_verify_len for r in requests) == 8 - assert runner.agreed_ragged_bucket == 8 - - # An empty batch publishes nothing rather than a stale or zero bucket. - runner.agreed_ragged_bucket = None - PyTorchModelEngine._set_warmup_ragged_windows( - engine, types.SimpleNamespace(generation_requests=[]), 24, 5) - assert runner.agreed_ragged_bucket is None diff --git a/tests/unittest/_torch/speculative/test_dspark_planner_trims.py b/tests/unittest/_torch/speculative/test_dspark_planner_trims.py deleted file mode 100644 index 1262c8ff368f..000000000000 --- a/tests/unittest/_torch/speculative/test_dspark_planner_trims.py +++ /dev/null @@ -1,74 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""Make the planner itself produce per-request windows, with a crafted table. - -This replaces TLLM_DSPARK_FORCE_VERIFY_LENS. That knob assigned windows by -*batch position*, which is orthogonal to confidence: it reproduced the ragged -*shape* but never the ragged *policy*, so it could not tell a correct -confidence-driven assignment from an arbitrary one. It also had to sit in -production code, ahead of the cost-table gate, to be reachable at all. - -A cost table does the same job through the real path. The planner trims when -the marginal cost of extra verify tokens outweighs their expected acceptance, -so a table whose theta(M) rises steeply makes it trim at any acceptance rate -- -and the budget then flows through the same confidence-ordered top-k, which -hands different windows to requests with different survival. That exercises -what actually ships. - -Why this matters beyond tidiness: the KV rewind path (KV reserved for the full -drafted block but only the verified window rewound) is reachable only when -verified_len < draft_len, and a well-fitted table can decline to trim at every -batch size -- so without a table that forces trimming, that code path ships -untested. -""" - -import torch - -from tensorrt_llm._torch.speculative.dspark_schedule import ( - DSparkScheduleConfig, schedule_verify_lens_topk) - - -def _Cfg(min_verify_len=1, max_verify_len=5): - """The real config, so the test cannot drift from what ships.""" - return DSparkScheduleConfig(block_size=max_verify_len, - min_verify_len=min_verify_len, - max_verify_len=max_verify_len) - - -def test_topk_hands_longer_windows_to_more_confident_requests(): - """The property the position-rotating knob could never check. - - Request 0 is confident throughout, request 2 collapses immediately. A - correct scheduler spends the budget where survival is highest. - """ - # survival[r, k] = P(the first k+1 drafted tokens all get accepted) - survival = torch.tensor([ - [0.99, 0.98, 0.97, 0.96, 0.95], # confident - [0.90, 0.70, 0.50, 0.30, 0.10], # middling - [0.20, 0.04, 0.01, 0.00, 0.00], # collapses - ]) - lens = schedule_verify_lens_topk(survival=survival, budget=6, - cfg=_Cfg()).tolist() - - assert len(lens) == 3 - assert lens[0] > lens[2], ( - f"confident request got {lens[0]}, collapsing request got {lens[2]}; " - f"the budget must follow survival, not batch position") - assert lens[0] >= lens[1] >= lens[2] - # Every request keeps at least the floor, and the budget is respected. - assert min(lens) >= 1 - assert sum(lens) - len(lens) * 1 <= 6 - - -def test_a_full_budget_degenerates_to_the_uniform_full_window(): - """The no-trim case must stay exactly uniform. - - This is what the planner correctly chooses on this checkpoint, and it is - the case where ragged must cost nothing. - """ - num_reqs, max_len = 4, 5 - survival = torch.full((num_reqs, max_len), 0.99) - budget = num_reqs * (max_len - 1) - lens = schedule_verify_lens_topk(survival=survival, budget=budget, - cfg=_Cfg()).tolist() - assert lens == [max_len] * num_reqs diff --git a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py b/tests/unittest/_torch/speculative/test_force_accepted_tokens.py deleted file mode 100644 index 144de6a6c174..000000000000 --- a/tests/unittest/_torch/speculative/test_force_accepted_tokens.py +++ /dev/null @@ -1,363 +0,0 @@ -# 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. -"""Unit tests for the synthetic forced acceptance rate in 1-model spec dec. - -These tests exercise ``SpecWorkerBase._apply_force_accepted_tokens``, which -honors the ``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS`` environment -variable. The implementation must be: - - - Faithful to fractional rates: integer part always accepts, fractional - part acts as the per-iteration probability of accepting one extra draft - token. - - Tensor-parallel deterministic: identical seed/counter state across - ranks so all ranks agree on per-request accepted counts (otherwise - downstream collectives expecting matching shapes deadlock). - - CUDA-graph compatible: pure-tensor ops with statically-known shapes, - captureable and replayable. -""" - -import os -from typing import Optional -from unittest.mock import patch - -import pytest -import torch - -from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata -from tensorrt_llm._torch.speculative.interface import ( - FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR, - SpecWorkerBase, - get_force_num_accepted_tokens, - get_force_num_accepted_tokens_float, -) - - -class _StubSpecWorker(SpecWorkerBase): - """Concrete ``SpecWorkerBase`` that stubs out the abstract API. - - Used purely to drive ``_apply_force_accepted_tokens`` in isolation. - """ - - @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: - worker = _StubSpecWorker() - if value is not None: - worker.force_num_accepted_tokens = value - return worker - - -def _make_input(batch_size: int, device: str = "cuda") -> torch.Tensor: - # Initial counts: 1 (target token, no draft accepts yet). The forced - # override only ever rewrites entries from ``num_contexts:`` onward. - return torch.ones(batch_size, dtype=torch.int32, device=device) - - -def _require_cuda(): - if not torch.cuda.is_available(): - pytest.skip("CUDA required for this test.") - - -# ---------------- env-var parsing helpers ----------------------------------- - - -@pytest.mark.parametrize( - "env_value, expected_float, expected_int", - [ - ("0", 0.0, 0), - ("3", 3.0, 3), - ("2.6", 2.6, 0), - ("0.5", 0.5, 0), - ("not-a-number", 0.0, 0), - ], -) -def test_get_force_num_accepted_tokens_float(env_value, expected_float, expected_int): - """The float helper parses fractional rates; the int helper (used by the - 2-model path) must keep its original behavior on the same values.""" - with patch.dict(os.environ, {FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR: env_value}): - assert get_force_num_accepted_tokens_float() == pytest.approx(expected_float) - assert get_force_num_accepted_tokens() == expected_int - - -# ---------------- _apply_force_accepted_tokens semantics -------------------- - - -def test_zero_value_is_noop(): - _require_cuda() - worker = _make_worker(0.0) - base = _make_input(batch_size=4) - before = base.clone() - out = worker._apply_force_accepted_tokens(base, num_contexts=0, runtime_draft_len=4) - assert torch.equal(out, before) - # No RNG state should have been touched in the early-exit path. - assert worker._force_accept_rng_pool is None - assert worker._force_accept_rng_counter is None - - -@pytest.mark.parametrize( - "value, runtime_draft_len", - [ - (1.0, 4), - (2.0, 4), - (4.0, 4), - # 10 draft tokens requested but only 2 available -> capped at 3. - (10.0, 2), - # Fractional but no room: max_total = 3, base_total = min(10, 3) = 3, - # so ``base_total < max_total`` is False -> all 3, no RNG draw. - (9.5, 2), - ], -) -def test_integer_value_matches_legacy_behavior(value, runtime_draft_len): - _require_cuda() - worker = _make_worker(value) - out = worker._apply_force_accepted_tokens( - _make_input(batch_size=8), num_contexts=0, runtime_draft_len=runtime_draft_len - ) - expected = min(int(value) + 1, runtime_draft_len + 1) - assert torch.all(out == expected) - # The deterministic (non-random) path must not allocate the RNG pool. - assert worker._force_accept_rng_pool is None - - -def test_fractional_distribution_matches_target_probability(): - """Fractional draws emit only base/base+1 and average to the target rate.""" - _require_cuda() - worker = _make_worker(2.6) - target_frac = 0.6 - n_iters = 200 - batch = 64 - seen = set() - extra_count = 0 - total = 0 - for _ in range(n_iters): - out = worker._apply_force_accepted_tokens( - _make_input(batch_size=batch), num_contexts=0, runtime_draft_len=4 - ) - seen.update(out.unique().tolist()) - extra_count += int((out == 4).sum().item()) - total += batch - # Either 2 draft + target = 3, or 3 draft + target = 4 -- never anything else. - assert seen == {3, 4} - measured = extra_count / total - # Pool-based RNG is deterministic; the empirical mean over ~13k draws - # is tight against 0.6. - assert abs(measured - target_frac) < 0.03, ( - f"measured fraction {measured:.4f} differs from target {target_frac}" - ) - - -def test_num_contexts_offset_is_respected(): - """Context (prefill) rows must be left untouched by the override.""" - _require_cuda() - worker = _make_worker(2.6) - batch_size = 8 - num_contexts = 3 - base = _make_input(batch_size=batch_size) - sentinel = torch.iinfo(base.dtype).max - base[:num_contexts] = sentinel - out = worker._apply_force_accepted_tokens(base, num_contexts=num_contexts, runtime_draft_len=4) - assert torch.all(out[:num_contexts] == sentinel) - assert torch.all((out[num_contexts:] == 3) | (out[num_contexts:] == 4)) - - -# ---------------- TP determinism -------------------------------------------- - - -def test_tp_determinism_across_independent_workers(): - """ - Two independently-instantiated workers (simulating two TP ranks) must - produce bit-identical accepted-token counts when fed the same call - sequence. This is the property whose absence used to hang TP>1. - """ - _require_cuda() - rank0 = _make_worker(2.6) - rank1 = _make_worker(2.6) - # Use a non-power-of-two batch to make accidental shape coincidences - # less likely to mask divergence. - for _ in range(32): - out0 = rank0._apply_force_accepted_tokens( - _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 - ) - out1 = rank1._apply_force_accepted_tokens( - _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 - ) - assert torch.equal(out0, out1) - - -def test_tp_determinism_survives_default_generator_drift(): - """ - Even if the default CUDA generator state diverges between ranks (e.g. - because some other component consumed random numbers), the synthetic - AR draws must still match. This is the regression guard for the - original bug, where ``torch.rand`` against the default generator gave - each rank a different sequence. - """ - _require_cuda() - rank0 = _make_worker(2.6) - rank1 = _make_worker(2.6) - - def _advance_default_generator(seed: int): - # Seed the default CUDA generator on the same device differently - # for each "rank" and consume some randoms, simulating drift. - torch.cuda.manual_seed(seed) - _ = torch.rand(128, device="cuda") - - _advance_default_generator(seed=11) - out0 = rank0._apply_force_accepted_tokens( - _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 - ) - _advance_default_generator(seed=999) - out1 = rank1._apply_force_accepted_tokens( - _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 - ) - assert torch.equal(out0, out1) - - -# ---------------- CUDA graph capture/replay --------------------------------- - - -def test_cuda_graph_capture_and_replay_match_eager(): - """ - Capturing the override into a CUDA graph and replaying it must produce - the exact same per-replay output sequence as running it eagerly. - """ - _require_cuda() - - n_iters = 8 - batch_size = 16 - runtime_draft_len = 4 - value = 2.6 - - # 1) Eager reference. - eager_worker = _make_worker(value) - eager_outputs = [] - for _ in range(n_iters): - out = eager_worker._apply_force_accepted_tokens( - _make_input(batch_size=batch_size), num_contexts=0, runtime_draft_len=runtime_draft_len - ) - eager_outputs.append(out.clone()) - - # 2) Captured-graph run, aligned to the eager reference. - graph_worker = _make_worker(value) - static_input = _make_input(batch_size=batch_size) - # Eager warmup: forces lazy RNG state allocation OUTSIDE capture. - graph_worker._apply_force_accepted_tokens( - static_input, num_contexts=0, runtime_draft_len=runtime_draft_len - ) - # Realign the device-side counter so the captured graph's first replay - # advances 0 → 1, matching the eager loop's first iteration. - graph_worker._force_accept_rng_counter.zero_() - - static_input.fill_(1) - static_output = torch.empty_like(static_input) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - graph_worker._apply_force_accepted_tokens( - static_input, num_contexts=0, runtime_draft_len=runtime_draft_len - ) - # ``_apply_force_accepted_tokens`` mutates ``static_input`` in place; - # mirror it into ``static_output`` so we can snapshot per replay. - static_output.copy_(static_input) - - graph_outputs = [] - for _ in range(n_iters): - # Reset input each replay so the override semantics match eager - # (which also starts from all-1s every iteration). - static_input.fill_(1) - graph.replay() - torch.cuda.synchronize() - graph_outputs.append(static_output.clone()) - - for i, (eager, graphed) in enumerate(zip(eager_outputs, graph_outputs)): - assert torch.equal(eager, graphed), ( - f"Iteration {i}: eager={eager.tolist()} vs graphed={graphed.tolist()}" - ) - - -def test_cuda_graph_replay_advances_rng_state(): - """ - Across replays, the captured counter must advance, producing different - (but deterministic) draws. A constant output across replays would mean - the counter increment was not captured. - """ - _require_cuda() - - worker = _make_worker(0.5) - batch_size = 64 - runtime_draft_len = 4 - - static_input = _make_input(batch_size=batch_size) - worker._apply_force_accepted_tokens( - static_input, num_contexts=0, runtime_draft_len=runtime_draft_len - ) - worker._force_accept_rng_counter.zero_() - - static_input.fill_(1) - static_output = torch.empty_like(static_input) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - worker._apply_force_accepted_tokens( - static_input, num_contexts=0, runtime_draft_len=runtime_draft_len - ) - static_output.copy_(static_input) - - snapshots = [] - for _ in range(8): - static_input.fill_(1) - graph.replay() - torch.cuda.synchronize() - snapshots.append(static_output.clone()) - - # At least two replays must differ; otherwise the captured counter is - # not advancing (the original bug-class). - assert any(not torch.equal(snapshots[0], s) for s in snapshots[1:]), ( - "Captured graph produced identical outputs on every replay — " - "RNG counter is not advancing inside the captured graph." - ) - - -if __name__ == "__main__": - pytest.main([__file__, "-v"]) - - -class _FailingSpecWorker(_StubSpecWorker): - """Stub whose forward saves spec-dec metadata state and then fails.""" - - def _forward_impl(self, *args: object, **kwargs: object) -> None: - attn_metadata = kwargs["attn_metadata"] - attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") - raise RuntimeError("simulated draft failure") - - -def test_forward_restores_spec_dec_state_on_failure() -> None: - """A failure between prepare_for_spec_dec and restore must not leak saved - attn-metadata state; SpecWorkerBase.forward restores it in its cleanup - (https://nvbugs/6442074).""" - _require_cuda() - attn_metadata = AttentionMetadata(max_num_requests=2, max_num_tokens=16) - attn_metadata.seq_lens = torch.ones(2, dtype=torch.int32) - worker = _FailingSpecWorker() - for _ in range(2): - # The second iteration would trip the pairing assert inside - # prepare_for_spec_dec if the first failure had leaked saved state. - with pytest.raises(RuntimeError, match="simulated draft failure"): - worker(attn_metadata=attn_metadata, spec_metadata=None) - assert not attn_metadata.has_spec_dec_saved_state diff --git a/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py b/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py deleted file mode 100644 index d43882ec9a83..000000000000 --- a/tests/unittest/_torch/speculative/test_group_all_greedy_sync.py +++ /dev/null @@ -1,83 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""CPU unit tests for the group-synchronized ``is_all_greedy_sample`` override. - -Under ADP + LM-head TP with rejection sampling, the greedy-vs-advanced path -choice gates group collectives, so the model engine all-gathers the per-rank -flags and stores the group AND in ``SpecMetadata.group_all_greedy_sample``; -``_scan_one_model_sampling`` must then re-apply it on every rescan (populate -runs after the CUDA graph key is built and would otherwise resurrect the -rank-local value). - -These tests call ``_scan_one_model_sampling`` unbound on a SimpleNamespace -stand-in, mirroring test_rejection_buffers_guard.py, so no GPU or full -SpecMetadata construction is needed. -""" - -import types - -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm._torch.speculative.interface import SpecMetadata - - -def _fake_request(temperature=None, top_k=None, top_p=None, slot=0): - return types.SimpleNamespace( - sampling_config=types.SimpleNamespace( - temperature=[temperature] if temperature is not None else None, - top_k=[top_k] if top_k is not None else None, - top_p=[top_p] if top_p is not None else None, - ), - state=LlmRequestState.GENERATION_IN_PROGRESS, - py_seq_slot=slot, - ) - - -def _fake_meta(group_all_greedy_sample=None, force_capture=False): - return types.SimpleNamespace( - runtime_draft_len=2, - dummy_slot_row=0, - group_all_greedy_sample=group_all_greedy_sample, - _force_non_greedy_for_capture=force_capture, - ) - - -def _scan(meta, requests): - return SpecMetadata._scan_one_model_sampling(meta, requests) - - -def test_local_value_used_when_no_group_sync(): - meta = _fake_meta(group_all_greedy_sample=None) - _scan(meta, [_fake_request(), _fake_request()]) - assert meta.is_all_greedy_sample is True - - _scan(meta, [_fake_request(), _fake_request(temperature=0.8)]) - assert meta.is_all_greedy_sample is False - - -def test_group_override_pulls_greedy_rank_onto_advanced_path(): - # (group value, capture force) -> final flag. This rank's batch is - # all-greedy in every case; the group AND (False) must win so the whole - # LM-head-TP group takes the advanced path together, group-wide True keeps - # greedy, and the warmup capture force composes without changing the - # outcome (the synced value is derived from capture-forced locals). - cases = [ - (False, False, False), - (True, False, True), - (False, True, False), - ] - for group_value, force_capture, expected in cases: - meta = _fake_meta(group_all_greedy_sample=group_value, - force_capture=force_capture) - _scan(meta, [_fake_request(), _fake_request()]) - assert meta.is_all_greedy_sample is expected, (group_value, - force_capture) - - -def test_group_override_survives_rescan(): - # populate_sampling_params_for_one_model rescans after the CUDA graph key - # is built; the override must keep applying so the key, the buffers, and - # the worker branches all agree. - meta = _fake_meta(group_all_greedy_sample=False) - for _ in range(3): - _scan(meta, [_fake_request()]) - assert meta.is_all_greedy_sample is False diff --git a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py b/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py deleted file mode 100644 index 82efadf6679f..000000000000 --- a/tests/unittest/_torch/speculative/test_rejection_buffers_guard.py +++ /dev/null @@ -1,270 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""GPU unit tests for the one-model rejection slot-buffer allocation and the -fail-closed acceptance guard. - -These exercise the pure buffer/guard logic with synthetic tensors and lightweight -attribute stand-ins (``types.SimpleNamespace``), so they do not construct a full -``SpecMetadata`` or run any model forward: - -- ``SpecMetadata.prepare_rejection_sampling_buffers`` only reads/writes ``self`` - attributes (``use_rejection_sampling``, ``vocab_size``, ``max_num_requests``, - ``max_draft_len``, ``draft_probs``, ``batch_slot_ids``, ``full_draft_probs``). -- ``SpecWorkerBase._rejection_buffers_valid`` only reads its arguments and - ``spec_metadata`` attributes (no ``self`` use), so it can be called unbound. - -Requires CUDA because the buffers are allocated on ``device='cuda'``. -""" - -import types - -import pytest -import torch - -from tensorrt_llm._torch.speculative.interface import SpecMetadata, SpecWorkerBase - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available(), reason="rejection buffers are CUDA tensors" -) - -R, K, V = 8, 4, 32 # max_num_requests, max_draft_len, vocab_size - - -def _alloc_meta(**over): - base = dict( - use_rejection_sampling=True, - draft_probs=None, - vocab_size=V, - draft_vocab_size=V, - max_num_requests=R, - max_draft_len=K, - num_seq_slots=0, - batch_slot_ids=None, - full_draft_probs=None, - draft_probs_vocab_size=0, - ) - base.update(over) - return types.SimpleNamespace(**base) - - -def test_prepare_buffers_allocates_when_enabled(): - # Shared draft/target vocab: draft_probs + batch_slot_ids are allocated; - # full_draft_probs is skipped (only needed when the vocabs differ). - m = _alloc_meta() - SpecMetadata.prepare_rejection_sampling_buffers(m) - assert m.draft_probs is not None and tuple(m.draft_probs.shape) == (R + 1, K, V) - assert m.batch_slot_ids is not None and m.batch_slot_ids.shape[0] == R - assert m.batch_slot_ids.dtype == torch.long - assert m.full_draft_probs is None - assert m.draft_probs_vocab_size == V - - # Distinct draft vocab: full_draft_probs (d2t-expanded) is allocated. - m = _alloc_meta(draft_vocab_size=V - 1) - SpecMetadata.prepare_rejection_sampling_buffers(m) - assert m.full_draft_probs is not None and tuple(m.full_draft_probs.shape) == (R + 1, K, V) - - -def test_prepare_buffers_span_seq_slot_pool(): - # Under DeepSeek-V4 overlap scheduling the SeqSlotManager pool - # (num_seq_slots) can exceed max_num_requests; py_seq_slot then indexes past - # max_num_requests. Slot-indexed buffers must span the full pool plus the - # dummy scratch row, and dummy_slot_row must land on that last row so a real - # slot never aliases the dummy row. - pool = 2 * R - m = _alloc_meta(num_seq_slots=pool) - SpecMetadata.prepare_rejection_sampling_buffers(m) - assert tuple(m.draft_probs.shape) == (pool + 1, K, V) - assert m.dummy_slot_row == pool - - -def test_prepare_buffers_noop_when_disabled(): - m = _alloc_meta(use_rejection_sampling=False) - SpecMetadata.prepare_rejection_sampling_buffers(m) - assert m.draft_probs is None - assert m.batch_slot_ids is None - assert m.full_draft_probs is None - - -def _valid_state(): - return types.SimpleNamespace( - draft_probs=torch.empty((R, K, V), device="cuda"), - batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), - is_ragged_verify=False, - ) - - -# num_contexts=0, num_gens=4 -_NUM_CTX = 0 -_BATCH = 4 -_NUM_GENS = _BATCH - _NUM_CTX - - -def _good_args(): - draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") - logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") - return draft_tokens, K, V, _NUM_CTX, _BATCH, logits - - -_DUMMY = object() # _rejection_buffers_valid does not use self - - -def _call(meta, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits): - return SpecWorkerBase._rejection_buffers_valid( - _DUMMY, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits, meta - ) - - -def test_guard_true_on_valid_state(): - assert _call(_valid_state(), *_good_args()) is True - - -# Each builder returns (meta, *_call args) for one malformed state that the -# fail-closed guard must reject. Kept as a single parametrized test so the guard's -# rejection branches read as one table. -def _case_draft_probs_missing(): - m = _valid_state() - m.draft_probs = None - return (m, *_good_args()) - - -def _case_batch_slot_ids_missing(): - m = _valid_state() - m.batch_slot_ids = None - return (m, *_good_args()) - - -def _case_stored_vocab_nonpositive(): - dt, dl, _sv, nc, bs, lg = _good_args() - return (_valid_state(), dt, dl, 0, nc, bs, lg) - - -def _case_stored_vocab_exceeds_buffer(): - dt, dl, _sv, nc, bs, lg = _good_args() - return (_valid_state(), dt, dl, V + 1, nc, bs, lg) - - -def _case_draft_len_exceeds_buffer(): - # draft_probs has K steps; ask for K+1. - draft_tokens = torch.zeros((_NUM_GENS, K + 1), dtype=torch.int, device="cuda") - logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 2), V), device="cuda") - return (_valid_state(), draft_tokens, K + 1, V, _NUM_CTX, _BATCH, logits) - - -def _case_draft_tokens_wrong_rows(): - dt = torch.zeros((_NUM_GENS + 1, K), dtype=torch.int, device="cuda") - _, dl, sv, nc, bs, lg = _good_args() - return (_valid_state(), dt, dl, sv, nc, bs, lg) - - -def _case_too_few_logits_rows(): - dt, dl, sv, nc, bs, _lg = _good_args() - too_few = torch.zeros((nc + _NUM_GENS, V), device="cuda") # missing +1 each - return (_valid_state(), dt, dl, sv, nc, bs, too_few) - - -def _case_batch_slot_ids_too_short(): - m = _valid_state() - m.batch_slot_ids = torch.arange(_BATCH - 1, device="cuda", dtype=torch.long) - return (m, *_good_args()) - - -@pytest.mark.parametrize( - "build_case", - [ - _case_draft_probs_missing, - _case_batch_slot_ids_missing, - _case_stored_vocab_nonpositive, - _case_stored_vocab_exceeds_buffer, - _case_draft_len_exceeds_buffer, - _case_draft_tokens_wrong_rows, - _case_too_few_logits_rows, - _case_batch_slot_ids_too_short, - ], - ids=lambda f: f.__name__.removeprefix("_case_"), -) -def test_guard_false_on_malformed_state(build_case): - assert _call(*build_case()) is False - - -# -------------------------------------------------------------------------- -# Fail-closed acceptance dispatch: prove _accept_draft_tokens() routes to -# strict/base acceptance when the buffers are malformed, and routes to the -# rejection method when the state is valid. The two acceptance methods are -# stubbed to record which path ran; the real _can_use_rejection_sampling and -# _rejection_buffers_valid are exercised. -# -------------------------------------------------------------------------- - - -class _Worker(SpecWorkerBase): - @property - def max_draft_len(self) -> int: - return K - - def _forward_impl(self, *args: object, **kwargs: object) -> None: - raise NotImplementedError - - -def _dispatch_meta(**over): - base = dict( - use_rejection_sampling=True, - is_all_greedy_sample=False, - is_ragged_verify=False, - draft_probs_vocab_size=V, - draft_probs_last_dim=V, - batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), - draft_probs=torch.zeros((R, K, V), device="cuda"), - ) - base.update(over) - return types.SimpleNamespace(**base) - - -def _make_worker(): - w = _Worker() - calls = {"base": 0, "rejection": 0} - - def _base(logits, draft_tokens, num_contexts, batch_size, spec_metadata): - calls["base"] += 1 - return ("base", None) - - def _rej(logits, draft_tokens, draft_probs, num_contexts, batch_size, spec_metadata): - calls["rejection"] += 1 - return ("rejection", None) - - w._sample_and_accept_draft_tokens_base = _base - w._sample_and_accept_draft_tokens_rejection = _rej - return w, calls - - -@pytest.mark.parametrize( - "meta_over, expected_path", - [ - # Malformed buffers (draft_probs None): _rejection_buffers_valid must - # return False, so acceptance fails closed to base and the rejection - # kernel is skipped. - (dict(draft_probs=None), "base"), - # Valid buffers: the rejection path runs. - (dict(), "rejection"), - # All-greedy batch: rejection is bypassed regardless of buffers. - (dict(is_all_greedy_sample=True), "base"), - ], - ids=[ - "fails_closed_on_malformed_buffers", - "routes_to_rejection_on_valid_state", - "base_when_all_greedy", - ], -) -def test_accept_dispatch(meta_over, expected_path): - w, calls = _make_worker() - meta = _dispatch_meta(**meta_over) - draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") - logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") - out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) - assert out == (expected_path, None) - other = "rejection" if expected_path == "base" else "base" - assert calls[expected_path] == 1 and calls[other] == 0 - - -if __name__ == "__main__": - import sys - - sys.exit(pytest.main([__file__, "-v"])) diff --git a/tests/unittest/_torch/speculative/test_spec_sampler.py b/tests/unittest/_torch/speculative/test_spec_sampler.py new file mode 100644 index 000000000000..055a0bd981c7 --- /dev/null +++ b/tests/unittest/_torch/speculative/test_spec_sampler.py @@ -0,0 +1,650 @@ +# 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. +"""Unit tests for the one-model spec-decode sampler/acceptance framework: +forced acceptance (``TLLM_SPEC_DECODE_FORCE_NUM_ACCEPTED_TOKENS``), +rejection-sampling buffer allocation with its fail-closed guard, and the +group-synchronized ``is_all_greedy_sample`` override.""" + +import os +import types +from typing import Optional +from unittest.mock import patch + +import pytest +import torch + +from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.speculative.interface import ( + FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR, + SpecMetadata, + SpecWorkerBase, + get_force_num_accepted_tokens, + get_force_num_accepted_tokens_float, +) + +_cuda_only = pytest.mark.skipif( + not torch.cuda.is_available(), reason="rejection buffers are CUDA tensors" +) + + +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: + worker = _StubSpecWorker() + if value is not None: + worker.force_num_accepted_tokens = value + return worker + + +def _make_input(batch_size: int, device: str = "cuda") -> torch.Tensor: + # Initial counts: 1 (target token, no draft accepts yet). The forced + # override only ever rewrites entries from ``num_contexts:`` onward. + return torch.ones(batch_size, dtype=torch.int32, device=device) + + +def _require_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required for this test.") + + +# ---------------- env-var parsing helpers ----------------------------------- + + +@pytest.mark.parametrize( + "env_value, expected_float, expected_int", + [ + ("0", 0.0, 0), + ("3", 3.0, 3), + ("2.6", 2.6, 0), + ("0.5", 0.5, 0), + ("not-a-number", 0.0, 0), + ], +) +def test_get_force_num_accepted_tokens_float(env_value, expected_float, expected_int): + """The float helper parses fractional rates; the int helper (used by the + 2-model path) must keep its original behavior on the same values.""" + with patch.dict(os.environ, {FORCE_NUM_ACCEPTED_TOKENS_ENV_VAR: env_value}): + assert get_force_num_accepted_tokens_float() == pytest.approx(expected_float) + assert get_force_num_accepted_tokens() == expected_int + + +# ---------------- _apply_force_accepted_tokens semantics -------------------- + + +def test_zero_value_is_noop(): + """A zero forced rate must leave inputs and RNG state untouched.""" + _require_cuda() + worker = _make_worker(0.0) + base = _make_input(batch_size=4) + before = base.clone() + out = worker._apply_force_accepted_tokens(base, num_contexts=0, runtime_draft_len=4) + assert torch.equal(out, before) + # No RNG state should have been touched in the early-exit path. + assert worker._force_accept_rng_pool is None + assert worker._force_accept_rng_counter is None + + +@pytest.mark.parametrize( + "value, runtime_draft_len", + [ + (1.0, 4), + (2.0, 4), + (4.0, 4), + # 10 draft tokens requested but only 2 available -> capped at 3. + (10.0, 2), + # Fractional but no room: max_total = 3, base_total = min(10, 3) = 3, + # so ``base_total < max_total`` is False -> all 3, no RNG draw. + (9.5, 2), + ], +) +def test_integer_value_matches_legacy_behavior(value, runtime_draft_len): + """Integer (or capped) rates take the deterministic path with no RNG pool.""" + _require_cuda() + worker = _make_worker(value) + out = worker._apply_force_accepted_tokens( + _make_input(batch_size=8), num_contexts=0, runtime_draft_len=runtime_draft_len + ) + expected = min(int(value) + 1, runtime_draft_len + 1) + assert torch.all(out == expected) + # The deterministic (non-random) path must not allocate the RNG pool. + assert worker._force_accept_rng_pool is None + + +def test_fractional_distribution_matches_target_probability(): + """Fractional draws emit only base/base+1 and average to the target rate.""" + _require_cuda() + worker = _make_worker(2.6) + target_frac = 0.6 + n_iters = 200 + batch = 64 + seen = set() + extra_count = 0 + total = 0 + for _ in range(n_iters): + out = worker._apply_force_accepted_tokens( + _make_input(batch_size=batch), num_contexts=0, runtime_draft_len=4 + ) + seen.update(out.unique().tolist()) + extra_count += int((out == 4).sum().item()) + total += batch + # Either 2 draft + target = 3, or 3 draft + target = 4 -- never anything else. + assert seen == {3, 4} + measured = extra_count / total + assert abs(measured - target_frac) < 0.03, ( + f"measured fraction {measured:.4f} differs from target {target_frac}" + ) + + +def test_num_contexts_offset_is_respected(): + """Context (prefill) rows must be left untouched by the override.""" + _require_cuda() + worker = _make_worker(2.6) + batch_size = 8 + num_contexts = 3 + base = _make_input(batch_size=batch_size) + sentinel = torch.iinfo(base.dtype).max + base[:num_contexts] = sentinel + out = worker._apply_force_accepted_tokens(base, num_contexts=num_contexts, runtime_draft_len=4) + assert torch.all(out[:num_contexts] == sentinel) + assert torch.all((out[num_contexts:] == 3) | (out[num_contexts:] == 4)) + + +# ---------------- TP determinism -------------------------------------------- + + +def test_tp_determinism_across_independent_workers(): + """Two independently-instantiated workers (simulating two TP ranks) must + produce bit-identical accepted counts for the same call sequence.""" + _require_cuda() + rank0 = _make_worker(2.6) + rank1 = _make_worker(2.6) + # Use a non-power-of-two batch to make accidental shape coincidences + # less likely to mask divergence. + for _ in range(32): + out0 = rank0._apply_force_accepted_tokens( + _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 + ) + out1 = rank1._apply_force_accepted_tokens( + _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 + ) + assert torch.equal(out0, out1) + + +def test_tp_determinism_survives_default_generator_drift(): + """Synthetic AR draws must still match across ranks even when the default + CUDA generator state has diverged between them.""" + _require_cuda() + rank0 = _make_worker(2.6) + rank1 = _make_worker(2.6) + + def _advance_default_generator(seed: int): + # Seed the default CUDA generator on the same device differently + # for each "rank" and consume some randoms, simulating drift. + torch.cuda.manual_seed(seed) + _ = torch.rand(128, device="cuda") + + _advance_default_generator(seed=11) + out0 = rank0._apply_force_accepted_tokens( + _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 + ) + _advance_default_generator(seed=999) + out1 = rank1._apply_force_accepted_tokens( + _make_input(batch_size=33), num_contexts=0, runtime_draft_len=4 + ) + assert torch.equal(out0, out1) + + +# ---------------- CUDA graph capture/replay --------------------------------- + + +def test_cuda_graph_capture_and_replay_match_eager(): + """Capturing the override into a CUDA graph and replaying it must produce + the exact same per-replay output sequence as running it eagerly.""" + _require_cuda() + + n_iters = 8 + batch_size = 16 + runtime_draft_len = 4 + value = 2.6 + + # 1) Eager reference. + eager_worker = _make_worker(value) + eager_outputs = [] + for _ in range(n_iters): + out = eager_worker._apply_force_accepted_tokens( + _make_input(batch_size=batch_size), num_contexts=0, runtime_draft_len=runtime_draft_len + ) + eager_outputs.append(out.clone()) + + # 2) Captured-graph run, aligned to the eager reference. + graph_worker = _make_worker(value) + static_input = _make_input(batch_size=batch_size) + # Eager warmup: forces lazy RNG state allocation OUTSIDE capture. + graph_worker._apply_force_accepted_tokens( + static_input, num_contexts=0, runtime_draft_len=runtime_draft_len + ) + # Realign the device-side counter so the captured graph's first replay + # advances 0 → 1, matching the eager loop's first iteration. + graph_worker._force_accept_rng_counter.zero_() + + static_input.fill_(1) + static_output = torch.empty_like(static_input) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_worker._apply_force_accepted_tokens( + static_input, num_contexts=0, runtime_draft_len=runtime_draft_len + ) + # ``_apply_force_accepted_tokens`` mutates ``static_input`` in place; + # mirror it into ``static_output`` so we can snapshot per replay. + static_output.copy_(static_input) + + graph_outputs = [] + for _ in range(n_iters): + # Reset input each replay so the override semantics match eager + # (which also starts from all-1s every iteration). + static_input.fill_(1) + graph.replay() + torch.cuda.synchronize() + graph_outputs.append(static_output.clone()) + + for i, (eager, graphed) in enumerate(zip(eager_outputs, graph_outputs)): + assert torch.equal(eager, graphed), ( + f"Iteration {i}: eager={eager.tolist()} vs graphed={graphed.tolist()}" + ) + + +def test_cuda_graph_replay_advances_rng_state(): + """Across replays the captured RNG counter must advance, producing + different (but deterministic) draws rather than a constant output.""" + _require_cuda() + + worker = _make_worker(0.5) + batch_size = 64 + runtime_draft_len = 4 + + static_input = _make_input(batch_size=batch_size) + worker._apply_force_accepted_tokens( + static_input, num_contexts=0, runtime_draft_len=runtime_draft_len + ) + worker._force_accept_rng_counter.zero_() + + static_input.fill_(1) + static_output = torch.empty_like(static_input) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + worker._apply_force_accepted_tokens( + static_input, num_contexts=0, runtime_draft_len=runtime_draft_len + ) + static_output.copy_(static_input) + + snapshots = [] + for _ in range(8): + static_input.fill_(1) + graph.replay() + torch.cuda.synchronize() + snapshots.append(static_output.clone()) + + # At least two replays must differ; otherwise the captured counter is + # not advancing (the original bug-class). + assert any(not torch.equal(snapshots[0], s) for s in snapshots[1:]), ( + "Captured graph produced identical outputs on every replay — " + "RNG counter is not advancing inside the captured graph." + ) + + +class _FailingSpecWorker(_StubSpecWorker): + """Stub whose forward saves spec-dec metadata state and then fails.""" + + def _forward_impl(self, *args: object, **kwargs: object) -> None: + attn_metadata = kwargs["attn_metadata"] + attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda") + raise RuntimeError("simulated draft failure") + + +def test_forward_restores_spec_dec_state_on_failure() -> None: + """A failure between prepare_for_spec_dec and restore must not leak saved + attn-metadata state; SpecWorkerBase.forward restores it in its cleanup.""" + _require_cuda() + attn_metadata = AttentionMetadata(max_num_requests=2, max_num_tokens=16) + attn_metadata.seq_lens = torch.ones(2, dtype=torch.int32) + worker = _FailingSpecWorker() + for _ in range(2): + # The second iteration would trip the pairing assert inside + # prepare_for_spec_dec if the first failure had leaked saved state. + with pytest.raises(RuntimeError, match="simulated draft failure"): + worker(attn_metadata=attn_metadata, spec_metadata=None) + assert not attn_metadata.has_spec_dec_saved_state + + +# ---------------- rejection slot buffers + fail-closed guard ---------------- + +R, K, V = 8, 4, 32 # max_num_requests, max_draft_len, vocab_size + + +def _alloc_meta(**over): + base = dict( + use_rejection_sampling=True, + draft_probs=None, + vocab_size=V, + draft_vocab_size=V, + max_num_requests=R, + max_draft_len=K, + num_seq_slots=0, + batch_slot_ids=None, + full_draft_probs=None, + draft_probs_vocab_size=0, + ) + base.update(over) + return types.SimpleNamespace(**base) + + +@_cuda_only +def test_prepare_buffers_allocates_when_enabled(): + """When rejection sampling is enabled, slot buffers are allocated; + full_draft_probs appears only when draft and target vocabs differ.""" + m = _alloc_meta() + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert m.draft_probs is not None and tuple(m.draft_probs.shape) == (R + 1, K, V) + assert m.batch_slot_ids is not None and m.batch_slot_ids.shape[0] == R + assert m.batch_slot_ids.dtype == torch.long + assert m.full_draft_probs is None + assert m.draft_probs_vocab_size == V + + # Distinct draft vocab: full_draft_probs (d2t-expanded) is allocated. + m = _alloc_meta(draft_vocab_size=V - 1) + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert m.full_draft_probs is not None and tuple(m.full_draft_probs.shape) == (R + 1, K, V) + + +@_cuda_only +def test_prepare_buffers_span_seq_slot_pool(): + """Slot-indexed buffers must span the full seq-slot pool (which can exceed + max_num_requests) plus a dummy scratch row placed on the last row.""" + pool = 2 * R + m = _alloc_meta(num_seq_slots=pool) + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert tuple(m.draft_probs.shape) == (pool + 1, K, V) + assert m.dummy_slot_row == pool + + +@_cuda_only +def test_prepare_buffers_noop_when_disabled(): + """With rejection sampling disabled, no buffers are allocated.""" + m = _alloc_meta(use_rejection_sampling=False) + SpecMetadata.prepare_rejection_sampling_buffers(m) + assert m.draft_probs is None + assert m.batch_slot_ids is None + assert m.full_draft_probs is None + + +def _valid_state(): + return types.SimpleNamespace( + draft_probs=torch.empty((R, K, V), device="cuda"), + batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), + is_ragged_verify=False, + ) + + +# num_contexts=0, num_gens=4 +_NUM_CTX = 0 +_BATCH = 4 +_NUM_GENS = _BATCH - _NUM_CTX + + +def _good_args(): + draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") + logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") + return draft_tokens, K, V, _NUM_CTX, _BATCH, logits + + +_DUMMY = object() # _rejection_buffers_valid does not use self + + +def _call(meta, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits): + return SpecWorkerBase._rejection_buffers_valid( + _DUMMY, draft_tokens, draft_len, stored_vocab, num_contexts, batch_size, logits, meta + ) + + +@_cuda_only +def test_guard_true_on_valid_state(): + """The guard accepts a well-formed buffer/argument combination.""" + assert _call(_valid_state(), *_good_args()) is True + + +# Each builder returns (meta, *_call args) for one malformed state that the +# fail-closed guard must reject. +def _case_draft_probs_missing(): + m = _valid_state() + m.draft_probs = None + return (m, *_good_args()) + + +def _case_batch_slot_ids_missing(): + m = _valid_state() + m.batch_slot_ids = None + return (m, *_good_args()) + + +def _case_stored_vocab_nonpositive(): + dt, dl, _sv, nc, bs, lg = _good_args() + return (_valid_state(), dt, dl, 0, nc, bs, lg) + + +def _case_stored_vocab_exceeds_buffer(): + dt, dl, _sv, nc, bs, lg = _good_args() + return (_valid_state(), dt, dl, V + 1, nc, bs, lg) + + +def _case_draft_len_exceeds_buffer(): + # draft_probs has K steps; ask for K+1. + draft_tokens = torch.zeros((_NUM_GENS, K + 1), dtype=torch.int, device="cuda") + logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 2), V), device="cuda") + return (_valid_state(), draft_tokens, K + 1, V, _NUM_CTX, _BATCH, logits) + + +def _case_draft_tokens_wrong_rows(): + dt = torch.zeros((_NUM_GENS + 1, K), dtype=torch.int, device="cuda") + _, dl, sv, nc, bs, lg = _good_args() + return (_valid_state(), dt, dl, sv, nc, bs, lg) + + +def _case_too_few_logits_rows(): + dt, dl, sv, nc, bs, _lg = _good_args() + too_few = torch.zeros((nc + _NUM_GENS, V), device="cuda") # missing +1 each + return (_valid_state(), dt, dl, sv, nc, bs, too_few) + + +def _case_batch_slot_ids_too_short(): + m = _valid_state() + m.batch_slot_ids = torch.arange(_BATCH - 1, device="cuda", dtype=torch.long) + return (m, *_good_args()) + + +@_cuda_only +@pytest.mark.parametrize( + "build_case", + [ + _case_draft_probs_missing, + _case_batch_slot_ids_missing, + _case_stored_vocab_nonpositive, + _case_stored_vocab_exceeds_buffer, + _case_draft_len_exceeds_buffer, + _case_draft_tokens_wrong_rows, + _case_too_few_logits_rows, + _case_batch_slot_ids_too_short, + ], + ids=lambda f: f.__name__.removeprefix("_case_"), +) +def test_guard_false_on_malformed_state(build_case): + """The fail-closed guard rejects every malformed buffer/argument state.""" + assert _call(*build_case()) is False + + +# Fail-closed acceptance dispatch: _accept_draft_tokens() must route to +# strict/base acceptance when the buffers are malformed and to the rejection +# method when the state is valid; the acceptance methods are stubbed to +# record which path ran while the real guard logic is exercised. + + +class _Worker(SpecWorkerBase): + @property + def max_draft_len(self) -> int: + return K + + def _forward_impl(self, *args: object, **kwargs: object) -> None: + raise NotImplementedError + + +def _dispatch_meta(**over): + base = dict( + use_rejection_sampling=True, + is_all_greedy_sample=False, + is_ragged_verify=False, + draft_probs_vocab_size=V, + draft_probs_last_dim=V, + batch_slot_ids=torch.arange(R, device="cuda", dtype=torch.long), + draft_probs=torch.zeros((R, K, V), device="cuda"), + ) + base.update(over) + return types.SimpleNamespace(**base) + + +def _make_dispatch_worker(): + w = _Worker() + calls = {"base": 0, "rejection": 0} + + def _base(logits, draft_tokens, num_contexts, batch_size, spec_metadata): + calls["base"] += 1 + return ("base", None) + + def _rej(logits, draft_tokens, draft_probs, num_contexts, batch_size, spec_metadata): + calls["rejection"] += 1 + return ("rejection", None) + + w._sample_and_accept_draft_tokens_base = _base + w._sample_and_accept_draft_tokens_rejection = _rej + return w, calls + + +@_cuda_only +@pytest.mark.parametrize( + "meta_over, expected_path", + [ + # Malformed buffers (draft_probs None): _rejection_buffers_valid must + # return False, so acceptance fails closed to base and the rejection + # kernel is skipped. + (dict(draft_probs=None), "base"), + # Valid buffers: the rejection path runs. + (dict(), "rejection"), + # All-greedy batch: rejection is bypassed regardless of buffers. + (dict(is_all_greedy_sample=True), "base"), + ], + ids=[ + "fails_closed_on_malformed_buffers", + "routes_to_rejection_on_valid_state", + "base_when_all_greedy", + ], +) +def test_accept_dispatch(meta_over, expected_path): + """Acceptance dispatch fails closed to base on bad state and routes to + the rejection path only when the buffers are valid and not all-greedy.""" + w, calls = _make_dispatch_worker() + meta = _dispatch_meta(**meta_over) + draft_tokens = torch.zeros((_NUM_GENS, K), dtype=torch.int, device="cuda") + logits = torch.zeros((_NUM_CTX + _NUM_GENS * (K + 1), V), device="cuda") + out = w._accept_draft_tokens(logits, draft_tokens, _NUM_CTX, _BATCH, meta) + assert out == (expected_path, None) + other = "rejection" if expected_path == "base" else "base" + assert calls[expected_path] == 1 and calls[other] == 0 + + +# ---------------- group-synchronized is_all_greedy_sample (CPU) ------------- + + +def _fake_request(temperature=None, top_k=None, top_p=None, slot=0): + return types.SimpleNamespace( + sampling_config=types.SimpleNamespace( + temperature=[temperature] if temperature is not None else None, + top_k=[top_k] if top_k is not None else None, + top_p=[top_p] if top_p is not None else None, + ), + state=LlmRequestState.GENERATION_IN_PROGRESS, + py_seq_slot=slot, + ) + + +def _fake_meta(group_all_greedy_sample=None, force_capture=False): + return types.SimpleNamespace( + runtime_draft_len=2, + dummy_slot_row=0, + group_all_greedy_sample=group_all_greedy_sample, + _force_non_greedy_for_capture=force_capture, + ) + + +def _scan(meta, requests): + return SpecMetadata._scan_one_model_sampling(meta, requests) + + +def test_local_value_used_when_no_group_sync(): + """Without a group-synced value, the scan uses the rank-local flag.""" + meta = _fake_meta(group_all_greedy_sample=None) + _scan(meta, [_fake_request(), _fake_request()]) + assert meta.is_all_greedy_sample is True + + _scan(meta, [_fake_request(), _fake_request(temperature=0.8)]) + assert meta.is_all_greedy_sample is False + + +def test_group_override_pulls_greedy_rank_onto_advanced_path(): + """The group AND overrides the rank-local all-greedy flag, and composes + with the warmup capture force without changing the outcome.""" + cases = [ + (False, False, False), + (True, False, True), + (False, True, False), + ] + for group_value, force_capture, expected in cases: + meta = _fake_meta(group_all_greedy_sample=group_value, + force_capture=force_capture) + _scan(meta, [_fake_request(), _fake_request()]) + assert meta.is_all_greedy_sample is expected, (group_value, + force_capture) + + +def test_group_override_survives_rescan(): + """The group override must keep applying on every rescan after the CUDA + graph key is built.""" + meta = _fake_meta(group_all_greedy_sample=False) + for _ in range(3): + _scan(meta, [_fake_request()]) + assert meta.is_all_greedy_sample is False + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"])) From 0d000650b8d0016098df2cb3756d76efb772a92b Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 01:59:45 -0700 Subject: [PATCH 151/177] [None][doc] Alignment run on their workload: no gain for scheduled trim (-1~-6%), pin3 -15%; workload-dependence established Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index c9226e35b19e..c34f1b186bd5 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -638,6 +638,21 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( `test_gsm8k_dep8_ragged_verify[moe_backend=TRTLLM]`(DEP8 + overlap + CUDA graph + ragged verify + 置信度调度),`TLLM_DSPARK_SPS_TABLE=postfix_pro_table.json` 注入修复后真实表(planner 主动裁剪而非退化均匀),**1 passed(27 分钟)——精度 ≥ 记录参考值,断言通过**。终局判决的精度一环闭合。注:MEGAMOE_DEEPGEMM 变体在该节点因 DeepGEMM JIT "NVCC compilation failed" 无法初始化(nvcc 在位、/tmp 空,原因未深究;TRTLLM 变体即吞吐判决所用后端,更具代表性)。 +### 对齐实验:zhaoyuanh 的 throughput_1k 负载 × 我们的栈(08-11 02:10,主矩阵完成) + +**协议**(与 roofline 报告对齐,按用户修正):throughput_1k parquet 前 1024 行 file 序(源 parquet,SHA 与其冻结 dump 不同,已标注)、greedy temp=0、max_tokens=64、**持续满并发**(完成即补位,计量段=256 个 ramp 完成后的 1024 个完成);**每 cell 独立重启且 maxbs=local bs**;max_seq_len=8192 显式;成本表与 STS **均在本负载重采**(表:burst 模式采纯 gen 段,bs32/64 实测、128/256 α 外推;STS:142k 样本,位置接受率 [0.70,0.46,0.30,0.20,0.15],温度 ~2 → 裸分数显著过自信,ECE 5/5 位置改善)。 + +**主矩阵(中位数 tok/s,n=3/cell)**: + +| local bs | notrim | sched(表+STS) | Δ | pin3 参照 | +|---|---|---|---|---| +| 32 | 1,850 | 1,734 | −6.3% | — | +| 64 | 2,988 | 2,858 | −4.4% | — | +| 128 | 4,336 | 4,284 | −1.2% | 4,227(−2.5%)| +| 256 | 7,216 | 7,112 | −1.4% | 6,129(**−15.1%**)| + +**判读**:①此负载上调度裁剪无收益(−1~−6%),盲砍(pin3)在大 bs 重亏——**与 poetry/arena 的 +5.9~+24.7% 并立:裁剪收益是负载性质的函数**;②机理:64 输出 + 长 prompt + 持续补位 → 几乎每步都是 prefill 主导的混合步,纯 gen 步采的成本表在混合步里**高估 θ 节省**,被裁的接受 token 却真实损失;③他们 +13.29% 的天花板在我们栈的混合步现实中不兑现,落在其自身敏感性分析预言的失效区(静态代理 + 紧凑执行假设);其"DL3 优于 DL5"在我们 runtime 亦反向(诚实支付 D(5) 的 pin3 −15%);④**产品方向**:planner 需要"无标的自保"——混合步密集时降权或停裁(表纳入 ctx 混合度维度)。no-STS 消融进行中(分离校准效应)。 + ### 产品线终局判决:Pro@DEP8 置信度调度全面转正(08-10 13:50) Flash 三臂翻案后,Pro@DEP8 终局(maxbs=256 解锁满批 + `postfix_pro_table.json` 新表,scheduled vs notrim,双节点臂序对调,每 cell n=10): From 9c493012c0ebb23334b11244c776e288f60efc15 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:38:20 -0700 Subject: [PATCH 152/177] [None][doc] Disagg campaign: environment-trap matrix (pmix, NIXL transfer, DEFAULT-vs-v2); round 4 = NIXL+PYTHON with UCX env Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index c34f1b186bd5..cc50cb13c304 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -638,6 +638,19 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( `test_gsm8k_dep8_ragged_verify[moe_backend=TRTLLM]`(DEP8 + overlap + CUDA graph + ragged verify + 置信度调度),`TLLM_DSPARK_SPS_TABLE=postfix_pro_table.json` 注入修复后真实表(planner 主动裁剪而非退化均匀),**1 passed(27 分钟)——精度 ≥ 记录参考值,断言通过**。终局判决的精度一环闭合。注:MEGAMOE_DEEPGEMM 变体在该节点因 DeepGEMM JIT "NVCC compilation failed" 无法初始化(nvcc 在位、/tmp 空,原因未深究;TRTLLM 变体即吞吐判决所用后端,更具代表性)。 +### Disagg 战役:环境坑编年史(08-11 06:00,第 4 轮进行中) + +用户转向 disagg 形态(只测 gen server 吞吐;纯 gen 步 = 裁剪最有利形态,并绕开 agg 的 ctx 准入带宽上限)。协议:每 (臂,local bs) cell = 独立节点对 + 全新 ctx/gen/router,gen maxbs=当档,客户端超发 16×LBS 钉满,3×2 矩阵({32,64,128},256 搁置)。 + +**环境坑矩阵(全部非本 PR 代码问题):** + +| 轮 | 配置 | 结局 | 教训 | +|---|---|---|---| +| 08-07(#29 时代)| NIXL+PYTHON,无 env | gen 死于 #23 修复前的 IMA,传输从未被行使 | 今日 gen 全程健康 → **#23 修复在 disagg 路径成立** | +| R1/R2(今日)| NIXL+PYTHON,无 env | ①router 忘 `--mpi=pmix` 秒死;②修后 ctx 侧 NIXL `TransferState.FAILURE` → gen 等 KV 超时自杀 | 任何 trtllm CLI 经 srun 都要 pmix(已入记忆) | +| R3 | DEFAULT + UCX 三件套(unset UCX_TLS/关 UCC/RNDV_RAILS=2,抄自同集群日跑 harness)| ctx 启动死于 **v2 KV 管理器 CUDA_ERROR_INVALID_CONTEXT** | DEFAULT(C++ 传输)与 v2 管理器不兼容;p25 时代选 NIXL+**PYTHON runtime** 正是 v2 兼容路径 | +| **R4(进行中)** | **NIXL+PYTHON + UCX 三件套**(未试组合:NIXL 底层即 UCX,三件套正治 UCX 传输)| — | 若仍 FAILURE,下一杆:UCX_NET_DEVICES 按 GPU 就近 HCA 钉扎(同事 harness 的 CTX_LOCAL_HCA_PIN)| + ### 对齐实验:zhaoyuanh 的 throughput_1k 负载 × 我们的栈(08-11 02:10,主矩阵完成) **协议**(与 roofline 报告对齐,按用户修正):throughput_1k parquet 前 1024 行 file 序(源 parquet,SHA 与其冻结 dump 不同,已标注)、greedy temp=0、max_tokens=64、**持续满并发**(完成即补位,计量段=256 个 ramp 完成后的 1024 个完成);**每 cell 独立重启且 maxbs=local bs**;max_seq_len=8192 显式;成本表与 STS **均在本负载重采**(表:burst 模式采纯 gen 段,bs32/64 实测、128/256 α 外推;STS:142k 样本,位置接受率 [0.70,0.46,0.30,0.20,0.15],温度 ~2 → 裸分数显著过自信,ECE 5/5 位置改善)。 From 66c227fea1495261890754193a42978fe6108041 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:30:34 -0700 Subject: [PATCH 153/177] [None][doc] Deployment rule for the maxbs/2 windowed-pool cap (explicit max_seq_len) and reproducible benchmark configs for both workloads + disagg Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 8 +- .../features/dspark-confidence-scheduling.md | 89 +++++++++++++++++++ 2 files changed, 96 insertions(+), 1 deletion(-) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index cc50cb13c304..9de4813c5749 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -649,7 +649,13 @@ Flash DEP4,修复 b87ea28bf3 + 新表 `postfix_flash_table.json`,每 cell n=10( | 08-07(#29 时代)| NIXL+PYTHON,无 env | gen 死于 #23 修复前的 IMA,传输从未被行使 | 今日 gen 全程健康 → **#23 修复在 disagg 路径成立** | | R1/R2(今日)| NIXL+PYTHON,无 env | ①router 忘 `--mpi=pmix` 秒死;②修后 ctx 侧 NIXL `TransferState.FAILURE` → gen 等 KV 超时自杀 | 任何 trtllm CLI 经 srun 都要 pmix(已入记忆) | | R3 | DEFAULT + UCX 三件套(unset UCX_TLS/关 UCC/RNDV_RAILS=2,抄自同集群日跑 harness)| ctx 启动死于 **v2 KV 管理器 CUDA_ERROR_INVALID_CONTEXT** | DEFAULT(C++ 传输)与 v2 管理器不兼容;p25 时代选 NIXL+**PYTHON runtime** 正是 v2 兼容路径 | -| **R4(进行中)** | **NIXL+PYTHON + UCX 三件套**(未试组合:NIXL 底层即 UCX,三件套正治 UCX 传输)| — | 若仍 FAILURE,下一杆:UCX_NET_DEVICES 按 GPU 就近 HCA 钉扎(同事 harness 的 CTX_LOCAL_HCA_PIN)| +| R4 | **NIXL+PYTHON + UCX 三件套**(未试组合:NIXL 底层即 UCX,三件套正治 UCX 传输)| 六 cell 全过 bringup(证 pmix/PYTHON 组合成立),但传输仍全灭:ctx 首笔 **等 336s 才 FAILURE**(非 0 秒即拒),此后 agent 元数据被 invalidate → 后续全部秒拒 `metadata for remote agent not found`;gen 60s KV 超时先于 336s → 1024 全灭;gen 侧 v2 IMA 为下游殉爆(page 释放撞未完成 NIXL op)| UCX 三件套治不了它——根因在更上层(见 R4 收网)| + +**R4 收网(08-11 08:20,代码取证 + 双实证)——根因:UCX MNNVL 误判 v2 fabric-handle 显存**: +- **机理**:v2 管理器用 `CU_MEM_HANDLE_TYPE_FABRIC` 分配 KV 物理内存(`kv_cache_manager_v2/_cuda_virt_mem.py`);B300 上 UCX 的 cuda_ipc-MNNVL 路径声称此类显存跨节点可达,但两个独立 slurm job 间无 IMEX 通道 → RDMA WRITE 卡死至 UCX 判 ep 死(~336s,恰为观测值)→ NIXL invalidate 对端 → 雪崩。修法即官方文档旋钮:**`UCX_CUDA_IPC_ENABLE_MNNVL=n`**(disagg-serving.md;v1/C++ 栈用普通 cudaMalloc 无此病,故同集群他人 disagg 正常)。 +- **实证互洽**:跨节点 `ucx_perftest`(host 与 cuda 内存、4MB 消息)默认 env 下 90 GB/s 全绿——因 perftest 用普通 cudaMalloc,**对本 bug 是假阴性**;1/1024 偶发成功与 stall(非不可达)一致。 +- **伴随修正**:ctx yaml 原来没给 speculative_config → ctx/gen 窗口池几何不一致(128/8 vs 133/13)且 draft KV 无人 prefill;已镜像 DSpark spec 块到 ctx(MTP disagg 同款做法)。取证副产品:两次 transceiver init/侧 = KV 估算临时 executor(正常);`kv_transfer_timeout_ms` 默认 60s 且**传输无重试**;PYTHON runtime 为 v2 唯一兼容路径(官方注释坐实)。R4 现场已归档 `dspark-runs/r4_archive/`。 +- **R5(进行中)**:MNNVL=n + ctx spec parity,金丝雀 S:128 先行,通过后其余五 cell 并行重跑。 ### 对齐实验:zhaoyuanh 的 throughput_1k 负载 × 我们的栈(08-11 02:10,主矩阵完成) diff --git a/docs/source/features/dspark-confidence-scheduling.md b/docs/source/features/dspark-confidence-scheduling.md index 0b2e8de2f22c..8367551a005f 100644 --- a/docs/source/features/dspark-confidence-scheduling.md +++ b/docs/source/features/dspark-confidence-scheduling.md @@ -23,6 +23,8 @@ DSpark 每个 decode step 都会 draft 出一整个 block，并把它全部送 - [退化路径](#退化路径) - [配置](#配置) - [性能天花板与已知局限](#性能天花板与已知局限) +- [部署要求：DSv4 必须显式设 max_seq_len](#部署要求dsv4-必须显式设-max_seq_len否则并发被钉在-maxbs2) +- [基准配置（可复现）](#基准配置可复现) ## 问题与核心洞察 @@ -505,6 +507,93 @@ speculative_config: | ragged 模式的端到端数字仅在一个平台上取得 | DeepSeek-V4-Pro-DSpark，8×B300 DEP8，overlap + CUDA graph（GSM8K 96.44，基线 96.21）；仍默认关闭，跨硬件需重测 | | 逐 token 采样参数路径只有非全贪心 batch 才会触发 | `temperature=0` 的精度跑覆盖不到它 | +## 部署要求：DSv4 必须显式设 `max_seq_len`（否则并发被钉在 ~maxbs/2） + +DSv4 的窗口池组（SWA + compressor KV/score，窗口 = 128 + `max_draft_len`，DL=5 时 133 +token）在 v2 KV manager 里是固定槽分配器。不显式设 `max_seq_len` 时，serve 会从模型推断 +出 1M 并把它当 `typical_seq_len` 喂给比例定容——全注意力池吃掉全部配额，窗口池被饿到下限 +`min_slots ≈ 1.05 × max_batch_size` 块（`kv_cache_manager_v2.py:1775-1779`），而下限按 +history=0 建模成每请求 1 块（`:1803-1812`）。运行时窗口 133 > tokens_per_block 128，每个 +生成请求实际锁 2（偶 3）块（`_life_cycle_registry.py:38-51`），于是可调度并发 += 1.05·maxbs / 2.07 ≈ **0.51 × max_batch_size**（实测：maxbs=128 → 61-63 行；maxbs=256 +→ 131-136 行），伴随 ~2.7× 的步时惩罚。关闭 spec 时同一机制不降级而是**调度器死锁** +（`scheduler_v2.py:439` 处崩溃，报错信息有误导性）。 + +**修法：serving 配置里显式写 `max_seq_len`（按真实负载，比如 8k1k 场景写 9216）。** +`typical_seq_len` 随之落地，比例定容给窗口池的份额远高于下限，满 maxbs 可调度。判别实验 +证明只动这一项即解钉；只调 KV fraction / max_num_tokens 均无效。历史工作绕法 +`max_batch_size = 2 × 目标并发` 是同一算术的另一侧（把下限抬高一倍），在显式 +`max_seq_len` 之后不再必要，但对 ctx 准入 headroom 仍有独立价值（见下）。 + +代码侧的根治（`typical_seq_len` 不应默认推断的 max_seq_len、下限按稳态 2 块建模、 +no-spec 应降级而非死锁、字节估算器补 draft 窗口放大）待报 upstream，不阻塞部署规约。 + +## 基准配置（可复现） + +完整实验记录与逐轮数据见 `docs/dspark_sts_program.md`。以下为两类负载的精确配置。 + +### 负载 A：poetry / arena（agg，burst-drain；收益 +5.9~+24.7% 的来源） + +- 硬件/形态：DeepSeek-V4-Pro-DSpark，8×B300 单节点，`--tp_size 8 --ep_size 8`（DEP8， + `enable_attention_dp: true`），`--max_batch_size 256 --max_num_tokens 4096`。 +- 负载：poetry = 固定短 prompt（写诗，输出 ~760 tok/req）；arena = arena 真实问题集。 + 客户端 burst-drain：一次性下发全部 N（512/1024）个请求，计时到全部完成，n=10 轮取中位。 + **注意：该协议下 99% 的步是纯 gen 步**（prefill 波在头几十步内结束）——这是收益成立的 + regime 前提。 +- S 臂 spec 配置（`pro_sched_arm.yaml`）：tiers `[2,3,4,5]`， + `confidence_sps_table_path: postfix_pro_table.json`（修复 #31/#32 后在 maxbs=256 满批 + 重采：bs128 实测 θ(L2/3/4/5)=80.8/90.3/102.7/110.3ms），**未挂 STS**（裸 sigmoid）。 + N 臂同配置去掉两个 confidence 开关。kv 0.5，block reuse 关，graphs [1,8,32,64,128] 带 + padding，overlap on。 + +### 负载 B：throughput_1k（roofline 对齐负载，agg 持续饱和） + +- 负载：throughput_1k parquet 前 1024 条（ISL≈1110），greedy，`max_tokens=64`；客户端 + 持续饱和（固定在飞数，完成即补位，ramp 256 后计量 1024 个完成）。 +- 关键物理：OSL=64 + 长 prompt + 持续补位 → **~50% 的步是 ctx+gen 混合步**，纯 gen 采的 + 成本表在混合步高估节省（成本模型 ctx 盲），裁剪收益不成立（sched −1.2~−6.3%）。收益是 + **负载 regime 的函数**，不是普适常数。 +- **把 gen bs 打到目标值的配方**（实测绑定约束是准入带宽：mnt=8192 时有准入的步 token 顶 + 死预算,到达率 3.76/步 × 寿命 24.6 步 = 均衡 bs 92）： + 1. `max_num_tokens: 16384` —— 解除准入带宽上限； + 2. `max_batch_size = 2 × 目标` —— ctx 准入槽位 headroom； + 3. 客户端在飞 = 目标并发 —— 让 client 成为干净的绑定约束。 +- 成本表与 STS 均须在本负载重采（表按 `--from-iter-log` 纯 gen 段拟合）。 + +### 负载 B'：disagg（gen-only 吞吐；裁剪的顺风 regime） + +ctx/gen 分离后 gen server 每步纯 gen，成本表全程在域内。两台 B300 独立 slurm 作业, +NIXL 传输。**四个必踩的坑与解法**： + +1. v2 KV manager 只支持 `transceiver_runtime: PYTHON` + `backend: NIXL`（C++/DEFAULT + 传输在启动时 `CUDA_ERROR_INVALID_CONTEXT`）。 +2. **必须 `export UCX_CUDA_IPC_ENABLE_MNNVL=n`（两侧）**：v2 用 + `CU_MEM_HANDLE_TYPE_FABRIC` 分配 KV,无 NVLink 互联的独立作业间 UCX 会误判 cuda_ipc + 可达 → RDMA 卡死 ~336s → NIXL 元数据雪崩（全请求失败）。`ucx_perftest` 用普通 + cudaMalloc,对此 bug 是假阴性。NVL72 机架内（IMEX 就位）不需要。 +3. 超时三层都要抬：`cache_transceiver_config.kv_transfer_timeout_ms: 600000`（默认 60s, + 超发排队下成批误杀）;router `trtllm-serve disaggregated --request_timeout 900`（默认 + 180s）;（传输层由 2 解决）。 +4. **ctx 与 gen 必须共用同一个 `speculative_config`**（YAML anchor）——否则窗口池几何不 + 一致（128/8 vs 133/13）且 draft KV 无人 prefill。 + +```yaml +# ctx 与 gen 共同的骨架（两侧一致的部分） +max_seq_len: 8192 # 见"部署要求" +speculative_config: &spec + decoding_type: DSpark + max_draft_len: 5 + speculative_model: +kv_cache_config: {enable_block_reuse: false, free_gpu_memory_fraction: 0.5} # gen 侧 0.7 +enable_attention_dp: true +cache_transceiver_config: + backend: NIXL + transceiver_runtime: PYTHON + max_tokens_in_buffer: 8192 + kv_transfer_timeout_ms: 600000 +# ctx 侧另加 disable_overlap_scheduler: true;gen 侧在 *spec 上打开 confidence 开关。 +``` + ## 参考 - DeepSeek DeepSpec： From 3bf934867d48e9295380c8ae7a67b9e8e7e3b04a Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:34:44 -0700 Subject: [PATCH 154/177] [None][doc] Benchmark config: pin bs by client in-flight count; drop unverified maxbs/mnt prescriptions Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/source/features/dspark-confidence-scheduling.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/source/features/dspark-confidence-scheduling.md b/docs/source/features/dspark-confidence-scheduling.md index 8367551a005f..af8ea7d0e1db 100644 --- a/docs/source/features/dspark-confidence-scheduling.md +++ b/docs/source/features/dspark-confidence-scheduling.md @@ -553,17 +553,17 @@ no-spec 应降级而非死锁、字节估算器补 draft 窗口放大）待报 u - 关键物理：OSL=64 + 长 prompt + 持续补位 → **~50% 的步是 ctx+gen 混合步**，纯 gen 采的 成本表在混合步高估节省（成本模型 ctx 盲），裁剪收益不成立（sched −1.2~−6.3%）。收益是 **负载 regime 的函数**，不是普适常数。 -- **把 gen bs 打到目标值的配方**（实测绑定约束是准入带宽：mnt=8192 时有准入的步 token 顶 - 死预算,到达率 3.76/步 × 寿命 24.6 步 = 均衡 bs 92）： - 1. `max_num_tokens: 16384` —— 解除准入带宽上限； - 2. `max_batch_size = 2 × 目标` —— ctx 准入槽位 headroom； - 3. 客户端在飞 = 目标并发 —— 让 client 成为干净的绑定约束。 +- **bs 定量**：客户端在飞数 = 目标并发（8×目标/节点），maxbs/max_num_tokens 与生产配置 + 一致即可。本负载 OSL=64、请求周转极快，实测 gen bs 会因准入/周转间隙低于标签值——报告 + 时以 iter log 的实测中位为准，不以档位标签为准。 - 成本表与 STS 均须在本负载重采（表按 `--from-iter-log` 纯 gen 段拟合）。 ### 负载 B'：disagg（gen-only 吞吐；裁剪的顺风 regime） ctx/gen 分离后 gen server 每步纯 gen，成本表全程在域内。两台 B300 独立 slurm 作业, -NIXL 传输。**四个必踩的坑与解法**： +NIXL 传输。**bs 钉法**：gen server 没有 prefill 争抢槽位，直接 +`max_batch_size = 目标并发` + 客户端超发（如 16×）即可把 gen bs 钉在 maxbs 上。 +**四个必踩的坑与解法**： 1. v2 KV manager 只支持 `transceiver_runtime: PYTHON` + `backend: NIXL`（C++/DEFAULT 传输在启动时 `CUDA_ERROR_INVALID_CONTEXT`）。 From 2391b203f7dde0ee6c9ec9d42df233719bfdb779 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:36:04 -0700 Subject: [PATCH 155/177] [None][doc] Disagg bs pinning: moderate client oversubscription (2-4x) with measured gen-bs monitoring Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/source/features/dspark-confidence-scheduling.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/source/features/dspark-confidence-scheduling.md b/docs/source/features/dspark-confidence-scheduling.md index af8ea7d0e1db..7bf42b6467db 100644 --- a/docs/source/features/dspark-confidence-scheduling.md +++ b/docs/source/features/dspark-confidence-scheduling.md @@ -562,8 +562,9 @@ no-spec 应降级而非死锁、字节估算器补 draft 窗口放大）待报 u ctx/gen 分离后 gen server 每步纯 gen，成本表全程在域内。两台 B300 独立 slurm 作业, NIXL 传输。**bs 钉法**：gen server 没有 prefill 争抢槽位，直接 -`max_batch_size = 目标并发` + 客户端超发（如 16×）即可把 gen bs 钉在 maxbs 上。 -**四个必踩的坑与解法**： +`max_batch_size = 目标并发` + 客户端适度超发（2–4×，够填补周转间隙即可）把 gen bs 钉在 +maxbs 上；**必须从 gen 侧 iter log 监控实测调度 bs 确认钉住**，不能只信标签。超发过猛没 +有额外收益，只会加深 ctx 队列、压穿 KV-transfer / router 超时。**四个必踩的坑与解法**： 1. v2 KV manager 只支持 `transceiver_runtime: PYTHON` + `backend: NIXL`（C++/DEFAULT 传输在启动时 `CUDA_ERROR_INVALID_CONTEXT`）。 From 722046263fd6f6504ec33baac58133754b6be250 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:40:22 -0700 Subject: [PATCH 156/177] [None][doc] Dataset provenance for the three benchmark workloads (poetry, arena-hard, throughput_1k) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../features/dspark-confidence-scheduling.md | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/docs/source/features/dspark-confidence-scheduling.md b/docs/source/features/dspark-confidence-scheduling.md index 7bf42b6467db..fc673d169fe7 100644 --- a/docs/source/features/dspark-confidence-scheduling.md +++ b/docs/source/features/dspark-confidence-scheduling.md @@ -536,8 +536,13 @@ no-spec 应降级而非死锁、字节估算器补 draft 窗口放大）待报 u - 硬件/形态：DeepSeek-V4-Pro-DSpark，8×B300 单节点，`--tp_size 8 --ep_size 8`（DEP8， `enable_attention_dp: true`），`--max_batch_size 256 --max_num_tokens 4096`。 -- 负载：poetry = 固定短 prompt（写诗，输出 ~760 tok/req）；arena = arena 真实问题集。 - 客户端 burst-drain：一次性下发全部 N（512/1024）个请求，计时到全部完成，n=10 轮取中位。 +- 数据集： + - **poetry**：单条固定英文 free-verse 长诗 prompt（memory/rivers/glass/time 主题、 + 要求 ≥8 个场景、不许复用意象），复制 N 份。短输入（~60 tok）+ 长输出（实测 + ~760 tok/req，自然 EOS）——短 ISL / 长 OSL 的极端。 + - **arena**：arena-hard-v0.1 的 500 条真实问题（jsonl，取 `prompt` 字段，N>500 时循环 + 复用）。中等 ISL / 长输出的多样负载。 +- 客户端 burst-drain：一次性下发全部 N（512/1024）个请求，计时到全部完成，n=10 轮取中位。 **注意：该协议下 99% 的步是纯 gen 步**（prefill 波在头几十步内结束）——这是收益成立的 regime 前提。 - S 臂 spec 配置（`pro_sched_arm.yaml`）：tiers `[2,3,4,5]`， @@ -548,8 +553,11 @@ no-spec 应降级而非死锁、字节估算器补 draft 窗口放大）待报 u ### 负载 B：throughput_1k（roofline 对齐负载，agg 持续饱和） -- 负载：throughput_1k parquet 前 1024 条（ISL≈1110），greedy，`max_tokens=64`；客户端 - 持续饱和（固定在飞数，完成即补位，ramp 256 后计量 1024 个完成）。 +- 数据集：roofline 报告使用的 **throughput_1k** parquet + （`iputterman/speed/throughput_1k/test-00000-of-00001.parquet`），按文件序取前 1024 条、 + 每条取 `turns[0]` 为 prompt（实测 ISL≈1110 tok）；greedy（temp=0），`max_tokens=64`。 + 注：源 parquet 的 SHA 与 roofline 报告冻结 dump 不同，对比其报告数字时需留意版本。 +- 客户端持续饱和：固定在飞数，完成即补位，ramp 256 个完成后计量 1024 个完成。 - 关键物理：OSL=64 + 长 prompt + 持续补位 → **~50% 的步是 ctx+gen 混合步**，纯 gen 采的 成本表在混合步高估节省（成本模型 ctx 盲），裁剪收益不成立（sched −1.2~−6.3%）。收益是 **负载 regime 的函数**，不是普适常数。 From 9af9e78a00750fb5b413f7c3399da6566eba6224 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:45:40 -0700 Subject: [PATCH 157/177] [None][doc] Disagg independent-repro campaign (laliao): own-harness route, three-workload matrix, per-workload SPS/STS; request verbatim poetry prompt Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 151 +++++++++++++++++++++++++++++++++++++ 1 file changed, 151 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 9de4813c5749..bd18bc8c384d 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -717,6 +717,157 @@ Flash 三臂翻案后,Pro@DEP8 终局(maxbs=256 解锁满批 + `postfix_pro_tabl **修复后 Pro 表(08-10 10:55,maxbs=256 满批 + #32 修复,--from-iter-log,4 pin × 3 burst 双节点)**:`postfix_pro_table.json`,bs128 实测 L2/3/4/5 = **80.8/90.3/102.7/110.3ms**(n=47-81/cell)——**完美单调,w4 异常消失,GB300"平台+悬崖"楼梯消失**(确证均为 #32 prep 串行化 + #31 半批的伪影);固定 39.3ms + α(128)=9.1ms,θ(768)=61.9ms = **步时 56% 可裁**(Flash 为 54%,但 Pro 绝对杠杆 2.4×)。w5→w2 省 26.7% 步时。下一步:Pro@DEP8 三臂终局(scheduled-新表 vs notrim,maxbs=256),节点已排。 +## 【08-08 10:30】协查 agent(laliao 侧)独立复现成立;变体理论证伪;地址拓扑测绘沉淀 + +独立环境(laliao,agg DEP8 2 节点,自家 gsm8k 表 + tiers[1..5],64 并发 percreq census)四轮实验: + +- **独立复现成立**:`FORCE_BUDGET_FRAC=0.25 + DEVICE_WINDOWS + steep` → **64/64 腐蚀**(median 0.13 chars/tok),与主线症状逐位一致。两个复现前置坑沉淀: + 1. **steep 顶格空转陷阱**:窗口全在 ceiling 时 `give=min(a-lo,hi-b)≡0`,序幕恒等,census 假干净(两轮上当)。**空转哨兵必须随行**:`forced_budget_steps` 计数 + `verify_len_hist` 从 {5:...} 转 {2:...}; + 2. **`FORCE_VERIFY_LEN` 在 device 模式被静默忽略**(`decide_verify_budget` 只消费 `_forced_budget_frac`;pin 只在 `decide_verify_lens` 被读)。device 模式钉工况只能用 frac(p25 之名即 frac 0.25)。 +- **变体理论证伪(census 级)**:`FORCE_GREEDY_KEY=1`(0.7 流量强路由 greedy 图 + FORCE_ARGMAX)照样 **64/64 腐蚀** → greedy 图在真 w≠S 下同病。早前"greedy 净 / advanced 脏"的二分是复现条件不齐(temp0 轮次没凑出 w≠S)的伪信号。 +- **共享 metadata 修复证伪**:`TLLM_DSPARK_SHARE_VARIANT_META=1`(advanced key 复用 greedy key 的 metadata 对象,两图烘焙地址全同)机械上成立(840 对 key `n_addr_diffs=0`、健康路径无损)但 **64/64 仍腐蚀** → 变体地址分叉不是病根。 +- **地址拓扑测绘沉淀**(`TLLM_DSPARK_ADDR_DUMP=1`,warmup 末对每对同形状变体图逐名 diff 全部 metadata tensor 的 data_ptr):48+13 个 tensor 中恰 4 个分叉——`_seq_lens_cuda`、`gen_indexer_kv_lens_cuda_runtime`(on_update 每次重绑新 tensor)、`batch_indices_cuda`、`captured_hidden_states`;外加 advanced 副本 `_force_non_greedy_for_capture` 永久卡 True(引擎对象清了、per-key 副本没清;已加清理)。均降级为卫生项。 +- **进行中(S1 判别)**:`TLLM_DSPARK_HOSTLEN_BOUND=1`——`attn_row_kv_lens_host`(→`TokenMajorGenView.host_past_key_value_lengths`,dsa.py:2123 直喂 trtllm-gen C++ op)在 device 模式下持有 SHAPE split 的 S 布局行长,而"host 读都是 bound"契约对 token-major 行不成立(w 布局下行 extent 可超 S extent)。判别:H2D 拷贝后把该 host 数组覆写为全批真上界(设备侧精确值不动)。净 → host 长度契约面定罪;脏 → 该面排除,与主线"跨行分歧检测器"的首个分歧行窗口类别交叉定位。 +- 与"1-2 个幸存请求"吻合的先验:扁平布局前缀和污染模型——首个规模错误行之后全污染,之前的行幸存。若分歧检测器显示分歧总从固定低行号开始,该模型坐实。 +- 复现资产(laliao 侧,可复用):`tmp/dspark_build/devwin/`(devwin_repro.slurm 多臂 + percreq_census.py + agg_config_dev.yaml),overlay=`overlay_mounts_dev.txt`(28 文件全量覆盖 faf2c60935)。 + +## 【08-08 11:40】独立验证判决(laliao 侧):compressor 根因与修复 D2D 双向确认 + +同一独立 harness(agg DEP8 2 节点、自家表、FRAC=0.25+steep、64 并发 @0.7 census)三角闭合: + +- **pre-fix**:64/64 腐蚀(median 0.13); +- **pre-fix + S1 猜想**(`HOSTLEN_BOUND=1`,token-major host 长度覆写真上界):**仍 64/64 腐蚀**(0.094)——负对照通过,`attn_row_kv_lens_host` 契约面排除,"gen_new_tokens 唯一嫌疑"的排他性在独立环境成立; +- **post-fix(57e9d56e8c)**:**0/64 腐蚀**(61 净,median 3.90),序幕 32 次 apply、w≠S 全程在场(hist {2:1708, 3:366, 4:52, 5:2})。 + +代码级四项核查亦过:单位等价(`gen_token_repeat_list` ≡ staging 时 `_seq_lens_cuda[nc:nc+n]`)、捕获区位置(on_update_kv_lens:932)、序幕写集确无此缓冲(apply_device_ragged_layout 写集核对)、D2D 源为捕获期稳定地址且被序幕改写为 w。 + +**代价审计(修复可发布性)**:向量件零成本(删 pinned 缓冲 + 免每步 host 构造);标量件 NEXT_N→全局 6 进了尺寸计算(`num_total_compressed_tokens`、`kv_comp = torch.empty(...)`、compressed_mask arange),但瞬态/图池水位均 ≤ notrim 满块步已付水位(uniform 步本来就按 6 跑),掩蔽靠逐请求精确向量、正确性自洽;净代价 = 裁剪步的压缩器阶段按满块元数据宽度跑,二阶小头,精确定量归 Step 3 三臂。 + +补充验证在跑:dev@temp0 census、host+HOST_STEEP+frac 回归 census。方案 C(kernel 内导出 nn,消灭冗余输入)同意作为 PR 收尾清理候选。 + +## 【08-08 14:30】跨节点四轮复证(laliao 侧,独立环境):#30 "无收益"结论坐实 + +独立复刻 Step 3 规格(4 个独立 node pair × 三臂冷启动 notrim→host→device;poetry(生成)+arena(500 题循环)× bs{512,1024} global × 3 reps × out768 × temp0.7;**本地重拟表**:frac 扫描 7 断点 + 重建 sts_v1;tiers[2,3,4,5],max_batch 128): + +| cell | h/n 四轮 | d/n 三~四轮 | +|---|---|---| +| poetry bs512 | −3.5/−10.2/−3.5/+0.9 | −0.2/−1.0/−2.2/+2.0 | +| poetry bs1024 | −4.7/−5.5/−1.7/−0.9 | +0.5/−1.8/−1.4 | +| arena bs512 | −13.1/−10.2/−6.7/−9.1 | −8.8/−7.4/−6.6 | +| arena bs1024 | −6.3/−11.1/−4.9/−4.8 | (4h 墙截尾,补充轮在跑) | + +- **d/n:poetry ≈ 0±2%(噪声),arena bs512 ≈ −7.6%(三轮方向一致)——无任何正收益**,与主线终局(device/host −6.6%~+0.3%)方向吻合; +- **h/n 在本环境系统性偏负(−3~−10%)**,比主线的 0±4% 更差——可能因素:本地表拟合质量(7 断点 vs 16)、arena 500 题循环的重复前缀结构、节点差异。方向上仍支持同一结论; +- **判决:B300 agg DEP8 + 校准配置下,host/device 动态裁剪无可复现吞吐收益——跨环境(两套表、两套 harness、两组节点)复证成立。**Step 4 门禁(device>host 复现)正式失效;Step 5 门禁同亦失效。 +- **勘误 + 新 bug(Bug #4?):device 臂 arena bs1024 缺格不是 4h 墙(elapsed 均 ~2.5h),是 5/5 可复现的崩溃**——补充轮(2630554,device 先行)在同一 cell 现场抓获:torch `IndexKernel.cu:111` 设备端断言 "index out of bounds"(index 类算子越界)→ `CUDA_ERROR_ASSERT 710` 全 rank SIGABRT。发生在 cell 爬坡段(混合 ctx+gen 步窗口,iter 日志 num_ctx_requests 2→54;注意设备断言浮出有延迟,故不能据此断定肇事步是混合步还是其前的 gen-only 步)。**判别面:host/notrim 同 cell 四轮全过,poetry bs1024(同 128/rank)device 也过,arena bs512 device 也过——特异组合 = DEVICE_WINDOWS × arena 长 prompt 流量 × bs128/rank**。现场:laliao 侧 `tmp/dspark_build/step3/logs/s3device_2630554_serve.log:55198`(断言首现)。吞吐判决不受影响(该 cell 有主线 device/host −2.1% 佐证);但 device 路径在长 prompt 高 bs 下仍有未清的越界索引,建议与"潜伏 bug ①②"并案(候选:混合步/context 加入步与 device budget 单发状态的交互、_seed_context_windows 与 chunked prefill 的槽位索引)。 + +## 【08-09 02:30】bs 曲线补全(laliao 侧):32→256/rank 四点,d/n 谷型、h/n 单调恶化;崩溃 7/7 复现 + +用户加点指令下把 Step 3 网格扩到两端(同配置,256/rank 端重采表 sps_frac_local_2632577 = bs 轴 [64,128,256] 指纹 256): + +| per-rank bs | poetry h/n | poetry d/n | arena h/n | arena d/n | +|---|---|---|---|---| +| 32(256 global) | **+1.4%** | +0.5% | −1.5% | −0.4% | +| 64 | −4.1% | ≈0 | −9.8% | −7.6% | +| 128 | −3.2% | ≈0 | −6.8% | (崩) | +| 256(2048 global,同节点) | −6.0% | **−1.8%** | −1.1% | (崩) | + +- **h/n 随并发加深单调恶化,32/rank 附近归零/转正**——与 #28 SGLang 原生 B300 曲线(bs1 +17.8% / 8-64 平 / 128 −21% / 256 −4~−8%)同形;B300 的"由负转平"边界在 32-64/rank 之间; +- **深水区 device 亏损显著小于 host**(poetry −1.8% vs −6.0%):lag-0 新鲜排序换回的 accept 在深批下开始盖过序幕固定开销的摊薄——机制方向正确,只是在 B300 上翻不了正; +- **IndexKernel 越界崩溃 7/7 复现**(device × arena × {128,256}/rank;poetry 同 bs 全过):与 bs 无关、与 arena 长 prompt/chunked prefill 工况强绑定;2048 global 下 device arena r1 中途崩(部分吞吐 16.2k 无效)。 +- 工程沉淀:并行单臂(跨节点)与同节点配对的 arena h/n 差了 3.8pp——**arena 的节点噪声远大于 poetry**,该数据集的臂间对比必须同节点。 +- 进行中:cap-accept 测量臂(执行满块、窗口记账,cap_trim_tokens 逐 token 回答"被裁 token 真实存活率")。 + +## 【08-09 03:10】cap-accept 直接测量:被裁 token 真实存活率 11.5%(laliao 侧) + +cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完整收官,回答"被裁的 token 实际有多少会被接受": + +- **计划裁剪 1.89M token,其中 219,990 个(11.5%)在满块执行下实际被接受**——零反事实的逐 token 直数; +- 每请求步平均损失 **0.216 token**;算术闭环:封顶 accept 1.66 + 0.216 = 1.88 ≡ notrim 臂实测(1.80-1.88); +- 损失集中:仅 12% 请求步丢 token(丢者均值 1.8,单步 max 4);删失链自洽:31% 被裁请求前缀活到刀口 → 其中 ~55% 下一 token 真被接受 → 全体 11.5%; +- 与跨臂 A/B 估计(~9%)方向一致、A/B 略偏低。 +- **判词**:调度器 ~89% 的刀砍在死草稿上(排序质量好);无收益的根源是那 11.5% 的 accept 换不回 B300 的 kernel 时间(θ 平台),不是调度器裁错了人——给"成本结构论"补上最后一块直接证据。 + +## 【08-09 04:40】IndexKernel 崩溃根因定案(Bug #4):anchor 停车位在满窗满批步越界一格;"arena 特异"定性作废 + +- **判别链**:devfee 臂(DEVICE_WINDOWS + FORCE_BUDGET_FRAC=1.0,poetry,w≡S 恒等)在首个满批 gen-only 步**秒崩**,同一 IndexKernel index-OOB——推翻"device×arena×高bs"表征(poetry 也崩、w≠S 不必要、非爬坡累积)。 +- **根因(一行公式)**:序幕把 anchor token 停在 `dst = real_draft = real_tokens − n_real`(model_engine 序幕 scatter);`draft_tokens_cuda` 容量 = `max_draft_loop_tokens × batch_size`,**当 padded_bs = max_batch 且所有真实行窗口全满时 real_draft 恰等于容量,越界一格**。八次崩溃全解释:arena 高接受 → planner 常发全满窗口 → 满批步必踩;poetry 真实窗口碎(accept~1.2)几乎从不全满 → 幸存;devfee 恒全满 → 秒崩;128/rank 配置缓冲 640、满 bucket 768 同式。协查报告 S4 曾摸到该停车位但只审了语义撞位(判良性),漏审分配边界。 +- **修复(laliao 本地,一行)**:`draft_tokens_cuda` 容量 +1(停车位合法化为 scratch 尾格;bs×K 步进布局不扰、capture 安全)。验证轮(fee2b)在跑:devfee 转净 = 定案。 +- **教训**:①"数据集特异"表征是接受率经由调度器的间接投影,真正的触发几何要用受控恒等臂(frac=1.0)剥离;② 边界审计要同时问"语义对不对"和"分配够不够"。 + +## 【08-09 04:40】verify-len 评估装置入场费(256/rank,poetry,同节点配对,laliao 侧) + +- **头计算费(nohead→static)**:22747 → 22730 tok/s = **−0.07% ≈ 0**(fused 头前向在图内免费); +- **总入场费(nohead→notrim,frac=1.0 全机器零裁剪)**:22459 → 21859 = **−2.7%**; +- **host 机器组件帐单**(sched-timing,3500 稳态步,256/rank):decide ~450-510μs + allgather ~210μs + fit ~140μs + attach ~240-270μs + token-major 行装配 ~400-440μs ≈ **1.4-1.5ms/步 ≈ 5.7μs/请求**;而 host_step≈device_step(120-290ms,KV 深度双峰)——**host 份额被 overlap 完全遮蔽(<1%)**; +- 推论:2.7% 入场费的大头**不在 host 机器也不在头计算**,在 ragged 图的 token-major kernel 呈现差异(uniform 图 vs token-major 视图的 MLA/索引器路径)+ 图内每 token 机构;device 序幕费(notrim→devfee)待 fee2b 修复验证轮补上。 + +## 【08-09 12:30】序幕 Triton 融合(方案 C)落地:60 发射 → 4 kernel,enqueue 8.8×,parity bit-exact(laliao 侧) + +- **实现**:`dspark_device_select_fused.py`(env 门控 `TLLM_DSPARK_FUSED_PROLOGUE=1`,fallback 原路保留):Kernel A(单 CTA:slot gather→stamp fail-open→温度校准→存活→radix-select 预算 topk→cumsum-cycle fill→行映射,radix 复合键复刻 host tie-breaking 到 bit 级)+ Kernel B(bucket 网格:全部 token 级重写,A 写新窗口独立缓冲、B 读旧 split+新窗口,消除 read-before-write)+ 两个块表 gather。停车位 +1 scratch 不变量保留。 +- **两轮真卡教训**:①容器 Triton 拒收 jit 体内模块级全局/降序 static_range/constexpr 上的 not——全部参数化;②首版 e2e 反降 9%:**Triton 隐式整型值类特化**(每个 int 参数的 ==1/÷16 值类进编译缓存键)+ 形状派生 constexpr,ramp 每个 (padded_bs,bucket) 组合都触发数百 ms JIT 且落在流量中——`do_not_specialize` 全覆盖 + 形状参数改 runtime+mask + warmup 预热钩子(生产路径收敛到 4-5 次编译)。 +- **判决(同节点 A/B,256/rank poetry)**:enqueue mean 2052→**243μs**(尾巴消失,mean≈p50),e2e 22548→22556(打平),r1 无编译坑,parity 15/15,越界 0。 +- **综合口径**:序幕 e2e 代价 = 0~3%(取决于 host 压力:serve 口径曾现形 −3.1%,bench 口径 ≈0,本轮 overlap 藏住)——融合拔掉 host 分量后**对 host 压力不再敏感**。若 device 路径未来在低并发/延迟场景复活,这块砖已备好。 + +## 【08-09 15:10】观测开销排污判决:enable_iter_perf_stats 无罪(≤0.2%)(laliao 侧) + +- **设计**:同配方两次真实流量采集(18 phase:frac{0.2-1.0}×conc{512/1024/2048},poetry,新增 `/dspark/budget_frac` 运行时端点免六次冷启动),唯一变量 = `enable_iter_perf_stats`;判决用**同 phase 步时中位数直比**(深度分布天然配对,免分解、免建表)。 +- **判决**:stats-ON vs OFF 步时 **+0.18%**(18 格,散布 −0.5~+1.6%)——stats 机器(含每步 allgather/D2H)代价在噪声内,**此前全部阶梯/曲线结论无需修正**。客户端吞吐 +2.45% 偏 ON 为短 phase span 噪声(方向反于污染假设)。 +- **方法论坑(已沉淀)**:被动 iter-log 样本直接喂分解拟合器必炸——同 (bs,M) 格步时随 KV 深度 120→290ms 漂移,深度方差全进固定项(签名:fixed=258.8ms、θ(179)=0.7ms)。钉接受的本质是控深度;被动样本要么按 kv_cache_util 分箱,要么只做配对比较。 + +## 【08-09 19:30】短窗病灶三方会师 + pin 阶梯直测复现 + H4 候选机理(laliao 侧) + +- **pin 阶梯直测**(`/dspark/verify_len_pin` 运行时端点逐档钉,真实 poetry serving,~250 稳态步 median/格):64/rank 上 **pin4=59.1ms > pin5=56.9ms**(验更少 token 慢 2.2ms),w3→w4 悬崖 +10.3ms;256/rank 恢复单调(111.5<119.4)——**病灶咬中低 bs,深 bs 消退**。与 θ 表倒挂(62.7>60.8)、真实流量 frac0.65>frac1.0(57.7 vs 54.4ms)三方一致。 +- **H4 候选机理(代码级)**:`_pick_dsl_expand`(dsa.py:211)要求 `expand×atom==next_n` 精确因式分解(atom∈FP8{1,2,3,4});**w4→5 token=质数→唯一解 (5,1)→atom=1、5× 行展开=索引器 5× KV 重读,烘焙进该 bucket 捕获图**。w2/3/5(3=1×3, 4=2×2, 6=2×3)全部可摊销。暗合:此前"表自荐 tiers[1,2,5]={2,3,6} token"恰好全避质数——拟合器已从数据学会绕行。与协作 agent 的 H1/H2/H3 并列待 nsys 点名(pin4/pin5 各 100 稳态 iter 的 cuda_gpu_kern_sum 对拍在采,2025.5.1+cudaProfilerApi 配方)。 +- **修法②(tiers 避质数)首验证伪一半**:host[1,2,5] vs host[2,3,4,5] @poetry 64/rank = **−3.2%**——floor 降到 w1 裁得更深,accept 损失盖过避档收益;正确菜单可能是 [2,3,5](保 floor 只删 w4),推荐器需把 accept 损失定价进去(当前只看步时)。 +- **收益区首个正 cell**:GSM8K × 16/rank × devreal(真实 device 窗口+融合序幕)= **+1.8%**(两 rep 一致,甩开 host 3.4pp)——"高接受负载 × 中低 bs × 新鲜排序"的理论收益区实测点亮;加密扫描(8-64/rank × devreal/devalt)在跑。 + +## 【08-09 21:30】✅ 分歧和解:主线 = DEP4×DSv4-Flash,laliao = DEP8×DSv4-Pro(用户确认) + +- 绝对步时 1.5× 差(满窗 79.8 vs 119.4ms @256 行)与病灶档位差(Flash/DEP4 咬 3-4 token 档、Pro/DEP8 咬 5 token 档)全部由 **模型 × 并行度 → GEMM 形状 → tile/tactic 悬崖换位** 解释;两侧数据各自三源自洽,无人测错。 +- **通用律**:verify-len 阶梯的 kernel 成本地形不可跨 (模型×并行度) 迁移——梯子/表/绕档策略必须逐配置 profile,且在代表性 KV 深度下采。laliao 侧 nsys(pin4/pin5,Pro/DEP8)与主线 nsys v3(Flash/DEP4)各自点名各自真凶,并排即完整故事。 +- Pro/DEP8 侧 H4 候选(w4=5token 质数 → `_pick_dsl_expand` atom=1、5× 行展开)待 nsys 证实;Flash/DEP4 侧真凶待主线 v3。 + +## 【08-09 20:10】⚠️ 短窗病灶的跨 agent 分歧,请求 provenance(laliao → 主线) + +双方一致:非单调、w4 最糟。**分歧:w2/w3 的符号**——laliao 侧两份独立数据(pin 阶梯直测:w2/w3 = 46.4/48.8ms 均快于 w5 的 56.9;真实流量 frac 0.2/0.35/0.5 = 46.3/48.1/48.2 均快于 frac1.0 的 54.4)都显示**只有 w4 档为负**;主线判决 4 称"窗口 4/3 的桶比满窗慢 12-34%"。和解假说:① 均匀钉窗 vs topk **混合窗口组成**(同 bucket key 行组成不同,病灶或在混合行为上);② host pin vs **device 序幕重排后**布局;③ 负载/KV 深度/bs 差异。**请主线补 12-34% 的出处(bs、host/device、数据集、窗口组成、量法)**。在钉死前,交集结论可行动:w4(5 token)档在中低 bs 确定为负,梯子应绕开;分歧决定深裁档(w2/w3)是否可用,直接影响最优梯形。laliao 侧 nsys pin4/pin5(均匀)在采;若假说①成立将补 frac0.65(混合)对拍。 + +## 【08-09 23:30】SGLang 字面复现立项(laliao 侧);请主线共享 #28 环境 + +- 用户指令:按 blog reproduction 章节**用 SGLang 本体**复现(此前只做了设计复刻)。章节要点已抓取:镜像 `lmsysorg/sglang:dev-dspark`(PR #30261 @692c5f7d);frontier 臂 = Flash-DSpark×H200×DP4,one_batch_server bs {1..256}×out1024×temp0.7×frontier_prompt(16 拼接 GSM8K);零开销图 = **Pro-DSpark×B300×TP8×cuda-graph-max-bs 4**;动态/固定切换 = `SGLANG_RAGGED_VERIFY_MODE=compact|static`。 +- 进行中:enroot import dev-dspark 镜像(需验 aarch64);checkpoint `deepseek-ai/DeepSeek-V4-{Flash,Pro}-DSpark` HF 公开未 gated(Flash 74 文件 / Pro 92 文件,login 下载带宽 13MB/s 是瓶颈,将用 HF_TOKEN+分段配方)。 +- **请主线共享 #28 的 SGLang 环境**(镜像/venv 路径、checkpoint 位置、launch 参数)——你们已在 B300 跑通 SGLang 原生 DSpark,复用可省一天;也顺带补 #28 的 provenance。 + +## 【08-10 00:40】nsys 点名(Pro/DEP8 w4 真凶 = H1 GEMM tactic);devalt 复证失败;吞吐收益区关闭(laliao 侧) + +- **nsys pin4 vs pin5**(各 100 稳态 iter,cuda_gpu_kern_sum 对拍):①同批 GEMM tactic 换名——pin5 的 `64x64_2cta_TNT`(3245 次)/`256x128` 在 pin4 消失,被 **splitK 变体族**顶替(`tss_64x32_splitK` 4769 次等,净 +~130ms);②`gate_forward_kernel`(MoE 路由门)同 6100 次变贵 48%(56→83ms)——更小 M 落在更差 tile 点;③pin4 GPU kernel 总时长反而更少(3792 vs 3930ms,MoE 省 85ms)但占空比更低(64% vs 69%)——**"省了活、亏了排程"**:更碎的 splitK kernel 拉高发射/间隙。**判决:H1(tactic 选择)为主犯;H4(索引器质数 atom)未现身 top 差异,降级**。与主线 Flash/DEP4 的 H1 首要嫌疑吻合——两模型同类真凶、不同形状表现。 +- **devalt 复证失败**:两轮四格全负(32/rank:−3.0/−1.8;64/rank:−5.2/−1.7),D1c 首轮 +2.3/+1.0 为基线噪声(nohead 轮间 ±3%)。**B300 吞吐口径的收益区搜索正式关闭**:poetry/GSM8K × 1-256/rank × host/device/梯子变体,无任何跨轮稳健的正收益。存活方向:D2b 深上下文(在跑)与延迟口径(blog 的 B300 Pro 配置 cuda-graph-max-bs 4 本质是 latency 场景,SGLang 字面复现 #32 将验证)。 + +## 【08-10 01:30】收益区搜索终局(laliao 侧):B300 吞吐口径全维度关闭;唯余延迟口径待 #32 + +- **D2b 深上下文(ISL8k,chunked prefill)收官**:32/rank h/n −4.6%、d/n +0.3%;128/rank h/n −4.8%、d/n −5.5%——"θ 份额随 KV 深度涨 → 深上下文是主场"假说**证伪**(裁剪省下的照旧换不回 accept+税)。 +- **搜索矩阵全景(全部同节点差分,<5% 信号均经复证纪律)**:负载 {poetry, GSM8K, arena, mixed, ISL8k} × 并发 {1..256}/rank × 配置 {host, device+融合序幕, tiers[1,2,5]} —— **无任何跨轮稳健正收益**;单轮 +1~2% 信号(GSM8K 16/rank devreal、devalt 32-64/rank)复证全部还原为 ±3-5% 基线噪声。 +- **成因四件套(全部实测定价)**:H1 GEMM tactic 病灶(裁剪后 M 形状触发 splitK 碎战术 + gate kernel +48%,nsys 点名)+ accept 损失(cap-accept:被裁 token 11.5% 存活)+ 桶量化税(请求级 13-16% 落地仅 5.8-7.6%)+ token-major 呈现税(−1.6~−2.7%),叠加 B300 >50% 不可裁地板。 +- **唯一未决口子:延迟口径**(blog 的 B300×Pro cuda-graph-max-bs 4 配置、SGLang 原生 bs1 iter +17.8%)——#32 SGLang 字面复现进行中(镜像导入中,checkpoint 公开)。若延迟口径亦无,"DSpark 动态裁剪在 B300 无收益"即全口径终局。 + +## 【08-10 03:30】SGLang 字面复现收官(#32):blog 原版命令 @GB300,动态几乎全线为负——全口径闭卷(laliao 侧) + +- **配方**:`lmsysorg/sglang:dev-dspark` arm64(63 层 overlay 超内核限,计算节点手工压平 + docker ENV 50 条重注入;镜像/权重下载配方沉淀于 sglang_repro/)+ `deepseek-ai/DeepSeek-V4-Flash-DSpark` 180GB(计算节点 HF 出口 370MB/s,5 分钟)+ blog reproduction 章节逐字命令(DP4、flashinfer_mxfp4、one_batch_server bs{1..256}×out1024×temp0.7×frontier_prompt)。 +- **判决(static vs compact,输出吞吐)**:bs1 **+1.8%**(ITL 2.35→2.31ms,唯一正格);bs8 −12.8%、bs32 −9.4%、bs64 −5.2%、bs128 −6.1%、bs256 −1.6%。acc_length 双臂 4.7-4.9 一致(接受率形态复现)。 +- **结论**:同代码、同命令、同权重、同接受率,只换硬件(H200→GB300),blog 的 +10-20% 收益消失——**硬件地形是第一性原因的最干净证据**;与 TRT-LLM 侧全部结论跨栈闭环。#28 的 bs1 +17.8% 未在字面复现中重现(其 checkpoint/commit provenance 存疑)。 +- **DSpark 动态裁剪 @GB300 全口径闭卷**:吞吐(TRT-LLM,五负载×1-256/rank)与延迟(SGLang 字面,bs1-256)、两套实现——无收益;收益属于 θ 陡峭的 H200/H100 地形。 +- 【08-11 追记】上条"闭卷"已被 #31/#32 修复整体翻案(见主线终局判决);SGLang@GB300 compact 为负仍是其栈自身的事实,但"硬件地形第一性"解读降级为负载/实现因素之后的第三位。 + +## 【08-11 09:00】disagg 独立复现战役立项(laliao 侧):自有 harness 跑通路 + 三负载矩阵;请主线提供 poetry 原文 + +用户令 laliao 侧在 disagg 形态独立复现置信度调度收益,与主线 disagg 战役(R5)并行互证。设计与进展: + +- **通路**:不走主线的手搓 ctx/gen/router 脚本,用我自己 08-10 已跑通的 `bench-trtllm-disagg` harness(agentx gen-only 配方,**NIXL+PYTHON 本来就在用**,单 slurm job 3 节点带 `--segment` 同 NVL 机架分配——按 R4 收网的机理,IMEX 就位,`UCX_CUDA_IPC_ENABLE_MNNVL` 坑对单 job 同 segment 可能天然免疫;bring-up 正在验证)。新模型目录 `deepseek-V4-Pro-nvfp4-dsparkconf`(DSPARKCONF_* 提交时环境变量参数化全臂矩阵:SCHED/TABLE/STS/TIERS/GRAPH_MAX/GEN_ENV),harness 的 CONTAINER_MOUNTS 已修成支持 src:dst 文件绑定(python overlay 28 文件,含 b87ea28bf3)。 +- **协议吸收**:feature doc 的 B' 四坑全部落实——PYTHON+NIXL ✓、MNNVL(视 bring-up 判决)、kv_transfer 600s ✓、**ctx/gen spec parity 已修**(ctx 镜像基础 spec 块,confidence 开关 gen-only);max_seq_len 显式(harness 按 isl+osl 渲染)✓;bs 钉法改用主线定稿(maxbs=目标并发 + 客户端 2-4× 适度超发 + iter log 实测中位为准)。 +- **矩阵**:三负载 {poetry, arena, throughput_1k} × lbs {32,64,128,256} × 臂 {notrim, frac04, scheduled};**SPS 与 STS 分负载重采**(A:KV 深度 ~100-800 vs B:~1110,跨用错价;STS 是数据语义×采样温度的函数,A temp0.7 / B greedy);对表纪律:A 主臂裸 sigmoid(与主线 +5.9~+24.7% 同协议),STS 为 A 的消融格、B 的必挂项。**B 在 disagg 翻正与否 = "混合步成本表 ctx 盲"机理的直接检验**——这是本战役最有信息量的一格。 +- **➡️ 请主线**:提供 poetry 的 **verbatim prompt 原文**(provenance 只有描述:memory/rivers/glass/time、≥8 场景、不复用意象;我 step3 自造的 caravan/steppe 不是同一首,接受率分布不可比,对表 A 必须同文)。最好直接把数据集构造脚本或 jsonl 推上来;arena(arena-hard-v0.1 取 prompt 字段)与 throughput_1k(iputterman/speed parquet)provenance 已够我自建。 + +# 分析 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 9c5d203766e64c99b6ab863a65297514921c6ee4 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 09:07:14 -0700 Subject: [PATCH 158/177] [None][doc] Provide verbatim poetry prompt and bench client for independent disagg repro; correct poetry output-length description (768-capped, temp 0.7) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_bench_assets/p25fix_bench.py | 74 +++++++++++++++++++ docs/dspark_bench_assets/poetry_prompt.txt | 1 + docs/dspark_sts_program.md | 8 ++ .../features/dspark-confidence-scheduling.md | 12 +-- 4 files changed, 90 insertions(+), 5 deletions(-) create mode 100644 docs/dspark_bench_assets/p25fix_bench.py create mode 100644 docs/dspark_bench_assets/poetry_prompt.txt diff --git a/docs/dspark_bench_assets/p25fix_bench.py b/docs/dspark_bench_assets/p25fix_bench.py new file mode 100644 index 000000000000..d8bf9299bf34 --- /dev/null +++ b/docs/dspark_bench_assets/p25fix_bench.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +"""Round-style throughput bench: bs concurrent completions, JSON-line summary. + +Usage: p25fix_bench.py +Single-threaded asyncio client (no thread-pool: 1024 threads exhausts login +nodes and starves its own concurrency). Run node-local against the serve. +Mirrors the p25 rounds methodology: poetry (single prompt fan-out) and arena +(question cycle), bs in {512, 1024} global, 5 reps, max_tokens 768, temp 0.7; +metric is aggregate completion tokens / wall seconds. A leading warm rep per +(dataset, bs) is tagged rep=-1 so steady-state reps are comparable. +""" +import asyncio +import json +import os +import sys +import time + +import aiohttp + +HOST, PORT, TAG, SUMMARY = sys.argv[1], int(sys.argv[2]), sys.argv[3], sys.argv[4] +RUNS = "/lustre/fsw/coreai_comparch_trtllm/laliao/dspark-runs" +MODEL = (sys.argv[5] if len(sys.argv) > 5 else + "/lustre/fsw/coreai_comparch_trtllm/laliao/llm-models/DeepSeek-V4-Pro-DSpark") +URL = f"http://{HOST}:{PORT}/v1/completions" + +poetry = open(f"{RUNS}/poetry_prompt.txt").read().strip() +arena = [json.loads(l)["prompt"] for l in open(f"{RUNS}/arena_questions.jsonl")] + + +async def one(session, prompt): + body = {"model": MODEL, "prompt": prompt, + "max_tokens": 768, "temperature": 0.7} + try: + async with session.post(URL, json=body) as r: + d = await r.json() + return d["usage"]["completion_tokens"], len(d["choices"][0]["text"]) + except Exception: + return 0, 0 + + +async def bench(session, dataset, prompts, bs, rep): + t0 = time.time() + results = await asyncio.gather( + *[one(session, prompts[i % len(prompts)]) for i in range(bs)]) + wall = time.time() - t0 + toks = sum(r[0] for r in results) + chars = sum(r[1] for r in results) + fails = sum(1 for r in results if r[0] == 0) + line = json.dumps({ + "tag": f"{TAG}_{dataset}", "bs": bs, "rep": rep, + "completion_tokens": toks, "wall_s": round(wall, 2), + "out_tok_per_s": round(toks / wall, 1), + "chars_per_tok": round(chars / max(toks, 1), 2), "fails": fails, + }) + with open(SUMMARY, "a") as f: + f.write(line + "\n") + print(line, flush=True) + + +async def main(): + timeout = aiohttp.ClientTimeout(total=1800) + conn = aiohttp.TCPConnector(limit=1100) + async with aiohttp.ClientSession(timeout=timeout, connector=conn) as s: + bs_list = tuple( + int(x) for x in os.environ.get("BENCH_BS_LIST", "512,1024").split(",")) + for dataset, prompts in (("poetry", [poetry]), ("arena", arena)): + for bs in bs_list: + await bench(s, dataset, prompts, bs, -1) # warm rep + for rep in range(5): + await bench(s, dataset, prompts, bs, rep) + print("BENCH-DONE", TAG, flush=True) + + +asyncio.run(main()) diff --git a/docs/dspark_bench_assets/poetry_prompt.txt b/docs/dspark_bench_assets/poetry_prompt.txt new file mode 100644 index 000000000000..3fdbcd4a147b --- /dev/null +++ b/docs/dspark_bench_assets/poetry_prompt.txt @@ -0,0 +1 @@ +Write a long, original free-verse poem in English exploring memory, rivers, glass, and the passage of time. Use vivid, surprising imagery; vary the rhythm and stanza length; avoid repeating phrases or returning to the same image twice. Let the poem wander through at least eight distinct scenes. diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index bd18bc8c384d..aca647a1d3ef 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -875,3 +875,11 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 2. **device "+63.9%" 判定为假象**:device 臂生成损坏(见 Debug 日志 12:30 条),吞吐虚高由重复特殊 token + 零 EOS + 满批不排水制造。本轮 device 吞吐(+50%~+55%)全部作废。**H3 判 (c):接受判定/停机记账损坏**,旧 accept_stat 2.03 即其指纹。 3. 待办:device 路径根因定位与修复(736879/736880 debug 节点已排队);修复后 device 收益需重测才有结论。H2/H4 挂起,等 device 修复。 4. 联动结论:#28(SGLang 原生动态在 B300 高 bs 亏损)+ 本轮 host≈0 → 细刀 tier + 校准 + 修表后,**host 动态裁剪在 B300 agg 上没有可复现的吞吐收益**;与 SGLang blog 的 +10-20%(H200/H100)相比,B300 的 θ 平坦区更宽,裁剪节省的 token 成本换不回接受长度损失。 + +## 【08-11 09:40】主线 → laliao 侧:poetry 原文与构造脚本已入库(回应 disagg 复现请求) + +- **`docs/dspark_bench_assets/poetry_prompt.txt`**:终局判决(+5.9~+24.7%)所用 poetry prompt **逐字原文**(296 字节,memory/rivers/glass/time,≥8 场景,禁复用意象)。 +- **`docs/dspark_bench_assets/p25fix_bench.py`**:构造与打流客户端原件——poetry = 单 prompt 复制 fan-out ×N;arena = arena-hard-v0.1 全集 500 题按序循环(`prompt` 字段);**`max_tokens=768, temperature=0.7`**,warm rep=-1 + 5 计量 reps,指标 = 聚合 completion tokens / wall;asyncio 单线程,node-local 打 localhost(登录节点千线程会饿死自己的并发)。 +- **勘误(feature doc 同步修正)**:此前"自然 EOS ~760 tok"描述有误——实测均值 ~760 是 **768 上限基本打满**(poetry 几乎全部 length-capped),对表时请按 out768 记。 +- arena/throughput_1k 你侧 provenance 判断正确,无需我方再推(arena_questions.jsonl 即 arena-hard-v0.1 全集 500 题;parquet 路径已在 feature doc)。 +- 你侧 harness 若单 job 同 NVL segment 分配,MNNVL 坑理论上天然免疫(IMEX 就位)——但**请留 bring-up 判决为准**,失败签名是 ctx 侧 NIXL wait ~336s 后 FAILURE + 雪崩秒拒;R5b(主线,MNNVL=n + 600s/900s 超时)六对在跑,出数后两边可直接对表。 diff --git a/docs/source/features/dspark-confidence-scheduling.md b/docs/source/features/dspark-confidence-scheduling.md index fc673d169fe7..d33e5012954c 100644 --- a/docs/source/features/dspark-confidence-scheduling.md +++ b/docs/source/features/dspark-confidence-scheduling.md @@ -537,11 +537,13 @@ no-spec 应降级而非死锁、字节估算器补 draft 窗口放大）待报 u - 硬件/形态：DeepSeek-V4-Pro-DSpark，8×B300 单节点，`--tp_size 8 --ep_size 8`（DEP8， `enable_attention_dp: true`），`--max_batch_size 256 --max_num_tokens 4096`。 - 数据集： - - **poetry**：单条固定英文 free-verse 长诗 prompt（memory/rivers/glass/time 主题、 - 要求 ≥8 个场景、不许复用意象），复制 N 份。短输入（~60 tok）+ 长输出（实测 - ~760 tok/req，自然 EOS）——短 ISL / 长 OSL 的极端。 - - **arena**：arena-hard-v0.1 的 500 条真实问题（jsonl，取 `prompt` 字段，N>500 时循环 - 复用）。中等 ISL / 长输出的多样负载。 + - **poetry**：单条固定英文 free-verse 长诗 prompt（逐字原文见 + `docs/dspark_bench_assets/poetry_prompt.txt`），复制 N 份。短输入（~60 tok）+ + `max_tokens=768, temperature=0.7`（实测均值 ~760 tok/req，基本打满上限）—— + 短 ISL / 长 OSL 的极端。 + - **arena**：arena-hard-v0.1 全集 500 条真实问题（jsonl，取 `prompt` 字段，N>500 时 + 循环复用），同 `max_tokens=768, temperature=0.7`。中等 ISL / 长输出的多样负载。 + - 打流客户端原件：`docs/dspark_bench_assets/p25fix_bench.py`。 - 客户端 burst-drain：一次性下发全部 N（512/1024）个请求，计时到全部完成，n=10 轮取中位。 **注意：该协议下 99% 的步是纯 gen 步**（prefill 波在头几十步内结束）——这是收益成立的 regime 前提。 From 6ac04b5eade0140416c911d3cc21789d0b719ba7 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 09:14:30 -0700 Subject: [PATCH 159/177] [None][doc] Disagg repro launched (poetry line); request throughput_1k extract asset (gated dataset, 401 on our side) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index aca647a1d3ef..648f1fa6b147 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -867,6 +867,12 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - **矩阵**:三负载 {poetry, arena, throughput_1k} × lbs {32,64,128,256} × 臂 {notrim, frac04, scheduled};**SPS 与 STS 分负载重采**(A:KV 深度 ~100-800 vs B:~1110,跨用错价;STS 是数据语义×采样温度的函数,A temp0.7 / B greedy);对表纪律:A 主臂裸 sigmoid(与主线 +5.9~+24.7% 同协议),STS 为 A 的消融格、B 的必挂项。**B 在 disagg 翻正与否 = "混合步成本表 ctx 盲"机理的直接检验**——这是本战役最有信息量的一格。 - **➡️ 请主线**:提供 poetry 的 **verbatim prompt 原文**(provenance 只有描述:memory/rivers/glass/time、≥8 场景、不复用意象;我 step3 自造的 caravan/steppe 不是同一首,接受率分布不可比,对表 A 必须同文)。最好直接把数据集构造脚本或 jsonl 推上来;arena(arena-hard-v0.1 取 prompt 字段)与 throughput_1k(iputterman/speed parquet)provenance 已够我自建。 +## 【08-11 09:40】disagg 复现开工(laliao 侧):poetry/arena 线已发射;➡️ 请主线推 throughput_1k 抽取件 + +- poetry 原文与 768-cap/temp0.7 修正已收,verbatim prompt 已铸成 trtllm_custom 数据集(poetry/arena 各 8192 行);客户端(benchmark_serving)已开关化:BENCH_KEEP_EOS=1(自然 EOS)+ BENCH_TEMPERATURE=0.7,对齐 burst 协议。 +- bring-up R1 尸检:非传输非代码——**gen DEP8 每 rank 85GB 权重装载撞 harness 1800s 健康等待超时**(ctx 健康、gen 差几分钟;DEP32 每 rank 21GB 所以旧 agentx 跑通)。wait_server 已抬 3600s。R2(poetry lbs64,ctx-spec-parity + 主线钉法 maxbs=lbs + 2× 超发)已在队。 +- **➡️ 请主线**:`iputterman/speed` 是私有 gated,我侧 401,贵区 dspark-runs 路径在本文件系统亦不可达。请把 throughput_1k 的 **前 1024 条抽取件**(file 序,turns[0],建议直接 trtllm_custom jsonl,~5MB)推到 `docs/dspark_bench_assets/`,负载 B 线即可开工——B 在 disagg 翻正与否是"混合步成本表 ctx 盲"机理的直接检验,是三负载矩阵里信息量最大的一格。 + # 分析 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 7374e6519b097d7d3ac8cf5b687e41f27ac4d00f Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 09:18:00 -0700 Subject: [PATCH 160/177] [None][feat] Experimental mixed-step trim guard (TLLM_DSPARK_SKIP_MIXED_TRIM): decline trimming on prefill-carrying steps; A/B shows it recovers the mixed-workload loss (-8.0% -> -1.0% at 32/rank) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 13 +++++++++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 7 ++++++- tensorrt_llm/_torch/speculative/dspark_verify.py | 10 ++++++++++ 3 files changed, 29 insertions(+), 1 deletion(-) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 648f1fa6b147..c771716e585d 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -889,3 +889,16 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - **勘误(feature doc 同步修正)**:此前"自然 EOS ~760 tok"描述有误——实测均值 ~760 是 **768 上限基本打满**(poetry 几乎全部 length-capped),对表时请按 out768 记。 - arena/throughput_1k 你侧 provenance 判断正确,无需我方再推(arena_questions.jsonl 即 arena-hard-v0.1 全集 500 题;parquet 路径已在 feature doc)。 - 你侧 harness 若单 job 同 NVL segment 分配,MNNVL 坑理论上天然免疫(IMEX 就位)——但**请留 bring-up 判决为准**,失败签名是 ctx 侧 NIXL wait ~336s 后 FAILURE + 雪崩秒拒;R5b(主线,MNNVL=n + 600s/900s 超时)六对在跑,出数后两边可直接对表。 + +## 【08-11 09:50】混合步裁剪 A/B 判决(mixab 三臂,agg throughput_1k):机理实锤,门控回收损失 + +用户质疑"混合步不进 ragged path"→ 代码取证(planner 无 ctx gate,混合步照裁、eager 执行,成本模型 ctx 盲)→ 加一行实验门控 `TLLM_DSPARK_SKIP_MIXED_TRIM=1`(rank 本地批含 ctx 即 decline,现有 all-or-nothing 门自动组一致)→ 三臂同协议实测(各自独立节点,3 reps 中位): + +| 目标 bs/rank | notrim | sched | sched+跳混合步 | +|---|---|---|---| +| 32 | 1,858 | 1,710(**−8.0%**) | 1,840(**−1.0%**) | +| 128 | 3,798 | 3,865(≈0) | 3,876(≈0) | + +- **机理坐实**:t32 门控命中 399/736 步(54%,恰为混合步密度),trim 30.5%→5.7%,accept_len 1.38→1.54,吞吐损失 −8.0%→−1.0% ——**损失几乎全部来自混合步上的裁剪**(纯 gen 步的裁剪在本负载 ≈ 中性); +- 产品方向确认:agg 形态需要混合步自保(默认门控或成本表加 ctx 维度);disagg(纯 gen)天然免疫; +- **bs 配方对照(用户质疑成立)**:t128 的 maxbs=2×+mnt=16384 协议没把 bs 钉到 128(gen 步中位 ~89),且 notrim 吞吐比旧协议(maxbs=128/mnt=8192,4,336)低 ~12%——mnt 抬升让准入步更重,纯亏。两旋钮处方已从 feature doc 撤下。 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 2d345cbedd4b..b1a058479d5a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3387,7 +3387,12 @@ def _dspark_confidence_draft_len(self, # would re-rank a batch it was never computed for. self.model_engine._dspark_device_budget = None if compute_windows: - if planner.device_windows and mode.trims_submitted_tokens: + if planner.skip_mixed_trim and scheduled_batch.context_requests: + # Mixed-step guard (experiment): prefill work shares this + # step, so the pure-decode cost table cannot price the trim. + # Declining here rides the normal all-or-nothing gate below. + fallback_reason = "mixed_step_declined" + elif planner.device_windows and mode.trims_submitted_tokens: # Device-window mode: the host decides only the CAPACITY, from # the lag-2 snapshot; the returned lens are a shape split used # for the cross-rank agreement and the fit. The true windows diff --git a/tensorrt_llm/_torch/speculative/dspark_verify.py b/tensorrt_llm/_torch/speculative/dspark_verify.py index 635483c566a2..d3b530cbb3fc 100644 --- a/tensorrt_llm/_torch/speculative/dspark_verify.py +++ b/tensorrt_llm/_torch/speculative/dspark_verify.py @@ -59,6 +59,13 @@ #: once at init on every rank. DEVICE_WINDOWS_ENV = "TLLM_DSPARK_DEVICE_WINDOWS" +#: Experiment-only mixed-step guard: when set, a rank whose scheduled batch +#: carries context (prefill) requests declines to trim that step, and the +#: cross-rank all-or-nothing gate reverts the whole group to the full block. +#: The cost table is profiled on pure-decode steps and cannot price a trim +#: whose step time is dominated by prefill compute; this isolates that effect. +SKIP_MIXED_TRIM_ENV = "TLLM_DSPARK_SKIP_MIXED_TRIM" + def device_windows_enabled() -> bool: """Whether device-side window selection is on (process-constant env).""" @@ -128,6 +135,9 @@ def __init__( # Two pinned buffers alternate; the pair we are NOT writing this step # is the older, guaranteed-landed one. self.device_windows = os.environ.get(DEVICE_WINDOWS_ENV, "0") == "1" + # Constant per process like the pin: every rank reads the same env at + # spawn, so the per-step decline needs no extra collective. + self.skip_mixed_trim = os.environ.get(SKIP_MIXED_TRIM_ENV, "0") == "1" if self.device_windows and self._forced_verify_len is not None: raise ValueError( f"{DEVICE_WINDOWS_ENV}=1 is incompatible with " From 348b01eea20fadbf8d418deb95f0e08ad0e80125 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 10:22:24 -0700 Subject: [PATCH 161/177] [None][doc] Same-node confirm locks mixed-step verdict (+7-8pp recovery); disagg transfer final recipe (UCX_TLS drop-IB); gen-starvation physics and the OSL lever Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index c771716e585d..93da003868bc 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -902,3 +902,10 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - **机理坐实**:t32 门控命中 399/736 步(54%,恰为混合步密度),trim 30.5%→5.7%,accept_len 1.38→1.54,吞吐损失 −8.0%→−1.0% ——**损失几乎全部来自混合步上的裁剪**(纯 gen 步的裁剪在本负载 ≈ 中性); - 产品方向确认:agg 形态需要混合步自保(默认门控或成本表加 ctx 维度);disagg(纯 gen)天然免疫; - **bs 配方对照(用户质疑成立)**:t128 的 maxbs=2×+mnt=16384 协议没把 bs 钉到 128(gen 步中位 ~89),且 notrim 吞吐比旧协议(maxbs=128/mnt=8192,4,336)低 ~12%——mnt 抬升让准入步更重,纯亏。两旋钮处方已从 feature doc 撤下。 + +## 【08-11 10:30】混合步判决同节点确认 + disagg 传输终局 + gen 饥饿物理(主线) + +- **混合步判决锁定**(同节点背靠背、双节点臂序对调):748120 sched 1721 vs skipmixed 1859(+8.0%);748121 反序 1739 vs 1855(+6.7%);skipmixed ≈ notrim(1858)打平。t32 的损失全部来自混合步裁剪,`TLLM_DSPARK_SKIP_MIXED_TRIM` 完全回收。 +- **disagg 传输六轮终局(R6 全绿,fails=0)**:最后一刀是 **`UCX_TLS=cuda_copy,cuda_ipc,sm,self,tcp`**(去 IB verbs;v2 manager 的 VMM/fabric 显存走 IB RDMA 注册有问题,MNNVL=n 必要但不充分——金丝雀"零失败"是提前杀进程的假阴性,R5b 同配置 40-80s 断开雪崩)。CI 对 GB300 的 UCX_TLS 处方与此完全一致(test_disaggregated.py)。完整配方五件套:pmix + NIXL/PYTHON + MNNVL=n + UCX_TLS 去 IB + 600s/900s 超时。 +- **gen 饥饿物理**:throughput_1k(ISL:OSL≈17:1)下 1 ctx 喂不饱 1 gen——gen 实测 bs 仅 1-5/rank,1:1 矩阵测的是 ctx 管道(lbs128 仅 ~1.1k tok/s)。QA 8k1k 用 ctx8:gen1;17:1 需要更高配比。**解法:客户端 OSL 旋钮**(`SPEEDBENCH_MAX_TOKENS`)拉长请求寿命——OSL=1024 时 gen bs 90 秒即爬到 22-25/rank 并继续上行,S/N 128 对的配对补测(同调用 3 稳态 reps)在收。 +- 副产品:OSL 旋钮同样解锁了 **agg 上实测 bs=128 的表采集**(v3 当年 128 档采不到的"早完成者"问题源于 OSL=64 寿命太短)——v4 采集(OSL 512、client 1024、pin 5/4/3/2)已在 mixab 空闲节点上跑。 From 99ab172203eb6ae6365260df8df2ece39069e402 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 10:24:41 -0700 Subject: [PATCH 162/177] [None][doc] First valid disagg cell: pure-gen at shallow bs still -8.8%; unified law = pure-gen density x batch depth Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 93da003868bc..142666458abe 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -909,3 +909,10 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - **disagg 传输六轮终局(R6 全绿,fails=0)**:最后一刀是 **`UCX_TLS=cuda_copy,cuda_ipc,sm,self,tcp`**(去 IB verbs;v2 manager 的 VMM/fabric 显存走 IB RDMA 注册有问题,MNNVL=n 必要但不充分——金丝雀"零失败"是提前杀进程的假阴性,R5b 同配置 40-80s 断开雪崩)。CI 对 GB300 的 UCX_TLS 处方与此完全一致(test_disaggregated.py)。完整配方五件套:pmix + NIXL/PYTHON + MNNVL=n + UCX_TLS 去 IB + 600s/900s 超时。 - **gen 饥饿物理**:throughput_1k(ISL:OSL≈17:1)下 1 ctx 喂不饱 1 gen——gen 实测 bs 仅 1-5/rank,1:1 矩阵测的是 ctx 管道(lbs128 仅 ~1.1k tok/s)。QA 8k1k 用 ctx8:gen1;17:1 需要更高配比。**解法:客户端 OSL 旋钮**(`SPEEDBENCH_MAX_TOKENS`)拉长请求寿命——OSL=1024 时 gen bs 90 秒即爬到 22-25/rank 并继续上行,S/N 128 对的配对补测(同调用 3 稳态 reps)在收。 - 副产品:OSL 旋钮同样解锁了 **agg 上实测 bs=128 的表采集**(v3 当年 128 档采不到的"早完成者"问题源于 OSL=64 寿命太短)——v4 采集(OSL 512、client 1024、pin 5/4/3/2)已在 mixab 空闲节点上跑。 + +## 【08-11 10:50】disagg 首个有效数据格(OSL1024,128 对配对稳态):纯 gen 浅 bs 仍 −8.8%;统一律成形 + +- **数字**(饱和管道 3 稳态 reps,0 fails):notrim 中位 17,294 vs sched 15,770 tok/s = **−8.8%**,三 rep 全序一致(跨对比较,各自全新 setup)。 +- **机理对齐**:bench 窗口 gen 实际 bs 仅 **29-49/rank**(1 ctx 喂 1 gen 的上限,即使 OSL=1024);这落在 B300 θ 平坦区(#28 的负转平边界 32-64/rank)。sched 裁了 12.3% token(32% 步选低档),regret 18.9%,accept 2.27(书文续写高接受,被裁 token 更贵)——**纯 gen 步不是充分条件,bs 深度才是**。 +- **统一律**:裁剪收益 = f(纯 gen 步密度 × bs 深度)。agg×本负载:混合步吃掉收益(门控可回收到 0);disagg×本负载:纯 gen 但浅 bs,照亏;poetry/arena×agg burst:纯 gen 且深 bs(128-256/rank 满批),+5.9~+24.7%。三个 regime 一张图。 +- **下一步(节点已到期,明日执行)**:disagg 要进"顺风局"需 **ctx≥2:gen1 + OSL 1024**(把 gen bs 喂到 128/rank)+ **v4 实测表**(OSL 杠杆解锁的 agg bs=128 采集在跑);那一格才是"disagg 能否翻正"的真判决。 From d835dd7807e640dd976544825eec883a24022dca Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 10:46:54 -0700 Subject: [PATCH 163/177] [None][doc] Independent agg verdict: frac04 mechanical trim negative across the board on Pro@DEP8; disagg ctx OOM root cause (tp8 fix) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 142666458abe..d8e0ca4a8873 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -873,6 +873,21 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - bring-up R1 尸检:非传输非代码——**gen DEP8 每 rank 85GB 权重装载撞 harness 1800s 健康等待超时**(ctx 健康、gen 差几分钟;DEP32 每 rank 21GB 所以旧 agentx 跑通)。wait_server 已抬 3600s。R2(poetry lbs64,ctx-spec-parity + 主线钉法 maxbs=lbs + 2× 超发)已在队。 - **➡️ 请主线**:`iputterman/speed` 是私有 gated,我侧 401,贵区 dspark-runs 路径在本文件系统亦不可达。请把 throughput_1k 的 **前 1024 条抽取件**(file 序,turns[0],建议直接 trtllm_custom jsonl,~5MB)推到 `docs/dspark_bench_assets/`,负载 B 线即可开工——B 在 disagg 翻正与否是"混合步成本表 ctx 盲"机理的直接检验,是三负载矩阵里信息量最大的一格。 +## 【08-11 11:30】agg 独立复证判决(laliao 侧):frac04 机械裁在 Pro@DEP8 全线为负;Pro 翻案的承重臂只能是 scheduled + +修复后代码(b87ea28bf3 已 overlay)、maxbs=256+显式 max_seq_len、到达性验证干净(iter-log mode=64/128 满档)、双独立节点对 × n=5: + +| 负载 | 全局 bs | 对A frac04 vs notrim | 对B | +|---|---|---|---| +| poetry | 1024 | −5.0% | −2.6% | +| poetry | 512 | −4.8% | −3.3% | +| arena | 1024 | **−16.9%** | **−14.9%** | +| arena | 512 | **−18.8%** | −15.6% | + +- 与主线不矛盾:frac04 转正(+5.3/+7.6%)是 **Flash DEP4** 的数据;Pro 终局从未报 frac04。本判决补上缺格:**Pro 上机械 40% 裁剪稳定亏钱**(accept 损失定价高),Pro 的收益只能靠"按请求该裁才裁"的 scheduled+修复后表。arena 亏得最狠(−15~19%)与 Flash 侧 arena≈0 的方向一致放大——多样负载对机械裁最不友好。 +- 复证下一步已在跑:我环境重采 postfix Pro 表(pin{2,3,4,5} × conc{512,1024,2048} poetry,--from-iter-log),然后 scheduled vs notrim 终局对表主线 +9.9~+24.7%。 +- disagg 进展:R2 死因=ctx 挂 spec 后 host OOM(tp4 单节点 4 rank 双份权重 staging 撑爆 host RAM;slurm oom_kill 实锤)——**ctx 已改 tp8×2 节点**(与 gen 同规格),R3 在队。坑登记:Pro 这个体量,disagg ctx 至少 tp8(或关 TRTLLM_PINNED_WEIGHT_STAGING 试 tp4,未验)。 + # 分析 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From eadd65d1f8d6032a2db3c1f54bc952d028dd4282 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:02:46 -0700 Subject: [PATCH 164/177] [None][doc] v4 measured table + mixed-step gate turn agg t128 from -1.2% to +0.4~+1.7% (same-node paired, both orders) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index d8e0ca4a8873..e455ea33eea3 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -931,3 +931,9 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - **机理对齐**:bench 窗口 gen 实际 bs 仅 **29-49/rank**(1 ctx 喂 1 gen 的上限,即使 OSL=1024);这落在 B300 θ 平坦区(#28 的负转平边界 32-64/rank)。sched 裁了 12.3% token(32% 步选低档),regret 18.9%,accept 2.27(书文续写高接受,被裁 token 更贵)——**纯 gen 步不是充分条件,bs 深度才是**。 - **统一律**:裁剪收益 = f(纯 gen 步密度 × bs 深度)。agg×本负载:混合步吃掉收益(门控可回收到 0);disagg×本负载:纯 gen 但浅 bs,照亏;poetry/arena×agg burst:纯 gen 且深 bs(128-256/rank 满批),+5.9~+24.7%。三个 regime 一张图。 - **下一步(节点已到期,明日执行)**:disagg 要进"顺风局"需 **ctx≥2:gen1 + OSL 1024**(把 gen bs 喂到 128/rank)+ **v4 实测表**(OSL 杠杆解锁的 agg bs=128 采集在跑);那一格才是"disagg 能否翻正"的真判决。 + +## 【08-11 11:40】v4 实测表 + 混合步门控:agg t128 翻回打平/微正(当日收官) + +- **v4 表**(`align_pro_table_v4.json`,与 v3 的 32/64 样本合并拟合):OSL=512 长寿命客户端在 agg 稳住 bs=128 纯 gen 步(1,385 iter 实测样本,"早完成者"问题根治),θ 实测到 768 token 工作点。**v3 线性外推在 M=768 高估 ~17%**(≈63 外推 vs 53.8ms 实测)——bs128 过度裁剪的定量来源;α(128)=5.0ms 补齐。 +- **终局配对判决**(同节点背靠背、双向臂序,原始对齐协议 maxbs=128/mnt=8192/client=1024):notrim 4,306/4,293 vs sched(v4+`SKIP_MIXED_TRIM`)4,322/4,365 = **+0.4%/+1.7%**;对照原矩阵同格 −1.2%(v2 表、无门控)。**修正表 + 门控把本负载 agg 从负拉回打平/微正**;上限受限于 50% 混合步不可裁 + 纯 gen 步等效 bs ~89 的浅 θ。 +- t32 档未复测(门控主导,预期 ≈0);disagg 深 bs(ctx≥2:gen1 + OSL1024 + v4)仍是唯一可能的显著正收益格,待明日节点。 From b20ce02aff25cb5e8c4c2b668dd626eab2989009 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 14:32:06 -0700 Subject: [PATCH 165/177] [None][doc] Disagg bring-up gates pass (77k full-batch steps, transfer clean); STS-violation found in first sched replication; own postfix table cross-validates Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index e455ea33eea3..62a7b2aa6100 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -888,6 +888,14 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - 复证下一步已在跑:我环境重采 postfix Pro 表(pin{2,3,4,5} × conc{512,1024,2048} poetry,--from-iter-log),然后 scheduled vs notrim 终局对表主线 +9.9~+24.7%。 - disagg 进展:R2 死因=ctx 挂 spec 后 host OOM(tp4 单节点 4 rank 双份权重 staging 撑爆 host RAM;slurm oom_kill 实锤)——**ctx 已改 tp8×2 节点**(与 gen 同规格),R3 在队。坑登记:Pro 这个体量,disagg ctx 至少 tp8(或关 TRTLLM_PINNED_WEIGHT_STAGING 试 tp4,未验)。 +## 【08-11 14:30】disagg bring-up 门禁全过 + agg 复证首轮的 STS 违规教训(laliao 侧) + +- **disagg R4 门禁全过**(poetry lbs64,ctx tp8×2 节点 + gen DEP8×2 + 两侧显式 max_seq_len 8192):**77,823 步 mode=64 逐字满档**、KV 传输全程行使零 FAILURE、无 IMA/OOM。单 job 同 NVL segment(--segment 分配)下 **MNNVL 免疫假设证实**——五件套里的 MNNVL=n / UCX_TLS 去 IB 是跨 job 跨机架的处方,同机架单 job 不需要。仅墙钟 2.5h 被截尾,已抬 4h。 +- **两个新坑入册**:①ctx(tp4 单节点)挂 spec 后 host OOM(双份权重 staging;Pro 体量 disagg ctx 至少 tp8);②harness 自动推导 ctx `max_seq_len=isl+40`,poetry 短 ISL 下= **104 < 窗口池几何 133**,稀疏 topk 被钳成非 4 倍数 → `SparseAttnTopK must be a multiple of 4` 在 ctx warmup 验证步断言(ctx+spec+短 ISL 的新角落,ISL≈1110 的负载永不触发)。修法都指回协议:**两侧显式 8192**。 +- **我环境的 postfix Pro 表出炉**(pin{2..5}×conc{512,1024,2048} poetry,--from-iter-log):fixed 35.1ms、α(256)=3.8ms、θ 单调(192→7.7 … 768→40.9 … 1536→78.9ms);**bs128 步时 54% 可裁 / w5→w2 省 27%,与主线表(56% / 26.7%)跨环境互证**。login 无 tensorrt_llm 的拟合配方:venv numpy + 包级 stub runner(绕开重 __init__)。 +- **agg scheduled 复证首轮为负(−4~−12%)——但发现协议违规**:我的 config 挂了旧 `sts_v1_reconstructed.json`,主线 Pro 的 +9.9~+24.7% 是**裸 sigmoid**;过时温度扭曲排序即乱裁。v2(去 STS,协议逐字对齐)已重排。⚠️ 若 v2 仍负,则分歧升级为真(候选:表绝对值差异/图集 [1,8,32,64,128] vs 全阶梯/客户端 burst 语义),将请主线共享 `pro_sched_arm.yaml` 逐键对表。 +- stage A 第一波在队:poetry lbs{64,128} × {pin2..5, notrim+STS 采集},disagg 表与 STS 原地重采。 + # 分析 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 1a487278678f4778ce19661f982695a92495bfa1 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 20:40:52 -0700 Subject: [PATCH 166/177] [None][doc] Publish exact profinal arm yamls and postfix table for key-by-key protocol alignment (collaborator agg replication v2) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- .../postfix_pro_table.json | 302 ++++++++++++++++++ docs/dspark_bench_assets/pro31_iterlog.yaml | 21 ++ docs/dspark_bench_assets/pro_sched_arm.yaml | 28 ++ 3 files changed, 351 insertions(+) create mode 100644 docs/dspark_bench_assets/postfix_pro_table.json create mode 100644 docs/dspark_bench_assets/pro31_iterlog.yaml create mode 100644 docs/dspark_bench_assets/pro_sched_arm.yaml diff --git a/docs/dspark_bench_assets/postfix_pro_table.json b/docs/dspark_bench_assets/postfix_pro_table.json new file mode 100644 index 000000000000..df7fe7807dae --- /dev/null +++ b/docs/dspark_bench_assets/postfix_pro_table.json @@ -0,0 +1,302 @@ +{ + "token_counts": [ + 96, + 128, + 192, + 256, + 320, + 384, + 640, + 768 + ], + "step_time_ms": [ + 5.332674, + 7.125582, + 19.072141, + 22.564258, + 30.341472, + 32.343115, + 54.299678, + 61.859097 + ], + "fixed_overhead_ms": 39.296781, + "batch_sizes": [ + 32, + 64, + 128 + ], + "batch_overhead_ms": [ + 0.0, + 3.221512, + 9.122014 + ], + "_meta": { + "lookup": "interp", + "encoding": "decomposed", + "note": "step_time_ms is theta(M) only -- the trimmable term. The non-trimmable floor is fixed_overhead_ms + batch_overhead_ms[bs]. Do not read step_time_ms as whole-step time.", + "token_unit": "total_verify_tokens = num_requests * (verify_len + 1); the bonus token is included", + "cells": [ + { + "batch_size": 32, + "verify_len": 2, + "total_verify_tokens": 96, + "step_time_ms": 44.629455, + "num_samples": 114, + "p10_ms": 42.656446, + "p90_ms": 45.921993 + }, + { + "batch_size": 32, + "verify_len": 3, + "total_verify_tokens": 128, + "step_time_ms": 46.422362, + "num_samples": 252, + "p10_ms": 44.540524, + "p90_ms": 48.05913 + }, + { + "batch_size": 32, + "verify_len": 5, + "total_verify_tokens": 192, + "step_time_ms": 58.368921, + "num_samples": 71, + "p10_ms": 56.812048, + "p90_ms": 60.264587 + }, + { + "batch_size": 64, + "verify_len": 2, + "total_verify_tokens": 192, + "step_time_ms": 61.590433, + "num_samples": 87, + "p10_ms": 60.002756, + "p90_ms": 63.626337 + }, + { + "batch_size": 64, + "verify_len": 3, + "total_verify_tokens": 256, + "step_time_ms": 65.08255, + "num_samples": 77, + "p10_ms": 63.558149, + "p90_ms": 67.810392 + }, + { + "batch_size": 64, + "verify_len": 4, + "total_verify_tokens": 320, + "step_time_ms": 72.859764, + "num_samples": 54, + "p10_ms": 71.078515, + "p90_ms": 74.275923 + }, + { + "batch_size": 64, + "verify_len": 5, + "total_verify_tokens": 384, + "step_time_ms": 74.861407, + "num_samples": 70, + "p10_ms": 72.710967, + "p90_ms": 76.689982 + }, + { + "batch_size": 128, + "verify_len": 2, + "total_verify_tokens": 384, + "step_time_ms": 80.761909, + "num_samples": 81, + "p10_ms": 78.371763, + "p90_ms": 82.598686 + }, + { + "batch_size": 128, + "verify_len": 3, + "total_verify_tokens": 512, + "step_time_ms": 90.330839, + "num_samples": 47, + "p10_ms": 87.460041, + "p90_ms": 92.462635 + }, + { + "batch_size": 128, + "verify_len": 4, + "total_verify_tokens": 640, + "step_time_ms": 102.718472, + "num_samples": 46, + "p10_ms": 99.931121, + "p90_ms": 105.115652 + }, + { + "batch_size": 128, + "verify_len": 5, + "total_verify_tokens": 768, + "step_time_ms": 110.277891, + "num_samples": 53, + "p10_ms": 108.377361, + "p90_ms": 113.493729 + } + ], + "fit": { + "max_abs_residual_ms": 0.0, + "max_rel_residual": 0.0, + "additive_tolerance": 0.02, + "riser_tolerance": 0.01, + "min_riser_ms": 0.728598, + "typical_step_ms": 72.859764, + "warnings": [] + }, + "block_size": 5, + "min_verify_len": 1, + "modal_batch_size": 128, + "warmup_steps_discarded": 128, + "statistic": "median", + "source": "iter_log", + "logs": [ + "/lustre/fsw/coreai_comparch_trtllm/laliao/dspark-runs/protbl_744692_filtered.log", + "/lustre/fsw/coreai_comparch_trtllm/laliao/dspark-runs/protbl_744693_filtered.log" + ], + "recommended_confidence_verify_len_tiers": [ + 1, + 3, + 5 + ], + "recommended_tiers_note": "Nothing reads this from the table file. Copy it into DSparkDecodingConfig.confidence_verify_len_tiers, which is also what sizes the captured CUDA graph set.", + "diagnostics": { + "step_time_spread_by_batch_size": { + "32": 0.307856, + "64": 0.507969, + "128": 0.553581 + }, + "best_step_time_spread": 0.553581, + "profitability_probe": [ + { + "batch_size": 32, + "acceptance_rate": 0.3, + "chosen_verify_len": 3, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 32, + "acceptance_rate": 0.5, + "chosen_verify_len": 3, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 32, + "acceptance_rate": 0.7, + "chosen_verify_len": 3, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 32, + "acceptance_rate": 0.8, + "chosen_verify_len": 3, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 32, + "acceptance_rate": 0.9, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 32, + "acceptance_rate": 0.95, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 64, + "acceptance_rate": 0.3, + "chosen_verify_len": 1, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 64, + "acceptance_rate": 0.5, + "chosen_verify_len": 1, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 64, + "acceptance_rate": 0.7, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 64, + "acceptance_rate": 0.8, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 64, + "acceptance_rate": 0.9, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 64, + "acceptance_rate": 0.95, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 128, + "acceptance_rate": 0.3, + "chosen_verify_len": 1, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 128, + "acceptance_rate": 0.5, + "chosen_verify_len": 1, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 128, + "acceptance_rate": 0.7, + "chosen_verify_len": 3, + "max_tier": 5, + "trims": true + }, + { + "batch_size": 128, + "acceptance_rate": 0.8, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 128, + "acceptance_rate": 0.9, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + }, + { + "batch_size": 128, + "acceptance_rate": 0.95, + "chosen_verify_len": 5, + "max_tier": 5, + "trims": false + } + ], + "trims_somewhere": true + } + } +} diff --git a/docs/dspark_bench_assets/pro31_iterlog.yaml b/docs/dspark_bench_assets/pro31_iterlog.yaml new file mode 100644 index 000000000000..36ef5e7d8911 --- /dev/null +++ b/docs/dspark_bench_assets/pro31_iterlog.yaml @@ -0,0 +1,21 @@ +speculative_config: + decoding_type: DSpark + max_draft_len: 5 + speculative_model: /lustre/fsw/coreai_comparch_trtllm/laliao/llm-models/DeepSeek-V4-Pro-DSpark + enable_confidence_scheduling: false +moe_config: + backend: TRTLLM +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.5 +enable_attention_dp: true +cuda_graph_config: + enable_padding: true + batch_sizes: + - 1 + - 8 + - 32 + - 64 + - 128 +disable_overlap_scheduler: false +print_iter_log: true diff --git a/docs/dspark_bench_assets/pro_sched_arm.yaml b/docs/dspark_bench_assets/pro_sched_arm.yaml new file mode 100644 index 000000000000..370a68531eb2 --- /dev/null +++ b/docs/dspark_bench_assets/pro_sched_arm.yaml @@ -0,0 +1,28 @@ +speculative_config: + decoding_type: DSpark + max_draft_len: 5 + speculative_model: /lustre/fsw/coreai_comparch_trtllm/laliao/llm-models/DeepSeek-V4-Pro-DSpark + enable_confidence_scheduling: true + enable_ragged_verify: true + confidence_verify_len_tiers: + - 2 + - 3 + - 4 + - 5 + confidence_sps_table_path: /lustre/fsw/coreai_comparch_trtllm/laliao/dspark-runs/postfix_pro_table.json +moe_config: + backend: TRTLLM +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.5 +enable_attention_dp: true +cuda_graph_config: + enable_padding: true + batch_sizes: + - 1 + - 8 + - 32 + - 64 + - 128 +disable_overlap_scheduler: false +print_iter_log: true From 6fce248053ac84a20f18db2f81dddd0449c9a7cb Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 22:46:02 -0700 Subject: [PATCH 167/177] [None][doc] Agg replication closed: client-semantics conviction, unified law verified 3 ways; disagg gen no-commit regression suspect (D1 bisect running) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 62a7b2aa6100..3ce67057c1bc 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -896,6 +896,26 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - **agg scheduled 复证首轮为负(−4~−12%)——但发现协议违规**:我的 config 挂了旧 `sts_v1_reconstructed.json`,主线 Pro 的 +9.9~+24.7% 是**裸 sigmoid**;过时温度扭曲排序即乱裁。v2(去 STS,协议逐字对齐)已重排。⚠️ 若 v2 仍负,则分歧升级为真(候选:表绝对值差异/图集 [1,8,32,64,128] vs 全阶梯/客户端 burst 语义),将请主线共享 `pro_sched_arm.yaml` 逐键对表。 - stage A 第一波在队:poetry lbs{64,128} × {pin2..5, notrim+STS 采集},disagg 表与 STS 原地重采。 +## 【08-11 22:30】agg 独立复证闭环:客户端语义定罪 + 统一律三面验证;disagg 发现 gen 不提交进度的疑似分支回归(laliao 侧) + +**agg 复证故事线完整闭环**(全部修复后代码、我自己重采的表、到达性逐字验证): + +| 轮 | 客户端 | scheduled vs notrim | 判读 | +|---|---|---|---| +| schedv2 | 3× 补位闭环(混合步 11%) | poetry −6.3/−6.5%,arena −7.7/−9.8% | 分歧真实但根因在负载形态 | +| **schedv3-burst** | **单波 burst-drain(混合步 ~1%,对齐主线协议)** | **poetry +4.6%/+7.8%**,arena +0.4/−2.4% | **定性复现主线翻案**;唯一变量是客户端 | +| schedv3-gate | 补位 + SKIP_MIXED_TRIM | 门控只回收 +0.5~+2.6pp | 补位损失大头 = 纯 gen 密度稀释(regime 本身),错误裁剪只占小头 | + +- 到达性与混合步密度均 iter-log 直测:两臂 gen 行数众数逐字打满(64/128),补位客户端 11% 的 gen 步带 prefill,burst ~1%。**统一律(收益=纯 gen 密度×bs 深度)在独立环境三面验证:同机同表,换负载形态即可让收益在 +7.8% 与 −10% 间摆动。** +- 遗留:poetry 幅度差(我 +7.8 vs 主线 +24.7 @1024;候选:表绝对步时 1.45× 差(76 vs 110ms bs128 满窗,KV 深度/采集负载不同)、图集 [1..256] 全阶梯 vs [1,8,32,64,128]、burst 计量细节);arena ≈0 vs 主线 +6.8/+5.9(请主线提供 arena 臂的 burst 语义与 OSL cap 细节对表)。 +- frac04 机械裁补格判决(前报)与 schedv3 合并看:**Pro 的收益只有 scheduled+修复后表+纯 gen regime 三者同时成立才兑现**。 + +**disagg 疑似分支回归(阻塞我方 disagg 战役,请主线协查)**: +- 症状:gen server 每步执行真实完整 verify+draft(102ms 设备时间、384 tok/步)但**永不提交进度**——KV util 0.005 平线 129k 步、active 恒 64/512、零完成;1024 请求完整走过 router→ctx→传输→gen,全链无错。pin/sched/FRAC 各臂均复现,与置信度调度开关无关。 +- 反证:我 08-10 旧镜像(b8bc42c3ae,同 harness 同 gen_only 模式,DSpark mtp6)完成流通;主线 Flash disagg cell 亦通。判别局 D1(同 poetry cell × 旧镜像无 overlay)在跑:通则回归定罪(faf2c60935+overlay 代码代的 disagg gen 接受提交路径),再二分 overlay 28 文件 vs sqsh 基底。 +- 附:到达性机制无懈可击(6 个僵尸 pin cell 的 mode 全部逐字满档),disagg 表/STS/三臂在 bug 后面排队。 +- ⚠️ 请主线看一眼:你们 disagg gen 侧用的 commit/镜像与 gen 第一步(消费 ctx 交接后)的提交路径,近期有无已知改动? + # 分析 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 0411ae1f7ca73588e04480ed00c71f339a3676d1 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 11 Aug 2026 23:56:57 -0700 Subject: [PATCH 168/177] [None][doc] Reply: arena burst semantics for protocol alignment; disagg gen commit-path provenance (branch HEAD proven working via osl1kb) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 3ce67057c1bc..3fd8623e3296 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -965,3 +965,16 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - **v4 表**(`align_pro_table_v4.json`,与 v3 的 32/64 样本合并拟合):OSL=512 长寿命客户端在 agg 稳住 bs=128 纯 gen 步(1,385 iter 实测样本,"早完成者"问题根治),θ 实测到 768 token 工作点。**v3 线性外推在 M=768 高估 ~17%**(≈63 外推 vs 53.8ms 实测)——bs128 过度裁剪的定量来源;α(128)=5.0ms 补齐。 - **终局配对判决**(同节点背靠背、双向臂序,原始对齐协议 maxbs=128/mnt=8192/client=1024):notrim 4,306/4,293 vs sched(v4+`SKIP_MIXED_TRIM`)4,322/4,365 = **+0.4%/+1.7%**;对照原矩阵同格 −1.2%(v2 表、无门控)。**修正表 + 门控把本负载 agg 从负拉回打平/微正**;上限受限于 50% 混合步不可裁 + 纯 gen 步等效 bs ~89 的浅 θ。 - t32 档未复测(门控主导,预期 ≈0);disagg 深 bs(ctx≥2:gen1 + OSL1024 + v4)仍是唯一可能的显著正收益格,待明日节点。 + +## 【08-12 00:10】主线 → laliao 侧:arena burst 语义对表 + disagg gen 提交路径 provenance(回应两问) + +**① arena 臂语义**(profinal,+6.8/+5.9% 的那轮)——原件都在 `docs/dspark_bench_assets/`: +- 客户端 `p25fix_bench.py` 逐字:arena-hard-v0.1 全集 500 题循环取 `prompt`;`max_tokens=768, temperature=0.7`,**无 ignore_eos**(自然停止允许,实测均值 ~760 ≈ 多数打满 cap);burst-drain = asyncio.gather 一次性 N 发全下、全完成计 wall;每 (dataset,bs) 首 rep 为 warm(rep=-1 丢弃)+ 5 计量 reps;指标 = Σcompletion_tokens/wall。 +- serve:`pro_sched_arm.yaml` 原件已推;graphs [1,8,32,64,128]+padding(非全阶梯),maxbs 256、mnt 4096。 +- 你 arena ≈0 vs 我 +6.8/+5.9 的候选还有一个:**臂间同节点纪律**——arena 的节点噪声你自己测过 3.8pp(08-09 工程沉淀),profinal 是同节点臂序对调;若你 arena 双臂跨节点,±4pp 即可吞掉信号。 + +**② disagg gen 提交路径 provenance**(你的 no-commit 回归协查): +- 我方 disagg gen = **sqsh `dspark-p0/trtllm.sqsh` + 全量仓库 mount**(非 overlay):R6/osl1k 轮 = 分支树 ≈ `9c5d203766`+门控补丁;今晚 deep 战役 = `1a48727867`。 +- **反证数据点:我方分支 HEAD 的 disagg gen 提交路径是通的**——08-11 上午 osl1kb 轮,gen 完整消费 ctx 交接并产出 1024-token 完成(每 rep ~1.04M token,fails=0,KV util 正常起落)。你的症状(verify+draft 真执行、active 恒 64/512、零完成、KV 平线)在我方全分支历史中未出现过。 +- 我方近期 commit 无有意改动 disagg gen 首步提交/接受路径(门控补丁只在 planner phase-1 加 decline)。你的 base 是 `faf2c60935`+28 文件 overlay——**怀疑点排序:overlay 组合(部分文件新、部分旧的 Franken 树)> sqsh 基底 > 分支本体**;D1 判别(旧镜像无 overlay)正是对的刀口。若 D1 通,建议直接换"全量 mount 我们分支 HEAD"复跑,跳过 28 文件二分。 +- 我方 deep 战役(3ctx:1gen,OSL1024,今晚)同路径在跑,出数即再添一个分支 HEAD 的通/不通数据点。 From 2754700e2035f5e0cb9bbdc4a87c22f6dd41fd6c Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 12 Aug 2026 02:06:38 -0700 Subject: [PATCH 169/177] [None][doc] Deep-bs disagg verdict: throughput_1k turns positive (+4.4% same-node paired) at pure-gen x deep-bs; unified law all four cells closed Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 3fd8623e3296..d024a19b1b5a 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -978,3 +978,12 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - **反证数据点:我方分支 HEAD 的 disagg gen 提交路径是通的**——08-11 上午 osl1kb 轮,gen 完整消费 ctx 交接并产出 1024-token 完成(每 rep ~1.04M token,fails=0,KV util 正常起落)。你的症状(verify+draft 真执行、active 恒 64/512、零完成、KV 平线)在我方全分支历史中未出现过。 - 我方近期 commit 无有意改动 disagg gen 首步提交/接受路径(门控补丁只在 planner phase-1 加 decline)。你的 base 是 `faf2c60935`+28 文件 overlay——**怀疑点排序:overlay 组合(部分文件新、部分旧的 Franken 树)> sqsh 基底 > 分支本体**;D1 判别(旧镜像无 overlay)正是对的刀口。若 D1 通,建议直接换"全量 mount 我们分支 HEAD"复跑,跳过 28 文件二分。 - 我方 deep 战役(3ctx:1gen,OSL1024,今晚)同路径在跑,出数即再添一个分支 HEAD 的通/不通数据点。 + +## 【08-12 02:00】深 bs disagg 终局(throughput_1k,3ctx:1gen,OSL1024,v4 表):统一律最后一格点亮,+4.4%(同节点配对) + +- **设计**:3 ctx(tp8)喂 1 gen(DEP8),OSL=1024 拉长寿命(稳态 gen bs 中位 116/rank,120-128 区间大量出现,吞吐为 1:1 disagg 的 10 倍),client 2048(2× 超发),两 quad 臂序对调,五件套 env,v4 实测表。 +- **数字**(稳态 3 reps,0 fails): + - 同节点配对(gen 749968,N→S):notrim med 37,588 vs sched med 39,247 = **+4.4%**; + - 跨 quad:quad2 sched med 43,242(其 notrim 臂死于 walltime)——表面 +15.0%,但同臂跨 quad 差 ~10%(43,242 vs 39,247)暴露强节点效应,**该数字不可引用,以同节点 +4.4% 为准**。 +- **判读**:throughput_1k 在"纯 gen 密度 × bs 深度"双高角落**首次转正**——统一律四格闭合(同负载:agg 混合步 −8%→门控 0%,disagg 浅 bs −8.8%,disagg 深 bs **+4.4%**)。幅度低于 poetry(+24.7%)符合预期:本负载 accept 更高(书文续写 2.3-2.6,被裁 token 更贵)且 sched 3 reps 内有爬升(36,965→40,773),更多 reps 可能收敛到更高中位。 +- 方法论注记:sched 臂 rep 间上行趋势 + 双臂 range 有重叠,+4.4% 的置信度为"方向可靠、幅度 ±3pp";加固需 n≥5 双向臂序(节点已到期,列为可选补测)。 From b3a116f8d28b59b5978a1fa09edc86a53d85a937 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 12 Aug 2026 04:48:48 -0700 Subject: [PATCH 170/177] [None][doc] Quantization economics verified (fp8 vs nvfp4 same-node A/B); freeze-bug bisect: overlay and ctx-spec exonerated, D4 running Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index d024a19b1b5a..003fab23d49b 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -916,6 +916,23 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - 附:到达性机制无懈可击(6 个僵尸 pin cell 的 mode 全部逐字满档),disagg 表/STS/三臂在 bug 后面排队。 - ⚠️ 请主线看一眼:你们 disagg gen 侧用的 commit/镜像与 gen 第一步(消费 ctx 交接后)的提交路径,近期有无已知改动? +## 【08-12 05:30】量化经济学实锤(fp8 量级轮)+ 冻结 bug 判别链收缩(laliao 侧) + +**fp8 量级复证**(发现我 07-30 已下完 fp8 原版检查点——.incomplete 残留是 hf-download 死垃圾,尺寸校验其实早通过;主线 pro_sched_arm.yaml 逐键 + 他们的表 + burst 客户端,同机同代码只换量化): + +| 负载/bs | nvfp4 | **fp8** | 主线 fp8 | +|---|---|---|---| +| poetry 512 | +4.6% | +3.9% | +9.9% | +| poetry 1024 | +7.8% | **+11.8%** | +24.7% | +| arena 512 | +0.4% | **+5.5%** | +6.8% | +| arena 1024 | −2.4% | +0.3% | +5.9% | + +- **verify 越贵收益越大的量化档位效应直接验证**(poetry@1024 +4pp、arena@512 从零翻正);逐键对表功不可没(头号差异就是 fp8 vs nvfp4 checkpoint + moe TRTLLM vs MEGAMOE_CUTEDSL)。产品含义:**报 DSpark 收益必须标注检查点量化;nvfp4 部署下天花板显著低于 fp8**。 +- 与主线剩余差距(11.8 vs 24.7)第二阶候选:N 臂定义(我 FRAC=1.0+SCHED_TIMING 有机器税,主线是关 confidence 开关——下轮用 CFG_NOHEAD 式 N 臂)、客户端实现(p25fix asyncio vs benchmark_serving)、warm-rep 剔除。 +- 注:本轮 arena 绝对吞吐显著高于 nvfp4 轮(63k vs 40k tok/s @1024)——burst 单波 completion-span 口径对 ISL 分布敏感,横向绝对值勿直接比,臂内 delta 有效。 + +**disagg 冻结 bug 判别链**:D2(去 overlay)= **ZOMBIE(20253 步)→ overlay 28 文件无罪**;D3(ctx 去 spec)= ctx 崩 v2 `dangling rawref` → 该路不通,反证 trap 4,ctx-spec 修复方向无罪;D4(旧镜像 b8bc42c3ae × ctx tp4 无 spec × 冻结组合 poetry/lbs64/benchmark_serving)在队——健康则回归区间锁定 b8bc42c3ae→faf2c60935 两构建之间,冻结则配置组合(客户端/streaming/lbs)定罪。 + # 分析 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 1c62fe02a89b1b91860ee06124601c3bb93ff67b Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 12 Aug 2026 06:21:47 -0700 Subject: [PATCH 171/177] [None][doc] 256-deep disagg verdict: +7.0/+15.6% both arm orders; v5 all-measured table to 1536 tokens; monotone depth trend; router backlog trap Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 003fab23d49b..296c9bfbe2f4 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -1004,3 +1004,14 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - 跨 quad:quad2 sched med 43,242(其 notrim 臂死于 walltime)——表面 +15.0%,但同臂跨 quad 差 ~10%(43,242 vs 39,247)暴露强节点效应,**该数字不可引用,以同节点 +4.4% 为准**。 - **判读**:throughput_1k 在"纯 gen 密度 × bs 深度"双高角落**首次转正**——统一律四格闭合(同负载:agg 混合步 −8%→门控 0%,disagg 浅 bs −8.8%,disagg 深 bs **+4.4%**)。幅度低于 poetry(+24.7%)符合预期:本负载 accept 更高(书文续写 2.3-2.6,被裁 token 更贵)且 sched 3 reps 内有爬升(36,965→40,773),更多 reps 可能收敛到更高中位。 - 方法论注记:sched 臂 rep 间上行趋势 + 双臂 range 有重叠,+4.4% 的置信度为"方向可靠、幅度 ±3pp";加固需 n≥5 双向臂序(节点已到期,列为可选补测)。 + +## 【08-12 06:00】256 档深 bs disagg 终局:+7.0%/+15.6%(双向同节点配对);v5 全实测表;bs 深度趋势完整 + +- **v5 表交付**(`align_pro_table_v5.json`,disagg 256 满批 pin 阶梯 + v4 样本合并):θ 实测覆盖 96→1536 token(θ(1536)=106.4ms),α(256)=14.4ms——32/64/128/256 全档无外推。 +- **数字**(4ctx:1gen,OSL1024,client 2048=1×,v5 表,0 fails,同 gen 节点背靠背、两组臂序对调): + - quint2(N→S):notrim 38,561 vs sched **44,561 = +15.6%**; + - quint1(S→N):sched 45,551 vs notrim 42,574 = **+7.0%**; + - sched 跨节点仅差 2%(45.6k/44.6k),notrim 波动较大(38.6k/42.6k,rep 内有爬升)→ 判读 **+7~+16%,中心 ~+11%**。 + - 实测 gen bs:notrim med 194(p90 236)、sched med 174——sched 排水快在飞更低,收益是在更浅 bs 下拿到的(锁同 bs 只会更高)。 +- **bs 深度趋势(throughput_1k,同代码同负载)**:29-49/rank **−8.8%** → ~116/rank **+4.4%** → ~174-236/rank **+7~+15.6%**——收益随深度单调走强,统一律的 bs 因子完整成像;256 档已接近 poetry 满批量级(+24.7%)。 +- **新坑入册(backlog)**:client 在飞 4096 > router(uvicorn)默认 backlog 2048 → 半数连接挂死、管道爬行(91% 失败,gen 只 40/rank)。128 档 client 恰好 2048 纯属侥幸。修法:client ≤2048,或给 `trtllm-serve disaggregated` 加 backlog 旋钮(上游项)。被污染的 4096 行已在 summary 以 tag 区分(*b = 干净复测)。 From 431925a1701f3feaf1255b44c119710aa72f13b4 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 12 Aug 2026 17:35:26 -0700 Subject: [PATCH 172/177] [None][doc] AA slo20 campaign: ragged path validated in production shape (trim 0.50); freeze convicted as short-ISL; slo20 SPS table v1 Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 296c9bfbe2f4..10353ea2bd5c 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -933,6 +933,15 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 **disagg 冻结 bug 判别链**:D2(去 overlay)= **ZOMBIE(20253 步)→ overlay 28 文件无罪**;D3(ctx 去 spec)= ctx 崩 v2 `dangling rawref` → 该路不通,反证 trap 4,ctx-spec 修复方向无罪;D4(旧镜像 b8bc42c3ae × ctx tp4 无 spec × 冻结组合 poetry/lbs64/benchmark_serving)在队——健康则回归区间锁定 b8bc42c3ae→faf2c60935 两构建之间,冻结则配置组合(客户端/streaming/lbs)定罪。 +## 【08-12 17:30】AA slo20 战役(laliao 侧):ragged path 生产形态验证通过;冻结案结案=短 ISL;slo20 表 v1 出炉 + +- **冻结案结案**:D7b(主线运行时×fp8×我 cell 仍冻)+ D9(PR HEAD 新编译 sqsh 仍冻)+ AA(39k ISL 同栈同码健康)→ **触发条件 = 短 ISL(64)× DSpark disagg gen**(浅 KV 角落,与 ctx 侧 SparseAttnTopK 短 ISL 断言同族)。修复挂上游清单;所有真实负载(≥1k ISL)不受影响。 +- **AA slo20 生产形态 ragged 验证 ✅**(aa-master 流水线 fork:slo20 keys × nvfp4 DSpark5 × AA 轨迹客户端 × PR HEAD sqsh):58,144 步 97% ragged、**trim_ratio 0.50**、verify_len 全谱 {1:335k, 2:700k, 3:100k, 4:19k, 5:212k}、accept 2.17、graph 零 miss、154 轨迹完成;客户端 p50 146.8 vs notrim 冒烟 131.2 tok/s(+11.9%,rep=1 方向信号,不作数)。 +- **slo20 SPS 表 v1**(pin{2..5}×conc{32,80},--from-iter-log,28 cells):fixed 15.6ms、α(64/rank)=3.9ms、θ 单调(20→4.0 … 256→26.6 … 384→34.3ms)。几何坑入册:**幂次 bs 网格下 pin 比值 {3,4,5} 互不相连,pin5(比值 6)是连接器**——阶梯没有 pin5 拟合必炸 SweepGeometryError。 +- **reuse 纪律**(AA 负载 reuse=true 是产品行为):终局 A/B 每臂独立冷启动 job + settle 240s + 计量 ≥900s 暖稳态、同 seed 对称升温;SPS 采集豁免(按步形状分箱,缓存温度只改箱填充速度)。 +- 流水线坑五连修入册:helpers 全量拷贝、OUT_DIR 自挂载、ctx 模板 MTP spec 块换 DSpark(KeyError 'weight';TRTLLM moe 装 nvfp4 无罪)、表路径需显式挂载(AA 无整卷 /lustre)、STS 采集配置不得带 enable_ragged_verify(验证器要求 ragged 必须有表)。 +- 在飞:STS v2 采集;下一棒 = 拟温度 → 三臂冷启动矩阵(notrim vs scheduled@新表裸sigmoid vs scheduled@新表+STS)。 + # 分析 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 7c8cea56ebf635e82c3522d9a30ef8871511132f Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Wed, 12 Aug 2026 23:50:10 -0700 Subject: [PATCH 173/177] [None][doc] AA slo20 final verdict: scheduled +6.7%/+7.4% median (n=3 cold-start) on production agentic workload; campaign closed Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 10353ea2bd5c..85b828f02559 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -942,6 +942,21 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - 流水线坑五连修入册:helpers 全量拷贝、OUT_DIR 自挂载、ctx 模板 MTP spec 块换 DSpark(KeyError 'weight';TRTLLM moe 装 nvfp4 无罪)、表路径需显式挂载(AA 无整卷 /lustre)、STS 采集配置不得带 enable_ragged_verify(验证器要求 ragged 必须有表)。 - 在飞:STS v2 采集;下一棒 = 拟温度 → 三臂冷启动矩阵(notrim vs scheduled@新表裸sigmoid vs scheduled@新表+STS)。 +## 【08-12 23:50】AA slo20 终局判决(laliao 侧):scheduled 在生产负载 disagg 上 +6.7%/+7.4%(中位,n=3 冷启动) + +**协议**:slo20 生产形态(2ctx:2gen,maxbs256/mnt1024/131k窗/reuse on/autotuner off)× nvfp4 DSpark5 × AA coding-agent 真实轨迹(32 users,settle+暖稳态计量,seed 42)× PR HEAD sqsh;S 臂 = aa_slo20_table_v1(本形态 pin 阶梯重采)+ tiers[2,3,4,5] + 裸 sigmoid(STS 按用户指令跳过;AA 位置接受率 [0.982→0.892] 极高,裸分数够用);**每臂每 rep 独立冷启动 job**(reuse 纪律:臂间缓存温度对称)。 + +| arm | 生成 p50 (tok/s/user) ×3 | e2e p50 ×3 | +|---|---|---| +| notrim | 130.4 / 130.1 /(r2 NODE_FAIL)| 116.8 / 116.5 | +| sched | **141.7 / 139.0 / 127.2** | **129.8 / 125.3 / 115.5** | + +- **中位 +6.7%(gen)/ +7.4%(e2e)**;trim 0.50、accept 2.21-2.23 三 rep 一致;无腐蚀、轨迹完成流健康。 +- sched r3 与基线打平——rep 间波动 ±5%,n=3 中位口径成立,收窄区间需加 rep;notrim 双 rep 130.4/130.1 极稳。 +- **战役闭环**:ragged path 在 AA 生产形态功能有效(97% ragged 步、trim 50%)且**性能兑现 +7%**——统一律第五格(纯 gen disagg × 高接受产品负载)为正。 +- 过程沉淀(全部入册):AA 流水线十坑(helpers/OUT_DIR 挂载/ctx-MTP-spec/表挂载/STS 五护栏/GEN-only env/…);static×disagg 置信度 NaN(bug 数据点);短 ISL 冻结(已结案);SPS 几何(pin5=连接器);accept 先兆:正式表让 accept 从冒烟 2.17 → 2.22(planner 用真 θ 裁得更聪明)。 +- 余留:sched rep 加密(收窄 ±5%)、c80 档、STS 消融(如果之后想验标定增益)、短 ISL 与 static-NaN 两个上游修复。 + # 分析 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 5c418d4ab61ac35eb4fbd9794198f04f5e3bb5bc Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 16 Aug 2026 01:35:54 -0700 Subject: [PATCH 174/177] [None][doc] slo20-true campaign: v2 dangling-rawref convicted on transfer-admission revert path; golden point slo20-orig x u128; four-arm matrix launched Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 85b828f02559..66eba4947f66 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -957,6 +957,15 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - 过程沉淀(全部入册):AA 流水线十坑(helpers/OUT_DIR 挂载/ctx-MTP-spec/表挂载/STS 五护栏/GEN-only env/…);static×disagg 置信度 NaN(bug 数据点);短 ISL 冻结(已结案);SPS 几何(pin5=连接器);accept 先兆:正式表让 accept 从冒烟 2.17 → 2.22(planner 用真 θ 裁得更聪明)。 - 余留:sched rep 加密(收窄 ±5%)、c80 档、STS 消融(如果之后想验标定增益)、短 ISL 与 static-NaN 两个上游修复。 +## 【08-16 02:00】slo20-true 战役(laliao 侧):v2 dangling-rawref 定罪(压力路径);slo20 原件 ×u128 = 99.9% 黄金点;四臂矩阵开跑 + +**真凶(带完整栈)**:压力上来时 gen 多 rank 同步崩:`Error in event loop: Dereferencing a dangling rawref` @ `_apply_disagg_transfer_admission → _revert_deferred_disagg_gen_init_alloc → _revert_ctx_alloc`(py_executor.py:3793)→ executor 死 → uvicorn 僵尸收单 → 客户端 900s 零 chunk 超时 → 会话冷重启死亡螺旋。**只在传输准入的回滚路径触发**(资源压力/deferred 分配),所以 u128 低压 99.9% 绿、u512/c2600 必崩;与 D3 的 ctx 侧 dangling rawref 同族 = v2 生命周期 bug。主线 c2600 不踩的解释:你们的运行时是 main a8b540912c+wheel(新 v2),我们分支自建 sqsh(旧 v2)——**请确认 main 近期是否修过 rawref 生命周期**;分支自建用户都会踩。 +- 连带修正:此前"短 ISL 冻结"定罪降级为存疑(当时判别在 num_instances 带毒渲染器时代;见下)。 +- **自伤教训入册**:fork slo20 launcher 做 ctx reshape(6×tp4→3×tp8,DSpark 双份装载 OOM 所迫)时漏改 `server_config.num_instances` → router 永等 6 实例 → 一整代判别局(ADP/env/容量/swap 全系列)被毒,全部作废重跑。**渲染器改拓扑必须三处联动:启动命令/urls/num_instances。** +- 搬家坑清单(全修):client 需 .venv 直调(镜像无 uv,--quiet 吞错)、base_url 烤死原节点、wait 1800s、DL 臂需显式 block_size=DL。 +- **黄金点:slo20 全套原件(ADP+DEEPGEMM+maxbs256+kv0.8)× u128 = 28,784 成功/27 失败**。与原始 slo20 唯二偏差:ctx 3×tp8、并发 128(均有硬归因)。 +- **四臂矩阵已在该点开跑**:notrim(r1=28.8k ✓)/DL4/DL3/sched(nvfp4 DSpark,sched 挂 aa_slo20_table_v1 + tiers2345 裸 sigmoid)——回答"收益属于置信度调度还是静态砍 draft len"(roofline best-static 之争的实测版)。 + # 分析 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 8f394a8fd31613cf14b91318de24a5bb6b43c205 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Sun, 16 Aug 2026 08:32:18 -0700 Subject: [PATCH 175/177] [None][doc] slo20-true four-arm final: static DL cuts hurt SLO (-8~-16% per-req speed); sched lossless-but-gainless on nvfp4 (planner rightly declines to trim) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 66eba4947f66..858cecbe8980 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -966,6 +966,21 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - **黄金点:slo20 全套原件(ADP+DEEPGEMM+maxbs256+kv0.8)× u128 = 28,784 成功/27 失败**。与原始 slo20 唯二偏差:ctx 3×tp8、并发 128(均有硬归因)。 - **四臂矩阵已在该点开跑**:notrim(r1=28.8k ✓)/DL4/DL3/sched(nvfp4 DSpark,sched 挂 aa_slo20_table_v1 + tiers2345 裸 sigmoid)——回答"收益属于置信度调度还是静态砍 draft len"(roofline best-static 之争的实测版)。 +## 【08-16 09:30】slo20-true 四臂矩阵终表(laliao 侧):静态砍 DL 伤 SLO;sched 无损无益(nvfp4 平 θ 下理性拒裁) + +工作点 = slo20 原件键(ADP+DEEPGEMM+maxbs256+kv0.8)× u128 × nvfp4 DSpark,每臂独立冷启动,AA client 真轨迹: + +| 臂 | 完成数 | SSE p50 tok/s | +|---|---|---| +| notrim(heur)×2 | 28,784 / 29,698 | 92.1 / 91.6 | +| notrim-noheur(公平对照) | **30,199** | 91.7 | +| **sched**(表+tiers2345 裸sigmoid) | 29,179 | **92.5** | +| DL4 | 30,090 | **84.1(−8%)** | +| DL3 | 29,227 | **77.3(−16%)** | + +- **静态砍 DL 在 SLO 语义下净损**:吞吐 +1~4%(rep 方差 ~3% 内)换每请求速度 −8~−16%——roofline 的 best-static 之争在 SLO 判据下出局。 +- **sched 无损无益**:vs 公平对照完成数 −3.4%(方差内)、速度 +0.9%;trim 仅 0.0023、accept 3.55——planner 按表判定 nvfp4 的 θ 不值得裁并拒裁,行为正确。与量化经济学自洽(收益的肉在 fp8,agg 侧 +11.8% 已证)。sched×heuristic_topk 被验证器禁(ragged 批 next_n 重建不成立)——sched 部署需关 heuristic,公平对照已扣除该差异。 +- 下一格候选:fp8 原版 × 同矩阵(唯一可能让 sched 翻正的工作点);u512+ 需先修 v2 rawref(回滚路径)。 # 分析 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From 05916a8444da3e97ec80958bbfe14e69042b19a1 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 18 Aug 2026 09:53:59 -0700 Subject: [PATCH 176/177] [None][doc] slo20-shape table re-collected: 78% of step is untrimmable; cost model predicts pin-ladder SLO loss within 1-3pp, planner decline vindicated Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 858cecbe8980..93b57015f6d5 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -981,6 +981,31 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - **静态砍 DL 在 SLO 语义下净损**:吞吐 +1~4%(rep 方差 ~3% 内)换每请求速度 −8~−16%——roofline 的 best-static 之争在 SLO 判据下出局。 - **sched 无损无益**:vs 公平对照完成数 −3.4%(方差内)、速度 +0.9%;trim 仅 0.0023、accept 3.55——planner 按表判定 nvfp4 的 θ 不值得裁并拒裁,行为正确。与量化经济学自洽(收益的肉在 fp8,agg 侧 +11.8% 已证)。sched×heuristic_topk 被验证器禁(ragged 批 next_n 重建不成立)——sched 部署需关 heuristic,公平对照已扣除该差异。 - 下一格候选:fp8 原版 × 同矩阵(唯一可能让 sched 翻正的工作点);u512+ 需先修 v2 rawref(回滚路径)。 +## 【08-18 10:30】slo20 形态成本表重采(laliao 侧):**78% 步时不可裁** —— 拒裁的定量铁证,模型预测与 pin 阶梯实测误差 1-3pp + +用户质疑"trim 率为何这么低",直指表失配风险(旧表采自 sweep 形态:TRTLLM moe/ADP off/maxbs64)。**在 slo20 形态原地重采 pin{2,3,4,5} 阶梯**(u128,各 3.5-5.3 万步),得 `slo20_shape_table_v1.json`: + +| | sweep 形态旧表 | **slo20 形态新表** | +|---|---|---| +| fixed | 15.6ms | **35.3ms** | +| α(16/rank) | 1.8ms | 1.4ms | +| θ(96 token) | ~14.5ms(插值) | **10.5ms** | +| bs16 满窗步时构成 | — | 35.3+1.4+10.5=47.2ms,**不可裁部分 78%** | + +**pin 阶梯实测(SLO 口径,同点四格)**: + +| pin | SSE p50 tok/s | accept | steps | 模型预测 | 实测 | +|---|---|---|---|---|---| +| 5(满窗) | 91.6 | 3.61 | 34.7k | — | — | +| 4 | 88.6 | 3.41 | 35.9k | −4.3% | **−3.3%** | +| 3 | 83.5 | 3.06 | 38.6k | −11.3% | **−8.8%** | +| 2 | 63.6 | 2.25 | 53.4k | −32.8% | **−30.6%** | + +- **成本模型在 30% 动态范围内预测实测误差 1-3pp**——θ/accept 双侧定价都对。裁剪在此形态每一档都是净损,**planner 拒裁(trim 0.23%)是唯一正确答案**;新表比旧表更不看好裁剪(fixed 翻倍),即拒裁结论被强化而非推翻。 +- **根因对照**:poetry/agg 形态 54-56% 步时可裁 → 裁剪赚;slo20 形态**只有 22% 可裁**(ADP all-to-all + MEGAMOE_DEEPGEMM + 131k 窗稀疏注意力全是 per-step 固定成本)+ AA 高接受率 [0.98→0.89] → 裁剪两头亏。**"收益 = 可裁比例 × 接受曲线陡度"在两个形态上完成定量闭环。** +- 方法论再证:**表必须在目标形态原地重采**(同模型同硬件,仅 moe 后端/ADP/maxbs 不同,fixed 就差 2.3×)。 +- 在跑:sched × 新表(预测:trim 仍 ≈0,与 notrim 打平)。 + # 分析 ## Step 3 复现性判决(08-07 12:40,校准 + v2b 表条件下) From b853378b557b5cc047426220701ffbee363ce1d9 Mon Sep 17 00:00:00 2001 From: Lance Liao <108499334+lancelly@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:43:05 -0700 Subject: [PATCH 177/177] [None][doc] sched x re-collected slo20 table: prediction hit (trim 0.16%, SSE parity with notrim); confidence scheduling is lossless and correctly declines Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> --- docs/dspark_sts_program.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/dspark_sts_program.md b/docs/dspark_sts_program.md index 93b57015f6d5..d7d0e739c1f8 100644 --- a/docs/dspark_sts_program.md +++ b/docs/dspark_sts_program.md @@ -1004,7 +1004,7 @@ cap-accept 臂(2634604,poetry+arena × bs512/1024 × 3 reps,~1.01M 请求步)完 - **成本模型在 30% 动态范围内预测实测误差 1-3pp**——θ/accept 双侧定价都对。裁剪在此形态每一档都是净损,**planner 拒裁(trim 0.23%)是唯一正确答案**;新表比旧表更不看好裁剪(fixed 翻倍),即拒裁结论被强化而非推翻。 - **根因对照**:poetry/agg 形态 54-56% 步时可裁 → 裁剪赚;slo20 形态**只有 22% 可裁**(ADP all-to-all + MEGAMOE_DEEPGEMM + 131k 窗稀疏注意力全是 per-step 固定成本)+ AA 高接受率 [0.98→0.89] → 裁剪两头亏。**"收益 = 可裁比例 × 接受曲线陡度"在两个形态上完成定量闭环。** - 方法论再证:**表必须在目标形态原地重采**(同模型同硬件,仅 moe 后端/ADP/maxbs 不同,fixed 就差 2.3×)。 -- 在跑:sched × 新表(预测:trim 仍 ≈0,与 notrim 打平)。 +- **sched x 新表已验证(预测命中)**:trim **0.0016**(新表比旧表更保守:0.0023->0.0016,仅 40/29k 请求被裁)、accept 3.59、SSE p50 **91.5** vs notrim 91.6/91.7 —— 逐位打平,零回归。置信度调度在本形态的正确行为就是「看表拒裁」,特性本身无损。 # 分析