[3/N][CP] Implement zigzag CP strategy by Fridge003 · Pull Request #28421 · sgl-project/sglang

Fridge003 · 2026-06-16T09:30:39Z

Summary

Implement the CP v2 zigzag strategy for Qwen3/FlashAttention, gated behind SGLANG_ENABLE_CP_V2=1.
Add model-runner CP v2 metadata preparation, first-layer sequence sharding, final hidden-state gathering, and FlashAttention KV gather/save dispatch without changing legacy cp_utils.py.
Add CP strategy unit tests plus a Qwen3 30B DeepEP e2e subtest with --attn-cp-size 4 --ep 4 --moe-a2a-backend deepep.

Notes

CP v2 disables fused RoPE set-KV writes while active so K/V cache materialization happens after zigzag all-gather/reordering.
Legacy CP behavior remains disabled unless SGLANG_ENABLE_CP_V2=1 is set.

CI States

Latest PR Test (Base): 🚫 Run #27791684107
Latest PR Test (Extra): ❌ Run #27791684005

gemini-code-assist

Code Review

This pull request implements the Zigzag Context Parallel (CP-v2) strategy, replacing previous placeholders with concrete implementations for metadata building, hidden state and KV cache sharding/gathering, and attention dispatch. It also integrates these changes into the model runner and adds corresponding unit and integration tests. The review feedback highlights two critical issues: first, a device fallback mechanism in zigzag.py that may fail on non-CUDA platforms (such as NPU or MUSA) and should instead query the global server arguments; second, a bug in _forward_extend_cp_v2 where the auxiliary hidden states tuple is not correctly returned on non-last pipeline parallel ranks when capture_aux_hidden_states is enabled.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Fridge003 · 2026-06-18T04:05:48Z

/rerun-test test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_gqa_preill_cp.py

github-actions · 2026-06-18T04:06:09Z

Results for /rerun-test test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_gqa_preill_cp.py:

🚀 ubuntu-latest (1 test): ❌ View workflow run

cd test/ && python3 registered/cp/test_cp_strategy_unit.py

🚀 4-gpu-h100 (1 test): ❌ View workflow run

cd test/ && python3 registered/cp/test_gqa_preill_cp.py

Fridge003 · 2026-06-18T04:10:37Z

/rerun-test test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_gqa_preill_cp.py

github-actions · 2026-06-18T04:11:03Z

Results for /rerun-test test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_gqa_preill_cp.py:

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_cp_strategy_unit.py

🚀 4-gpu-h100 (1 test): ❌ View workflow run

cd test/ && python3 registered/cp/test_gqa_preill_cp.py

gemini-code-assist · 2026-06-18T06:42:41Z

Warning

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

Fridge003 · 2026-06-18T06:45:56Z

/rerun-test test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_gqa_preill_cp.py

github-actions · 2026-06-18T06:46:21Z

Results for /rerun-test test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_gqa_preill_cp.py:

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_cp_strategy_unit.py

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_gqa_preill_cp.py

Fridge003 · 2026-06-18T08:32:16Z

/rerun-test test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_gqa_preill_cp.py test/registered/cp/test_gqa_prefill_cp_legacy.py

github-actions · 2026-06-18T08:32:39Z

Results for /rerun-test test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_gqa_preill_cp.py test/registered/cp/test_gqa_prefill_cp_legacy.py:

🚀 ubuntu-latest (1 test): ❌ View workflow run

cd test/ && python3 registered/cp/test_cp_strategy_unit.py

🚀 4-gpu-h100 (2 tests): ❌ View workflow run

cd test/ && python3 registered/cp/test_gqa_preill_cp.py
cd test/ && python3 registered/cp/test_gqa_prefill_cp_legacy.py

Fridge003 · 2026-06-18T08:37:41Z

/rerun-test test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_gqa_preill_cp.py test/registered/cp/test_gqa_prefill_cp_legacy.py

github-actions · 2026-06-18T08:38:06Z

Results for /rerun-test test/registered/cp/test_cp_strategy_unit.py test/registered/cp/test_gqa_preill_cp.py test/registered/cp/test_gqa_prefill_cp_legacy.py:

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_cp_strategy_unit.py

🚀 4-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/cp/test_gqa_preill_cp.py
cd test/ && python3 registered/cp/test_gqa_prefill_cp_legacy.py

Fridge003 · 2026-06-18T19:58:38Z

/rerun-test test/registered/unit/server_args/test_server_args.py

github-actions · 2026-06-18T19:59:05Z

Results for /rerun-test test/registered/unit/server_args/test_server_args.py:

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/server_args/test_server_args.py

Fridge003 · 2026-06-18T21:39:03Z

/rerun-test test/registered/cp/test_deepseek_v32_cp_single_node.py test/registered/cp/test_deepseek_v3_cp_single_node.py test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py

github-actions · 2026-06-18T21:39:27Z

Results for /rerun-test test/registered/cp/test_deepseek_v32_cp_single_node.py test/registered/cp/test_deepseek_v3_cp_single_node.py test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py:

🚀 8-gpu-h200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/cp/test_deepseek_v32_cp_single_node.py
cd test/ && python3 registered/cp/test_deepseek_v3_cp_single_node.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py

Fridge003 · 2026-06-18T21:57:40Z

/rerun-test registered/cp/test_gqa_preill_cp.py

github-actions · 2026-06-18T21:58:00Z

