model: support baidu unlimited-ocr (#29186)

Co-authored-by: Mick <mickjagger19@icloud.com>

test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py

@@ -0,0 +1,48 @@
+"""Unit tests for all-SWA ChunkCache release semantics."""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.mem_cache.chunk_cache import PureSWAChunkCache
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=3, suite="base-a-test-cpu")
+
+
+class _FakeAllocator:
+    def __init__(self):
+        self.freed = []
+
+    def free(self, indices):
+        self.freed.append(indices.detach().cpu().clone())
+
+
+class _FakeReq:
+    req_pool_idx = 0
+    swa_evict_floor = 3
+    swa_evicted_seqlen = 6
+
+    def pop_committed_kv_cache(self):
+        return 8
+
+
+class TestPureSWAChunkCache(CustomTestCase):
+    def test_finished_req_skips_already_evicted_swa_range(self):
+        cache = PureSWAChunkCache.__new__(PureSWAChunkCache)
+        cache.req_to_token_pool = SimpleNamespace(
+            req_to_token=torch.arange(10, dtype=torch.int64).unsqueeze(0)
+        )
+        cache.token_to_kv_pool_allocator = _FakeAllocator()
+
+        cache.cache_finished_req(_FakeReq())
+
+        self.assertEqual(len(cache.token_to_kv_pool_allocator.freed), 1)
+        freed = cache.token_to_kv_pool_allocator.freed[0]
+        self.assertTrue(torch.equal(freed, torch.tensor([0, 1, 2, 6, 7])))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_registry.py

@@ -29,6 +29,7 @@ def _make_ctx(
     enable_hierarchical_cache=False,
     disable_radix_cache=False,
     effective_chunked_prefill_size=None,
+    full_tokens_per_layer=None,
 ):
     server_args = MagicMock()
     server_args.radix_cache_backend = backend
@@ -47,6 +48,7 @@ def _make_ctx(
         tp_size=1,
         tp_rank=0,
         tp_group=MagicMock(),
+        full_tokens_per_layer=full_tokens_per_layer,
     )
 
 
@@ -172,6 +174,21 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
             SWAChunkCache.assert_called_once_with(ctx.params)
             self.assertIs(result, SWAChunkCache.return_value)
 
+    def test_pure_swa_chunk_cache_when_chunked_prefill_disable_and_all_swa(self):
+        ctx = _make_ctx(
+            effective_chunked_prefill_size=512,
+            disable_radix_cache=True,
+            is_hybrid_swa=True,
+            full_tokens_per_layer=0,
+        )
+        with patch(
+            "sglang.srt.mem_cache.chunk_cache.PureSWAChunkCache"
+        ) as PureSWAChunkCache:
+            PureSWAChunkCache.return_value = MagicMock()
+            result = default_radix_cache_factory(ctx)
+            PureSWAChunkCache.assert_called_once_with(ctx.params)
+            self.assertIs(result, PureSWAChunkCache.return_value)
+
     def test_cpp_radix_cache_when_env_flag_set(self):
         ctx = _make_ctx()
         # `radix_cache_cpp` requires ninja + C++ extension to import, so
@@ -281,6 +298,16 @@ class TestDefaultRadixCacheFactory(CustomTestCase):
             SWA.assert_called_once_with(params=ctx.params)
             self.assertIs(result, SWA.return_value)
 
+    def test_pure_swa_radix_cache_when_all_swa(self):
+        ctx = _make_ctx(is_hybrid_swa=True, full_tokens_per_layer=0)
+        with patch(
+            "sglang.srt.mem_cache.pure_swa_radix_cache.PureSWARadixCache"
+        ) as PureSWA:
+            PureSWA.return_value = MagicMock()
+            result = default_radix_cache_factory(ctx)
+            PureSWA.assert_called_once_with(params=ctx.params)
+            self.assertIs(result, PureSWA.return_value)
+
     def test_mamba_radix_cache_when_hybrid_ssm(self):
         ctx = _make_ctx(is_hybrid_ssm=True)
         with patch("sglang.srt.mem_cache.mamba_radix_cache.MambaRadixCache") as Mamba: