diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 19ea8e235..977d4c31e 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py @@ -3,6 +3,8 @@ from __future__ import annotations import logging from typing import TYPE_CHECKING +from sglang.srt.environ import envs + if TYPE_CHECKING: from sglang.srt.server_args import ServerArgs @@ -93,6 +95,11 @@ def validate_deepseek_v4_cp(server_args: ServerArgs) -> None: assert ( server_args.tp_size <= 8 ), "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues." + logger.warning( + "Disabling SGLANG_OPT_FLASHMLA_SPARSE_PREFILL because DeepSeekV4 " + "context parallelism is enabled." + ) + envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.set(False) logger.warning( f"Enable Context Parallel for DeepSeekV4, " f"dp_size={server_args.dp_size}, moe_dense_tp_size={server_args.moe_dense_tp_size}, " diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index fefa853c4..3938c10a6 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -874,7 +874,7 @@ class Envs: SGLANG_OPT_USE_COMPRESSOR_V2 = EnvBool(True) SGLANG_FP8_PAGED_MQA_LOGITS_TORCH = EnvBool(False) SGLANG_TOPK_TRANSFORM_512_TORCH = EnvBool(False) - SGLANG_OPT_FLASHMLA_SPARSE_PREFILL = EnvBool(False) + SGLANG_OPT_FLASHMLA_SPARSE_PREFILL = EnvBool(True) # SWA radix cache # TODO(DSV4): @ispobock this has bug on main branch when retract diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index f9571fe77..c6c21a1b6 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -44,6 +44,7 @@ from sglang.srt.layers.attention.dsv4.quant_k_cache import ( ) from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import ( SparsePrefillChunkCache, + SparsePrefillWorkspace, ) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode @@ -363,8 +364,8 @@ class DSV4Metadata: c128_compress_metadata: Optional[FusedCompressMetadata] = None # Lazily populated on the first call to ``_forward_prefill_sparse`` and - # reused across every layer in the chunk. Reset to ``None`` on copy_ so - # cuda-graph replay rebuilds it for the next forward. + # reused across every layer in the chunk. Reset to ``None`` when graph + # metadata is refreshed so replay rebuilds it from the live batch. sparse_prefill_cache: Optional[SparsePrefillChunkCache] = None @property @@ -397,6 +398,7 @@ class DSV4Metadata: self.c128_compress_metadata, src=static_metadata.c128_compress_metadata, ) + self.sparse_prefill_cache = None @dataclass @@ -506,6 +508,7 @@ class DeepseekV4AttnBackend( DSV4RawDecodeMetadata, ] = None self.online_c128_mtp = OnlineC128MTPController(self) + self.sparse_prefill_workspace = SparsePrefillWorkspace(self.device) def _move_to_device(self, x: List[int]) -> torch.Tensor: pin_tensor = torch.tensor(x, dtype=torch.int32, pin_memory=True) @@ -1470,6 +1473,8 @@ class DeepseekV4AttnBackend( cache = self.forward_metadata.sparse_prefill_cache if cache is None: + seq_lens_cpu = forward_batch.seq_lens_cpu + assert seq_lens_cpu is not None # ``swa_window_size`` on the pool is its storage page size, not # the model's SWA window — pass both explicitly. cache = SparsePrefillChunkCache.build( @@ -1481,6 +1486,7 @@ class DeepseekV4AttnBackend( swa_window_size=SWA_WINDOW, swa_page_size=token_to_kv_pool.swa_window_size, num_qo_tokens=q_flat.shape[0], + max_seq_len=int(seq_lens_cpu.max().item()), ) self.forward_metadata.sparse_prefill_cache = cache @@ -1491,7 +1497,7 @@ class DeepseekV4AttnBackend( extra_page_size = None flat_token_ids = None if compress_ratio == 0: - workspace = cache.c0_workspace + workspace = self.sparse_prefill_workspace.get(cache.swa_token_ids.shape[0]) combined_indices = cache.c0_combined_indices combined_lens = cache.c0_combined_lens swa_slice = workspace @@ -1502,7 +1508,6 @@ class DeepseekV4AttnBackend( assert core_attn_metadata.c128_page_indices is not None cache.ensure_c128(core_attn_metadata.c128_page_indices) flat_token_ids = cache.c128_flat_token_ids - workspace = cache.c128_workspace combined_indices = cache.c128_combined_indices combined_lens = cache.c128_combined_lens else: @@ -1512,11 +1517,15 @@ class DeepseekV4AttnBackend( ) cache.ensure_c4(core_attn_metadata.page_table, extra_page_size) flat_token_ids = cache.c4_flat_token_ids - workspace = cache.c4_workspace combined_indices, combined_lens = cache.combine_c4_layer( - c4_sparse_raw_indices=core_attn_metadata.c4_sparse_raw_indices, + c4_sparse_raw_indices=core_attn_metadata.c4_sparse_raw_indices[ + : cache.num_qo_tokens + ], ) n_compressed = flat_token_ids.shape[0] + workspace = self.sparse_prefill_workspace.get( + n_compressed + cache.swa_token_ids.shape[0] + ) compressed_slice = workspace[:n_compressed] swa_slice = workspace[n_compressed:] 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 a84c80347..be5ef1915 100644 --- a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py +++ b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py @@ -50,6 +50,31 @@ SPARSE_PREFILL_TOPK_ALIGNMENT = 128 WORKSPACE_DIM = DIM_NOPE + DIM_ROPE +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. + """ + + def __init__(self, device: torch.device): + self.device = device + self._buffer: Optional[torch.Tensor] = None + + def get(self, num_tokens: int) -> torch.Tensor: + assert num_tokens > 0 + current_capacity = self._buffer.shape[0] if self._buffer is not None else 0 + if num_tokens > current_capacity: + self._buffer = torch.empty( + (num_tokens, 1, WORKSPACE_DIM), + dtype=torch.bfloat16, + device=self.device, + ) + return self._buffer[:num_tokens] + + def combined_topk_width(topk: int, window_size: int) -> int: """Width of the padded combined_indices last dim that ``combine_topk_swa_indices`` would produce for these args.""" @@ -341,6 +366,10 @@ class SparsePrefillChunkCache: # Geometry computed once per chunk. num_reqs: int num_qo_tokens: int + # Actual maximum sequence length in this forward. CUDA-graph metadata may + # have a much wider page table sized for the capture limit; gather only the + # live sequence extent instead of materializing that padded capacity. + max_seq_len: int # Model's SWA window — the per-query attention range. Used by # combine_topk_swa_indices' WINDOW_SIZE and by build_swa_token_ids's # gather_lens. Must match SWA_WINDOW from the backend (e.g. 128), NOT @@ -361,24 +390,16 @@ class SparsePrefillChunkCache: # c0 pre-computed combine output (entire input set is chunk-invariant). c0_combined_indices: torch.Tensor = field(default=None) c0_combined_lens: torch.Tensor = field(default=None) - # Preallocated workspace reused across layers — avoids per-layer - # ``torch.cat`` and bf16 allocations. Shape (total_swa, 1, 512) bf16 for - # c0, (total_compressed + total_swa, 1, 512) for c4/c128. Dequant kernels - # write directly via ``out=workspace[slice]``. - c0_workspace: torch.Tensor = field(default=None) - # c128: positional layout of the c128 cache + pre-computed combine. c128_flat_token_ids: Optional[torch.Tensor] = None # (num_reqs * c128_max,) int32 c128_combined_indices: Optional[torch.Tensor] = None c128_combined_lens: Optional[torch.Tensor] = None - c128_workspace: Optional[torch.Tensor] = None # c4: positional layout of the c4 cache (combine output is per-layer). c4_flat_token_ids: Optional[torch.Tensor] = None # (num_reqs * c4_max,) int32 c4_page_size: Optional[int] = None c4_compressed_base: Optional[torch.Tensor] = None # (num_reqs,) int32 c4_swa_base: Optional[torch.Tensor] = None # (num_reqs,) int32 - c4_workspace: Optional[torch.Tensor] = None # Tail stays at the -1 sentinel because the valid prefix length is # chunk-invariant per request — subsequent layers only overwrite that prefix. c4_combined_indices: Optional[torch.Tensor] = None @@ -395,6 +416,7 @@ class SparsePrefillChunkCache: swa_window_size: int, swa_page_size: int, num_qo_tokens: int, + max_seq_len: int, ) -> "SparsePrefillChunkCache": device = seq_lens.device num_reqs = seq_lens.shape[0] @@ -416,6 +438,7 @@ class SparsePrefillChunkCache: cache = cls( num_reqs=num_reqs, num_qo_tokens=num_qo_tokens, + max_seq_len=max_seq_len, swa_window_size=swa_window_size, swa_page_size=swa_page_size, seq_lens=seq_lens, @@ -442,11 +465,6 @@ class SparsePrefillChunkCache: compress_ratio=1, topk=0, ) - cache.c0_workspace = torch.empty( - (swa_token_ids.shape[0], 1, WORKSPACE_DIM), - dtype=torch.bfloat16, - device=device, - ) return cache def ensure_c128(self, c128_page_indices: torch.Tensor) -> None: @@ -465,9 +483,15 @@ class SparsePrefillChunkCache: if self.c128_flat_token_ids is not None: return device = self.seq_lens.device - c128_max = c128_page_indices.shape[-1] + c128_max = max(self.max_seq_len // 128, 1) + assert c128_max <= c128_page_indices.shape[-1], ( + f"live c128 extent {c128_max} exceeds metadata capacity " + f"{c128_page_indices.shape[-1]}" + ) last_q_per_req = (self.query_start_loc[1:] - 1).long() - per_req_c128 = c128_page_indices[last_q_per_req] + per_req_c128 = c128_page_indices.narrow(1, 0, c128_max).index_select( + 0, last_q_per_req + ) # Clamp -1 -> 0 so dequant doesn't OOB; combine masks the invalid # tail via topk_len. flat_c128_ids = per_req_c128.reshape(-1).clamp_min(0).to(torch.int32) @@ -499,11 +523,6 @@ class SparsePrefillChunkCache: self.c128_flat_token_ids = flat_c128_ids self.c128_combined_indices = combined_indices self.c128_combined_lens = combined_lens - self.c128_workspace = torch.empty( - (total_compressed + self.swa_token_ids.shape[0], 1, WORKSPACE_DIM), - dtype=torch.bfloat16, - device=device, - ) def ensure_c4( self, @@ -520,10 +539,17 @@ class SparsePrefillChunkCache: if self.c4_flat_token_ids is not None: return device = self.seq_lens.device - max_blocks = page_table.shape[-1] - c4_max = max_blocks * c4_page_size + c4_max = max(self.max_seq_len // 4, 1) + c4_capacity = page_table.shape[-1] * c4_page_size + assert ( + c4_max <= c4_capacity + ), f"live c4 extent {c4_max} exceeds metadata capacity {c4_capacity}" first_q_per_req = self.query_start_loc[:-1].long() - per_req_page_table = page_table[first_q_per_req] + num_blocks = (c4_max + c4_page_size - 1) // c4_page_size + assert num_blocks <= page_table.shape[1] + per_req_page_table = page_table.narrow(1, 0, num_blocks).index_select( + 0, first_q_per_req + ) k_arange = torch.arange(c4_max, dtype=torch.int32, device=device) block_idx = (k_arange // c4_page_size).long() @@ -542,11 +568,6 @@ class SparsePrefillChunkCache: self.c4_page_size = c4_page_size self.c4_compressed_base = compressed_base self.c4_swa_base = swa_base - self.c4_workspace = torch.empty( - (total_compressed + self.swa_token_ids.shape[0], 1, WORKSPACE_DIM), - dtype=torch.bfloat16, - device=device, - ) def combine_c4_layer( self, diff --git a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py index b52e32bad..dbd949869 100644 --- a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py +++ b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py @@ -17,6 +17,7 @@ import sys import unittest from pathlib import Path from types import SimpleNamespace +from unittest import mock import torch @@ -266,6 +267,27 @@ class TestDSV4AttentionBackendCorrectness(CustomTestCase): class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase): """CPU-only checks for the DSV4 BCG metadata replay contract.""" + @staticmethod + def _make_sparse_prefill_cache(max_seq_len): + from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import ( + SparsePrefillChunkCache, + ) + + int32 = dict(dtype=torch.int32) + return SparsePrefillChunkCache( + num_reqs=2, + num_qo_tokens=2, + max_seq_len=max_seq_len, + swa_window_size=128, + swa_page_size=128, + seq_lens=torch.tensor([max_seq_len, max_seq_len], **int32), + query_start_loc=torch.tensor([0, 1, 2], **int32), + swa_token_ids=torch.empty(0, **int32), + swa_first_pos=torch.zeros(2, **int32), + swa_gather_lens=torch.zeros(2, **int32), + swa_offsets=torch.zeros(3, **int32), + ) + def _make_core_metadata(self, base: int): from sglang.srt.layers.attention.deepseek_v4_backend import DSV4AttnMetadata @@ -400,6 +422,7 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase): capture_metadata = DSV4Metadata( self._make_core_metadata(0), indexer_metadata=None ) + capture_metadata.sparse_prefill_cache = object() replay_metadata = DSV4Metadata( self._make_core_metadata(1000), indexer_metadata=None ) @@ -427,6 +450,7 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase): self.assertEqual(calls[0][1], backend.MAX_SEQ_LEN_FOR_CAPTURE) self.assertTrue(calls[0][2]) self.assertIs(backend.forward_metadata, capture_metadata) + self.assertIsNone(capture_metadata.sparse_prefill_cache) self.assertTrue( torch.equal( capture_metadata.core_attn_metadata.seq_lens_casual, @@ -434,6 +458,60 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase): ) ) + def test_sparse_prefill_workspace_reuses_and_grows(self): + from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import ( + SparsePrefillWorkspace, + ) + + workspace = SparsePrefillWorkspace(torch.device("cpu")) + first = workspace.get(3) + reused = workspace.get(2) + grown = workspace.get(7) + + self.assertEqual(first.shape, (3, 1, 512)) + self.assertEqual(reused.data_ptr(), first.data_ptr()) + self.assertEqual(grown.shape, (7, 1, 512)) + self.assertNotEqual(grown.data_ptr(), first.data_ptr()) + self.assertEqual(workspace._buffer.data_ptr(), grown.data_ptr()) + + def test_sparse_prefill_c4_uses_live_extent(self): + page_table = torch.zeros((2, 4096), dtype=torch.int32) + for max_seq_len in (3, 4, 255, 256, 259, 260): + with self.subTest(max_seq_len=max_seq_len): + cache = self._make_sparse_prefill_cache(max_seq_len) + cache.ensure_c4(page_table, c4_page_size=64) + expected_extent = max(max_seq_len // 4, 1) + self.assertEqual(cache.c4_flat_token_ids.numel(), 2 * expected_extent) + self.assertEqual( + cache.c4_compressed_base.tolist(), [0, expected_extent] + ) + + def test_sparse_prefill_c128_uses_live_extent(self): + from sglang.srt.layers.attention.dsv4 import sparse_prefill_utils + + page_indices = torch.full((2, 8192), -1, dtype=torch.int32) + for max_seq_len in (127, 128, 255, 256): + with self.subTest(max_seq_len=max_seq_len): + cache = self._make_sparse_prefill_cache(max_seq_len) + expected_extent = max(max_seq_len // 128, 1) + combined = ( + torch.empty((2, 256), dtype=torch.int32), + torch.empty(2, dtype=torch.int32), + ) + with mock.patch.object( + sparse_prefill_utils, + "combine_topk_swa_indices", + return_value=combined, + ) as combine: + cache.ensure_c128(page_indices) + + self.assertEqual(cache.c128_flat_token_ids.numel(), 2 * expected_extent) + self.assertEqual(combine.call_args.kwargs["topk"], expected_extent) + self.assertEqual( + combine.call_args.kwargs["topk_indices"].shape, + (2, expected_extent), + ) + class TestDSV4SwaOutCacheLocResolution(CustomTestCase): """`get_swa_out_cache_loc`: cached fast path vs store-time fallback.