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,
)