Skip to content

[Feat][NVIDIA] DeepSeek-V4.1: Prefill support CP with multistream forward prepare (4/N) - #40179

Open
SYChen123 wants to merge 27 commits into
sgl-project:mainfrom
SYChen123:v4.1/cp-multistream
Open

[Feat][NVIDIA] DeepSeek-V4.1: Prefill support CP with multistream forward prepare (4/N)#40179
SYChen123 wants to merge 27 commits into
sgl-project:mainfrom
SYChen123:v4.1/cp-multistream

Conversation

@SYChen123

@SYChen123 SYChen123 commented Sep 18, 2026 •
edited by github-actions Bot
Loading

Copy link
Copy Markdown
Contributor

Motivation

Issue #39700. When using CP on deepseek v4.1, _forward_prepare_low_ratio_multi_stream will be disabled. This PR supports their compatibility.

Modifications

Developed based on #40548 and need #40548 to be merged first. The implementaion is similar with #39588.

  1. Do CP allgather in advance on the default stream. And the multistream forward prepare accepts the pre-gathered kv score as input. No CP allgather will be performed inside the multistream forward prepare.
  2. After cp materialization, store_swa_k, compressor and indexer(operations before mqa_logits computation) are overlapped in multiple streams.
  3. In BCG for DSV4.1, enable the indexer and compressor to be captured into graph.

Accuracy Tests

image

Speed Tests and Profiling

image

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ❌ Run #38354933523
Latest PR Test (Extra): ❌ Run #38354933180
Latest PR Test (AMD ROCm 10): ❌ Run #38354933330

@BJWang-ant BJWang-ant mentioned this pull request Sep 20, 2026
Open
2 tasks done
@BJWang-ant BJWang-ant mentioned this pull request Sep 20, 2026
Merged
@SYChen123 SYChen123 mentioned this pull request Sep 21, 2026
Open
5 tasks
@github-actions github-actions Bot added Multi-modal multi-modal language model blackwell SM100/SM120 labels Sep 21, 2026
BJWang-ant and others added 24 commits October 13, 2026 12:18
@BJWang-ant @SYChen123
fix bug
d0d8ae1
@BJWang-ant @SYChen123
run pass cp+megamoe+bcg
a7d835c
@BJWang-ant @SYChen123
Reuse shared compressed KV dequantization in V4.1 CP prefill
dd3a279 
V4.1 consumer layers currently dequantize the same compressed KV on every attention call. Keep a separate workspace per compression ratio and reuse its compressed prefix within one forward, keyed by source layer and allocation address. Refresh on producer execution or allocation replacement, while updating SWA and top-k for every layer.

Limit reuse to CP sparse BF16 prefill with compression ratios 1 and 2. Per-forward gather recreation invalidates the keys across batches and BCG replays. Add a GPU regression covering source changes, source re-execution, ratio interleaving, new batches, and workspace growth/shrink against forced dequantization.

Validation: B200 TP4/EP4 with CP, BCG, MegaMoE and DSPARK; GPU regression passed; 16/16 long/short arithmetic outputs match baseline. Six-step traces reduce compressed dequant calls from 228 to 24. Three unchanged test.sh trials improve median output throughput by 3.20% at concurrency64 and 4.11% at96. Concurrency32 has overlapping timing ranges and a smaller median change.
@SYChen123
Support V4.1 DP attention in DP-only DSpark PD
a0a0893
@BJWang-ant @SYChen123
Revert "Reuse shared compressed KV dequantization in V4.1 CP prefill"
20a1ffd 
This reverts commit 007f7c8.
@SYChen123
Support V4.1 decode vision MegaMoE
8ad77b3
@BJWang-ant @SYChen123
fix pd-tp-conflict
6864ce3
@SYChen123
CP BCG capture indexer and compressor into graph
696bb19
@SYChen123
Indexer BCG path uses topk_v2
ab18a2b
@SYChen123
Add SGLANG_DSV41_BCG_DENSE_INDEXER to support dense indexer on cp bcg…
830338b 
… path.
@SYChen123
Optimize CP BCG dense indexer path
5c98dcb
@SYChen123
Use sparse candidate consumers in CP BCG
e8e01bb
@SYChen123
Port CP BCG candidate indexing to V4.1 indexer refactor
fa75cbf
@SYChen123
Support multistream forward prepare on CP.
e25f134
@SYChen123
Overlap CP BCG indexer preparation
e5ee13a
@SYChen123
Fix CP multistream indexer import and formatting
b47fc37
@SYChen123
Fix V4.1 multistream indexer dispatch and MegaMoE compatibility
5e5e080
@SYChen123
fix CP flag used in deepseek_v4
565c272
@SYChen123
current stream waits stream_kv after indexer being launched.
79f287d
@SYChen123
perf: overlap DSV4.1 CP Q projection with allgather
b4b972b
@SYChen123
perf: submit DSV4.1 CP allgathers before Q projection
0bf0ffe
@SYChen123
remove stream_indexer and replace it with current stream
15aaa88
@SYChen123
chore(dsv41): lint multistream rebase resolutions
a31abfc
@SYChen123
test(dsv41): align PD fixtures with current cache contracts
e44e29a
@SYChen123
SYChen123 force-pushed the v4.1/cp-multistream branch from f458a0d to e44e29a Compare October 13, 2026 06:06
@SYChen123 SYChen123 changed the title [DSV4.1][PD][4/N] Prefill support CP with multistream forward prepare [Feat][NVIDIA] Prefill support CP with multistream forward prepare (4/N) Oct 20, 2026
@SYChen123 SYChen123 changed the title [Feat][NVIDIA] Prefill support CP with multistream forward prepare (4/N) [Feat][NVIDIA] DeepSeek-V4.1: Prefill support CP with multistream forward prepare (4/N) Oct 20, 2026
@github-actions

