From d15a2dc72c81a0c2dd72074f3e208bfa38df831b Mon Sep 17 00:00:00 2001
From: Thomas Wang <thomawan@amd.com>
Date: Tue, 2 Jun 2026 13:37:07 +0800
Subject: [PATCH] [AMD] dpsk-v4 swa loc cache support (#26931)

---
 .../srt/mem_cache/deepseek_v4_memory_pool.py  | 21 ++++++++-----------
 python/sglang/srt/models/deepseek_v4.py       |  8 +++----
 2 files changed, 13 insertions(+), 16 deletions(-)

diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
index b44388f37..94902677e 100644
--- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
+++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py
@@ -512,6 +512,13 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
 
         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)
+
     def get_contiguous_buf_infos(self) -> Tuple[List[int], List[int], List[int]]:
         data_ptrs: List[int] = []
         data_lens: List[int] = []
@@ -758,12 +765,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         raw_loc: torch.Tensor,
         cache_k: torch.Tensor,
     ) -> None:
-        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)
-            swa_loc = self.cached_loc
-        else:
-            swa_loc = self.translate_loc_from_full_to_swa(raw_loc)
+        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
         )
@@ -778,12 +780,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool):
         freqs_cis: torch.Tensor,
         positions: torch.Tensor,
     ) -> None:
-        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)
-            swa_loc = self.cached_loc
-        else:
-            swa_loc = self.translate_loc_from_full_to_swa(raw_loc)
+        swa_loc = self.get_cached_swa_loc(raw_loc, layer_id)
         fused_k_norm_rope_flashmla(
             kv=kv,
             kv_weight=kv_weight,
diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index d0c04d953..953a36b23 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -646,8 +646,8 @@ class MQALayer(nn.Module):
             )
 
             token_to_kv_pool = get_token_to_kv_pool()
-            swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(
-                forward_batch.out_cache_loc
+            swa_loc = token_to_kv_pool.get_cached_swa_loc(
+                forward_batch.out_cache_loc, self.layer_id
             )
             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
@@ -736,8 +736,8 @@ class MQALayer(nn.Module):
             )
 
             token_to_kv_pool = get_token_to_kv_pool()
-            swa_loc = token_to_kv_pool.translate_loc_from_full_to_swa(
-                forward_batch.out_cache_loc
+            swa_loc = token_to_kv_pool.get_cached_swa_loc(
+                forward_batch.out_cache_loc, self.layer_id
             )
             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