[Mamba] Support configurable conv-window layouts (#31059)

2026-07-14 14:41:10 -07:00
parent 08c46e1f1a
commit 463a3f4248
2 changed files with 113 additions and 32 deletions
@@ -18,7 +18,11 @@ from sglang.srt.mem_cache.cache_init_params import CacheInitParams
 from sglang.srt.mem_cache.common import available_and_evictable_str
 from sglang.srt.mem_cache.hi_mamba_radix_cache import HiMambaRadixCache
 from sglang.srt.mem_cache.mamba_radix_cache import LRUList, MambaRadixCache, TreeNode
-from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool
+from sglang.srt.mem_cache.memory_pool import (
+    HybridLinearKVPool,
+    HybridReqToTokenPool,
+    MambaPool,
+)
 from sglang.srt.mem_cache.radix_cache import RadixKey
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
@@ -154,6 +158,61 @@ class TestMamba(unittest.TestCase):
            req_to_token_pool.mamba_allocator.available_size() == mamba_cache_size - 1
        )

+    def test_mamba_pool_deduplicated_conv_window_axis(self):
+        class WindowFirstMambaPool(MambaPool):
+            conv_window_axis = 0
+
+        num_mamba_layers = 2
+        spec_state_size = 3
+        speculative_num_draft_tokens = 4
+        window_size = 3
+        conv_dim = 5
+
+        pool = object.__new__(WindowFirstMambaPool)
+        physical, view = pool._allocate_deduplicated_conv_window(
+            conv_shape=(window_size, conv_dim),
+            num_mamba_layers=num_mamba_layers,
+            spec_state_size=spec_state_size,
+            speculative_num_draft_tokens=speculative_num_draft_tokens,
+            conv_dtype=torch.float32,
+        )
+
+        shared_window_size = speculative_num_draft_tokens + window_size - 1
+        self.assertEqual(
+            physical.shape,
+            (
+                num_mamba_layers,
+                spec_state_size + 1,
+                shared_window_size,
+                conv_dim,
+            ),
+        )
+        self.assertEqual(
+            view.shape,
+            (
+                num_mamba_layers,
+                spec_state_size + 1,
+                speculative_num_draft_tokens,
+                window_size,
+                conv_dim,
+            ),
+        )
+
+        physical.copy_(
+            torch.arange(
+                physical.numel(), dtype=physical.dtype, device=physical.device
+            ).reshape_as(physical)
+        )
+        for step in range(speculative_num_draft_tokens):
+            torch.testing.assert_close(
+                view[:, :, step],
+                physical[:, :, step : step + window_size],
+            )
+        torch.testing.assert_close(view[:, :, :-1, 1:], view[:, :, 1:, :-1])
+
+        view[0, 0, 0, 1, 0] = -1
+        self.assertEqual(view[0, 0, 1, 0, 0].item(), -1)
+
    def test_mamba_radix_cache_1(self):
        tree, allocator, req_to_token_pool, make_dummy_req = (
            self._setup_tree_and_allocator()