[CP] Add CP-v2 support for Kimi-Linear by Fridge003 · Pull Request #31661 · sgl-project/sglang

Fridge003 · 2026-07-18T07:21:41Z

Summary

add a Kimi-specific CP-v2 layer communicator under layers/cp
gather model-entry shards before KDA, split full TP batches before MLA, and gather MLA outputs/residuals before post-attention norm and MoE/MLP
shard KDA projection, recurrent-cache, and state heads over attention CP ranks; all-reduce KDA outputs over the CP group, with global-TP fallback when CP is disabled
support zigzag CP-v2 in the GB300-default FlashInfer MLA backend using two paged MLA plans and materialized full latent KV
enable the Kimi-Linear model in the CP-v2 eager path and expose its input embeddings

PR #31619 has merged; this PR is rebased onto its CP-v2 implementation on main.

Layout transitions

Boundary	Operation
model entry → first KDA	gather CP shard into full token order
full TP MLP output → MLA	split into the rank-local zigzag shard
MLA attention → post-attention norm/MLP	gather hidden state and residual into full token order
final MLP → generic CP-v2 model exit	reshard for the existing model-exit gather

Gathering immediately after MLA is necessary because the following TP MoE/MLP all-reduces require identical token dimensions on all ranks.

Verification

Focused CPU tests on current main:

python3 -m pytest -q \
  test/registered/cp/test_kimi_linear_cp_v2.py \
  test/registered/cp/test_cp_strategy_unit.py

27 passed, 16 warnings, 18 subtests passed

All pre-commit hooks pass for the changed files, and the final diff passes git diff --check.

Latest targeted /rerun-test checks on commit fd7edf561e:

CP-v2 CPU coverage: 23 tests passed (run)
Kimi Linear GSM8K e2e on 2×H100: passed, score 0.915 (run)

The TP4/CP4 GB300 result below predates the KDA head-sharding update. Repeating it on baizhou-dev-2 is currently blocked because the only advertised free GB300 nodes are NodeStatusUnknown.

Four-GB300 launch on baizhou-dev-2:

python3 -m sglang.launch_server \
  --model-path /mirror/models/Kimi-Linear-48B-A3B-Instruct \
  --tp 4 \
  --attn-cp-size 4 \
  --enable-prefill-cp \
  --cp-strategy zigzag \
  --trust-remote-code

GSM8K, 200 examples, 32 concurrent requests, default decode CUDA graphs:

Score: 0.920
Total latency: 21.471 s
Output throughput: 869.357 token/s

The TP4 non-CP accuracy reference on the same model and dataset subset was 0.890.

CI States

Latest PR Test (Base): ❌ Run #29684396607
Latest PR Test (Extra): ❌ Run #29684396500

gemini-code-assist

Code Review

This pull request implements CP-v2 context-parallel prefill support for Kimi-Linear, alternating between Kimi Delta Attention (KDA) on complete token batches and Multi-head Latent Attention (MLA) on zigzag token shards. It introduces KimiLinearCPV2LayerCommunicator to handle layout transitions, updates FlashInfer MLA to support CP-v2 prefill, ensures KDA state parameters and cache shapes use global tensor parallelism, and adds comprehensive unit tests. Feedback suggests removing the unused stream parameter from the postprocess_layer method signature in the communicator.

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.

gemini-code-assist · 2026-07-18T07:24:11Z

