minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
YAMY
2026-07-01 13:50:05 -07:00
committed by GitHub
co-authored by Baizhou Zhang
parent 8f0d320d31
commit c865347b98
5 changed files with 150 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/deepseek_v4_hook.py
+7
View File
@@ -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
+1 -1
View File
@@ -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
+15 -6
View File
@@ -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:]
python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py
+49 -28
View File
@@ -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,