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

[AMD] Cache unified_kv swa_loc once per step instead of per layer (#27978)

Browse Source 
Co-authored-by: Thomas Wang <1am9trash@gmail.com>
This commit is contained in:
Xinyi Song
2026-06-11 23:13:15 -07:00
committed by GitHub
co-authored by Thomas Wang
parent ca17bd8347
commit 36d61613a1
2 changed files with 40 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py
+37
View File
@@ -123,6 +123,9 @@ class DSV4AttnMetadata:
c128_topk_lengths_raw: Optional[torch.Tensor] = None
# unified_kv: per-forward prebuilt ragged decode index
# SWA ring write target (req_slot*ring + pos%ring), computed once per
# forward in _attach_unified_kv_decode_streams, read by every layer's store.
unified_swa_loc: Optional[torch.Tensor] = None
unified_swa_indices: Optional[torch.Tensor] = None
unified_swa_indptr: Optional[torch.Tensor] = None
unified_hca_indices: Optional[torch.Tensor] = None
@@ -196,6 +199,7 @@ class DSV4AttnMetadata:
# Recomputed by the recorded init_forward_metadata_in_graph op
# each forward; not copied across replays.
"swa_out_cache_loc",
"unified_swa_loc",
"c1_flashmla_metadata",
"c4_flashmla_metadata",
"c128_flashmla_metadata",
@@ -1028,6 +1032,13 @@ class DeepseekV4HipRadixBackend(
ring_stride=pool.unified_swa_ring_size,
swa_pages=pool.unified_swa_pages,
)
# SWA ring write target, same value for every layer this forward.
# Decode: N tokens == N reqs, positions already aligned (no repeat).
req_slot = req_pool_indices[:N].to(torch.int64)
core.unified_swa_loc = (
req_slot * pool.unified_swa_ring_size
+ core.positions_casual.to(torch.int64) % pool.unified_swa_ring_size
).to(torch.int32)
def _attach_unified_kv_prefill_meta(
self,
@@ -1206,6 +1217,32 @@ class DeepseekV4HipRadixBackend(
torch.int32
)
def get_unified_swa_loc(self, forward_batch: ForwardBatch) -> torch.Tensor:
"""SWA ring write target for unified_kv, shared by all layers.
Fast path: the per-forward value cached in _attach_unified_kv_decode_streams
(recorded inside cuda graphs, so replay re-reads live buffers). Fallback:
recompute at store time, matching the pre-cache per-layer behavior, for
paths that never ran the decode-stream init (eager prefill/extend, idle,
or a batch re-padded after init -> shape mismatch).
"""
positions = forward_batch.positions
core = getattr(self.forward_metadata, "core_attn_metadata", None)
cached = core.unified_swa_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] == positions.shape[0]
):
return cached
ring = self.token_to_kv_pool.unified_swa_ring_size
req_slot = forward_batch.req_pool_indices.to(torch.int64)
if req_slot.shape[0] != positions.shape[0]:
req_slot = req_slot.repeat_interleave(
positions.shape[0] // req_slot.shape[0]
)
return (req_slot * ring + positions.to(torch.int64) % ring).to(torch.int32)
def store_cache(
self, layer_id: int, swa_k: torch.Tensor, forward_batch: ForwardBatch
) -> None:
python/sglang/srt/models/deepseek_v4.py
+3 -11
View File
@@ -765,18 +765,10 @@ class MQALayer(nn.Module):
token_to_kv_pool = get_token_to_kv_pool()
if unified:
swa_ring_size = token_to_kv_pool.unified_swa_ring_size
swa_cache = token_to_kv_pool.get_unified_kv(self.layer_id)
# ring slot = req_slot * ring + pos % ring, per token.
# positions is per-token; req_pool_indices is per-req.
req_slot = forward_batch.req_pool_indices.to(torch.int64)
if req_slot.shape[0] != positions.shape[0]:
req_slot = req_slot.repeat_interleave(
positions.shape[0] // req_slot.shape[0]
)
swa_loc = (
req_slot * swa_ring_size + positions.to(torch.int64) % swa_ring_size
).to(torch.int32)
# swa_loc is layer-independent; computed once per forward by the
# backend and cached on the metadata (read here by every layer).
swa_loc = attn_backend.get_unified_swa_loc(forward_batch)
swa_page_size, bf16_store = 1, True
else:
swa_cache = token_to_kv_pool.swa_kv_pool.kv_buffer[self.layer_id]