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

Unify full→SWA index translation in init_forward_metadata; drop pool caches (#27091)

Browse Source 
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-06-03 16:12:27 -07:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 8980eb82de
commit c9ca56da8c
29 changed files with 274 additions and 814 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
-1
View File
@@ -702,7 +702,6 @@ class Envs:
SGLANG_OPT_FLASHMLA_SPARSE_PREFILL = EnvBool(False)
# SWA radix cache
SGLANG_OPT_CACHE_SWA_TRANSLATION = EnvBool(True)
# TODO(DSV4): @ispobock this has bug on main branch when retract
SGLANG_OPT_SWA_RADIX_CACHE_COMPACT = EnvBool(False)
SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT = EnvBool(False)
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+78 -4
View File
@@ -116,6 +116,9 @@ class DSV4AttnMetadata:
swa_topk_lengths: torch.Tensor
c4_sparse_topk: int
# SWA KV-store write target (out_cache_loc translated to SWA space), computed
# once per iteration in make_core_attn_metadata and read by the store path.
swa_out_cache_loc: Optional[torch.Tensor] = None
c4_out_loc: Optional[torch.Tensor] = None
c4_topk_lengths_raw: Optional[torch.Tensor] = None
c4_topk_lengths_clamp1: Optional[torch.Tensor] = None
@@ -172,6 +175,9 @@ class DSV4AttnMetadata:
"c4_sparse_raw_indices",
],
assign_fields=[
# Recomputed by the recorded init_forward_metadata_in_graph op
# each forward; not copied across replays.
"swa_out_cache_loc",
"c1_flashmla_metadata",
"c4_flashmla_metadata",
"c128_flashmla_metadata",
@@ -218,6 +224,7 @@ class DSV4AttnMetadata:
]
_CP_GLOBAL_FIELDS = [
"raw_out_loc",
"swa_out_cache_loc",
"c4_out_loc",
"c128_out_loc",
]
@@ -699,6 +706,36 @@ class DeepseekV4AttnBackend(
raw_metadata=self.forward_metadata,
)
# Compute the SWA KV-store write target once per forward and cache it on
# the metadata for every layer's store. This is recorded inside the cuda
# graph, so replay re-reads the live out_cache_loc buffer (spec-v2 and DP
# padding rebind out_cache_loc after out-graph metadata prep). flash_mla
# kernels require int32 indices.
metadata = self.forward_metadata
if (
isinstance(metadata, DSV4Metadata)
and forward_batch.out_cache_loc is not None
):
out_cache_loc = forward_batch.out_cache_loc
if (
forward_batch.forward_mode.is_decode_or_idle()
and self.topk > 0
and self.speculative_num_steps > 1
):
# Multi-step draft decode shares one out_cache_loc buffer across
# steps; mirror the eager init's per-step slice.
out_cache_loc = per_step_draft_out_cache_loc(
out_cache_loc,
forward_batch.batch_size,
self.topk,
self.speculative_num_steps,
)[self.speculative_step_id]
metadata.core_attn_metadata.swa_out_cache_loc = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to(
torch.int32
)
)
def init_forward_metadata_out_graph(
self,
forward_batch: ForwardBatch,
@@ -789,12 +826,22 @@ class DeepseekV4AttnBackend(
)
elif bucket == _GraphBucket.DRAFT_EXTEND:
num_tokens_per_bs = self.draft_extend_num_tokens_per_bs
if out_cache_loc is not None:
# Pad the real write locations to the captured token count so
# raw_out_loc reflects the actual replay out_cache_loc.
out_cache_loc = torch.nn.functional.pad(
out_cache_loc,
pad=(0, num_tokens_per_bs * bs - len(out_cache_loc)),
mode="constant",
value=0,
)
temp_metadata = self.init_forward_metadata_draft_extend(
max_seq_len=chosen_max_seq_len,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
seq_lens_cpu=seq_lens_cpu.tolist(),
num_tokens_per_bs=num_tokens_per_bs,
out_cache_loc=out_cache_loc,
use_prefill_cuda_graph=True,
)
else:
@@ -934,21 +981,47 @@ class DeepseekV4AttnBackend(
if current_raw is not None:
self.forward_metadata = current_raw
def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor:
"""Resolve the SWA KV-store write target for the current forward.
Fast path: the per-forward value cached by init_forward_metadata_in_graph
(recorded inside cuda graphs, so replay re-reads live buffers). Fallback:
translate at store time, matching the pre-cache behavior, for paths that
never run the in-graph init — eager idle (forward_idle skips attn init),
runners that only run the out-graph prep (e.g.
EAGLEDraftExtendCudaGraphRunner) — or whose batch was re-padded after
init (shape mismatch). Idle always falls back: its metadata is absent or
left over from a previous forward, and translating the zero-padded
out_cache_loc writes to the dummy slot.
"""
out_cache_loc = forward_batch.out_cache_loc
core = getattr(self.forward_metadata, "core_attn_metadata", None)
cached = core.swa_out_cache_loc if core is not None else None
if (
cached is not None
and not forward_batch.forward_mode.is_idle()
and cached.shape[0] == out_cache_loc.shape[0]
):
return cached
return self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to(
torch.int32
)
def store_cache(
self, layer_id: int, swa_k: torch.Tensor, forward_batch: ForwardBatch
) -> None:
raw_loc = forward_batch.out_cache_loc
swa_loc = self.get_swa_out_cache_loc(forward_batch)
if envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get():
self.token_to_kv_pool.set_swa_key_buffer_radix_fused(
layer_id=layer_id,
raw_loc=raw_loc,
swa_loc=swa_loc,
cache_k=swa_k,
)
else:
swa_k_pack = quant_to_nope_fp8_rope_bf16_pack_triton(swa_k)
self.token_to_kv_pool.set_swa_key_buffer_radix(
layer_id=layer_id,
raw_loc=raw_loc,
swa_loc=swa_loc,
cache_nope_fp8_rope_bf16_pack=swa_k_pack,
)
@@ -1322,7 +1395,8 @@ class DeepseekV4AttnBackend(
assert raw_indices.shape == (num_qo_tokens, SWA_WINDOW)
raw_indices.masked_fill_(invalid_offset_mask, -1)
swa_indices = self.token_to_kv_pool.translate_loc_from_full_to_swa(raw_indices)
return swa_indices
# flash_mla attention requires int32 page indices.
return swa_indices.to(torch.int32)
class DeepseekV4MultiStepBackend(DeepseekV4AttnBackend):
python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py
+78 -4
View File
@@ -106,6 +106,9 @@ class DSV4AttnMetadata:
swa_topk_lengths: torch.Tensor
c4_sparse_topk: int
# SWA KV-store write target (out_cache_loc translated to SWA space), computed
# once per iteration in make_core_attn_metadata and read by the store path.
swa_out_cache_loc: Optional[torch.Tensor] = None
c4_out_loc: Optional[torch.Tensor] = None
c4_topk_lengths_raw: Optional[torch.Tensor] = None
c4_topk_lengths_clamp1: Optional[torch.Tensor] = None
@@ -160,6 +163,9 @@ class DSV4AttnMetadata:
"c4_sparse_page_indices",
],
assign_fields=[
# Recomputed by the recorded init_forward_metadata_in_graph op
# each forward; not copied across replays.
"swa_out_cache_loc",
"c1_flashmla_metadata",
"c4_flashmla_metadata",
"c128_flashmla_metadata",
@@ -206,6 +212,7 @@ class DSV4AttnMetadata:
]
_CP_GLOBAL_FIELDS = [
"raw_out_loc",
"swa_out_cache_loc",
"c4_out_loc",
"c128_out_loc",
]
@@ -672,6 +679,36 @@ class DeepseekV4HipRadixBackend(
raw_metadata=self.forward_metadata,
)
# Compute the SWA KV-store write target once per forward and cache it on
# the metadata for every layer's store. This is recorded inside the cuda
# graph, so replay re-reads the live out_cache_loc buffer (spec-v2 and DP
# padding rebind out_cache_loc after out-graph metadata prep). flash_mla
# kernels require int32 indices.
metadata = self.forward_metadata
if (
isinstance(metadata, DSV4Metadata)
and forward_batch.out_cache_loc is not None
):
out_cache_loc = forward_batch.out_cache_loc
if (
forward_batch.forward_mode.is_decode_or_idle()
and self.topk > 0
and self.speculative_num_steps > 1
):
# Multi-step draft decode shares one out_cache_loc buffer across
# steps; mirror the eager init's per-step slice.
out_cache_loc = per_step_draft_out_cache_loc(
out_cache_loc,
forward_batch.batch_size,
self.topk,
self.speculative_num_steps,
)[self.speculative_step_id]
metadata.core_attn_metadata.swa_out_cache_loc = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to(
torch.int32
)
)
def init_forward_metadata_out_graph(
self,
forward_batch: ForwardBatch,
@@ -760,12 +797,22 @@ class DeepseekV4HipRadixBackend(
)
elif bucket == _GraphBucket.DRAFT_EXTEND:
num_tokens_per_bs = self.draft_extend_num_tokens_per_bs
if out_cache_loc is not None:
# Pad the real write locations to the captured token count so
# raw_out_loc reflects the actual replay out_cache_loc.
out_cache_loc = torch.nn.functional.pad(
out_cache_loc,
pad=(0, num_tokens_per_bs * bs - len(out_cache_loc)),
mode="constant",
value=0,
)
temp_metadata = self.init_forward_metadata_draft_extend(
max_seq_len=chosen_max_seq_len,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
seq_lens_cpu=seq_lens_cpu.tolist(),
num_tokens_per_bs=num_tokens_per_bs,
out_cache_loc=out_cache_loc,
use_prefill_cuda_graph=True,
)
else:
@@ -905,21 +952,47 @@ class DeepseekV4HipRadixBackend(
if current_raw is not None:
self.forward_metadata = current_raw
def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor:
"""Resolve the SWA KV-store write target for the current forward.
Fast path: the per-forward value cached by init_forward_metadata_in_graph
(recorded inside cuda graphs, so replay re-reads live buffers). Fallback:
translate at store time, matching the pre-cache behavior, for paths that
never run the in-graph init — eager idle (forward_idle skips attn init),
runners that only run the out-graph prep (e.g.
EAGLEDraftExtendCudaGraphRunner) — or whose batch was re-padded after
init (shape mismatch). Idle always falls back: its metadata is absent or
left over from a previous forward, and translating the zero-padded
out_cache_loc writes to the dummy slot.
"""
out_cache_loc = forward_batch.out_cache_loc
core = getattr(self.forward_metadata, "core_attn_metadata", None)
cached = core.swa_out_cache_loc if core is not None else None
if (
cached is not None
and not forward_batch.forward_mode.is_idle()
and cached.shape[0] == out_cache_loc.shape[0]
):
return cached
return self.token_to_kv_pool.translate_loc_from_full_to_swa(out_cache_loc).to(
torch.int32
)
def store_cache(
self, layer_id: int, swa_k: torch.Tensor, forward_batch: ForwardBatch
) -> None:
raw_loc = forward_batch.out_cache_loc
swa_loc = self.get_swa_out_cache_loc(forward_batch)
if envs.SGLANG_OPT_USE_FUSED_STORE_CACHE.get():
self.token_to_kv_pool.set_swa_key_buffer_radix_fused(
layer_id=layer_id,
raw_loc=raw_loc,
swa_loc=swa_loc,
cache_k=swa_k,
)
else:
swa_k_pack = quant_to_nope_fp8_rope_bf16_pack_triton(swa_k)
self.token_to_kv_pool.set_swa_key_buffer_radix(
layer_id=layer_id,
raw_loc=raw_loc,
swa_loc=swa_loc,
cache_nope_fp8_rope_bf16_pack=swa_k_pack,
)
@@ -1165,7 +1238,8 @@ class DeepseekV4HipRadixBackend(
assert raw_indices.shape == (num_qo_tokens, SWA_WINDOW)
raw_indices.masked_fill_(invalid_offset_mask, -1)
swa_indices = self.token_to_kv_pool.translate_loc_from_full_to_swa(raw_indices)
return swa_indices
# flash_mla attention requires int32 page indices.
return swa_indices.to(torch.int32)
class DeepseekV4MultiStepBackend(DeepseekV4HipRadixBackend):
python/sglang/srt/layers/attention/dsa_backend.py
+2 -1
View File
@@ -1486,9 +1486,10 @@ class DeepseekSparseAttnBackend(
# todo hisparse: to cover more backends
if self.hisparse_coordinator is not None:
# flash_mla_sparse_fwd / tilelang require int32 page indices.
page_table_1 = self.token_to_kv_pool.translate_loc_to_hisparse_device(
page_table_1
)
).to(torch.int32)
if dsa_impl == "tilelang":
if q_rope is not None:
python/sglang/srt/layers/attention/dsv4/compressor_v2.py
+4 -1
View File
@@ -509,7 +509,10 @@ class CompressorBackendMixin:
if hasattr(compress_kv_pool, "translate_loc_to_hisparse_device"):
# The v2 compressor writes directly into the raw C4 KV tensor.
# HiSparse C4 therefore needs the physical C4 location here.
out_loc = compress_kv_pool.translate_loc_to_hisparse_device(out_loc)
# The compress kernel requires an int32 write location.
out_loc = compress_kv_pool.translate_loc_to_hisparse_device(
out_loc
).to(torch.int32)
self._forward_compress_all_in_one(
kv_score_buffer=state_pool.kv_score_buffer.kv_score,
kv_score_input=kv_score_input,
python/sglang/srt/layers/attention/dsv4/indexer.py
+2 -1
View File
@@ -600,10 +600,11 @@ class C4IndexerBackendMixin:
)
)
else:
# flash_mla C4 attention requires int32 page indices.
core_metadata.c4_sparse_page_indices = (
token_to_kv_pool.c4_kv_pool.translate_loc_to_hisparse_device(
core_metadata.c4_sparse_page_indices
)
).to(torch.int32)
)
if capture_enabled:
python/sglang/srt/layers/attention/flashattention_backend.py
+8 -7
View File
@@ -693,10 +693,11 @@ class FlashAttentionBackend(AttentionBackend):
]
if self.use_sliding_window_kv_pool:
# FA3 requires an int32 page_table.
metadata.swa_page_table = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
metadata.page_table
)
).to(torch.int32)
)
# Convert the page table to a strided format which is needed by FA3 API
@@ -886,7 +887,7 @@ class FlashAttentionBackend(AttentionBackend):
else:
page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
metadata.page_table
)
).to(torch.int32)
cu_seqlens_q = metadata.cu_seqlens_q
cache_seqlens = metadata.cache_seqlens_int32
max_seqlen_q = metadata.max_seq_len_q
@@ -1365,7 +1366,7 @@ class FlashAttentionBackend(AttentionBackend):
page_table = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
metadata.page_table
)
).to(torch.int32)
)
cache_seqlens = metadata.cache_seqlens_int32
max_seqlen_q = metadata.max_seq_len_q
@@ -2424,7 +2425,7 @@ class FlashAttentionBackend(AttentionBackend):
if self.use_sliding_window_kv_pool:
page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
metadata.page_table
)
).to(torch.int32)
else:
page_table = metadata.page_table
if cu_seqlens_q is None or cache_seqlens_int32 is None or page_table is None:
@@ -2551,7 +2552,7 @@ class FlashAttentionBackend(AttentionBackend):
if self.use_sliding_window_kv_pool:
sliced_page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa(
metadata.page_table[:bs, :max_seq_len]
)
).to(torch.int32)
else:
sliced_page_table = metadata.page_table[:bs, :max_seq_len]
@@ -2628,10 +2629,10 @@ class FlashAttentionBackend(AttentionBackend):
if self.use_sliding_window_kv_pool:
page_table_a = self.token_to_kv_pool.translate_loc_from_full_to_swa(
page_table_a
)
).to(torch.int32)
page_table_b = self.token_to_kv_pool.translate_loc_from_full_to_swa(
page_table_b
)
).to(torch.int32)
prepare_swa_spec_page_table_triton(
page_table,
python/sglang/srt/layers/attention/triton_backend.py
-3
View File
@@ -1519,12 +1519,9 @@ def update_sliding_window_buffer(
)
if hasattr(token_to_kv_pool, "translate_loc_from_full_to_swa"):
kv_last_index = window_kv_indptr[-1]
# Flush before+after: window_kv_indices is a different tensor than out_cache_loc.
token_to_kv_pool.invalidate_loc_cache()
window_kv_indices[:kv_last_index] = (
token_to_kv_pool.translate_loc_from_full_to_swa(
window_kv_indices[:kv_last_index]
)
)
token_to_kv_pool.invalidate_loc_cache()
return window_kv_indptr, window_kv_indices, window_kv_lens, window_kv_start_idx
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+6 -3
View File
@@ -164,9 +164,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
if self._swa_kv_pool is None:
return None
shape = token_indices.shape
return self._swa_kv_pool.translate_loc_from_full_to_swa(
token_indices.reshape(-1)
).reshape(shape)
# trtllm-gen SWA attention kernels require int32 page indices.
return (
self._swa_kv_pool.translate_loc_from_full_to_swa(token_indices.reshape(-1))
.reshape(shape)
.to(torch.int32)
)
def _alloc_swa_page_table(
self, max_bs: int, max_num_pages: int
python/sglang/srt/mem_cache/base_swa_memory_pool.py
-3
View File
@@ -16,9 +16,6 @@ class BaseSWAKVPool(KVCache):
swa_kv_pool: KVCache
def invalidate_loc_cache(self) -> None:
pass
@abc.abstractmethod
def register_mapping(self, full_to_swa_index_mapping: torch.Tensor) -> None:
raise NotImplementedError()
python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+4 -22
View File
@@ -513,15 +513,8 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
else:
self._init_paged_compress_states(enable_memory_saver)
self._should_cache_swa = envs.SGLANG_OPT_CACHE_SWA_TRANSLATION.get()
self.cached_loc = None
def register_mapping(self, full_to_swa_index_mapping: torch.Tensor):
self.full_to_swa_index_mapping = full_to_swa_index_mapping
self.cached_loc = None # mapping replaced; discard any cached translation
def invalidate_loc_cache(self) -> None:
self.cached_loc = None
def get_ring_size(self, compress_ratio: int) -> int:
server_args = get_global_server_args()
@@ -530,15 +523,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor):
assert self.full_to_swa_index_mapping is not None
return self.full_to_swa_index_mapping[kv_indices].to(torch.int32)
def get_cached_swa_loc(self, raw_loc: torch.Tensor, layer_id: int) -> torch.Tensor:
if self._should_cache_swa:
if layer_id == self.start_layer or self.cached_loc is None:
self.cached_loc = self.translate_loc_from_full_to_swa(raw_loc)
return self.cached_loc
return self.translate_loc_from_full_to_swa(raw_loc)
return self.full_to_swa_index_mapping[kv_indices]
def get_contiguous_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
data_ptrs: List[int] = []
@@ -768,10 +753,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
def set_swa_key_buffer_radix(
self,
layer_id: int,
raw_loc: torch.Tensor,
swa_loc: torch.Tensor,
cache_nope_fp8_rope_bf16_pack: NopeFp8RopeBf16Pack,
) -> None:
swa_loc = self.translate_loc_from_full_to_swa(raw_loc)
self.swa_kv_pool.set_key_buffer(
self._swa_local_layer_id(layer_id), swa_loc, cache_nope_fp8_rope_bf16_pack
)
@@ -783,10 +767,9 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
def set_swa_key_buffer_radix_fused(
self,
layer_id: int,
raw_loc: torch.Tensor,
swa_loc: torch.Tensor,
cache_k: torch.Tensor,
) -> None:
swa_loc = self.get_cached_swa_loc(raw_loc, layer_id)
return self.swa_kv_pool.set_key_buffer_fused(
self._swa_local_layer_id(layer_id), swa_loc, cache_k
)
@@ -794,14 +777,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
def set_swa_key_buffer_radix_fused_norm_rope(
self,
layer_id: int,
raw_loc: torch.Tensor,
swa_loc: torch.Tensor,
kv: torch.Tensor,
kv_weight: torch.Tensor,
eps: float,
freqs_cis: torch.Tensor,
positions: torch.Tensor,
) -> None:
swa_loc = self.get_cached_swa_loc(raw_loc, layer_id)
fused_k_norm_rope_flashmla(
kv=kv,
kv_weight=kv_weight,
python/sglang/srt/mem_cache/hisparse_memory_pool.py
+1 -3
View File
@@ -77,9 +77,7 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool):
)
def translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor):
return self.full_to_hisparse_device_index_mapping[compressed_indices].to(
torch.int32
)
return self.full_to_hisparse_device_index_mapping[compressed_indices]
def _translate_loc_to_hisparse_device(self, compressed_indices: torch.Tensor):
return self.full_to_hisparse_device_index_mapping[compressed_indices]
python/sglang/srt/mem_cache/swa_memory_pool.py
+1 -28
View File
@@ -92,8 +92,6 @@ class SWAKVPool(BaseSWAKVPool):
for swa_layer_id, global_layer_id in enumerate(swa_attention_layer_ids):
self.layers_mapping[global_layer_id] = (swa_layer_id, True)
self.full_to_swa_index_mapping: Optional[torch.Tensor] = None
self._cached_swa_loc: Optional[torch.Tensor] = None
self._cached_loc_key: Optional[tuple] = None
k_size, v_size = self.get_kv_size_bytes()
self.mem_usage = (k_size + v_size) / GB
@@ -103,11 +101,6 @@ class SWAKVPool(BaseSWAKVPool):
def register_mapping(self, full_to_swa_index_mapping: torch.Tensor):
self.full_to_swa_index_mapping = full_to_swa_index_mapping
self.invalidate_loc_cache()
def invalidate_loc_cache(self) -> None:
self._cached_swa_loc = None
self._cached_loc_key = None
def register_layer_transfer_counter(self, layer_transfer_counter):
# Wait happens at this wrapper. Inner pools must not wait again.
@@ -167,21 +160,8 @@ class SWAKVPool(BaseSWAKVPool):
def translate_loc_from_full_to_swa(self, kv_indices: torch.Tensor) -> torch.Tensor:
assert self.full_to_swa_index_mapping is not None
# data_ptr() (not untyped_storage().data_ptr()) encodes the offset, so
# views at different positions within the same storage get distinct keys.
# -1 in kv_indices maps to -1 via the sentinel appended to the mapping.
key = (kv_indices.data_ptr(), kv_indices.numel())
if key != self._cached_loc_key:
if self._cached_loc_key is not None:
logger.debug(
"translate_loc_from_full_to_swa: loc tensor changed mid-forward "
"without invalidate_loc_cache() — possible missing call site"
)
self._cached_swa_loc = self.full_to_swa_index_mapping[kv_indices].to(
torch.int32
)
self._cached_loc_key = key
return self._cached_swa_loc
return self.full_to_swa_index_mapping[kv_indices]
def set_kv_buffer(
self,
@@ -425,7 +405,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
return self._kvcache.translate_loc_from_full_to_swa(kv_indices)
def alloc(self, need_size: int):
self._kvcache.invalidate_loc_cache()
assert self.page_size == 1
if need_size > self.full_attn_allocator.available_size():
return None
@@ -454,7 +433,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
last_loc: torch.Tensor, # last_loc for full layers
extend_num_tokens: int,
):
self._kvcache.invalidate_loc_cache()
assert self.page_size > 1
num_new_pages = get_num_new_pages(
@@ -507,7 +485,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
extend_num_tokens: int,
swa_tail_len: int,
):
self._kvcache.invalidate_loc_cache()
"""Allocate full KV for the whole extend and SWA KV only for the tail.
This is used by disaggregated decode preallocation: decode receives full
@@ -571,7 +548,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
seq_lens_cpu: torch.Tensor,
last_loc: torch.Tensor, # last_loc for full layers
):
self._kvcache.invalidate_loc_cache()
assert self.page_size > 1
swa_last_loc = self.translate_loc_from_full_to_swa(last_loc)
@@ -619,7 +595,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
if full_indices.numel() == 0:
return
assert full_indices.numel() == swa_indices.numel()
self._kvcache.invalidate_loc_cache()
if _is_npu:
self.full_to_swa_index_mapping[full_indices.to(torch.int64)] = (
swa_indices.to(torch.int64)
@@ -628,7 +603,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
self.full_to_swa_index_mapping[full_indices] = swa_indices
def free_swa(self, free_index: torch.Tensor):
self._kvcache.invalidate_loc_cache()
swa_indices = self.full_to_swa_index_mapping[free_index]
swa_indices = swa_indices[swa_indices > 0]
self.swa_attn_allocator.free(swa_indices)
@@ -646,7 +620,6 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
self.swa_attn_allocator.restore_state(state[1])
def clear(self):
self._kvcache.invalidate_loc_cache()
self.swa_attn_allocator.clear()
self.full_attn_allocator.clear()
# Note: the last item is -1, we don't clear it, see the comment in __init__
python/sglang/srt/model_executor/breakable_cuda_graph_runner.py
-3
View File
@@ -414,9 +414,6 @@ class BreakableCudaGraphRunner:
self.model_runner.attn_backend.init_forward_metadata(forward_batch)
def run_once():
# Invalidate SWA loc cache — same fix as in cuda_graph_runner.run_once.
if self.model_runner.is_hybrid_swa:
self.model_runner.token_to_kv_pool.invalidate_loc_cache()
return self._run_forward(forward_batch, num_tokens)
with forward_context(
python/sglang/srt/model_executor/cuda_graph_runner.py
-6
View File
@@ -1047,12 +1047,6 @@ class CudaGraphRunner:
attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True)
def run_once():
# Without this, warmup-1 caches the translation; the capture
# run hits the cache, skips the gather, and replay reuses
# stale SWA locations.
if self.model_runner.is_hybrid_swa:
self.model_runner.token_to_kv_pool.invalidate_loc_cache()
# Must run inside the capture block: warmup mutations here are
# undone by on_after_cuda_graph_warmup so capture starts clean.
attn_backend.init_forward_metadata_in_graph(forward_batch)
python/sglang/srt/model_executor/model_runner.py
-3
View File
@@ -3362,9 +3362,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.hisparse_coordinator.wait_for_pending_backup()
self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size)
if self.is_hybrid_swa:
self.token_to_kv_pool.invalidate_loc_cache()
# Replay cuda graph if applicable
if can_run_graph:
ret = self.graph_runner.replay(
python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py
-4
View File
@@ -623,10 +623,6 @@ class PiecewiseCudaGraphRunner:
# Run and capture
def run_once():
# Invalidate SWA loc cache — same fix as in cuda_graph_runner.run_once.
if self.model_runner.is_hybrid_swa:
self.model_runner.token_to_kv_pool.invalidate_loc_cache()
# Clean intermediate result cache for DP attention
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = (
None
python/sglang/srt/models/deepseek_v4.py
+11 -10
View File
@@ -471,6 +471,7 @@ class MQALayer(nn.Module):
x: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
attn_backend,
qkv_a: Optional[torch.Tensor] = None,
) -> None:
"""Fused: rmsnorm + RoPE + write directly to FlashMLA paged cache.
@@ -487,7 +488,7 @@ class MQALayer(nn.Module):
assert isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
token_to_kv_pool.set_swa_key_buffer_radix_fused_norm_rope(
layer_id=self.layer_id,
raw_loc=forward_batch.out_cache_loc,
swa_loc=attn_backend.get_swa_out_cache_loc(forward_batch),
kv=kv,
kv_weight=self.kv_norm.weight.data,
eps=self.eps,
@@ -562,7 +563,9 @@ class MQALayer(nn.Module):
if qkv_a_ready is not None:
stream_kv.wait_event(qkv_a_ready)
# Fused norm + rope + cache write -- no bf16 KV intermediate.
self._compute_kv_to_cache(x_linear, positions, forward_batch, qkv_a=qkv_a)
self._compute_kv_to_cache(
x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
)
del qkv_a
@@ -646,9 +649,7 @@ class MQALayer(nn.Module):
)
token_to_kv_pool = get_token_to_kv_pool()
swa_loc = token_to_kv_pool.get_cached_swa_loc(
forward_batch.out_cache_loc, self.layer_id
)
swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch)
swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id]
swa_page_size = token_to_kv_pool.swa_kv_pool.page_size
@@ -672,7 +673,9 @@ class MQALayer(nn.Module):
else:
q_lora = self.q_norm(q_lora)
q = self._compute_q_b(q_lora, positions, q_out)
self._compute_kv_to_cache(x_linear, positions, forward_batch, qkv_a=qkv_a)
self._compute_kv_to_cache(
x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
)
del qkv_a
@@ -736,9 +739,7 @@ class MQALayer(nn.Module):
)
token_to_kv_pool = get_token_to_kv_pool()
swa_loc = token_to_kv_pool.get_cached_swa_loc(
forward_batch.out_cache_loc, self.layer_id
)
swa_loc = attn_backend.get_swa_out_cache_loc(forward_batch)
swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id]
swa_page_size = token_to_kv_pool.swa_kv_pool.page_size
@@ -790,7 +791,7 @@ class MQALayer(nn.Module):
)
else:
self._compute_kv_to_cache(
x_linear, positions, forward_batch, qkv_a=qkv_a
x_linear, positions, forward_batch, attn_backend, qkv_a=qkv_a
)
kv = None
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
-3
View File
@@ -352,9 +352,6 @@ class EAGLEDraftCudaGraphRunner:
)
def run_once():
if self.model_runner.is_hybrid_swa:
self.model_runner.token_to_kv_pool.invalidate_loc_cache()
self.draft_attn_backend.init_forward_metadata_in_graph(forward_batch)
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
-3
View File
@@ -376,9 +376,6 @@ class EAGLEDraftExtendCudaGraphRunner:
)
def run_once():
if self.model_runner.is_hybrid_swa:
self.model_runner.token_to_kv_pool.invalidate_loc_cache()
# Clean intermediate result cache for DP attention
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
set_dp_buffer_len(
python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
-3
View File
@@ -270,9 +270,6 @@ class FrozenKVMTPCudaGraphRunner:
)
def run_once():
if self.model_runner.is_hybrid_swa:
self.model_runner.token_to_kv_pool.invalidate_loc_cache()
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
set_dp_buffer_len(
global_dp_buffer_len,
python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
-3
View File
@@ -401,9 +401,6 @@ class MultiLayerEagleDraftExtendCudaGraphRunner:
attn_backend = self.eagle_worker.draft_extend_attn_backend_list[self.step]
def run_once():
if self.model_runner.is_hybrid_swa:
self.model_runner.token_to_kv_pool.invalidate_loc_cache()
# Clean intermediate result cache for DP attention
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
set_dp_buffer_len(
python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py
+8 -3
View File
@@ -518,9 +518,12 @@ class ProjectedDSV4Attention(nn.Module):
# `[num_tokens, 1, hidden_dim]`.
k_flat = k.reshape(k.shape[0], -1).to(torch.bfloat16)
pack = quant_to_nope_fp8_rope_bf16_pack_triton(k_flat)
attn_backend.token_to_kv_pool.set_swa_key_buffer_radix(
pool = attn_backend.token_to_kv_pool
pool.set_swa_key_buffer_radix(
layer_id=self.attn.layer_id,
raw_loc=forward_batch.out_cache_loc.to(torch.int64),
swa_loc=pool.translate_loc_from_full_to_swa(
forward_batch.out_cache_loc.to(torch.int64)
),
cache_nope_fp8_rope_bf16_pack=pack,
)
out = attn_backend.forward(
@@ -546,7 +549,9 @@ def _write_swa_cache(
pack = quant_to_nope_fp8_rope_bf16_pack_triton(k_bf16.to(torch.bfloat16))
runner.token_to_kv_pool.set_swa_key_buffer_radix(
layer_id=layer_id,
raw_loc=loc.to(torch.int64),
swa_loc=runner.token_to_kv_pool.translate_loc_from_full_to_swa(
loc.to(torch.int64)
),
cache_nope_fp8_rope_bf16_pack=pack,
)
test/manual/core/test_dsv4_cached_loc_invalidation.py
-180
View File
@@ -1,180 +0,0 @@
"""Regression test for PR #25889: DeepSeekV4TokenToKVPool.register_mapping()
must clear cached_loc.
Bug scenario (pre-fix):
During a forward pass, the first SWA layer (layer_id == start_layer) computes
and caches `cached_loc` via translate_loc_from_full_to_swa(). If HiCache then
loads back SWA KV from host, it calls register_mapping() to install the new
full->swa index mapping. Before the fix, register_mapping() only stored the
new tensor but did NOT clear cached_loc. Subsequent SWA layers (layer_id >
start_layer) saw `cached_loc is not None` and returned the stale translation,
silently writing KV to wrong SWA pool slots.
Fix (PR #25889): add `self.cached_loc = None` in register_mapping().
Test structure:
- test_stale_without_fix: shows the stale-return bug using a replica
of the pre-fix logic
- test_correct_with_fix: verifies the fixed logic returns fresh values
- test_actual_pool_register_mapping: exercises the real production method
directly (bypassing the full __init__)
Run with:
python -m pytest test/manual/core/test_dsv4_cached_loc_invalidation.py -v
"""
import unittest
import torch
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.utils import get_device
from sglang.test.test_utils import CustomTestCase
# ---------------------------------------------------------------------------
# Minimal stub that replicates DeepSeekV4TokenToKVPool's caching pattern.
# Used to demonstrate both sides of the bug without constructing the full pool.
# ---------------------------------------------------------------------------
class _DSV4CacheStub:
"""Stripped-down replica of the caching logic in DeepSeekV4TokenToKVPool."""
start_layer = 3
def __init__(self, device):
self.device = device
self.cached_loc = None
self.full_to_swa_index_mapping = None
# --- pre-fix version ---
def register_mapping_buggy(self, mapping: torch.Tensor) -> None:
self.full_to_swa_index_mapping = mapping
# BUG: cached_loc not cleared → stale on next mid-forward call
# --- post-fix version (PR #25889) ---
def register_mapping_fixed(self, mapping: torch.Tensor) -> None:
self.full_to_swa_index_mapping = mapping
self.cached_loc = None # THE FIX
def _translate(self, raw_loc: torch.Tensor) -> torch.Tensor:
return self.full_to_swa_index_mapping[raw_loc]
def get_swa_loc(self, layer_id: int, raw_loc: torch.Tensor) -> torch.Tensor:
"""Exact replica of set_swa_key_buffer_radix_fused caching branch."""
if layer_id == self.start_layer or self.cached_loc is None:
self.cached_loc = self._translate(raw_loc)
return self.cached_loc
def _make_mapping(indices, values, size=32, device="cpu"):
m = torch.zeros(size, dtype=torch.int64, device=device)
m[indices] = torch.tensor(values, dtype=torch.int64, device=device)
return m
class TestDSV4CachedLocBugAndFix(CustomTestCase):
"""Shows the pre-fix bug and verifies the post-fix behaviour."""
def setUp(self):
self.device = get_device()
self.raw_loc = torch.tensor([0, 1, 2, 3], device=self.device)
self.mapping_v1 = _make_mapping(
[0, 1, 2, 3], [10, 11, 12, 13], device=self.device
)
self.mapping_v2 = _make_mapping(
[0, 1, 2, 3], [20, 21, 22, 23], device=self.device
)
def test_stale_without_fix(self):
"""Without the fix, register_mapping() mid-forward leaves a stale cached_loc.
Sequence:
1. start_layer: cached_loc = translate(mapping_v1) = [10..13]
2. HiCache load-back: register_mapping(mapping_v2) ← buggy, no clear
3. start_layer+1: layer_id != start_layer, cached_loc is not None
→ returns stale [10..13], NOT the correct [20..23]
"""
stub = _DSV4CacheStub(self.device)
stub.register_mapping_buggy(self.mapping_v1)
# Forward pass — first SWA layer primes the cache.
loc = stub.get_swa_loc(stub.start_layer, self.raw_loc)
self.assertEqual(loc.tolist(), [10, 11, 12, 13])
# HiCache load-back installs new mapping (buggy path).
stub.register_mapping_buggy(self.mapping_v2)
self.assertIsNotNone(stub.cached_loc, "Bug: cached_loc not cleared")
# Next SWA layer — should use mapping_v2 but returns mapping_v1.
loc_next = stub.get_swa_loc(stub.start_layer + 1, self.raw_loc)
self.assertEqual(
loc_next.tolist(),
[10, 11, 12, 13],
"Confirms the bug: stale cached_loc [10..13] returned instead of [20..23]",
)
def test_correct_with_fix(self):
"""With the fix, register_mapping() clears cached_loc; next layer recomputes.
Same sequence as test_stale_without_fix but using the fixed register_mapping.
"""
stub = _DSV4CacheStub(self.device)
stub.register_mapping_fixed(self.mapping_v1)
# Forward pass — first SWA layer primes the cache.
loc = stub.get_swa_loc(stub.start_layer, self.raw_loc)
self.assertEqual(loc.tolist(), [10, 11, 12, 13])
# HiCache load-back installs new mapping (fixed path).
stub.register_mapping_fixed(self.mapping_v2)
self.assertIsNone(
stub.cached_loc, "Fix: cached_loc cleared by register_mapping"
)
# Next SWA layer — recomputes with mapping_v2.
loc_next = stub.get_swa_loc(stub.start_layer + 1, self.raw_loc)
self.assertEqual(
loc_next.tolist(),
[20, 21, 22, 23],
"Fix works: fresh translation [20..23] from new mapping",
)
class TestDSV4ActualPoolRegisterMapping(CustomTestCase):
"""Exercises the production DeepSeekV4TokenToKVPool.register_mapping() directly.
Uses __new__ to bypass the complex __init__ (which needs full GPU pool setup)
and tests only the register_mapping / cached_loc contract.
"""
def test_register_mapping_clears_cached_loc(self):
device = get_device()
# Bypass full __init__ — only the fields register_mapping touches matter.
pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
old_loc = torch.tensor([10, 11, 12, 13], device=device)
pool.cached_loc = old_loc
pool.full_to_swa_index_mapping = None
new_mapping = torch.arange(64, dtype=torch.int64, device=device)
pool.register_mapping(new_mapping)
self.assertIsNone(pool.cached_loc, "register_mapping must clear cached_loc")
self.assertIs(pool.full_to_swa_index_mapping, new_mapping)
def test_register_mapping_clears_none_cached_loc(self):
"""Idempotent when cached_loc is already None."""
pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
pool.cached_loc = None
pool.full_to_swa_index_mapping = None
mapping = torch.arange(16, dtype=torch.int64, device=get_device())
pool.register_mapping(mapping)
self.assertIsNone(pool.cached_loc)
self.assertIs(pool.full_to_swa_index_mapping, mapping)
if __name__ == "__main__":
unittest.main()
test/manual/core/test_dsv4_hicache_swa_translation_cache.py
-114
View File
@@ -1,114 +0,0 @@
"""E2E regression for PR #25889: DSV4 cached_loc stale after HiCache load-back.
Bug (pre-fix):
DeepSeekV4TokenToKVPool.register_mapping() replaces full_to_swa_index_mapping
but does NOT clear self.cached_loc. When SGLANG_OPT_CACHE_SWA_TRANSLATION=True,
set_swa_key_buffer_radix_fused() caches the full→SWA translation across SWA
layers. After a HiCache commit/load-back that calls register_mapping() with a
new mapping, subsequent SWA layers reuse the stale cached_loc and write KV to
wrong SWA pool slots — producing divergent logprobs.
Fix (PR #25889):
register_mapping() sets self.cached_loc = None so the first SWA layer in the
next forward pass recomputes the translation from the fresh mapping.
Test strategy:
Subclass the existing DSV4 Flash HiCache KL suite and activate the SWA
translation cache via SGLANG_OPT_CACHE_SWA_TRANSLATION=1. The mixin tests
(test_multiturn_logprobs_match, test_multiturn_prefill_cache_hit_branching,
test_multiturn_decode_cache_hit_branching) compare logprobs from cold and
warm radix-cache hits. Without the fix the stale translation corrupts SWA KV
data and the KL divergence exceeds the threshold; with the fix it stays within.
"""
import unittest
from test_unified_radix_cache_kl_hicache import (
DSV4_FLASH_LAUNCH_TIMEOUT,
DSV4_FLASH_MODEL,
_assert_dsv4_decode_cached_tokens,
)
from sglang.srt.utils import kill_process_tree
from sglang.test.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin
from sglang.test.kl_multiturn_utils import get_input_ids
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
class TestDSV4HiCacheSWATranslationCache(UnifiedRadixTreeTestMixin, CustomTestCase):
"""DSV4 Flash FP8 + HiCache + SWA translation cache enabled.
Identical server config to TestUnifiedDeepSeekV4FlashHiCache but with
SGLANG_OPT_CACHE_SWA_TRANSLATION=1 to activate the cached_loc path.
Without PR #25889 the KL tests fail; with the fix they pass.
"""
kl_threshold = 0.005
sampling_temperature = 0
decode_hit_request_batch_size = 3
decode_hit_inter_batch_delay_s = 0.5
decode_cache_assert = staticmethod(_assert_dsv4_decode_cached_tokens)
gsm8k_threshold = 0.90
num_gsm8k_questions = 100
@unittest.skipIf(True, "Covered by test_multiturn_prefill_cache_hit_branching.")
def test_multiturn_logprobs_match(self):
pass
@classmethod
def setUpClass(cls):
cls.model = DSV4_FLASH_MODEL
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DSV4_FLASH_LAUNCH_TIMEOUT,
other_args=[
"--trust-remote-code",
"--tp-size",
"4",
"--attention-backend",
"compressed",
"--page-size",
"256",
"--chunked-prefill-size",
"8192",
"--mem-fraction-static",
"0.9",
"--disable-shared-experts-fusion",
"--enable-hierarchical-cache",
"--hicache-ratio",
"4",
"--hicache-write-policy",
"write_through",
"--hicache-io-backend",
"direct",
"--hicache-mem-layout",
"page_first_direct",
"--swa-full-tokens-ratio",
"0.25",
"--max-total-tokens",
"20000",
"--max-running-requests",
"4",
],
env={
"SGLANG_DSV4_FP4_EXPERTS": "0",
"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1",
# Activate the SWA translation cache — the flag that exposes the bug.
"SGLANG_OPT_CACHE_SWA_TRANSLATION": "1",
},
)
cls.input_ids = get_input_ids(cls.model, num_samples=18)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
if __name__ == "__main__":
unittest.main()
test/manual/core/test_dsv4_stale_loc_crash.py
-164
View File
@@ -1,164 +0,0 @@
"""Crash regression for PR #25889: stale cached_loc after register_mapping().
Bug:
DeepSeekV4TokenToKVPool caches the full→SWA index translation in
self.cached_loc (when SGLANG_OPT_CACHE_SWA_TRANSLATION=True).
register_mapping() was replacing full_to_swa_index_mapping without
clearing cached_loc, so a subsequent set_swa_key_buffer_radix_fused
call would use stale SWA indices and, if those indices exceed the
current pool size, raise a RuntimeError (OOB tensor access).
Crash scenario reproduced here:
Pass 1 (large SWA pool, size=8): first SWA layer primes
cached_loc = [4, 5, 6, 7].
register_mapping() called with new mapping (pre-fix: cache not cleared).
Pass 2 (smaller SWA pool, size=4): same SWA layer finds cached_loc is
not None → uses stale [4, 5, 6, 7] → OOB write on size-4 pool →
RuntimeError.
Fix: register_mapping() sets self.cached_loc = None so the next call
recomputes with the fresh mapping.
Run with:
python -m pytest test/registered/unit/mem_cache/test_dsv4_stale_loc_crash.py -v
"""
import unittest
import torch
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
_NUM_HEADS = 2
_HEAD_DIM = 8
_SWA_LARGE = 8 # size-8 pool used during pass 1
_SWA_SMALL = 4 # size-4 pool used during pass 2 (simulates post-HiCache loadback)
class _SWAPoolMock:
"""Minimal SWA pool mock whose set_key_buffer_fused does a real tensor write.
A write with OOB swa_loc raises RuntimeError, reproducing the crash.
"""
def __init__(self, size: int):
self.buf = torch.zeros(size, _NUM_HEADS, _HEAD_DIM)
def set_key_buffer_fused(
self, local_layer_id: int, swa_loc: torch.Tensor, cache_k: torch.Tensor
) -> None:
n = swa_loc.numel()
self.buf[swa_loc.long()] = cache_k[:n].reshape(n, _NUM_HEADS, _HEAD_DIM)
def _build_pool(
mapping: torch.Tensor,
swa_pool: _SWAPoolMock,
start_layer: int = 0,
) -> DeepSeekV4TokenToKVPool:
"""Create a minimal DeepSeekV4TokenToKVPool via __new__, bypassing __init__."""
pool = DeepSeekV4TokenToKVPool.__new__(DeepSeekV4TokenToKVPool)
pool.cached_loc = None
pool._should_cache_swa = True
pool.start_layer = start_layer
pool.full_to_swa_index_mapping = mapping
pool.swa_kv_pool = swa_pool
# _swa_local_layer_id: map global SWA layer id → local index 0 for this test.
pool._swa_local_layer_id = lambda lid: 0
return pool
def _mapping(indices: list, size: int = 16) -> torch.Tensor:
m = torch.full((size,), -1, dtype=torch.int64)
for i, v in enumerate(indices):
m[i] = v
return m
class TestDSV4StaleLocCrash(CustomTestCase):
"""
Two paired tests that together constitute the crash regression for #25889.
test_crash_without_fix: reproduces the RuntimeError that occurs when
register_mapping() does NOT clear cached_loc.
test_fix_prevents_crash: same sequence with the fixed register_mapping()
→ no error, correct SWA slots written.
"""
def setUp(self):
# raw_loc: full-pool indices for 4 tokens.
self.raw_loc = torch.tensor([0, 1, 2, 3], dtype=torch.int64)
# cache_k: synthetic key data for 4 tokens.
self.cache_k = torch.ones(4, _NUM_HEADS, _HEAD_DIM)
# SWA layer id > start_layer (0), so cached_loc is only reset by the
# "is None" branch, NOT by the "layer_id == start_layer" branch.
self.swa_layer_id = 1
# mapping_v1: raw_loc [0,1,2,3] → SWA slots [4,5,6,7] (valid for size-8 pool).
self.mapping_v1 = _mapping([4, 5, 6, 7])
# mapping_v2: same raw_loc → SWA slots [0,1,2,3] (valid for size-4 pool).
self.mapping_v2 = _mapping([0, 1, 2, 3])
def test_crash_without_fix(self):
"""Without the fix, stale cached_loc [4,5,6,7] causes RuntimeError on
the size-4 pool used in pass 2."""
large_pool = _SWAPoolMock(_SWA_LARGE)
pool = _build_pool(self.mapping_v1, large_pool, start_layer=0)
# Pass 1: swa_layer_id=1, cached_loc is None → compute and cache [4,5,6,7].
pool.set_swa_key_buffer_radix_fused(
self.swa_layer_id, self.raw_loc, self.cache_k
)
self.assertEqual(pool.cached_loc.tolist(), [4, 5, 6, 7], "Pass 1: cache primed")
# PRE-FIX register_mapping: replace mapping WITHOUT clearing cached_loc.
pool.full_to_swa_index_mapping = self.mapping_v2
# Pass 2 on a smaller pool (size=4). cached_loc is still [4,5,6,7].
# OOB write → RuntimeError.
pool.swa_kv_pool = _SWAPoolMock(_SWA_SMALL)
with self.assertRaises((RuntimeError, IndexError)):
pool.set_swa_key_buffer_radix_fused(
self.swa_layer_id, self.raw_loc, self.cache_k
)
def test_fix_prevents_crash(self):
"""With the fix, register_mapping() clears cached_loc. Pass 2 recomputes
[0,1,2,3] from mapping_v2 and writes to the correct size-4 pool slots."""
large_pool = _SWAPoolMock(_SWA_LARGE)
pool = _build_pool(self.mapping_v1, large_pool, start_layer=0)
# Pass 1: prime cache → cached_loc = [4,5,6,7].
pool.set_swa_key_buffer_radix_fused(
self.swa_layer_id, self.raw_loc, self.cache_k
)
self.assertEqual(pool.cached_loc.tolist(), [4, 5, 6, 7])
# FIXED register_mapping: clears cached_loc.
pool.register_mapping(self.mapping_v2)
self.assertIsNone(
pool.cached_loc, "Fix: cached_loc cleared by register_mapping"
)
# Pass 2 on the smaller pool: cached_loc is None → recompute with mapping_v2.
small_pool = _SWAPoolMock(_SWA_SMALL)
pool.swa_kv_pool = small_pool
pool.set_swa_key_buffer_radix_fused(
self.swa_layer_id, self.raw_loc, self.cache_k
)
self.assertEqual(
pool.cached_loc.tolist(), [0, 1, 2, 3], "Fresh indices after fix"
)
# Verify data landed in the correct SWA slots [0-3], not the stale [4-7].
self.assertTrue(
small_pool.buf[0:4].abs().sum().item() > 0,
"Correct SWA slots 0-3 received data",
)
if __name__ == "__main__":
unittest.main()
test/manual/core/test_swa_loc_translation_cache.py
-233
View File
@@ -1,233 +0,0 @@
"""Manual tests for SWAKVPool.translate_loc_from_full_to_swa cache behaviour.
These tests cover three properties introduced by PR #25824:
1. Cache key uses data_ptr() — correctly distinguishes views at different
offsets within the same storage (untyped_storage().data_ptr() would not).
2. Allocator mutations invalidate the cache — alloc/free/clear/
set_full_to_swa_mapping each call invalidate_loc_cache() so the next
translation sees the fresh mapping.
3. BaseSWAKVPool.invalidate_loc_cache is a no-op default — subclasses that
don't cache (e.g. DSV4) can be called safely without AttributeError.
Run with:
python -m pytest test/manual/core/test_swa_loc_translation_cache.py -v
"""
import unittest
import torch
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
from sglang.srt.utils import get_device
from sglang.test.test_utils import CustomTestCase
def _build_pool(
kv_size: int = 32,
kv_size_swa: int = 32,
page_size: int = 1,
):
device = get_device()
num_layers = 8
full_layer_ids = [0, 4]
swa_layer_ids = [i for i in range(num_layers) if i not in set(full_layer_ids)]
pool = SWAKVPool(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=torch.bfloat16,
head_num=4,
head_dim=64,
swa_attention_layer_ids=swa_layer_ids,
full_attention_layer_ids=full_layer_ids,
enable_kvcache_transpose=False,
device=device,
)
allocator = SWATokenToKVPoolAllocator(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=torch.bfloat16,
device=device,
kvcache=pool,
need_sort=False,
)
return pool, allocator, device
class TestCacheKeyDataPtr(CustomTestCase):
"""Cache key uses data_ptr(), which encodes the storage offset."""
def test_same_offset_view_is_cache_hit(self):
"""Two different Python objects pointing to the same base are a hit."""
pool, allocator, device = _build_pool()
loc = allocator.alloc(4)
self.assertIsNotNone(loc)
# Create two slice objects at offset 0 — same data_ptr, same numel.
view_a = loc[:4]
view_b = loc[:4]
self.assertIsNot(view_a, view_b) # different Python objects
self.assertEqual(view_a.data_ptr(), view_b.data_ptr())
result_a = pool.translate_loc_from_full_to_swa(view_a)
result_b = pool.translate_loc_from_full_to_swa(view_b)
# Both should return the identical tensor (cache hit).
self.assertIs(result_a, result_b)
def test_different_offset_view_is_cache_miss(self):
"""Views at different offsets produce different data_ptr → cache miss."""
pool, allocator, device = _build_pool(kv_size=32, kv_size_swa=32)
loc = allocator.alloc(10)
self.assertIsNotNone(loc)
self.assertGreaterEqual(loc.numel(), 10)
view_lo = loc[0:5]
view_hi = loc[5:10]
self.assertEqual(view_lo.numel(), view_hi.numel()) # same numel
# Different data_ptr (different storage offset).
self.assertNotEqual(view_lo.data_ptr(), view_hi.data_ptr())
# Prime the cache with view_lo.
result_lo = pool.translate_loc_from_full_to_swa(view_lo)
# view_hi should be a cache miss and produce a distinct translation.
result_hi = pool.translate_loc_from_full_to_swa(view_hi)
# They should NOT be the same object (different cache entries).
self.assertIsNot(result_lo, result_hi)
# And the content must differ (different full indices → different swa).
self.assertFalse(torch.equal(result_lo, result_hi))
def test_storage_base_ptr_would_collide(self):
"""Demonstrate that untyped_storage().data_ptr() WOULD collide for the
two views above — confirming data_ptr() is the right key."""
t = torch.arange(20, device=get_device())
a, b = t[0:10], t[5:15]
# Same storage base — old key would collide.
self.assertEqual(a.untyped_storage().data_ptr(), b.untyped_storage().data_ptr())
self.assertEqual(a.numel(), b.numel())
# But data_ptr differs — new key is safe.
self.assertNotEqual(a.data_ptr(), b.data_ptr())
class TestAllocatorMutationInvalidation(CustomTestCase):
"""Each allocator method that writes the mapping calls invalidate_loc_cache."""
def _prime_and_check_invalidation(self, pool, allocator, mutate_fn):
"""Helper: prime cache, mutate, assert fresh translation."""
loc = allocator.alloc(4)
self.assertIsNotNone(loc)
# Prime the cache.
first = pool.translate_loc_from_full_to_swa(loc)
self.assertIsNotNone(pool._cached_loc_key)
# Mutate — should invalidate.
mutate_fn(allocator, loc)
# Cache must be cleared after mutation.
self.assertIsNone(pool._cached_loc_key)
self.assertIsNone(pool._cached_swa_loc)
def test_alloc_invalidates(self):
pool, allocator, _ = _build_pool()
loc = allocator.alloc(4)
pool.translate_loc_from_full_to_swa(loc)
self.assertIsNotNone(pool._cached_loc_key)
# Another alloc should invalidate.
allocator.alloc(4)
self.assertIsNone(pool._cached_loc_key)
def test_free_swa_invalidates(self):
pool, allocator, _ = _build_pool()
loc = allocator.alloc(4)
pool.translate_loc_from_full_to_swa(loc)
self.assertIsNotNone(pool._cached_loc_key)
allocator.free_swa(loc)
self.assertIsNone(pool._cached_loc_key)
def test_clear_invalidates(self):
pool, allocator, _ = _build_pool()
loc = allocator.alloc(4)
pool.translate_loc_from_full_to_swa(loc)
self.assertIsNotNone(pool._cached_loc_key)
allocator.clear()
self.assertIsNone(pool._cached_loc_key)
def test_set_full_to_swa_mapping_invalidates(self):
"""HiCache load-back path: set_full_to_swa_mapping must invalidate."""
pool, allocator, device = _build_pool(kv_size=32, kv_size_swa=32)
loc = allocator.alloc(4)
pool.translate_loc_from_full_to_swa(loc)
self.assertIsNotNone(pool._cached_loc_key)
# Simulate HiCache rebuild with new swa indices.
new_swa = torch.arange(4, dtype=torch.int64, device=device)
allocator.set_full_to_swa_mapping(loc, new_swa)
self.assertIsNone(pool._cached_loc_key)
# Translation after rebuild should reflect the new mapping.
result = pool.translate_loc_from_full_to_swa(loc)
self.assertEqual(result.tolist(), new_swa.tolist())
class TestBaseClassNoOp(CustomTestCase):
"""BaseSWAKVPool.invalidate_loc_cache is a no-op default — must not raise."""
def test_noop_does_not_raise(self):
# BaseSWAKVPool is abstract; instantiate via SWAKVPool which inherits.
pool, _, _ = _build_pool()
# Calling on the concrete class uses the override — that's fine.
pool.invalidate_loc_cache() # must not raise
pool.invalidate_loc_cache() # idempotent
def test_base_class_noop_directly(self):
"""Call the base-class method directly to verify it's a true no-op."""
pool, _, _ = _build_pool()
# Prime the cache first.
loc = pool.full_to_swa_index_mapping # any tensor
pool._cached_loc_key = ("dummy", 1)
pool._cached_swa_loc = torch.zeros(1)
# Call the BASE class method directly — should not clear the cache
# (it's a no-op; the concrete override is what clears).
BaseSWAKVPool.invalidate_loc_cache(pool)
# base no-op: cache untouched
self.assertIsNotNone(pool._cached_loc_key)
class TestExplicitInvalidationCycle(CustomTestCase):
"""Simulates the per-forward-pass invalidation done by model_runner."""
def test_fresh_translation_after_explicit_invalidation(self):
"""After invalidate_loc_cache(), a new alloc produces the right mapping."""
pool, allocator, device = _build_pool(kv_size=32, kv_size_swa=32)
# First "forward pass": alloc 4 tokens, translate.
loc1 = allocator.alloc(4)
trans1 = pool.translate_loc_from_full_to_swa(loc1).clone()
# Simulate start of next forward pass: model_runner calls invalidate.
pool.invalidate_loc_cache()
self.assertIsNone(pool._cached_loc_key)
# Alloc 4 more (mapping changes), translate loc1 again.
loc2 = allocator.alloc(4)
# Alloc already invalidated; translate loc1 with fresh mapping.
trans1_after = pool.translate_loc_from_full_to_swa(loc1)
# loc1's SWA mapping hasn't changed (same full→swa assignment),
# so result should be equal — but it must have been recomputed
# (cache key was None before this call).
self.assertEqual(trans1.tolist(), trans1_after.tolist())
# loc2 should have different translation than loc1.
trans2 = pool.translate_loc_from_full_to_swa(loc2)
# They have different indices, so translation differs.
self.assertFalse(torch.equal(trans1_after, trans2))
if __name__ == "__main__":
unittest.main()
test/registered/attention/unittests/dsv4/test_deepseek_v4.py
+71
View File
@@ -16,6 +16,7 @@ import importlib.util
import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
import torch
@@ -337,5 +338,75 @@ class TestDSV4AttentionBackendCorrectness(CustomTestCase):
run_dsv4_eagle_draft_extend_cuda_graph_runner_case(self, case)
class TestDSV4SwaOutCacheLocResolution(CustomTestCase):
"""`get_swa_out_cache_loc`: cached fast path vs store-time fallback.
The KV-store consumers run in paths that never invoke
`init_forward_metadata_in_graph` (eager idle, runners that only run the
out-graph prep) or whose batch is re-padded after init (DP attention).
The resolver must use the per-forward cached value only when it is
provably current and fall back to translating `out_cache_loc` otherwise.
"""
def _make_backend(self, mapping: torch.Tensor):
from sglang.srt.layers.attention.deepseek_v4_backend import (
DeepseekV4AttnBackend,
)
backend = object.__new__(DeepseekV4AttnBackend)
backend.forward_metadata = None
backend.token_to_kv_pool = SimpleNamespace(
translate_loc_from_full_to_swa=lambda loc: mapping[loc]
)
return backend
@staticmethod
def _make_fb(out_cache_loc: torch.Tensor, forward_mode: ForwardMode):
return SimpleNamespace(out_cache_loc=out_cache_loc, forward_mode=forward_mode)
@staticmethod
def _set_cached(backend, cached: torch.Tensor):
backend.forward_metadata = SimpleNamespace(
core_attn_metadata=SimpleNamespace(swa_out_cache_loc=cached)
)
def test_no_metadata_falls_back_to_translate(self):
mapping = torch.arange(10, dtype=torch.int64) * 2
backend = self._make_backend(mapping)
fb = self._make_fb(torch.tensor([3, 4]), ForwardMode.DECODE)
out = backend.get_swa_out_cache_loc(fb)
self.assertEqual(out.dtype, torch.int32)
self.assertEqual(out.tolist(), [6, 8])
def test_current_cached_value_is_used(self):
mapping = torch.arange(10, dtype=torch.int64) * 2
backend = self._make_backend(mapping)
cached = torch.tensor([6, 8], dtype=torch.int32)
self._set_cached(backend, cached)
fb = self._make_fb(torch.tensor([3, 4]), ForwardMode.DECODE)
self.assertIs(backend.get_swa_out_cache_loc(fb), cached)
def test_stale_shape_falls_back_to_translate(self):
# DP padding rebinds out_cache_loc to a longer tensor after init;
# a pre-pad cached value must not be used.
mapping = torch.arange(10, dtype=torch.int64) * 2
backend = self._make_backend(mapping)
self._set_cached(backend, torch.tensor([6, 8], dtype=torch.int32))
fb = self._make_fb(torch.tensor([3, 4, 0, 0]), ForwardMode.DRAFT_EXTEND_V2)
out = backend.get_swa_out_cache_loc(fb)
self.assertEqual(out.tolist(), [6, 8, 0, 0])
def test_idle_never_uses_cached_value(self):
# Idle forwards skip attn init, so any metadata is left over from a
# previous forward; writing dummy tokens to its locations would
# corrupt live KV. Idle must translate the zero-padded out_cache_loc.
mapping = torch.arange(10, dtype=torch.int64) * 2
backend = self._make_backend(mapping)
self._set_cached(backend, torch.tensor([6, 8], dtype=torch.int32))
fb = self._make_fb(torch.tensor([0, 0]), ForwardMode.IDLE)
out = backend.get_swa_out_cache_loc(fb)
self.assertEqual(out.tolist(), [0, 0])
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_swa_alloc_extend_page_estimation.py
-1
View File
@@ -33,7 +33,6 @@ def _make_self(*, page_size: int, full_available: int, swa_available: int):
),
translate_loc_from_full_to_swa=lambda last_loc: last_loc,
full_to_swa_index_mapping=torch.zeros(64, dtype=torch.int64),
_kvcache=SimpleNamespace(invalidate_loc_cache=lambda: None),
)