Repository navigation
Conversation
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>
…nfidence 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>
…path PR NVIDIA#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 <lancelly@users.noreply.github.com>
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 <lancelly@users.noreply.github.com>
…e 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 <laliao@login-bia02.bia.clusters.nvidia.com>
…k path 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 <laliao@login-bia02.bia.clusters.nvidia.com>
… 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 <laliao@login-bia02.bia.clusters.nvidia.com>
…DP shape gate 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 <laliao@login-bia02.bia.clusters.nvidia.com>
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 <laliao@login-bia02.bia.clusters.nvidia.com>
…ations 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 <laliao@login-bia02.bia.clusters.nvidia.com>
…of silent 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 <laliao@login-bia02.bia.clusters.nvidia.com>
…pressor differential 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 <laliao@login-bia02.bia.clusters.nvidia.com>
…ion 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 <laliao@login-bia02.bia.clusters.nvidia.com>
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 <laliao@login-bia02.bia.clusters.nvidia.com>
…ckends 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 <laliao@login-bia02.bia.clusters.nvidia.com>
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 <laliao@login-bia02.bia.clusters.nvidia.com>
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 <laliao@login-bia02.bia.clusters.nvidia.com>
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 e1c28ee 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 <laliao@login-bia02.bia.clusters.nvidia.com>
…lent through capture Signed-off-by: Lance Liao <laliao@login-bia02.bia.clusters.nvidia.com>
Signed-off-by: Lance Liao <laliao@login-bia02.bia.clusters.nvidia.com>
…sses
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 <laliao@login-bia02.bia.clusters.nvidia.com>
Signed-off-by: Lance Liao <laliao@login-bia02.bia.clusters.nvidia.com>
…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 <laliao@login-bia02.bia.clusters.nvidia.com>
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 <laliao@login-bia02.bia.clusters.nvidia.com>
`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 <laliao@login-bia02.bia.clusters.nvidia.com>
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=<n>` 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 <laliao@login-bia02.bia.clusters.nvidia.com>
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 <laliao@login-bia02.bia.clusters.nvidia.com>
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 <laliao@login-bia02.bia.clusters.nvidia.com>
…im-regret metric Signed-off-by: Lance Liao <laliao@login-bia02.bia.clusters.nvidia.com>
…im (-1~-6%), pin3 -15%; workload-dependence established Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…sfer, DEFAULT-vs-v2); round 4 = NIXL+PYTHON with UCX env Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…it max_seq_len) and reproducible benchmark configs for both workloads + disagg Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…unverified maxbs/mnt prescriptions Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…) with measured gen-bs monitoring Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…try, arena-hard, throughput_1k) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…oute, three-workload matrix, per-workload SPS/STS; request verbatim poetry prompt Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…endent disagg repro; correct poetry output-length description (768-capped, temp 0.7) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…k extract asset (gated dataset, 401 on our side) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…ED_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>
…ry); 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>
…8%; unified law = pure-gen density x batch depth Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…across the board on Pro@DEP8; disagg ctx OOM root cause (tp8 fix) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
….2% to +0.4~+1.7% (same-node paired, both orders) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…r clean); STS-violation found in first sched replication; own postfix table cross-validates Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…y-by-key protocol alignment (collaborator agg replication v2) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…ied law verified 3 ways; disagg gen no-commit regression suspect (D1 bisect running) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…gg gen commit-path provenance (branch HEAD proven working via osl1kb) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…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>
…/B); freeze-bug bisect: overlay and ctx-spec exonerated, D4 running Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…all-measured table to 1536 tokens; monotone depth trend; router backlog trap Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…ape (trim 0.50); freeze convicted as short-ISL; slo20 SPS table v1 Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
… cold-start) on production agentic workload; campaign closed Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…sfer-admission revert path; golden point slo20-orig x u128; four-arm matrix launched Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…6% 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>
…e; cost model predicts pin-ladder SLO loss within 1-3pp, planner decline vindicated Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
…16%, SSE parity with notrim); confidence scheduling is lossless and correctly declines Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com>
This was referenced Oct 28, 2026
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
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: its only consumer
computed a prefix length with the wrong algorithm, synced to the host, and
assumed batch size 1.
This PR 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
the acceptance rule — and therefore the output distribution — is untouched.
What's here
written in place; rebinding it would be invisible to an already-captured graph.
probabilities; the batch budget is allocated by global survival rank with a
value-independent tie-break, so the path is sync-free and CUDA-graph safe.
τ/costover a profiledstep-cost table. It is a global argmax rather than a greedy first descent,
because the cost curve is a staircase and Θ is not unimodal on it. Without a
profiled table the planner refuses to trim rather than optimize against a model
in which every extra token looks free.
enable_ragged_verify: eachrequest gets its own window, packed onto a flat token axis with a
qo_indptrand padded up to a captured token bucket.
DSparkDecodingConfig; golden manifest regenerated.Three pre-existing bugs fixed along the way
dspark.pybaked in a stride ofmax_draft_len + 1when reading capturedtarget hidden states. Under any dynamic draft length that reads the wrong
rows. This one is independent of this feature and would be reasonable to
take on its own.
kv_len_offsetssubtracted abatch-wide token count from a per-request quantity, and the overlap branch of
_prepare_tp_inputsappended a fixed number ofposition_idsper request.Both are silent wrong-answer paths once request lengths differ.
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: the derived totals are algebraically identical to the previous
multiplicative formulas, and
py_verify_lenis unset on every non-DSpark path.Test Coverage
New:
tests/unittest/_torch/speculative/hw_agnostic/test_dspark_confidence_schedule.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_ragged.pytests/unittest/_torch/speculative/hw_agnostic/test_dspark_overlap_layout.pytest_gsm8k_dep8_megamoe_deepgemm_confidence_schedulingintests/integration/defs/accuracy/test_llm_api_pytorch.py, asserting the samereference accuracy as the baseline — the scheduler changes only how many
drafted tokens are verified, never the acceptance rule, so any delta is a
broken invariant rather than a throughput tradeoff.
Extended:
test_dspark_heads.py,test_dspark_draft.py,test_dspark_worker.py,test_dspark_cuda_graph.py,tests/unittest/llmapi/test_llm_args.py.208 unit tests pass locally, plus the existing non-DSpark speculative tests as a
regression check.
Accuracy, CUDA-graph and attention-DP acceptance all still need a checkpoint
with trained confidence-head weights, which is not available yet (the model
directory is still an empty placeholder). The GSM8K test above is the intended
gate. Two specific gaps a reviewer should know about:
len(tiers)× more graphs than theuniform ladder, and captured-graph memory comes out of KV cache. 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 inside the padding machinery — not something to do
without being able to run it. Ragged is off by default (
enable_ragged_verify).all-greedy, so a
temperature=0accuracy run cannot cover it.Also worth flagging: the feature's headroom is bounded by the fact that the draft
pass is not trimmable. Only the target's verify-token term is recoverable, so
gains are a high-concurrency effect and are break-even or negative at small batch
sizes. This matches the cost model SGLang published for the same feature.
PR Checklist