A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 22:30 UTC
[mamba] unify KDA conv states into one cache to match 2-state SSM layout (#44539)
Jiangyun Zhu committed
Jun 5, 2026 at 02:38 UTC
a55fccfc7cefc2d085d3557bace0e345cef67961
3 files changed
+16
-30
vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py
+6
-6
@@ -85,7 +85,7 @@ direct_register_custom_op(
85
class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
86
def get_state_dtype(
87
self,
88
- ) -> tuple[torch.dtype, torch.dtype, torch.dtype, torch.dtype]:
88
+ ) -> tuple[torch.dtype, torch.dtype]:
89
if self.model_config is None or self.cache_config is None:
90
raise ValueError("model_config and cache_config must be set")
91
return MambaStateDtypeCalculator.kda_state_dtype(
@@ -94,7 +94,7 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
94
95
def get_state_shape(
96
self,
97
- ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
97
+ ) -> tuple[tuple[int, ...], tuple[int, ...]]:
98
return MambaStateShapeCalculator.kda_state_shape(
99
self.tp_size, self.num_heads, self.head_dim, conv_kernel_size=self.conv_size
100
)
@@ -300,13 +300,13 @@ class KimiGatedDeltaNetAttention(GatedDeltaNetAttention):
300
g1 = g1[:, :num_actual_tokens]
301
beta = beta[:, :num_actual_tokens]
302
303
- (conv_state_q, conv_state_k, conv_state_v, recurrent_state) = constant_caches
303
+ (conv_state, recurrent_state) = constant_caches
304
# conv_state must be (..., dim, width-1) for the conv kernels.
305
# DS layout stores it that way directly; SD layout needs a transpose.
306
if not is_conv_state_dim_first():
307
- conv_state_q = conv_state_q.transpose(-1, -2)
308
- conv_state_k = conv_state_k.transpose(-1, -2)
309
- conv_state_v = conv_state_v.transpose(-1, -2)
307
+ conv_state = conv_state.transpose(-1, -2)
308
+
309
+ conv_state_q, conv_state_k, conv_state_v = conv_state.chunk(3, dim=-2)
310
311
q_conv_weights = self.q_conv1d.weight.view(
312
self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2)
vllm/model_executor/layers/mamba/mamba_utils.py
+7
-19
@@ -120,9 +120,9 @@ class MambaStateDtypeCalculator:
120
cls,
121
model_dtype: ModelDType | torch.dtype,
122
mamba_cache_dtype: MambaDType,
123
- ):
123
+ ) -> tuple[torch.dtype, torch.dtype]:
124
state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
125
- return (state_dtype, state_dtype, state_dtype, torch.float32)
125
+ return (state_dtype, torch.float32)
126
127
128
class MambaStateShapeCalculator:
@@ -243,7 +243,7 @@ class MambaStateShapeCalculator:
243
head_k_dim: int | None = None,
244
conv_kernel_size: int = 4,
245
num_spec: int = 0,
246
- ) -> tuple[tuple[int, int], tuple[int, int], tuple[int, int], tuple[int, int, int]]:
246
+ ) -> tuple[tuple[int, int], tuple[int, int, int]]:
247
if num_k_heads is None:
248
num_k_heads = num_heads
249
if head_k_dim is None:
@@ -252,19 +252,12 @@ class MambaStateShapeCalculator:
252
proj_size = num_heads * head_dim
253
proj_k_size = num_k_heads * head_k_dim
254
255
+ conv_dim = proj_size + 2 * proj_k_size
256
conv_state_shape = cls._orient_conv_shape(
256
- divide(proj_size, tp_world_size), conv_kernel_size - 1
257
- )
258
- conv_state_k_shape = cls._orient_conv_shape(
259
- divide(proj_k_size, tp_world_size), conv_kernel_size - 1
257
+ divide(conv_dim, tp_world_size), conv_kernel_size - 1
258
)
259
recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim)
262
- return (
263
- conv_state_shape,
264
- conv_state_k_shape,
265
- conv_state_k_shape,
266
- recurrent_state_shape,
267
- )
260
+ return (conv_state_shape, recurrent_state_shape)
261
262
263
@dataclass
@@ -365,9 +358,4 @@ class MambaStateCopyFuncCalculator:
358
359
@classmethod
360
def kda_state_copy_func(cls):
368
- return (
369
- get_conv_copy_spec,
370
- get_conv_copy_spec,
371
- get_conv_copy_spec,
372
- get_temporal_copy_spec,
373
- )
361
+ return (get_conv_copy_spec, get_temporal_copy_spec)
vllm/model_executor/models/kimi_linear.py
+3
-5
@@ -600,7 +600,7 @@ class KimiLinearForCausalLM(
600
def get_mamba_state_dtype_from_config(
601
cls,
602
vllm_config: "VllmConfig",
603
- ) -> tuple[torch.dtype, torch.dtype, torch.dtype, torch.dtype]:
603
+ ) -> tuple[torch.dtype, torch.dtype]:
604
return MambaStateDtypeCalculator.kda_state_dtype(
605
vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
606
)
@@ -608,7 +608,7 @@ class KimiLinearForCausalLM(
608
@classmethod
609
def get_mamba_state_shape_from_config(
610
cls, vllm_config: "VllmConfig"
611
- ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
611
+ ) -> tuple[tuple[int, ...], tuple[int, ...]]:
612
parallel_config = vllm_config.parallel_config
613
hf_config = vllm_config.model_config.hf_config
614
tp_size = parallel_config.tensor_parallel_size
@@ -628,9 +628,7 @@ class KimiLinearForCausalLM(
628
@classmethod
629
def get_mamba_state_copy_func(
630
cls,
631
- ) -> tuple[
632
- MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc
633
- ]:
631
+ ) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
632
return MambaStateCopyFuncCalculator.kda_state_copy_func()
633
634
def compute_logits(