Repository navigation
Conversation
…ject#48414 Copy the DCP decode-runtime hunks of the Kimi-K3.1 Day-0 branch onto main as they are: index-K sharded by token, the local top-k plus global merge in the indexer, the DSA backend's DCP prefill gather and decode partials, and the Kimi-K3.1 DCP override. Follow-up commits adapt it to main's APIs and narrow its scope. Co-authored-by: Khoa Pham <khoa.pham@radixark.ai> Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> Co-authored-by: Po-Han Huang <53919306+nvpohanh@users.noreply.github.com> Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>
- dsa_dcp_sharded: build the merge with DcpTopKMerge.create (sgl-project#48588), once per process even when it declines; call the candidate, paged and packed top-k with dcp_size / dcp_rank / deterministic and the default cluster choice (sgl-project#48501). - deepseek_v4_backend: pass the FlashMLA schedule as schedule= (sgl-project#48501). - dsa_backend: run FlashMLA through flash_mla_interface and the backend's _call_* hooks instead of sgl_kernel.flash_mla. - Drop the native-head FlashMLA branches; their kernel is not on main.
- The override rejects speculative decoding, PD disaggregation, HiCache, HiSparse and prefill CP under DCP; their DCP paths stay on sgl-project#48414. - The DSA backend requires flashmla_kv for prefill and decode under DCP; trtllm and Prims-TS DCP are not ported. - The candidate all-gather uses the multimem all-gather or the group all-gather; the prefix-only gather kernel is not on main.
…oject#48414 Copy the speculative hunks of the Kimi-K3.1 Day-0 branch's DCP runtime onto the DCP decode port as they are: DSA target verify returns per-rank partials, the verify rows' local lengths (DG-native on SM100), the DSPARK and DFLASH branches of the Kimi-K3.1 DCP override, and DFLASH MLA drafts on the MLA-only draft backends under DCP. Follow-up commits adapt it to main.
…sl_mla - The draft-backend check loads the draft config with get_config and reads parse_dflash_draft_config(...).is_mla, as main's trtllm_mha check does, instead of the source's is_dflash_mla_draft_config. - Only cutedsl_mla is an MLA-only DFLASH draft backend here; trtllm_mla needs the sliding-window draft table work that is not ported. - The DFlashMLAAttention comment points at mla_attention_dcp_size, which keeps a DCP MLA draft on per-rank owner rows on main.
5 of 21 tasks
kpham-sgl
marked this pull request as ready for review
November 5, 2026 22:43
kpham-sgl
requested review from
1am9trash,
Fridge003,
Qiaolin-Yu,
Ying1123,
alphabetc1,
ch-wan,
fzyzcjy,
hanming-lu,
hnyls2002,
huangtingwei9988,
hubertlu-tw,
hzh0425,
ispobock,
kkHuang-amd,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
November 5, 2026 22:43
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
kpham-sgl
requested review from
BBuf,
Edwardf0t1,
hebiao064 and
rainj-me
as code owners
November 5, 2026 22:43
The target's pool budget charged the DFLASH draft KV at attn_dcp_size times its cell size, as if the draft pool always spanned the widened locs. On main an MLA draft that executes DCP (mla_attention_dcp_size) keeps per-rank rows, so its loc_space_scale is 1 and the charge was dcp_size times too high. - draft_loc_space_scale(hf_config=, use_mla_backend=, dcp_size=) holds KVCacheConfigurator.loc_space_scale's draft rule; the draft property calls it, and the target-side configurator resolves dflash_draft_loc_space_scale from the draft model config with it. - _dflash_draft_cell_size and the layer-count fallback scale the draft by dflash_draft_loc_space_scale instead of attn_dcp_size. - Drop the source's mla_draft_rows_are_dcp_collapsed predicate and its resolver; main's rule replaces them. - The unit test covers both scales and draft_loc_space_scale itself.
Collaborator
Author
|
/rerun-test test/registered/unit/model_executor/test_pool_configurator.py |
Contributor
|
- Target verify is rejected under DCP for now, so the indexer and the backend take only decode rows: remove the verify branches, the verify wording, and the draft-token factor in the merge row count. - Build the candidate all-gather only when the fused merge declines; init_dcp_topk_merge returns whether it is in use. Drop the unused all_gather parameter of dcp_global_topk. - Assert that the FlashMLA head count matches the one its schedule metadata was built for. - Condense the multi-line comments the port added, and note that the FlashMLA schedule constants also cover the head64x2 (h128/d576) path.
Comment and docstring changes only (AST-identical apart from docstrings). Adds a TODO to budget the dense merge scratch, and restores main's index-K sizing comment with its wording updated for K3.1.
Handle --dcp-size > 1 inside _kimi_k31_overrides, like Kimi-K3's DCP branch: reject speculative decoding, PD and HiCache, and default the replicated q projection. Drop the power-of-two rule and the checks the shared hooks already make (DCP divides TP, no prefill CP, no HiSparse for K3.1), and restore Kimi-K3's own block to main.
Use the same predicate the DSA backend checks instead of getattr on index_topk.
Drop the scattered K3.1 one-line comments, restore main's wording on the comments the port had edited, and keep one NOTE on the flag in layers/dcp/layout.py.
save_kv_cache is always True here, the RadixAttention default; the source needs it only for its side-stream KV write, which is not ported.
…ackend Follow the MLA DCP split of ownership with DSA-specific entry points: - prepare_decode_context_parallel_metadata_dsa (planner) builds this rank's prefix rows and a buffer laid out request after request, the order the ragged top-k indices address; - all_gather_kv_cache_for_extend_dsa (comm) all-gathers the cached prefix rows and quantizes the extend tokens as the pool writes them; - bcg_dcp_extend_kv_gather_dsa (forward_mla) and the MHA one-shot path call it; the DSA backend only reads attn_dcp_metadata.dcp_kv_buffer. Kimi-K3 picks the DSA planner when index-K is sharded, so K3.1 no longer allocates the unused MLA gather buffer.
…he DSA backend" This reverts commit 2d3e598; the change goes to the source branch first so the port stays a port.
Pass an empty FlashMLASchedMeta so FlashMLA builds the split-KV schedule from topk_length, instead of precomputing a windowed one through the DeepSeek-V4 backend's private helper. deepseek_v4_backend.py is back to main; the precomputed schedule is a decode speedup for a later port.
Restore main's padding ladder for the decode head count, add the same for the DCP prefill (local) head count, and pick between them in _forward_flashmla_kv instead of recomputing from the tensor and asserting.
Drop _validate_dcp_backend from the DSA backend. Its backend check is already enforced by the Kimi-K3.1 DSA backend pin and its model check cannot fail; the K3.1 DCP override now rejects a non-fused or non-v2 sgl-kernel top-k before the model loads.
Replace the backend's dcp_size/dcp_rank (replicated index-K DCP) and the port's attn_dcp_size/attn_dcp_rank (sharded index-K DCP) with one attn_dcp_size/attn_dcp_rank, the DCP width this backend runs at, and a dsa_shards_index_k flag for the mode. Sites of one mode test the flag; non-K3.1 models evaluate every site as before.
…o kpham/k31-dcp-spec-verify # Conflicts: # python/sglang/srt/arg_groups/model_overrides/kimi_k3.py
sgl-project#48721 dropped its target-verify branches and moved the sharded-mode checks to dsa_shards_index_k. Restore verify on top of that shape: - the paged top-k and the k-only path take target-verify rows; - prefill gathers KV only outside target verify; - the merge sizes rows by draft tokens; - the verify-only sites test dsa_shards_index_k instead of the DCP width. The Kimi-K3.1 DCP override admits DSPARK (static ragged verify) and DFLASH (cutedsl_mla draft) inline, with decode attention mode for verify.
Three tests build DeepseekSparseAttnBackend with __new__; set the attn_dcp_size/attn_dcp_rank, dsa_shards_index_k, real_page_span and prefill head fields that __init__ now sets for a non-DCP backend.
…o kpham/k31-dcp-spec-verify
The remaining tests and the dsa_metadata_kit helper that build DeepseekSparseAttnBackend without __init__ now set dsa_shards_index_k, attn_dcp_size/attn_dcp_rank and, where they set a page size, real_page_span, as __init__ does for a non-DCP backend.
…am-sgl/sglang into kpham/k31-dcp-spec-verify
…am-sgl/sglang into kpham/k31-dcp-spec-verify
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.
Original authors: @yhyang201
Stacks on #48721. The diff includes #48721 until it merges; the commits of this PR start at "Raw port of DSPARK and DFLASH under DCP".
Port of #48414
On top of the Kimi-K3.1 DCP decode runtime (#48721), this ports the branch's speculative decoding under DCP. DSA target verify attends over each rank's shard and returns
(o, lse)partials for the existing DCP merge, as decode does. The indexer scores each verify row's local share of its causal length; on SM100, DG-native paged MQA logits read a request's index-K once for all of its draft rows. The Kimi-K3.1 DCP override admits DSPARK (static ragged verify, decode attention mode) and DFLASH (cutedsl_mla draft backend, decode attention mode). DFLASH MLA drafts are no longer rejected under DCP: they run the target's DCP decode path. The target's pool budget charges the DFLASH draft KV once per draft pool row per target row, not per widened location (#64 on the branch).Against the raw port (first commit), the draft-backend check reads the draft config through main's
get_config/parse_dflash_draft_config(...).is_mlainstead of the branch'sis_dflash_mla_draft_config. Onlycutedsl_mlais admitted as an MLA-only DFLASH draft backend; the branch also admitstrtllm_mla, which needs the sliding-window draft table work. For the draft KV budget, the branch'smla_draft_rows_are_dcp_collapsedpredicate and its resolver are replaced by main's rule.draft_loc_space_scale(hf_config=, use_mla_backend=, dcp_size=)now holdsKVCacheConfigurator.loc_space_scale's draft rule; the target-side configurator resolvesdflash_draft_loc_space_scalefrom the draft model config with it._dflash_draft_cell_sizeand the layer-count fallback scale the draft by that instead ofattn_dcp_size. API changes for the branch's callers:_dflash_draft_loc_scale/dflash_draft_rows_dcp_collapsedbecomeKVCacheConfigurator.dflash_draft_loc_space_scale.Not taken, because main's design covers it: the branch's DCP-collapsed draft row layout (
draft_mla_rows_are_physical,mla_draft_rows_are_dcp_collapsed). On main,mla_attention_dcp_sizealready keeps a DCP-executing MLA draft on per-rank rows (loc_space_scale == 1).Left for later ports:
splitmode; their only caller is the megakernel verify path.Modifications
dsa_backend.py: target verify under DCP skips the prefill KV gather and asks_forward_flashmla_kvfor a partial. Verify rows get this rank's local lengths, laid out[B, next_n]for DG-native. The captured-verify DG-native refresh stays off under DCP.dsa_indexer.py: DG-native paged MQA logits for DCP verify on SM100.kimi_k3.pyoverride: DSPARK and DFLASH rules in the DCP branch of_kimi_k31_overrides. PD, HiCache, HiSparse and prefill CP stay rejected.speculative_hook.py,draft_worker_common.py:cutedsl_mlais a DFLASH draft backend for MLA drafts.models/dflash.py: drop the--dcp-size > 1rejection of MLA drafts.kv_cache_configurator.py,pool_configurator.py: the DFLASH draft KV is charged atcell_size * dflash_draft_loc_space_scale: 1 for an MLA draft that executes DCP,attn_dcp_sizefor a draft that spans the widened locs. Before, it was alwaysattn_dcp_size.Tests
Static checks:
pre-commitandruff --select Fon the touched files.test_pool_configurator.pycovers the draft charge at loc_space_scale 1 and 8 anddraft_loc_space_scaleitself; it has not been run yet (pending a/rerun-test). GPU validation owed: Kimi-K3.1 TP8 / DCP8 with DSPARK (SGLANG_RAGGED_VERIFY_MODE=static) and with an MLA DFLASH draft on B300, checking accept length and GSM8K (sgl-eval) against DCP1.CI States
Latest PR Test (Base): ❌ Run #42156554976
Latest PR Test (Extra): ❌ Run #42156554981
Latest PR Test (AMD ROCm 10): ❌ Run #42156555020
Tip
Base CI only covers the core cases, and its workload may not exercise your change. Use
/rerun-testto run the tests your change touches, or copy this prompt for an agent: