minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
paulzhang-tm
2026-07-14 14:41:10 -07:00
committed by GitHub
parent 08c46e1f1a
commit 463a3f4248
2 changed files with 113 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/memory_pool.py
+53 -31
View File
@@ -311,6 +311,10 @@ class ReqToTokenPool:
class MambaPool:
# Axis of each two-dimensional conv state that represents the sliding window.
# Upstream states use (dim, K-1); subclasses may preserve another layout.
conv_window_axis = -1
@dataclass(frozen=True, kw_only=True)
class State:
conv: List[torch.Tensor]
@@ -352,6 +356,46 @@ class MambaPool:
intermediate_ssm: torch.Tensor
intermediate_conv_window: List[torch.Tensor]
def _allocate_deduplicated_conv_window(
self,
*,
conv_shape: Tuple[int, int],
num_mamba_layers: int,
spec_state_size: int,
speculative_num_draft_tokens: int,
conv_dtype: torch.dtype,
) -> Tuple[torch.Tensor, torch.Tensor]:
window_axis = self.conv_window_axis % len(conv_shape)
win = conv_shape[window_axis]
physical_conv_shape = list(conv_shape)
physical_conv_shape[window_axis] = speculative_num_draft_tokens + win - 1
phys = torch.zeros(
(
num_mamba_layers,
spec_state_size + 1,
*physical_conv_shape,
),
dtype=conv_dtype,
device="cuda",
)
physical_conv_strides = phys.stride()[2:]
window_stride = physical_conv_strides[window_axis]
view = phys.as_strided(
(
phys.shape[0],
phys.shape[1],
speculative_num_draft_tokens,
*conv_shape,
),
(
phys.stride(0),
phys.stride(1),
window_stride,
*physical_conv_strides,
),
)
return phys, view
def __init__(
self,
*,
@@ -538,36 +582,12 @@ class MambaPool:
if dedup_conv_window:
intermediate_conv_window_cache = []
for conv_shape in conv_state_shape:
conv_dim, win = conv_shape # win == conv_kernel - 1 == K-1
shared_win = (
speculative_num_draft_tokens + win - 1
) # D + (K-1) - 1
phys = torch.zeros(
size=(
num_mamba_layers,
spec_state_size + 1,
conv_dim,
shared_win,
),
dtype=conv_dtype,
device="cuda",
)
# view[l, s, step, d, w] = phys[l, s, d, step + w]
view = phys.as_strided(
(
phys.shape[0],
phys.shape[1],
speculative_num_draft_tokens,
conv_dim,
win,
),
(
phys.stride(0),
phys.stride(1),
phys.stride(3), # step -> shared-win axis (stride 1)
phys.stride(2), # dim
phys.stride(3), # win -> shared-win axis (stride 1)
),
phys, view = self._allocate_deduplicated_conv_window(
conv_shape=conv_shape,
num_mamba_layers=num_mamba_layers,
spec_state_size=spec_state_size,
speculative_num_draft_tokens=speculative_num_draft_tokens,
conv_dtype=conv_dtype,
)
self._intermediate_conv_window_phys.append(phys)
intermediate_conv_window_cache.append(view)
@@ -818,6 +838,8 @@ class MambaPool:
class HybridReqToTokenPool(ReqToTokenPool):
"""A memory pool that maps a request to its token locations."""
mamba_pool_cls = MambaPool
def __init__(
self,
*,
@@ -880,7 +902,7 @@ class HybridReqToTokenPool(ReqToTokenPool):
linear_replayssm_cache_len: int = 16,
mamba_envelope_layout: bool = False,
):
self.mamba_pool = MambaPool(
self.mamba_pool = self.mamba_pool_cls(
size=mamba_size,
spec_state_size=mamba_spec_state_size,
cache_params=cache_params,
test/registered/unit/mem_cache/test_mamba_unittest.py
+60 -1
View File
@@ -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()