[BUG FIX]Fix DSA CPU offload mamba indices signature (#27645)

This commit is contained in:
ziang663
2026-06-09 14:03:38 -07:00
committed by GitHub
parent eb8dceda44
commit 42322947aa
2 changed files with 15 additions and 4 deletions
@@ -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():