From 274d1425ebe4d0ac529eec178c2d587587edaee5 Mon Sep 17 00:00:00 2001
From: Mohammad Angkad <mohammad.angkad@radixark.ai>
Date: Sat, 19 Sep 2026 13:09:35 +0000
Subject: [PATCH] Forward SWA prealloc reclaim through the DSV4 HiSparse
 allocator

---
 .../srt/mem_cache/allocator/hisparse.py       |  8 ++++++++
 .../unit/mem_cache/test_hisparse_allocator.py | 19 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py
index 706df84bf50a..9508aa1502ea 100644
--- a/python/sglang/srt/mem_cache/allocator/hisparse.py
+++ b/python/sglang/srt/mem_cache/allocator/hisparse.py
@@ -376,6 +376,14 @@ def full_available_size(self):
     def swa_available_size(self):
         return self.logical_attn_allocator.swa_available_size()
 
+    def reclaim_for_prealloc(
+        self, tree_cache, full_tokens: int, swa_tokens: int
+    ) -> str | None:
+        # C4 needs no reclaim here: full_available_size prices it into the budget.
+        return self.logical_attn_allocator.reclaim_for_prealloc(
+            tree_cache, full_tokens, swa_tokens
+        )
+
     def free_swa(self, free_indices: torch.Tensor):
         self.logical_attn_allocator.free_swa(free_indices)
 
diff --git a/test/registered/unit/mem_cache/test_hisparse_allocator.py b/test/registered/unit/mem_cache/test_hisparse_allocator.py
index d0fc8615635e..f2049b8d3315 100644
--- a/test/registered/unit/mem_cache/test_hisparse_allocator.py
+++ b/test/registered/unit/mem_cache/test_hisparse_allocator.py
@@ -60,6 +60,25 @@ def test_forwards_swa_tail_allocation_to_logical_allocator(self):
         self.assertEqual(kwargs["extend_num_tokens"], 512)
         self.assertEqual(kwargs["swa_tail_len"], 128)
 
+    def test_forwards_prealloc_reclaim_to_logical_allocator(self):
+        """PD decode preallocation must not crash on the HiSparse composite."""
+        allocator = object.__new__(DeepSeekV4HiSparseTokenToKVPoolAllocator)
+        logical_allocator = MagicMock(spec=["reclaim_for_prealloc"])
+        allocator.logical_attn_allocator = logical_allocator
+        logical_allocator.reclaim_for_prealloc.return_value = None
+
+        tree_cache = object()
+        self.assertIsNone(allocator.reclaim_for_prealloc(tree_cache, 512, 256))
+        logical_allocator.reclaim_for_prealloc.assert_called_once_with(
+            tree_cache, 512, 256
+        )
+
+        logical_allocator.reclaim_for_prealloc.return_value = "SWA eviction short"
+        self.assertEqual(
+            allocator.reclaim_for_prealloc(tree_cache, 512, 256),
+            "SWA eviction short",
+        )
+
     def test_hisparse_budget_uses_full_logical_capacity_for_swa_tail(self):
         from sglang.srt.disaggregation.decode import DecodePreallocQueue