[Perf] dsv4: Dequantize each shared sparse-prefill KV source once per forward

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

@@ -75,7 +75,7 @@
 from sglang.srt.layers.attention.dsv4.online_c128_mtp import OnlineC128MTPController
 from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
     SparsePrefillChunkCache,
-    SparsePrefillWorkspace,
+    SparsePrefillKVBuffer,
     use_dsv4_q8kv8_sparse_prefill,
 )
 from sglang.srt.layers.attention.dsv4.v41_indexer import (
@@ -1302,7 +1302,7 @@ def __init__(
             DSV4RawDecodeMetadata,
         ] = None
         self.online_c128_mtp = OnlineC128MTPController(self)
-        self.sparse_prefill_workspace = SparsePrefillWorkspace(self.device)
+        self.sparse_prefill_kv_buffer = SparsePrefillKVBuffer(self.device)
         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()
@@ -2758,6 +2758,7 @@ def _build_sparse_prefill_chunk_cache(
         *,
         num_qo_tokens: int,
     ) -> SparsePrefillChunkCache:
+        self.sparse_prefill_kv_buffer.forget_head()
         seq_lens_cpu = forward_batch.seq_lens_cpu
         assert seq_lens_cpu is not None
         # Decoder tails require explicit SWA indices to exclude unwritten history.
@@ -4198,7 +4199,7 @@ def _forward_prefill_sparse(
         extra_page_size = None
         flat_token_ids = None
         if compress_ratio == 0:
-            workspace = self.sparse_prefill_workspace.get(cache.swa_token_ids.shape[0])
+            workspace = self.sparse_prefill_kv_buffer.get(cache.swa_token_ids.shape[0])
             combined_indices = cache.c0_combined_indices
             combined_lens = cache.c0_combined_lens
             swa_slice = workspace
@@ -4209,10 +4210,18 @@ 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]
+            # Layers sharing a kv_source reuse its dequant: no later layer writes
+            # that cache, and graph replay keeps the layer order.
+            workspace, compressed_filled = self.sparse_prefill_kv_buffer.get_with_head(
+                n_compressed + cache.swa_token_ids.shape[0],
+                head_source=(
+                    extra_k_cache.data_ptr(),
+                    flat_token_ids.data_ptr(),
+                    n_compressed,
+                ),
             )
-            compressed_slice = workspace[:n_compressed]
+            if not compressed_filled:
+                compressed_slice = workspace[:n_compressed]
             swa_slice = workspace[n_compressed:]
 
         if compressed_slice is not None:
@@ -4369,7 +4378,7 @@ def _forward_prefill_sparse_q8kv8(
         flat_token_ids = None
 
         if compress_ratio == 0:
-            workspace = self.sparse_prefill_workspace.get(
+            workspace = self.sparse_prefill_kv_buffer.get(
                 cache.swa_token_ids.shape[0] + 1,
                 dtype=fp8_dtype,
             )
@@ -4384,7 +4393,7 @@ def _forward_prefill_sparse_q8kv8(
             )
 
             n_compressed = flat_token_ids.shape[0]
-            workspace = self.sparse_prefill_workspace.get(
+            workspace = self.sparse_prefill_kv_buffer.get(
                 n_compressed + cache.swa_token_ids.shape[0] + 1,
                 dtype=fp8_dtype,
             )

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

@@ -91,7 +91,7 @@
 )
 from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
     SparsePrefillChunkCache,
-    SparsePrefillWorkspace,
+    SparsePrefillKVBuffer,
 )
 from sglang.srt.layers.attention.hip_flash_mla import (
     _apply_inverse_rope,
@@ -1075,9 +1075,9 @@ def __init__(
             DSV4RawVerifyMetadata,
             DSV4RawDecodeMetadata,
         ] = None
-        self.opus_prefill_workspace = SparsePrefillWorkspace(self.device)
+        self.opus_prefill_workspace = SparsePrefillKVBuffer(self.device)
         self.opus_compressed_workspaces = {
-            ratio: SparsePrefillWorkspace(self.device) for ratio in (1, 2)
+            ratio: SparsePrefillKVBuffer(self.device) for ratio in (1, 2)
         }
         # Reset with the metadata: its inputs belong to one prefill forward.
         self._opus_prefill_state: Optional[_OpusPrefillState] = None

python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py

@@ -34,7 +34,7 @@
 
 import os
 from dataclasses import dataclass, field
-from typing import Dict, Optional
+from typing import Dict, Optional, Tuple
 
 import torch
 
@@ -72,24 +72,43 @@ def use_dsv4_q8kv8_sparse_prefill(dsv4_prefill_backend: str = "auto") -> bool:
     return dsv4_prefill_backend == "flashmla_sparse_q8"
 
 
-class SparsePrefillWorkspace:
+class SparsePrefillKVBuffer:
     """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 overwrite the rows they read, except head rows ``get_with_head``
+    reports as filled, 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):
         self.device = device
         self._buffer: Optional[torch.Tensor] = None
+        self._head_source: Optional[tuple] = None
 
     def get(
         self,
         num_tokens: int,
         dtype: torch.dtype = torch.bfloat16,
     ) -> torch.Tensor:
+        self._head_source = None
+        return self._reserve(num_tokens, dtype)
+
+    def get_with_head(
+        self, num_tokens: int, head_source: tuple
+    ) -> Tuple[torch.Tensor, bool]:
+        """Returns the rows and whether the head already holds ``head_source``."""
+        buffer = self._reserve(num_tokens, torch.bfloat16)
+        head_filled = head_source == self._head_source
+        self._head_source = head_source
+        return buffer, head_filled
+
+    def forget_head(self) -> None:
+        """Per forward: the next forward's tensors may reuse freed addresses."""
+        self._head_source = None
+
+    def _reserve(self, num_tokens: int, dtype: torch.dtype) -> torch.Tensor:
         assert num_tokens > 0
         current_capacity = self._buffer.shape[0] if self._buffer is not None else 0
         current_dtype = self._buffer.dtype if self._buffer is not None else None
@@ -99,6 +118,7 @@ def get(
                 dtype=dtype,
                 device=self.device,
             )
+            self._head_source = None
         return self._buffer[:num_tokens]
 
 

test/manual/test_dsv41_ragged_prefill.py

@@ -7,7 +7,7 @@
 from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
     CompressedGather,
     SparsePrefillChunkCache,
-    SparsePrefillWorkspace,
+    SparsePrefillKVBuffer,
 )
 
 
@@ -175,9 +175,9 @@ def test_00_mixed_batch_under_memory_budget(self):
             cache = make_cache([44198, 577152, 95723, 51427, 6144], [1] * 5, "cuda")
             table = torch.arange(2255, device="cuda", dtype=torch.int32).expand(5, -1)
             with self.assertRaises(torch.OutOfMemoryError):
-                SparsePrefillWorkspace(torch.device("cuda")).get(2902511)
+                SparsePrefillKVBuffer(torch.device("cuda")).get(2902511)
             gather = cache.ensure_compressed(1, table, 256)
-            workspace = SparsePrefillWorkspace(torch.device("cuda")).get(
+            workspace = SparsePrefillKVBuffer(torch.device("cuda")).get(
                 gather.flat_token_ids.numel() + 16751
             )
             self.assertEqual(workspace.shape[0], 791395)
@@ -200,7 +200,7 @@ def test_reported_workspace_sizes(self):
             new_rows = gather.flat_token_ids.numel() + swa
             self.assertEqual(old_rows, expected_old)
             self.assertEqual(new_rows, expected_new)
-            ws = SparsePrefillWorkspace(torch.device("cuda"))
+            ws = SparsePrefillKVBuffer(torch.device("cuda"))
             before = torch.cuda.memory_allocated()
             buf = ws.get(new_rows)
             allocated = torch.cuda.memory_allocated() - before

test/registered/attention/unittests/dsv4/test_deepseek_v4.py

@@ -723,12 +723,12 @@ def test_trtllm_semaphore_capacity_covers_configured_query_rows(self):
             with self.assertRaisesRegex(RuntimeError, "exceeds the persistent"):
                 trtllm._check_trtllm_query_rows(65)
 
-    def test_sparse_prefill_workspace_reuses_and_grows(self):
+    def test_sparse_prefill_kv_buffer_reuses_and_grows(self):
         from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
-            SparsePrefillWorkspace,
+            SparsePrefillKVBuffer,
         )
 
-        workspace = SparsePrefillWorkspace(torch.device("cpu"))
+        workspace = SparsePrefillKVBuffer(torch.device("cpu"))
         first = workspace.get(3)
         reused = workspace.get(2)
         grown = workspace.get(7)

test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py

@@ -26,7 +26,7 @@
 from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
     CompressedGather,
     SparsePrefillChunkCache,
-    SparsePrefillWorkspace,
+    SparsePrefillKVBuffer,
     use_dsv4_q8kv8_sparse_prefill,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
@@ -184,7 +184,7 @@ def _make_backend(
         sparse_prefill_cache=None, late_layer_tail=None
     )
     backend.req_to_token = req_to_token
-    backend.sparse_prefill_workspace = SparsePrefillWorkspace(device)
+    backend.sparse_prefill_kv_buffer = SparsePrefillKVBuffer(device)
     backend.softmax_scale = 512**-0.5
     backend.head_dim_v = 512
     backend.dsv4_prefill_backend = dsv4_prefill_backend

test/registered/unit/mem_cache/test_dsv4_compressed_pools.py

@@ -18,6 +18,7 @@
 from sglang.srt.layers.attention.deepseek_v4_backend import DeepseekV4AttnBackend
 from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import (
     SparsePrefillChunkCache,
+    SparsePrefillKVBuffer,
 )
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
     DeepSeekV4SingleKVPool,
@@ -656,6 +657,7 @@ def test_sparse_prefill_chunk_cache_builds_without_mapping(self):
         backend.token_to_kv_pool = pool
         backend.req_to_token = torch.zeros((3, 2048), dtype=torch.int32)
         backend.forward_metadata = SimpleNamespace(late_layer_tail=None)
+        backend.sparse_prefill_kv_buffer = SparsePrefillKVBuffer(torch.device("cpu"))
         forward_batch = SimpleNamespace(
             forward_mode=ForwardMode.EXTEND,
             seq_lens=torch.tensor([rows]),