[CP] Migrate MLA prefill CP (DeepSeek V3) to CP-v2 zigzag strategy by kpham-sgl · Pull Request #31619 · sgl-project/sglang

kpham-sgl · 2026-07-17T22:45:05Z

Motivation

Part of the Prefill CP refactor roadmap (#27252) and the Context Parallelism roadmap (#21788). #28421 landed the CP-v2 ZigzagCPStrategy for GQA/MHA (Qwen3) and the model-agnostic eager-runner boundary. This PR ports MLA prefill CP for DeepSeek V3/R1 (originally #23292) onto that abstraction and makes the strategy own the MLA latent all-gather.

DeepseekV3ForCausalLM is added to CP_V2_DEFAULT_MODEL_CLASSES, so CP-v2 is its default (like Qwen3MoeForCausalLM). The legacy cp_utils v1 path stays intact and reachable via SGLANG_ENABLE_CP_V2=0.

Modifications

Boundary migration (v1↔v2 coexist):

layers/cp/utils.py — add DeepseekV3ForCausalLM to CP_V2_DEFAULT_MODEL_CLASSES.
models/deepseek_v2.py / models/deepseek_nextn.py — gate the in-model CP entry split + exit gather to v1-only (use_cp_v1); the eager runner owns them under v2.
layers/attention/flashattention_backend.py — MLA CP attention dispatches through cp_strategy.run_attention under v2, else the v1 cp_attn_forward_extend.
model_executor/runner/eager_runner.py — CP-v2 extend is a single self-contained _execute_extend_cp_v2 (shard at the boundary, run the body on the rank-local slice, gather before logits).

Decision A — strategy owns the MLA latent all-gather (v2):

models/.../forward_mla.py — skip the upstream rebuild_cp_kv_cache under v2 (k_nope/k_pe stay rank-local).
flashattention_backend.py — the MLA KV write dispatches to cp_strategy.materialize_full_mla_kv under v2 (gather rank-local latent → full, then set_mla_kv_buffer); v1 keeps the upstream-rebuilt full write.
layers/cp/zigzag.py — add materialize_full_mla_kv (fuse compressed-KV + rope, one all-gather via gather_kv_cache, split back for the paged write).

Already handled (no change needed): the seq_lens_k padding flagged on #28421 — zigzag build_metadata bakes pad_len (#28421) and the FA backend widens page_table by pad_delta (#23292). Decision A's gather preserves the same padded length.

Accuracy Tests

Existing test/registered/cp/test_deepseek_v3_cp_single_node.py (tp=8, dp=2, attn-cp=4, DeepSeek-V3-0324, GSM8k ≥ 0.935) now runs CP-v2 via the default-class enablement. Verified on 8×H200:

Config	GSM8k
v1 (`SGLANG_ENABLE_CP_V2=0`)	0.966
v2 + Decision A (strategy latent gather)	0.96 / 0.962
v2, final (eager-runner extraction)	0.964

Non-regression: v1 MLA CP unchanged (0.966); NSA/DSA and MHA CP paths are no-ops under default SGLANG_ENABLE_CP_V2=0.

Out of scope (follow-ups)

Kimi K2.5 — CP-v2 for the multimodal wrapper (resolving its inner language_model) is deferred to a follow-up PR.
MTP (nextn) + CP-v2 and DSA + CP-v2 remain out of scope (v1-only use_cp_v1 gates make v2 a no-op there).

Checklist

Format with pre-commit
Accuracy verified for DeepSeek V3 (see table)
Speed/profiling (TODO)

CI States

Latest PR Test (Base): ❌ Run #29629808219
Latest PR Test (Extra): ❌ Run #29629808052

gemini-code-assist · 2026-07-17T22:45:08Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

Route MLA prefill context parallelism (DeepSeek V3/R1) through the layers/cp ZigzagCPStrategy abstraction, mirroring the GQA path from #28421. DeepSeek V3 is added to CP_V2_DEFAULT_MODEL_CLASSES so CP-v2 is the default for it (as with Qwen3Moe); the legacy cp_utils v1 path stays intact and reachable via SGLANG_ENABLE_CP_V2=0. Under v2 the eager runner owns the entry shard (cp_split_before_forward) and exit gather (cp_gather_after_forward) and drives DeepseekV2Model.forward directly, so: - deepseek_v2.py / deepseek_nextn.py: gate the in-model CP entry split and exit gather to the v1-only path (`and not is_cp_v2_active`). - flashattention_backend.py: MLA CP attention dispatches through cp_strategy.run_attention under v2, else the v1 cp_attn_forward_extend. The MLA latent all-gather (rebuild_cp_kv_cache) stays upstream in forward_absorb_prepare for both paths; zigzag metadata is field-compatible with the v1 ContextParallelMetadata, so no KV-write change is needed. Pure-MLA arch only: DSA archs (DeepseekV32/DeepseekV4) are excluded (CP-v2 does not yet cover DSA), and Kimi K2.5 needs separate wiring for its .language_model wrapper. The existing test_deepseek_v3_cp_single_node.py (tp8 dp2 attn-cp4) now exercises CP-v2 via the default-class enablement. Part of the CP refactor roadmap (#27252) / context parallelism roadmap (#21788). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Decision A — move the MLA latent all-gather into the CP strategy for v2: - forward_mla.py: skip the upstream rebuild_cp_kv_cache under CP-v2 (k_nope/k_pe stay rank-local through forward_absorb_core). - flashattention_backend.py: the MLA KV write dispatches to cp_strategy.materialize_full_mla_kv under v2 (gather rank-local latent to full, then set_mla_kv_buffer); v1 keeps the upstream-rebuilt full write. - zigzag.py: add materialize_full_mla_kv (fuse compressed-KV + rope, one all-gather via gather_kv_cache, split back for the paged write). Kimi K2.5 wiring: - eager_runner.py: drive resolve_cp_forward_model(model) instead of the raw model, so multimodal wrappers expose their inner CausalLM for the CP-v2 path. - kimi_k25.py: KimiK25ForConditionalGeneration.get_cp_model() returns language_model. - cp/utils.py: add resolve_cp_forward_model + KimiK25ForConditionalGeneration to CP_V2_DEFAULT_MODEL_CLASSES. The seq_lens_k padding flagged on #28421 is already handled (zigzag build_metadata bakes pad_len; the FA backend widens page_table by pad_delta); Decision A's gather preserves the same padded length. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

deepseek_v2.py: compute use_cp_v1 once and reuse it for the entry split and exit gather instead of repeating the (dsa/mla prefill CP) and not is_cp_v2_active condition. deepseek_nextn.py: rename use_cp -> use_cp_v1 to match. Pure refactor, no behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Replace the bespoke get_cp_model() hook with the standard multimodal get_language_model() accessor (already implemented by Kimi's siblings: pixtral/voxtral/deepseek_ocr/...). resolve_cp_forward_model now unwraps via get_language_model when present, else returns the model itself. Behavior is identical for flat CausalLMs (DeepseekV3/Qwen3 have no get_language_model -> resolve to self); Kimi gains the standard accessor it was missing. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Remove the resolve_cp_forward_model helper and inline the standard get_language_model unwrap at its single call site in the eager runner, matching the getattr-for-optional-capability style already used there. Flat CausalLMs resolve to themselves (no get_language_model); multimodal wrappers (Kimi) unwrap to their inner CausalLM. No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Move the scattered CP-v2 prep/shard/gather/logits out of _execute_extend into a single self-contained _execute_extend_cp_v2 method, so the main dispatch is just `elif cp_v2_active: ret = self._execute_extend_cp_v2(...)` with no leaked forward_positions / cp_model. Drop the multimodal language-model resolution (and Kimi from CP_V2_DEFAULT_MODEL_CLASSES) — deferred to a follow-up PR; this PR targets DeepSeek V3 only. Validated: test_deepseek_v3_cp_single_node.py GSM8k 0.964. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

kpham-sgl · 2026-07-18T03:56:07Z

I will defer Kimi changes to another PR. Need to come up with a clean way for MM models

gemini-code-assist · 2026-07-18T03:56:12Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Fridge003 · 2026-07-18T05:37:39Z

MLA/DSA tests passed here
https://github.com/sgl-project/sglang/actions/runs/29629808052/job/88041219215

…gl-project#31619) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

github-actions Bot added the deepseek label Jul 17, 2026

kpham-sgl force-pushed the cp-v2-mla-prefill branch from 83c9bfa to 9c1d125 Compare July 18, 2026 00:05

kpham-sgl changed the title ~~[CP] Migrate MLA prefill CP to CP-v2 zigzag strategy (v1 + v2 coexist)~~ [CP] Migrate MLA prefill CP to CP-v2 zigzag strategy Jul 18, 2026

kpham-sgl commented Jul 18, 2026

View reviewed changes

Comment thread python/sglang/srt/models/deepseek_v2.py Outdated

Comment thread python/sglang/srt/models/deepseek_nextn.py

kpham-sgl and others added 5 commits July 18, 2026 01:26

[CP] Remove stale comment on CP_V2_DEFAULT_MODEL_CLASSES

ec9f915

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

kpham-sgl changed the title ~~[CP] Migrate MLA prefill CP to CP-v2 zigzag strategy~~ [CP] Migrate MLA prefill CP (DeepSeek V3) to CP-v2 zigzag strategy Jul 18, 2026

kpham-sgl marked this pull request as ready for review July 18, 2026 03:56

kpham-sgl requested review from BBuf, Edwardf0t1, Fridge003, HaiShaw, Qiaolin-Yu, Ying1123, ch-wan, fzyzcjy, hebiao064, hnyls2002, ispobock and merrymercy as code owners July 18, 2026 03:56

kpham-sgl added run-ci run-ci-extra labels Jul 18, 2026

[CP] Drop docstring from materialize_full_mla_kv

f0cba34

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Fridge003 approved these changes Jul 18, 2026

View reviewed changes

Fridge003 merged commit 7a89621 into main Jul 18, 2026
111 of 132 checks passed

Fridge003 deleted the cp-v2-mla-prefill branch July 18, 2026 05:37

Fridge003 mentioned this pull request Jul 18, 2026

[CP] Add CP-v2 support for Kimi-Linear #31661

Closed

vincentzed mentioned this pull request Jul 19, 2026

[Feature] Decode context parallelism (DCP) for DSA models (DeepSeek V3.2, GLM-5.x) vincentzed/sglang#13

Open

Fridge003 mentioned this pull request Jul 19, 2026

[Roadmap]Prefill Context Parallel Refactor #27252

Open

19 tasks

Zhylkaaa pushed a commit to Zhylkaaa/sglang that referenced this pull request Jul 29, 2026

[CP] Migrate MLA prefill CP (DeepSeek V3) to CP-v2 zigzag strategy (s…

4cb2baa

…gl-project#31619) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026

[CP] Migrate MLA prefill CP (DeepSeek V3) to CP-v2 zigzag strategy (s…

82559fa

…gl-project#31619) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Navigation Menu

[CP] Migrate MLA prefill CP (DeepSeek V3) to CP-v2 zigzag strategy - #31619

[CP] Migrate MLA prefill CP (DeepSeek V3) to CP-v2 zigzag strategy#31619
Fridge003 merged 8 commits into
mainfrom
cp-v2-mla-prefill

kpham-sgl commented Jul 17, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot commented Jul 17, 2026

Uh oh!

Uh oh!

Uh oh!

kpham-sgl commented Jul 18, 2026

Uh oh!

gemini-code-assist Bot commented Jul 18, 2026

Uh oh!

Fridge003 commented Jul 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

kpham-sgl commented Jul 17, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Out of scope (follow-ups)

Checklist

CI States

Uh oh!

gemini-code-assist Bot commented Jul 17, 2026

Uh oh!

Uh oh!

Uh oh!

kpham-sgl commented Jul 18, 2026

Uh oh!

gemini-code-assist Bot commented Jul 18, 2026

Uh oh!

Fridge003 commented Jul 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kpham-sgl commented Jul 17, 2026 •

edited by github-actions Bot

Loading