Reuse shared compressed KV dequantization in DeepSeek V4.1 CP prefill

(cherry picked from commit 44eb378e94610e4ccc6b990bf09c105f8c8f9ee9)

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -1296,6 +1296,11 @@ class DeepseekV4AttnBackend(
         ] = 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()
@@ -4034,20 +4039,38 @@ class DeepseekV4AttnBackend(
                 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,

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):
@@ -275,14 +276,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.
     """
 

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:]]))