From 42322947aa32b99e3d8d98731a070507fe67b535 Mon Sep 17 00:00:00 2001
From: ziang663 <119752791+ziang663@users.noreply.github.com>
Date: Wed, 10 Jun 2026 05:03:38 +0800
Subject: [PATCH] [BUG FIX]Fix DSA CPU offload mamba indices signature (#27645)

---
 python/sglang/srt/mem_cache/memory_pool.py             | 10 ++++++----
 .../unit/mem_cache/test_dsa_pool_host_unit.py          |  9 +++++++++
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py
index 8efe9aae9..dc57f9bb1 100644
--- a/python/sglang/srt/mem_cache/memory_pool.py
+++ b/python/sglang/srt/mem_cache/memory_pool.py
@@ -2360,13 +2360,13 @@ class DSATokenToKVPool(MLATokenToKVPool):
             pool=self, buf=buf, loc=loc, index_k=index_k, index_k_scale=index_k_scale
         )
 
-    def get_cpu_copy(self, indices):
+    def get_cpu_copy(self, indices, mamba_indices=None):
         # DSA keeps a page-indexed index_k_with_scale_buffer alongside kv_buffer.
         # Retract frees the slots/pages and they get reused by other reqs'
         # set_index_k_scale_buffer, so we must offload it here too -- otherwise
         # resume restores kv_buffer but leaves foreign index/scale in place and
         # DSA attention reads garbage at those token positions.
-        kv_cache_cpu = super().get_cpu_copy(indices)
+        kv_cache_cpu = super().get_cpu_copy(indices, mamba_indices=mamba_indices)
 
         page_indices = indices[:: self.page_size] // self.page_size
         torch.cuda.synchronize()
@@ -2385,8 +2385,10 @@ class DSATokenToKVPool(MLATokenToKVPool):
 
         return {"kv": kv_cache_cpu, "index_k": index_k_cpu}
 
-    def load_cpu_copy(self, kv_cache_cpu_dict, indices):
-        super().load_cpu_copy(kv_cache_cpu_dict["kv"], indices)
+    def load_cpu_copy(self, kv_cache_cpu_dict, indices, mamba_indices=None):
+        super().load_cpu_copy(
+            kv_cache_cpu_dict["kv"], indices, mamba_indices=mamba_indices
+        )
 
         page_indices = indices[:: self.page_size] // self.page_size
         index_k_cpu = kv_cache_cpu_dict["index_k"]
diff --git a/test/registered/unit/mem_cache/test_dsa_pool_host_unit.py b/test/registered/unit/mem_cache/test_dsa_pool_host_unit.py
index 939e40f5e..27b102539 100644
--- a/test/registered/unit/mem_cache/test_dsa_pool_host_unit.py
+++ b/test/registered/unit/mem_cache/test_dsa_pool_host_unit.py
@@ -1,3 +1,4 @@
+import inspect
 import unittest
 
 import torch
@@ -15,6 +16,14 @@ from sglang.test.ci.ci_register import register_cuda_ci
 register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small")
 
 
+class TestDSAOffloadSignatures(unittest.TestCase):
+    def test_cpu_copy_methods_accept_mamba_indices(self):
+        for method_name in ("get_cpu_copy", "load_cpu_copy"):
+            with self.subTest(method_name=method_name):
+                signature = inspect.signature(getattr(DSATokenToKVPool, method_name))
+                self.assertIn("mamba_indices", signature.parameters)
+
+
 class TestDSAHiCacheTransfer(unittest.TestCase):
     def setUp(self):
         if not torch.cuda.is_available():