Results for /rerun-test registered/cp/test_gqa_preill_cp.py:

🚀 4-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/cp/test_gqa_preill_cp.py

kpham-sgl

Mostly self reminders, thanks for the refactor

kpham-sgl · 2026-06-18T23:26:53Z

    ) -> Any:
-        raise NotImplementedError(
-            "Zigzag attention dispatch will land in a follow-up PR"
+        assert (


Need similar logic here, as currently zigzag's build_metadata does not modify / is not aware of seq_lens_k
Without this padding, for CP=4 in GQA you may run into IMA

kpham-sgl · 2026-06-18T23:41:27Z

Be careful with the padding happens inside

sglang/python/sglang/srt/model_executor/model_runner.py

Line 3311 in 62ab09a

forward_batch.prepare_mlp_sync_batch(self)

kpham-sgl · 2026-06-18T23:44:07Z

            self.assertIsNotNone(get_cp_strategy())


+class TestCPZigzagStrategy(CustomTestCase):


May need to fuse / add attention computation unit tests (like the new attention test kit)

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>

Implement CP v2 zigzag for Qwen3

29345f0

Fridge003 changed the title ~~[codex] Implement CP v2 zigzag for Qwen3~~ [3/N][CP] Implement zigzag strategy Jun 16, 2026

gemini-code-assist Bot reviewed Jun 16, 2026

View reviewed changes

Comment thread python/sglang/srt/layers/cp/zigzag.py Outdated

Comment thread python/sglang/srt/model_executor/model_runner.py Outdated

Fridge003 mentioned this pull request Jun 16, 2026

[Roadmap]Prefill Context Parallel Refactor #27252

Open

19 tasks

Fridge003 changed the title ~~[3/N][CP] Implement zigzag strategy~~ [3/N][CP] Implement zigzag CP strategy Jun 16, 2026

Fridge003 commented Jun 18, 2026

View reviewed changes

Fridge003 force-pushed the codex/cp-v2-zigzag-qwen3 branch from 7971196 to 3bf9568 Compare June 18, 2026 04:10

Fridge003 commented Jun 18, 2026

View reviewed changes

Fridge003 marked this pull request as ready for review June 18, 2026 06:42

Fridge003 requested review from BBuf, Edwardf0t1, HaiShaw, Qiaolin-Yu, Ying1123, ch-wan, hebiao064, hnyls2002, ispobock and merrymercy as code owners June 18, 2026 06:42

Address CP v2 zigzag review comments

f6d6201

Fridge003 force-pushed the codex/cp-v2-zigzag-qwen3 branch from 3bf9568 to f6d6201 Compare June 18, 2026 06:45

Merge branch 'main' into codex/cp-v2-zigzag-qwen3

85f3263

Fridge003 commented Jun 18, 2026

View reviewed changes

Address CP v2 review comments

0030086

Fridge003 force-pushed the codex/cp-v2-zigzag-qwen3 branch from 706f221 to 0030086 Compare June 18, 2026 08:37

Fridge003 added the run-ci label Jun 18, 2026

kpham-sgl self-assigned this Jun 18, 2026

Fix server args CP unit test

02c77b7

Use CP v2 flags in GQA CP test

58819f8

Fridge003 merged commit e3026ef into main Jun 18, 2026
96 of 136 checks passed

Fridge003 deleted the codex/cp-v2-zigzag-qwen3 branch June 18, 2026 22:10

kpham-sgl reviewed Jun 18, 2026

View reviewed changes

Fridge003 mentioned this pull request Jul 7, 2026

Fix FA3 prefill CP NaNs #30439

Merged

mmangkad mentioned this pull request Jul 11, 2026

Fix DeepEP CI test registration #30873

Merged

Fridge003 mentioned this pull request Jul 16, 2026

Support MiMo V2.5 with zigzag context parallelism #29972

Merged

kpham-sgl mentioned this pull request Jul 17, 2026

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

Merged

3 tasks

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

[3/N][CP] Implement zigzag CP strategy (sgl-project#28421)

09bf49b

		self.assertIsNotNone(get_cp_strategy())


		class TestCPZigzagStrategy(CustomTestCase):

Conversation

Fridge003 commented Jun 16, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Notes

CI States

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Fridge003 commented Jun 18, 2026

Uh oh!

github-actions Bot commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fridge003 commented Jun 18, 2026

Uh oh!

github-actions Bot commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot commented Jun 18, 2026

Uh oh!

Fridge003 commented Jun 18, 2026

Uh oh!

github-actions Bot commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Fridge003 commented Jun 18, 2026

Uh oh!

github-actions Bot commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fridge003 commented Jun 18, 2026

Uh oh!

github-actions Bot commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fridge003 commented Jun 18, 2026

Uh oh!

github-actions Bot commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fridge003 commented Jun 18, 2026

Uh oh!

github-actions Bot commented Jun 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Fridge003 commented Jun 16, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Jun 18, 2026 •

edited

Loading

github-actions Bot commented Jun 18, 2026 •

edited

Loading

github-actions Bot commented Jun 18, 2026 •

edited

Loading

github-actions Bot commented Jun 18, 2026 •

edited

Loading

github-actions Bot commented Jun 18, 2026 •

edited

Loading

github-actions Bot commented Jun 18, 2026 •

edited

Loading

github-actions Bot commented Jun 18, 2026 •

edited

Loading

github-actions Bot commented Jun 18, 2026 •

edited

Loading