From 36d61613a1659ee2d4d50d6bc1fdd3c1837fb332 Mon Sep 17 00:00:00 2001 From: Xinyi Song Date: Thu, 11 Jun 2026 23:13:15 -0700 Subject: [PATCH] [AMD] Cache unified_kv swa_loc once per step instead of per layer (#27978) Co-authored-by: Thomas Wang <1am9trash@gmail.com> --- .../deepseek_v4_backend_hip_radix.py | 37 +++++++++++++++++++ python/sglang/srt/models/deepseek_v4.py | 14 ++----- 2 files changed, 40 insertions(+), 11 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 426c5749b..9db320671 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -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: diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 127d743b9..9387ac88d 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -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]