This comment has been minimized.

Sign in to view
@github-actions

github-actions Bot commented Oct 20, 2026

Copy link
Copy Markdown
Contributor

Required checks failed on e44e29a; this PR cannot be merged until they pass.

PR title (run)

PR title must use "[Type] scope: title" or "[Type] [Hardware] scope: title". Type must be one of: Feat, Fix, Refactor, Perf, Doc, Test, CI, Chore. Scope is required and may be any module or model name, with unrestricted letter case (letters, numbers, dots, underscores, slashes, or hyphens). Hardware is optional; use a single tag without spaces. Use uppercase hardware tags (AMD, NVIDIA, NPU, MI300X), except for exactly these four spellings: Intel, Apple, Huawei, Ascend. Reverts must use GitHub's default format: Revert "original title". Use one space between fields and after the colon, with a non-empty title and no leading or trailing whitespace. Examples: "[Fix] attention: Fix KV cache indexing", "[Feat] qwen3: Add model support", "[Perf] [AMD] flash-attn: Optimize decode performance".
Edit the title; the check reruns on its own.

@airMeng airMeng mentioned this pull request Nov 4, 2026
Open

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner
@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner
@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 is a code owner
@Fridge003 Fridge003 Awaiting requested review from Fridge003 Fridge003 is a code owner
@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner
@xiezhq-hermann xiezhq-hermann Awaiting requested review from xiezhq-hermann
@hanming-lu hanming-lu Awaiting requested review from hanming-lu
@yizhang2077 yizhang2077 Awaiting requested review from yizhang2077
@hzh0425 hzh0425 Awaiting requested review from hzh0425
@alphabetc1 alphabetc1 Awaiting requested review from alphabetc1
@huangtingwei9988 huangtingwei9988 Awaiting requested review from huangtingwei9988
@Qiaolin-Yu Qiaolin-Yu Awaiting requested review from Qiaolin-Yu Qiaolin-Yu is a code owner
@hebiao064 hebiao064 Awaiting requested review from hebiao064 hebiao064 is a code owner
@HaiShaw HaiShaw Awaiting requested review from HaiShaw HaiShaw is a code owner
@ByronHsu ByronHsu Awaiting requested review from ByronHsu ByronHsu is a code owner
@ShangmingCai ShangmingCai Awaiting requested review from ShangmingCai ShangmingCai is a code owner
@Duyi-Wang Duyi-Wang Awaiting requested review from Duyi-Wang Duyi-Wang is a code owner
@sogalin sogalin Awaiting requested review from sogalin sogalin is a code owner
@ch-wan ch-wan Awaiting requested review from ch-wan ch-wan is a code owner
@BBuf BBuf Awaiting requested review from BBuf BBuf is a code owner
@Edwardf0t1 Edwardf0t1 Awaiting requested review from Edwardf0t1 Edwardf0t1 is a code owner
@DarkSharpness DarkSharpness Awaiting requested review from DarkSharpness DarkSharpness is a code owner
@celve celve Awaiting requested review from celve celve is a code owner
@HydraQYH HydraQYH Awaiting requested review from HydraQYH HydraQYH is a code owner
@yuan-luo yuan-luo Awaiting requested review from yuan-luo yuan-luo is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

blackwell SM100/SM120 deepseek jit-kernel Multi-modal multi-modal language model

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@SYChen123 @BJWang-ant