[DeepSeek V4] Enable FlashMLA sparse prefill by default (#29775)

Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>

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}, "

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

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
+    # reused across every layer in the chunk. Reset to ``None`` when graph
-    # cuda-graph replay rebuilds it for the next forward.
+    # 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:]
 

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(self.max_seq_len // 4, 1)
-        c4_max = max_blocks * c4_page_size
+        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,

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.