+    def postprocess_layer(
+        self,
+        hidden_states: torch.Tensor,
+        residual: Optional[torch.Tensor],
+        forward_batch: ForwardBatch,
+        stream: Optional[Any] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:


The stream parameter in postprocess_layer is unused. It is recommended to remove it from the method signature to keep the API clean and maintainable.

def postprocess_layer( self, hidden_states: torch.Tensor, residual: Optional[torch.Tensor], forward_batch: ForwardBatch, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:

Fridge003 · 2026-07-18T09:25:38Z

Remove the assertion of deepep?

Fridge003 · 2026-07-19T10:48:25Z

/rerun-test

github-actions · 2026-07-19T10:48:47Z

⛔ Please specify a test: /rerun-test <file>::<TestClass.test_method>

Examples:

/rerun-test test/registered/core/test_srt_endpoint.py::TestSRTEndpoint.test_simple_decode
/rerun-test registered/core/test_srt_endpoint.py::TestSRTEndpoint
/rerun-test test_srt_endpoint.py
/rerun-test test_a.py test_b.py test_c.py (multiple tests)
/rerun-test test_*backend*.py (wildcard — reruns every matching file; wrap the pattern in backticks so GitHub keeps the * literal)

Fridge003 · 2026-07-19T10:49:49Z

/rerun-test test/registered/models_e2e/test_kimi_linear_models.py::TestKimiLinear.test_gsm8k

github-actions · 2026-07-19T10:50:09Z

Results for /rerun-test test/registered/models_e2e/test_kimi_linear_models.py::TestKimiLinear.test_gsm8k:

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

cd test/ && python3 registered/models_e2e/test_kimi_linear_models.py TestKimiLinear.test_gsm8k

Fridge003 · 2026-07-19T10:51:04Z

/rerun-test test/registered/cp/test_kimi_linear_cp_v2.py

github-actions · 2026-07-19T10:51:28Z

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

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

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

Fridge003 · 2026-07-19T10:55:47Z

/rerun-test test/registered/cp/test_kimi_linear_cp_v2.py test/registered/models_e2e/test_kimi_linear_models.py::TestKimiLinear.test_gsm8k

github-actions · 2026-07-19T10:56:09Z

Results for /rerun-test test/registered/cp/test_kimi_linear_cp_v2.py test/registered/models_e2e/test_kimi_linear_models.py::TestKimiLinear.test_gsm8k:

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

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

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

cd test/ && python3 registered/models_e2e/test_kimi_linear_models.py TestKimiLinear.test_gsm8k

Fridge003 · 2026-07-19T11:01:42Z

/rerun-test test/registered/cp/test_kimi_linear_cp_v2.py test/registered/models_e2e/test_kimi_linear_models.py::TestKimiLinear.test_gsm8k

github-actions · 2026-07-19T11:02:07Z

Results for /rerun-test test/registered/cp/test_kimi_linear_cp_v2.py test/registered/models_e2e/test_kimi_linear_models.py::TestKimiLinear.test_gsm8k:

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

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

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

cd test/ && python3 registered/models_e2e/test_kimi_linear_models.py TestKimiLinear.test_gsm8k

Fridge003 added 30 commits July 18, 2026 00:03

docs: design Kimi-Linear CP-v2 transitions

340cc02

docs: plan Kimi-Linear CP-v2 implementation

f3f6bf4

test: specify Kimi-Linear CP-v2 entry gather

1d54fc2

feat: gather Kimi-Linear CP input for KDA

407354f

test: specify KDA to MLA CP split

3f452cd

feat: shard Kimi-Linear state before MLA

630f45d

test: specify MLA to KDA CP gather

d5921b9

feat: gather Kimi-Linear residual before KDA

3c213d6

test: specify Kimi decoder CP communicator wiring

7b24010

feat: wire CP-v2 transitions into Kimi layers

2d7e74d

test: complete Kimi decoder fixture

7f9362e

test: isolate Kimi MLA layer wiring

92ce1af

test: require CP-v2 for Kimi-Linear

70a8767

feat: enable CP-v2 for Kimi-Linear

898ac51

test: require Kimi input embedding accessor

cd71780

feat: expose Kimi input embeddings to CP-v2

8754d7e

test: cover Kimi CP-v2 no-op and zigzag round trip

78c2083

test: require KDA heads to use global TP

02f464b

fix: partition KDA backend heads over global TP

4f6220a

test: require unfused KDA projection to use global TP

bf733d9

fix: shard unfused KDA projections over global TP

f0364b5

test: require KDA cache state to use global TP

01b92d0

fix: shard KDA state cache over global TP

e2ecec6

test: cover Kimi CP-v2 embedding keyword

55f2a64

fix: accept CP-v2 input embeddings in Kimi

b41d63a

test: cover FlashInfer MLA CP-v2 dispatch

5af9dd8

feat: support FlashInfer MLA in CP-v2

34ced26

test: match FlashInfer MLA output layout

5ff97cc

fix: run Kimi KDA and MLP on global TP batches

37c791f

docs: target merged CP-v2 base

697fff4

Fridge003 requested review from BBuf, Edwardf0t1, HaiShaw, Qiaolin-Yu, Ying1123, ch-wan, hebiao064, ispobock, sogalin, wisclmy0611 and zijiexia as code owners July 18, 2026 07:21

github-actions Bot added the documentation Improvements or additions to documentation label Jul 18, 2026

gemini-code-assist Bot reviewed Jul 18, 2026

View reviewed changes

docs: remove implementation planning files

8e9d306

Fridge003 marked this pull request as draft July 18, 2026 07:41

feat: shard KDA heads over CP ranks

c0f5526

test: avoid CUDA stream in Kimi CP CPU test

44e30b2

test: mock optional FlashInfer MLA wrapper

fd7edf5

Fridge003 closed this Jul 20, 2026

Navigation Menu

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

[CP] Add CP-v2 support for Kimi-Linear#31661
Fridge003 wants to merge 35 commits into
sgl-project:mainfrom
Fridge003:codex/kimi-linear-cp-v2-main

Fridge003 commented Jul 18, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 18, 2026

Uh oh!

Fridge003 commented Jul 18, 2026

Uh oh!

Fridge003 commented Jul 19, 2026

Uh oh!

github-actions Bot commented Jul 19, 2026

Uh oh!

Fridge003 commented Jul 19, 2026

Uh oh!

github-actions Bot commented Jul 19, 2026 •

edited

Loading

Uh oh!

Fridge003 commented Jul 19, 2026

Uh oh!

github-actions Bot commented Jul 19, 2026 •

edited

Loading

Uh oh!

Fridge003 commented Jul 19, 2026

Uh oh!

github-actions Bot commented Jul 19, 2026 •

edited

Loading

Uh oh!

Fridge003 commented Jul 19, 2026

Uh oh!

github-actions Bot commented Jul 19, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

Fridge003 commented Jul 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Layout transitions

Verification

CI States

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 18, 2026

Choose a reason for hiding this comment

Uh oh!

Fridge003 commented Jul 18, 2026

Uh oh!

Fridge003 commented Jul 19, 2026

Uh oh!

github-actions Bot commented Jul 19, 2026

Uh oh!

Fridge003 commented Jul 19, 2026

Uh oh!

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

Uh oh!

Uh oh!

Fridge003 commented Jul 19, 2026

Uh oh!

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

Uh oh!

Uh oh!

Fridge003 commented Jul 19, 2026

Uh oh!

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

Uh oh!

Uh oh!

Fridge003 commented Jul 19, 2026

Uh oh!

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

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Fridge003 commented Jul 18, 2026 •

edited

Loading

github-actions Bot commented Jul 19, 2026 •

edited

Loading

github-actions Bot commented Jul 19, 2026 •

edited

Loading

github-actions Bot commented Jul 19, 2026 •

edited

Loading

github-actions Bot commented Jul 19, 2026 •

edited

Loading