From c8257db4fae5e2f61241ee38ec9ffbab2b76119c Mon Sep 17 00:00:00 2001
From: abing <wangbingjia.wbj@ant-intl.com>
Date: Tue, 15 Sep 2026 23:41:20 -0700
Subject: [PATCH 01/27] run pass llm cp

---
 python/sglang/srt/disaggregation/common/conn.py     |  1 -
 .../srt/model_executor/runner/eager_runner.py       | 13 ++++++++++---
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py
index db2b6d14ac86..23d9c7e99ffb 100644
--- a/python/sglang/srt/disaggregation/common/conn.py
+++ b/python/sglang/srt/disaggregation/common/conn.py
@@ -1214,7 +1214,6 @@ def try_ensure_parallel_info(self, bootstrap_addr: str) -> bool:
                     "enable DSpark with the same block size and target/draft KV "
                     "layout. Upgrade both servers together."
                 )
-
             decode_tp_size_aligns_prefill_world_size = (
                 info.attn_cp_size > 1
                 and (self.is_mla_backend or self.is_hybrid_mla_backend)
diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py
index b1f21a26dad0..271578f12f29 100644
--- a/python/sglang/srt/model_executor/runner/eager_runner.py
+++ b/python/sglang/srt/model_executor/runner/eager_runner.py
@@ -363,14 +363,21 @@ def _execute_extend_cp(
         """
         model = self.model_runner.model
 
+        input_ids = forward_batch.input_ids
         if prepare_inputs := getattr(model, "prepare_cp_inputs", None):
             input_embeds, positions, model_kwargs = prepare_inputs(
                 forward_batch, **kwargs
             )
         else:
             input_embeds = kwargs.get("input_embeds")
+            # Vision preparation needs global token order before CP sharding.
+            prepare_language = getattr(model, "prepare_language_model_inputs", None)
+            if prepare_language is not None:
+                input_ids, input_embeds = prepare_language(
+                    input_ids, forward_batch, input_embeds
+                )
             if input_embeds is None:
-                input_embeds = model.get_input_embeddings()(forward_batch.input_ids)
+                input_embeds = model.get_input_embeddings()(input_ids)
             positions = forward_batch.positions
             model_kwargs = {}
             if (pp_proxy_tensors := kwargs.get("pp_proxy_tensors")) is not None:
@@ -380,7 +387,7 @@ def _execute_extend_cp(
             input_embeds,
             positions,
             forward_batch,
-            forward_batch.input_ids,
+            input_ids,
         ) as (sharded_input_embeds, sharded_positions, model_input_ids):
             model_kwargs["input_embeds"] = sharded_input_embeds
             hidden_states = model.model(
@@ -423,7 +430,7 @@ def _execute_extend_cp(
             if aux_hidden_states is None:
                 logits_kwargs["hidden_states_before_norm"] = hidden_states_before_norm
         return model.logits_processor(
-            forward_batch.input_ids,
+            input_ids,
             hidden_states,
             model.lm_head,
             forward_batch,

From e8ecdd8851e87bbcae30f315e36bd15b797a6058 Mon Sep 17 00:00:00 2001
From: abing <wangbingjia.wbj@ant-intl.com>
Date: Wed, 16 Sep 2026 00:26:02 -0700
Subject: [PATCH 02/27] add test

---
 python/sglang/test/dsv41_cp_test_utils.py     |  79 +++++
 .../disaggregation/test_dsv41_dspark_pd.py    | 286 ++++++++++++++++
 .../test_fake_transfer_chunk_progress.py      |  46 +++
 .../models/test_deepseek_v41_multimodal_cp.py | 316 ++++++++++++++++++
 .../unit/models/test_deepseek_v41_text_cp.py  | 219 ++++++++++++
 .../test_deepseek_v41_text_only_pd.py         |  73 ++++
 6 files changed, 1019 insertions(+)
 create mode 100644 python/sglang/test/dsv41_cp_test_utils.py
 create mode 100644 test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
 create mode 100644 test/registered/unit/mem_cache/test_fake_transfer_chunk_progress.py
 create mode 100644 test/registered/unit/models/test_deepseek_v41_multimodal_cp.py
 create mode 100644 test/registered/unit/models/test_deepseek_v41_text_cp.py
 create mode 100644 test/registered/unit/server_args/test_deepseek_v41_text_only_pd.py

diff --git a/python/sglang/test/dsv41_cp_test_utils.py b/python/sglang/test/dsv41_cp_test_utils.py
new file mode 100644
index 000000000000..5c6430078fa5
--- /dev/null
+++ b/python/sglang/test/dsv41_cp_test_utils.py
@@ -0,0 +1,79 @@
+"""Small CPU tensors; production CP slicing/gather, mocked collective transport."""
+
+from contextlib import ExitStack, contextmanager, nullcontext
+from types import SimpleNamespace as NS
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.layers.cp.interleave import InterleaveCPStrategy
+from sglang.srt.layers.cp.padding import pad_logical_token_to_physical
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+
+CP = "sglang.srt.layers.cp"
+
+
+@contextmanager
+def cp_context(size, rank, lengths=(3, 6), prefix_lengths=(7, 13)):
+    """Keep real interleave indexing/padding; replace only runtime context."""
+    strategy = InterleaveCPStrategy(size)
+    parallel = NS(attn_cp_size=size, attn_cp_rank=rank, attn_cp_group=None)
+    batch = NS(
+        forward_mode=ForwardMode.EXTEND,
+        input_ids=torch.arange(1, sum(lengths) + 1),
+        positions=torch.cat(
+            [
+                torch.arange(prefix, prefix + length)
+                for prefix, length in zip(prefix_lengths, lengths)
+            ]
+        ),
+        extend_seq_lens_cpu=list(lengths),
+        extend_prefix_lens_cpu=list(prefix_lengths),
+        mm_inputs=None,
+        spec_info=None,
+    )
+    batch.attn_cp_metadata = strategy.build_metadata(
+        sum(lengths), [p + n for p, n in zip(prefix_lengths, lengths)], list(lengths)
+    )
+    with ExitStack() as stack:
+        for module in ("base", "utils", "padding", "interleave"):
+            stack.enter_context(
+                patch(CP + "." + module + ".get_parallel", return_value=parallel)
+            )
+        stack.enter_context(patch(CP + ".utils.get_cp_strategy", return_value=strategy))
+        stack.enter_context(
+            patch(CP + ".padding.get_cp_padding_align_size", return_value=size)
+        )
+        stack.enter_context(
+            patch(
+                CP + ".utils.get_moe_a2a_backend", return_value=NS(is_none=lambda: True)
+            )
+        )
+        pad_logical_token_to_physical(batch.attn_cp_metadata)
+        yield strategy, batch
+
+
+@contextmanager
+def simulated_collective(strategy, batch, global_tensor):
+    """Inject peer buffers into all-gather; retain production unpadding/reordering."""
+    physical = max(batch.attn_cp_metadata.per_rank_actual_token)
+    buffers = []
+    for rank in range(strategy.cp_size):
+        buf = global_tensor.new_zeros((physical, *global_tensor.shape[1:]))
+        local = global_tensor[rank :: strategy.cp_size]
+        buf[: len(local)] = local
+        buffers.append(buf)
+
+    def gather(output, local):
+        torch.testing.assert_close(local, buffers[strategy.cp_rank], rtol=0, atol=0)
+        output.copy_(torch.cat(buffers))
+
+    with (
+        patch(
+            CP + ".interleave.use_symmetric_memory",
+            side_effect=lambda *a, **k: nullcontext(),
+        ),
+        patch(CP + ".interleave.is_allocation_symmetric", return_value=False),
+        patch(CP + ".interleave.attn_cp_all_gather_into_tensor", side_effect=gather),
+    ):
+        yield
diff --git a/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py b/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
new file mode 100644
index 000000000000..69e58c50579b
--- /dev/null
+++ b/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
@@ -0,0 +1,286 @@
+import asyncio
+import copy
+import json
+import unittest
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock, patch
+
+from sglang.srt.disaggregation.base.conn import StateType
+from sglang.srt.disaggregation.common.conn import (
+    CommonKVBootstrapServer,
+    CommonKVManager,
+)
+from sglang.srt.disaggregation.decode import DecodePreallocQueue
+from sglang.srt.disaggregation.utils import get_dsv41_spec_layout
+from sglang.srt.mem_cache.common import retraction_backup
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+from sglang.srt.runtime_context import get_context
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+def make_layout():
+    args = SimpleNamespace(
+        mla_compression_ratios=[0, 2, 1],
+        kv_layer_ids=[1, 2],
+        kv_item_lens=[512, 1024],
+        state_types=[StateType.SWA, StateType.DSV4_REQUEST_STATE, StateType.SWA],
+        state_item_lens=[[512], [32768], [512]],
+    )
+    with get_context().override_server_args(
+        speculative_algorithm="DSPARK", speculative_num_draft_tokens=6
+    ):
+        return get_dsv41_spec_layout(args)
+
+
+class TestDSV41DSparkPD(CustomTestCase):
+    def test_bootstrap_validates_before_caching(self):
+        layout = make_layout()
+        cases = [("matching", layout, layout, 4, True), ("legacy", None, None, 2, True)]
+        for key, value in (
+            ("num_draft_tokens", 5),
+            ("kv_layer_ids", [2, 1]),
+            ("kv_item_lens", [256, 1024]),
+            ("state_types", ["swa", "c128_state"]),
+            ("state_item_lens", [[512], [8192], [512]]),
+        ):
+            different = copy.deepcopy(layout)
+            different[key] = value
+            cases.append((key, layout, different, 4, False))
+        cases += [
+            ("prefill_only", None, layout, 4, False),
+            ("decode_only_or_old_prefill", layout, None, 4, False),
+            ("tp_mismatch", layout, layout, 2, False),
+        ]
+        for name, local, peer, tp_size, supported in cases:
+            with self.subTest(name=name):
+                manager = object.__new__(CommonKVManager)
+                manager.prefill_info_table = {}
+                manager.kv_args = SimpleNamespace(page_size=256)
+                manager.kv_cache_dtype_str = "fp8_e4m3"
+                manager.dsv41_spec_layout = local
+                manager.attn_tp_size = 4
+                manager.dcp_size = 1
+                manager._resolve_rank_mapping = Mock()
+                response = Mock(status_code=200)
+                response.json.return_value = dict(
+                    attn_tp_size=tp_size,
+                    attn_cp_size=1,
+                    dp_size=1,
+                    pp_size=1,
+                    page_size=256,
+                    kv_cache_dtype="fp8_e4m3",
+                    follow_bootstrap_room=True,
+                    dsv41_spec_layout=peer,
+                )
+                with patch(
+                    "sglang.srt.disaggregation.common.conn.requests.get",
+                    return_value=response,
+                ) as fetch:
+                    if supported:
+                        self.assertTrue(
+                            manager.try_ensure_parallel_info("prefill:8998")
+                        )
+                        self.assertTrue(
+                            manager.try_ensure_parallel_info("prefill:8998")
+                        )
+                        fetch.assert_called_once()
+                    else:
+                        with self.assertRaisesRegex(
+                            RuntimeError, "DeepSeek-V4.1 DSpark PD"
+                        ):
+                            manager.try_ensure_parallel_info("prefill:8998")
+                        self.assertFalse(manager.prefill_info_table)
+                        manager._resolve_rank_mapping.assert_not_called()
+
+    def test_python_bootstrap_preserves_layout_and_rejects_mixed_ranks(self):
+        with patch.object(CommonKVBootstrapServer, "run"):
+            server = CommonKVBootstrapServer("127.0.0.1", 8998)
+        layout = make_layout()
+        payload = dict(
+            attn_tp_size=1,
+            attn_tp_rank=0,
+            attn_cp_size=1,
+            attn_cp_rank=0,
+            attn_dp_size=1,
+            attn_dp_rank=0,
+            pp_size=1,
+            pp_rank=0,
+            system_dp_size=1,
+            system_dp_rank=0,
+            rank_ip="127.0.0.1",
+            rank_port=1234,
+            page_size=256,
+            kv_cache_dtype="fp8_e4m3",
+            dsv41_spec_layout=layout,
+        )
+        request = Mock(json=AsyncMock(return_value=payload))
+        self.assertEqual(asyncio.run(server._handle_route_put(request)).status, 200)
+        query = Mock(
+            query={
+                key: "-1"
+                for key in (
+                    "prefill_dp_rank",
+                    "prefill_cp_rank",
+                    "target_tp_rank",
+                    "target_pp_rank",
+                )
+            }
+        )
+        response = asyncio.run(server._handle_route_get(query))
+        self.assertEqual(json.loads(response.text)["dsv41_spec_layout"], layout)
+        payload["dsv41_spec_layout"] = None
+        self.assertEqual(asyncio.run(server._handle_route_put(request)).status, 400)
+        self.assertEqual(server._registered_count, 1)
+        self.assertEqual(server.dsv41_spec_layout, layout)
+
+    def test_retraction_recomputes_from_prefill_and_replays_boundary_token(self):
+        pool = object.__new__(DeepSeekV4TokenToKVPool)
+        pool.compression_ratios = [0, 2, 1]
+        pool.device = "cuda"
+        allocator = Mock(get_kvcache=Mock(return_value=pool))
+        for algorithm in (None, "DSPARK"):
+            with (
+                self.subTest(algorithm=algorithm),
+                get_context().override_server_args(speculative_algorithm=algorithm),
+                patch("torch.get_device_module") as device_module,
+            ):
+                req = SimpleNamespace(
+                    output_ids=[7, 8],
+                    bootstrap_host="prefill",
+                    time_stats=Mock(),
+                    offload_kv_cache=Mock(),
+                )
+                request_pool = Mock()
+                self.assertTrue(
+                    retraction_backup(
+                        req, Mock(), request_pool, allocator, "cpu_tensor"
+                    )
+                )
+                queue = SimpleNamespace(
+                    token_to_kv_pool_allocator=allocator,
+                    _check_if_req_exceed_kv_capacity=Mock(return_value=False),
+                    _create_receiver_and_enqueue=Mock(),
+                    _resolve_prefill_dp_rank=Mock(return_value=0),
+                    retracted_queue=[],
+                    pending_reqs=[],
+                )
+                DecodePreallocQueue.add(queue, req, is_retracted=True)
+                if algorithm == "DSPARK":
+                    req.offload_kv_cache.assert_not_called()
+                    device_module.return_value.synchronize.assert_called_once_with(
+                        "cuda"
+                    )
+                    self.assertEqual(req.output_ids, [7])
+                    self.assertEqual(req.pd_rebootstrap_forced_output_id, 8)
+                    self.assertTrue(req.pd_rebootstrap_in_progress)
+                    queue._create_receiver_and_enqueue.assert_called_once_with(
+                        req, is_rebootstrap=True
+                    )
+                    self.assertFalse(queue.retracted_queue)
+                else:
+                    req.offload_kv_cache.assert_called_once_with(
+                        request_pool, allocator
+                    )
+                    device_module.assert_not_called()
+                    self.assertEqual(req.output_ids, [7, 8])
+                    self.assertEqual(queue.retracted_queue, [req])
+
+
+class TestDSV41CPPDHandshake(CustomTestCase):
+    def make(self, rank=0, hybrid=True):
+        m = object.__new__(CommonKVManager)
+        m.prefill_info_table = {}
+        m.kv_args = SimpleNamespace(page_size=256, engine_rank=rank)
+        m.kv_cache_dtype_str = "fp8_e4m3"
+        m.dsv41_spec_layout = {"kv_item_lens": [512], "state_item_lens": [[32768]]}
+        m.attn_tp_size = 4
+        m.attn_cp_size = 1
+        m.attn_cp_rank = 0
+        m.dcp_size = 1
+        m.is_mla_backend = False
+        m.is_hybrid_mla_backend = hybrid
+        m.enable_all_cp_ranks_for_transfer = True
+        m.pp_size = 1
+        m.pp_rank = 0
+        return m
+
+    def fetch(self, m, tp, cp, layout=None):
+        response = Mock(status_code=200)
+        response.json.return_value = dict(
+            attn_tp_size=tp,
+            attn_cp_size=cp,
+            dp_size=1,
+            pp_size=1,
+            page_size=256,
+            kv_cache_dtype="fp8_e4m3",
+            follow_bootstrap_room=True,
+            dsv41_spec_layout=layout or m.dsv41_spec_layout,
+        )
+        with patch(
+            "sglang.srt.disaggregation.common.conn.requests.get", return_value=response
+        ):
+            return m.try_ensure_parallel_info("prefill:8761")
+
+    def test_cp4_maps_all_shards_to_each_decode_rank(self):
+        for rank in range(4):
+            m = self.make(rank)
+            self.assertTrue(self.fetch(m, 1, 4))
+            info = m.prefill_info_table["prefill:8761"]
+            self.assertEqual(info.target_tp_ranks, [0])
+            self.assertEqual(info.target_cp_ranks, [0, 1, 2, 3])
+            self.assertEqual(info.required_prefill_response_num, 4)
+            self.assertEqual(info.required_dst_info_num, 4)
+
+    def test_dsv4_pool_is_classified_as_mla(self):
+        from sglang.srt.disaggregation.utils import is_mla_backend
+        from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+
+        pool = object.__new__(DeepSeekV4TokenToKVPool)
+        self.assertTrue(is_mla_backend(pool))
+        m = self.make(hybrid=False)
+        m.is_mla_backend = is_mla_backend(pool)
+        self.assertTrue(self.fetch(m, 1, 4))
+        self.assertEqual(
+            m.prefill_info_table["prefill:8761"].required_prefill_response_num, 4
+        )
+
+    def test_cp2_tp2_maps_corresponding_tp_and_both_cp_ranks(self):
+        for rank in range(4):
+            m = self.make(rank)
+            self.assertTrue(self.fetch(m, 2, 2))
+            info = m.prefill_info_table["prefill:8761"]
+            self.assertEqual(info.target_tp_ranks, [rank // 2])
+            self.assertEqual(info.target_cp_ranks, [0, 1])
+            self.assertEqual(info.required_prefill_response_num, 2)
+
+    def test_plain_tp4_unchanged(self):
+        m = self.make(3)
+        self.assertTrue(self.fetch(m, 4, 1))
+        info = m.prefill_info_table["prefill:8761"]
+        self.assertEqual(info.target_tp_ranks, [3])
+        self.assertEqual(info.target_cp_ranks, [0])
+
+    def test_unequal_model_tp_rejected(self):
+        for tp, cp in [(2, 1), (1, 2), (1, 8)]:
+            m = self.make()
+            with self.assertRaisesRegex(RuntimeError, "same TP size"):
+                self.fetch(m, tp, cp)
+            self.assertFalse(m.prefill_info_table)
+
+    def test_nonhybrid_cp_mismatch_rejected(self):
+        m = self.make(hybrid=False)
+        with self.assertRaisesRegex(RuntimeError, "same TP size"):
+            self.fetch(m, 1, 4)
+
+    def test_layout_mismatch_still_rejected(self):
+        m = self.make()
+        with self.assertRaisesRegex(RuntimeError, "layout mismatch"):
+            self.fetch(m, 1, 4, {"kv_item_lens": [1024], "state_item_lens": [[32768]]})
+        self.assertFalse(m.prefill_info_table)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/registered/unit/mem_cache/test_fake_transfer_chunk_progress.py b/test/registered/unit/mem_cache/test_fake_transfer_chunk_progress.py
new file mode 100644
index 000000000000..840d77a5a469
--- /dev/null
+++ b/test/registered/unit/mem_cache/test_fake_transfer_chunk_progress.py
@@ -0,0 +1,46 @@
+"""Chunk continuation when fake PD transfer skips shared radix insertion."""
+
+import unittest
+from types import SimpleNamespace as NS
+from unittest.mock import Mock
+
+import torch
+
+from sglang.srt.mem_cache.common import maybe_cache_unfinished_req
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestFakeTransferChunkProgress(CustomTestCase):
+    def test_chunk_progress_without_shared_insert(self):
+        slots = torch.arange(16385, dtype=torch.int32).reshape(1, -1)
+        cache = NS(
+            req_to_token_pool=NS(req_to_token=slots), cache_unfinished_req=Mock()
+        )
+        req = NS(
+            skip_radix_cache_insert=True,
+            kv=NS(req_pool_idx=0, cache_protected_len=0),
+            get_fill_ids=lambda: range(16384),
+        )
+        maybe_cache_unfinished_req(req, cache, chunked=True)
+        self.assertEqual(16385 - len(req.prefix_indices), 1)
+        self.assertEqual(req.kv.cache_protected_len, 0)
+        cache.cache_unfinished_req.assert_not_called()
+        self.assertEqual(req.prefix_indices.dtype, torch.int64)
+        slots[0, 0] = -1
+        self.assertEqual(req.prefix_indices[0].item(), 0)
+        req.get_fill_ids = lambda: range(16385)
+        maybe_cache_unfinished_req(req, cache, chunked=True)
+        self.assertEqual(len(req.prefix_indices), 16385)
+
+    def test_real_transfer_preserves_cache_path(self):
+        req = NS(skip_radix_cache_insert=False)
+        cache = NS(cache_unfinished_req=Mock())
+        maybe_cache_unfinished_req(req, cache, chunked=True)
+        cache.cache_unfinished_req.assert_called_once_with(req, chunked=True)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/registered/unit/models/test_deepseek_v41_multimodal_cp.py b/test/registered/unit/models/test_deepseek_v41_multimodal_cp.py
new file mode 100644
index 000000000000..d0a52b632690
--- /dev/null
+++ b/test/registered/unit/models/test_deepseek_v41_multimodal_cp.py
@@ -0,0 +1,316 @@
+"""V4.1 image/text CP input contracts; vision and model compute are mocked."""
+
+import unittest
+from contextlib import contextmanager
+from types import SimpleNamespace as NS
+from unittest.mock import Mock, patch
+
+import torch
+
+from sglang.srt.managers.schedule_batch import MultimodalInputs
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.model_executor.runner.eager_runner import EagerRunner
+from sglang.srt.models.deepseek_v4 import MM_PAD_SHIFT_VALUE, DeepseekV4ForCausalLM
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.dsv41_cp_test_utils import cp_context, simulated_collective
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+MODEL = "sglang.srt.models.deepseek_v4"
+RUNNER = "sglang.srt.model_executor.runner.eager_runner"
+IMAGE_ID = 129264
+
+
+class TestDSV41MultimodalCP(CustomTestCase):
+    def test_image_spans_cross_ranks_before_shard_and_gather(self):
+        # Two image spans with distinct cache hashes; first request is text-only.
+        original = torch.tensor(
+            [
+                7,
+                8,
+                9,
+                MM_PAD_SHIFT_VALUE + 11,
+                MM_PAD_SHIFT_VALUE + 11,
+                10,
+                MM_PAD_SHIFT_VALUE + 23,
+                MM_PAD_SHIFT_VALUE + 23,
+                12,
+            ]
+        )
+        normalized = torch.tensor(
+            [7, 8, 9, IMAGE_ID, IMAGE_ID, 10, IMAGE_ID, IMAGE_ID, 12]
+        )
+        full = torch.arange(36, dtype=torch.float32).reshape(9, 4)
+        # Distinct image features expose using text embeddings or wrong row order.
+        full[3:5] += 1000
+        full[6:8] += 2000
+        for size in (2, 4):
+            for rank in range(size):
+                with (
+                    self.subTest(size=size, rank=rank),
+                    cp_context(size, rank) as (strategy, batch),
+                ):
+                    batch.input_ids = original.clone()
+                    batch.mm_inputs = [None, MultimodalInputs(mm_items=[])]
+                    model = NS(
+                        vision=object(),
+                        config=NS(image_token_id=IMAGE_ID),
+                        get_input_embeddings=Mock(
+                            side_effect=AssertionError(
+                                "Raw image hashes entered text embeddings"
+                            )
+                        ),
+                        _prepare_mm_embeddings=Mock(return_value=full),
+                        capture_aux_hidden_states=False,
+                        pp_group=NS(is_last_rank=True),
+                        lm_head=object(),
+                        logits_processor=Mock(return_value="ok"),
+                    )
+                    model.prepare_language_model_inputs = lambda ids, fb, emb: (
+                        DeepseekV4ForCausalLM.prepare_language_model_inputs(
+                            model, ids, fb, emb
+                        )
+                    )
+
+                    def body(ids, positions, fb, input_embeds):
+                        model._prepare_mm_embeddings.assert_called_once_with(
+                            batch.input_ids, batch
+                        )
+                        n = len(normalized[rank::size])
+                        torch.testing.assert_close(ids[:n], normalized[rank::size])
+                        torch.testing.assert_close(input_embeds[:n], full[rank::size])
+                        torch.testing.assert_close(
+                            positions[:n], batch.positions[rank::size]
+                        )
+                        self.assertFalse(
+                            (fb.input_ids_global >= MM_PAD_SHIFT_VALUE).any().item()
+                        )
+                        return input_embeds
+
+                    model.model = body
+                    with (
+                        simulated_collective(strategy, batch, full),
+                        patch(RUNNER + ".torch.cuda.current_stream", return_value=None),
+                    ):
+                        result = EagerRunner._execute_extend_cp(
+                            NS(model_runner=NS(model=model)), batch, {}
+                        )
+                    self.assertEqual(result, "ok")
+                    torch.testing.assert_close(
+                        model.logits_processor.call_args.args[0], normalized
+                    )
+                    torch.testing.assert_close(
+                        model.logits_processor.call_args.args[1], full
+                    )
+                    torch.testing.assert_close(batch.input_ids, original)
+                    self.assertFalse(hasattr(batch, "input_ids_global"))
+
+    def test_chunk_prefix_metadata_and_scheduler_hashes_survive_embedder(self):
+        for prefixes, lengths in (([0, 0], [3, 6]), ([16384, 127], [3, 6])):
+            with self.subTest(prefixes=prefixes):
+                ids = torch.tensor([7, 8, 9] + [MM_PAD_SHIFT_VALUE + 17] * 6)
+                original = ids.clone()
+                image = MultimodalInputs(mm_items=[])
+                batch = NS(
+                    mm_inputs=[None, image],
+                    extend_prefix_lens_cpu=prefixes,
+                    extend_seq_lens_cpu=lengths,
+                )
+                full = torch.arange(27, dtype=torch.float32).reshape(9, 3)
+                embedding = object()
+                model = NS(get_input_embeddings=lambda: embedding)
+
+                def embed(**kwargs):
+                    self.assertEqual(kwargs["extend_prefix_lens"], prefixes)
+                    self.assertEqual(kwargs["extend_seq_lens"], lengths)
+                    self.assertIs(kwargs["mm_inputs_list"][1], image)
+                    self.assertEqual(kwargs["mm_inputs_list"][0].mm_items, [])
+                    self.assertIs(kwargs["input_embedding"], embedding)
+                    self.assertNotEqual(kwargs["input_ids"].data_ptr(), ids.data_ptr())
+                    kwargs["input_ids"].zero_()
+                    return full, {}
+
+                with patch(MODEL + ".embed_mm_inputs", side_effect=embed) as mocked:
+                    result = DeepseekV4ForCausalLM._prepare_mm_embeddings(
+                        model, ids, batch
+                    )
+                mocked.assert_called_once()
+                self.assertIs(result, full)
+                self.assertIs(batch.mm_input_embeds, full)
+                torch.testing.assert_close(ids, original)
+
+    def test_vision_enabled_text_batch_skips_image_encoder(self):
+        for mm_inputs in (None, [None, None], []):
+            with self.subTest(mm_inputs=mm_inputs):
+                ids = torch.tensor([4, 5, 6])
+                model = NS(
+                    vision=object(),
+                    config=NS(image_token_id=IMAGE_ID),
+                    _prepare_mm_embeddings=Mock(),
+                )
+                batch = NS(forward_mode=ForwardMode.EXTEND, mm_inputs=mm_inputs)
+                result, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs(
+                    model, ids, batch
+                )
+                torch.testing.assert_close(result, ids)
+                self.assertIsNone(embeds)
+                model._prepare_mm_embeddings.assert_not_called()
+
+    def test_decode_idle_and_verify_preserve_vocab_ids(self):
+        for mode in (ForwardMode.DECODE, ForwardMode.IDLE, ForwardMode.TARGET_VERIFY):
+            with self.subTest(mode=mode):
+                ids = torch.tensor([4, IMAGE_ID, 6])
+                model = NS(
+                    vision=object(),
+                    config=NS(image_token_id=IMAGE_ID),
+                    _prepare_mm_embeddings=Mock(),
+                )
+                batch = NS(forward_mode=mode, mm_inputs=None)
+                result, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs(
+                    model, ids, batch
+                )
+                self.assertIs(result, ids)
+                self.assertIsNone(embeds)
+                model._prepare_mm_embeddings.assert_not_called()
+
+    def test_image_embedding_failure_does_not_mutate_scheduler_ids(self):
+        ids = torch.tensor([7, MM_PAD_SHIFT_VALUE + 12, 8])
+        original = ids.clone()
+        batch = NS(
+            mm_inputs=[MultimodalInputs(mm_items=[])],
+            extend_prefix_lens_cpu=[0],
+            extend_seq_lens_cpu=[3],
+        )
+        model = NS(get_input_embeddings=lambda: object())
+
+        def fail(**kwargs):
+            kwargs["input_ids"].zero_()
+            raise RuntimeError("vision failure")
+
+        with patch(MODEL + ".embed_mm_inputs", side_effect=fail):
+            with self.assertRaisesRegex(RuntimeError, "vision failure"):
+                DeepseekV4ForCausalLM._prepare_mm_embeddings(model, ids, batch)
+        torch.testing.assert_close(ids, original)
+        self.assertFalse(hasattr(batch, "mm_input_embeds"))
+
+
+class TestDSV41MultimodalInputs(CustomTestCase):
+    def setUp(self):
+        self.ids = torch.tensor(
+            [7, MM_PAD_SHIFT_VALUE + 12, MM_PAD_SHIFT_VALUE + 12, 9, 10]
+        )
+        self.original = self.ids.clone()
+        self.embeds = torch.arange(15, dtype=torch.float32).reshape(5, 3)
+        self.model = NS(
+            vision=object(),
+            config=NS(image_token_id=129264),
+            _prepare_mm_embeddings=Mock(return_value=self.embeds),
+        )
+        self.batch = NS(
+            input_ids=self.ids,
+            forward_mode=ForwardMode.EXTEND,
+            mm_inputs=[MultimodalInputs(mm_items=[])],
+        )
+
+    def test_prepare_global_embeddings_and_normalized_ids(self):
+        ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs(
+            self.model, self.ids, self.batch
+        )
+        self.assertEqual(ids.tolist(), [7, 129264, 129264, 9, 10])
+        self.assertIs(embeds, self.embeds)
+        self.model._prepare_mm_embeddings.assert_called_once_with(self.ids, self.batch)
+        self.assertTrue(torch.equal(self.ids, self.original))
+
+    def test_reject_preembedded_images(self):
+        with self.assertRaisesRegex(ValueError, "Cannot combine"):
+            DeepseekV4ForCausalLM.prepare_language_model_inputs(
+                self.model, self.ids, self.batch, self.embeds
+            )
+
+    def test_text_only_model_keeps_existing_embeddings(self):
+        self.model.vision = None
+        ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs(
+            self.model, self.ids, self.batch, self.embeds
+        )
+        self.assertIs(ids, self.ids)
+        self.assertIs(embeds, self.embeds)
+        self.model._prepare_mm_embeddings.assert_not_called()
+
+    def test_text_subclass_without_vision_module(self):
+        del self.model.vision
+        ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs(
+            self.model, self.ids, self.batch, self.embeds
+        )
+        self.assertIs(ids, self.ids)
+        self.assertIs(embeds, self.embeds)
+
+    def test_decode_keeps_vocabulary_ids(self):
+        self.batch.forward_mode = ForwardMode.DECODE
+        ids, embeds = DeepseekV4ForCausalLM.prepare_language_model_inputs(
+            self.model, self.ids, self.batch
+        )
+        self.assertIs(ids, self.ids)
+        self.assertIsNone(embeds)
+        self.model._prepare_mm_embeddings.assert_not_called()
+
+    def test_embedding_does_not_mutate_scheduler_hashes(self):
+        self.batch.extend_prefix_lens_cpu = [0]
+        self.batch.extend_seq_lens_cpu = [5]
+        self.model.get_input_embeddings = lambda: None
+
+        def embed(**kwargs):
+            kwargs["input_ids"].zero_()
+            return (self.embeds, {})
+
+        with patch("sglang.srt.models.deepseek_v4.embed_mm_inputs", side_effect=embed):
+            result = DeepseekV4ForCausalLM._prepare_mm_embeddings(
+                self.model, self.ids, self.batch
+            )
+        self.assertTrue(torch.equal(self.ids, self.original))
+        self.assertIs(result, self.batch.mm_input_embeds)
+
+    def test_cp_runner_prepares_before_sharding_and_uses_model_ids_for_logits(self):
+        normalized = torch.tensor([7, 129264, 129264, 9, 10])
+        self.batch.positions = torch.arange(5)
+        self.model.prepare_language_model_inputs = lambda ids, batch, emb: (
+            DeepseekV4ForCausalLM.prepare_language_model_inputs(
+                self.model, ids, batch, emb
+            )
+        )
+        self.model.get_input_embeddings = Mock(
+            side_effect=AssertionError("Raw hashes must not enter text embedding")
+        )
+        self.model.model = Mock(return_value=self.embeds[1::4])
+        self.model.capture_aux_hidden_states = False
+        self.model.pp_group = NS(is_last_rank=True)
+        self.model.lm_head = object()
+        self.model.logits_processor = Mock(return_value="ok")
+
+        @contextmanager
+        def shard(embeds, positions, batch, ids):
+            self.assertIs(embeds, self.embeds)
+            self.assertTrue(torch.equal(ids, normalized))
+            yield (embeds[1::4], positions[1::4], ids[1::4])
+
+        runner = NS(model_runner=NS(model=self.model))
+        module = "sglang.srt.model_executor.runner.eager_runner"
+        with (
+            patch(module + ".cp_shard_model_inputs", side_effect=shard),
+            patch(module + ".cp_gather_after_forward", return_value=self.embeds),
+            patch(module + ".torch.cuda.current_stream", return_value=None),
+        ):
+            result = EagerRunner._execute_extend_cp(runner, self.batch, {})
+        self.assertEqual(result, "ok")
+        args, kwargs = self.model.model.call_args
+        self.assertEqual(args[0].tolist(), [129264])
+        self.assertTrue(torch.equal(kwargs["input_embeds"], self.embeds[1::4]))
+        self.assertTrue(
+            torch.equal(self.model.logits_processor.call_args.args[0], normalized)
+        )
+        self.assertTrue(torch.equal(self.batch.input_ids, self.original))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/registered/unit/models/test_deepseek_v41_text_cp.py b/test/registered/unit/models/test_deepseek_v41_text_cp.py
new file mode 100644
index 000000000000..7689b118c01b
--- /dev/null
+++ b/test/registered/unit/models/test_deepseek_v41_text_cp.py
@@ -0,0 +1,219 @@
+"""Pure-language V4.1 CP input, padding and DSpark state regressions."""
+
+import unittest
+from types import SimpleNamespace as NS
+from unittest.mock import Mock, patch
+
+import torch
+
+from sglang.srt.layers.cp.utils import (
+    cp_gather_after_forward,
+    cp_shard_model_inputs,
+    is_cp_active,
+)
+from sglang.srt.model_executor.runner.eager_runner import EagerRunner
+from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.dsv41_cp_test_utils import cp_context, simulated_collective
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+RUNNER = "sglang.srt.model_executor.runner.eager_runner"
+
+
+class TestDSV41TextCP(CustomTestCase):
+    def test_interleave_roundtrip_mixed_lengths_prefix_and_padding(self):
+        for size in (2, 4):
+            for length in (4, 5, 9, 127, 128, 129):
+                for rank in range(size):
+                    with (
+                        self.subTest(size=size, length=length, rank=rank),
+                        cp_context(size, rank, (1, length - 1), (0, 16384)) as (
+                            strategy,
+                            batch,
+                        ),
+                    ):
+                        embeddings = torch.arange(
+                            length * 3, dtype=torch.float32
+                        ).reshape(length, 3)
+                        original_ids = batch.input_ids.clone()
+                        with cp_shard_model_inputs(
+                            embeddings, batch.positions, batch, batch.input_ids
+                        ) as (local, positions, ids):
+                            count = len(embeddings[rank::size])
+                            torch.testing.assert_close(
+                                local[:count], embeddings[rank::size]
+                            )
+                            torch.testing.assert_close(
+                                positions[:count], batch.positions[rank::size]
+                            )
+                            torch.testing.assert_close(
+                                ids[:count], batch.input_ids[rank::size]
+                            )
+                            self.assertEqual(
+                                torch.count_nonzero(local[count:]).item(), 0
+                            )
+                            self.assertEqual(torch.count_nonzero(ids[count:]).item(), 0)
+                            with simulated_collective(strategy, batch, embeddings):
+                                restored = cp_gather_after_forward(local, batch)
+                            torch.testing.assert_close(
+                                restored, embeddings, rtol=0, atol=0
+                            )
+                        torch.testing.assert_close(batch.input_ids, original_ids)
+                        self.assertFalse(hasattr(batch, "input_ids_global"))
+
+    def test_speculative_state_and_global_ids_restored_on_exception(self):
+        for size in (2, 4):
+            for rank in range(size):
+                for had_global in (False, True):
+                    with (
+                        self.subTest(size=size, rank=rank, had_global=had_global),
+                        cp_context(size, rank) as (_, batch),
+                    ):
+                        full = torch.arange(36, dtype=torch.float32).reshape(9, 4)
+                        batch.spec_info = NS(hidden_states=full)
+                        previous = object()
+                        if had_global:
+                            batch.input_ids_global = previous
+                        with self.assertRaisesRegex(RuntimeError, "injected"):
+                            with cp_shard_model_inputs(
+                                full, batch.positions, batch, batch.input_ids
+                            ):
+                                n = len(full[rank::size])
+                                torch.testing.assert_close(
+                                    batch.spec_info.hidden_states[:n], full[rank::size]
+                                )
+                                # Global MoE IDs are in rank-major order, with padding.
+                                physical = sum(
+                                    batch.attn_cp_metadata.per_rank_actual_token
+                                )
+                                padded = batch.input_ids.new_zeros(physical)
+                                padded[:9] = batch.input_ids
+                                expected = torch.cat(
+                                    [padded[r::size] for r in range(size)]
+                                )
+                                torch.testing.assert_close(
+                                    batch.input_ids_global, expected
+                                )
+                                raise RuntimeError("injected")
+                        self.assertIs(batch.spec_info.hidden_states, full)
+                        if had_global:
+                            self.assertIs(batch.input_ids_global, previous)
+                        else:
+                            self.assertFalse(hasattr(batch, "input_ids_global"))
+
+    def test_short_prompt_falls_back_from_cp(self):
+        with cp_context(4, 0, (1, 2), (0, 0)) as (_, batch):
+            self.assertFalse(is_cp_active(batch))
+
+    def test_runner_text_embedding_and_preembedded_paths(self):
+        for preembedded in (False, True):
+            for rank in range(4):
+                with (
+                    self.subTest(preembedded=preembedded, rank=rank),
+                    cp_context(4, rank) as (strategy, batch),
+                ):
+                    full = torch.arange(27, dtype=torch.float32).reshape(9, 3)
+                    embedding = Mock(return_value=full)
+                    model = NS(
+                        vision=None,
+                        _prepare_mm_embeddings=Mock(
+                            side_effect=AssertionError("Text must not invoke vision")
+                        ),
+                        get_input_embeddings=Mock(return_value=embedding),
+                        pp_group=NS(is_last_rank=True),
+                        lm_head=object(),
+                        capture_aux_hidden_states=False,
+                        logits_processor=Mock(return_value="ok"),
+                    )
+                    model.prepare_language_model_inputs = lambda ids, fb, emb: (
+                        DeepseekV4ForCausalLM.prepare_language_model_inputs(
+                            model, ids, fb, emb
+                        )
+                    )
+
+                    def body(ids, positions, fb, input_embeds):
+                        n = len(full[rank::4])
+                        torch.testing.assert_close(ids[:n], batch.input_ids[rank::4])
+                        torch.testing.assert_close(
+                            positions[:n], batch.positions[rank::4]
+                        )
+                        torch.testing.assert_close(input_embeds[:n], full[rank::4])
+                        return input_embeds
+
+                    model.model = body
+                    with (
+                        simulated_collective(strategy, batch, full),
+                        patch(RUNNER + ".torch.cuda.current_stream", return_value=None),
+                    ):
+                        result = EagerRunner._execute_extend_cp(
+                            NS(model_runner=NS(model=model)),
+                            batch,
+                            {"input_embeds": full} if preembedded else {},
+                        )
+                    self.assertEqual(result, "ok")
+                    if preembedded:
+                        model.get_input_embeddings.assert_not_called()
+                    else:
+                        embedding.assert_called_once_with(batch.input_ids)
+                    model._prepare_mm_embeddings.assert_not_called()
+                    args = model.logits_processor.call_args.args
+                    torch.testing.assert_close(args[0], batch.input_ids)
+                    torch.testing.assert_close(args[1], full)
+
+    def test_dspark_aux_tensor_and_list_gathered_without_pre_norm_override(self):
+        for as_list in (False, True):
+            with self.subTest(as_list=as_list), cp_context(4, 2) as (strategy, batch):
+                full = torch.arange(27, dtype=torch.float32).reshape(9, 3)
+                local = strategy.shard_hidden_states(full, batch)
+                aux = [local.clone(), local.clone()] if as_list else local.clone()
+                model = NS(
+                    get_input_embeddings=lambda: lambda ids: full,
+                    model=Mock(return_value=((local, local.clone()), aux)),
+                    capture_aux_hidden_states=True,
+                    pp_group=NS(is_last_rank=True),
+                    lm_head=object(),
+                    logits_processor=Mock(return_value="ok"),
+                )
+                with (
+                    simulated_collective(strategy, batch, full),
+                    patch(RUNNER + ".torch.cuda.current_stream", return_value=None),
+                ):
+                    EagerRunner._execute_extend_cp(
+                        NS(model_runner=NS(model=model)), batch, {}
+                    )
+                args, kwargs = model.logits_processor.call_args
+                torch.testing.assert_close(args[1], full)
+                for tensor in args[4] if as_list else [args[4]]:
+                    torch.testing.assert_close(tensor, full)
+                self.assertNotIn("hidden_states_before_norm", kwargs)
+
+    def test_target_hidden_states_before_norm_preserved_without_dspark_aux(self):
+        with cp_context(4, 1) as (strategy, batch):
+            full = torch.arange(27, dtype=torch.float32).reshape(9, 3)
+            local = strategy.shard_hidden_states(full, batch)
+            model = NS(
+                get_input_embeddings=lambda: lambda ids: full,
+                model=Mock(return_value=(local, local.clone())),
+                capture_aux_hidden_states=False,
+                pp_group=NS(is_last_rank=True),
+                lm_head=object(),
+                logits_processor=Mock(return_value="ok"),
+            )
+            with (
+                simulated_collective(strategy, batch, full),
+                patch(RUNNER + ".torch.cuda.current_stream", return_value=None),
+            ):
+                EagerRunner._execute_extend_cp(
+                    NS(model_runner=NS(model=model)), batch, {}
+                )
+            torch.testing.assert_close(
+                model.logits_processor.call_args.kwargs["hidden_states_before_norm"],
+                full,
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/registered/unit/server_args/test_deepseek_v41_text_only_pd.py b/test/registered/unit/server_args/test_deepseek_v41_text_only_pd.py
new file mode 100644
index 000000000000..e07a4a6c55ff
--- /dev/null
+++ b/test/registered/unit/server_args/test_deepseek_v41_text_only_pd.py
@@ -0,0 +1,73 @@
+"""V4.1 language-model-only PD configuration validation."""
+
+import unittest
+from types import SimpleNamespace as NS
+from unittest.mock import patch
+
+from sglang.srt.arg_groups import model_hook
+from sglang.srt.server_args import ServerArgs
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestDSV41TextOnlyPDPolicy(CustomTestCase):
+    def validate(
+        self,
+        model_type="deepseek_v41",
+        mode="prefill",
+        arch="DeepseekV4ForCausalLM",
+        **flags,
+    ):
+        cfg = NS(
+            language_model_only=True,
+            encoder_only=False,
+            language_only=False,
+            enable_prefix_mm_cache=False,
+            enable_broadcast_mm_inputs_process=False,
+            mm_enable_dp_encoder=False,
+            disaggregation_mode=mode,
+        )
+        for name, value in flags.items():
+            setattr(cfg, name, value)
+        model = NS(hf_config=NS(model_type=model_type, architectures=[arch]))
+        args = NS(
+            LANGUAGE_MODEL_ONLY_ARCHITECTURES=ServerArgs.LANGUAGE_MODEL_ONLY_ARCHITECTURES
+        )
+        with (
+            patch.object(model_hook, "resolving_view", return_value=cfg),
+            patch.object(model_hook, "model_config_of", return_value=model),
+        ):
+            model_hook.handle_language_model_only(args)
+
+    def test_v41_modes(self):
+        for mode in ("null", "prefill", "decode"):
+            with self.subTest(mode=mode):
+                self.validate(mode=mode)
+
+    def test_other_models_still_reject_pd(self):
+        with self.assertRaisesRegex(ValueError, "incompatible"):
+            self.validate(model_type="cosmos3", arch="Cosmos3ForConditionalGeneration")
+
+    def test_encoder_options_still_rejected(self):
+        for flag in (
+            "encoder_only",
+            "language_only",
+            "enable_prefix_mm_cache",
+            "enable_broadcast_mm_inputs_process",
+            "mm_enable_dp_encoder",
+        ):
+            with (
+                self.subTest(flag=flag),
+                self.assertRaisesRegex(ValueError, "cannot be combined"),
+            ):
+                self.validate(**{flag: True})
+
+    def test_unknown_arch_rejected(self):
+        with self.assertRaisesRegex(ValueError, "does not support"):
+            self.validate(arch="UnknownArchitecture")
+
+
+if __name__ == "__main__":
+    unittest.main()

From 24b7c7e2550884e2570b08109459ed6e76231672 Mon Sep 17 00:00:00 2001
From: abing <wangbingjia.wbj@ant-intl.com>
Date: Thu, 17 Sep 2026 19:43:41 -0700
Subject: [PATCH 03/27] update code

---
 python/sglang/srt/arg_groups/deepseek_v4_hook.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py
index 4b4a2af18eec..c38008526c40 100644
--- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py
+++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py
@@ -324,6 +324,10 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
         and prefill_graph.max_seq_len is None
         and not cp_breakable_prefill
     ):
+        # The non-CP captured low-ratio indexer scores a static context width.
+        # CP BCG runs these sources eagerly with live prefix metadata, so this
+        # default would only force long-prefix CP batches back to eager.
+        # Explicit max_seq_len values still constrain both paths.
         declare_resolution(
             server_args,
             "validate_deepseek_v41_features",

From d0d8ae1dfd1fea96e641af095616185ae37418d2 Mon Sep 17 00:00:00 2001
From: abing <wangbingjia.wbj@ant-intl.com>
Date: Sat, 19 Sep 2026 02:32:27 -0700
Subject: [PATCH 04/27] fix bug

---
 python/sglang/srt/models/deepseek_v4.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index d5bc001a3e14..df860be21d5d 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -4599,16 +4599,19 @@ def __init__(
             and not getattr(config, "language_model_only", False)
         ):
             if (
-                get_parallel().attn_cp_size != 1
-                or get_parallel().pp_group.world_size != 1
+                get_parallel().pp_group.world_size != 1
                 or not (
                     get_moe_a2a_backend().is_none()
-                    or (_is_cuda and get_moe_a2a_backend().is_megamoe())
+                    or (
+                        _is_cuda
+                        and get_moe_a2a_backend().is_megamoe()
+                        and get_parallel().attn_cp_size == 1
+                    )
                 )
             ):
                 raise ValueError(
-                    "V4.1 vision supports TP/EP/DP without CP or PP; "
-                    "MoE A2A requires CUDA MegaMoE"
+                    "V4.1 vision supports TP/EP/DP and prefill CP without PP; "
+                    "MoE A2A requires CUDA MegaMoE without CP"
                 )
 
             args = SimpleNamespace(**vars(config), dim=config.hidden_size)

From a7d835c475064e45a3dd25e6a968e5f4cf497ce0 Mon Sep 17 00:00:00 2001
From: abing <wangbingjia.wbj@ant-intl.com>
Date: Thu, 17 Sep 2026 04:04:38 -0700
Subject: [PATCH 05/27] run pass cp+megamoe+bcg

---
 .../layers/attention/deepseek_v4_backend.py   | 30 +++++++-
 python/sglang/srt/layers/cp/bcg.py            | 68 ++++++++++++++++---
 .../runner/prefill_cuda_graph_runner.py       |  7 +-
 python/sglang/srt/models/deepseek_v4.py       | 20 ++++--
 4 files changed, 110 insertions(+), 15 deletions(-)

diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
index e0c400c28071..6931a5c34701 100644
--- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py
+++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
@@ -697,7 +697,14 @@ def refresh_for_breakable_cuda_graph_replay_(self, other: DSV4AttnMetadata) -> N
             if src_val is None and dst_val is None:
                 continue
             assert dst_val is not None, f"{field_name=} {src_val=} {dst_val=}"
-            dst_val.copy_(src_val)
+            shape_mismatch = dst_val.shape != src_val.shape
+            assert not shape_mismatch or field_name in self._CP_GLOBAL_FIELDS, (
+                f"Only CP-global replay metadata may use a shorter live prefix, "
+                f"got {field_name=} {src_val.shape=} {dst_val.shape=}"
+            )
+            _copy_tensor_allowing_storage_alias(
+                dst_val, src_val, pad_value=0 if shape_mismatch else None
+            )
 
         # Safe to replace: captured kernels read only the per-replay objects, or
         # the field is produced in-graph before the attention graph break reads it.
@@ -947,6 +954,27 @@ def _prefill_graph_max_seq_len() -> Optional[int]:
     return get_exec().graph.cuda_graph_config.prefill.max_seq_len
 
 
+def _copy_tensor_allowing_storage_alias(
+    dst: torch.Tensor, src: torch.Tensor, *, pad_value: Optional[int] = None
+) -> None:
+    """Copy replay metadata while preserving capture-stable destination addresses."""
+    if dst is src:
+        return
+    if dst.untyped_storage().data_ptr() == src.untyped_storage().data_ptr():
+        src = src.clone()
+    if dst.shape == src.shape:
+        dst.copy_(src)
+        return
+    assert (
+        pad_value is not None
+        and dst.ndim == src.ndim
+        and dst.shape[0] >= src.shape[0]
+        and dst.shape[1:] == src.shape[1:]
+    ), f"Cannot copy replay metadata from {src.shape=} to {dst.shape=}"
+    dst.fill_(pad_value)
+    dst[: src.shape[0]].copy_(src)
+
+
 @dataclass
 class DSV4Metadata:
     core_attn_metadata: DSV4AttnMetadata
diff --git a/python/sglang/srt/layers/cp/bcg.py b/python/sglang/srt/layers/cp/bcg.py
index 4d1f1a5d77b0..3dedf48eab5e 100644
--- a/python/sglang/srt/layers/cp/bcg.py
+++ b/python/sglang/srt/layers/cp/bcg.py
@@ -23,18 +23,22 @@
 
 from sglang.srt.arg_groups.overrides import (
     attention_backends_of,
+    model_config_of,
     resolved_view,
     resolving_view,
 )
+from sglang.srt.configs.model_config import is_deepseek_v4
 from sglang.srt.layers.cp.base import get_cp_strategy
 from sglang.srt.layers.cp.interleave import InterleaveCPStrategy
 from sglang.srt.layers.cp.padding import get_cp_padding_align_size
 from sglang.srt.layers.cp.utils import (
     cp_gather_after_forward,
+    cp_shard_hidden_states,
     cp_split_before_forward,
     prepare_cp_forward,
 )
 from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy
+from sglang.srt.layers.logits_processor import LogitsMetadata
 from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
 
 if TYPE_CHECKING:
@@ -51,14 +55,18 @@ def supports_prefill_cp_bcg(server_args: ServerArgs) -> bool:
     cfg = resolving_view(server_args)
     resolved = resolved_view(server_args)
     prefill_attention_backend, _ = attention_backends_of(resolved_view(server_args))
+    supports_layout = (
+        cfg.cp_strategy == "zigzag" and prefill_attention_backend == "trtllm_mha"
+    ) or (
+        cfg.cp_strategy == "interleave"
+        and prefill_attention_backend == "dsv4"
+        and is_deepseek_v4(model_config_of(server_args).hf_config)
+    )
     return (
         cfg.enable_prefill_cp
         and cfg.pp_size == 1
         and resolved.attn_cp_size == cfg.tp_size
-        and (
-            (cfg.cp_strategy == "zigzag" and prefill_attention_backend == "trtllm_mha")
-            or (cfg.cp_strategy == "interleave" and prefill_attention_backend == "dsv4")
-        )
+        and supports_layout
     )
 
 
@@ -101,6 +109,8 @@ class PrefillCPBCGInput:
 
     input_embeds: torch.Tensor
     positions: torch.Tensor
+    input_ids: Optional[torch.Tensor] = None
+    num_token_non_padded: Optional[torch.Tensor] = None
     bucket_local_tokens: Dict[int, int] = field(default_factory=dict)
     live_local_tokens: int = 0
 
@@ -119,6 +129,8 @@ def create(cls, runner: PrefillCudaGraphRunner) -> PrefillCPBCGInput:
                     (runner.max_num_tokens,),
                     dtype=torch.int64,
                 ),
+                input_ids=torch.zeros((runner.max_num_tokens,), dtype=torch.int64),
+                num_token_non_padded=torch.zeros((), dtype=torch.int32),
             )
 
     def required_local_tokens(self, extend_seq_lens: Any) -> Optional[int]:
@@ -231,6 +243,7 @@ def prepare(
         raw_tokens = int(forward_batch.extend_num_tokens)
         global_input_ids = forward_batch.input_ids[:raw_tokens]
         global_positions = forward_batch.positions[:raw_tokens]
+        local_input_ids = cp_shard_hidden_states(global_input_ids, forward_batch)
         global_input_embeds = runner.model_runner.model.get_input_embeddings()(
             global_input_ids
         )
@@ -261,12 +274,31 @@ def prepare(
 
         input_embeds = self.input_embeds[:captured_local_tokens]
         positions = self.positions[:captured_local_tokens]
+        assert self.input_ids is not None
+        input_ids = self.input_ids[:captured_local_tokens]
         input_embeds.zero_()
         positions.zero_()
+        input_ids.zero_()
         input_embeds[:live_local_tokens].copy_(local_input_embeds)
         positions[:live_local_tokens].copy_(local_positions)
+        input_ids[:live_local_tokens].copy_(local_input_ids)
         forward_batch.input_embeds = input_embeds
-        forward_batch.positions = positions
+        forward_batch._cp_positions = positions
+        # Keep the global input_ids field intact: the runner uses its length to
+        # select the global capture bucket. The DSV4 body consumes this fixed,
+        # rank-local view for hash routing and MegaMoE.
+        forward_batch._cp_input_ids = input_ids
+        forward_batch.input_ids_global = input_ids
+        if forward_batch.num_token_non_padded is not None:
+            assert self.num_token_non_padded is not None
+            metadata = forward_batch.attn_cp_metadata
+            logical_tokens = (
+                metadata.per_rank_logical_token or metadata.per_rank_actual_token
+            )
+            strategy = get_cp_strategy()
+            assert strategy is not None
+            self.num_token_non_padded.fill_(logical_tokens[strategy.cp_rank])
+            forward_batch.num_token_non_padded = self.num_token_non_padded
         self.live_local_tokens = live_local_tokens
 
         if isinstance(get_cp_strategy(), InterleaveCPStrategy):
@@ -352,11 +384,31 @@ def gather_aux(aux):
             hidden_states, hidden_states_before_norm = hidden_states
             if aux_hidden_states is None:
                 logits_kwargs["hidden_states_before_norm"] = hidden_states_before_norm
-        return model.logits_processor(
-            forward_batch.input_ids,
+        input_ids = forward_batch.input_ids
+        logits_metadata = forward_batch
+        tail = None
+        language_model = getattr(model, "model", None)
+        if (
+            capture_aux_hidden_states
+            and getattr(language_model, "late_layer_start", None) is not None
+            and forward_batch.forward_mode.is_extend_without_speculative()
+        ):
+            tail_metadata = runner.model_runner.attn_backend.tail_forward_metadata
+            tail = tail_metadata.late_layer_tail
+            input_ids = tail.rows(input_ids)
+            logits_metadata = LogitsMetadata.from_forward_batch(forward_batch)
+            logits_metadata.extend_seq_lens = tail.extend_seq_lens
+            logits_metadata.extend_seq_lens_cpu = tail.extend_seq_lens_cpu
+            logits_metadata.extend_logprob_start_lens_cpu = tail.extend_seq_lens_cpu
+
+        output = model.logits_processor(
+            input_ids,
             hidden_states,
             model.lm_head,
-            forward_batch,
+            logits_metadata,
             aux_hidden_states,
             **logits_kwargs,
         )
+        if tail is not None:
+            output.hidden_states_token_indices = tail.token_indices
+        return output
diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
index 242d0b23e6a6..80873bc8f0ea 100644
--- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
@@ -755,6 +755,9 @@ def _capture_num_token_non_padded(self, num_tokens: int) -> Optional[torch.Tenso
 
     def _get_layer_model_positions(self, forward_batch: ForwardBatch) -> torch.Tensor:
         """Mirror outer multimodal wrappers when BCG captures layer_model directly."""
+        cp_positions = getattr(forward_batch, "_cp_positions", None)
+        if cp_positions is not None:
+            return cp_positions
         if forward_batch.mrope_positions is None:
             return forward_batch.positions
 
@@ -837,7 +840,9 @@ def _run_forward(self, forward_batch: ForwardBatch, num_tokens: int):
             if self._uses_eager_prefill_tail():
                 # BCG / Full: capture the transformer body only.
                 positions = self._get_layer_model_positions(forward_batch)
-                input_ids = forward_batch.input_ids
+                input_ids = getattr(
+                    forward_batch, "_cp_input_ids", forward_batch.input_ids
+                )
                 kwargs = _build_layer_model_forward_kwargs(
                     self.layer_model, forward_batch, pp_proxy_tensors
                 )
diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index df860be21d5d..ee1edd71afb9 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -2111,7 +2111,10 @@ def _forward_prepare(
             if (
                 forward_batch.forward_mode.is_extend()
                 and is_in_breakable_cuda_graph()
-                and not getattr(attn_backend, "low_ratio_prefill_graph", False)
+                and (
+                    dsa_use_prefill_cp(forward_batch)
+                    or not getattr(attn_backend, "low_ratio_prefill_graph", False)
+                )
             ):
                 bcg_deepseek_v4_low_ratio_sources(self, x, q_lora, positions)
             else:
@@ -4137,11 +4140,18 @@ def _forward_layers_hc_pre_from_prev(
         )
         if self.engram_hasher is not None:
             if cp_extend:
-                # n-gram hashing needs each token's predecessors: hash the whole prompt
+                # N-gram hashing needs each token's predecessors, so hash the
+                # whole prompt before selecting this CP rank's interleaved rows.
+                # The hasher builds request-to-token indices dynamically; keep
+                # that work at an eager break during breakable graph capture.
                 total = int(forward_batch.attn_cp_metadata.total_seq_lens)
-                hash_ids = self.engram_hasher(
-                    forward_batch.input_ids[:total], forward_batch
-                )
+                global_input_ids = forward_batch.input_ids[:total]
+                if is_in_breakable_cuda_graph():
+                    hash_ids = bcg_deepseek_v4_engram_hash_ids(
+                        self.engram_hasher, global_input_ids
+                    )
+                else:
+                    hash_ids = self.engram_hasher(global_input_ids, forward_batch)
                 parallel = get_parallel()
                 hash_ids = hash_ids[parallel.attn_cp_rank :: parallel.attn_cp_size]
                 pad_rows = hidden_states.shape[0] - hash_ids.shape[0]

From dd3a27963f1b146a9c7f622905516cecea5a4309 Mon Sep 17 00:00:00 2001
From: abing <wangbingjia.wbj@ant-intl.com>
Date: Sat, 19 Sep 2026 10:24:01 -0700
Subject: [PATCH 06/27] Reuse shared compressed KV dequantization in V4.1 CP
 prefill

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.
---
 .../layers/attention/deepseek_v4_backend.py   |  41 +++--
 .../attention/dsv4/sparse_prefill_utils.py    |  17 +-
 .../attention/test_dsv41_shared_prefill_kv.py | 151 ++++++++++++++++++
 3 files changed, 194 insertions(+), 15 deletions(-)
 create mode 100644 test/registered/kernels/ops/attention/test_dsv41_shared_prefill_kv.py

diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
index 6931a5c34701..2123cd7bfa94 100644
--- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py
+++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
@@ -1244,6 +1244,11 @@ def __init__(
         ] = None
         self.online_c128_mtp = OnlineC128MTPController(self)
         self.sparse_prefill_workspace = SparsePrefillWorkspace(self.device)
+        # CP V4.1 consumers share compressed KV across layers. Separate ratio
+        # workspaces keep those prefixes intact while each layer refreshes SWA.
+        self.shared_compressed_prefill_workspaces = {
+            ratio: SparsePrefillWorkspace(self.device) for ratio in (1, 2)
+        }
         spec_alg = model_runner.spec_algorithm
         self.needs_cpu_seq_lens = not spec_alg.is_dspark() and (
             not _is_cuda or self.online_c128_mtp.enabled()
@@ -3700,20 +3705,38 @@ def _forward_prefill_sparse(
                 compress_ratio, core_attn_metadata, extra_page_size
             )
             n_compressed = flat_token_ids.shape[0]
-            workspace = self.sparse_prefill_workspace.get(
-                n_compressed + cache.swa_token_ids.shape[0]
+            reuse_compressed = compress_ratio in (1, 2) and is_cp_active(forward_batch)
+            workspace_pool = (
+                self.shared_compressed_prefill_workspaces[compress_ratio]
+                if reuse_compressed
+                else self.sparse_prefill_workspace
             )
+            workspace = workspace_pool.get(n_compressed + cache.swa_token_ids.shape[0])
             compressed_slice = workspace[:n_compressed]
             swa_slice = workspace[n_compressed:]
 
         if compressed_slice is not None:
-            dequantize_k_cache_paged(
-                extra_k_cache,
-                flat_token_ids,
-                page_size=extra_page_size,
-                out=compressed_slice,
-                layout=token_to_kv_pool.get_extra_key_layout(layer_id),
-            )
+            source_key = None
+            if reuse_compressed:
+                source_layer = token_to_kv_pool.source_layer_of(layer_id)
+                source_key = (source_layer, workspace.data_ptr())
+                gather = cache.compressed[compress_ratio]
+            # A source layer may have just updated its cache in place. Consumer
+            # layers only reuse the compressed prefix; their top-k and SWA stay live.
+            if (
+                source_key is None
+                or layer_id == source_key[0]
+                or gather.dequantized_source != source_key
+            ):
+                dequantize_k_cache_paged(
+                    extra_k_cache,
+                    flat_token_ids,
+                    page_size=extra_page_size,
+                    out=compressed_slice,
+                    layout=token_to_kv_pool.get_extra_key_layout(layer_id),
+                )
+                if source_key is not None:
+                    gather.dequantized_source = source_key
         dequantize_k_cache_paged(
             token_to_kv_pool.get_swa_key_buffer_radix(layer_id),
             cache.swa_token_ids,
diff --git a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py
index 4266493ad02a..eac2cd24b28c 100644
--- a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py
+++ b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py
@@ -78,10 +78,11 @@ def use_dsv4_q8kv8_sparse_prefill(dsv4_prefill_backend: str = "auto") -> bool:
 class SparsePrefillWorkspace:
     """Backend-owned scratch storage for sparse prefill KV dequantization.
 
-    The workspace contents are fully overwritten before every attention call,
-    so token buckets and compression ratios can safely share one buffer. Sparse
-    prefill executes eagerly and serially on the supported paths, which makes it
-    safe to replace the scratch allocation when a larger extent is needed.
+    Callers normally overwrite the entire workspace. Shared compressed-KV
+    callers keep separate workspaces per ratio and track prefix validity in the
+    per-forward gather cache, including the allocation address. Sparse prefill
+    executes eagerly and serially on the supported paths, so the allocation can
+    be replaced when a larger extent is needed.
     """
 
     def __init__(self, device: torch.device):
@@ -294,14 +295,18 @@ class CompressedGather:
     # chunk-invariant per request; subsequent layers only overwrite that prefix.
     combined_indices: Optional[torch.Tensor] = None
     combined_lens: Optional[torch.Tensor] = None
+    # Valid only for this forward's gather layout. Each ratio has its own
+    # workspace; its compressed prefix survives consumer layers' SWA writes.
+    dequantized_source: Optional[tuple[int, int]] = None
 
 
 @dataclass
 class SparsePrefillChunkCache:
     """Cache prefill-chunk metadata shared across layers.
 
-    Fields depend on request/token mappings and compressed page tables, not
-    per-layer k_cache; per-layer top-k combinations are recomputed into reused
+    Gather layouts depend on request/token mappings and compressed page tables.
+    Shared-source dequantization keys live only for this forward; per-layer
+    top-k combinations are recomputed into reused
     buffers.
     """
 
diff --git a/test/registered/kernels/ops/attention/test_dsv41_shared_prefill_kv.py b/test/registered/kernels/ops/attention/test_dsv41_shared_prefill_kv.py
new file mode 100644
index 000000000000..d8b9de75dcf0
--- /dev/null
+++ b/test/registered/kernels/ops/attention/test_dsv41_shared_prefill_kv.py
@@ -0,0 +1,151 @@
+"""Shared compressed-KV workspace validity across layers and forwards."""
+import sys
+from types import SimpleNamespace as NS
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from sglang.srt.layers.attention import deepseek_v4_backend as backend
+from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
+    CompressedGather,
+    SparsePrefillWorkspace,
+    WORKSPACE_DIM,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+def test_shared_compressed_dequant_lifetime():
+    device = "cuda"
+    obj = object.__new__(backend.DeepseekV4AttnBackend)
+    obj.sparse_prefill_workspace = SparsePrefillWorkspace(device)
+    obj.shared_compressed_prefill_workspaces = {
+        ratio: SparsePrefillWorkspace(device) for ratio in (1, 2)
+    }
+    obj.softmax_scale = 0.1
+    obj.head_dim_v = WORKSPACE_DIM
+    sources = {0: 0, 1: 1, 2: 0, 3: 1, 4: 4, 5: 4, 6: 6, 7: 7}
+    ratios = {0: 1, 1: 2, 2: 1, 3: 2, 4: 1, 5: 1, 6: 0, 7: 4}
+    compressed = {
+        source: torch.full(
+            (128, 1, WORKSPACE_DIM),
+            float(source + 1),
+            device=device,
+            dtype=torch.bfloat16,
+        )
+        for source in (0, 1, 4, 7)
+    }
+    swa = torch.empty((16, 1, WORKSPACE_DIM), device=device, dtype=torch.bfloat16)
+    pool = NS(
+        source_layer_of=lambda layer: sources[layer],
+        get_extra_key_page_size=lambda layer: 1,
+        get_extra_key_buffer=lambda layer: compressed[sources[layer]],
+        get_extra_key_layout=lambda layer: None,
+        get_swa_key_buffer_radix=lambda layer: swa,
+        get_swa_key_layout=lambda: None,
+    )
+    calls = []
+    active = [True]
+
+    def dequant(src, indices, *, out, **kwargs):
+        calls.append("swa" if src is swa else sources_by_ptr[src.data_ptr()])
+        out.copy_(src.index_select(0, indices.long()))
+
+    sources_by_ptr = {v.data_ptr(): k for k, v in compressed.items()}
+
+    def make_cache(n):
+        gathers = {
+            ratio: CompressedGather(
+                flat_token_ids=torch.arange(n // ratio, device=device, dtype=torch.int32),
+                compressed_base=torch.zeros(1, device=device, dtype=torch.int32),
+                swa_base=torch.zeros(1, device=device, dtype=torch.int32),
+            )
+            for ratio in (1, 2, 4)
+        }
+        indices = torch.zeros((4, 128), device=device, dtype=torch.int32)
+        lengths = torch.full((4,), 3, device=device, dtype=torch.int32)
+        cache = NS(
+            compressed=gathers,
+            swa_token_ids=torch.arange(3, device=device),
+            swa_page_size=1,
+            c0_combined_indices=indices,
+            c0_combined_lens=lengths,
+        )
+        cache.layer_inputs = lambda ratio, core, page: (
+            gathers[ratio].flat_token_ids,
+            indices,
+            lengths,
+        )
+        return cache
+
+    def forward(layer):
+        return obj._forward_prefill_sparse(
+            torch.empty((4, 1, 1, WORKSPACE_DIM), device=device, dtype=torch.bfloat16),
+            layer,
+            ratios[layer],
+            NS(),
+            pool,
+            NS(),
+            torch.zeros(1, device=device),
+        )
+
+    with (
+        patch.object(backend, "is_cp_active", side_effect=lambda _: active[0]),
+        patch.object(backend, "dequantize_k_cache_paged", side_effect=dequant),
+        patch(
+            "sgl_kernel.flash_mla.flash_mla_sparse_fwd",
+            side_effect=lambda **kw: (kw["kv"].clone(), None, None),
+        ),
+    ):
+        # Same-size replay, then growth and shrink must all read freshly written KV.
+        for step, n in enumerate((8, 8, 40, 4)):
+            obj.forward_metadata = NS(sparse_prefill_cache=make_cache(n))
+            for source, tensor in compressed.items():
+                tensor.fill_(source + 1 + step * 10)
+            for layer, should_dequant in (
+                (0, True),
+                (1, True),
+                (2, False),
+                (3, False),
+                (4, True),
+                (5, False),
+                (2, True),
+                (3, False),
+                (6, False),
+                (7, True),
+                (7, True),
+            ):
+                swa.fill_(100 + layer + step)
+                before = len(calls)
+                active[0] = True
+                actual = forward(layer)
+                actual_calls = calls[before:]
+                assert actual_calls.count("swa") == 1
+                assert len(actual_calls) == 1 + int(should_dequant)
+                active[0] = False
+                expected = forward(layer)
+                torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+            # Workspace replacement invalidates even an unchanged source identity.
+            obj.shared_compressed_prefill_workspaces[1].get(256 + step * 256)
+            active[0] = True
+            before = len(calls)
+            actual = forward(2)
+            assert calls[before:] == [0, "swa"]
+            active[0] = False
+            torch.testing.assert_close(actual, forward(2), rtol=0, atol=0)
+
+            # Re-executing a producer may mutate the same cache address in place.
+            compressed[0].add_(1)
+            active[0] = True
+            before = len(calls)
+            actual = forward(0)
+            assert calls[before:] == [0, "swa"]
+            active[0] = False
+            torch.testing.assert_close(actual, forward(0), rtol=0, atol=0)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, *sys.argv[1:]]))

From a0a08932508830fce5595e9bc20299aec2e3abcd Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Fri, 18 Sep 2026 19:11:24 +0800
Subject: [PATCH 07/27] Support V4.1 DP attention in DP-only DSpark PD

---
 .../sglang/srt/arg_groups/deepseek_v4_hook.py |   6 +-
 .../sglang/srt/disaggregation/common/conn.py  |  15 ++-
 .../disaggregation/test_dsv41_dspark_pd.py    | 121 +++++++++++++++++-
 3 files changed, 135 insertions(+), 7 deletions(-)

diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py
index c38008526c40..b21fbe1bdf58 100644
--- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py
+++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py
@@ -289,11 +289,15 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
             read_ragged_verify_mode,
         )
 
-        # CP size may still be implicit here; validate_deepseek_v4_cp resolves it.
+        # This validation runs before validate_deepseek_v4_cp resolves
+        # attn_cp_size from TP. Keep the existing interleave prefill CP path
+        # alongside DP-only attention; CP+DP remains unsupported.
         supported_prefill_cp = (
             cfg.disaggregation_mode == "prefill"
             and cfg.enable_prefill_cp
             and cfg.cp_strategy == "interleave"
+            and cfg.dp_size == 1
+            and cfg.attn_cp_size in (1, cfg.tp_size)
         )
         if (
             read_ragged_verify_mode() is not RaggedVerifyMode.STATIC
diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py
index 23d9c7e99ffb..3c7277696fd8 100644
--- a/python/sglang/srt/disaggregation/common/conn.py
+++ b/python/sglang/srt/disaggregation/common/conn.py
@@ -1220,13 +1220,20 @@ def try_ensure_parallel_info(self, bootstrap_addr: str) -> bool:
                 and self.attn_cp_size == 1
                 and info.attn_tp_size * info.attn_cp_size == self.attn_tp_size
             )
-            if (
-                info.attn_tp_size != self.attn_tp_size
-                and not decode_tp_size_aligns_prefill_world_size
+            non_cp_layout = info.attn_cp_size == self.attn_cp_size == 1
+            non_cp_mla_layout = non_cp_layout and (
+                self.is_mla_backend or self.is_hybrid_mla_backend
+            )
+            if not (
+                decode_tp_size_aligns_prefill_world_size
+                or non_cp_mla_layout
+                or (non_cp_layout and info.attn_tp_size == self.attn_tp_size)
             ):
                 raise RuntimeError(
                     "DeepSeek-V4.1 DSpark PD requires the same TP size on both "
-                    "servers (including prefill CP ranks for MLA)"
+                    "servers (including prefill CP ranks for MLA), or CP=1 on "
+                    "both servers with an MLA KV layout for unequal attention TP. "
+                    "Decode CP is not supported."
                 )
 
         if self.dcp_size > 1:
diff --git a/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py b/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
index 69e58c50579b..43cd378fa9ee 100644
--- a/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
+++ b/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
@@ -62,6 +62,9 @@ def test_bootstrap_validates_before_caching(self):
                 manager.kv_cache_dtype_str = "fp8_e4m3"
                 manager.dsv41_spec_layout = local
                 manager.attn_tp_size = 4
+                manager.attn_cp_size = 1
+                manager.is_mla_backend = False
+                manager.is_hybrid_mla_backend = False
                 manager.dcp_size = 1
                 manager._resolve_rank_mapping = Mock()
                 response = Mock(status_code=200)
@@ -263,13 +266,50 @@ def test_plain_tp4_unchanged(self):
         self.assertEqual(info.target_tp_ranks, [3])
         self.assertEqual(info.target_cp_ranks, [0])
 
-    def test_unequal_model_tp_rejected(self):
-        for tp, cp in [(2, 1), (1, 2), (1, 8)]:
+    def test_unequal_model_tp_with_cp_rejected(self):
+        for tp, cp in [(1, 2), (1, 8), (4, 2)]:
             m = self.make()
             with self.assertRaisesRegex(RuntimeError, "same TP size"):
                 self.fetch(m, tp, cp)
             self.assertFalse(m.prefill_info_table)
 
+    def test_dp_only_unequal_attention_tp(self):
+        for mla in (True, False):
+            for prefill_tp, decode_tp in ((4, 1), (4, 2), (2, 4)):
+                for rank in range(decode_tp):
+                    with self.subTest(
+                        mla=mla, prefill_tp=prefill_tp, decode_tp=decode_tp, rank=rank
+                    ):
+                        m = self.make(rank, hybrid=not mla)
+                        m.is_mla_backend = mla
+                        m.attn_tp_size = decode_tp
+                        self.assertTrue(self.fetch(m, prefill_tp, 1))
+                        info = m.prefill_info_table["prefill:8761"]
+                        self.assertEqual(info.target_cp_ranks, [0])
+                        if decode_tp >= prefill_tp:
+                            self.assertEqual(
+                                info.target_tp_ranks, [rank // (decode_tp // prefill_tp)]
+                            )
+                        else:
+                            width = prefill_tp // decode_tp
+                            self.assertEqual(
+                                info.target_tp_ranks,
+                                list(range(rank * width, (rank + 1) * width)),
+                            )
+
+    def test_dp_only_non_mla_tp_mismatch_rejected(self):
+        m = self.make(hybrid=False)
+        with self.assertRaisesRegex(RuntimeError, "same TP size"):
+            self.fetch(m, 2, 1)
+        self.assertFalse(m.prefill_info_table)
+
+    def test_decode_cp_rejected(self):
+        m = self.make()
+        m.attn_cp_size = 2
+        with self.assertRaisesRegex(RuntimeError, "Decode CP"):
+            self.fetch(m, 4, 1)
+        self.assertFalse(m.prefill_info_table)
+
     def test_nonhybrid_cp_mismatch_rejected(self):
         m = self.make(hybrid=False)
         with self.assertRaisesRegex(RuntimeError, "same TP size"):
@@ -282,5 +322,82 @@ def test_layout_mismatch_still_rejected(self):
         self.assertFalse(m.prefill_info_table)
 
 
+class TestDSV41PDFeatureValidation(CustomTestCase):
+    def test_cp_and_dp_feature_matrix(self):
+        from sglang.srt.arg_groups import deepseek_v4_hook as hook
+        from sglang.srt.model_executor.cuda_graph_config import Backend
+        from sglang.srt.speculative.ragged_verify import RaggedVerifyMode
+
+        defaults = dict(
+            enable_encoder_swa_bounded_replay=False,
+            enable_decoder_swa_bounded_replay=False,
+            speculative_algorithm="DSPARK",
+            enable_hisparse=False,
+            dsv4_attn_backend="auto",
+            enable_two_batch_overlap=False,
+            pp_size=1,
+            disaggregation_mode="decode",
+            disaggregation_transfer_backend="mooncake",
+            enable_prefill_cp=False,
+            enable_dp_attention=False,
+            cp_strategy="interleave",
+            tp_size=4,
+            dp_size=1,
+            attn_cp_size=1,
+            dcp_size=1,
+            cuda_graph_config=SimpleNamespace(
+                prefill=SimpleNamespace(backend=Backend.DISABLED, max_seq_len=None)
+            ),
+        )
+        prefill_cp = dict(disaggregation_mode="prefill", enable_prefill_cp=True)
+        cases = [
+            ("plain_tp", {}, True),
+            ("decode_dp", dict(enable_dp_attention=True, dp_size=4), True),
+            (
+                "prefill_dp",
+                dict(disaggregation_mode="prefill", enable_dp_attention=True, dp_size=4),
+                True,
+            ),
+            ("prefill_cp_before_resolution", prefill_cp, True),
+            (
+                "prefill_cp_after_resolution",
+                dict(prefill_cp, attn_cp_size=4, enable_dp_attention=True),
+                True,
+            ),
+            ("cp_plus_dp", dict(prefill_cp, dp_size=2), False),
+            ("cp_zigzag", dict(prefill_cp, cp_strategy="zigzag"), False),
+            ("decode_prefill_cp_flag", dict(enable_prefill_cp=True), False),
+            ("decode_attention_cp", dict(attn_cp_size=2), False),
+            ("decode_dcp", dict(dcp_size=2), False),
+            ("unsupported_transfer", dict(disaggregation_transfer_backend="nixl"), False),
+        ]
+        for name, changes, supported in cases:
+            cfg = SimpleNamespace(**(defaults | changes))
+            with (
+                self.subTest(name=name),
+                patch.object(hook, "resolving_view", return_value=cfg),
+                patch.object(
+                    hook,
+                    "model_config_of",
+                    return_value=SimpleNamespace(
+                        hf_config=SimpleNamespace(model_type="deepseek_v41")
+                    ),
+                ),
+                patch(
+                    "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate.is_unified_kv_triton",
+                    return_value=False,
+                ),
+                patch(
+                    "sglang.srt.speculative.ragged_verify.read_ragged_verify_mode",
+                    return_value=RaggedVerifyMode.STATIC,
+                ),
+            ):
+                if supported:
+                    hook.validate_deepseek_v41_features(cfg)
+                else:
+                    with self.assertRaisesRegex(ValueError, "DeepSeek-V4.1 DSpark PD"):
+                        hook.validate_deepseek_v41_features(cfg)
+
+
 if __name__ == "__main__":
     unittest.main()

From 20a1ffd6d0f6a81148f822c213e981cf70ee1e20 Mon Sep 17 00:00:00 2001
From: abing <wangbingjia.wbj@ant-intl.com>
Date: Sat, 19 Sep 2026 20:08:01 -0700
Subject: [PATCH 08/27] Revert "Reuse shared compressed KV dequantization in
 V4.1 CP prefill"

This reverts commit 007f7c8ebb5a889f95418dc9a7e2d3deaf149f54.
---
 .../layers/attention/deepseek_v4_backend.py   |  41 ++---
 .../attention/dsv4/sparse_prefill_utils.py    |  17 +-
 .../attention/test_dsv41_shared_prefill_kv.py | 151 ------------------
 3 files changed, 15 insertions(+), 194 deletions(-)
 delete mode 100644 test/registered/kernels/ops/attention/test_dsv41_shared_prefill_kv.py

diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
index 2123cd7bfa94..6931a5c34701 100644
--- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py
+++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
@@ -1244,11 +1244,6 @@ def __init__(
         ] = None
         self.online_c128_mtp = OnlineC128MTPController(self)
         self.sparse_prefill_workspace = SparsePrefillWorkspace(self.device)
-        # CP V4.1 consumers share compressed KV across layers. Separate ratio
-        # workspaces keep those prefixes intact while each layer refreshes SWA.
-        self.shared_compressed_prefill_workspaces = {
-            ratio: SparsePrefillWorkspace(self.device) for ratio in (1, 2)
-        }
         spec_alg = model_runner.spec_algorithm
         self.needs_cpu_seq_lens = not spec_alg.is_dspark() and (
             not _is_cuda or self.online_c128_mtp.enabled()
@@ -3705,38 +3700,20 @@ def _forward_prefill_sparse(
                 compress_ratio, core_attn_metadata, extra_page_size
             )
             n_compressed = flat_token_ids.shape[0]
-            reuse_compressed = compress_ratio in (1, 2) and is_cp_active(forward_batch)
-            workspace_pool = (
-                self.shared_compressed_prefill_workspaces[compress_ratio]
-                if reuse_compressed
-                else self.sparse_prefill_workspace
+            workspace = self.sparse_prefill_workspace.get(
+                n_compressed + cache.swa_token_ids.shape[0]
             )
-            workspace = workspace_pool.get(n_compressed + cache.swa_token_ids.shape[0])
             compressed_slice = workspace[:n_compressed]
             swa_slice = workspace[n_compressed:]
 
         if compressed_slice is not None:
-            source_key = None
-            if reuse_compressed:
-                source_layer = token_to_kv_pool.source_layer_of(layer_id)
-                source_key = (source_layer, workspace.data_ptr())
-                gather = cache.compressed[compress_ratio]
-            # A source layer may have just updated its cache in place. Consumer
-            # layers only reuse the compressed prefix; their top-k and SWA stay live.
-            if (
-                source_key is None
-                or layer_id == source_key[0]
-                or gather.dequantized_source != source_key
-            ):
-                dequantize_k_cache_paged(
-                    extra_k_cache,
-                    flat_token_ids,
-                    page_size=extra_page_size,
-                    out=compressed_slice,
-                    layout=token_to_kv_pool.get_extra_key_layout(layer_id),
-                )
-                if source_key is not None:
-                    gather.dequantized_source = source_key
+            dequantize_k_cache_paged(
+                extra_k_cache,
+                flat_token_ids,
+                page_size=extra_page_size,
+                out=compressed_slice,
+                layout=token_to_kv_pool.get_extra_key_layout(layer_id),
+            )
         dequantize_k_cache_paged(
             token_to_kv_pool.get_swa_key_buffer_radix(layer_id),
             cache.swa_token_ids,
diff --git a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py
index eac2cd24b28c..4266493ad02a 100644
--- a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py
+++ b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py
@@ -78,11 +78,10 @@ def use_dsv4_q8kv8_sparse_prefill(dsv4_prefill_backend: str = "auto") -> bool:
 class SparsePrefillWorkspace:
     """Backend-owned scratch storage for sparse prefill KV dequantization.
 
-    Callers normally overwrite the entire workspace. Shared compressed-KV
-    callers keep separate workspaces per ratio and track prefix validity in the
-    per-forward gather cache, including the allocation address. Sparse prefill
-    executes eagerly and serially on the supported paths, so the allocation can
-    be replaced when a larger extent is needed.
+    The workspace contents are fully overwritten before every attention call,
+    so token buckets and compression ratios can safely share one buffer. Sparse
+    prefill executes eagerly and serially on the supported paths, which makes it
+    safe to replace the scratch allocation when a larger extent is needed.
     """
 
     def __init__(self, device: torch.device):
@@ -295,18 +294,14 @@ class CompressedGather:
     # chunk-invariant per request; subsequent layers only overwrite that prefix.
     combined_indices: Optional[torch.Tensor] = None
     combined_lens: Optional[torch.Tensor] = None
-    # Valid only for this forward's gather layout. Each ratio has its own
-    # workspace; its compressed prefix survives consumer layers' SWA writes.
-    dequantized_source: Optional[tuple[int, int]] = None
 
 
 @dataclass
 class SparsePrefillChunkCache:
     """Cache prefill-chunk metadata shared across layers.
 
-    Gather layouts depend on request/token mappings and compressed page tables.
-    Shared-source dequantization keys live only for this forward; per-layer
-    top-k combinations are recomputed into reused
+    Fields depend on request/token mappings and compressed page tables, not
+    per-layer k_cache; per-layer top-k combinations are recomputed into reused
     buffers.
     """
 
diff --git a/test/registered/kernels/ops/attention/test_dsv41_shared_prefill_kv.py b/test/registered/kernels/ops/attention/test_dsv41_shared_prefill_kv.py
deleted file mode 100644
index d8b9de75dcf0..000000000000
--- a/test/registered/kernels/ops/attention/test_dsv41_shared_prefill_kv.py
+++ /dev/null
@@ -1,151 +0,0 @@
-"""Shared compressed-KV workspace validity across layers and forwards."""
-import sys
-from types import SimpleNamespace as NS
-from unittest.mock import patch
-
-import pytest
-import torch
-
-from sglang.srt.layers.attention import deepseek_v4_backend as backend
-from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
-    CompressedGather,
-    SparsePrefillWorkspace,
-    WORKSPACE_DIM,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
-
-
-def test_shared_compressed_dequant_lifetime():
-    device = "cuda"
-    obj = object.__new__(backend.DeepseekV4AttnBackend)
-    obj.sparse_prefill_workspace = SparsePrefillWorkspace(device)
-    obj.shared_compressed_prefill_workspaces = {
-        ratio: SparsePrefillWorkspace(device) for ratio in (1, 2)
-    }
-    obj.softmax_scale = 0.1
-    obj.head_dim_v = WORKSPACE_DIM
-    sources = {0: 0, 1: 1, 2: 0, 3: 1, 4: 4, 5: 4, 6: 6, 7: 7}
-    ratios = {0: 1, 1: 2, 2: 1, 3: 2, 4: 1, 5: 1, 6: 0, 7: 4}
-    compressed = {
-        source: torch.full(
-            (128, 1, WORKSPACE_DIM),
-            float(source + 1),
-            device=device,
-            dtype=torch.bfloat16,
-        )
-        for source in (0, 1, 4, 7)
-    }
-    swa = torch.empty((16, 1, WORKSPACE_DIM), device=device, dtype=torch.bfloat16)
-    pool = NS(
-        source_layer_of=lambda layer: sources[layer],
-        get_extra_key_page_size=lambda layer: 1,
-        get_extra_key_buffer=lambda layer: compressed[sources[layer]],
-        get_extra_key_layout=lambda layer: None,
-        get_swa_key_buffer_radix=lambda layer: swa,
-        get_swa_key_layout=lambda: None,
-    )
-    calls = []
-    active = [True]
-
-    def dequant(src, indices, *, out, **kwargs):
-        calls.append("swa" if src is swa else sources_by_ptr[src.data_ptr()])
-        out.copy_(src.index_select(0, indices.long()))
-
-    sources_by_ptr = {v.data_ptr(): k for k, v in compressed.items()}
-
-    def make_cache(n):
-        gathers = {
-            ratio: CompressedGather(
-                flat_token_ids=torch.arange(n // ratio, device=device, dtype=torch.int32),
-                compressed_base=torch.zeros(1, device=device, dtype=torch.int32),
-                swa_base=torch.zeros(1, device=device, dtype=torch.int32),
-            )
-            for ratio in (1, 2, 4)
-        }
-        indices = torch.zeros((4, 128), device=device, dtype=torch.int32)
-        lengths = torch.full((4,), 3, device=device, dtype=torch.int32)
-        cache = NS(
-            compressed=gathers,
-            swa_token_ids=torch.arange(3, device=device),
-            swa_page_size=1,
-            c0_combined_indices=indices,
-            c0_combined_lens=lengths,
-        )
-        cache.layer_inputs = lambda ratio, core, page: (
-            gathers[ratio].flat_token_ids,
-            indices,
-            lengths,
-        )
-        return cache
-
-    def forward(layer):
-        return obj._forward_prefill_sparse(
-            torch.empty((4, 1, 1, WORKSPACE_DIM), device=device, dtype=torch.bfloat16),
-            layer,
-            ratios[layer],
-            NS(),
-            pool,
-            NS(),
-            torch.zeros(1, device=device),
-        )
-
-    with (
-        patch.object(backend, "is_cp_active", side_effect=lambda _: active[0]),
-        patch.object(backend, "dequantize_k_cache_paged", side_effect=dequant),
-        patch(
-            "sgl_kernel.flash_mla.flash_mla_sparse_fwd",
-            side_effect=lambda **kw: (kw["kv"].clone(), None, None),
-        ),
-    ):
-        # Same-size replay, then growth and shrink must all read freshly written KV.
-        for step, n in enumerate((8, 8, 40, 4)):
-            obj.forward_metadata = NS(sparse_prefill_cache=make_cache(n))
-            for source, tensor in compressed.items():
-                tensor.fill_(source + 1 + step * 10)
-            for layer, should_dequant in (
-                (0, True),
-                (1, True),
-                (2, False),
-                (3, False),
-                (4, True),
-                (5, False),
-                (2, True),
-                (3, False),
-                (6, False),
-                (7, True),
-                (7, True),
-            ):
-                swa.fill_(100 + layer + step)
-                before = len(calls)
-                active[0] = True
-                actual = forward(layer)
-                actual_calls = calls[before:]
-                assert actual_calls.count("swa") == 1
-                assert len(actual_calls) == 1 + int(should_dequant)
-                active[0] = False
-                expected = forward(layer)
-                torch.testing.assert_close(actual, expected, rtol=0, atol=0)
-
-            # Workspace replacement invalidates even an unchanged source identity.
-            obj.shared_compressed_prefill_workspaces[1].get(256 + step * 256)
-            active[0] = True
-            before = len(calls)
-            actual = forward(2)
-            assert calls[before:] == [0, "swa"]
-            active[0] = False
-            torch.testing.assert_close(actual, forward(2), rtol=0, atol=0)
-
-            # Re-executing a producer may mutate the same cache address in place.
-            compressed[0].add_(1)
-            active[0] = True
-            before = len(calls)
-            actual = forward(0)
-            assert calls[before:] == [0, "swa"]
-            active[0] = False
-            torch.testing.assert_close(actual, forward(0), rtol=0, atol=0)
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, *sys.argv[1:]]))

From 8ad77b3562ca9fb0571a568c05ee7b31ad0cd855 Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Fri, 18 Sep 2026 19:13:10 +0800
Subject: [PATCH 09/27] Support V4.1 decode vision MegaMoE

---
 python/sglang/srt/layers/moe/mega_moe.py | 1 +
 python/sglang/srt/models/deepseek_v4.py  | 1 +
 2 files changed, 2 insertions(+)

diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py
index 20d03d8cc7f7..8e318a7517f3 100644
--- a/python/sglang/srt/layers/moe/mega_moe.py
+++ b/python/sglang/srt/layers/moe/mega_moe.py
@@ -34,6 +34,7 @@
 from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.models.deepseek_common.utils import _device_sm
 from sglang.srt.multimodal.dsv41.vl_routing import vision_topk
+from sglang.srt.multimodal.dsv41.vl_routing import vision_topk
 from sglang.srt.utils import is_hip, is_sm100_supported
 
 if TYPE_CHECKING:
diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index ee1edd71afb9..1967ee2b46ff 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -187,6 +187,7 @@
 )
 from sglang.srt.runtime_context import (
     get_device,
+    get_disagg,
     get_exec,
     get_forward,
     get_parallel,

From 6864ce352a079b3addb7049ae21b400692e1587a Mon Sep 17 00:00:00 2001
From: abing <wangbingjia.wbj@ant-intl.com>
Date: Sat, 19 Sep 2026 20:38:10 -0700
Subject: [PATCH 10/27] fix pd-tp-conflict

---
 .../sglang/srt/disaggregation/common/conn.py  |  4 ++-
 .../disaggregation/test_dsv41_dspark_pd.py    | 31 ++++++++++++++++---
 2 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py
index 3c7277696fd8..2654d0458047 100644
--- a/python/sglang/srt/disaggregation/common/conn.py
+++ b/python/sglang/srt/disaggregation/common/conn.py
@@ -1227,7 +1227,9 @@ def try_ensure_parallel_info(self, bootstrap_addr: str) -> bool:
             if not (
                 decode_tp_size_aligns_prefill_world_size
                 or non_cp_mla_layout
-                or (non_cp_layout and info.attn_tp_size == self.attn_tp_size)
+                # Equal attention TP also supports prefill CP -> decode DP:
+                # the receiver gathers every CP shard into one DP replica.
+                or (self.attn_cp_size == 1 and info.attn_tp_size == self.attn_tp_size)
             ):
                 raise RuntimeError(
                     "DeepSeek-V4.1 DSpark PD requires the same TP size on both "
diff --git a/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py b/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
index 43cd378fa9ee..4c95615b3751 100644
--- a/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
+++ b/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
@@ -267,12 +267,28 @@ def test_plain_tp4_unchanged(self):
         self.assertEqual(info.target_cp_ranks, [0])
 
     def test_unequal_model_tp_with_cp_rejected(self):
-        for tp, cp in [(1, 2), (1, 8), (4, 2)]:
+        for tp, cp in [(1, 2), (1, 8)]:
             m = self.make()
             with self.assertRaisesRegex(RuntimeError, "same TP size"):
                 self.fetch(m, tp, cp)
             self.assertFalse(m.prefill_info_table)
 
+    def test_cp4_prefill_to_dp4_decode(self):
+        for mla in (True, False):
+            for rank in range(4):
+                with self.subTest(mla=mla, dp_rank=rank):
+                    m = self.make(rank, hybrid=not mla)
+                    m.is_mla_backend = mla
+                    m.attn_tp_size = 1
+                    m.attn_dp_size = 4
+                    m.attn_dp_rank = rank
+                    self.assertTrue(self.fetch(m, 1, 4))
+                    info = m.prefill_info_table["prefill:8761"]
+                    self.assertEqual(info.target_tp_ranks, [0])
+                    self.assertEqual(info.target_cp_ranks, [0, 1, 2, 3])
+                    self.assertEqual(info.required_dst_info_num, 1)
+                    self.assertEqual(info.required_prefill_response_num, 4)
+
     def test_dp_only_unequal_attention_tp(self):
         for mla in (True, False):
             for prefill_tp, decode_tp in ((4, 1), (4, 2), (2, 4)):
@@ -288,7 +304,8 @@ def test_dp_only_unequal_attention_tp(self):
                         self.assertEqual(info.target_cp_ranks, [0])
                         if decode_tp >= prefill_tp:
                             self.assertEqual(
-                                info.target_tp_ranks, [rank // (decode_tp // prefill_tp)]
+                                info.target_tp_ranks,
+                                [rank // (decode_tp // prefill_tp)],
                             )
                         else:
                             width = prefill_tp // decode_tp
@@ -355,7 +372,9 @@ def test_cp_and_dp_feature_matrix(self):
             ("decode_dp", dict(enable_dp_attention=True, dp_size=4), True),
             (
                 "prefill_dp",
-                dict(disaggregation_mode="prefill", enable_dp_attention=True, dp_size=4),
+                dict(
+                    disaggregation_mode="prefill", enable_dp_attention=True, dp_size=4
+                ),
                 True,
             ),
             ("prefill_cp_before_resolution", prefill_cp, True),
@@ -369,7 +388,11 @@ def test_cp_and_dp_feature_matrix(self):
             ("decode_prefill_cp_flag", dict(enable_prefill_cp=True), False),
             ("decode_attention_cp", dict(attn_cp_size=2), False),
             ("decode_dcp", dict(dcp_size=2), False),
-            ("unsupported_transfer", dict(disaggregation_transfer_backend="nixl"), False),
+            (
+                "unsupported_transfer",
+                dict(disaggregation_transfer_backend="nixl"),
+                False,
+            ),
         ]
         for name, changes, supported in cases:
             cfg = SimpleNamespace(**(defaults | changes))

From 696bb19080d1c9e25ac5e91216e3bcc77fc723a8 Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Sun, 20 Sep 2026 18:32:16 +0800
Subject: [PATCH 11/27] CP BCG capture indexer and compressor into graph

---
 .../sglang/srt/arg_groups/cuda_graph_hook.py  | 31 +++++++-
 .../sglang/srt/arg_groups/deepseek_v4_hook.py | 14 +---
 .../layers/attention/deepseek_v4_backend.py   | 78 ++++++++++++-------
 python/sglang/srt/layers/moe/mega_moe.py      |  1 -
 .../runner/prefill_cuda_graph_runner.py       | 28 ++++++-
 python/sglang/srt/models/deepseek_v4.py       |  5 +-
 6 files changed, 107 insertions(+), 50 deletions(-)

diff --git a/python/sglang/srt/arg_groups/cuda_graph_hook.py b/python/sglang/srt/arg_groups/cuda_graph_hook.py
index 2c325b6c3dcc..fc751d5c76c2 100644
--- a/python/sglang/srt/arg_groups/cuda_graph_hook.py
+++ b/python/sglang/srt/arg_groups/cuda_graph_hook.py
@@ -419,11 +419,38 @@ def finalize_cuda_graph_prefill_max_context(server_args: Any) -> None:
         "page_size must be resolved before prefill CUDA graph max context size"
     )
 
+    model_config = model_config_of(server_args)
     max_context_size = _resolve_max_context_size(
         requested_size=requested_size,
         page_size=page_size,
-        model_context_len=model_config_of(server_args).context_len,
+        model_context_len=model_config.context_len,
     )
+    is_v41 = (
+        getattr(getattr(model_config, "hf_config", None), "model_type", None)
+        == "deepseek_v41"
+    )
+    prefill_graph = cfg.cuda_graph_config.prefill
+    if (
+        is_v41
+        and prefill_graph.backend != Backend.DISABLED
+        and (cfg.enable_prefill_cp or cfg.attn_cp_size > 1)
+        and prefill_graph.max_seq_len is not None
+        and prefill_graph.max_seq_len > max_context_size
+    ):
+        raise ValueError(
+            "DeepSeek-V4.1 CP requires --cuda-graph-max-seq-len-prefill "
+            "to be at most --cuda-graph-prefill-max-context; CP uses "
+            "max_seq_len for graph admission."
+        )
+    changes = {"max_context_size": max_context_size}
+    if (
+        is_v41
+        and prefill_graph.backend != Backend.DISABLED
+        and prefill_graph.max_seq_len is None
+    ):
+        # CP metadata cannot use the fixed-context override; the paged indexer
+        # and graph admission both use max_seq_len instead.
+        changes["max_seq_len"] = max_context_size
     logger.info(
         "Prefill CUDA graph max context size: %d; graph keys remain token-only.",
         max_context_size,
@@ -434,7 +461,7 @@ def finalize_cuda_graph_prefill_max_context(server_args: Any) -> None:
         cuda_graph_config=with_phase(
             cfg.cuda_graph_config,
             Phase.PREFILL,
-            max_context_size=max_context_size,
+            **changes,
         ),
     )
 
diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py
index b21fbe1bdf58..6bb6ea789dd4 100644
--- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py
+++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py
@@ -317,21 +317,13 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None:
             )
 
     prefill_graph = cfg.cuda_graph_config.prefill
-    cp_breakable_prefill = (
-        cfg.enable_prefill_cp
-        and cfg.cp_strategy == "interleave"
-        and cfg.tp_size > 1
-        and prefill_graph.backend == Backend.BREAKABLE
-    )
     if (
         prefill_graph.backend != Backend.DISABLED
         and prefill_graph.max_seq_len is None
-        and not cp_breakable_prefill
+        and prefill_graph.max_context_size is None
     ):
-        # The non-CP captured low-ratio indexer scores a static context width.
-        # CP BCG runs these sources eagerly with live prefix metadata, so this
-        # default would only force long-prefix CP batches back to eager.
-        # Explicit max_seq_len values still constrain both paths.
+        # Captured ratio-1/2 indexers use a static context width for both CP
+        # and non-CP; an explicit max-context is copied to max_seq_len later.
         declare_resolution(
             server_args,
             "validate_deepseek_v41_features",
diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
index 6931a5c34701..09e44c6de5c4 100644
--- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py
+++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
@@ -1559,11 +1559,11 @@ def _low_ratio_prefill_indexer_metadata(
     @property
     def low_ratio_prefill_graph(self) -> bool:
         """Whether ratio-1/2 sources use captured projections and indexer metadata."""
-        return (
-            bool(self.low_ratios)
-            and has_dense_fp4_indexer()
-            and is_sm100_or_newer()
-            and get_parallel().attn_cp_size == 1
+        supported = (
+            bool(self.low_ratios) and has_dense_fp4_indexer() and is_sm100_or_newer()
+        )
+        return supported and (
+            get_parallel().attn_cp_size == 1 or _prefill_graph_max_seq_len() is not None
         )
 
     def can_run_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool:
@@ -2843,36 +2843,56 @@ def _forward_low_ratio_sources_cp(
     ) -> None:
         # Every rank writes the whole prompt's compressed state, scoring its own rows.
         cp_meta = forward_batch.attn_cp_metadata
-        total = int(cp_meta.total_seq_lens)
         tail = self.forward_metadata.late_layer_tail
-        if tail is not None:
-            q_lens_cpu = tail.local_lens_cpu
-            req_global, pos_global = tail.req_global, tail.pos_global
-        else:
-            q_lens_cpu = interleave_rows_per_request(
-                _as_int_list(forward_batch.extend_seq_lens_cpu),
-                get_parallel().attn_cp_rank,
-                get_parallel().attn_cp_size,
-            )
-            req_global = token_req_indices(forward_batch, num_tokens=total)
-            pos_global = forward_batch.positions[:total].to(torch.int64)
-        num_local = sum(q_lens_cpu)
         if run_compressor and layer.compressor is not None:
             x_global = cp_materialize_global_token_order(
                 x.contiguous(), forward_batch, torch.cuda.current_stream()
-            )[:total]
+            )
+            if self._low_ratio_in_prefill_graph():
+                assert tail is None, "bounded SWA replay cannot enter prefill BCG"
+                core = self.forward_metadata.core_metadata
+                bucket = core.raw_out_loc.shape[0]
+                req_global = self.forward_metadata.low_ratio_req_indices
+                assert req_global is not None and req_global.shape[0] >= bucket
+                assert x_global.shape[0] >= bucket
+                assert forward_batch.positions.shape[0] >= bucket
+                req_global = req_global[:bucket]
+                pos_global = forward_batch.positions[:bucket].to(torch.int64)
+                x_global = x_global[:bucket]
+            else:
+                total = int(cp_meta.total_seq_lens)
+                if tail is not None:
+                    req_global, pos_global = tail.req_global, tail.pos_global
+                else:
+                    req_global = token_req_indices(forward_batch, num_tokens=total)
+                    pos_global = forward_batch.positions[:total].to(torch.int64)
+                x_global = x_global[:total]
             self._low_ratio_compress_torch(layer, x_global, req_global, pos_global)
         if run_indexer and layer.indexer is not None:
-            # A rank's local rows are not the batch's, so no req_rows to pass.
-            self._low_ratio_index_topk(
-                layer,
-                x[:num_local],
-                q_lora[:num_local],
-                None,
-                positions[:num_local].to(torch.int64),
-                forward_batch,
-                rows_per_request=q_lens_cpu,
-            )
+            if self._low_ratio_in_prefill_graph():
+                q = layer.indexer.queries(q_lora, layer.freqs_cis[positions])
+                w = layer.indexer.head_weights(x)
+                self._low_ratio_index_topk_captured(layer, q, w)
+            else:
+                if tail is not None:
+                    q_lens_cpu = tail.local_lens_cpu
+                else:
+                    q_lens_cpu = interleave_rows_per_request(
+                        _as_int_list(forward_batch.extend_seq_lens_cpu),
+                        get_parallel().attn_cp_rank,
+                        get_parallel().attn_cp_size,
+                    )
+                num_local = sum(q_lens_cpu)
+                # A rank's local rows are not the batch's, so no req_rows to pass.
+                self._low_ratio_index_topk(
+                    layer,
+                    x[:num_local],
+                    q_lora[:num_local],
+                    None,
+                    positions[:num_local].to(torch.int64),
+                    forward_batch,
+                    rows_per_request=q_lens_cpu,
+                )
 
     def _low_ratio_compress(self, layer, x, req, pos, forward_batch) -> None:
         if forward_batch.forward_mode.is_decode():
diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py
index 8e318a7517f3..20d03d8cc7f7 100644
--- a/python/sglang/srt/layers/moe/mega_moe.py
+++ b/python/sglang/srt/layers/moe/mega_moe.py
@@ -34,7 +34,6 @@
 from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.models.deepseek_common.utils import _device_sm
 from sglang.srt.multimodal.dsv41.vl_routing import vision_topk
-from sglang.srt.multimodal.dsv41.vl_routing import vision_topk
 from sglang.srt.utils import is_hip, is_sm100_supported
 
 if TYPE_CHECKING:
diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
index 80873bc8f0ea..9d1481e474fc 100644
--- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
@@ -586,9 +586,31 @@ def __init__(self, model_runner: ModelRunner):
         ) and should_enable_cp_bcg_capture(server_args)
         if self.enable_cp_bcg_capture:
             if self.max_context_size is not None:
-                # TODO(SYChen123): Preserve max_seq_len_override through CP's padded
-                # metadata preparation before enabling the fixed context limit.
-                self._ignore_max_context_size("CP breakable prefill CUDA graph")
+                is_v41 = (
+                    getattr(
+                        getattr(model_runner.model_config, "hf_config", None),
+                        "model_type",
+                        None,
+                    )
+                    == "deepseek_v41"
+                )
+                if (
+                    is_v41
+                    and prefill_config.max_seq_len is not None
+                    and prefill_config.max_seq_len <= self.max_context_size
+                ):
+                    # CP's padded metadata cannot use max_seq_len_override;
+                    # max_seq_len still bounds admission and indexer width.
+                    logger.info(
+                        "CP prefill graph uses max_seq_len=%d for admission "
+                        "and indexer width (configured max-context=%d); "
+                        "skipping fixed metadata override.",
+                        prefill_config.max_seq_len,
+                        self.max_context_size,
+                    )
+                    self.max_context_size = None
+                else:
+                    self._ignore_max_context_size("CP breakable prefill CUDA graph")
             self.capture_num_tokens = filter_prefill_cp_bcg_capture_num_tokens(
                 self.capture_num_tokens, server_args
             )
diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index 1967ee2b46ff..56108033b618 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -2112,10 +2112,7 @@ def _forward_prepare(
             if (
                 forward_batch.forward_mode.is_extend()
                 and is_in_breakable_cuda_graph()
-                and (
-                    dsa_use_prefill_cp(forward_batch)
-                    or not getattr(attn_backend, "low_ratio_prefill_graph", False)
-                )
+                and not getattr(attn_backend, "low_ratio_prefill_graph", False)
             ):
                 bcg_deepseek_v4_low_ratio_sources(self, x, q_lora, positions)
             else:

From ab18a2b1603268a64309fae40d4da997fcfd58b1 Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Mon, 21 Sep 2026 14:38:06 +0800
Subject: [PATCH 12/27] Indexer BCG path uses topk_v2

---
 .../srt/layers/attention/dsv4/metadata.py     | 40 ++++++++++++++---
 .../unit/layers/test_dsv4_nonpaged_indexer.py | 43 +++++++++++++++++++
 2 files changed, 76 insertions(+), 7 deletions(-)

diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py
index bd850d38b342..2900df1050ad 100644
--- a/python/sglang/srt/layers/attention/dsv4/metadata.py
+++ b/python/sglang/srt/layers/attention/dsv4/metadata.py
@@ -277,15 +277,31 @@ def __post_init__(self):
         if self.use_topk_v2:
             from sglang.kernels.ops.attention.dsv4 import plan_topk_v2
 
-            self.topk_metadata = plan_topk_v2(self.compressed_seq_lens)
-            if self.rows_per_chunk is not None:
-                self.topk_metadata_chunks = [
-                    plan_topk_v2(self.compressed_seq_lens[rows])
-                    for rows in iter_row_chunks(
-                        num_rows=self.compressed_seq_lens.shape[0],
-                        rows_per_chunk=self.rows_per_chunk,
+            if self.row_chunk > 0:
+                plans = [
+                    plan_topk_v2(self.compressed_seq_lens[_s : _s + self.row_chunk])
+                    for _s in range(
+                        0, self.compressed_seq_lens.shape[0], self.row_chunk
                     )
                 ]
+                # Keep one capture-stable tensor so replay can refresh every
+                # chunk plan with a single copy_. The last plan may have fewer
+                # rows; callers slice its valid prefix before launching top-k.
+                self.topk_metadata = self.compressed_seq_lens.new_zeros(
+                    (len(plans), self.row_chunk + 1, 2)
+                )
+                for i, plan in enumerate(plans):
+                    self.topk_metadata[i, : plan.shape[0]].copy_(plan)
+            else:
+                self.topk_metadata = plan_topk_v2(self.compressed_seq_lens)
+                if self.rows_per_chunk is not None:
+                    self.topk_metadata_chunks = [
+                        plan_topk_v2(self.compressed_seq_lens[rows])
+                        for rows in iter_row_chunks(
+                            num_rows=self.compressed_seq_lens.shape[0],
+                            rows_per_chunk=self.rows_per_chunk,
+                        )
+                    ]
         else:
             self.topk_metadata = torch.empty((0,))
 
@@ -349,6 +365,16 @@ def row_chunks(self):
         )
         return chunks
 
+    def topk_plan_for_chunk(self, chunk_index: int, rows: slice) -> torch.Tensor:
+        assert self.use_topk_v2, "top-k v2 plan requested from v1 metadata"
+        if self.row_chunk > 0:
+            start = 0 if rows.start is None else rows.start
+            stop = self.compressed_seq_lens.shape[0] if rows.stop is None else rows.stop
+            return self.topk_metadata[chunk_index, : stop - start + 1]
+        if self.topk_metadata_chunks is not None:
+            return self.topk_metadata_chunks[chunk_index]
+        return self.topk_metadata
+
     def copy_(self, other: PagedIndexerMetadata):
         # A chunked schedule list has no in-place copy; rebind it instead.
         chunked = isinstance(self.deep_gemm_metadata, list) or isinstance(
diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
index 59e2d1b53297..924531cbffc4 100644
--- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
+++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
@@ -134,6 +134,49 @@ def test_physical_page_size_controls_metadata_and_replay(self):
             destination.compressed_seq_lens, source.compressed_seq_lens
         )
 
+    def test_chunked_topk_v2_plans_are_capture_stable_and_replayable(self):
+        def fake_plan(lengths):
+            value = int(lengths[0])
+            return torch.full((lengths.numel() + 1, 2), value, dtype=torch.int32)
+
+        def make_metadata(lengths):
+            with (
+                envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.override(True),
+                patch(
+                    "sglang.kernels.ops.attention.dsv4.plan_topk_v2",
+                    side_effect=fake_plan,
+                ),
+            ):
+                return PagedIndexerMetadata(
+                    page_size=256,
+                    compressed_page_size=64,
+                    page_table=torch.zeros((len(lengths), 2), dtype=torch.int32),
+                    compressed_seq_lens=torch.tensor(lengths, dtype=torch.int32),
+                    use_topk_v2=True,
+                    row_chunk=2,
+                )
+
+        destination = make_metadata([10, 20, 30, 40, 50])
+        self.assertEqual(destination.topk_metadata.shape, (3, 3, 2))
+        torch.testing.assert_close(
+            destination.topk_plan_for_chunk(0, slice(0, 2)),
+            torch.full((3, 2), 10, dtype=torch.int32),
+        )
+        torch.testing.assert_close(
+            destination.topk_plan_for_chunk(1, slice(2, 4)),
+            torch.full((3, 2), 30, dtype=torch.int32),
+        )
+        torch.testing.assert_close(
+            destination.topk_plan_for_chunk(2, slice(4, 5)),
+            torch.full((2, 2), 50, dtype=torch.int32),
+        )
+
+        source = make_metadata([11, 21, 31, 41, 51])
+        plan_ptr = destination.topk_metadata.data_ptr()
+        destination.copy_(source)
+        self.assertEqual(destination.topk_metadata.data_ptr(), plan_ptr)
+        torch.testing.assert_close(destination.topk_metadata, source.topk_metadata)
+
 
 class TestDSV4FlashInferTopK(CustomTestCase):
     def test_compact_page_transform_respects_fuse_topk(self):

From 830338bef634e51ec80304c717831f18bde42e85 Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Mon, 21 Sep 2026 18:01:54 +0800
Subject: [PATCH 13/27] Add SGLANG_DSV41_BCG_DENSE_INDEXER to support dense
 indexer on cp bcg path.

---
 python/sglang/srt/environ.py                  |   3 +
 .../layers/attention/deepseek_v4_backend.py   |  93 ++++++++++++
 .../unit/layers/test_dsv4_nonpaged_indexer.py | 132 ++++++++++++++++++
 3 files changed, 228 insertions(+)

diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py
index d4cc184b3083..af1abcd2fd26 100644
--- a/python/sglang/srt/environ.py
+++ b/python/sglang/srt/environ.py
@@ -1626,6 +1626,9 @@ class Envs:
     # Run the DeepSeek-V4.1 ratio-1/2 prefill indexer on the torch path instead
     # of the DeepGEMM dense fp4 logits kernel (test oracle / fallback).
     SGLANG_DSV41_TORCH_PREFILL_INDEXER = EnvBool(False)
+    # Pack up to eight CP prefill requests into contiguous FP4 index K, then
+    # use dense DeepGEMM and ragged top-k v2 inside the breakable CUDA graph.
+    SGLANG_DSV41_BCG_DENSE_INDEXER = EnvBool(False)
     SGLANG_FP8_PAGED_MQA_LOGITS_TORCH = EnvBool(False)
     SGLANG_OPT_FLASHMLA_SPARSE_PREFILL = EnvBool(True)
     # gfx950 DeepSeek-V4.1 prefill: attend with aiter's OPUS sparse kernel over a bf16
diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
index 09e44c6de5c4..cc4a5d794370 100644
--- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py
+++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
@@ -250,6 +250,49 @@ def _maybe_precompute_flashmla_sched_meta(
     flashmla_metadata.num_splits = num_splits
 
 
+def _dense_fp4_mqa_logits(
+    q_fp4: Tuple[torch.Tensor, torch.Tensor],
+    kv_fp4: Tuple[torch.Tensor, torch.Tensor],
+    weights: torch.Tensor,
+    ks: torch.Tensor,
+    ke: torch.Tensor,
+    max_seqlen_k: int,
+) -> torch.Tensor:
+    from deep_gemm import fp8_fp4_mqa_logits
+
+    return fp8_fp4_mqa_logits(
+        q_fp4,
+        kv_fp4,
+        weights,
+        ks.to(torch.int32),
+        ke.to(torch.int32),
+        False,
+        max_seqlen_k,
+    )
+
+
+def _prefill_graph_dense_k_layout(
+    req_to_token: torch.Tensor,
+    pool,
+    layer_id: int,
+    ratio: int,
+    width: int,
+    local_req_ids: torch.Tensor,
+    req_ids: torch.Tensor,
+    req_lens: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+    """Pack up to eight replayed requests into capture-stable contiguous FP4 K."""
+    req_ord = (local_req_ids[:, None] == req_ids[None, :]).to(torch.int32).argmax(1)
+    ks = req_ord.to(torch.int32) * width
+    logical_k = torch.arange(width, device=req_to_token.device)
+    valid_k = logical_k[None, :] < (req_lens[:, None] // ratio)
+    logical_pos = (logical_k[None, :] * ratio).expand_as(valid_k)
+    logical_pos = logical_pos.masked_fill(~valid_k, 0)
+    k_slots = req_to_token[req_ids[:, None], logical_pos].to(torch.int64) // ratio
+    k_slots = k_slots.masked_fill(~valid_k, 0).clamp_min_(0).flatten()
+    return ks, k_slots, pool.get_low_ratio_index_k_fp4(layer_id, k_slots)
+
+
 def _low_ratio_source_projections(layer, x, q_lora, positions, bufs):
     from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
         get_tc_piecewise_forward_context,
@@ -946,6 +989,7 @@ def _tail_rows(
 # Rows per logits chunk for the ratio-1/2 indexer inside the prefill CUDA graph;
 # its width is the graph's max_seq_len, and longer contexts replay eagerly.
 _PREFILL_GRAPH_INDEXER_ROW_CHUNK = 2048
+_PREFILL_GRAPH_DENSE_INDEXER_MAX_REQUESTS = 8
 
 
 def _prefill_graph_max_seq_len() -> Optional[int]:
@@ -990,6 +1034,10 @@ class DSV4Metadata:
     # Shared by all low-ratio source layers; graph replay refreshes them live.
     low_ratio_req_indices: Optional[torch.Tensor] = None
     low_ratio_pos_i64: Optional[torch.Tensor] = None
+    low_ratio_local_req_indices: Optional[torch.Tensor] = None
+    low_ratio_dense_req_indices: Optional[torch.Tensor] = None
+    low_ratio_dense_seq_lens: Optional[torch.Tensor] = None
+    prefill_graph_dense_indexer: bool = False
 
     # Per-step scratch for TP-padded query heads, zeroed by the first user.
     # Later layers overwrite real heads and preserve the zero padding.
@@ -1040,6 +1088,19 @@ def refresh_for_breakable_cuda_graph_replay_(self, static_metadata: DSV4Metadata
         maybe_copy_inplace(
             self.low_ratio_pos_i64, src=static_metadata.low_ratio_pos_i64
         )
+        if self.prefill_graph_dense_indexer:
+            maybe_copy_inplace(
+                self.low_ratio_local_req_indices,
+                src=static_metadata.low_ratio_local_req_indices,
+            )
+            maybe_copy_inplace(
+                self.low_ratio_dense_req_indices,
+                src=static_metadata.low_ratio_dense_req_indices,
+            )
+            maybe_copy_inplace(
+                self.low_ratio_dense_seq_lens,
+                src=static_metadata.low_ratio_dense_seq_lens,
+            )
         maybe_copy_inplace(
             self.c4_compress_metadata, src=static_metadata.c4_compress_metadata
         )
@@ -1526,6 +1587,31 @@ def create(compress_ratio: Literal[4, 128]):
             metadata.low_ratio_pos_i64 = core_attn_metadata.positions_casual.to(
                 torch.int64
             )
+            metadata.prefill_graph_dense_indexer = (
+                cp_active
+                and envs.SGLANG_DSV41_BCG_DENSE_INDEXER.get()
+                and req_pool_indices.numel()
+                <= _PREFILL_GRAPH_DENSE_INDEXER_MAX_REQUESTS
+            )
+            if metadata.prefill_graph_dense_indexer:
+                local_index = cp_metadata.local_index
+                if local_index is None:
+                    parallel = get_parallel()
+                    local_index = slice(
+                        parallel.attn_cp_rank, None, parallel.attn_cp_size
+                    )
+                metadata.low_ratio_local_req_indices = req_pool_indices_repeated[
+                    local_index
+                ].contiguous()
+                pad = (
+                    _PREFILL_GRAPH_DENSE_INDEXER_MAX_REQUESTS - req_pool_indices.numel()
+                )
+                metadata.low_ratio_dense_req_indices = torch.nn.functional.pad(
+                    req_pool_indices.to(torch.int64), (0, pad)
+                )
+                metadata.low_ratio_dense_seq_lens = torch.nn.functional.pad(
+                    seq_lens.to(torch.int64), (0, pad)
+                )
         return metadata
 
     def _low_ratio_prefill_indexer_metadata(
@@ -1567,6 +1653,13 @@ def low_ratio_prefill_graph(self) -> bool:
         )
 
     def can_run_prefill_cuda_graph(self, forward_batch: ForwardBatch) -> bool:
+        if (
+            envs.SGLANG_DSV41_BCG_DENSE_INDEXER.get()
+            and is_cp_active(forward_batch)
+            and forward_batch.req_pool_indices.numel()
+            > _PREFILL_GRAPH_DENSE_INDEXER_MAX_REQUESTS
+        ):
+            return False
         max_seq_len = _prefill_graph_max_seq_len()
         seq_lens_cpu = forward_batch.seq_lens_cpu
         if max_seq_len is None or seq_lens_cpu is None or seq_lens_cpu.numel() == 0:
diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
index 924531cbffc4..a9ecc446bad2 100644
--- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
+++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
@@ -43,6 +43,138 @@
 
 
 class TestDSV4PagedIndexerMetadata(CustomTestCase):
+    def test_dsv41_bcg_dense_indexer_keeps_paged_fallback_chunked(self):
+        from sglang.srt.layers.attention.deepseek_v4_backend import (
+            DeepseekV4AttnBackend,
+        )
+
+        backend = SimpleNamespace(
+            page_size=256,
+            token_to_kv_pool=SimpleNamespace(get_index_k_page_size=lambda ratio: 64),
+        )
+        core = SimpleNamespace(
+            page_table=torch.zeros((2049, 256), dtype=torch.int32),
+            seq_lens_casual=torch.full((2049,), 256, dtype=torch.int32),
+        )
+        module = "sglang.srt.layers.attention.deepseek_v4_backend"
+        with (
+            patch(f"{module}._prefill_graph_max_seq_len") as max_context,
+            patch(f"{module}.PagedIndexerMetadata") as metadata_ctor,
+        ):
+            with envs.SGLANG_DSV41_BCG_DENSE_INDEXER.override(False):
+                for context in (16384, 32768, 65536):
+                    max_context.return_value = context
+                    DeepseekV4AttnBackend._low_ratio_prefill_indexer_metadata(
+                        backend, core, 1
+                    )
+                    self.assertEqual(metadata_ctor.call_args.kwargs["row_chunk"], 2048)
+
+            with envs.SGLANG_DSV41_BCG_DENSE_INDEXER.override(True):
+                max_context.return_value = 32768
+                DeepseekV4AttnBackend._low_ratio_prefill_indexer_metadata(
+                    backend, core, 1
+                )
+                self.assertEqual(metadata_ctor.call_args.kwargs["row_chunk"], 2048)
+                self.assertTrue(metadata_ctor.call_args.kwargs["use_topk_v2"])
+
+    def test_dsv41_bcg_dense_k_layout_uses_live_request_order(self):
+        from sglang.srt.layers.attention.deepseek_v4_backend import (
+            _prefill_graph_dense_k_layout,
+        )
+
+        req_to_token = torch.arange(8)[:, None] * 1000 + torch.arange(10)[None, :] * 2
+        pool = SimpleNamespace(
+            get_low_ratio_index_k_fp4=lambda layer, slots: (
+                slots[:, None].expand(-1, 64).contiguous(),
+                slots.clone(),
+            )
+        )
+        ks, slots, packed = _prefill_graph_dense_k_layout(
+            req_to_token,
+            pool,
+            1,
+            2,
+            5,
+            torch.tensor([7, 3, 7, 3]),
+            torch.tensor([7, 3, 0, 0, 0, 0, 0, 0]),
+            torch.tensor([8, 6, 0, 0, 0, 0, 0, 0]),
+        )
+        torch.testing.assert_close(ks, torch.tensor([0, 5, 0, 5], dtype=torch.int32))
+        torch.testing.assert_close(
+            slots[:10],
+            torch.tensor([3500, 3502, 3504, 3506, 0, 1500, 1502, 1504, 0, 0]),
+        )
+        self.assertEqual(slots.shape, (40,))
+        torch.testing.assert_close(packed[0][:, 0], slots)
+        torch.testing.assert_close(packed[1], slots)
+
+    def test_dsv41_dense_logits_offsets_are_int32(self):
+        from sglang.srt.layers.attention.deepseek_v4_backend import (
+            _dense_fp4_mqa_logits,
+        )
+
+        gemm = MagicMock(return_value=torch.empty((2, 16)))
+        with patch.dict(
+            sys.modules,
+            {"deep_gemm": SimpleNamespace(fp8_fp4_mqa_logits=gemm)},
+        ):
+            _dense_fp4_mqa_logits(
+                (None, None),
+                (None, None),
+                torch.empty((2, 1)),
+                torch.tensor([0, 16], dtype=torch.int64),
+                torch.tensor([8, 20], dtype=torch.int64),
+                16,
+            )
+        self.assertEqual(gemm.call_args.args[3].dtype, torch.int32)
+        self.assertEqual(gemm.call_args.args[4].dtype, torch.int32)
+
+    def test_dsv41_bcg_dense_request_metadata_replays_in_place(self):
+        from sglang.srt.layers.attention.deepseek_v4_backend import DSV4Metadata
+
+        core = SimpleNamespace(
+            refresh_for_breakable_cuda_graph_replay_=lambda other: None
+        )
+        captured = DSV4Metadata(core, None)
+        captured.prefill_graph_dense_indexer = True
+        captured.low_ratio_local_req_indices = torch.zeros(4, dtype=torch.int64)
+        captured.low_ratio_dense_req_indices = torch.zeros(8, dtype=torch.int64)
+        captured.low_ratio_dense_seq_lens = torch.zeros(8, dtype=torch.int64)
+        live = DSV4Metadata(core, None)
+        live.prefill_graph_dense_indexer = True
+        live.low_ratio_local_req_indices = torch.tensor([7, 3, 7, 3])
+        live.low_ratio_dense_req_indices = torch.tensor([7, 3, 0, 0, 0, 0, 0, 0])
+        live.low_ratio_dense_seq_lens = torch.tensor([8, 6, 0, 0, 0, 0, 0, 0])
+
+        addresses = (
+            captured.low_ratio_local_req_indices.data_ptr(),
+            captured.low_ratio_dense_req_indices.data_ptr(),
+            captured.low_ratio_dense_seq_lens.data_ptr(),
+        )
+        captured.refresh_for_breakable_cuda_graph_replay_(live)
+        self.assertEqual(
+            addresses,
+            (
+                captured.low_ratio_local_req_indices.data_ptr(),
+                captured.low_ratio_dense_req_indices.data_ptr(),
+                captured.low_ratio_dense_seq_lens.data_ptr(),
+            ),
+        )
+        torch.testing.assert_close(
+            captured.low_ratio_local_req_indices, live.low_ratio_local_req_indices
+        )
+        torch.testing.assert_close(
+            captured.low_ratio_dense_req_indices, live.low_ratio_dense_req_indices
+        )
+        torch.testing.assert_close(
+            captured.low_ratio_dense_seq_lens, live.low_ratio_dense_seq_lens
+        )
+
+        # A large capture may use paged logits even if a smaller replay batch
+        # would be eligible for dense logits.
+        paged_capture = DSV4Metadata(core, None)
+        paged_capture.refresh_for_breakable_cuda_graph_replay_(live)
+
     def test_sm120_fp4_forces_deep_gemm_metadata(self):
         expected = torch.tensor([[0, 0], [1, 0]], dtype=torch.int32)
         deep_gemm = SimpleNamespace(

From 5c98dcb109076f4590f81572aca28d769c0fd288 Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Wed, 23 Sep 2026 17:01:44 +0800
Subject: [PATCH 14/27] Optimize CP BCG dense indexer path

---
 .../deepseek_v4/prefill_candidate_topk.cuh    | 137 ++++++++++++++++++
 .../sgl_kernel/deepseek_v4/topk_impl.cuh      |  20 ++-
 .../ops/attention/dsv4/prefill_candidates.py  |  42 ++++++
 .../attention/dsv4/sparse_prefill_utils.py    |  18 ++-
 4 files changed, 205 insertions(+), 12 deletions(-)
 create mode 100644 python/sglang/kernels/jit/csrc/deepseek_v4/prefill_candidate_topk.cuh
 create mode 100644 python/sglang/kernels/ops/attention/dsv4/prefill_candidates.py

diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/prefill_candidate_topk.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/prefill_candidate_topk.cuh
new file mode 100644
index 000000000000..bed4d9e82829
--- /dev/null
+++ b/python/sglang/kernels/jit/csrc/deepseek_v4/prefill_candidate_topk.cuh
@@ -0,0 +1,137 @@
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <sgl_kernel/deepseek_v4/topk_impl.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace sglang {
+namespace prefill_candidate {
+namespace impl = device::topk;
+
+struct Params {
+  const float* scores;
+  const int32_t* lengths;
+  const bool* candidates;
+  const int32_t* offsets;
+  int32_t* output;
+  int64_t score_stride;
+  int64_t candidate_stride;
+  uint32_t block_size;
+  uint32_t topk;
+};
+
+struct ScoreFilter {
+  const bool* blocks;
+  uint32_t block_size;
+
+  SGL_DEVICE float operator()(float value, uint32_t index) const {
+    return blocks[index / block_size] ? value : -impl::infinity_value();
+  }
+};
+
+template <bool kPDL>
+__global__ __launch_bounds__(impl::TopKConfig::kBlockSize, impl::TopKConfig::kOccupancy) void select(
+    const __grid_constant__ Params params) {
+  device::enable_smem_spilling();
+  const auto row = blockIdx.x;
+  const auto length = static_cast<uint32_t>(params.lengths[row]);
+  const auto offset = params.offsets[row];
+  const auto scores = params.scores + row * params.score_stride;
+  const auto candidates = params.candidates + row * params.candidate_stride;
+  const auto output = params.output + row * static_cast<int64_t>(params.topk);
+  const ScoreFilter filter{candidates, params.block_size};
+
+  if (length <= params.topk) {
+    device::PDLWaitPrimary<kPDL>();
+    for (uint32_t i = threadIdx.x; i < params.topk; i += blockDim.x) {
+      const bool valid = i < length && candidates[i / params.block_size] && scores[i] > -impl::infinity_value();
+      output[i] = valid ? static_cast<int32_t>(i) + offset : -1;
+    }
+    return;
+  }
+
+  using Register2 = impl::TopKRegister<2>;
+  using Register4 = impl::TopKRegister<4>;
+  using Streaming = impl::TopKStreaming;
+  const impl::TopKProblem problem{
+      .in = scores,
+      .out = output,
+      .topk = params.topk,
+      .seq_len = length,
+      .bias = impl::broadcast(offset),
+  };
+  __shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
+  if (length <= Register2::kMaxSeqLen) {
+    Register2::forward<kPDL>(problem, &smem, filter);
+  } else if (length <= Register4::kMaxSeqLen) {
+    Register4::forward<kPDL>(problem, &smem, filter);
+  } else {
+    Streaming::forward<kPDL>(problem, &smem, filter);
+  }
+  // Selection can underfill with -inf positions. Match mask_topk_scores in the
+  // same CTA, after every thread has finished publishing its selected indices.
+  __syncthreads();
+  for (uint32_t i = threadIdx.x; i < params.topk; i += blockDim.x) {
+    const int64_t index = static_cast<int64_t>(output[i]) - offset;
+    const bool valid = index >= 0 && index < length && candidates[index / params.block_size] &&
+                       scores[index] > -impl::infinity_value();
+    if (!valid) output[i] = -1;
+  }
+}
+
+template <bool kPDL>
+struct Kernel {
+  static void
+  run(const tvm::ffi::TensorView scores,
+      const tvm::ffi::TensorView lengths,
+      const tvm::ffi::TensorView candidates,
+      const tvm::ffi::TensorView offsets,
+      const tvm::ffi::TensorView output,
+      int64_t block_size) {
+    using namespace host;
+    auto B = SymbolicSize{"rows"};
+    auto W = SymbolicSize{"width"};
+    auto S = SymbolicSize{"score_stride"};
+    auto C = SymbolicSize{"candidate_width"};
+    auto CS = SymbolicSize{"candidate_stride"};
+    auto K = SymbolicSize{"topk"};
+    auto dev = SymbolicDevice{};
+    dev.set_options<kDLGPU>();
+    TensorMatcher({B, W}).with_strides({S, 1}).with_dtype<float>().with_device(dev).verify(scores);
+    TensorMatcher({B}).with_dtype<int32_t>().with_device(dev).verify(lengths);
+    TensorMatcher({B, C}).with_strides({CS, 1}).with_device(dev).verify(candidates);
+    RuntimeCheck(
+        candidates.dtype().code == kDLBool && candidates.dtype().bits == 8 && candidates.dtype().lanes == 1,
+        "candidate mask must have dtype bool");
+    TensorMatcher({B}).with_dtype<int32_t>().with_device(dev).verify(offsets);
+    TensorMatcher({B, K}).with_dtype<int32_t>().with_device(dev).verify(output);
+    RuntimeCheck(block_size > 0, "candidate block size must be positive");
+    RuntimeCheck(C.unwrap() >= (W.unwrap() + block_size - 1) / block_size, "candidate mask is too narrow");
+    RuntimeCheck(S.unwrap() % 4 == 0, "score rows must be 16-byte aligned");
+    RuntimeCheck(K.unwrap() > 0 && K.unwrap() <= impl::TopKConfig::kMaxTopK, "topk must be in (0, 2048]");
+    if (B.unwrap() == 0) return;
+    const Params params{
+        .scores = static_cast<const float*>(scores.data_ptr()),
+        .lengths = static_cast<const int32_t*>(lengths.data_ptr()),
+        .candidates = static_cast<const bool*>(candidates.data_ptr()),
+        .offsets = static_cast<const int32_t*>(offsets.data_ptr()),
+        .output = static_cast<int32_t*>(output.data_ptr()),
+        .score_stride = S.unwrap(),
+        .candidate_stride = CS.unwrap(),
+        .block_size = static_cast<uint32_t>(block_size),
+        .topk = static_cast<uint32_t>(K.unwrap()),
+    };
+    LaunchKernel(static_cast<uint32_t>(B.unwrap()), impl::TopKConfig::kBlockSize, dev.unwrap())
+        .config({.use_pdl = kPDL})
+        .launch(select<kPDL>, params);
+  }
+};
+}  // namespace prefill_candidate
+}  // namespace sglang
diff --git a/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh
index b0b208854097..b83764b5b3c2 100644
--- a/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh
+++ b/python/sglang/kernels/jit/include/sgl_kernel/deepseek_v4/topk_impl.cuh
@@ -533,14 +533,20 @@ struct TopKRadixBase : TopKConfig {
 // wastes work on shorter sequences.
 // ---------------------------------------------------------------------------
 
+struct IdentityScoreFilter {
+  SGL_DEVICE float operator()(float value, uint32_t) const {
+    return value;
+  }
+};
+
 template <uint32_t kLocalVecs_>
 struct TopKRegister : TopKRadixBase<12> {
   static constexpr uint32_t kLocalVecs = kLocalVecs_;
   static constexpr uint32_t kMaxSeqLen = kBlockSize * kVecSize * kLocalVecs;
   using Smem = typename TopKRadixBase<12>::Smem;
 
-  template <bool kUsePDL>
-  SGL_DEVICE static void forward(const TopKProblem& problem, void* _smem) {
+  template <bool kUsePDL, typename ScoreFilter = IdentityScoreFilter>
+  SGL_DEVICE static void forward(const TopKProblem& problem, void* _smem, ScoreFilter filter = {}) {
     const auto tx = threadIdx.x;
     const auto smem = static_cast<Smem*>(_smem);
 
@@ -575,6 +581,8 @@ struct TopKRegister : TopKRadixBase<12> {
       }
 #pragma unroll
       for (uint32_t j = 0; j < kVecSize; ++j) {
+        const auto idx = vi * kVecSize + j;
+        if (idx < problem.seq_len) local_vecs[i][j] = filter(local_vecs[i][j], idx);
         atomicAdd(&smem->histogram[extract_coarse_bin<kHistBits>(local_vecs[i][j])], 1);
       }
     }
@@ -640,8 +648,8 @@ struct TopKStreaming : TopKRadixBase<12> {
  public:
   static constexpr uint32_t kMaxSeqLen = std::numeric_limits<uint32_t>::max();
 
-  template <bool kUsePDL>
-  SGL_DEVICE static void forward(TopKProblem problem, void* _smem) {
+  template <bool kUsePDL, typename ScoreFilter = IdentityScoreFilter>
+  SGL_DEVICE static void forward(TopKProblem problem, void* _smem, ScoreFilter filter = {}) {
     const auto tx = threadIdx.x;
     const auto smem = static_cast<Smem*>(_smem);
 
@@ -654,7 +662,8 @@ struct TopKStreaming : TopKRadixBase<12> {
     PDLWaitPrimary<kUsePDL>();
 
     // Phase 1: Load and build histogram
-    for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t) {
+    for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t idx) {
+      val = filter(val, idx);
       const auto bin = extract_coarse_bin<kHistBits>(val);
       atomicAdd(&smem->histogram[bin], 1);
     });
@@ -682,6 +691,7 @@ struct TopKStreaming : TopKRadixBase<12> {
     const auto v_lo = smem->v_lo;
     const auto topk = problem.topk;
     for_each_input(problem.in, problem.seq_len, [&](float val, uint32_t idx) {
+      val = filter(val, idx);
       if (val >= v_hi) {
         const auto pos = atomicAdd(&smem->count_gt, 1);
         if (pos < topk) [[likely]] {
diff --git a/python/sglang/kernels/ops/attention/dsv4/prefill_candidates.py b/python/sglang/kernels/ops/attention/dsv4/prefill_candidates.py
new file mode 100644
index 000000000000..6c573aa65ebb
--- /dev/null
+++ b/python/sglang/kernels/ops/attention/dsv4/prefill_candidates.py
@@ -0,0 +1,42 @@
+"""Fused candidate-aware top-k for the dense FP4 prefill indexer."""
+
+import torch
+
+from sglang.kernels.jit.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+from .utils import make_name
+
+
+@cache_once
+def _candidate_topk_module():
+    args = make_cpp_args(is_arch_support_pdl())
+    return load_jit(
+        make_name("prefill_candidate_topk"),
+        *args,
+        cuda_files=["deepseek_v4/prefill_candidate_topk.cuh"],
+        cuda_wrappers=[("run", f"prefill_candidate::Kernel<{args}>::run")],
+    )
+
+
+def topk_prefill_candidates(
+    scores: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_mask: torch.Tensor,
+    block_size: int,
+    out_offsets: torch.Tensor,
+    out_indices: torch.Tensor,
+) -> None:
+    """Select within each query's candidate blocks, without modifying scores.
+
+    Columns are request-local, seq_lens are nonnegative and bounded by the score
+    width. Output indices include out_offsets; invalid slots are -1. Row strides
+    must be multiples of four floats for the vectorized top-k reads.
+    """
+    _candidate_topk_module().run(
+        scores, seq_lens, block_mask, out_offsets, out_indices, block_size
+    )
diff --git a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py
index 4266493ad02a..528626b5644a 100644
--- a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py
+++ b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py
@@ -78,10 +78,11 @@ def use_dsv4_q8kv8_sparse_prefill(dsv4_prefill_backend: str = "auto") -> bool:
 class SparsePrefillWorkspace:
     """Backend-owned scratch storage for sparse prefill KV dequantization.
 
-    The workspace contents are fully overwritten before every attention call,
-    so token buckets and compression ratios can safely share one buffer. Sparse
-    prefill executes eagerly and serially on the supported paths, which makes it
-    safe to replace the scratch allocation when a larger extent is needed.
+    Callers normally overwrite the entire workspace. Shared compressed-KV
+    callers keep separate workspaces per ratio and track prefix validity in the
+    per-forward gather cache, including the allocation address. Sparse prefill
+    executes eagerly and serially on the supported paths, so the allocation can
+    be replaced when a larger extent is needed.
     """
 
     def __init__(self, device: torch.device):
@@ -294,15 +295,18 @@ class CompressedGather:
     # chunk-invariant per request; subsequent layers only overwrite that prefix.
     combined_indices: Optional[torch.Tensor] = None
     combined_lens: Optional[torch.Tensor] = None
+    # Valid only for this forward's gather layout. Each ratio has its own
+    # workspace; its compressed prefix survives consumer layers' SWA writes.
+    dequantized_source: Optional[tuple[int, int]] = None
 
 
 @dataclass
 class SparsePrefillChunkCache:
     """Cache prefill-chunk metadata shared across layers.
 
-    Fields depend on request/token mappings and compressed page tables, not
-    per-layer k_cache; per-layer top-k combinations are recomputed into reused
-    buffers.
+    Gather layouts depend on request/token mappings and compressed page tables.
+    Shared-source dequantization keys live only for this forward; per-layer
+    top-k combinations are recomputed into reused buffers.
     """
 
     # Geometry computed once per chunk.

From e8e01bb8776e600b2cd5323d26b0490811fcdf0c Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Wed, 23 Sep 2026 18:54:42 +0800
Subject: [PATCH 15/27] Use sparse candidate consumers in CP BCG

---
 .../kernels/ops/attention/dsv4/fp4_indexer.py | 114 ++++++++++++++++++
 1 file changed, 114 insertions(+)

diff --git a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py
index 213d97c2b9c9..e8c3b05a54da 100644
--- a/python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py
+++ b/python/sglang/kernels/ops/attention/dsv4/fp4_indexer.py
@@ -399,6 +399,69 @@ def _store_fp4_index_k_cache_kernel(
     )
 
 
+@triton.jit
+def _gather_fp4_index_k_cache_masked_kernel(
+    cache,
+    req_to_token,
+    req_ids,
+    req_lens,
+    out_payload,
+    out_scales,
+    req_stride: tl.constexpr,
+    cache_stride: tl.constexpr,
+    width: tl.constexpr,
+    compress_ratio: tl.constexpr,
+    page_size: tl.constexpr,
+    block_rows: tl.constexpr,
+):
+    """Gather only the valid compressed-K prefix of each fixed graph slot."""
+    request = tl.program_id(0)
+    row_block = tl.program_id(1)
+    logical_k = row_block * block_rows + tl.arange(0, block_rows)
+    valid = logical_k < tl.load(req_lens + request) // compress_ratio
+
+    req_id = tl.load(req_ids + request)
+    full_slot = tl.load(
+        req_to_token + req_id * req_stride + logical_k.to(tl.int64) * compress_ratio,
+        mask=valid,
+        other=0,
+    )
+    compressed_slot = full_slot // compress_ratio
+    page = compressed_slot // page_size
+    page_offset = compressed_slot - page * page_size
+    out_row = request * width + logical_k
+
+    payload_cols = tl.arange(0, INDEX_K_PAYLOAD_BYTES)
+    payload = tl.load(
+        cache
+        + page[:, None] * cache_stride
+        + page_offset[:, None] * INDEX_K_PAYLOAD_BYTES
+        + payload_cols[None, :],
+        mask=valid[:, None],
+        other=0,
+    )
+    tl.store(
+        out_payload
+        + out_row[:, None].to(tl.int64) * INDEX_K_PAYLOAD_BYTES
+        + payload_cols[None, :],
+        payload,
+        mask=valid[:, None],
+    )
+
+    scale_cols = tl.arange(0, INDEX_K_SCALE_BYTES)
+    scale_bytes = tl.load(
+        cache
+        + page[:, None] * cache_stride
+        + page_size * INDEX_K_PAYLOAD_BYTES
+        + page_offset[:, None] * INDEX_K_SCALE_BYTES
+        + scale_cols[None, :],
+        mask=valid[:, None],
+        other=0,
+    ).to(tl.int32)
+    packed_scale = tl.sum(scale_bytes << (scale_cols[None, :] * 8), axis=1)
+    tl.store(out_scales + out_row, packed_scale, mask=valid)
+
+
 def quantize_fp4_indexer_tensor(
     x: torch.Tensor, rne: bool = False
 ) -> tuple[torch.Tensor, torch.Tensor]:
@@ -461,6 +524,57 @@ def store_fp4_index_k_cache(
     )
 
 
+def gather_fp4_index_k_cache_masked(
+    cache: torch.Tensor,
+    req_to_token: torch.Tensor,
+    req_ids: torch.Tensor,
+    req_lens: torch.Tensor,
+    *,
+    width: int,
+    compress_ratio: int,
+    page_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack graph-stable contiguous FP4 K without touching padded K slots.
+
+    The result reserves ``width`` rows for every request slot because
+    DeepGEMM needs fixed offsets under CUDA graph replay. Only
+    ``req_lens[r] // compress_ratio`` rows are populated. DeepGEMM and top-k
+    consume the same real lengths, so the untouched suffix is never read.
+    """
+    assert cache.ndim == 2 and cache.element_size() == 1
+    assert cache.shape[1] == page_size * INDEX_K_SLOT_BYTES
+    assert req_to_token.ndim == 2 and req_to_token.dtype == torch.int32
+    assert req_ids.ndim == 1 and req_lens.shape == req_ids.shape
+    assert width > 0 and compress_ratio > 0
+
+    num_rows = req_ids.shape[0] * width
+    payload = torch.empty(
+        (num_rows, INDEX_K_PAYLOAD_BYTES.value),
+        dtype=torch.int8,
+        device=cache.device,
+    )
+    scales = torch.empty((num_rows,), dtype=torch.int32, device=cache.device)
+    block_rows = 32
+    _gather_fp4_index_k_cache_masked_kernel[
+        (req_ids.shape[0], triton.cdiv(width, block_rows))
+    ](
+        cache,
+        req_to_token,
+        req_ids,
+        req_lens,
+        payload,
+        scales,
+        req_to_token.stride(0),
+        cache.stride(0),
+        width,
+        compress_ratio,
+        page_size,
+        block_rows,
+        num_warps=8,
+    )
+    return payload, scales
+
+
 @triton.jit
 def _gather_fp4_index_k_kernel(
     cache,

From fa75cbfcc67cc6ec904209c1eddc804e8be53b3f Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Thu, 8 Oct 2026 12:08:20 +0800
Subject: [PATCH 16/27] Port CP BCG candidate indexing to V4.1 indexer refactor

---
 .../layers/attention/deepseek_v4_backend.py   | 273 ++++++++++++++++--
 .../unit/layers/test_dsv4_nonpaged_indexer.py |  34 +--
 2 files changed, 259 insertions(+), 48 deletions(-)

diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
index cc4a5d794370..532bb2de49af 100644
--- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py
+++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
@@ -19,6 +19,14 @@
 import torch
 import torch.nn.functional as F
 
+from sglang.kernels.ops.attention.dsv4.candidate_blocks import (
+    amax_topk_blocks,
+    candidate_row_lens,
+)
+from sglang.kernels.ops.attention.dsv4.candidate_table import (
+    build_sparse_indexer_schedule,
+    sort_candidate_blocks,
+)
 from sglang.kernels.ops.attention.dsv4.decode_attention_sm100 import (
     can_use_swapab_attention,
 )
@@ -29,6 +37,13 @@
     gather_dequant_requant_fp8_paged,
     q8kv8_padded_num_heads,
 )
+from sglang.kernels.ops.attention.dsv4.fp4_indexer import (
+    gather_fp4_index_k_cache_masked,
+)
+from sglang.kernels.ops.attention.dsv4.index_logits import (
+    deep_gemm_fp4_paged_mqa_logits,
+    sparse_logits,
+)
 from sglang.kernels.ops.attention.dsv4.kv_layout import KVLayout
 from sglang.kernels.ops.attention.dsv4.metadata_kernel import (
     fill_all_compressed_indices,
@@ -37,6 +52,14 @@
     init_compression_metadata as _init_compression_metadata_triton,
 )
 from sglang.kernels.ops.attention.dsv4.online_c128_mtp import OnlineC128MTPController
+from sglang.kernels.ops.attention.dsv4.prefill_candidates import (
+    topk_prefill_candidates,
+)
+from sglang.kernels.ops.attention.dsv4.topk import (
+    topk_transform_paged_torch,
+    topk_transform_ragged_v2,
+    topk_transform_sparse,
+)
 from sglang.kernels.ops.attention.dsv4_attn_metadata_kernels import (
     BuildCausalSwaPageIndices,
     BuildPageTablePositions,
@@ -87,6 +110,10 @@
     make_candidate_indexer,
     make_full_topk_indexer,
 )
+from sglang.srt.layers.attention.dsv4.v41_indexer.scoring import (
+    get_index_k_cache,
+    quantize_index_q,
+)
 from sglang.srt.layers.attention.verify_mask import (
     VerifyMask,
     maybe_create_verify_mask,
@@ -271,26 +298,25 @@ def _dense_fp4_mqa_logits(
     )
 
 
-def _prefill_graph_dense_k_layout(
-    req_to_token: torch.Tensor,
-    pool,
-    layer_id: int,
-    ratio: int,
+def _prefill_graph_dense_k_offsets(
     width: int,
     local_req_ids: torch.Tensor,
     req_ids: torch.Tensor,
-    req_lens: torch.Tensor,
-) -> Tuple[torch.Tensor, torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-    """Pack up to eight replayed requests into capture-stable contiguous FP4 K."""
+) -> torch.Tensor:
     req_ord = (local_req_ids[:, None] == req_ids[None, :]).to(torch.int32).argmax(1)
-    ks = req_ord.to(torch.int32) * width
-    logical_k = torch.arange(width, device=req_to_token.device)
-    valid_k = logical_k[None, :] < (req_lens[:, None] // ratio)
-    logical_pos = (logical_k[None, :] * ratio).expand_as(valid_k)
-    logical_pos = logical_pos.masked_fill(~valid_k, 0)
-    k_slots = req_to_token[req_ids[:, None], logical_pos].to(torch.int64) // ratio
-    k_slots = k_slots.masked_fill(~valid_k, 0).clamp_min_(0).flatten()
-    return ks, k_slots, pool.get_low_ratio_index_k_fp4(layer_id, k_slots)
+    return req_ord * width
+
+
+@dataclass
+class _CapturedCandidateMask(CandidateMetadata):
+    block_mask: torch.Tensor
+
+
+@dataclass
+class _CapturedSparseTable(CandidateMetadata):
+    blocks: torch.Tensor
+    schedule: torch.Tensor
+    valid_lens: torch.Tensor
 
 
 def _low_ratio_source_projections(layer, x, q_lora, positions, bufs):
@@ -1038,6 +1064,7 @@ class DSV4Metadata:
     low_ratio_dense_req_indices: Optional[torch.Tensor] = None
     low_ratio_dense_seq_lens: Optional[torch.Tensor] = None
     prefill_graph_dense_indexer: bool = False
+    low_ratio_dense_k_offsets: Dict[int, torch.Tensor] = field(default_factory=dict)
 
     # Per-step scratch for TP-padded query heads, zeroed by the first user.
     # Later layers overwrite real heads and preserve the zero padding.
@@ -3378,12 +3405,7 @@ def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> Non
             else self.forward_metadata.c2_indexer_metadata
         )
         assert metadata is not None, f"no prefill graph indexer metadata for {ratio = }"
-        if indexer.uses_candidates or indexer.is_candidate_source:
-            # Inside the window every block is a candidate: the full top-k applies.
-            width = metadata.max_compressed_seq_len
-            assert (
-                width <= indexer.candidate_topk_blocks * indexer.candidate_block_size
-            ), f"prefill graph indexer width {width} exceeds the candidate window"
+        width = metadata.max_compressed_seq_len
         inputs = CapturedPrefillInputs(
             indexer=indexer,
             layer_id=layer.layer_id,
@@ -3393,7 +3415,212 @@ def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> Non
             out_page_indices=core.sparse_page_indices(ratio),
             out_raw_indices=core.sparse_raw_indices(ratio),
         )
-        self.full_topk_indexer.topk_prefill_captured(inputs)
+        dense = (
+            self.forward_metadata.prefill_graph_dense_indexer
+            and projected_q.shape[0]
+            >= _PREFILL_GRAPH_DENSE_INDEXER_MAX_REQUESTS
+            * _PREFILL_GRAPH_INDEXER_ROW_CHUNK
+        )
+        two_level = (
+            width > indexer.candidate_topk_blocks * indexer.candidate_block_size
+            and (indexer.is_candidate_source or indexer.uses_candidates)
+        )
+        if not dense and not two_level:
+            self.full_topk_indexer.topk_prefill_captured(inputs)
+            return
+        if two_level:
+            assert indexer.candidate_block_size == 8
+
+        q_fp4, q_sf = quantize_index_q(projected_q)
+        rows, heads = q_fp4.shape[:2]
+        weights = projected_w.float()
+        lens = metadata.compressed_seq_lens
+        page_table = metadata.page_table
+        page_size = metadata.compressed_page_size
+        out_page = inputs.out_page_indices
+        out_raw = inputs.out_raw_indices
+        topk = min(indexer.index_topk, width)
+        published = self.forward_metadata.candidate_metadata
+        local_req_ids = None
+        ks = None
+        if dense:
+            req_ids = self.forward_metadata.low_ratio_dense_req_indices
+            req_lens = self.forward_metadata.low_ratio_dense_seq_lens
+            local_req_ids = self.forward_metadata.low_ratio_local_req_indices
+            assert req_ids is not None and req_lens is not None
+            assert local_req_ids is not None and local_req_ids.shape[0] == rows
+            ks = self.forward_metadata.low_ratio_dense_k_offsets.get(ratio)
+            if ks is None:
+                ks = _prefill_graph_dense_k_offsets(width, local_req_ids, req_ids)
+                self.forward_metadata.low_ratio_dense_k_offsets[ratio] = ks
+
+        sparse_consumer = (
+            dense
+            and two_level
+            and indexer.uses_candidates
+            and not indexer.is_candidate_source
+        )
+        if sparse_consumer:
+            assert isinstance(published, _CapturedSparseTable)
+        elif two_level and indexer.uses_candidates:
+            assert isinstance(published, _CapturedCandidateMask)
+
+        k_cache = get_index_k_cache(
+            token_to_kv_pool=self.token_to_kv_pool,
+            layer_id=layer.layer_id,
+            page_size=page_size,
+        )
+        if dense and not sparse_consumer:
+            k_fp4 = gather_fp4_index_k_cache_masked(
+                self.token_to_kv_pool.get_index_k_with_scale_buffer(layer.layer_id),
+                self.req_to_token,
+                req_ids,
+                req_lens,
+                width=width,
+                compress_ratio=ratio,
+                page_size=page_size,
+            )
+
+        block_masks = []
+        for chunk, plan in [(slice(0, rows), None)] if dense else metadata.row_chunks():
+            chunk_lens = lens[chunk]
+            if sparse_consumer:
+                scores = sparse_logits(
+                    q_fp4[chunk].view(-1, 1, heads, 64),
+                    q_sf[chunk].view(-1, 1, heads),
+                    k_cache,
+                    weights[chunk].to(torch.bfloat16),
+                    published.schedule,
+                    published.blocks.shape[1],
+                )
+            elif dense:
+                scores = _dense_fp4_mqa_logits(
+                    (q_fp4[chunk], q_sf[chunk]),
+                    k_fp4,
+                    weights[chunk],
+                    ks[chunk],
+                    ks[chunk] + chunk_lens,
+                    width,
+                )
+            else:
+                scores = deep_gemm_fp4_paged_mqa_logits(
+                    (
+                        q_fp4[chunk].view(-1, 1, heads, 64),
+                        q_sf[chunk].view(-1, 1, heads),
+                    ),
+                    k_cache,
+                    weights[chunk],
+                    chunk_lens,
+                    page_table[chunk],
+                    plan,
+                    width,
+                )
+
+            if two_level and indexer.is_candidate_source:
+                nblocks, valid_lens = candidate_row_lens(
+                    chunk_lens, indexer.candidate_topk_blocks
+                )
+                blocks = amax_topk_blocks(
+                    scores, chunk_lens, nblocks, indexer.candidate_topk_blocks
+                )
+                if dense:
+                    sort_candidate_blocks(
+                        blocks, chunk_lens, page_table[chunk], page_size
+                    )
+                    schedule = build_sparse_indexer_schedule(
+                        blocks,
+                        chunk_lens,
+                        page_table[chunk],
+                        page_size,
+                        q_fp4.dtype,
+                        local_req_ids[chunk].to(torch.int32).contiguous(),
+                    )
+                    self.forward_metadata.candidate_metadata = _CapturedSparseTable(
+                        blocks, schedule, valid_lens
+                    )
+                else:
+                    ncols = (width + 7) // 8
+                    mask = torch.zeros(
+                        (scores.shape[0], ncols + 1),
+                        dtype=torch.bool,
+                        device=scores.device,
+                    )
+                    mask.scatter_(
+                        1, blocks.to(torch.int64).masked_fill(blocks < 0, ncols), True
+                    )
+                    block_masks.append(mask[:, :ncols])
+
+            if sparse_consumer:
+                selected = torch.empty(
+                    (scores.shape[0], topk), dtype=torch.int32, device=scores.device
+                )
+                topk_transform_sparse(
+                    scores,
+                    published.valid_lens[chunk],
+                    published.blocks[chunk],
+                    selected,
+                )
+            elif two_level and indexer.uses_candidates:
+                selected = torch.empty(
+                    (scores.shape[0], topk), dtype=torch.int32, device=scores.device
+                )
+                topk_prefill_candidates(
+                    scores,
+                    chunk_lens,
+                    published.block_mask[chunk],
+                    indexer.candidate_block_size,
+                    ks[chunk] if dense else torch.zeros_like(chunk_lens),
+                    selected,
+                )
+                if dense:
+                    selected = torch.where(
+                        selected >= 0, selected - ks[chunk, None], selected
+                    )
+            elif dense:
+                selected = torch.empty(
+                    (scores.shape[0], topk), dtype=torch.int32, device=scores.device
+                )
+                topk_transform_ragged_v2(
+                    scores, chunk_lens, out_offsets=ks[chunk], out_indices=selected
+                )
+                selected = torch.where(
+                    selected >= 0, selected - ks[chunk, None], selected
+                )
+            else:
+                topk_transform_paged_torch(
+                    scores,
+                    chunk_lens,
+                    page_table[chunk],
+                    out_page[chunk, :topk],
+                    page_size,
+                    out_raw[chunk, :topk],
+                )
+                continue
+
+            valid = selected >= 0
+            sentinel = torch.iinfo(torch.int32).max
+            selected = selected.masked_fill(~valid, sentinel).sort(dim=-1).values
+            valid = selected != sentinel
+            idx = selected.clamp_max(width - 1).to(torch.int64)
+            if dense:
+                slots = (
+                    self.req_to_token[local_req_ids[chunk, None], idx * ratio].to(
+                        torch.int64
+                    )
+                    // ratio
+                )
+            else:
+                slots = (
+                    page_table[chunk].gather(-1, idx // page_size) * page_size
+                    + idx % page_size
+                )
+            out_page[chunk, :topk] = torch.where(valid, slots, -1).to(torch.int32)
+            out_raw[chunk, :topk] = torch.where(valid, selected, -1).to(torch.int32)
+
+        if block_masks:
+            self.forward_metadata.candidate_metadata = _CapturedCandidateMask(
+                torch.cat(block_masks, dim=0)
+            )
 
     def _make_low_ratio_decode_indexer_inputs(self, layer, x, q_lora, req, pos, mode):
         ratio = layer.compress_ratio
diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
index a9ecc446bad2..2749ac769b65 100644
--- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
+++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
@@ -67,46 +67,30 @@ def test_dsv41_bcg_dense_indexer_keeps_paged_fallback_chunked(self):
                     DeepseekV4AttnBackend._low_ratio_prefill_indexer_metadata(
                         backend, core, 1
                     )
-                    self.assertEqual(metadata_ctor.call_args.kwargs["row_chunk"], 2048)
+                    self.assertEqual(
+                        metadata_ctor.call_args.kwargs["row_chunk"],
+                        min(2048, (2048 * 16384) // context),
+                    )
 
             with envs.SGLANG_DSV41_BCG_DENSE_INDEXER.override(True):
                 max_context.return_value = 32768
                 DeepseekV4AttnBackend._low_ratio_prefill_indexer_metadata(
                     backend, core, 1
                 )
-                self.assertEqual(metadata_ctor.call_args.kwargs["row_chunk"], 2048)
-                self.assertTrue(metadata_ctor.call_args.kwargs["use_topk_v2"])
+                self.assertEqual(metadata_ctor.call_args.kwargs["row_chunk"], 1024)
+                self.assertFalse(metadata_ctor.call_args.kwargs["use_topk_v2"])
 
-    def test_dsv41_bcg_dense_k_layout_uses_live_request_order(self):
+    def test_dsv41_bcg_dense_offsets_use_live_request_order(self):
         from sglang.srt.layers.attention.deepseek_v4_backend import (
-            _prefill_graph_dense_k_layout,
+            _prefill_graph_dense_k_offsets,
         )
 
-        req_to_token = torch.arange(8)[:, None] * 1000 + torch.arange(10)[None, :] * 2
-        pool = SimpleNamespace(
-            get_low_ratio_index_k_fp4=lambda layer, slots: (
-                slots[:, None].expand(-1, 64).contiguous(),
-                slots.clone(),
-            )
-        )
-        ks, slots, packed = _prefill_graph_dense_k_layout(
-            req_to_token,
-            pool,
-            1,
-            2,
+        ks = _prefill_graph_dense_k_offsets(
             5,
             torch.tensor([7, 3, 7, 3]),
             torch.tensor([7, 3, 0, 0, 0, 0, 0, 0]),
-            torch.tensor([8, 6, 0, 0, 0, 0, 0, 0]),
         )
         torch.testing.assert_close(ks, torch.tensor([0, 5, 0, 5], dtype=torch.int32))
-        torch.testing.assert_close(
-            slots[:10],
-            torch.tensor([3500, 3502, 3504, 3506, 0, 1500, 1502, 1504, 0, 0]),
-        )
-        self.assertEqual(slots.shape, (40,))
-        torch.testing.assert_close(packed[0][:, 0], slots)
-        torch.testing.assert_close(packed[1], slots)
 
     def test_dsv41_dense_logits_offsets_are_int32(self):
         from sglang.srt.layers.attention.deepseek_v4_backend import (

From e25f1342eefd39010b286b975c53aecc340fd9ec Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Mon, 21 Sep 2026 16:08:55 +0800
Subject: [PATCH 17/27] Support multistream forward prepare on CP.

---
 .../layers/attention/deepseek_v4_backend.py   |  53 +++-
 .../layers/attention/dsv4/compressor_v2.py    |   7 +-
 .../srt/layers/attention/dsv4/indexer.py      |  12 +
 python/sglang/srt/models/deepseek_v4.py       | 237 ++++++++++++++----
 4 files changed, 257 insertions(+), 52 deletions(-)

diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
index 532bb2de49af..f8f64b1b2b47 100644
--- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py
+++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
@@ -1675,6 +1675,8 @@ def low_ratio_prefill_graph(self) -> bool:
         supported = (
             bool(self.low_ratios) and has_dense_fp4_indexer() and is_sm100_or_newer()
         )
+        # CP's paged indexer needs a bounded static context width. Captured
+        # multistream prepare joins its source workers before scoring these rows.
         return supported and (
             get_parallel().attn_cp_size == 1 or _prefill_graph_max_seq_len() is not None
         )
@@ -2761,11 +2763,35 @@ def init_forward_metadata_for_breakable_cuda_graph_capture(
         )
         if self.low_ratio_prefill_graph and forward_batch.forward_mode.is_extend():
             for ratio in self.low_ratios:
-                self._source_projection_buffers(
-                    forward_batch.out_cache_loc.shape[0], ratio
-                )
+                if dsa_use_prefill_cp(forward_batch):
+                    self._cp_indexer_projection_buffers(
+                        forward_batch._cp_positions.shape[0], ratio
+                    )
+                else:
+                    self._source_projection_buffers(
+                        forward_batch.out_cache_loc.shape[0], ratio
+                    )
         return self.forward_metadata
 
+    def _cp_indexer_projection_buffers(self, num_tokens: int, ratio: int) -> dict:
+        """Capture-stable CP-local Q and weight rows for the paged indexer."""
+        sets = getattr(self, "_cp_indexer_proj_bufs", None)
+        if sets is None:
+            sets = self._cp_indexer_proj_bufs = {}
+        key = (ratio, num_tokens)
+        if key not in sets:
+            cfg = self.model_runner.model_config.hf_text_config
+            heads, dim = int(cfg.index_n_heads), int(cfg.index_head_dim)
+            sets[key] = {
+                "q": torch.zeros(
+                    num_tokens, heads, dim, dtype=torch.bfloat16, device=self.device
+                ),
+                "w": torch.zeros(
+                    num_tokens, heads, dtype=torch.bfloat16, device=self.device
+                ),
+            }
+        return sets[key]
+
     def _source_projection_buffers(self, num_tokens: int, ratio: int) -> dict:
         cfg = self.model_runner.model_config.hf_text_config
         heads, dim = int(cfg.index_n_heads), int(cfg.index_head_dim)
@@ -2909,6 +2935,7 @@ def forward_low_ratio_sources(
         forward_batch: ForwardBatch,
         run_compressor: bool = True,
         run_indexer: bool = True,
+        precomputed_x_global: Optional[torch.Tensor] = None,
     ) -> None:
         """Runs on every ratio 1/2 layer before its attention."""
         if forward_batch.forward_mode.is_idle():
@@ -2924,6 +2951,7 @@ def forward_low_ratio_sources(
                 forward_batch=forward_batch,
                 run_compressor=run_compressor,
                 run_indexer=run_indexer,
+                precomputed_x_global=precomputed_x_global,
             )
             return
         meta = self.forward_metadata
@@ -2959,15 +2987,26 @@ def forward_low_ratio_sources(
             self._low_ratio_index_topk(layer, x, q_lora, req, pos, forward_batch)
 
     def _forward_low_ratio_sources_cp(
-        self, *, layer, x, q_lora, positions, forward_batch, run_compressor, run_indexer
+        self,
+        *,
+        layer,
+        x,
+        q_lora,
+        positions,
+        forward_batch,
+        run_compressor,
+        run_indexer,
+        precomputed_x_global: Optional[torch.Tensor] = None,
     ) -> None:
         # Every rank writes the whole prompt's compressed state, scoring its own rows.
         cp_meta = forward_batch.attn_cp_metadata
         tail = self.forward_metadata.late_layer_tail
         if run_compressor and layer.compressor is not None:
-            x_global = cp_materialize_global_token_order(
-                x.contiguous(), forward_batch, torch.cuda.current_stream()
-            )
+            x_global = precomputed_x_global
+            if x_global is None:
+                x_global = cp_materialize_global_token_order(
+                    x.contiguous(), forward_batch, torch.cuda.current_stream()
+                )
             if self._low_ratio_in_prefill_graph():
                 assert tail is None, "bounded SWA replay cannot enter prefill BCG"
                 core = self.forward_metadata.core_metadata
diff --git a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py
index 3afb09f3ff46..8368915159d4 100644
--- a/python/sglang/srt/layers/attention/dsv4/compressor_v2.py
+++ b/python/sglang/srt/layers/attention/dsv4/compressor_v2.py
@@ -158,13 +158,18 @@ def forward_unified(
         forward_batch: ForwardBatch,
         layer_id: int,
         compressor: Compressor,
+        precomputed_kv_score: Optional[torch.Tensor] = None,
     ) -> None:
         if forward_batch.forward_mode.is_idle():
             return
 
         token_to_kv_pool = self.token_to_kv_pool
         token_to_kv_pool = cast("DeepSeekV4TokenToKVPool", token_to_kv_pool)
-        kv_score_input = compressor.compute_kv_score(x, forward_batch)
+        kv_score_input = (
+            precomputed_kv_score
+            if precomputed_kv_score is not None
+            else compressor.compute_kv_score(x, forward_batch)
+        )
 
         state_pool = compressor.get_state_pool(self)
         from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py
index f0bac3d4108a..e1560e1b6f45 100644
--- a/python/sglang/srt/layers/attention/dsv4/indexer.py
+++ b/python/sglang/srt/layers/attention/dsv4/indexer.py
@@ -505,6 +505,7 @@ def _forward_prepare_multi_stream(
         forward_batch: ForwardBatch,
         alt_streams: Optional[List[torch.cuda.Stream]] = None,
         q_lora_ready: Optional[torch.cuda.Event] = None,
+        precomputed_kv_score: Optional[torch.Tensor] = None,
     ) -> Tuple[IndexerQuery, torch.Tensor]:
         if TYPE_CHECKING:
             assert isinstance(self, CompressorBackendMixin)
@@ -523,6 +524,7 @@ def _forward_prepare_multi_stream(
             forward_batch=forward_batch,
             layer_id=c4_indexer.layer_id,
             compressor=c4_indexer.compressor,
+            precomputed_kv_score=precomputed_kv_score,
         )
 
         # The weight projection is small and fast; compute it on its own
@@ -550,6 +552,7 @@ def _forward_prepare_normal(
         positions: torch.Tensor,
         forward_batch: ForwardBatch,
         skip_compressor: bool = False,
+        precomputed_kv_score: Optional[torch.Tensor] = None,
     ) -> Tuple[IndexerQuery, torch.Tensor]:
         if TYPE_CHECKING:
             assert isinstance(self, CompressorBackendMixin)
@@ -562,6 +565,7 @@ def _forward_prepare_normal(
                 forward_batch=forward_batch,
                 layer_id=c4_indexer.layer_id,
                 compressor=c4_indexer.compressor,
+                precomputed_kv_score=precomputed_kv_score,
             )
         return q, weights
 
@@ -749,6 +753,7 @@ def forward_c4_indexer(
         enable_multi_stream: bool = False,
         q_lora_ready: Optional[torch.cuda.Event] = None,
         skip_compressor: bool = False,
+        precomputed_kv_score: Optional[torch.Tensor] = None,
     ) -> None:
         if forward_batch.forward_mode.is_idle():
             return
@@ -787,6 +792,7 @@ def forward_c4_indexer(
                 forward_batch=forward_batch,
                 alt_streams=alt_streams,
                 q_lora_ready=q_lora_ready,
+                precomputed_kv_score=precomputed_kv_score,
             )
         else:
             assert q_lora_ready is None
@@ -797,6 +803,7 @@ def forward_c4_indexer(
                 positions=positions,
                 forward_batch=forward_batch,
                 skip_compressor=skip_compressor,
+                precomputed_kv_score=precomputed_kv_score,
             )
 
         use_fp4_indexer = c4_indexer.use_fp4_indexer
@@ -1213,7 +1220,11 @@ def forward(
         enable_multi_stream: bool = False,
         q_lora_ready: Optional[torch.cuda.Event] = None,
         skip_compressor: bool = False,
+        precomputed_kv_score: Optional[torch.Tensor] = None,
     ) -> None:
+        kwargs = {}
+        if precomputed_kv_score is not None:
+            kwargs["precomputed_kv_score"] = precomputed_kv_score
         return attn_backend.forward_c4_indexer(
             x=x,
             q_lora=q_lora,
@@ -1223,4 +1234,5 @@ def forward(
             enable_multi_stream=enable_multi_stream,
             q_lora_ready=q_lora_ready,
             skip_compressor=skip_compressor,
+            **kwargs,
         )
diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index 56108033b618..7f831c6512cc 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -1457,6 +1457,45 @@ def _compute_kv_bf16(
         )
         return kv
 
+    def _materialize_cp_swa_k(
+        self,
+        x: torch.Tensor,
+        forward_batch: ForwardBatch,
+        qkv_a: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if qkv_a is not None:
+            kv = qkv_a[..., self.q_lora_rank :]
+        else:
+            kv, _ = self.wkv(x)
+        return cp_materialize_global_token_order(
+            kv.contiguous(), forward_batch, torch.cuda.current_stream()
+        )
+
+    def _store_cp_swa_k(
+        self,
+        kv: torch.Tensor,
+        forward_batch: ForwardBatch,
+        attn_backend,
+    ) -> None:
+        # Eager prefill writes the live prefix. BCG captures the full bucket;
+        # replay refreshes padding locations to the dummy cache slot.
+        kv = kv[: forward_batch.attn_cp_metadata.total_seq_lens]
+        tail = attn_backend.forward_metadata.late_layer_tail
+        global_positions = (
+            tail.pos_global
+            if tail is not None
+            else forward_batch.positions[: kv.shape[0]]
+        )
+        get_token_to_kv_pool().set_swa_key_buffer_radix_fused_norm_rope(
+            layer_id=self.layer_id,
+            swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch),
+            kv=kv,
+            kv_weight=self.kv_norm.weight.data,
+            eps=self.eps,
+            freqs_cis=self.freqs_cis,
+            positions=global_positions,
+        )
+
     def _forward_prepare_multi_stream(
         self,
         x: torch.Tensor,
@@ -1474,20 +1513,40 @@ def _forward_prepare_multi_stream(
         stream_compressor = self.alt_streams[1]
         stream_indexer = self.alt_streams[2]
 
-        stream_kv.wait_stream(current_stream)
-        stream_compressor.wait_stream(current_stream)
-        stream_indexer.wait_stream(current_stream)
+        use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)
+        if not use_cp:
+            stream_kv.wait_stream(current_stream)
+            stream_compressor.wait_stream(current_stream)
+            stream_indexer.wait_stream(current_stream)
 
         x_linear = x_quant if x_quant is not None else x
         qkv_a: Optional[torch.Tensor] = None
         qkv_a_ready: Optional[torch.cuda.Event] = None
         if self.fuse_wqa_wkv:
             qkv_a, _ = self.wqkv_a(x_linear)
-            qkv_a_ready = current_stream.record_event()
+            if not use_cp:
+                qkv_a_ready = current_stream.record_event()
 
         q_lora, q_for_wqb = self._compute_q_a(x_linear, qkv_a=qkv_a)
         q_lora_ready = current_stream.record_event()
 
+        swa_k: Optional[torch.Tensor] = None
+        indexer_kv_score: Optional[torch.Tensor] = None
+        compressor_kv_score: Optional[torch.Tensor] = None
+        if use_cp:
+            # Keep CP collectives in parent-stream order, then release all
+            # worker streams after their gathered inputs are ready.
+            swa_k = self._materialize_cp_swa_k(x_linear, forward_batch, qkv_a=qkv_a)
+            if self.indexer is not None:
+                indexer_kv_score = self.indexer.compressor.compute_kv_score(
+                    x, forward_batch
+                )
+            if self.compressor is not None:
+                compressor_kv_score = self.compressor.compute_kv_score(x, forward_batch)
+            stream_kv.wait_stream(current_stream)
+            stream_compressor.wait_stream(current_stream)
+            stream_indexer.wait_stream(current_stream)
+
         if self.indexer is not None:
             with torch.cuda.stream(stream_indexer):
                 self.indexer(
@@ -1497,20 +1556,29 @@ def _forward_prepare_multi_stream(
                     attn_backend=attn_backend,
                     enable_multi_stream=True,
                     q_lora_ready=q_lora_ready,
+                    precomputed_kv_score=indexer_kv_score,
                 )
 
         with torch.cuda.stream(stream_kv):
-            if qkv_a_ready is not None:
-                stream_kv.wait_event(qkv_a_ready)
-            # Fused norm + rope + cache write -- no bf16 KV intermediate.
-            self._compute_kv_to_cache(
-                x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
-            )
+            if use_cp:
+                assert swa_k is not None
+                self._store_cp_swa_k(swa_k, forward_batch, attn_backend)
+            else:
+                if qkv_a_ready is not None:
+                    stream_kv.wait_event(qkv_a_ready)
+                # Fused norm + rope + cache write -- no bf16 KV intermediate.
+                self._compute_kv_to_cache(
+                    x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
+                )
 
         if self.compressor is not None:
             with torch.cuda.stream(stream_compressor):
                 attn_backend.forward_core_compressor(
-                    x, forward_batch, self.layer_id, self.compressor
+                    x,
+                    forward_batch,
+                    self.layer_id,
+                    self.compressor,
+                    precomputed_kv_score=compressor_kv_score,
                 )
 
         q = self._compute_q_b(q_for_wqb, positions, q_out)
@@ -1521,6 +1589,78 @@ def _forward_prepare_multi_stream(
 
         return q
 
+    def _forward_prepare_low_ratio_cp_multi_stream(
+        self,
+        x: torch.Tensor,
+        positions: torch.Tensor,
+        forward_batch: ForwardBatch,
+        attn_backend,
+        q_out: Optional[torch.Tensor] = None,
+        x_quant=None,
+    ) -> torch.Tensor:
+        """Overlap CP SWA/cache sources with the Q projection after allgather."""
+        assert self.alt_streams is not None
+        current_stream = torch.cuda.current_stream()
+        stream_kv, stream_sources = self.alt_streams[0], self.alt_streams[-1]
+        x_linear = x_quant if x_quant is not None else x
+
+        qkv_a: Optional[torch.Tensor] = None
+        if self.fuse_wqa_wkv:
+            qkv_a, _ = self.wqkv_a(x_linear)
+        q_lora, q_for_wqb = self._compute_q_a(x_linear, qkv_a=qkv_a)
+
+        # Collectives must launch in identical parent-stream order on every CP
+        # rank. Neither worker may start consuming until both gathers finish.
+        swa_k = self._materialize_cp_swa_k(x_linear, forward_batch, qkv_a=qkv_a)
+        x_global = None
+        if self.compressor is not None and not forward_batch.encoder_swa_replay:
+            x_global = cp_materialize_global_token_order(
+                x.contiguous(), forward_batch, current_stream
+            )
+
+        captured = is_in_breakable_cuda_graph()
+        indexer_buffers = None
+        if captured and self.indexer is not None:
+            assert attn_backend.low_ratio_prefill_graph
+            indexer_buffers = attn_backend._cp_indexer_projection_buffers(
+                x.shape[0], self.compress_ratio
+            )
+
+        stream_kv.wait_stream(current_stream)
+        stream_sources.wait_stream(current_stream)
+        with torch.cuda.stream(stream_kv):
+            self._store_cp_swa_k(swa_k, forward_batch, attn_backend)
+        with torch.cuda.stream(stream_sources):
+            if self.compressor is not None or (
+                self.indexer is not None and not captured
+            ):
+                attn_backend.forward_low_ratio_sources(
+                    layer=self,
+                    x=x,
+                    q_lora=None if captured else q_lora,
+                    positions=positions,
+                    forward_batch=forward_batch,
+                    run_indexer=not captured,
+                    precomputed_x_global=x_global,
+                )
+            if indexer_buffers is not None:
+                # Capture the full CP-local bucket; the paged metadata masks
+                # padding when the indexer runs after the worker join.
+                indexer_buffers["q"].copy_(
+                    self.indexer.queries(q_lora, self.freqs_cis[positions])
+                )
+                indexer_buffers["w"].copy_(self.indexer.head_weights(x))
+
+        q = self._compute_q_b(q_for_wqb, positions, q_out)
+        current_stream.wait_stream(stream_kv)
+        current_stream.wait_stream(stream_sources)
+        if indexer_buffers is not None:
+            attn_backend._low_ratio_index_topk_captured(
+                self, indexer_buffers["q"], indexer_buffers["w"]
+            )
+        del qkv_a
+        return q
+
     def _forward_prepare_low_ratio_multi_stream(
         self,
         x: torch.Tensor,
@@ -1530,6 +1670,16 @@ def _forward_prepare_low_ratio_multi_stream(
         q_out: Optional[torch.Tensor] = None,
         x_quant=None,
     ) -> torch.Tensor:
+        if self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch):
+            return self._forward_prepare_low_ratio_cp_multi_stream(
+                x,
+                positions,
+                forward_batch,
+                attn_backend,
+                q_out,
+                x_quant=x_quant,
+            )
+
         # Both side streams are joined before returning, and nothing they read is
         # released before the join.
         assert self.alt_streams is not None
@@ -2067,31 +2217,8 @@ def _forward_prepare(
                     forward_batch=forward_batch,
                 )
             elif use_cp:
-                # every rank writes the whole chunk's window KV with the fused fp32 store
-                if qkv_a is not None:
-                    kv = qkv_a[..., self.q_lora_rank :]
-                else:
-                    kv, _ = self.wkv(x_linear)
-                kv = cp_materialize_global_token_order(
-                    kv.contiguous(),
-                    forward_batch,
-                    torch.cuda.current_stream(),
-                )
-                tail = attn_backend.forward_metadata.late_layer_tail
-                global_positions = (
-                    tail.pos_global
-                    if tail is not None
-                    else forward_batch.positions[: kv.shape[0]]
-                )
-                get_token_to_kv_pool().set_swa_key_buffer_radix_fused_norm_rope(
-                    layer_id=self.layer_id,
-                    swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch),
-                    kv=kv,
-                    kv_weight=self.kv_norm.weight.data,
-                    eps=self.eps,
-                    freqs_cis=self.freqs_cis,
-                    positions=global_positions,
-                )
+                kv = self._materialize_cp_swa_k(x_linear, forward_batch, qkv_a=qkv_a)
+                self._store_cp_swa_k(kv, forward_batch, attn_backend)
                 kv = None
             else:
                 self._compute_kv_to_cache(
@@ -2185,6 +2312,15 @@ def forward(
                 (DeepseekV4AttnBackend, DeepseekV4HipRadixBackend),
             )
 
+        from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
+            is_unified_kv_fp8,
+            is_unified_kv_triton,
+        )
+
+        unified = is_unified_kv_triton()
+        use_prefill_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(
+            forward_batch
+        )
         enable_multi_stream = (
             envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get()
             and self.alt_streams is not None
@@ -2204,6 +2340,24 @@ def forward(
             and not forward_batch.forward_mode.is_extend_or_draft_extend_or_mixed()
         )
 
+        low_ratio_cp_multi_stream = (
+            self.is_dsv41
+            and use_prefill_cp
+            and forward_batch.forward_mode.is_extend()
+            and envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get()
+            and self.alt_streams is not None
+            and not unified
+            and (
+                not get_is_capture_mode()
+                or (
+                    (
+                        is_in_breakable_cuda_graph()
+                        or x.shape[0] <= self._multi_stream_bs_limit
+                    )
+                    and getattr(attn_backend, "low_ratio_prefill_graph", False)
+                )
+            )
+        )
         low_ratio_multi_stream = (
             ((_is_cuda and get_platform().is_blackwell) or _is_gfx95_supported)
             and self.compress_ratio in (1, 2)
@@ -2219,14 +2373,9 @@ def forward(
                         == Mxfp8DenseGemmBackend.FLASHINFER_CUTEDSL
                     )
                 )
+                or low_ratio_cp_multi_stream
             )
         )
-        from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
-            is_unified_kv_fp8,
-            is_unified_kv_triton,
-        )
-
-        unified = is_unified_kv_triton()
         unified_fp8_verify = (
             unified
             and is_unified_kv_fp8()
@@ -2332,8 +2481,8 @@ def forward(
         attn_sink = self._local_attn_sink(kernel_num_heads)
 
         if enable_multi_stream:
-            # Multi-stream path always fuses cache write into the K kernel,
-            # so the bf16 KV intermediate is gone.
+            # CP gathers raw KV before its side-stream store. Other paths fuse
+            # the store into the K kernel; neither returns a KV intermediate.
             if _is_hip:
                 q = self._forward_prepare_multi_stream_hip(
                     x,

From e5ee13a556638eb4c91d931ba8ead703f5de61a8 Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Wed, 23 Sep 2026 20:48:59 +0800
Subject: [PATCH 18/27] Overlap CP BCG indexer preparation

---
 .../layers/attention/deepseek_v4_backend.py   | 67 +++++++++++++++++--
 python/sglang/srt/models/deepseek_v4.py       | 27 ++++++--
 2 files changed, 82 insertions(+), 12 deletions(-)

diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
index f8f64b1b2b47..294e03bee90e 100644
--- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py
+++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
@@ -3434,7 +3434,55 @@ def _low_ratio_index_topk(
             else:
                 self.full_topk_indexer.topk_prefill(inputs)
 
-    def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> None:
+    @staticmethod
+    def _low_ratio_quantize_q_prefill_graph(q):
+        return quantize_index_q(q)
+
+    def _low_ratio_gather_k_prefill_graph(self, layer):
+        metadata = (
+            self.forward_metadata.c1_indexer_metadata
+            if layer.compress_ratio == 1
+            else self.forward_metadata.c2_indexer_metadata
+        )
+        assert metadata is not None
+        if not (
+            self.forward_metadata.prefill_graph_dense_indexer
+            and metadata.compressed_seq_lens.shape[0]
+            >= _PREFILL_GRAPH_DENSE_INDEXER_MAX_REQUESTS
+            * _PREFILL_GRAPH_INDEXER_ROW_CHUNK
+        ):
+            return None
+        indexer = layer.indexer
+        width = metadata.max_compressed_seq_len
+        two_level = (
+            width > indexer.candidate_topk_blocks * indexer.candidate_block_size
+            and (indexer.is_candidate_source or indexer.uses_candidates)
+        )
+        if two_level and indexer.uses_candidates and not indexer.is_candidate_source:
+            return None
+        req_ids = self.forward_metadata.low_ratio_dense_req_indices
+        req_lens = self.forward_metadata.low_ratio_dense_seq_lens
+        local_req_ids = self.forward_metadata.low_ratio_local_req_indices
+        assert req_ids is not None and req_lens is not None and local_req_ids is not None
+        ratio = layer.compress_ratio
+        ks = self.forward_metadata.low_ratio_dense_k_offsets.get(ratio)
+        if ks is None:
+            ks = _prefill_graph_dense_k_offsets(width, local_req_ids, req_ids)
+            self.forward_metadata.low_ratio_dense_k_offsets[ratio] = ks
+        k_fp4 = gather_fp4_index_k_cache_masked(
+            self.token_to_kv_pool.get_index_k_with_scale_buffer(layer.layer_id),
+            self.req_to_token,
+            req_ids,
+            req_lens,
+            width=width,
+            compress_ratio=ratio,
+            page_size=metadata.compressed_page_size,
+        )
+        return k_fp4, ks
+
+    def _low_ratio_index_topk_captured(
+        self, layer, projected_q, projected_w, *, prepared_q=None, prepared_dense_k=None
+    ) -> None:
         ratio = layer.compress_ratio
         indexer = layer.indexer
         core = self.forward_metadata.core_metadata
@@ -3470,7 +3518,9 @@ def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> Non
         if two_level:
             assert indexer.candidate_block_size == 8
 
-        q_fp4, q_sf = quantize_index_q(projected_q)
+        q_fp4, q_sf = (
+            quantize_index_q(projected_q) if prepared_q is None else prepared_q
+        )
         rows, heads = q_fp4.shape[:2]
         weights = projected_w.float()
         lens = metadata.compressed_seq_lens
@@ -3488,10 +3538,13 @@ def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> Non
             local_req_ids = self.forward_metadata.low_ratio_local_req_indices
             assert req_ids is not None and req_lens is not None
             assert local_req_ids is not None and local_req_ids.shape[0] == rows
-            ks = self.forward_metadata.low_ratio_dense_k_offsets.get(ratio)
-            if ks is None:
-                ks = _prefill_graph_dense_k_offsets(width, local_req_ids, req_ids)
-                self.forward_metadata.low_ratio_dense_k_offsets[ratio] = ks
+            if prepared_dense_k is not None:
+                k_fp4, ks = prepared_dense_k
+            else:
+                ks = self.forward_metadata.low_ratio_dense_k_offsets.get(ratio)
+                if ks is None:
+                    ks = _prefill_graph_dense_k_offsets(width, local_req_ids, req_ids)
+                    self.forward_metadata.low_ratio_dense_k_offsets[ratio] = ks
 
         sparse_consumer = (
             dense
@@ -3509,7 +3562,7 @@ def _low_ratio_index_topk_captured(self, layer, projected_q, projected_w) -> Non
             layer_id=layer.layer_id,
             page_size=page_size,
         )
-        if dense and not sparse_consumer:
+        if dense and not sparse_consumer and prepared_dense_k is None:
             k_fp4 = gather_fp4_index_k_cache_masked(
                 self.token_to_kv_pool.get_index_k_with_scale_buffer(layer.layer_id),
                 self.req_to_token,
diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index 7f831c6512cc..a9aec2fecdd2 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -1600,8 +1600,9 @@ def _forward_prepare_low_ratio_cp_multi_stream(
     ) -> torch.Tensor:
         """Overlap CP SWA/cache sources with the Q projection after allgather."""
         assert self.alt_streams is not None
+        assert len(self.alt_streams) >= 3
         current_stream = torch.cuda.current_stream()
-        stream_kv, stream_sources = self.alt_streams[0], self.alt_streams[-1]
+        stream_kv, stream_compressor, stream_indexer = self.alt_streams[:3]
         x_linear = x_quant if x_quant is not None else x
 
         qkv_a: Optional[torch.Tensor] = None
@@ -1627,10 +1628,12 @@ def _forward_prepare_low_ratio_cp_multi_stream(
             )
 
         stream_kv.wait_stream(current_stream)
-        stream_sources.wait_stream(current_stream)
+        stream_compressor.wait_stream(current_stream)
+        stream_indexer.wait_stream(current_stream)
         with torch.cuda.stream(stream_kv):
             self._store_cp_swa_k(swa_k, forward_batch, attn_backend)
-        with torch.cuda.stream(stream_sources):
+        prepared_dense_k = None
+        with torch.cuda.stream(stream_compressor):
             if self.compressor is not None or (
                 self.indexer is not None and not captured
             ):
@@ -1643,6 +1646,12 @@ def _forward_prepare_low_ratio_cp_multi_stream(
                     run_indexer=not captured,
                     precomputed_x_global=x_global,
                 )
+            if indexer_buffers is not None:
+                # The compressor above writes this layer's index-K cache. Gather
+                # its live dense prefixes on the same stream after that write.
+                prepared_dense_k = attn_backend._low_ratio_gather_k_prefill_graph(self)
+        prepared_q = None
+        with torch.cuda.stream(stream_indexer):
             if indexer_buffers is not None:
                 # Capture the full CP-local bucket; the paged metadata masks
                 # padding when the indexer runs after the worker join.
@@ -1650,13 +1659,21 @@ def _forward_prepare_low_ratio_cp_multi_stream(
                     self.indexer.queries(q_lora, self.freqs_cis[positions])
                 )
                 indexer_buffers["w"].copy_(self.indexer.head_weights(x))
+                prepared_q = attn_backend._low_ratio_quantize_q_prefill_graph(
+                    indexer_buffers["q"]
+                )
 
         q = self._compute_q_b(q_for_wqb, positions, q_out)
         current_stream.wait_stream(stream_kv)
-        current_stream.wait_stream(stream_sources)
+        current_stream.wait_stream(stream_compressor)
+        current_stream.wait_stream(stream_indexer)
         if indexer_buffers is not None:
             attn_backend._low_ratio_index_topk_captured(
-                self, indexer_buffers["q"], indexer_buffers["w"]
+                self,
+                indexer_buffers["q"],
+                indexer_buffers["w"],
+                prepared_q=prepared_q,
+                prepared_dense_k=prepared_dense_k,
             )
         del qkv_a
         return q

From b47fc37cf776b8f27708efe4f5eb1e9fd79d8169 Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Thu, 8 Oct 2026 12:23:51 +0800
Subject: [PATCH 19/27] Fix CP multistream indexer import and formatting

---
 python/sglang/srt/layers/attention/deepseek_v4_backend.py | 4 +++-
 python/sglang/srt/models/deepseek_v4.py                   | 1 +
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
index 294e03bee90e..e086c05d0ebb 100644
--- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py
+++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
@@ -3463,7 +3463,9 @@ def _low_ratio_gather_k_prefill_graph(self, layer):
         req_ids = self.forward_metadata.low_ratio_dense_req_indices
         req_lens = self.forward_metadata.low_ratio_dense_seq_lens
         local_req_ids = self.forward_metadata.low_ratio_local_req_indices
-        assert req_ids is not None and req_lens is not None and local_req_ids is not None
+        assert (
+            req_ids is not None and req_lens is not None and local_req_ids is not None
+        )
         ratio = layer.compress_ratio
         ks = self.forward_metadata.low_ratio_dense_k_offsets.get(ratio)
         if ks is None:
diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index a9aec2fecdd2..ed56968e4682 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -59,6 +59,7 @@
     use_npu_arch35_mxfp8_wo_a,
 )
 from sglang.srt.layers.attention.dsa.utils import (
+    dsa_use_prefill_cp,
     is_dsa_enable_prefill_cp,
 )
 from sglang.srt.layers.attention.dsv4.compressor import Compressor

From 5e5e0803a4d58336cd37f5217ae053d0649285bc Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Thu, 8 Oct 2026 16:32:11 +0800
Subject: [PATCH 20/27] Fix V4.1 multistream indexer dispatch and MegaMoE
 compatibility

---
 .../layers/attention/deepseek_v4_backend.py   |  78 ++---
 .../attention/dsv4/v41_indexer/full_topk.py   |  25 +-
 python/sglang/srt/layers/moe/mega_moe.py      |  13 +-
 .../disaggregation/test_dsv41_dspark_pd.py    |   7 +-
 .../layers/moe/test_mega_moe_deepgemm_api.py  |  47 +++
 .../unit/layers/test_dsv4_nonpaged_indexer.py | 267 ++++++++++++++++--
 6 files changed, 340 insertions(+), 97 deletions(-)

diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
index e086c05d0ebb..44e78f766c25 100644
--- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py
+++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
@@ -4,16 +4,7 @@
 import functools
 import logging
 from dataclasses import dataclass, field
-from typing import (
-    TYPE_CHECKING,
-    Dict,
-    List,
-    Literal,
-    Optional,
-    Tuple,
-    TypeVar,
-    Union,
-)
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, TypeVar, Union
 
 import msgspec
 import torch
@@ -52,11 +43,8 @@
     init_compression_metadata as _init_compression_metadata_triton,
 )
 from sglang.kernels.ops.attention.dsv4.online_c128_mtp import OnlineC128MTPController
-from sglang.kernels.ops.attention.dsv4.prefill_candidates import (
-    topk_prefill_candidates,
-)
+from sglang.kernels.ops.attention.dsv4.prefill_candidates import topk_prefill_candidates
 from sglang.kernels.ops.attention.dsv4.topk import (
-    topk_transform_paged_torch,
     topk_transform_ragged_v2,
     topk_transform_sparse,
 )
@@ -83,11 +71,11 @@
     FusedCompressMetadata,
     create_paged_compressor_data,
 )
-from sglang.srt.layers.attention.dsv4.dsv41_sparse import (
-    _rope_fq4,
-    token_req_indices,
+from sglang.srt.layers.attention.dsv4.dsv41_sparse import _rope_fq4, token_req_indices
+from sglang.srt.layers.attention.dsv4.indexer import (
+    C4IndexerBackendMixin,
+    topk_transform_paged_from_metadata,
 )
-from sglang.srt.layers.attention.dsv4.indexer import C4IndexerBackendMixin
 from sglang.srt.layers.attention.dsv4.metadata import (
     _LARGE_INDEXER_QUERY_THRESHOLD,
     PagedIndexerMetadata,
@@ -114,18 +102,12 @@
     get_index_k_cache,
     quantize_index_q,
 )
-from sglang.srt.layers.attention.verify_mask import (
-    VerifyMask,
-    maybe_create_verify_mask,
-)
+from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
 from sglang.srt.layers.cp.interleave import (
     InterleaveContextParallelMetadata,
     interleave_rows_per_request,
 )
-from sglang.srt.layers.cp.utils import (
-    cp_materialize_global_token_order,
-    is_cp_active,
-)
+from sglang.srt.layers.cp.utils import cp_materialize_global_token_order, is_cp_active
 from sglang.srt.layers.dp_attention import (
     get_local_dp_buffer_len,
     set_local_dp_buffer_len,
@@ -133,12 +115,7 @@
 from sglang.srt.mem_cache.deepseek_v4_compress_state import KVAndScore
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
-from sglang.srt.runtime_context import (
-    get_exec,
-    get_parallel,
-    get_platform,
-    get_spec,
-)
+from sglang.srt.runtime_context import get_exec, get_parallel, get_platform, get_spec
 from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc
 from sglang.srt.speculative.ragged_verify import (
     RaggedVerifyMode,
@@ -304,7 +281,7 @@ def _prefill_graph_dense_k_offsets(
     req_ids: torch.Tensor,
 ) -> torch.Tensor:
     req_ord = (local_req_ids[:, None] == req_ids[None, :]).to(torch.int32).argmax(1)
-    return req_ord * width
+    return (req_ord * width).to(torch.int32)
 
 
 @dataclass
@@ -1663,7 +1640,7 @@ def _low_ratio_prefill_indexer_metadata(
             compressed_page_size=index_page_size,
             page_table=page_table,
             compressed_seq_lens=c_seq_lens,
-            use_topk_v2=False,
+            use_topk_v2=self.dsa_topk_backend.should_use_topk_v2() and not _is_xpu,
             use_prefill_cuda_graph=True,
             compress_ratio=compress_ratio,
             row_chunk=row_chunk if row_chunk < c_seq_lens.shape[0] else 0,
@@ -3445,12 +3422,7 @@ def _low_ratio_gather_k_prefill_graph(self, layer):
             else self.forward_metadata.c2_indexer_metadata
         )
         assert metadata is not None
-        if not (
-            self.forward_metadata.prefill_graph_dense_indexer
-            and metadata.compressed_seq_lens.shape[0]
-            >= _PREFILL_GRAPH_DENSE_INDEXER_MAX_REQUESTS
-            * _PREFILL_GRAPH_INDEXER_ROW_CHUNK
-        ):
+        if not self.forward_metadata.prefill_graph_dense_indexer:
             return None
         indexer = layer.indexer
         width = metadata.max_compressed_seq_len
@@ -3504,12 +3476,7 @@ def _low_ratio_index_topk_captured(
             out_page_indices=core.sparse_page_indices(ratio),
             out_raw_indices=core.sparse_raw_indices(ratio),
         )
-        dense = (
-            self.forward_metadata.prefill_graph_dense_indexer
-            and projected_q.shape[0]
-            >= _PREFILL_GRAPH_DENSE_INDEXER_MAX_REQUESTS
-            * _PREFILL_GRAPH_INDEXER_ROW_CHUNK
-        )
+        dense = self.forward_metadata.prefill_graph_dense_indexer
         two_level = (
             width > indexer.candidate_topk_blocks * indexer.candidate_block_size
             and (indexer.is_candidate_source or indexer.uses_candidates)
@@ -3576,7 +3543,8 @@ def _low_ratio_index_topk_captured(
             )
 
         block_masks = []
-        for chunk, plan in [(slice(0, rows), None)] if dense else metadata.row_chunks():
+        chunks = [(slice(0, rows), None)] if dense else metadata.row_chunks()
+        for chunk_idx, (chunk, plan) in enumerate(chunks):
             chunk_lens = lens[chunk]
             if sparse_consumer:
                 scores = sparse_logits(
@@ -3681,13 +3649,17 @@ def _low_ratio_index_topk_captured(
                     selected >= 0, selected - ks[chunk, None], selected
                 )
             else:
-                topk_transform_paged_torch(
+                topk_transform_paged_from_metadata(
                     scores,
-                    chunk_lens,
-                    page_table[chunk],
-                    out_page[chunk, :topk],
-                    page_size,
-                    out_raw[chunk, :topk],
+                    metadata,
+                    out_page[:, :topk],
+                    out_raw[:, :topk],
+                    rows=chunk,
+                    topk_metadata=(
+                        metadata.topk_plan_for_chunk(chunk_idx, chunk)
+                        if metadata.use_topk_v2
+                        else None
+                    ),
                 )
                 continue
 
diff --git a/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py
index a2f317a9d392..3fa94d0b6039 100644
--- a/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py
+++ b/python/sglang/srt/layers/attention/dsv4/v41_indexer/full_topk.py
@@ -10,7 +10,6 @@
 from sglang.kernels.ops.attention.dsv4.index_logits import (
     deep_gemm_fp4_paged_mqa_logits,
 )
-from sglang.kernels.ops.attention.dsv4.topk import topk_transform_paged_torch
 from sglang.srt.layers.attention.dsv4.indexer import topk_transform_paged_from_metadata
 
 from .scoring import (
@@ -24,11 +23,7 @@
     select_decode,
     write_prefill,
 )
-from .types import (
-    CapturedPrefillInputs,
-    DecodeInputs,
-    PrefillInputs,
-)
+from .types import CapturedPrefillInputs, DecodeInputs, PrefillInputs
 
 if TYPE_CHECKING:
     from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
@@ -156,7 +151,7 @@ def _deep_gemm_prefill_captured(self, inputs: CapturedPrefillInputs) -> None:
         lens = metadata.compressed_seq_lens
         page_table = metadata.page_table
         topk = min(indexer.index_topk, width)
-        for rows, plan in metadata.row_chunks():
+        for chunk_idx, (rows, plan) in enumerate(metadata.row_chunks()):
             logits = deep_gemm_fp4_paged_mqa_logits(
                 (q_fp4[rows], q_sf[rows]),
                 k_cache,
@@ -166,11 +161,15 @@ def _deep_gemm_prefill_captured(self, inputs: CapturedPrefillInputs) -> None:
                 plan,
                 width,
             )
-            topk_transform_paged_torch(
+            topk_transform_paged_from_metadata(
                 logits,
-                lens[rows],
-                page_table[rows],
-                inputs.out_page_indices[rows, :topk],
-                page_size,
-                inputs.out_raw_indices[rows, :topk],
+                metadata,
+                inputs.out_page_indices[:, :topk],
+                inputs.out_raw_indices[:, :topk],
+                rows=rows,
+                topk_metadata=(
+                    metadata.topk_plan_for_chunk(chunk_idx, rows)
+                    if metadata.use_topk_v2
+                    else None
+                ),
             )
diff --git a/python/sglang/srt/layers/moe/mega_moe.py b/python/sglang/srt/layers/moe/mega_moe.py
index 20d03d8cc7f7..6fc11e7880b3 100644
--- a/python/sglang/srt/layers/moe/mega_moe.py
+++ b/python/sglang/srt/layers/moe/mega_moe.py
@@ -99,12 +99,11 @@ def _configure_mega_moe_deep_gemm_num_sms(deep_gemm):
 
 def check_mega_moe_shapes(hidden: int, intermediate: int, mma_type: str) -> None:
     # SGLang's pre-dispatch requires 16-byte input SF rows (hidden % 512).
-    # The pinned DeepGEMM FP8/FP4 path copies intermediate scales separately
-    # from token TMA loads, so intermediate % 128 suffices, including V4.1
-    # Flash's 2304-wide intermediate dimension.
-    # Keep the existing restrictions on the other MMA paths.
+    # Retain upstream input/NVFP4 alignment, but also admit the verified W4A4
+    # MXFP4 Flash shape (5120, 2304). Both FP8/FP4 and MXFP4 use 128-column
+    # intermediate tiles; a 512-channel gate would reject the latter.
     scale_group = 16 if mma_type == "nvfp4xnvfp4" else 32
-    align = 128 if mma_type == "fp8xfp4" else 16 * scale_group
+    align = 128 if mma_type in ("fp8xfp4", "mxf4xmxf4") else 16 * scale_group
     hidden_align = 16 * scale_group
     if hidden % hidden_align != 0 or intermediate % align != 0:
         raise ValueError(
@@ -523,9 +522,7 @@ def build_mega_moe_experts_weights(experts) -> None:
         build_flydsl_mega_moe_weights(experts)
         return
 
-    from deep_gemm import (
-        transform_sf_into_required_layout,
-    )
+    from deep_gemm import transform_sf_into_required_layout
 
     if getattr(experts, "_mega_moe_weights_built", False):
         return
diff --git a/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py b/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
index 4c95615b3751..d3711bb2d230 100644
--- a/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
+++ b/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
@@ -12,7 +12,7 @@
 )
 from sglang.srt.disaggregation.decode import DecodePreallocQueue
 from sglang.srt.disaggregation.utils import get_dsv41_spec_layout
-from sglang.srt.mem_cache.common import retraction_backup
+from sglang.srt.mem_cache.common import backup_kv_cache
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.runtime_context import get_context
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -158,14 +158,13 @@ def test_retraction_recomputes_from_prefill_and_replays_boundary_token(self):
                 )
                 request_pool = Mock()
                 self.assertTrue(
-                    retraction_backup(
-                        req, Mock(), request_pool, allocator, "cpu_tensor"
-                    )
+                    backup_kv_cache(req, Mock(), request_pool, allocator, "cpu_tensor")
                 )
                 queue = SimpleNamespace(
                     token_to_kv_pool_allocator=allocator,
                     _check_if_req_exceed_kv_capacity=Mock(return_value=False),
                     _create_receiver_and_enqueue=Mock(),
+                    _init_receiver=Mock(),
                     _resolve_prefill_dp_rank=Mock(return_value=0),
                     retracted_queue=[],
                     pending_reqs=[],
diff --git a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py
index 5965ea59ba38..57adef60297a 100644
--- a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py
+++ b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py
@@ -480,10 +480,57 @@ def test_shape_check_uses_pinned_fp8_fp4_alignment(self):
             mega_moe.check_mega_moe_shapes(256, 128, "fp8xfp4")
         mega_moe.check_mega_moe_shapes(4096, 1536, "fp8xfp4")
         mega_moe.check_mega_moe_shapes(4096, 1024, "mxf4xmxf4")
+        mega_moe.check_mega_moe_shapes(5120, 2304, "mxf4xmxf4")
+        with self.assertRaisesRegex(ValueError, "intermediate_size.*multiple of 128"):
+            mega_moe.check_mega_moe_shapes(5120, 2240, "mxf4xmxf4")
         # 768 is a multiple of 256, so the NVFP4 (g16) rule accepts it.
         mega_moe.check_mega_moe_shapes(2048, 768, "nvfp4xnvfp4")
         with self.assertRaisesRegex(ValueError, "intermediate_size.*multiple of 256"):
             mega_moe.check_mega_moe_shapes(2048, 384, "nvfp4xnvfp4")
+    def test_dsv4_weight_prep_preserves_packed_byte_abi(self):
+        """DeepGEMM's FP4 payload uses int8, shared with the normal fallback."""
+        experts = SimpleNamespace()
+        for name, shape in (
+            ("w13_weight", (1, 256, 128)),
+            ("w2_weight", (1, 256, 64)),
+        ):
+            setattr(
+                experts,
+                name,
+                torch.nn.Parameter(
+                    torch.zeros(shape, dtype=torch.int8), requires_grad=False
+                ),
+            )
+        for name, shape in (
+            ("w13_weight_scale_inv", (1, 256, 8)),
+            ("w2_weight_scale_inv", (1, 256, 4)),
+        ):
+            setattr(
+                experts,
+                name,
+                torch.nn.Parameter(torch.ones(shape), requires_grad=False),
+            )
+        deep_gemm = self.deep_gemm
+        deep_gemm.transform_sf_into_required_layout = MagicMock(
+            side_effect=lambda _sf, mn, k, **kw: torch.zeros(
+                (kw["num_groups"], mn, k // 128), dtype=torch.int32
+            )
+        )
+        with (
+            patch.dict(sys.modules, {"deep_gemm": deep_gemm}),
+            patch.object(mega_moe, "_mega_moe_mma_type", return_value="mxf4xmxf4"),
+            # Isolate packed-byte ABI from the separate dimension gate above.
+            patch.object(mega_moe, "check_mega_moe_shapes"),
+        ):
+            mega_moe.build_mega_moe_experts_weights(experts)
+        for parameter, prepared in (
+            (experts.w13_weight, experts.mega_l1_weights[0]),
+            (experts.w2_weight, experts.mega_l2_weights[0]),
+        ):
+            self.assertEqual(parameter.dtype, torch.int8)
+            self.assertEqual(prepared.dtype, torch.int8)
+            self.assertEqual(parameter.data_ptr(), prepared.data_ptr())
+            torch.testing.assert_close(prepared.view(torch.int8), parameter)
 
     def test_mxf4_l1_uses_packed_gate_up_interleave(self):
         source = torch.arange(32).reshape(1, 32)
diff --git a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
index 2749ac769b65..9124d23b52f3 100644
--- a/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
+++ b/test/registered/unit/layers/test_dsv4_nonpaged_indexer.py
@@ -50,6 +50,7 @@ def test_dsv41_bcg_dense_indexer_keeps_paged_fallback_chunked(self):
 
         backend = SimpleNamespace(
             page_size=256,
+            dsa_topk_backend=DSATopKBackend.SGL_KERNEL,
             token_to_kv_pool=SimpleNamespace(get_index_k_page_size=lambda ratio: 64),
         )
         core = SimpleNamespace(
@@ -59,26 +60,32 @@ def test_dsv41_bcg_dense_indexer_keeps_paged_fallback_chunked(self):
         module = "sglang.srt.layers.attention.deepseek_v4_backend"
         with (
             patch(f"{module}._prefill_graph_max_seq_len") as max_context,
-            patch(f"{module}.PagedIndexerMetadata") as metadata_ctor,
+            patch(
+                f"{module}.expand_index_page_table",
+                side_effect=lambda table, **kwargs: table.repeat_interleave(4, dim=1),
+            ),
+            envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.override(True),
+            envs.SGLANG_OPT_USE_TOPK_V2.override(True),
+            patch(
+                "sglang.kernels.ops.attention.dsv4.plan_topk_v2",
+                side_effect=lambda lens: torch.zeros(
+                    (lens.numel() + 1, 2), dtype=torch.int32
+                ),
+            ),
         ):
-            with envs.SGLANG_DSV41_BCG_DENSE_INDEXER.override(False):
-                for context in (16384, 32768, 65536):
-                    max_context.return_value = context
-                    DeepseekV4AttnBackend._low_ratio_prefill_indexer_metadata(
-                        backend, core, 1
-                    )
-                    self.assertEqual(
-                        metadata_ctor.call_args.kwargs["row_chunk"],
-                        min(2048, (2048 * 16384) // context),
-                    )
-
-            with envs.SGLANG_DSV41_BCG_DENSE_INDEXER.override(True):
-                max_context.return_value = 32768
-                DeepseekV4AttnBackend._low_ratio_prefill_indexer_metadata(
-                    backend, core, 1
-                )
-                self.assertEqual(metadata_ctor.call_args.kwargs["row_chunk"], 1024)
-                self.assertFalse(metadata_ctor.call_args.kwargs["use_topk_v2"])
+            for dense in (False, True):
+                with envs.SGLANG_DSV41_BCG_DENSE_INDEXER.override(dense):
+                    for context in (16384, 32768, 65536):
+                        max_context.return_value = context
+                        metadata = (
+                            DeepseekV4AttnBackend._low_ratio_prefill_indexer_metadata(
+                                backend, core, 1
+                            )
+                        )
+                        self.assertEqual(metadata.row_chunk, 2048)
+                        self.assertEqual(metadata.max_compressed_seq_len, context)
+                        self.assertTrue(metadata.use_topk_v2)
+                        self.assertEqual(metadata.topk_metadata.shape, (2, 2049, 2))
 
     def test_dsv41_bcg_dense_offsets_use_live_request_order(self):
         from sglang.srt.layers.attention.deepseek_v4_backend import (
@@ -294,6 +301,228 @@ def make_metadata(lengths):
         torch.testing.assert_close(destination.topk_metadata, source.topk_metadata)
 
 
+class TestCapturedLowRatioIndexer(CustomTestCase):
+    def test_multistream_dense_k_preparation_keeps_small_cp_buckets(self):
+        from sglang.srt.layers.attention import deepseek_v4_backend as mod
+
+        gathered = torch.ones((16, 64), dtype=torch.int8)
+        metadata = SimpleNamespace(max_compressed_seq_len=16, compressed_page_size=64)
+        backend = SimpleNamespace(
+            req_to_token=None,
+            token_to_kv_pool=SimpleNamespace(
+                get_index_k_with_scale_buffer=lambda layer_id: None
+            ),
+            forward_metadata=SimpleNamespace(
+                c1_indexer_metadata=metadata,
+                prefill_graph_dense_indexer=True,
+                low_ratio_dense_req_indices=torch.tensor([7, 3, 0, 0, 0, 0, 0, 0]),
+                low_ratio_dense_seq_lens=torch.tensor([16, 16, 0, 0, 0, 0, 0, 0]),
+                low_ratio_local_req_indices=torch.tensor([3, 7, 3, 7]),
+                low_ratio_dense_k_offsets={},
+            ),
+        )
+        layer = SimpleNamespace(
+            compress_ratio=1,
+            layer_id=0,
+            indexer=SimpleNamespace(
+                candidate_topk_blocks=2,
+                candidate_block_size=8,
+                is_candidate_source=False,
+                uses_candidates=False,
+            ),
+        )
+        with patch.object(
+            mod, "gather_fp4_index_k_cache_masked", return_value=gathered
+        ):
+            k, offsets = mod.DeepseekV4AttnBackend._low_ratio_gather_k_prefill_graph(
+                backend, layer
+            )
+        self.assertIs(k, gathered)
+        torch.testing.assert_close(
+            offsets, torch.tensor([16, 0, 16, 0], dtype=torch.int32)
+        )
+
+    def test_paged_capture_writes_each_chunk_with_its_own_v2_plan(self):
+        from sglang.srt.layers.attention.dsv4.v41_indexer import CapturedPrefillInputs
+        from sglang.srt.layers.attention.dsv4.v41_indexer import full_topk as mod
+
+        lens = torch.tensor([0, 2, 5, 9, 16], dtype=torch.int32)
+        pages = torch.arange(10, 15, dtype=torch.int32).view(5, 1)
+        planner = SimpleNamespace(
+            get_num_sms=lambda: 1,
+            get_paged_mqa_logits_metadata=lambda lengths, *args: torch.zeros(
+                (1, 2), dtype=torch.int32
+            ),
+        )
+
+        def plan(lengths):
+            return torch.stack(
+                (
+                    torch.cat((lengths, lengths.new_tensor([-1]))),
+                    torch.zeros(lengths.numel() + 1, dtype=torch.int32),
+                ),
+                dim=1,
+            )
+
+        with (
+            patch.dict(sys.modules, {"deep_gemm": planner}),
+            envs.SGLANG_OPT_USE_JIT_INDEXER_METADATA.override(False),
+            envs.SGLANG_FP8_PAGED_MQA_LOGITS_TORCH.override(False),
+            patch("sglang.kernels.ops.attention.dsv4.plan_topk_v2", side_effect=plan),
+        ):
+            metadata = PagedIndexerMetadata(
+                page_size=256,
+                compressed_page_size=64,
+                page_table=pages,
+                compressed_seq_lens=lens,
+                use_topk_v2=True,
+                use_prefill_cuda_graph=True,
+                row_chunk=2,
+            )
+        inputs = CapturedPrefillInputs(
+            indexer=SimpleNamespace(n_heads=2, n_local_heads=2, index_topk=4),
+            layer_id=0,
+            q=torch.zeros((5, 2, 128)),
+            weights=torch.zeros((5, 2)),
+            paged_metadata=metadata,
+            out_page_indices=torch.full((5, 4), -99, dtype=torch.int32),
+            out_raw_indices=torch.full((5, 4), -99, dtype=torch.int32),
+        )
+
+        def topk(logits, lengths, table, out, page_size, topk_plan, raw):
+            # The GPU boundary is replaced, not the metadata dispatcher.
+            # A full-batch or wrong chunk plan must fail here, including the
+            # final one-row chunk.
+            torch.testing.assert_close(topk_plan[:-1, 0], lengths)
+            self.assertEqual(topk_plan.shape, (lengths.numel() + 1, 2))
+            out.fill_(-1)
+            raw.fill_(-1)
+            for row, length in enumerate(lengths.tolist()):
+                k = min(length, out.shape[1])
+                selected = logits[row, :length].topk(k).indices.to(torch.int32)
+                raw[row, :k] = selected
+                out[row, :k] = (
+                    table[row, selected.long() // page_size] * page_size
+                    + selected % page_size
+                )
+
+        with (
+            patch.object(
+                mod,
+                "quantize_index_q",
+                return_value=(
+                    torch.zeros((5, 2, 64), dtype=torch.int8),
+                    torch.zeros((5, 2), dtype=torch.int32),
+                ),
+            ),
+            patch.object(mod, "get_index_k_cache", return_value=None),
+            patch.object(
+                mod,
+                "deep_gemm_fp4_paged_mqa_logits",
+                side_effect=lambda q, *args: torch.arange(
+                    64, dtype=torch.float32
+                ).expand(q[0].shape[0], -1),
+            ),
+            patch(f"{_INDEXER}.topk_transform_paged_v2", side_effect=topk),
+        ):
+            indexer = mod.FullTopKIndexer(
+                token_to_kv_pool=None,
+                req_to_token=None,
+                use_deep_gemm_prefill=True,
+                use_deep_gemm_decode=True,
+            )
+            indexer.topk_prefill_captured(inputs)
+
+        for row, length in enumerate(lens.tolist()):
+            k = min(length, 4)
+            expected = torch.full((4,), -1, dtype=torch.int32)
+            expected[:k] = torch.arange(
+                length - 1, length - k - 1, -1, dtype=torch.int32
+            )
+            torch.testing.assert_close(inputs.out_raw_indices[row], expected)
+            expected_page = torch.where(
+                expected >= 0, pages[row, 0] * 64 + expected, -1
+            )
+            torch.testing.assert_close(inputs.out_page_indices[row], expected_page)
+
+    def test_dense_capture_is_selected_for_cp_local_rows_below_row_chunk(self):
+        from sglang.srt.layers.attention import deepseek_v4_backend as mod
+
+        lens = torch.tensor([8, 9, 10, 11], dtype=torch.int32)
+        metadata = SimpleNamespace(
+            max_compressed_seq_len=16,
+            compressed_seq_lens=lens,
+            compressed_page_size=64,
+            page_table=torch.zeros((4, 1), dtype=torch.int32),
+        )
+        out_page = torch.full((4, 4), -99, dtype=torch.int32)
+        out_raw = torch.full_like(out_page, -99)
+        requests = torch.tensor([0, 1, 0, 1])
+        offsets = requests.to(torch.int32) * 16
+        backend = SimpleNamespace(
+            req_to_token=torch.arange(32).view(2, 16) * 3,
+            token_to_kv_pool=SimpleNamespace(
+                get_index_k_with_scale_buffer=lambda layer_id: None
+            ),
+            forward_metadata=SimpleNamespace(
+                c1_indexer_metadata=metadata,
+                c2_indexer_metadata=None,
+                prefill_graph_dense_indexer=True,
+                candidate_metadata=None,
+                core_metadata=SimpleNamespace(
+                    sparse_page_indices=lambda ratio: out_page,
+                    sparse_raw_indices=lambda ratio: out_raw,
+                ),
+                low_ratio_dense_req_indices=torch.tensor([0, 1, 0, 0, 0, 0, 0, 0]),
+                low_ratio_dense_seq_lens=torch.tensor([16, 16, 0, 0, 0, 0, 0, 0]),
+                low_ratio_local_req_indices=requests,
+                low_ratio_dense_k_offsets={1: offsets},
+            ),
+        )
+        layer = SimpleNamespace(
+            compress_ratio=1,
+            layer_id=0,
+            indexer=SimpleNamespace(
+                candidate_topk_blocks=2,
+                candidate_block_size=8,
+                is_candidate_source=False,
+                uses_candidates=False,
+                index_topk=4,
+            ),
+        )
+
+        def select(logits, lengths, *, out_offsets, out_indices):
+            out_indices.copy_(
+                out_offsets[:, None] + lengths[:, None] - 1 - torch.arange(4)
+            )
+
+        with (
+            patch.object(
+                mod,
+                "quantize_index_q",
+                return_value=(
+                    torch.zeros((4, 2, 64), dtype=torch.int8),
+                    torch.zeros((4, 2), dtype=torch.int32),
+                ),
+            ),
+            patch.object(mod, "get_index_k_cache", return_value=None),
+            patch.object(mod, "gather_fp4_index_k_cache_masked", return_value=None),
+            patch.object(
+                mod, "_dense_fp4_mqa_logits", return_value=torch.zeros((4, 16))
+            ),
+            patch.object(mod, "topk_transform_ragged_v2", side_effect=select),
+        ):
+            mod.DeepseekV4AttnBackend._low_ratio_index_topk_captured(
+                backend, layer, torch.zeros((4, 2, 128)), torch.zeros((4, 2))
+            )
+        expected = lens[:, None] - 4 + torch.arange(4, dtype=torch.int32)
+        torch.testing.assert_close(out_raw, expected)
+        torch.testing.assert_close(
+            out_page,
+            backend.req_to_token[requests[:, None], expected.long()].to(torch.int32),
+        )
+
+
 class TestDSV4FlashInferTopK(CustomTestCase):
     def test_compact_page_transform_respects_fuse_topk(self):
         score_storage = torch.arange(160, dtype=torch.float32).reshape(2, 80)

From 565c272cc410d07606ae9f199de978f7063ea223 Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Fri, 9 Oct 2026 14:21:46 +0800
Subject: [PATCH 21/27] fix CP flag used in deepseek_v4

---
 python/sglang/srt/models/deepseek_v4.py | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index ed56968e4682..9148d5b87e78 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -58,10 +58,7 @@
 from sglang.srt.hardware_backend.npu.utils import (
     use_npu_arch35_mxfp8_wo_a,
 )
-from sglang.srt.layers.attention.dsa.utils import (
-    dsa_use_prefill_cp,
-    is_dsa_enable_prefill_cp,
-)
+from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
 from sglang.srt.layers.attention.dsv4.compressor import Compressor
 from sglang.srt.layers.attention.dsv4.dsv41_sparse import (
     DeepseekV41Compressor,
@@ -1514,7 +1511,7 @@ def _forward_prepare_multi_stream(
         stream_compressor = self.alt_streams[1]
         stream_indexer = self.alt_streams[2]
 
-        use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)
+        use_cp = is_cp_active(forward_batch)
         if not use_cp:
             stream_kv.wait_stream(current_stream)
             stream_compressor.wait_stream(current_stream)
@@ -1688,7 +1685,7 @@ def _forward_prepare_low_ratio_multi_stream(
         q_out: Optional[torch.Tensor] = None,
         x_quant=None,
     ) -> torch.Tensor:
-        if self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch):
+        if is_cp_active(forward_batch):
             return self._forward_prepare_low_ratio_cp_multi_stream(
                 x,
                 positions,
@@ -2336,9 +2333,7 @@ def forward(
         )
 
         unified = is_unified_kv_triton()
-        use_prefill_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(
-            forward_batch
-        )
+        use_prefill_cp = is_cp_active(forward_batch)
         enable_multi_stream = (
             envs.SGLANG_OPT_USE_MULTI_STREAM_OVERLAP.get()
             and self.alt_streams is not None

From 79f287dea0edeef123f4b95fea6a3013badfacbe Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Fri, 9 Oct 2026 19:02:21 +0800
Subject: [PATCH 22/27] current stream waits stream_kv after indexer being
 launched.

---
 .../layers/attention/deepseek_v4_backend.py   |  95 +++++++++----
 python/sglang/srt/models/deepseek_v4.py       |  41 +++---
 .../attention/dsv4/test_c2_prefill_ring.py    | 128 ++++++++++++++++++
 .../unit/models/test_deepseek_v41_text_cp.py  |  52 ++++++-
 4 files changed, 263 insertions(+), 53 deletions(-)

diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
index 44e78f766c25..ab474335d8a8 100644
--- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py
+++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py
@@ -1037,6 +1037,10 @@ class DSV4Metadata:
     # Shared by all low-ratio source layers; graph replay refreshes them live.
     low_ratio_req_indices: Optional[torch.Tensor] = None
     low_ratio_pos_i64: Optional[torch.Tensor] = None
+    # Global request offsets, not CP-local rows. A token bucket can replay with
+    # more requests than it captured, so keep the request axis pool-sized.
+    low_ratio_cp_extend_start_loc: Optional[torch.Tensor] = None
+    low_ratio_cp_extend_seq_lens: Optional[torch.Tensor] = None
     low_ratio_local_req_indices: Optional[torch.Tensor] = None
     low_ratio_dense_req_indices: Optional[torch.Tensor] = None
     low_ratio_dense_seq_lens: Optional[torch.Tensor] = None
@@ -1068,6 +1072,14 @@ def copy_(self, other: DSV4Metadata):
         maybe_copy_inplace(self.indexer_metadata, src=other.indexer_metadata)
         maybe_copy_inplace(self.c1_indexer_metadata, src=other.c1_indexer_metadata)
         maybe_copy_inplace(self.c2_indexer_metadata, src=other.c2_indexer_metadata)
+        maybe_copy_inplace(
+            self.low_ratio_cp_extend_start_loc,
+            src=other.low_ratio_cp_extend_start_loc,
+        )
+        maybe_copy_inplace(
+            self.low_ratio_cp_extend_seq_lens,
+            src=other.low_ratio_cp_extend_seq_lens,
+        )
         maybe_copy_inplace(self.c4_compress_metadata, src=other.c4_compress_metadata)
         maybe_copy_inplace(
             self.c128_compress_metadata, src=other.c128_compress_metadata
@@ -1092,6 +1104,14 @@ def refresh_for_breakable_cuda_graph_replay_(self, static_metadata: DSV4Metadata
         maybe_copy_inplace(
             self.low_ratio_pos_i64, src=static_metadata.low_ratio_pos_i64
         )
+        maybe_copy_inplace(
+            self.low_ratio_cp_extend_start_loc,
+            src=static_metadata.low_ratio_cp_extend_start_loc,
+        )
+        maybe_copy_inplace(
+            self.low_ratio_cp_extend_seq_lens,
+            src=static_metadata.low_ratio_cp_extend_seq_lens,
+        )
         if self.prefill_graph_dense_indexer:
             maybe_copy_inplace(
                 self.low_ratio_local_req_indices,
@@ -1588,6 +1608,13 @@ def create(compress_ratio: Literal[4, 128]):
                 else None
             )
             metadata.low_ratio_req_indices = req_pool_indices_repeated.to(torch.int64)
+            if cp_active and 2 in low:
+                (
+                    metadata.low_ratio_cp_extend_start_loc,
+                    metadata.low_ratio_cp_extend_seq_lens,
+                ) = self._low_ratio_cp_compressor_offsets(
+                    extend_start_loc, extend_seq_lens, self.req_to_token_pool.size
+                )
             metadata.low_ratio_pos_i64 = core_attn_metadata.positions_casual.to(
                 torch.int64
             )
@@ -2738,37 +2765,17 @@ def init_forward_metadata_for_breakable_cuda_graph_capture(
             max_seq_len_override=max_seq_len,
             use_prefill_cuda_graph=True,
         )
-        if self.low_ratio_prefill_graph and forward_batch.forward_mode.is_extend():
+        if (
+            self.low_ratio_prefill_graph
+            and forward_batch.forward_mode.is_extend()
+            and not dsa_use_prefill_cp(forward_batch)
+        ):
             for ratio in self.low_ratios:
-                if dsa_use_prefill_cp(forward_batch):
-                    self._cp_indexer_projection_buffers(
-                        forward_batch._cp_positions.shape[0], ratio
-                    )
-                else:
-                    self._source_projection_buffers(
-                        forward_batch.out_cache_loc.shape[0], ratio
-                    )
+                self._source_projection_buffers(
+                    forward_batch.out_cache_loc.shape[0], ratio
+                )
         return self.forward_metadata
 
-    def _cp_indexer_projection_buffers(self, num_tokens: int, ratio: int) -> dict:
-        """Capture-stable CP-local Q and weight rows for the paged indexer."""
-        sets = getattr(self, "_cp_indexer_proj_bufs", None)
-        if sets is None:
-            sets = self._cp_indexer_proj_bufs = {}
-        key = (ratio, num_tokens)
-        if key not in sets:
-            cfg = self.model_runner.model_config.hf_text_config
-            heads, dim = int(cfg.index_n_heads), int(cfg.index_head_dim)
-            sets[key] = {
-                "q": torch.zeros(
-                    num_tokens, heads, dim, dtype=torch.bfloat16, device=self.device
-                ),
-                "w": torch.zeros(
-                    num_tokens, heads, dtype=torch.bfloat16, device=self.device
-                ),
-            }
-        return sets[key]
-
     def _source_projection_buffers(self, num_tokens: int, ratio: int) -> dict:
         cfg = self.model_runner.model_config.hf_text_config
         heads, dim = int(cfg.index_n_heads), int(cfg.index_head_dim)
@@ -3003,7 +3010,24 @@ def _forward_low_ratio_sources_cp(
                     req_global = token_req_indices(forward_batch, num_tokens=total)
                     pos_global = forward_batch.positions[:total].to(torch.int64)
                 x_global = x_global[:total]
-            self._low_ratio_compress_torch(layer, x_global, req_global, pos_global)
+            metadata = self.forward_metadata
+            if (
+                self._low_ratio_in_prefill_graph()
+                and layer.compressor.use_fused_compress
+            ):
+                offsets = None
+                if layer.compress_ratio == 2:
+                    assert metadata.low_ratio_cp_extend_start_loc is not None
+                    assert metadata.low_ratio_cp_extend_seq_lens is not None
+                    offsets = (
+                        metadata.low_ratio_cp_extend_start_loc,
+                        metadata.low_ratio_cp_extend_seq_lens,
+                    )
+                self._low_ratio_compress_fused(
+                    layer, x_global, req_global, pos_global, extend_offsets=offsets
+                )
+            else:
+                self._low_ratio_compress_torch(layer, x_global, req_global, pos_global)
         if run_indexer and layer.indexer is not None:
             if self._low_ratio_in_prefill_graph():
                 q = layer.indexer.queries(q_lora, layer.freqs_cis[positions])
@@ -3084,6 +3108,19 @@ def _low_ratio_in_prefill_graph(self) -> bool:
 
         return self.low_ratio_prefill_graph and is_in_breakable_cuda_graph()
 
+    @staticmethod
+    def _low_ratio_cp_compressor_offsets(start_loc, seq_lens, capacity):
+        # The compact ring write-back launches over this static request axis.
+        # Zero lengths make unused entries no-ops, including after B shrinks.
+        assert start_loc is not None
+        count = start_loc.numel()
+        assert count <= capacity and seq_lens.numel() >= count
+        starts = start_loc.new_zeros(capacity, dtype=torch.int32)
+        lengths = seq_lens.new_zeros(capacity, dtype=torch.int32)
+        starts[:count].copy_(start_loc)
+        lengths[:count].copy_(seq_lens[:count])
+        return starts, lengths
+
     def _low_ratio_compress_decode(self, layer, x, req, pos) -> None:
         # Projection layout and fused-write support are fixed together at load time.
         if layer.compressor.use_fused_compress:
diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index 9148d5b87e78..182f8ab81423 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -39,9 +39,7 @@
     wo_a_bf16_small_batch,
     wo_a_bf16_small_batch_mxfp8,
 )
-from sglang.kernels.ops.quantization.fp8_kernel import (
-    sglang_per_token_group_quant_fp8,
-)
+from sglang.kernels.ops.quantization.fp8_kernel import sglang_per_token_group_quant_fp8
 from sglang.srt.compilation.compilation_config import register_split_op
 from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -55,9 +53,7 @@
     prime_rope_cos_sin,
     rope_cos_sin,
 )
-from sglang.srt.hardware_backend.npu.utils import (
-    use_npu_arch35_mxfp8_wo_a,
-)
+from sglang.srt.hardware_backend.npu.utils import use_npu_arch35_mxfp8_wo_a
 from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
 from sglang.srt.layers.attention.dsv4.compressor import Compressor
 from sglang.srt.layers.attention.dsv4.dsv41_sparse import (
@@ -1618,12 +1614,9 @@ def _forward_prepare_low_ratio_cp_multi_stream(
             )
 
         captured = is_in_breakable_cuda_graph()
-        indexer_buffers = None
-        if captured and self.indexer is not None:
+        capture_indexer = captured and self.indexer is not None
+        if capture_indexer:
             assert attn_backend.low_ratio_prefill_graph
-            indexer_buffers = attn_backend._cp_indexer_projection_buffers(
-                x.shape[0], self.compress_ratio
-            )
 
         stream_kv.wait_stream(current_stream)
         stream_compressor.wait_stream(current_stream)
@@ -1644,35 +1637,37 @@ def _forward_prepare_low_ratio_cp_multi_stream(
                     run_indexer=not captured,
                     precomputed_x_global=x_global,
                 )
-            if indexer_buffers is not None:
+            if capture_indexer:
                 # The compressor above writes this layer's index-K cache. Gather
                 # its live dense prefixes on the same stream after that write.
                 prepared_dense_k = attn_backend._low_ratio_gather_k_prefill_graph(self)
-        prepared_q = None
+        projected_q = projected_w = prepared_q = None
         with torch.cuda.stream(stream_indexer):
-            if indexer_buffers is not None:
+            if capture_indexer:
                 # Capture the full CP-local bucket; the paged metadata masks
                 # padding when the indexer runs after the worker join.
-                indexer_buffers["q"].copy_(
-                    self.indexer.queries(q_lora, self.freqs_cis[positions])
-                )
-                indexer_buffers["w"].copy_(self.indexer.head_weights(x))
+                # These projections are graph-owned intermediates, so consume
+                # them directly after the join without persistent staging copies.
+                projected_q = self.indexer.queries(q_lora, self.freqs_cis[positions])
+                projected_w = self.indexer.head_weights(x)
                 prepared_q = attn_backend._low_ratio_quantize_q_prefill_graph(
-                    indexer_buffers["q"]
+                    projected_q
                 )
 
         q = self._compute_q_b(q_for_wqb, positions, q_out)
-        current_stream.wait_stream(stream_kv)
         current_stream.wait_stream(stream_compressor)
         current_stream.wait_stream(stream_indexer)
-        if indexer_buffers is not None:
+        if capture_indexer:
             attn_backend._low_ratio_index_topk_captured(
                 self,
-                indexer_buffers["q"],
-                indexer_buffers["w"],
+                projected_q,
+                projected_w,
                 prepared_q=prepared_q,
                 prepared_dense_k=prepared_dense_k,
             )
+        # Indexer logits need index-K, not SWA KV. Join the SWA writer only at
+        # the attention boundary so it can overlap logits and top-k as well.
+        current_stream.wait_stream(stream_kv)
         del qkv_a
         return q
 
diff --git a/test/registered/kernels/ops/attention/dsv4/test_c2_prefill_ring.py b/test/registered/kernels/ops/attention/dsv4/test_c2_prefill_ring.py
index 598011ca314a..7074bfc045c2 100644
--- a/test/registered/kernels/ops/attention/dsv4/test_c2_prefill_ring.py
+++ b/test/registered/kernels/ops/attention/dsv4/test_c2_prefill_ring.py
@@ -76,6 +76,134 @@ def _row_tensors(rows: List[Tuple[int, int]], pad: int):
 
 
 class TestC2PrefillRing(CustomTestCase):
+    @unittest.skipUnless(_sm100(), "the c2 JIT kernels are the SM100 build")
+    def test_graph_replay_global_offsets_padding_and_ring_carry(self):
+        # CP materializes request-major global rows before compression. The
+        # request axis must stay fixed even when a token bucket replays B=1/5/1.
+        for layout in (KVLayout.V41, KVLayout.V41_FP4):
+            with self.subTest(layout=layout.name):
+                self._run_graph_replay(layout)
+
+    def _run_graph_replay(self, layout):
+        gen = torch.Generator(device="cuda").manual_seed(19)
+        bucket, capacity, ring_size = 32, 8, 8
+        inputs = _make_inputs(gen, bucket)
+        weight = torch.randn(
+            HEAD_DIM, generator=gen, device="cuda", dtype=torch.bfloat16
+        )
+        angles = torch.randn(128, ROPE_DIM // 2, generator=gen, device="cuda")
+        freqs = torch.view_as_real(
+            torch.polar(torch.ones_like(angles), angles)
+        ).flatten(-2)
+        state = _make_inputs(gen, capacity * ring_size + 1)
+        eager_state = state.clone()
+        cache = torch.zeros(
+            NUM_PAGES, layout.page_bytes(PAGE_SIZE), device="cuda", dtype=torch.uint8
+        )
+        eager_cache = cache.clone()
+        req = torch.zeros(bucket, device="cuda", dtype=torch.int64)
+        pos = torch.zeros(bucket, device="cuda", dtype=torch.int64)
+        loc = torch.zeros(bucket, device="cuda", dtype=torch.int64)
+        starts = torch.zeros(capacity, device="cuda", dtype=torch.int32)
+        lengths = torch.zeros_like(starts)
+        latent = torch.full(
+            (bucket, HEAD_DIM), 123, device="cuda", dtype=torch.bfloat16
+        )
+
+        def run(ring, dst_cache, offsets, out=None):
+            return c2_prefill_norm_rope_store(
+                inputs,
+                ring,
+                weight,
+                pos,
+                req,
+                loc,
+                EPS,
+                freqs,
+                dst_cache,
+                *offsets,
+                page_size=PAGE_SIZE,
+                ring_size=ring_size,
+                layout=layout,
+                out=out,
+            )
+
+        # Warm on a side stream before capturing; all rows initially pad.
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            run(state, cache, (starts, lengths), latent)
+        torch.cuda.current_stream().wait_stream(stream)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            run(state, cache, (starts, lengths), latent)
+
+        request_ids = [4, 1, 6, 2, 0]
+        next_pos = {4: 1, 1: 8, 6: 7, 2: 0, 0: 3}
+        for sizes in ([4], [9, 0, 3, 4, 2], [5], []):
+            rows, begin = [], []
+            for rid, n in zip(request_ids, sizes):
+                begin.append(len(rows))
+                rows.extend((rid, next_pos[rid] + i) for i in range(n))
+                next_pos[rid] += n
+            n = len(rows)
+            req.zero_()
+            pos.zero_()
+            loc.zero_()
+            starts.zero_()
+            lengths.zero_()
+            if n:
+                req[:n].copy_(torch.tensor([r for r, _ in rows], device="cuda"))
+                pos[:n].copy_(torch.tensor([p for _, p in rows], device="cuda"))
+                loc[:n].copy_(
+                    torch.tensor(
+                        [SLOT_BASE * (1 + r) + p for r, p in rows], device="cuda"
+                    )
+                )
+            m = len(sizes)
+            starts[:m].copy_(torch.tensor(begin, device="cuda", dtype=torch.int32))
+            lengths[:m].copy_(torch.tensor(sizes, device="cuda", dtype=torch.int32))
+            inputs.copy_(_make_inputs(gen, bucket))
+            before = eager_state.clone()
+            latent.fill_(123)
+            graph.replay()
+            eager = run(eager_state, eager_cache, (starts, lengths))
+            self.assertTrue(torch.equal(cache, eager_cache))
+            self.assertTrue(torch.equal(state, eager_state))
+            odd = [i for i, (_, p) in enumerate(rows) if p % 2]
+            self.assertTrue(torch.equal(latent[odd], eager[odd]))
+            untouched = [i for i in range(bucket) if i not in odd]
+            self.assertTrue(torch.all(latent[untouched] == 123).item())
+
+            # Independently check Torch's pair selection/softmax/finish, not
+            # only another invocation of the fused kernel. FP32 reduction order
+            # can cross a bf16 rounding boundary, so do not require byte equality.
+            for i in odd:
+                rid, p = rows[i]
+                previous = (
+                    inputs[i - 1]
+                    if i and rows[i - 1] == (rid, p - 1)
+                    else before[rid * ring_size + (p - 1) % ring_size]
+                )
+                pair = torch.stack((previous, inputs[i]))
+                kv, score = pair[:, :HEAD_DIM], pair[:, HEAD_DIM:]
+                pooled = (kv * score.softmax(0)).sum(0).to(torch.bfloat16).float()
+                expected = (
+                    pooled * torch.rsqrt(pooled.square().mean() + EPS) * weight.float()
+                ).to(torch.bfloat16)
+                torch.testing.assert_close(latent[i], expected, rtol=0.016, atol=0.016)
+            # The ring contains exactly each request's most recent rows, even
+            # when a request has more new rows than its ring capacity.
+            for start, size in zip(begin, sizes):
+                for i in range(start + max(0, size - ring_size), start + size):
+                    rid, p = rows[i]
+                    torch.testing.assert_close(
+                        state[rid * ring_size + p % ring_size],
+                        inputs[i],
+                        rtol=0,
+                        atol=0,
+                    )
+
     def _run(self, stream: Stream, layout: KVLayout, ring_size: int, seed: int):
         """The fused chunked path and the row-at-a-time decode path over the same
         stream; returns (cache, latents) for each."""
diff --git a/test/registered/unit/models/test_deepseek_v41_text_cp.py b/test/registered/unit/models/test_deepseek_v41_text_cp.py
index 7689b118c01b..d781f704b626 100644
--- a/test/registered/unit/models/test_deepseek_v41_text_cp.py
+++ b/test/registered/unit/models/test_deepseek_v41_text_cp.py
@@ -1,6 +1,7 @@
 """Pure-language V4.1 CP input, padding and DSpark state regressions."""
 
 import unittest
+from contextlib import nullcontext
 from types import SimpleNamespace as NS
 from unittest.mock import Mock, patch
 
@@ -12,7 +13,7 @@
     is_cp_active,
 )
 from sglang.srt.model_executor.runner.eager_runner import EagerRunner
-from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM
+from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM, MQALayer
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.dsv41_cp_test_utils import cp_context, simulated_collective
 from sglang.test.test_utils import CustomTestCase
@@ -24,6 +25,55 @@
 
 
 class TestDSV41TextCP(CustomTestCase):
+    def test_cp_multistream_joins_swa_after_indexer_before_attention(self):
+        # SWA KV is not an input of logits/top-k, but must be complete before
+        # attention starts. A moved/removed join silently changes that boundary.
+        order = []
+        worker_streams = [object(), object(), object()]
+        parent = NS(
+            wait_stream=lambda stream: order.append(worker_streams.index(stream))
+        )
+        workers = [NS(wait_stream=lambda stream: None) for _ in range(3)]
+        worker_streams[:] = workers
+        x = torch.zeros(4, 2)
+        positions = torch.arange(4)
+        indexer = NS(
+            queries=lambda q, freqs: x,
+            head_weights=lambda hidden: x,
+        )
+        layer = NS(
+            alt_streams=workers,
+            fuse_wqa_wkv=False,
+            compressor=object(),
+            indexer=indexer,
+            freqs_cis=torch.zeros(4),
+            _compute_q_a=lambda *args, **kwargs: (x, x),
+            _materialize_cp_swa_k=lambda *args, **kwargs: x,
+            _store_cp_swa_k=lambda *args: None,
+            _compute_q_b=lambda *args: x,
+        )
+        backend = NS(
+            low_ratio_prefill_graph=True,
+            forward_low_ratio_sources=lambda **kwargs: None,
+            _low_ratio_gather_k_prefill_graph=lambda layer: None,
+            _low_ratio_quantize_q_prefill_graph=lambda q: None,
+            _low_ratio_index_topk_captured=lambda *args, **kwargs: order.append("topk"),
+        )
+        module = "sglang.srt.models.deepseek_v4"
+        with (
+            patch(module + ".torch.cuda.current_stream", return_value=parent),
+            patch(
+                module + ".torch.cuda.stream", side_effect=lambda stream: nullcontext()
+            ),
+            patch(module + ".cp_materialize_global_token_order", return_value=x),
+            patch(module + ".is_in_breakable_cuda_graph", return_value=True),
+        ):
+            result = MQALayer._forward_prepare_low_ratio_cp_multi_stream(
+                layer, x, positions, NS(encoder_swa_replay=False), backend
+            )
+        self.assertIs(result, x)
+        self.assertEqual(order, [1, 2, "topk", 0])
+
     def test_interleave_roundtrip_mixed_lengths_prefix_and_padding(self):
         for size in (2, 4):
             for length in (4, 5, 9, 127, 128, 129):

From b4b972b92e7911fe6bb8ac9ef916bc8928db9549 Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Sat, 10 Oct 2026 13:46:10 +0800
Subject: [PATCH 23/27] perf: overlap DSV4.1 CP Q projection with allgather

---
 python/sglang/srt/models/deepseek_v4.py       |  15 +-
 .../unit/models/test_deepseek_v41_text_cp.py  | 156 +++++++++++++++++-
 2 files changed, 165 insertions(+), 6 deletions(-)

diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index 182f8ab81423..822049e65bad 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -1592,7 +1592,7 @@ def _forward_prepare_low_ratio_cp_multi_stream(
         q_out: Optional[torch.Tensor] = None,
         x_quant=None,
     ) -> torch.Tensor:
-        """Overlap CP SWA/cache sources with the Q projection after allgather."""
+        """Overlap Q with CP gathers, then prepare cache/indexer sources."""
         assert self.alt_streams is not None
         assert len(self.alt_streams) >= 3
         current_stream = torch.cuda.current_stream()
@@ -1604,8 +1604,14 @@ def _forward_prepare_low_ratio_cp_multi_stream(
             qkv_a, _ = self.wqkv_a(x_linear)
         q_lora, q_for_wqb = self._compute_q_a(x_linear, qkv_a=qkv_a)
 
+        # Q needs only the local q_lora. Start it before the gathers, rather
+        # than alongside the compressor GEMM, on the existing SWA writer stream.
+        stream_kv.wait_stream(current_stream)
+        with torch.cuda.stream(stream_kv):
+            q = self._compute_q_b(q_for_wqb, positions, q_out)
+
         # Collectives must launch in identical parent-stream order on every CP
-        # rank. Neither worker may start consuming until both gathers finish.
+        # rank. Workers that read gathered inputs wait for both gathers below.
         swa_k = self._materialize_cp_swa_k(x_linear, forward_batch, qkv_a=qkv_a)
         x_global = None
         if self.compressor is not None and not forward_batch.encoder_swa_replay:
@@ -1654,7 +1660,6 @@ def _forward_prepare_low_ratio_cp_multi_stream(
                     projected_q
                 )
 
-        q = self._compute_q_b(q_for_wqb, positions, q_out)
         current_stream.wait_stream(stream_compressor)
         current_stream.wait_stream(stream_indexer)
         if capture_indexer:
@@ -1665,8 +1670,8 @@ def _forward_prepare_low_ratio_cp_multi_stream(
                 prepared_q=prepared_q,
                 prepared_dense_k=prepared_dense_k,
             )
-        # Indexer logits need index-K, not SWA KV. Join the SWA writer only at
-        # the attention boundary so it can overlap logits and top-k as well.
+        # Indexer logits need index-K, not attention Q or SWA KV. Join the Q/SWA
+        # stream at the attention boundary so it can overlap logits/top-k too.
         current_stream.wait_stream(stream_kv)
         del qkv_a
         return q
diff --git a/test/registered/unit/models/test_deepseek_v41_text_cp.py b/test/registered/unit/models/test_deepseek_v41_text_cp.py
index d781f704b626..97fccf75b9ea 100644
--- a/test/registered/unit/models/test_deepseek_v41_text_cp.py
+++ b/test/registered/unit/models/test_deepseek_v41_text_cp.py
@@ -1,7 +1,8 @@
 """Pure-language V4.1 CP input, padding and DSpark state regressions."""
 
 import unittest
-from contextlib import nullcontext
+from contextlib import contextmanager, nullcontext
+from itertools import product
 from types import SimpleNamespace as NS
 from unittest.mock import Mock, patch
 
@@ -25,6 +26,159 @@
 
 
 class TestDSV41TextCP(CustomTestCase):
+    def test_cp_multistream_q_starts_before_gathers_and_is_joined(self):
+        """Q must depend on q_lora, not the CP gathers or compressor. The
+        cache writer must still wait for gathered KV, and attention must see
+        completed Q. Emulate stream frontiers to detect missing dependencies.
+        """
+        module = "sglang.srt.models.deepseek_v4"
+        for captured, compressor, fused, encoder_replay, has_q_out in product(
+            (False, True), repeat=5
+        ):
+            with self.subTest(
+                captured=captured,
+                compressor=compressor,
+                fused=fused,
+                encoder_replay=encoder_replay,
+                has_q_out=has_q_out,
+            ):
+                self._check_early_cp_q(
+                    module, captured, compressor, fused, encoder_replay, has_q_out
+                )
+
+    def _check_early_cp_q(
+        self, module, captured, compressor, fused, encoder_replay, has_q_out
+    ):
+        class Stream:
+            def __init__(self):
+                self.ready = set()
+
+            def wait_stream(self, source):
+                self.ready.update(source.ready)
+
+        parent = Stream()
+        workers = [Stream() for _ in range(3)]
+        active = [parent]
+        events = {}
+
+        @contextmanager
+        def use_stream(stream):
+            previous, active[0] = active[0], stream
+            try:
+                yield
+            finally:
+                active[0] = previous
+
+        def enqueue(name, value, *inputs):
+            ready = active[0].ready
+            for operand in inputs:
+                self.assertIn(operand.name, ready, f"{name} reads unfinished input")
+            events[name] = (active[0], ready.copy())
+            ready.add(name)
+            return NS(name=name, value=value)
+
+        x = NS(name="input", value=2)
+        x.contiguous = lambda: x
+        parent.ready.add(x.name)
+        positions = 0
+        q_out = NS(value=-1) if has_q_out else None
+
+        def compute_q_a(hidden, qkv_a=None):
+            q = enqueue("q_a", 2 * hidden.value, qkv_a or hidden)
+            return q, q
+
+        def compute_q_b(q, pos, out):
+            self.assertIs(pos, positions)
+            self.assertIs(out, q_out)
+            result = enqueue("q_b", 3 * q.value, q)
+            if out is None:
+                return result
+            out.name, out.value = result.name, result.value
+            return out
+
+        def sources(**kwargs):
+            inputs = [kwargs["precomputed_x_global"] or x]
+            if kwargs["q_lora"] is not None:
+                inputs.append(kwargs["q_lora"])
+            return enqueue("sources", 5, *inputs)
+
+        def gather_index_k(layer):
+            # With no compressor this layer borrows an already-written cache.
+            inputs = (NS(name="sources"),) if compressor else ()
+            return enqueue("index_k", 7, *inputs)
+
+        def topk(layer, q, w, *, prepared_q, prepared_dense_k):
+            enqueue("topk", 0, prepared_q, prepared_dense_k, w)
+
+        layer = NS(
+            alt_streams=workers,
+            fuse_wqa_wkv=fused,
+            wqkv_a=lambda hidden: (enqueue("qkv_a", hidden.value, hidden), None),
+            compressor=object() if compressor else None,
+            indexer=NS(
+                queries=lambda q, freqs: enqueue("index_q", q.value, q),
+                head_weights=lambda hidden: enqueue("index_w", hidden.value, hidden),
+            ),
+            freqs_cis=(0,),
+            _compute_q_a=compute_q_a,
+            _compute_q_b=compute_q_b,
+            _materialize_cp_swa_k=lambda hidden, *args, **kw: enqueue(
+                "swa_gather", hidden.value, hidden
+            ),
+            _store_cp_swa_k=lambda kv, *args: enqueue("swa_store", kv.value, kv),
+        )
+        backend = NS(
+            low_ratio_prefill_graph=True,
+            forward_low_ratio_sources=sources,
+            _low_ratio_gather_k_prefill_graph=gather_index_k,
+            _low_ratio_quantize_q_prefill_graph=lambda q: enqueue(
+                "index_quant", q.value, q
+            ),
+            _low_ratio_index_topk_captured=topk,
+        )
+        with (
+            patch(module + ".torch.cuda.current_stream", side_effect=lambda: active[0]),
+            patch(module + ".torch.cuda.stream", side_effect=use_stream),
+            patch(
+                module + ".cp_materialize_global_token_order",
+                side_effect=lambda hidden, *args: enqueue(
+                    "main_gather", hidden.value, hidden
+                ),
+            ),
+            patch(module + ".is_in_breakable_cuda_graph", return_value=captured),
+        ):
+            result = MQALayer._forward_prepare_low_ratio_cp_multi_stream(
+                layer,
+                x,
+                positions,
+                NS(encoder_swa_replay=encoder_replay),
+                backend,
+                q_out,
+            )
+
+        self.assertEqual(result.value, 12)
+        if has_q_out:
+            self.assertIs(result, q_out)
+        self.assertIn(result.name, parent.ready, "attention reads unfinished Q")
+        q_stream, q_dependencies = events["q_b"]
+        self.assertIs(q_stream, workers[0])
+        self.assertIn("q_a", q_dependencies)
+        self.assertTrue(
+            q_dependencies.isdisjoint({"swa_gather", "main_gather", "sources"})
+        )
+        self.assertIs(events["swa_gather"][0], parent)
+        store_stream, store_dependencies = events["swa_store"]
+        self.assertIs(store_stream, workers[0])
+        self.assertTrue({"swa_gather", "q_b"}.issubset(store_dependencies))
+        if compressor and not encoder_replay:
+            self.assertIs(events["main_gather"][0], parent)
+            self.assertIn("swa_gather", events["main_gather"][1])
+            self.assertIn("main_gather", events["sources"][1])
+        else:
+            self.assertNotIn("main_gather", events)
+        if captured:
+            self.assertNotIn("swa_store", events["topk"][1])
+
     def test_cp_multistream_joins_swa_after_indexer_before_attention(self):
         # SWA KV is not an input of logits/top-k, but must be complete before
         # attention starts. A moved/removed join silently changes that boundary.

From 0bf0ffe2462b1063243973b24ac909dcb7d8ec2c Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Sat, 10 Oct 2026 14:59:04 +0800
Subject: [PATCH 24/27] perf: submit DSV4.1 CP allgathers before Q projection

---
 python/sglang/srt/models/deepseek_v4.py            | 14 +++++++++-----
 .../unit/models/test_deepseek_v41_text_cp.py       | 14 ++++++++++----
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index 822049e65bad..0c7525816ac6 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -1592,7 +1592,7 @@ def _forward_prepare_low_ratio_cp_multi_stream(
         q_out: Optional[torch.Tensor] = None,
         x_quant=None,
     ) -> torch.Tensor:
-        """Overlap Q with CP gathers, then prepare cache/indexer sources."""
+        """Submit CP gathers before Q, then prepare cache/indexer sources."""
         assert self.alt_streams is not None
         assert len(self.alt_streams) >= 3
         current_stream = torch.cuda.current_stream()
@@ -1604,11 +1604,9 @@ def _forward_prepare_low_ratio_cp_multi_stream(
             qkv_a, _ = self.wqkv_a(x_linear)
         q_lora, q_for_wqb = self._compute_q_a(x_linear, qkv_a=qkv_a)
 
-        # Q needs only the local q_lora. Start it before the gathers, rather
-        # than alongside the compressor GEMM, on the existing SWA writer stream.
+        # Snapshot q_lora readiness before the gathers. Waiting here rather than
+        # at the Q launch below keeps Q independent of collective completion.
         stream_kv.wait_stream(current_stream)
-        with torch.cuda.stream(stream_kv):
-            q = self._compute_q_b(q_for_wqb, positions, q_out)
 
         # Collectives must launch in identical parent-stream order on every CP
         # rank. Workers that read gathered inputs wait for both gathers below.
@@ -1619,6 +1617,12 @@ def _forward_prepare_low_ratio_cp_multi_stream(
                 x.contiguous(), forward_batch, current_stream
             )
 
+        # Submit the long-running gathers before the SM-heavy Q GEMM. Q still
+        # waits only for q_lora and runs on the existing SWA writer stream, not
+        # the collective stream.
+        with torch.cuda.stream(stream_kv):
+            q = self._compute_q_b(q_for_wqb, positions, q_out)
+
         captured = is_in_breakable_cuda_graph()
         capture_indexer = captured and self.indexer is not None
         if capture_indexer:
diff --git a/test/registered/unit/models/test_deepseek_v41_text_cp.py b/test/registered/unit/models/test_deepseek_v41_text_cp.py
index 97fccf75b9ea..d153522ee1e1 100644
--- a/test/registered/unit/models/test_deepseek_v41_text_cp.py
+++ b/test/registered/unit/models/test_deepseek_v41_text_cp.py
@@ -26,10 +26,11 @@
 
 
 class TestDSV41TextCP(CustomTestCase):
-    def test_cp_multistream_q_starts_before_gathers_and_is_joined(self):
-        """Q must depend on q_lora, not the CP gathers or compressor. The
-        cache writer must still wait for gathered KV, and attention must see
-        completed Q. Emulate stream frontiers to detect missing dependencies.
+    def test_cp_multistream_gathers_submit_before_q_without_serializing(self):
+        """An early Q launch can starve CP collectives of SM resources. Submit
+        gathers first without making Q wait for them to finish. The cache writer
+        must still wait for KV, and attention must see completed Q. Emulate
+        stream frontiers to detect accidental collective/Q dependencies.
         """
         module = "sglang.srt.models.deepseek_v4"
         for captured, compressor, fused, encoder_replay, has_q_out in product(
@@ -88,6 +89,9 @@ def compute_q_a(hidden, qkv_a=None):
             return q, q
 
         def compute_q_b(q, pos, out):
+            self.assertIn("swa_gather", events, "Q submitted before SWA gather")
+            if compressor and not encoder_replay:
+                self.assertIn("main_gather", events, "Q submitted before main gather")
             self.assertIs(pos, positions)
             self.assertIs(out, q_out)
             result = enqueue("q_b", 3 * q.value, q)
@@ -167,12 +171,14 @@ def topk(layer, q, w, *, prepared_q, prepared_dense_k):
             q_dependencies.isdisjoint({"swa_gather", "main_gather", "sources"})
         )
         self.assertIs(events["swa_gather"][0], parent)
+        self.assertNotIn("q_b", events["swa_gather"][1])
         store_stream, store_dependencies = events["swa_store"]
         self.assertIs(store_stream, workers[0])
         self.assertTrue({"swa_gather", "q_b"}.issubset(store_dependencies))
         if compressor and not encoder_replay:
             self.assertIs(events["main_gather"][0], parent)
             self.assertIn("swa_gather", events["main_gather"][1])
+            self.assertNotIn("q_b", events["main_gather"][1])
             self.assertIn("main_gather", events["sources"][1])
         else:
             self.assertNotIn("main_gather", events)

From 15aaa88762a86b4fbc3023b9f42320097768c43c Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Sat, 10 Oct 2026 17:27:27 +0800
Subject: [PATCH 25/27] remove stream_indexer and replace it with current
 stream

---
 python/sglang/srt/models/deepseek_v4.py       | 33 ++++++++-----------
 .../unit/models/test_deepseek_v41_text_cp.py  | 31 +++++++++++++----
 2 files changed, 39 insertions(+), 25 deletions(-)

diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index 0c7525816ac6..7852d6595957 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -1592,11 +1592,11 @@ def _forward_prepare_low_ratio_cp_multi_stream(
         q_out: Optional[torch.Tensor] = None,
         x_quant=None,
     ) -> torch.Tensor:
-        """Submit CP gathers before Q, then prepare cache/indexer sources."""
+        """Prepare indexer queries on the parent alongside CP source workers."""
         assert self.alt_streams is not None
-        assert len(self.alt_streams) >= 3
+        assert len(self.alt_streams) >= 2
         current_stream = torch.cuda.current_stream()
-        stream_kv, stream_compressor, stream_indexer = self.alt_streams[:3]
+        stream_kv, stream_compressor = self.alt_streams[:2]
         x_linear = x_quant if x_quant is not None else x
 
         qkv_a: Optional[torch.Tensor] = None
@@ -1617,9 +1617,9 @@ def _forward_prepare_low_ratio_cp_multi_stream(
                 x.contiguous(), forward_batch, current_stream
             )
 
-        # Submit the long-running gathers before the SM-heavy Q GEMM. Q still
-        # waits only for q_lora and runs on the existing SWA writer stream, not
-        # the collective stream.
+        # Register gathers before Q during capture; replay can schedule these
+        # independent branches in either order. Q waits only for q_lora and
+        # runs on the SWA writer stream, not the collective stream.
         with torch.cuda.stream(stream_kv):
             q = self._compute_q_b(q_for_wqb, positions, q_out)
 
@@ -1630,7 +1630,6 @@ def _forward_prepare_low_ratio_cp_multi_stream(
 
         stream_kv.wait_stream(current_stream)
         stream_compressor.wait_stream(current_stream)
-        stream_indexer.wait_stream(current_stream)
         with torch.cuda.stream(stream_kv):
             self._store_cp_swa_k(swa_k, forward_batch, attn_backend)
         prepared_dense_k = None
@@ -1652,20 +1651,16 @@ def _forward_prepare_low_ratio_cp_multi_stream(
                 # its live dense prefixes on the same stream after that write.
                 prepared_dense_k = attn_backend._low_ratio_gather_k_prefill_graph(self)
         projected_q = projected_w = prepared_q = None
-        with torch.cuda.stream(stream_indexer):
-            if capture_indexer:
-                # Capture the full CP-local bucket; the paged metadata masks
-                # padding when the indexer runs after the worker join.
-                # These projections are graph-owned intermediates, so consume
-                # them directly after the join without persistent staging copies.
-                projected_q = self.indexer.queries(q_lora, self.freqs_cis[positions])
-                projected_w = self.indexer.head_weights(x)
-                prepared_q = attn_backend._low_ratio_quantize_q_prefill_graph(
-                    projected_q
-                )
+        if capture_indexer:
+            # Q/W and Q quantization need only local inputs, not compressed K or
+            # SWA KV. Keep them on the parent before joining the compressor so
+            # they overlap both workers without a separate indexer fork/join.
+            # The full CP-local bucket is graph-owned; metadata masks padding.
+            projected_q = self.indexer.queries(q_lora, self.freqs_cis[positions])
+            projected_w = self.indexer.head_weights(x)
+            prepared_q = attn_backend._low_ratio_quantize_q_prefill_graph(projected_q)
 
         current_stream.wait_stream(stream_compressor)
-        current_stream.wait_stream(stream_indexer)
         if capture_indexer:
             attn_backend._low_ratio_index_topk_captured(
                 self,
diff --git a/test/registered/unit/models/test_deepseek_v41_text_cp.py b/test/registered/unit/models/test_deepseek_v41_text_cp.py
index d153522ee1e1..1b68ace34318 100644
--- a/test/registered/unit/models/test_deepseek_v41_text_cp.py
+++ b/test/registered/unit/models/test_deepseek_v41_text_cp.py
@@ -27,10 +27,10 @@
 
 class TestDSV41TextCP(CustomTestCase):
     def test_cp_multistream_gathers_submit_before_q_without_serializing(self):
-        """An early Q launch can starve CP collectives of SM resources. Submit
-        gathers first without making Q wait for them to finish. The cache writer
-        must still wait for KV, and attention must see completed Q. Emulate
-        stream frontiers to detect accidental collective/Q dependencies.
+        """Gathers and Q must stay independent, while indexer query preparation
+        must not wait for compressed K. Emulate stream frontiers to detect a
+        premature compressor/SWA join, an unfinished top-k input or missing Q.
+        Submission order alone does not guarantee GPU execution order.
         """
         module = "sglang.srt.models.deepseek_v4"
         for captured, compressor, fused, encoder_replay, has_q_out in product(
@@ -112,7 +112,13 @@ def gather_index_k(layer):
             return enqueue("index_k", 7, *inputs)
 
         def topk(layer, q, w, *, prepared_q, prepared_dense_k):
-            enqueue("topk", 0, prepared_q, prepared_dense_k, w)
+            backend.topk_result = enqueue(
+                "topk",
+                prepared_q.value + prepared_dense_k.value + w.value,
+                prepared_q,
+                prepared_dense_k,
+                w,
+            )
 
         layer = NS(
             alt_streams=workers,
@@ -183,7 +189,20 @@ def topk(layer, q, w, *, prepared_q, prepared_dense_k):
         else:
             self.assertNotIn("main_gather", events)
         if captured:
+            for name in ("index_q", "index_w", "index_quant"):
+                stream, dependencies = events[name]
+                self.assertIs(stream, parent)
+                self.assertTrue(
+                    dependencies.isdisjoint({"q_b", "swa_store", "sources", "index_k"}),
+                    "Indexer query preparation waits for unrelated worker inputs",
+                )
+            self.assertEqual(backend.topk_result.value, 13)
             self.assertNotIn("swa_store", events["topk"][1])
+            if compressor:
+                self.assertIs(events["sources"][0], workers[1])
+        self.assertFalse(
+            workers[2].ready, "Unused indexer worker acquired dependencies"
+        )
 
     def test_cp_multistream_joins_swa_after_indexer_before_attention(self):
         # SWA KV is not an input of logits/top-k, but must be complete before
@@ -232,7 +251,7 @@ def test_cp_multistream_joins_swa_after_indexer_before_attention(self):
                 layer, x, positions, NS(encoder_swa_replay=False), backend
             )
         self.assertIs(result, x)
-        self.assertEqual(order, [1, 2, "topk", 0])
+        self.assertEqual(order, [1, "topk", 0])
 
     def test_interleave_roundtrip_mixed_lengths_prefix_and_padding(self):
         for size in (2, 4):

From a31abfc7974ae9248fcd2c94939cd7c65defc933 Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Tue, 13 Oct 2026 12:59:34 +0800
Subject: [PATCH 26/27] chore(dsv41): lint multistream rebase resolutions

---
 python/sglang/srt/models/deepseek_v4.py          | 16 ++++++----------
 .../layers/moe/test_mega_moe_deepgemm_api.py     |  1 +
 2 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index 7852d6595957..71c2b258cbe3 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -181,7 +181,6 @@
 )
 from sglang.srt.runtime_context import (
     get_device,
-    get_disagg,
     get_exec,
     get_forward,
     get_parallel,
@@ -4767,15 +4766,12 @@ def __init__(
             and config.vision_n_layers > 0
             and not getattr(config, "language_model_only", False)
         ):
-            if (
-                get_parallel().pp_group.world_size != 1
-                or not (
-                    get_moe_a2a_backend().is_none()
-                    or (
-                        _is_cuda
-                        and get_moe_a2a_backend().is_megamoe()
-                        and get_parallel().attn_cp_size == 1
-                    )
+            if get_parallel().pp_group.world_size != 1 or not (
+                get_moe_a2a_backend().is_none()
+                or (
+                    _is_cuda
+                    and get_moe_a2a_backend().is_megamoe()
+                    and get_parallel().attn_cp_size == 1
                 )
             ):
                 raise ValueError(
diff --git a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py
index 57adef60297a..9da4687b2f45 100644
--- a/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py
+++ b/test/registered/unit/layers/moe/test_mega_moe_deepgemm_api.py
@@ -487,6 +487,7 @@ def test_shape_check_uses_pinned_fp8_fp4_alignment(self):
         mega_moe.check_mega_moe_shapes(2048, 768, "nvfp4xnvfp4")
         with self.assertRaisesRegex(ValueError, "intermediate_size.*multiple of 256"):
             mega_moe.check_mega_moe_shapes(2048, 384, "nvfp4xnvfp4")
+
     def test_dsv4_weight_prep_preserves_packed_byte_abi(self):
         """DeepGEMM's FP4 payload uses int8, shared with the normal fallback."""
         experts = SimpleNamespace()

From e44e29a3ac72e413b5b5d79c499783a0f5a78f14 Mon Sep 17 00:00:00 2001
From: SYChen123 <csy512356@ant-intl.com>
Date: Tue, 13 Oct 2026 14:03:38 +0800
Subject: [PATCH 27/27] test(dsv41): align PD fixtures with current cache
 contracts

---
 .../disaggregation/test_dsv41_dspark_pd.py    | 33 ++++++--------
 .../test_fake_transfer_chunk_progress.py      | 45 +++++++++----------
 2 files changed, 35 insertions(+), 43 deletions(-)

diff --git a/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py b/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
index d3711bb2d230..6f7206a334aa 100644
--- a/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
+++ b/test/registered/unit/disaggregation/test_dsv41_dspark_pd.py
@@ -153,6 +153,7 @@ def test_retraction_recomputes_from_prefill_and_replays_boundary_token(self):
                 req = SimpleNamespace(
                     output_ids=[7, 8],
                     bootstrap_host="prefill",
+                    multimodal_inputs=None,
                     time_stats=Mock(),
                     offload_kv_cache=Mock(),
                 )
@@ -170,25 +171,16 @@ def test_retraction_recomputes_from_prefill_and_replays_boundary_token(self):
                     pending_reqs=[],
                 )
                 DecodePreallocQueue.add(queue, req, is_retracted=True)
-                if algorithm == "DSPARK":
-                    req.offload_kv_cache.assert_not_called()
-                    device_module.return_value.synchronize.assert_called_once_with(
-                        "cuda"
-                    )
-                    self.assertEqual(req.output_ids, [7])
-                    self.assertEqual(req.pd_rebootstrap_forced_output_id, 8)
-                    self.assertTrue(req.pd_rebootstrap_in_progress)
-                    queue._create_receiver_and_enqueue.assert_called_once_with(
-                        req, is_rebootstrap=True
-                    )
-                    self.assertFalse(queue.retracted_queue)
-                else:
-                    req.offload_kv_cache.assert_called_once_with(
-                        request_pool, allocator
-                    )
-                    device_module.assert_not_called()
-                    self.assertEqual(req.output_ids, [7, 8])
-                    self.assertEqual(queue.retracted_queue, [req])
+                # All V4 pools recompute on prefill; none keep a CPU KV backup.
+                req.offload_kv_cache.assert_not_called()
+                device_module.return_value.synchronize.assert_called_once_with("cuda")
+                self.assertEqual(req.output_ids, [7])
+                self.assertEqual(req.pd_rebootstrap_forced_output_id, 8)
+                self.assertTrue(req.pd_rebootstrap_in_progress)
+                queue._create_receiver_and_enqueue.assert_called_once_with(
+                    req, is_rebootstrap=True
+                )
+                self.assertFalse(queue.retracted_queue)
 
 
 class TestDSV41CPPDHandshake(CustomTestCase):
@@ -387,9 +379,10 @@ def test_cp_and_dp_feature_matrix(self):
             ("decode_prefill_cp_flag", dict(enable_prefill_cp=True), False),
             ("decode_attention_cp", dict(attn_cp_size=2), False),
             ("decode_dcp", dict(dcp_size=2), False),
+            ("nixl_transfer", dict(disaggregation_transfer_backend="nixl"), True),
             (
                 "unsupported_transfer",
-                dict(disaggregation_transfer_backend="nixl"),
+                dict(disaggregation_transfer_backend="fake"),
                 False,
             ),
         ]
diff --git a/test/registered/unit/mem_cache/test_fake_transfer_chunk_progress.py b/test/registered/unit/mem_cache/test_fake_transfer_chunk_progress.py
index 840d77a5a469..ff28c7720c77 100644
--- a/test/registered/unit/mem_cache/test_fake_transfer_chunk_progress.py
+++ b/test/registered/unit/mem_cache/test_fake_transfer_chunk_progress.py
@@ -4,9 +4,7 @@
 from types import SimpleNamespace as NS
 from unittest.mock import Mock
 
-import torch
-
-from sglang.srt.mem_cache.common import maybe_cache_unfinished_req
+from sglang.srt.mem_cache.common import checkpoint_kv_cache
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -15,31 +13,32 @@
 
 class TestFakeTransferChunkProgress(CustomTestCase):
     def test_chunk_progress_without_shared_insert(self):
-        slots = torch.arange(16385, dtype=torch.int32).reshape(1, -1)
-        cache = NS(
-            req_to_token_pool=NS(req_to_token=slots), cache_unfinished_req=Mock()
-        )
+        cache = NS(checkpoint=Mock())
         req = NS(
             skip_radix_cache_insert=True,
-            kv=NS(req_pool_idx=0, cache_protected_len=0),
-            get_fill_ids=lambda: range(16384),
+            finished=lambda: False,
+            extend_end=16384,
+            prefix_len=0,
         )
-        maybe_cache_unfinished_req(req, cache, chunked=True)
-        self.assertEqual(16385 - len(req.prefix_indices), 1)
-        self.assertEqual(req.kv.cache_protected_len, 0)
-        cache.cache_unfinished_req.assert_not_called()
-        self.assertEqual(req.prefix_indices.dtype, torch.int64)
-        slots[0, 0] = -1
-        self.assertEqual(req.prefix_indices[0].item(), 0)
-        req.get_fill_ids = lambda: range(16385)
-        maybe_cache_unfinished_req(req, cache, chunked=True)
-        self.assertEqual(len(req.prefix_indices), 16385)
+        checkpoint_kv_cache(req, cache)
+        self.assertEqual(16385 - req.prefix_len, 1)
+        cache.checkpoint.assert_not_called()
+        req.extend_end = 16385
+        checkpoint_kv_cache(req, cache)
+        self.assertEqual(req.prefix_len, 16385)
+        cache.checkpoint.assert_not_called()
 
     def test_real_transfer_preserves_cache_path(self):
-        req = NS(skip_radix_cache_insert=False)
-        cache = NS(cache_unfinished_req=Mock())
-        maybe_cache_unfinished_req(req, cache, chunked=True)
-        cache.cache_unfinished_req.assert_called_once_with(req, chunked=True)
+        req = NS(
+            skip_radix_cache_insert=False,
+            finished=lambda: False,
+            extend_end=16384,
+            prefix_len=0,
+        )
+        cache = NS(checkpoint=Mock())
+        checkpoint_kv_cache(req, cache)
+        cache.checkpoint.assert_called_once_with(req, up_to=16384)
+        self.assertEqual(req.prefix_len, 16384)
 
 
 if __name__ == "__main__":