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

[DSV4] Generalize attention metadata, sparse prefill, and KV pool over compress ratios (#39921)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-09-17 15:55:19 -07:00
committed by GitHub
parent 4f52a27563
commit 1f0c73e9bd
18 changed files with 956 additions and 409 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/dsv4/sparse_prefill_kernels.py
+5 -8
View File
@@ -44,6 +44,7 @@ def _combine_topk_swa_indices_kernel(
topk_indices_ptr,
topk_indices_stride,
query_start_loc_ptr,
query_pos_ptr,
seq_lens_ptr,
gather_lens_ptr,
compressed_base_ptr,
@@ -62,23 +63,19 @@ def _combine_topk_swa_indices_kernel(
base = tl.load(query_start_loc_ptr)
query_start = tl.load(query_start_loc_ptr + batch_idx) - base
query_end = tl.load(query_start_loc_ptr + batch_idx + 1) - base
query_len = query_end - query_start
seq_len = tl.load(seq_lens_ptr + batch_idx)
gather_len = tl.load(gather_lens_ptr + batch_idx)
compressed_base = tl.load(compressed_base_ptr + batch_idx)
swa_base = tl.load(swa_base_ptr + batch_idx)
start_pos = seq_len - query_len
# SWA portion of the gathered buffer starts from position
# (seq_len - gather_len), not 0. The +pos-gather_start formula maps a
# query's window back into the workspace's SWA region.
gather_start = seq_len - gather_len
for token_idx in range(query_start + worker_id, query_end, num_workers):
token_idx_in_query = token_idx - query_start
pos = start_pos + token_idx_in_query
# Both the C4 indexer and the C128 metadata builder emit
# min((pos+1)//compress_ratio, topk_tokens) valid entries. Caller
# passes top_k=0 for SWA-only layers to zero this out.
pos = tl.load(query_pos_ptr + token_idx)
# -1 entries inside the top-k span stay -1 (attention skips them).
# top_k=0 disables the compressed portion for SWA-only layers.
topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, top_k)
swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
@@ -93,7 +90,7 @@ def _combine_topk_swa_indices_kernel(
)
tl.store(
combined_indices_ptr + combined_row + offset,
topk_vals + compressed_base,
tl.where(topk_vals >= 0, topk_vals + compressed_base, -1),
mask=mask,
)
python/sglang/srt/disaggregation/decode.py
+8 -15
View File
@@ -61,10 +61,8 @@ from sglang.srt.disaggregation.utils import (
build_kv_layer_ids,
build_staging_slot_metadata,
get_dsa_tail_state_indices,
get_dsv4_c128_state_indices,
get_kv_class,
get_qsa_pending_state_indices,
is_dsv4_c128_online_enabled,
is_mla_backend,
is_unadmitted_reject,
poll_and_all_reduce,
@@ -1498,23 +1496,18 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
ring_rows = state_slot * ring_stride + (positions % ring_stride)
return ring_rows.astype(np.int32)
def _c128_state_payload():
online = is_dsv4_c128_online_enabled()
ring_size = 1 if online else self.token_to_kv_pool.get_ring_size(128)
return get_dsv4_c128_state_indices(
int(decode_req.req.kv.req_pool_idx),
seq_len,
online=online,
ring_size=ring_size,
def _request_state_payload():
return self.token_to_kv_pool.request_state_transfer_indices(
int(decode_req.req.kv.req_pool_idx), seq_len
)
state_types = self.kv_manager.kv_args.state_types
if StateType.DSV4_REQUEST_STATE in state_types:
clear_c128_state = getattr(
self.token_to_kv_pool, "clear_c128_req_state", None
clear_request_state = getattr(
self.token_to_kv_pool, "clear_request_scoped_state", None
)
if clear_c128_state is not None:
clear_c128_state(int(decode_req.req.kv.req_pool_idx))
if clear_request_state is not None:
clear_request_state(int(decode_req.req.kv.req_pool_idx))
payloads = {
StateType.MAMBA: _mamba_payload,
StateType.QSA_PENDING: _qsa_pending_payload,
@@ -1524,7 +1517,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
StateType.DSA_TAIL: _dsa_tail_payload,
StateType.MINIMAX_INDEX_K: _full_kv_pages_payload,
StateType.SWA_RING: _swa_ring_payload,
StateType.DSV4_REQUEST_STATE: _c128_state_payload,
StateType.DSV4_REQUEST_STATE: _request_state_payload,
StateType.BLOCK_SCALE: _full_kv_pages_payload,
StateType.BLOCK_SCALE_SWA: _swa_payload,
}
python/sglang/srt/disaggregation/prefill.py
+5 -25
View File
@@ -52,11 +52,9 @@ from sglang.srt.disaggregation.utils import (
build_kv_layer_ids,
build_staging_slot_metadata,
get_dsa_tail_state_indices,
get_dsv4_c128_state_indices,
get_kv_class,
get_qsa_pending_state_indices,
is_aborted,
is_dsv4_c128_online_enabled,
is_mla_backend,
is_unadmitted_reject,
poll_and_all_reduce_attn_cp_tp_group,
@@ -257,7 +255,6 @@ class PrefillBootstrapQueue:
hf_text_config=self.scheduler.model_config.hf_text_config,
)
)
kv_args.mla_compression_ratios = None
kv_data_ptrs, kv_data_lens, kv_item_lens = (
self.token_to_kv_pool.get_contiguous_buf_infos()
)
@@ -316,13 +313,6 @@ class PrefillBootstrapQueue:
req_to_token_pool=req_to_token_pool,
)
if isinstance(self.token_to_kv_pool, DeepSeekV4TokenToKVPool):
# V4's KVCache is organized by compression-ratio
# buckets rather than by layer.
kv_args.mla_compression_ratios = list(
self.token_to_kv_pool.compression_ratios
)
kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)
kv_manager = kv_manager_class(
kv_args,
@@ -1415,20 +1405,10 @@ class SchedulerDisaggregationPrefillMixin:
ring_rows = state_slot * ring_stride + (positions % ring_stride)
return ring_rows.astype(np.int32)
def _c128_state_payload():
online = is_dsv4_c128_online_enabled()
ring_size = (
1
if online
else self.token_to_kv_pool_allocator.get_kvcache().get_ring_size(
128
)
)
return get_dsv4_c128_state_indices(
int(req.kv.req_pool_idx),
c128_seq_len,
online=online,
ring_size=ring_size,
def _request_state_payload():
kvcache = self.token_to_kv_pool_allocator.get_kvcache()
return kvcache.request_state_transfer_indices(
int(req.kv.req_pool_idx), c128_seq_len
)
state_types = (
@@ -1443,7 +1423,7 @@ class SchedulerDisaggregationPrefillMixin:
StateType.DSA_TAIL: _dsa_tail_payload,
StateType.MINIMAX_INDEX_K: _full_kv_pages_payload,
StateType.SWA_RING: _swa_ring_payload,
StateType.DSV4_REQUEST_STATE: _c128_state_payload,
StateType.DSV4_REQUEST_STATE: _request_state_payload,
StateType.BLOCK_SCALE: _full_kv_pages_payload,
StateType.BLOCK_SCALE_SWA: _swa_payload,
}
python/sglang/srt/disaggregation/utils.py
+7 -46
View File
@@ -25,7 +25,7 @@ from sglang.srt.environ import envs
from sglang.srt.runtime_context import (
get_disagg,
)
from sglang.srt.utils import is_hip, is_npu
from sglang.srt.utils import is_npu
if TYPE_CHECKING:
from sglang.srt.disaggregation.base.conn import KVArgs, StateType
@@ -46,7 +46,6 @@ if is_npu():
# Constants & Enums
#########################
FAKE_BOOTSTRAP_HOST = "2.2.2.2"
_IS_HIP = is_hip()
def poll_and_all_reduce_pp(
@@ -78,50 +77,6 @@ def get_dsa_seed_metadata_dim(hf_config) -> int:
return get_dsa_mtp_topk_width(hf_config)
def is_dsv4_c128_online_enabled() -> bool:
"""Return whether DSV4 C128 uses request-scoped online state."""
return not _IS_HIP and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()
def get_dsv4_c4_state_indices(
req_pool_idx: int,
seq_len: int,
*,
ring_size: int,
) -> np.ndarray:
# Prefill and decode can have different ring sizes (8 or 16 with EAGLE/MTP);
# pair the overlap compressor's live rows by logical token position.
if ring_size < 8 or ring_size % 4 != 0:
raise ValueError(
f"C4 ring_size must be a multiple of 4 and at least 8, got {ring_size}"
)
seq_len = max(0, int(seq_len))
state_len = seq_len % 4 + 4
positions = np.arange(max(0, seq_len - state_len), seq_len, dtype=np.int64)
rows = int(req_pool_idx) * int(ring_size) + positions % int(ring_size)
return rows.astype(np.int32)
def get_dsv4_c128_state_indices(
req_pool_idx: int,
seq_len: int,
*,
online: bool,
ring_size: int,
) -> np.ndarray:
"""Return the PD transfer row/page indices for DSV4 C128 state."""
if seq_len == 0 or seq_len % 128 == 0:
return np.empty((0,), dtype=np.int32)
if online:
return np.array([int(req_pool_idx)], dtype=np.int32)
assert ring_size % 128 == 0, f"C128 ring_size must be 128-aligned, got {ring_size}"
pages_per_req = ring_size // 128
page = int(req_pool_idx) * pages_per_req + ((seq_len - 1) % ring_size) // 128
return np.array([page], dtype=np.int32)
def get_qsa_pending_state_indices(req: Req) -> np.ndarray:
"""Return the request-pool row that owns a QSA pending-state ring."""
req_pool_idx = req.kv.req_pool_idx
@@ -1380,6 +1335,12 @@ def setup_state_kv_args(
kv_args.state_layer_ids = []
kv_args.is_hybrid_mla_backend = False
kv_args.state_conv_shard_groups = []
# V4's KVCache is organized by compression-ratio buckets rather than by layer.
kv_args.mla_compression_ratios = (
list(token_to_kv_pool.compression_ratios)
if isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
else None
)
def append_dsa_tail(pool) -> None:
if not pool.kpool_use_compress:
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_allocator.py
+1 -1
View File
@@ -544,7 +544,7 @@ class DSV4NPUTokenToKVPoolAllocator(SWATokenToKVPoolAllocator):
row = req_to_token_pool.req_to_c128_sidecar[int(req_pool_idx)]
self.release_c128_pages(row[row > 0])
row.zero_()
self.get_kvcache().clear_c128_req_state(int(req_pool_idx))
self.get_kvcache().clear_request_scoped_state(int(req_pool_idx))
def available_size(self):
return min(
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_common_hooks.py
+4 -2
View File
@@ -88,8 +88,10 @@ def dsv4_state_payloads(
import numpy as np
from sglang.srt.disaggregation.ascend.conn import AscendStateType
from sglang.srt.disaggregation.utils import get_dsv4_c4_state_indices
from sglang.srt.hardware_backend.npu.utils import is_npu_arch35
from sglang.srt.mem_cache.deepseek_v4_compress_state import (
c4_state_transfer_indices,
)
seq_len = max(0, int(seq_len))
prefix_len = max(0, min(int(prefix_len), seq_len))
@@ -113,7 +115,7 @@ def dsv4_state_payloads(
if is_npu_arch35():
def c4_state_indices():
return get_dsv4_c4_state_indices(
return c4_state_transfer_indices(
req_pool_idx,
seq_len,
ring_size=req_to_token_pool.get_dsv4_c4_state_ring_size(),
python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py
+3
View File
@@ -118,6 +118,7 @@ class NPUCompressStatePool(CompressStatePool):
enable_memory_saver: bool,
ratio: int,
ring_size: int,
request_scoped: bool,
swa_page_size: int,
):
assert ratio in (
@@ -139,6 +140,7 @@ class NPUCompressStatePool(CompressStatePool):
enable_memory_saver=enable_memory_saver,
ratio=ratio,
online=False,
request_scoped=request_scoped,
swa_page_size=swa_page_size,
state_cache_page_size=ring_size,
)
@@ -352,6 +354,7 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool):
device=self.device,
enable_memory_saver=enable_memory_saver,
ratio=ratio,
request_scoped=ratio == 128,
swa_page_size=self.swa_page_size,
)
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+203 -115
View File
@@ -25,6 +25,7 @@ from sglang.kernels.ops.attention.dsv4.dequant_k_cache import (
gather_dequant_requant_fp8_paged,
q8kv8_padded_num_heads,
)
from sglang.kernels.ops.attention.dsv4.kv_layout import KVLayout
from sglang.kernels.ops.attention.dsv4.metadata_kernel import (
init_compression_metadata as _init_compression_metadata_triton,
)
@@ -98,7 +99,7 @@ _is_xpu = is_xpu()
logger = logging.getLogger(__name__)
SWA_WINDOW = 128
C4_TOPK = 512
DEFAULT_INDEX_TOPK = 512
PAGE_INDEX_ALIGNED_SIZE = 64
@@ -178,7 +179,10 @@ class DSV4AttnMetadata:
swa_page_indices: torch.Tensor
swa_topk_lengths: torch.Tensor
c4_sparse_topk: int
index_topk: int
# Sorted compress ratios present in this stage; absent ratios keep no
# buffers or schedules.
present_ratios: Tuple[int, ...]
# Shared by all layer stores; locations are in SWA space.
swa_out_cache_loc: Optional[torch.Tensor] = None
c4_out_loc: Optional[torch.Tensor] = None
@@ -212,6 +216,14 @@ class DSV4AttnMetadata:
def positions(self) -> torch.Tensor:
return self.positions_casual
@property
def has_c4(self) -> bool:
return 4 in self.present_ratios
@property
def has_c128(self) -> bool:
return 128 in self.present_ratios
def get_flashmla_metadata(self, compress_ratio: Literal[0, 4, 128]):
if compress_ratio == 0:
return self.c0_flashmla_metadata
@@ -222,14 +234,63 @@ class DSV4AttnMetadata:
else:
raise ValueError(f"invalid {compress_ratio=}")
# Per-ratio extra-cache metadata is stored as flat fields; these accessors
# unify the read and write paths over the ratio.
def sparse_page_indices(self, compress_ratio: int) -> torch.Tensor:
"""Slots into the ratio's extra cache, -1 padded: the indexer's top-k for
c4, every compressed block up to the position for c128."""
if compress_ratio == 4:
return self.c4_sparse_page_indices
if compress_ratio == 128:
return self.c128_page_indices
raise ValueError(f"invalid {compress_ratio=}")
def sparse_topk_lengths(self, compress_ratio: int) -> torch.Tensor:
if compress_ratio == 4:
return self.c4_sparse_topk_lengths
if compress_ratio == 128:
return self.c128_topk_lengths_clamp1
raise ValueError(f"invalid {compress_ratio=}")
def sparse_raw_indices(self, compress_ratio: int) -> Optional[torch.Tensor]:
"""The top-k as request-local compressed positions, for the sparse
prefill workspace; allocated for prefill metadata only. Only the indexer
ratios have one (c128 remaps its page indices instead)."""
if compress_ratio == 4:
return self.c4_sparse_raw_indices
raise ValueError(f"invalid {compress_ratio=}")
def set_sparse_topk(
self,
compress_ratio: int,
*,
page_indices: torch.Tensor,
topk_lengths: torch.Tensor,
raw_indices: Optional[torch.Tensor] = None,
) -> None:
"""Writer counterpart of the accessors above."""
if compress_ratio == 4:
self.c4_sparse_page_indices = page_indices
self.c4_sparse_topk_lengths = topk_lengths
if raw_indices is not None:
self.c4_sparse_raw_indices = raw_indices
elif compress_ratio == 128:
assert raw_indices is None, "c128 has no raw top-k"
self.c128_page_indices = page_indices
self.c128_topk_lengths_clamp1 = topk_lengths
else:
raise ValueError(f"invalid {compress_ratio=}")
def copy_(self, other: DSV4AttnMetadata) -> None:
copy_metadata(
src=other,
dst=self,
check_eq_fields=[
"c4_sparse_topk",
"index_topk",
"page_size",
"cuda_int32_kwargs",
"present_ratios",
],
copy_fields=[
"raw_out_loc",
@@ -269,9 +330,10 @@ class DSV4AttnMetadata:
)
def refresh_for_breakable_cuda_graph_replay_(self, other: DSV4AttnMetadata) -> None:
assert self.c4_sparse_topk == other.c4_sparse_topk
assert self.index_topk == other.index_topk
assert self.page_size == other.page_size
assert self.cuda_int32_kwargs == other.cuda_int32_kwargs
assert self.present_ratios == other.present_ratios
tensor_copy_fields = [
"raw_out_loc",
@@ -331,26 +393,36 @@ class DSV4AttnMetadata:
f"{self.raw_out_loc.shape=}, {num_tokens=}"
)
(
self.c4_out_loc,
_,
self.c4_topk_lengths_raw,
self.c4_topk_lengths_clamp1,
self.c128_out_loc,
_,
_,
self.c128_topk_lengths_clamp1,
self.c128_page_indices,
) = _init_compression_metadata_triton(
self.seq_lens_casual,
self.positions_casual,
self.raw_out_loc,
self.page_table,
self.page_size,
compute_page_indices=True,
)
if self.has_c4 or self.has_c128:
# One kernel produces both ratios; compute_page_indices=False only
# drops the [T, max_c128_len] table, which is c128-only.
(
c4_out_loc,
_,
c4_topk_lengths_raw,
c4_topk_lengths_clamp1,
c128_out_loc,
_,
_,
c128_topk_lengths_clamp1,
c128_page_indices,
) = _init_compression_metadata_triton(
self.seq_lens_casual,
self.positions_casual,
self.raw_out_loc,
self.page_table,
self.page_size,
compute_page_indices=self.has_c128,
)
if self.has_c4:
self.c4_out_loc = c4_out_loc
self.c4_topk_lengths_raw = c4_topk_lengths_raw
self.c4_topk_lengths_clamp1 = c4_topk_lengths_clamp1
if self.has_c128:
self.c128_out_loc = c128_out_loc
self.c128_topk_lengths_clamp1 = c128_topk_lengths_clamp1
self.c128_page_indices = _pad_last_dim(c128_page_indices)
self.c128_page_indices = _pad_last_dim(self.c128_page_indices)
self.swa_page_indices = _pad_last_dim(self.swa_page_indices)
# Cache-write locations stay in global logical order and are intentionally
@@ -361,6 +433,9 @@ class DSV4AttnMetadata:
"swa_page_indices",
"swa_topk_lengths",
"page_table",
]
# Same treatment, None for stages without that compress ratio.
_CP_REINDEX_OPTIONAL_FIELDS = [
"c4_topk_lengths_raw",
"c4_topk_lengths_clamp1",
"c128_page_indices",
@@ -385,15 +460,18 @@ class DSV4AttnMetadata:
expected_local_len = pre_global_len // cp_size
if num_tokens is None:
num_tokens = pre_global_len
for field_name in self._CP_REINDEX_FIELDS:
for field_name in self._CP_REINDEX_FIELDS + self._CP_REINDEX_OPTIONAL_FIELDS:
val = getattr(self, field_name, None)
if val is None:
assert field_name in self._CP_REINDEX_OPTIONAL_FIELDS, (
f"CP reindex: {field_name} is None"
)
continue
assert isinstance(val, torch.Tensor), (
f"CP reindex: {field_name} is {type(val)}, expected Tensor"
)
setattr(self, field_name, val[idx].contiguous())
for field_name in self._CP_REINDEX_FIELDS:
val = getattr(self, field_name)
val = val[idx].contiguous()
setattr(self, field_name, val)
assert val.shape[0] == expected_local_len, (
f"apply_cp_reindex post-condition: {field_name}.shape[0]={val.shape[0]} "
f"!= expected_local_len={expected_local_len} (cp_size={cp_size})"
@@ -408,28 +486,37 @@ class DSV4AttnMetadata:
)
def init_flashmla_related(self, is_prefill: bool = False):
# c4_sparse_topk is set from model_config.index_topk per-model
# index_topk is set from model_config.index_topk per-model
# (small model: 512, large model: 1024).
assert self.c4_sparse_topk in (512, 1024), (
f"unexpected c4_sparse_topk={self.c4_sparse_topk}; "
assert self.index_topk in (512, 1024), (
f"unexpected index_topk={self.index_topk}; "
"supported: 512 (small) or 1024 (large)"
)
assert self.c4_topk_lengths_clamp1 is not None
self.c4_sparse_topk_lengths = torch.clamp(
self.c4_topk_lengths_clamp1, max=self.c4_sparse_topk
)
self.c4_sparse_page_indices = torch.full(
(self.c4_topk_lengths_clamp1.size(0), self.c4_sparse_topk),
-1,
dtype=torch.int32,
device=self.c4_topk_lengths_clamp1.device,
)
self.c4_sparse_page_indices = _pad_last_dim(self.c4_sparse_page_indices)
if is_prefill:
self.c4_sparse_raw_indices = torch.empty_like(self.c4_sparse_page_indices)
if self.has_c4:
assert self.c4_topk_lengths_clamp1 is not None
self.c4_sparse_topk_lengths = torch.clamp(
self.c4_topk_lengths_clamp1, max=self.index_topk
)
self.c4_sparse_page_indices = torch.full(
(self.c4_topk_lengths_clamp1.size(0), self.index_topk),
-1,
dtype=torch.int32,
device=self.c4_topk_lengths_clamp1.device,
)
self.c4_sparse_page_indices = _pad_last_dim(self.c4_sparse_page_indices)
if is_prefill:
self.c4_sparse_raw_indices = torch.empty_like(
self.c4_sparse_page_indices
)
else:
self.c4_sparse_topk_lengths = None
self.c4_sparse_page_indices = None
self.c4_sparse_raw_indices = None
self.c0_flashmla_metadata = _create_flashmla_metadata()
self.c4_flashmla_metadata = _create_flashmla_metadata()
self.c128_flashmla_metadata = _create_flashmla_metadata()
self.c4_flashmla_metadata = _create_flashmla_metadata() if self.has_c4 else None
self.c128_flashmla_metadata = (
_create_flashmla_metadata() if self.has_c128 else None
)
def init_trtllm_sparse_buffers(self) -> None:
"""Build decode tables with 128 SWA columns followed by compressed KV.
@@ -638,11 +725,15 @@ class DeepseekV4AttnBackend(
self.token_to_kv_pool: DeepSeekV4TokenToKVPool = model_runner.token_to_kv_pool
self.hisparse_coordinator = model_runner.hisparse_coordinator
self.req_to_token = model_runner.req_to_token_pool.req_to_token
# Nothing is built for a compress ratio outside the pool's set.
self.present_ratios: Tuple[int, ...] = self.token_to_kv_pool.present_ratios
self.has_c4: bool = 4 in self.present_ratios
self.has_c128: bool = 128 in self.present_ratios
self.MAX_SEQ_LEN_FOR_CAPTURE = self.req_to_token.shape[1]
assert isinstance(self.token_to_kv_pool, DeepSeekV4TokenToKVPool)
self.c4_topk = getattr(
model_runner.model_config.hf_text_config, "index_topk", C4_TOPK
self.index_topk = getattr(
model_runner.model_config.hf_text_config, "index_topk", DEFAULT_INDEX_TOPK
)
kernel = get_exec().kernel
@@ -750,7 +841,7 @@ class DeepseekV4AttnBackend(
use_prefill_cuda_graph: bool,
online_c128_state_slot_offset: int,
) -> Optional[FusedCompressMetadata]:
if not self.online_c128_mtp.enabled():
if not self.has_c128 or not self.online_c128_mtp.enabled():
return None
assert seq_lens_cpu is not None
@@ -861,7 +952,7 @@ class DeepseekV4AttnBackend(
core_attn_metadata,
use_prefill_cuda_graph=use_prefill_cuda_graph,
)
if need_compress
if need_compress and self.has_c4
else None
)
if not need_compress:
@@ -904,13 +995,13 @@ class DeepseekV4AttnBackend(
online_state_slot_offset=online_c128_state_slot_offset,
)
c4_compress_metadata = create(compress_ratio=4)
c128_compress_metadata = create(compress_ratio=128)
return DSV4Metadata(
core_attn_metadata,
indexer_metadata,
c4_compress_metadata=c4_compress_metadata,
c128_compress_metadata=c128_compress_metadata,
c4_compress_metadata=create(compress_ratio=4) if self.has_c4 else None,
c128_compress_metadata=(
create(compress_ratio=128) if self.has_c128 else None
),
)
def init_forward_metadata_target_verify(
@@ -1045,7 +1136,11 @@ class DeepseekV4AttnBackend(
out_loc=out_cache_loc,
need_compress=True,
)
indexer_metadata = self.init_forward_metadata_indexer(core_attn_metadata)
indexer_metadata = (
self.init_forward_metadata_indexer(core_attn_metadata)
if self.has_c4
else None
)
create = functools.partial(
create_paged_compressor_data,
is_prefill=True,
@@ -1061,12 +1156,12 @@ class DeepseekV4AttnBackend(
online_state_slot_offset=online_c128_state_slot_offset,
)
c128_compress_metadata = raw_metadata.c128_compress_metadata
if c128_compress_metadata is None:
if c128_compress_metadata is None and self.has_c128:
c128_compress_metadata = create(compress_ratio=128)
return DSV4Metadata(
core_attn_metadata,
indexer_metadata,
c4_compress_metadata=create(compress_ratio=4),
c4_compress_metadata=create(compress_ratio=4) if self.has_c4 else None,
c128_compress_metadata=c128_compress_metadata,
)
@@ -1086,7 +1181,11 @@ class DeepseekV4AttnBackend(
out_loc=out_cache_loc,
need_compress=True,
)
indexer_metadata = self.init_forward_metadata_indexer(core_attn_metadata)
indexer_metadata = (
self.init_forward_metadata_indexer(core_attn_metadata)
if self.has_c4
else None
)
create = functools.partial(
create_paged_compressor_data,
@@ -1100,8 +1199,10 @@ class DeepseekV4AttnBackend(
return DSV4Metadata(
core_attn_metadata,
indexer_metadata,
c4_compress_metadata=create(compress_ratio=4),
c128_compress_metadata=create(compress_ratio=128),
c4_compress_metadata=create(compress_ratio=4) if self.has_c4 else None,
c128_compress_metadata=(
create(compress_ratio=128) if self.has_c128 else None
),
)
def init_forward_metadata_draft_extend(
@@ -1464,17 +1565,22 @@ class DeepseekV4AttnBackend(
)
if use_sparse_prefill:
metadata.sparse_prefill_cache = self._build_sparse_prefill_chunk_cache(
forward_batch, num_qo_tokens=num_qo_tokens
forward_batch, metadata.core_attn_metadata, num_qo_tokens=num_qo_tokens
)
# Marked for dense prefill too: that path reads only core_attn_metadata,
# which init_forward_metadata already snapshotted.
metadata.prefill_shared_reads_snapshotted = True
def _build_sparse_prefill_chunk_cache(
self, forward_batch: ForwardBatch, *, num_qo_tokens: int
self,
forward_batch: ForwardBatch,
core_attn_metadata: DSV4AttnMetadata,
*,
num_qo_tokens: int,
) -> SparsePrefillChunkCache:
seq_lens_cpu = forward_batch.seq_lens_cpu
assert seq_lens_cpu is not None
extend_seq_lens = forward_batch.extend_seq_lens
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
assert extend_seq_lens_cpu is not None
seq_lens_cpu_list = seq_lens_cpu.tolist()
@@ -1484,9 +1590,13 @@ class DeepseekV4AttnBackend(
seq_lens_cpu_list, extend_seq_lens_cpu, strict=True
)
)
# The rows this forward runs are the extend, one per causal position.
query_pos = core_attn_metadata.seq_lens_casual[:num_qo_tokens] - 1
return SparsePrefillChunkCache.build(
seq_lens=forward_batch.seq_lens.to(torch.int32),
extend_seq_lens=forward_batch.extend_seq_lens.to(torch.int32),
extend_seq_lens=extend_seq_lens.to(torch.int32),
query_lens=extend_seq_lens.to(torch.int32),
query_pos=query_pos,
req_pool_indices=forward_batch.req_pool_indices.to(torch.int32),
req_to_token=self.req_to_token,
full_to_swa=self.token_to_kv_pool.full_to_swa_index_mapping,
@@ -1713,8 +1823,10 @@ class DeepseekV4AttnBackend(
):
core = metadata.core_attn_metadata
core.c0_flashmla_metadata = _create_flashmla_metadata()
core.c4_flashmla_metadata = _create_flashmla_metadata()
core.c128_flashmla_metadata = _create_flashmla_metadata()
if core.has_c4:
core.c4_flashmla_metadata = _create_flashmla_metadata()
if core.has_c128:
core.c128_flashmla_metadata = _create_flashmla_metadata()
# PREP_IN_CUDA_GRAPH=True: warmup upgraded raw->full on the host;
# restore raw so capture re-runs the upgrade inside the graph.
@@ -1777,31 +1889,33 @@ class DeepseekV4AttnBackend(
swa_k_cache = token_to_kv_pool.get_swa_key_buffer_radix(layer_id)
extra_k_cache, extra_indices, extra_topk_lengths = None, None, None
if compress_ratio == 4:
if compress_ratio != 0:
extra_k_cache = token_to_kv_pool.get_extra_key_buffer(layer_id)
extra_indices = core_attn_metadata.c4_sparse_page_indices
extra_topk_lengths = core_attn_metadata.c4_sparse_topk_lengths
elif compress_ratio == 128:
extra_k_cache = token_to_kv_pool.get_extra_key_buffer(layer_id)
extra_indices = core_attn_metadata.c128_page_indices
extra_topk_lengths = core_attn_metadata.c128_topk_lengths_clamp1
extra_indices = core_attn_metadata.sparse_page_indices(compress_ratio)
extra_topk_lengths = core_attn_metadata.sparse_topk_lengths(
compress_ratio
)
swa_page_size = token_to_kv_pool.swa_page_size
assert swa_k_cache.ndim == 2
k_cache_total_dim = token_to_kv_pool.swa_kv_pool.kv_cache_total_dim
# The kernel detects each cache's format from the last dim of this view.
k_cache_total_dim = token_to_kv_pool.get_swa_key_bytes_per_token()
swa_k_cache = swa_k_cache[:, : swa_page_size * k_cache_total_dim].view(
swa_k_cache.shape[0], swa_page_size, 1, k_cache_total_dim
)
if extra_k_cache is not None:
extra_page_size = token_to_kv_pool.get_extra_key_page_size(layer_id)
extra_total_dim = token_to_kv_pool.get_extra_key_bytes_per_token(
layer_id
)
extra_k_cache = extra_k_cache[
:, : extra_page_size * k_cache_total_dim
:, : extra_page_size * extra_total_dim
].view(
extra_k_cache.shape[0],
extra_page_size,
1,
k_cache_total_dim,
extra_total_dim,
)
swa_page_indices = core_attn_metadata.swa_page_indices
swa_topk_lengths = core_attn_metadata.swa_topk_lengths
@@ -1966,7 +2080,7 @@ class DeepseekV4AttnBackend(
cache = self.forward_metadata.sparse_prefill_cache
if cache is None:
cache = self._build_sparse_prefill_chunk_cache(
forward_batch, num_qo_tokens=q_flat.shape[0]
forward_batch, core_attn_metadata, num_qo_tokens=q_flat.shape[0]
)
self.forward_metadata.sparse_prefill_cache = cache
@@ -1984,24 +2098,9 @@ class DeepseekV4AttnBackend(
else:
extra_page_size = token_to_kv_pool.get_extra_key_page_size(layer_id)
extra_k_cache = token_to_kv_pool.get_extra_key_buffer(layer_id)
if compress_ratio == 128:
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
combined_indices = cache.c128_combined_indices
combined_lens = cache.c128_combined_lens
else:
assert core_attn_metadata.c4_sparse_raw_indices is not None, (
"sparse-prefill c4 path requires c4_sparse_raw_indices "
"(allocated in init_flashmla_related when is_prefill=True)"
)
cache.ensure_c4(core_attn_metadata.page_table, extra_page_size)
flat_token_ids = cache.c4_flat_token_ids
combined_indices, combined_lens = cache.combine_c4_layer(
c4_sparse_raw_indices=core_attn_metadata.c4_sparse_raw_indices[
: cache.num_qo_tokens
],
)
flat_token_ids, combined_indices, combined_lens = cache.layer_inputs(
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]
@@ -2015,12 +2114,14 @@ class DeepseekV4AttnBackend(
flat_token_ids,
page_size=extra_page_size,
out=compressed_slice,
layout=token_to_kv_pool.get_extra_key_layout(layer_id),
)
dequantize_k_cache_paged(
token_to_kv_pool.get_swa_key_buffer_radix(layer_id),
cache.swa_token_ids,
page_size=cache.swa_page_size,
out=swa_slice,
layout=token_to_kv_pool.get_swa_key_layout(),
)
kv = workspace
@@ -2141,7 +2242,7 @@ class DeepseekV4AttnBackend(
cache = self.forward_metadata.sparse_prefill_cache
if cache is None:
cache = self._build_sparse_prefill_chunk_cache(
forward_batch, num_qo_tokens=q_flat.shape[0]
forward_batch, core_attn_metadata, num_qo_tokens=q_flat.shape[0]
)
self.forward_metadata.sparse_prefill_cache = cache
@@ -2161,25 +2262,9 @@ class DeepseekV4AttnBackend(
else:
extra_page_size = token_to_kv_pool.get_extra_key_page_size(layer_id)
extra_k_cache = token_to_kv_pool.get_extra_key_buffer(layer_id)
if compress_ratio == 128:
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
combined_indices = cache.c128_combined_indices
combined_lens = cache.c128_combined_lens
else:
assert core_attn_metadata.c4_sparse_raw_indices is not None, (
"Q8KV8 sparse-prefill c4 path requires c4_sparse_raw_indices "
"(allocated in init_flashmla_related when is_prefill=True)"
)
cache.ensure_c4(core_attn_metadata.page_table, extra_page_size)
flat_token_ids = cache.c4_flat_token_ids
combined_indices, combined_lens = cache.combine_c4_layer(
c4_sparse_raw_indices=core_attn_metadata.c4_sparse_raw_indices[
: cache.num_qo_tokens
],
)
flat_token_ids, combined_indices, combined_lens = cache.layer_inputs(
compress_ratio, core_attn_metadata, extra_page_size
)
n_compressed = flat_token_ids.shape[0]
workspace = self.sparse_prefill_workspace.get(
@@ -2189,6 +2274,8 @@ class DeepseekV4AttnBackend(
compressed_slice = workspace[:n_compressed]
swa_slice = workspace[n_compressed:]
# The Q8KV8 gather reads the 584-byte V4 layout only (its kernel is SM90).
assert token_to_kv_pool.get_swa_key_layout() is KVLayout.V4
if compressed_slice is not None:
gather_dequant_requant_fp8_paged(
extra_k_cache,
@@ -2352,7 +2439,8 @@ class DeepseekV4AttnBackend(
page_table=page_table,
swa_page_indices=swa_page_indices,
swa_topk_lengths=swa_topk_lengths,
c4_sparse_topk=self.c4_topk,
index_topk=self.index_topk,
present_ratios=self.present_ratios,
)
if need_compress:
python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py
+12 -14
View File
@@ -58,7 +58,7 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
SWA_WINDOW = 128
C4_TOPK = 512
DEFAULT_INDEX_TOPK = 512
PAGE_INDEX_ALIGNED_SIZE = 64
@@ -186,7 +186,7 @@ class DSV4AttnMetadata:
swa_page_indices: torch.Tensor
swa_topk_lengths: torch.Tensor
c4_sparse_topk: int
index_topk: int
# Shared by all layer stores; locations are in SWA space.
swa_out_cache_loc: Optional[torch.Tensor] = None
c4_out_loc: Optional[torch.Tensor] = None
@@ -228,7 +228,7 @@ class DSV4AttnMetadata:
src=other,
dst=self,
check_eq_fields=[
"c4_sparse_topk",
"index_topk",
"page_size",
"cuda_int32_kwargs",
],
@@ -263,7 +263,7 @@ class DSV4AttnMetadata:
)
def refresh_for_breakable_cuda_graph_replay_(self, other: DSV4AttnMetadata) -> None:
assert self.c4_sparse_topk == other.c4_sparse_topk
assert self.index_topk == other.index_topk
assert self.page_size == other.page_size
assert self.cuda_int32_kwargs == other.cuda_int32_kwargs
@@ -393,22 +393,20 @@ class DSV4AttnMetadata:
)
def init_flashmla_related(self, is_prefill: bool = False):
# c4_sparse_topk is set from model_config.index_topk per-model
# (small model: 512, large model: 1024).
assert self.c4_sparse_topk in (512, 1024), (
f"unexpected c4_sparse_topk={self.c4_sparse_topk}; "
assert self.index_topk in (512, 1024), (
f"unexpected index_topk={self.index_topk}; "
"supported: 512 (small) or 1024 (large)"
)
assert self.c4_topk_lengths_clamp1 is not None
self.c4_sparse_topk_lengths = torch.clamp(
self.c4_topk_lengths_clamp1, max=self.c4_sparse_topk
self.c4_topk_lengths_clamp1, max=self.index_topk
)
assert self.c4_topk_lengths_raw is not None
self.c4_sparse_topk_lengths_raw = torch.clamp(
self.c4_topk_lengths_raw, max=self.c4_sparse_topk
self.c4_topk_lengths_raw, max=self.index_topk
)
self.c4_sparse_page_indices = torch.full(
(self.c4_topk_lengths_clamp1.size(0), self.c4_sparse_topk),
(self.c4_topk_lengths_clamp1.size(0), self.index_topk),
-1,
dtype=torch.int32,
device=self.c4_topk_lengths_clamp1.device,
@@ -574,8 +572,8 @@ class DeepseekV4HipRadixBackend(
self.MAX_SEQ_LEN_FOR_CAPTURE = self.req_to_token.shape[1]
assert isinstance(self.token_to_kv_pool, DeepSeekV4TokenToKVPool)
self.c4_topk = getattr(
model_runner.model_config.hf_text_config, "index_topk", C4_TOPK
self.index_topk = getattr(
model_runner.model_config.hf_text_config, "index_topk", DEFAULT_INDEX_TOPK
)
self.enable_deepseek_v4_fp4_indexer: bool = (
get_exec().kernel.enable_deepseek_v4_fp4_indexer
@@ -2047,7 +2045,7 @@ class DeepseekV4HipRadixBackend(
page_table=page_table,
swa_page_indices=swa_page_indices,
swa_topk_lengths=swa_topk_lengths,
c4_sparse_topk=self.c4_topk,
index_topk=self.index_topk,
)
if need_compress:
python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py
+147 -115
View File
@@ -34,7 +34,7 @@ compressed branch becomes a no-op) and any ``compress_ratio >= 1``.
import os
from dataclasses import dataclass, field
from typing import Optional
from typing import Dict, Optional
import torch
import triton
@@ -114,6 +114,7 @@ def combined_topk_width(topk: int, window_size: int) -> int:
def combine_topk_swa_indices(
topk_indices: torch.Tensor,
query_start_loc: torch.Tensor,
query_pos: torch.Tensor,
seq_lens: torch.Tensor,
gather_lens: torch.Tensor,
compressed_base: torch.Tensor,
@@ -124,43 +125,22 @@ def combine_topk_swa_indices(
out_indices: Optional[torch.Tensor] = None,
out_lens: Optional[torch.Tensor] = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Combine topk + SWA indices into a single ``flash_mla_sparse_fwd`` row.
"""Combine top-k and SWA indices for flash_mla_sparse_fwd.
Args:
topk_indices: (num_tokens, K) int32. Per-query indices into the
compressed-cache region, **already in request-local space** —
i.e. in ``[0, compressed_gather_len[r])`` for the request that
owns each token. Pad entries can be any value; they are ignored
beyond ``topk_len``.
query_start_loc: (num_reqs+1,) int32. Cumulative query lengths; may
be in global (cross-chunk) space — kernel rebases by subtracting
``query_start_loc[0]``.
seq_lens: (num_reqs,) int32. Each request's full sequence length.
gather_lens: (num_reqs,) int32. Trailing tokens dequanted into the
SWA region for that request.
compressed_base: (num_reqs,) int32. Flat workspace offset where
request r's compressed region begins. Pass all-zeros (or any
value) for SWA-only layers since topk=0 disables this branch.
swa_base: (num_reqs,) int32. Flat workspace offset where request
r's SWA region begins.
window_size: SWA window size.
compress_ratio: must be ``>= 1`` even when topk==0.
topk: configured topk; pass 0 for SWA-only layers.
out_indices: optional preallocated ``(num_tokens, combined_topk)``
int32 buffer. If provided, the kernel writes the per-query prefix
``[0, topk_len + swa_len)``; positions beyond are not touched.
Caller must pre-fill with ``-1`` sentinels (and the chunk-invariant
valid-prefix length must hold across reuses).
out_lens: optional preallocated ``(num_tokens,)`` int32 buffer; the
kernel fully overwrites it, so any dtype-correct buffer works.
Top-k indices are int32 [num_tokens, K] in each request's compressed region.
Invalid entries within topk_len must be -1; later entries are ignored.
query_start_loc can include a cross-chunk offset; query_pos is absolute.
compressed_base and swa_base address the flat workspace. SWA-only layers use
topk == 0, but compress_ratio must still be positive.
Returns:
combined_indices: (num_tokens, padded_topk_swa) int32, padded to a
multiple of 128 with -1 sentinels.
combined_lens: (num_tokens,) int32, valid prefix length per token.
Returns int32 indices [num_tokens, padded_topk_swa] and per-token scanned-prefix
lengths, including -1 entries skipped by attention. Width is padded to 128.
Preallocated out_indices must contain -1 outside the written prefix;
reuse requires chunk-invariant scanned-prefix lengths. out_lens is overwritten.
"""
assert topk_indices.dtype == torch.int32
assert query_start_loc.dtype == torch.int32
assert query_pos.dtype == torch.int32
assert seq_lens.dtype == torch.int32
assert gather_lens.dtype == torch.int32
assert compressed_base.dtype == torch.int32
@@ -201,6 +181,7 @@ def combine_topk_swa_indices(
topk_indices,
topk_indices.stride(0),
query_start_loc,
query_pos,
seq_lens,
gather_lens,
compressed_base,
@@ -284,13 +265,25 @@ def build_swa_token_ids(
@dataclass
class SparsePrefillChunkCache:
"""Chunk-invariant scaffolding for ``_forward_prefill_sparse``.
class CompressedGather:
"""Positional layout of one compressed cache inside the workspace."""
The fields here depend only on the prefill chunk (forward_batch,
req_to_token, full_to_swa_index_mapping, and the c4/c128 page tables)
and not on the per-layer k_cache. Reused across every layer in the
chunk to avoid rebuilding tiny tensors 61 times per forward pass.
flat_token_ids: torch.Tensor # (num_reqs * c_max,) int32
compressed_base: torch.Tensor # (num_reqs,) int32
swa_base: torch.Tensor # (num_reqs,) int32
# Tail stays at the -1 sentinel because the valid prefix length is
# chunk-invariant per request; subsequent layers only overwrite that prefix.
combined_indices: Optional[torch.Tensor] = None
combined_lens: Optional[torch.Tensor] = None
@dataclass
class SparsePrefillChunkCache:
"""Cache prefill-chunk metadata shared across layers.
Fields depend on request/token mappings and compressed page tables, not
per-layer k_cache; per-layer top-k combinations are recomputed into reused
buffers.
"""
# Geometry computed once per chunk.
@@ -309,7 +302,8 @@ class SparsePrefillChunkCache:
# ``page_size`` so that ``slot // page_size`` recovers the right page.
swa_page_size: int
seq_lens: torch.Tensor # (num_reqs,) int32
query_start_loc: torch.Tensor # (num_reqs+1,) int32
query_start_loc: torch.Tensor # (num_reqs+1,) int32, query rows per request
query_pos: torch.Tensor # (num_qo_tokens,) int32, sequence position per row
# SWA-side (every layer needs these, all chunk-invariant).
swa_token_ids: torch.Tensor # (total_swa,) int32
@@ -320,26 +314,17 @@ 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)
# 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
# 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
# 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
c4_combined_lens: Optional[torch.Tensor] = None
# Compressed caches keyed by compress ratio: c128 (every block, combined once
# per chunk) and the top-k ratios (combined per layer).
compressed: Dict[int, CompressedGather] = field(default_factory=dict)
@classmethod
def build(
cls,
seq_lens: torch.Tensor,
extend_seq_lens: torch.Tensor,
query_lens: torch.Tensor,
query_pos: torch.Tensor,
req_pool_indices: torch.Tensor,
req_to_token: torch.Tensor,
full_to_swa: torch.Tensor,
@@ -349,11 +334,13 @@ class SparsePrefillChunkCache:
max_seq_len: int,
total_swa: int,
) -> "SparsePrefillChunkCache":
"""``query_lens`` / ``query_pos``: the rows this forward runs (the extend, or
a CP rank's interleaved share of it); the SWA gather spans the whole extend."""
device = seq_lens.device
num_reqs = seq_lens.shape[0]
query_start_loc = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device)
query_start_loc[1:] = torch.cumsum(extend_seq_lens, dim=0).to(torch.int32)
query_start_loc[1:] = torch.cumsum(query_lens, dim=0).to(torch.int32)
swa_token_ids, swa_first_pos, swa_gather_lens, swa_offsets = (
build_swa_token_ids(
@@ -375,6 +362,7 @@ class SparsePrefillChunkCache:
swa_page_size=swa_page_size,
seq_lens=seq_lens,
query_start_loc=query_start_loc,
query_pos=query_pos,
swa_token_ids=swa_token_ids,
swa_first_pos=swa_first_pos,
swa_gather_lens=swa_gather_lens,
@@ -389,6 +377,7 @@ class SparsePrefillChunkCache:
cache.c0_combined_indices, cache.c0_combined_lens = combine_topk_swa_indices(
topk_indices=zero_topk,
query_start_loc=query_start_loc,
query_pos=query_pos,
seq_lens=seq_lens,
gather_lens=swa_gather_lens,
compressed_base=zero_compressed_base,
@@ -399,7 +388,45 @@ class SparsePrefillChunkCache:
)
return cache
def ensure_c128(self, c128_page_indices: torch.Tensor) -> None:
def _workspace_bases(self, c_max: int) -> tuple[torch.Tensor, torch.Tensor]:
"""Flat workspace offsets of each request's compressed and SWA regions:
``c_max`` compressed slots per request, then the SWA gather."""
device = self.seq_lens.device
compressed_base = (
torch.arange(self.num_reqs, dtype=torch.int32, device=device) * c_max
).to(torch.int32)
swa_base = (self.num_reqs * c_max + self.swa_offsets[:-1]).to(torch.int32)
return compressed_base, swa_base
def layer_inputs(
self,
compress_ratio: int,
core_attn_metadata,
c_page_size: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""``(flat_token_ids, combined_indices, combined_lens)`` for one compressed
layer. c128 gathers every block and combines once per chunk; a top-k ratio
gathers from the page table once and combines per layer from the indexer's
raw top-k."""
if compress_ratio == 128:
page_indices = core_attn_metadata.sparse_page_indices(128)
assert page_indices is not None
gather = self.ensure_c128(page_indices)
return gather.flat_token_ids, gather.combined_indices, gather.combined_lens
raw_indices = core_attn_metadata.sparse_raw_indices(compress_ratio)
assert raw_indices is not None, (
f"sparse-prefill c{compress_ratio} path requires the raw top-k indices "
"(allocated in init_flashmla_related when is_prefill=True)"
)
gather = self.ensure_compressed(
compress_ratio, core_attn_metadata.page_table, c_page_size
)
combined_indices, combined_lens = self.combine_compressed(
compress_ratio, raw_indices[: self.num_qo_tokens]
)
return gather.flat_token_ids, combined_indices, combined_lens
def ensure_c128(self, c128_page_indices: torch.Tensor) -> CompressedGather:
"""Populate c128-side fields from per-query c128 page indices.
``c128_page_indices[q, j]`` carries slot ids derived from
@@ -412,25 +439,24 @@ class SparsePrefillChunkCache:
clamp_min(0) collapses to slot 0, sending dequant to a polluted
slot and producing garbage c128 entries.
"""
if self.c128_flat_token_ids is not None:
return
gather = self.compressed.get(128)
if gather is not None:
return gather
device = self.seq_lens.device
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()
# A request without rows on this rank gathers into a region nothing reads.
last_q_per_req = (self.query_start_loc[1:] - 1).clamp_min(0).long()
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)
compressed_base = (
torch.arange(self.num_reqs, dtype=torch.int32, device=device) * c128_max
).to(torch.int32)
total_compressed = self.num_reqs * c128_max
compressed_base, swa_base = self._workspace_bases(c128_max)
# Pre-compute the c128 combine output. topk_indices[q, j] = j is the
# arange-broadcast pattern; we materialize it once here so the
# combine kernel can read it like any other topk tensor.
@@ -439,10 +465,10 @@ class SparsePrefillChunkCache:
.expand(self.num_qo_tokens, -1)
.contiguous()
)
swa_base = (total_compressed + self.swa_offsets[:-1]).to(torch.int32)
combined_indices, combined_lens = combine_topk_swa_indices(
topk_indices=topk_indices,
query_start_loc=self.query_start_loc,
query_pos=self.query_pos,
seq_lens=self.seq_lens,
gather_lens=self.swa_gather_lens,
compressed_base=compressed_base,
@@ -452,89 +478,95 @@ class SparsePrefillChunkCache:
topk=c128_max,
)
self.c128_flat_token_ids = flat_c128_ids
self.c128_combined_indices = combined_indices
self.c128_combined_lens = combined_lens
gather = CompressedGather(
flat_token_ids=flat_c128_ids,
compressed_base=compressed_base,
swa_base=swa_base,
combined_indices=combined_indices,
combined_lens=combined_lens,
)
self.compressed[128] = gather
return gather
def ensure_c4(
def ensure_compressed(
self,
compress_ratio: int,
page_table: torch.Tensor,
c4_page_size: int,
) -> None:
"""Populate c4-side fields from the per-query page table.
``page_table`` is (num_qo_tokens, max_blocks); rows within a request
c_page_size: int,
) -> CompressedGather:
"""``page_table`` is (num_qo_tokens, max_blocks); rows within a request
are duplicates. The combine output is per-layer (depends on the
layer's remapped topk_indices), so we only cache the gather-side
scaffolding plus compressed/swa bases.
"""
if self.c4_flat_token_ids is not None:
return
gather = self.compressed.get(compress_ratio)
if gather is not None:
return gather
device = self.seq_lens.device
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}"
c_max = max(self.max_seq_len // compress_ratio, 1)
c_capacity = page_table.shape[-1] * c_page_size
assert c_max <= c_capacity, (
f"live c{compress_ratio} extent {c_max} exceeds metadata capacity {c_capacity}"
)
first_q_per_req = self.query_start_loc[:-1].long()
num_blocks = (c4_max + c4_page_size - 1) // c4_page_size
first_q_per_req = (
self.query_start_loc[:-1].clamp_max(self.num_qo_tokens - 1).long()
)
num_blocks = (c_max + c_page_size - 1) // c_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()
in_page = (k_arange % c4_page_size).to(torch.int32)
c4_token_ids_2d = (
per_req_page_table.index_select(1, block_idx) * c4_page_size + in_page
k_arange = torch.arange(c_max, dtype=torch.int32, device=device)
block_idx = (k_arange // c_page_size).long()
in_page = (k_arange % c_page_size).to(torch.int32)
token_ids_2d = (
per_req_page_table.index_select(1, block_idx) * c_page_size + in_page
).to(torch.int32)
flat_c4_ids = c4_token_ids_2d.reshape(-1).clamp_min(0)
total_compressed = self.num_reqs * c4_max
compressed_base = (
torch.arange(self.num_reqs, dtype=torch.int32, device=device) * c4_max
).to(torch.int32)
swa_base = (total_compressed + self.swa_offsets[:-1]).to(torch.int32)
flat_ids = token_ids_2d.reshape(-1).clamp_min(0)
compressed_base, swa_base = self._workspace_bases(c_max)
self.c4_flat_token_ids = flat_c4_ids
self.c4_page_size = c4_page_size
self.c4_compressed_base = compressed_base
self.c4_swa_base = swa_base
gather = CompressedGather(
flat_token_ids=flat_ids,
compressed_base=compressed_base,
swa_base=swa_base,
)
self.compressed[compress_ratio] = gather
return gather
def combine_c4_layer(
def combine_compressed(
self,
c4_sparse_raw_indices: torch.Tensor,
compress_ratio: int,
sparse_raw_indices: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Per-layer combine for c4. ``c4_sparse_raw_indices`` is the topk
kernel's positional output (``block_in_seq * c_page_size + in_page``)
— already in the request-local workspace coordinate that
``combine_topk_swa_indices`` expects, so no remap is needed.
Reuses preallocated ``c4_combined_indices`` / ``c4_combined_lens``
buffers across layers — the kernel only overwrites the valid prefix.
"""``sparse_raw_indices`` is the top-k as request-local compressed
positions, already the workspace coordinate ``combine_topk_swa_indices``
expects, so no remap is needed.
"""
topk = c4_sparse_raw_indices.shape[-1]
if self.c4_combined_indices is None:
gather = self.compressed[compress_ratio]
topk = sparse_raw_indices.shape[-1]
if gather.combined_indices is None:
device = self.seq_lens.device
self.c4_combined_indices = torch.full(
gather.combined_indices = torch.full(
(self.num_qo_tokens, combined_topk_width(topk, self.swa_window_size)),
-1,
dtype=torch.int32,
device=device,
)
self.c4_combined_lens = torch.zeros(
gather.combined_lens = torch.zeros(
self.num_qo_tokens, dtype=torch.int32, device=device
)
return combine_topk_swa_indices(
topk_indices=c4_sparse_raw_indices,
topk_indices=sparse_raw_indices,
query_start_loc=self.query_start_loc,
query_pos=self.query_pos,
seq_lens=self.seq_lens,
gather_lens=self.swa_gather_lens,
compressed_base=self.c4_compressed_base,
swa_base=self.c4_swa_base,
compressed_base=gather.compressed_base,
swa_base=gather.swa_base,
window_size=self.swa_window_size,
compress_ratio=4,
compress_ratio=compress_ratio,
topk=topk,
out_indices=self.c4_combined_indices,
out_lens=self.c4_combined_lens,
out_indices=gather.combined_indices,
out_lens=gather.combined_lens,
)
python/sglang/srt/mem_cache/deepseek_v4_compress_state.py
+62
View File
@@ -4,6 +4,7 @@ import dataclasses
from contextlib import nullcontext
from math import gcd
import numpy as np
import torch
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
@@ -80,6 +81,51 @@ class KVAndScore:
return KVAndScore(torch.cat([v.kv_score for v in tensors], dim=dim))
def c4_state_transfer_indices(
req_pool_idx: int,
seq_len: int,
*,
ring_size: int,
) -> np.ndarray:
"""PD transfer rows of the overlap C4 state: the live tail of the request's ring."""
# Prefill and decode can have different ring sizes (8 or 16 with EAGLE/MTP);
# pair the overlap compressor's live rows by logical token position.
if ring_size < 8 or ring_size % 4 != 0:
raise ValueError(
f"C4 ring_size must be a multiple of 4 and at least 8, got {ring_size}"
)
seq_len = max(0, int(seq_len))
state_len = seq_len % 4 + 4
positions = np.arange(max(0, seq_len - state_len), seq_len, dtype=np.int64)
rows = int(req_pool_idx) * int(ring_size) + positions % int(ring_size)
return rows.astype(np.int32)
def request_scoped_state_transfer_indices(
req_pool_idx: int,
seq_len: int,
*,
ratio: int,
online: bool,
ring_size: int,
) -> np.ndarray:
"""PD transfer indices of a request-scoped compress state: the one pending
partial block of ``ratio`` tokens, as the request's single online row or the
ring page that holds it. Nothing pends at a block boundary."""
if seq_len == 0 or seq_len % ratio == 0:
return np.empty((0,), dtype=np.int32)
if online:
return np.array([int(req_pool_idx)], dtype=np.int32)
assert ring_size % ratio == 0, (
f"ring_size must be a multiple of {ratio}, got {ring_size}"
)
pages_per_req = ring_size // ratio
page = int(req_pool_idx) * pages_per_req + ((seq_len - 1) % ring_size) // ratio
return np.array([page], dtype=np.int32)
class CompressStatePool:
def __init__(
self,
@@ -92,11 +138,16 @@ class CompressStatePool:
enable_memory_saver: bool,
ratio: int,
online: bool = False,
request_scoped: bool = False,
swa_page_size: int = 0,
online_mtp_max_draft_tokens: int = 0,
state_cache_page_size: int = 1,
):
self.ratio = ratio
# Request-scoped state is addressed by req_pool_idx (one ring per request
# slot) and travels on the PD request-state component; page-scoped state
# follows the SWA pages. The pool factory decides which ratios are which.
self.request_scoped = request_scoped
self.ring_size = ring_size
self.swa_page_size = swa_page_size
self.page_size = state_cache_page_size
@@ -143,6 +194,17 @@ class CompressStatePool:
else:
self.kv_score_buffer[-1].clear()
def transfer_indices(self, req_pool_idx: int, seq_len: int) -> np.ndarray:
"""PD transfer indices of this pool's state for one request."""
assert self.request_scoped, "page-scoped state travels with the SWA pages"
return request_scoped_state_transfer_indices(
req_pool_idx,
seq_len,
ratio=self.ratio,
online=self.online,
ring_size=self.ring_size,
)
def _alloc_kv_score_buffer(
self, *, dtype: torch.dtype, device: str, enable_memory_saver: bool
) -> None:
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+51 -7
View File
@@ -16,6 +16,7 @@ from sglang.kernels.ops.attention.dsv4 import (
index_buf_accessor as dsv4_index_buf_accessor,
)
from sglang.kernels.ops.attention.dsv4.index_buf_accessor import NopeFp8RopeBf16Pack
from sglang.kernels.ops.attention.dsv4.kv_layout import KVLayout
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels import layout
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.environ import envs
@@ -69,6 +70,9 @@ def get_swa_ring_size(sliding_window: int, is_speculative: bool = False) -> int:
class DeepSeekV4SingleKVPool(KVCache):
# Paged FlashMLA main-KV format of this pool's rows.
kv_layout: KVLayout = KVLayout.V4
def __init__(
self,
size: int,
@@ -834,6 +838,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
kv_pool_cls=kv_pool_cls,
)
# The distinct compress ratios this stage has, sorted. Registry pools kept
# for a ratio the model lacks (wire-layout alignment) do not count.
model_ratios = set(self.compression_ratios)
self.present_ratios: Tuple[int, ...] = tuple(
ratio for ratio in sorted(self.kv_pools) if ratio in model_ratios
)
self._init_compressed_layer_mapping()
self._init_paged_compress_states(enable_memory_saver)
@@ -1012,9 +1023,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
self.indexer_compress_state_pools,
]:
for pool in pools:
if pool is None:
continue
if pool.ratio == 128:
if pool is None or pool.request_scoped:
continue
t = pool.kv_score_buffer.kv_score
assert t.ndim == 2, f"expected 2D buffer, got {t.ndim}D"
@@ -1031,7 +1040,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
data_lens: List[int] = []
item_lens: List[int] = []
for pool in self.compress_state_pools:
if pool is None or pool.ratio != 128:
if pool is None or not pool.request_scoped:
continue
t = pool.kv_score_buffer.kv_score
assert t.ndim == 2, f"expected 2D buffer, got {t.ndim}D"
@@ -1163,6 +1172,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
enable_memory_saver=enable_memory_saver,
ratio=ratio,
online=(ratio == 128 and ONLINE_C128),
request_scoped=ratio == 128,
swa_page_size=self.swa_page_size,
online_mtp_max_draft_tokens=(
self.online_mtp_max_draft_tokens if ratio == 128 else 0
@@ -1267,10 +1277,24 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
state[state_locs, :half] = 0
state[state_locs, half:] = float("-inf")
def clear_c128_req_state(self, req_pool_idx: int) -> None:
"""Reset request-scoped C128 state for one req slot."""
def request_state_transfer_indices(self, req_pool_idx: int, seq_len: int):
"""PD transfer indices of the request-state component for one request."""
pools = [
p for p in self.compress_state_pools if p is not None and p.request_scoped
]
assert pools, "no request-scoped state pool"
# One index list addresses every request-state buffer (one per layer), so
# the request-scoped pools must share a ring layout.
layout = (pools[0].ratio, pools[0].online, pools[0].ring_size)
assert all((p.ratio, p.online, p.ring_size) == layout for p in pools), (
"request-scoped state pools must share one ring layout"
)
return pools[0].transfer_indices(req_pool_idx, seq_len)
def clear_request_scoped_state(self, req_pool_idx: int) -> None:
"""Reset request-scoped state for one req slot."""
for pool in self.compress_state_pools:
if pool is None or pool.ratio != 128:
if pool is None or not pool.request_scoped:
continue
state = pool.kv_score_buffer.kv_score
@@ -1332,6 +1356,26 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
assert compress_kv_pool is not None
return compress_kv_pool.page_size
def get_extra_key_layout(self, layer_id: int) -> KVLayout:
_, _, compress_kv_pool = self.layer_mapping[layer_id]
assert compress_kv_pool is not None
return compress_kv_pool.kv_layout
def get_extra_key_bytes_per_token(self, layer_id: int) -> int:
"""Last dim of the ``(pages, page_size, 1, bytes)`` view the attention
kernel detects the extra cache's format from."""
_, _, compress_kv_pool = self.layer_mapping[layer_id]
assert compress_kv_pool is not None
return compress_kv_pool.kv_cache_total_dim
def get_swa_key_layout(self) -> KVLayout:
return self.swa_kv_pool.kv_layout
def get_swa_key_bytes_per_token(self) -> int:
"""Last dim of the ``(pages, page_size, 1, bytes)`` view the attention
kernel detects the SWA cache's format from."""
return self.swa_kv_pool.kv_cache_total_dim
def get_extra_key_buffer(self, layer_id: int) -> torch.Tensor | None:
self.wait_layer_transfer(layer_id)
_, compress_layer_id, compress_kv_pool = self.layer_mapping[layer_id]
python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py
+23 -20
View File
@@ -1256,14 +1256,10 @@ def _extra_metadata_indices(
the upgraded `DSV4AttnMetadata`. Mirrors the dispatch in
`DeepseekV4AttnBackend.forward(compress_ratio=...)`.
"""
if compress_ratio == 4:
return (
core_metadata.c4_sparse_page_indices,
core_metadata.c4_sparse_topk_lengths,
)
if compress_ratio == 128:
return core_metadata.c128_page_indices, core_metadata.c128_topk_lengths_clamp1
raise ValueError(f"unsupported compress_ratio={compress_ratio}")
return (
core_metadata.sparse_page_indices(compress_ratio),
core_metadata.sparse_topk_lengths(compress_ratio),
)
def _pure_torch_dsv4_combined_reference(
@@ -1401,7 +1397,8 @@ def _seed_c4_sparse_indices(
non-trivial extra contribution.
"""
md = fixture.backend.forward_metadata.core_metadata
sparse_indices = md.c4_sparse_page_indices
ratio = fixture.case.compress_ratio
sparse_indices = md.sparse_page_indices(ratio)
num_q, sparse_topk = sparse_indices.shape
seed = torch.full(
(num_q, sparse_topk),
@@ -1412,12 +1409,13 @@ def _seed_c4_sparse_indices(
seed[:, :num_entries] = torch.arange(
num_entries, dtype=sparse_indices.dtype, device=sparse_indices.device
)
md.c4_sparse_page_indices = seed
md.c4_sparse_topk_lengths = torch.full(
(num_q,),
num_entries,
dtype=md.c4_sparse_topk_lengths.dtype,
device=md.c4_sparse_topk_lengths.device,
lengths = md.sparse_topk_lengths(ratio)
md.set_sparse_topk(
ratio,
page_indices=seed,
topk_lengths=torch.full(
(num_q,), num_entries, dtype=lengths.dtype, device=lengths.device
),
)
@@ -1438,10 +1436,11 @@ def _seed_c4_sparse_prefill_indices(
asserted below.
"""
md = fixture.backend.forward_metadata.core_metadata
raw_indices = md.c4_sparse_raw_indices
ratio = fixture.case.compress_ratio
raw_indices = md.sparse_raw_indices(ratio)
assert raw_indices is not None, "requires init_flashmla_related(is_prefill=True)"
num_q, width = raw_indices.shape
lens = (md.positions_casual + 1) // 4
lens = (md.positions_casual + 1) // ratio
max_len = int(lens.max().item())
pool = fixture.runner.token_to_kv_pool
c4_page_size = pool.get_extra_key_page_size(layer_id=0)
@@ -1457,9 +1456,13 @@ def _seed_c4_sparse_prefill_indices(
.expand(num_q, -1)
)
seeded = torch.where(seq < lens.unsqueeze(1), seq, seq.new_full((), -1))
md.c4_sparse_raw_indices = seeded
md.c4_sparse_page_indices = seeded.clone()
md.c4_sparse_topk_lengths = lens.to(md.c4_sparse_topk_lengths.dtype)
lengths = md.sparse_topk_lengths(ratio)
md.set_sparse_topk(
ratio,
page_indices=seeded.clone(),
topk_lengths=lens.to(lengths.dtype),
raw_indices=seeded,
)
def run_dsv4_target_verify_attention_case(