diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index e7602f5b287f..0ee3eb5ab91c 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/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, ) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index f6ce39ae12f9..18e26c5672a5 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/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 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 5478a8290def..8410e728f043 100644 --- a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py +++ b/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] diff --git a/test/manual/test_dsv41_ragged_prefill.py b/test/manual/test_dsv41_ragged_prefill.py index c0e734b3b76e..04b466b51e71 100644 --- a/test/manual/test_dsv41_ragged_prefill.py +++ b/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 diff --git a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py index 194a0ed5f2bd..98f563d8d30e 100644 --- a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py +++ b/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) diff --git a/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py b/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py index 836699af6c8b..357c7f95fa1d 100644 --- a/test/registered/kernels/ops/attention/test_q8kv8_sparse_prefill_backend.py +++ b/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 diff --git a/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py b/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py index eaa2843011a5..f3ba97165bac 100644 --- a/test/registered/unit/mem_cache/test_dsv4_compressed_pools.py +++ b/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]),