refactor linear memory pool (#13004)

python/sglang/srt/configs/mamba_utils.py

@@ -13,6 +13,7 @@
 """Common config utils for mamba2 - NemotronH, FalconH1, Qwen3Next, etc."""
 
 import os
+from abc import ABC
 from dataclasses import dataclass, field
 from typing import List, Optional
 
@@ -34,9 +35,45 @@ def extra_groups_for_head_shards(ngroups: int, tp_size: int):
     return tp_size - ngroups
 
 
+@dataclass(kw_only=True, frozen=True)
+class Mamba2StateDType:
+    conv: torch.dtype
+    temporal: torch.dtype
+
+
+CONV_DTYPE = torch.bfloat16
+
+
+def mamba2_state_dtype() -> Mamba2StateDType:
+    dtype_map = {
+        "float32": torch.float32,
+        "bfloat16": torch.bfloat16,
+    }
+    ssm_dtype = dtype_map[os.environ["SGLANG_MAMBA_SSM_DTYPE"]]
+    return Mamba2StateDType(conv=CONV_DTYPE, temporal=ssm_dtype)
+
+
+@dataclass(kw_only=True, frozen=True)
+class BaseLinearStateParams(ABC):
+    dtype: Mamba2StateDType = field(default_factory=mamba2_state_dtype)
+    layers: list[int]
+
+    @property
+    def mamba_cache_per_req(self) -> int:
+        conv_numel = int(
+            np.sum([np.prod(conv_shape) for conv_shape in self.shape.conv])
+        )
+
+        ssm_numel = int(np.prod(self.shape.temporal))
+        return (
+            conv_numel * self.dtype.conv.itemsize
+            + ssm_numel * self.dtype.temporal.itemsize
+        ) * len(self.layers)
+
+
 @dataclass(kw_only=True, frozen=True)
 class Mamba2StateShape:
-    conv: tuple[int, int]
+    conv: list[tuple[int, int]]
     temporal: tuple[int, int, int]
 
     intermediate_size: int
@@ -74,7 +111,7 @@ class Mamba2StateShape:
         #   e.g., QWen3-Next: (32, 128, 128)
         temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, state_size)
         return Mamba2StateShape(
-            conv=conv_state_shape,
+            conv=[conv_state_shape],
             temporal=temporal_state_shape,
             intermediate_size=intermediate_size,
             conv_dim=conv_dim,
@@ -87,35 +124,8 @@ class Mamba2StateShape:
 
 
 @dataclass(kw_only=True, frozen=True)
-class Mamba2StateDType:
-    conv: torch.dtype
-    temporal: torch.dtype
-
-
-CONV_DTYPE = torch.bfloat16
-
-
-def mamba2_state_dtype() -> Mamba2StateDType:
-    dtype_map = {
-        "float32": torch.float32,
-        "bfloat16": torch.bfloat16,
-    }
-    ssm_dtype = dtype_map[os.environ["SGLANG_MAMBA_SSM_DTYPE"]]
-    return Mamba2StateDType(conv=CONV_DTYPE, temporal=ssm_dtype)
-
-
-@dataclass(kw_only=True, frozen=True)
-class Mamba2CacheParams:
+class Mamba2CacheParams(BaseLinearStateParams):
     shape: Mamba2StateShape
-    dtype: Mamba2StateDType = field(default_factory=mamba2_state_dtype)
-    layers: list[int]
-
-    @property
-    def mamba_cache_per_req(self) -> int:
-        return (
-            int(np.prod(self.shape.conv)) * self.dtype.conv.itemsize
-            + int(np.prod(self.shape.temporal)) * self.dtype.temporal.itemsize
-        ) * len(self.layers)
 
 
 @dataclass(kw_only=True, frozen=True)
@@ -169,15 +179,5 @@ class KimiLinearStateShape:
 
 
 @dataclass(kw_only=True, frozen=True)
-class KimiLinearCacheParams:
+class KimiLinearCacheParams(BaseLinearStateParams):
     shape: KimiLinearStateShape
-    dtype: Mamba2StateDType = field(default_factory=mamba2_state_dtype)
-    layers: list[int]
-
-    @property
-    def mamba_cache_per_req(self) -> int:
-        return (
-            int(np.sum([np.prod(conv_shape) for conv_shape in self.shape.conv]))
-            * self.dtype.conv.itemsize
-            + int(np.prod(self.shape.temporal)) * self.dtype.temporal.itemsize
-        ) * len(self.layers)