[None][feat] Kimi K3: DSpark speculative decoding under Helix CP - #18166
Draft
lancelly wants to merge 11 commits into
Draft
[None][feat] Kimi K3: DSpark speculative decoding under Helix CP#18166lancelly wants to merge 11 commits into
lancelly wants to merge 11 commits into
Conversation
lancelly
force-pushed
the
user/laliao/kimi-k3-helix-dspark
branch
8 times, most recently
from
August 28, 2026 07:19
390aa69 to
e2c3fd2
Compare
Speculative verify groups (1 target + k draft tokens) under helix MLA-DCP + KDA-TP, overlap-scheduler-safe by construction: - Per-token primitive: global positions -> owner rank, rank-local KV write slot (-1 = not owned) and per-token attention bound local_len(pos+1). A group may straddle a ledger-page boundary onto two CP ranks; every consumer below is per-token, so no group-affinity placement is needed and the stateless round-robin ledger stays intact. - model_engine: extend-request packing emits provisional helix values (stale base / full-acceptance convention, mirroring the non-helix spec path); _preprocess_inputs applies the overlap accepted-count correction to helix_position_offsets and re-derives slots/bounds/ rank-local kv_lens on device (recompute_helix_spec_buffers) -- CUDA graph safe, exact under overlap. - trtllm attention metadata: helix_local_slots / helix_kv_bounds buffers, per-seq owned-new-token counts for the kv_lens math, vectorized helix_local_len_vec. - mla_rope_generation (cpp): optional per-token slot table supersedes the per-sequence inactive-rank gate for KV appends (third helix_tensor_params entry; kernels index slots per token). - CuTe DSL MLA decode (fp16/bf16): optional per-token kv_bounds replaces the implicit causal bound under helix; masked-phase span widened by one; stats epilogue emits the (-inf, 0) identity per token; gate now admits seq_len_q > 1 with helix when the spec buffers are live (fp8 KV stays rejected). - Guards: helix spec allowlist = standalone DSpark linear chains only, loud rejection otherwise; the drafter's paged KV manager is built on the repurposed CP-free mapping (helix ledger applies only to the target KV). - DFlash disagg gen-worker slot bootstrap (standalone twin of the embedded-DSpark NVIDIA#16767 fix): transferred requests get isolated context slots instead of aliasing the shared scratch row; acceptance is degraded until a ctx->gen window transfer exists, correctness is carried by verify. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com> (cherry picked from commit fdb9594b21ff84bc45ed65bf94742646867d793d) (cherry picked from commit dbd5dc6)
…ompute, loud fallback rejection Adversarial review of fdb9594b21 confirmed nine findings with four roots: - CRITICAL: both extend packing branches used max_beam_num_tokens-1 as the global position base, but a helix gen worker's token list is the rank-LOCAL round-robin subset (merge_helix_requests). Reconstruct the global base as total_input_len_cp + (max_beam_num_tokens - py_prompt_len) - 1 and override the extend tokens' position_ids with the global values (same convention as the non-spec helix loop). - CRITICAL: recompute_helix_spec_buffers only ran inside the overlap-enabled branch of _preprocess_inputs while every consumer arms on enable_spec_decode alone — overlap-off would read uninitialized slot/bound buffers. The helix recompute now runs on every spec step; the stale-base position correction stays overlap-only. - FallbackFmha now rejects helix verify groups (per-token ownership is inexpressible in the fused thop path) so a CuTe DSL rejection surfaces as a loud no-library error instead of silently wrong attention. - Autotuner profiling: re-derive a size-consistent kv_bounds dummy for bucketed batches (input 9 has no dynamic-dim spec). - CuTe DSL fold_sq padding rows: clamp the per-token bounds index to stay in gmem bounds. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com> (cherry picked from commit 37b561565da500c7f38dd9268a3ed3dd0b85fa3b) (cherry picked from commit fb9b82c)
… token External review finding: a CP rank owning only the tail page of a verify group has zero visible KV for the group's leading tokens while its per-sequence kv_len (bound of the LAST token) is nonzero, so the per-sequence _helix_zero_kv_mask never sanitizes those partial_o rows. The decode kernel fills fully-masked rows with a finite sentinel, making them an average over arbitrary pool bytes; the stats identity gives them corr = 0, but 0 * NaN = NaN would poison the token on every rank. Sanitize by the per-token bound (helix_kv_bounds == 0) on the spec path. Reachable only when a rank holds zero prompt KV (global prompt shorter than cp * tokens_per_block), which GSM8K's ~1k prompts never hit. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com> (cherry picked from commit 75565d3416de411310fa21370461f4e207daa11e) (cherry picked from commit aa08062)
… sanitize zero-bound rows per token, drop leaked dead code Three external-review findings on the draft: - The class-level Kimi K3 spec-mode assert predates DSpark on this base and made the helix allowlist unreachable (helix+DSpark passed the allowlist then died on the assert); extend it with is_dspark(). The embedded-flavour rejection keeps a getattr that is a documented no-op until draft_is_embedded_in_target lands with the embedded DSv4 flavour. - Sanitize zero-bound verify-group rows per token (cherry-picked from the integration branch): a rank owning only a group's tail page has nonzero per-seq kv_len but zero visible KV for leading tokens; their finite-sentinel decode rows average arbitrary pool bytes and 0 * NaN in the combine would poison the token. Reachable when the global prompt is shorter than cp * tokens_per_block. - Remove build_ugpu_block_offsets: leaked from the integration branch by a conflict resolution; references attributes that do not exist on this base and has no callers. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com> (cherry picked from commit 8d10f80)
Two capture-path defects in the helix speculative-decoding bookkeeping, both of which only bite once CUDA graphs are enabled together with the overlap scheduler. 1. `_preprocess_inputs()` corrects `helix_position_offsets` in place with `+= previous_pos_id_offsets_cuda`, but `_postprocess_inputs()` had no matching `-=`. Every other in-place correction on that path is symmetric (see `position_ids` a few lines above) precisely because the buffer is not rewritten between graph replays, so the missing reversal lets the offset accumulate step over step and the derived per-token owner/slot/bound values drift. Add the mirror, guarded by the same `_helix_spec_tokens_valid` condition. The recompute's OVERWRITES (slots/bounds/kv_lens) need no reversal: those buffers are rewritten from host state at the next step's prepare. 2. `recompute_helix_spec_buffers()` built the `-1` sentinel with `local_before.new_full((), -1)`, allocating a scalar tensor on every step of a routine that is captured into the graph. Use the scalar overload of `torch.where` instead so the captured region stays allocation-free. Validated on Kimi-K3 disaggregated GSM8K with helix cp8 + DSpark draft_len=7 and both CUDA graphs and the overlap scheduler enabled. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
… tighten metadata guards Two review findings on the helix speculative-decoding guards. 1. The in-target-drafter rejection keyed on `draft_is_embedded_in_target`, an attribute that does not exist anywhere in the tree; the `getattr` default made the check a permanent no-op, and there is no guarantee the attribute would arrive under that name. Key it on `spec_config.speculative_model` instead, a `DecodingBaseConfig` field: a standalone drafter always carries its own checkpoint, so its absence identifies the embedded flavour. The check now actually fires. The rejection matters because the helix draft KV manager is built on the repurposed (CP-as-TP) mapping, where every rank keeps the FULL drafter KV and the round-robin ownership ledger applies to the target KV alone. An in-target drafter instead reads and writes the sharded target KV, which the per-token bookkeeping does not cover, and would run silently wrong. 2. Replace every `getattr(..., '_helix_spec_tokens_valid', False)` with a direct read. Where the object is statically `TrtllmAttentionMetadata` (`self` inside the metadata class, the annotated `metadata` parameter of `mla_rope_generation`, and the three FMHA entry points, all declared `TrtllmAttentionMetadata` in `fmha/interface.py`) the attribute is a dataclass field with a default and always exists. Where the object comes from generic code (`modules/mla.py`, `pyexecutor/model_engine.py`) guard with `isinstance(..., TrtllmAttentionMetadata)`, matching the existing idiom in those same functions, and read the fields directly after. Behaviour is unchanged: the helix fields and the `update_helix_param` override live only on `TrtllmAttentionMetadata` (the base-class hook is a no-op and neither FlashInfer nor Vanilla overrides it), so non-TRTLLM metadata could never have armed the flag. The guards are stricter, though: a `getattr` default silently degrades to False -- disabling the helix speculative path outright -- if a field is renamed, and would accept a same-named field on a backend that has no helix support. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
Trim the comments introduced by this series to the constraints the code cannot state itself: drop restatements of adjacent code, of the error messages right below them and of the field declarations they repeat, drop one pure navigation note, and correct one claim (allocations during CUDA graph capture come from the graph pool; the scalar overload saves an allocation on the eager path, it is not a capture requirement). No functional change. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
The two helix blocks in the extend loop shared six of nine lines, including the global-base reconstruction that is the easiest part to get wrong (max_beam_num_tokens is rank-local on a helix gen worker, so it cannot serve as a global base). Move that shared part into _helix_pack_extend so the reconstruction exists once; each branch keeps only its three real differences (group size, whether the cached length is taken before or after the group, and whether the owned-token count is exact or a placeholder). No functional change. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
Both DSpark flavours carry a speculative_model path (the embedded one resolves it to the target checkpoint and probes the weight index), so a missing path does not identify the embedded flavour. The embedded case is already rejected downstream by draft_is_embedded_in_target. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
lancelly
force-pushed
the
user/laliao/kimi-k3-helix-dspark
branch
from
August 28, 2026 07:50
e2c3fd2 to
52ecd02
Compare
… mapping The KV budget split for a separate DSpark draft cache computed both cost terms with the raw helix mapping (tp_size=1): - MambaKVCacheParams.get_states_bytes_per_layer counted the KDA conv/ssm state unsharded, withholding 29.2 GiB/rank from the budget on a helix16 gen worker whose real pool is 1/16-sharded and costs 1.7 GiB. The estimator now shares the allocator's effective-TP rule (mamba_effective_tp_size, moved to config_utils so the two can never diverge again). - The drafter's per-token cost used the unrepurposed mapping; the value (20480 B/token) was right only because this drafter's kv-head count happens to equal cp_size. It is now computed per GLOBAL token with the repurposed mapping runtime construction uses, then scaled by cp_size into the target's rank-local-token unit. Non-helix paths are byte-identical (scale 1, same mapping). On helix16 mb64 the target KV capacity grows ~28% and the gen-only fill that previously died with 'Insufficient KV cache' has room to complete. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
lancelly
force-pushed
the
user/laliao/kimi-k3-helix-dspark
branch
from
August 28, 2026 07:51
52ecd02 to
c26623d
Compare
Every concrete decoding config carries decoding_type as a required Literal field, so the defensive getattr was dead weight; match the file's existing spec_dec_mode.is_dspark() idiom instead. Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
lancelly
force-pushed
the
user/laliao/kimi-k3-helix-dspark
branch
from
August 28, 2026 07:58
972db06 to
9e1d61f
Compare
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
1 participant
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.
Summary
Enables DSpark speculative decoding under Helix context parallelism for Kimi K3 (MLA-DCP + KDA-TP, disaggregated generation servers), including overlap-scheduler support.
The core mechanism is a per-token primitive derived from global token positions: for each verify-group token (1 target + k draft), we derive on device its owner CP rank, its rank-local KV write slot (
-1when another rank owns the position), and its per-token attention boundlocal_len(pos+1). A verify group may straddle a ledger-page boundary onto two CP ranks; every consumer is per-token, so the V2 superblock ledger's stateless round-robin ownership stays intact (no group-affinity placement, no reserve/rewind FIFO bookkeeping).Changes
total_input_len_cp + generated - 1; the request token list is rank-local under helix)._preprocess_inputsapplies the overlap accepted-count correction tohelix_position_offsetsand re-derives slots/bounds/rank-local kv_lens on device (recompute_helix_spec_buffers) — exact under overlap, CUDA-graph-safe, and identical for overlap-off.helix_local_slots/helix_kv_boundsper-token buffers, per-seq owned-new-token counts for the kv_lens math, vectorizedhelix_local_len_vec.mla_rope_generation(C++): an optional per-token slot table (thirdhelix_tensor_paramsentry, op schema unchanged) supersedes the per-sequence inactive-rank gate for KV appends.kv_boundsreplaces the implicit causal bound under helix (masked-phase span widened by one; padding-row index clamped); the stats epilogue emits the(-inf, 0)identity per token; the wrapper gate admitsseq_len_q > 1with helix when the spec buffers are live (fp8 KV stays rejected).FallbackFmharejects helix verify groups (per-token ownership is inexpressible in the fused thop path) so misconfiguration surfaces as a hard error instead of silently-wrong attention.DFlashSpecMetadata.preparebootstraps per-request context slots on disaggregated generation workers (the standalone twin of the embedded-DSpark [Bug] DSpark speculative decoding: accept length collapses to ~1 at generation batch size > 1 in disaggregated serving #16767 fix).Stacked on
Rebased onto main past the two merged dependencies; the diff is now #18131 plus this PR's own four commits (implementation, two review-fix rounds, external-review fixes: DSpark enabled in the K3 spec-mode assert, per-token zero-bound row sanitization, dead-code removal).
Validation
Validated end-to-end on the Rubin integration branch (
user/lizhiz/rubin-advancebase; same code modulo integration deltas):tests/unittest/_torch/modules/test_mla_helix.py: 24/24 (non-spec helix numeric regression gate).Known gaps (why this is a draft)
CuteDslMlaFmha.is_availabledefer gate: on main, an explicitly-setattn.flashinfer_mla_backendmakes the standalone CuTe DSL backend defer (from [TRTLLM-15033][feat] Upstream Kimi K3 MLA decode backend selection to main #17800). K3 always sets it, so on pure main semantics helix verify groups would hit the loud no-library rejection instead of the CuTe DSL kernel. The Rubin branch has no such gate. Needs a decision: exempt helix from the defer gate, or route K3 helix selection differently.🤖 Generated with Claude Code