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 5, 14:45 UTC
[Kernel] Mamba support different layout for Conv state (#37416)
Nicolò Lucchesi committed
Apr 3, 2026 at 01:50 UTC
66e86f1dbd565292a253e7d2d6851f65dc4f14ba
11 files changed
+169
-39
tests/models/language/generation/test_hybrid.py
+23
@@ -60,6 +60,14 @@ MAX_NUM_SEQS = 4
60
ATTN_BACKEND = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
61
62
63
+def _set_conv_state_layout(monkeypatch, layout: str) -> None:
64
+ """Set conv state layout env var and clear cache to pick up new value."""
65
+ from vllm.model_executor.layers.mamba import mamba_utils
66
+
67
+ monkeypatch.setenv("VLLM_SSM_CONV_STATE_LAYOUT", layout)
68
+ mamba_utils.get_conv_state_layout.cache_clear()
69
+
70
+
71
@pytest.mark.parametrize("model", SSM_MODELS + HYBRID_MODELS)
72
@pytest.mark.parametrize("max_tokens", [64])
73
@pytest.mark.parametrize("num_logprobs", [5])
@@ -102,12 +110,15 @@ def test_models(
110
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
111
@pytest.mark.parametrize("max_tokens", [64])
112
@pytest.mark.parametrize("num_logprobs", [5])
113
+@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"])
114
def test_batching(
115
vllm_runner,
116
example_prompts,
117
+ monkeypatch,
118
model: str,
119
max_tokens: int,
120
num_logprobs: int,
121
+ conv_state_layout: str,
122
) -> None:
123
try:
124
model_info = HF_EXAMPLE_MODELS.find_hf_info(model)
@@ -116,6 +127,8 @@ def test_batching(
127
except ValueError:
128
pass
129
130
+ _set_conv_state_layout(monkeypatch, conv_state_layout)
131
+
132
for_loop_outputs = []
133
with vllm_runner(model, max_num_seqs=MAX_NUM_SEQS) as vllm_model:
134
for prompt in example_prompts:
@@ -138,11 +151,14 @@ def test_batching(
151
152
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
153
@pytest.mark.parametrize("max_tokens", [10])
154
+@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"])
155
def test_chunked_prefill_with_parallel_sampling(
156
vllm_runner,
157
example_prompts,
158
+ monkeypatch,
159
model: str,
160
max_tokens: int,
161
+ conv_state_layout: str,
162
) -> None:
163
"""
164
Tests chunked prefill in conjunction with n > 1.
@@ -154,6 +170,8 @@ def test_chunked_prefill_with_parallel_sampling(
170
decoding steps inside a chunked prefill forward pass
171
(where we have both prefill and decode together)
172
"""
173
+ _set_conv_state_layout(monkeypatch, conv_state_layout)
174
+
175
sampling_params = SamplingParams(n=3, temperature=1, seed=0, max_tokens=max_tokens)
176
with vllm_runner(
177
model,
@@ -168,17 +186,22 @@ def test_chunked_prefill_with_parallel_sampling(
186
187
@pytest.mark.parametrize("model", [SSM_MODELS[0], HYBRID_MODELS[0]])
188
@pytest.mark.parametrize("max_tokens", [20])
189
+@pytest.mark.parametrize("conv_state_layout", ["SD", "DS"])
190
def test_mamba_cache_cg_padding(
191
vllm_runner,
192
example_prompts,
193
+ monkeypatch,
194
model: str,
195
max_tokens: int,
196
+ conv_state_layout: str,
197
) -> None:
198
"""
199
This test is for verifying that mamba cache is padded to CG captured
200
batch size. If it's not, a torch RuntimeError will be raised because
201
tensor dimensions aren't compatible.
202
"""
203
+ _set_conv_state_layout(monkeypatch, conv_state_layout)
204
+
205
vllm_config = EngineArgs(model=model, trust_remote_code=True).create_engine_config()
206
cudagraph_dispatcher = CudagraphDispatcher(vllm_config)
207
cudagraph_dispatcher.initialize_cudagraph_keys(
vllm/envs.py
+8
@@ -191,6 +191,7 @@ if TYPE_CHECKING:
191
VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16
192
VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300
193
VLLM_KV_CACHE_LAYOUT: Literal["NHD", "HND"] | None = None
194
+ VLLM_SSM_CONV_STATE_LAYOUT: Literal["SD", "DS"] | None = None
195
VLLM_COMPUTE_NANS_IN_LOGITS: bool = False
196
VLLM_USE_NVFP4_CT_EMULATIONS: bool = False
197
VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: Literal[
@@ -1409,6 +1410,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
1410
"VLLM_KV_CACHE_LAYOUT": env_with_choices(
1411
"VLLM_KV_CACHE_LAYOUT", None, ["NHD", "HND"]
1412
),
1413
+ # SSM conv state layout used for Mamba models.
1414
+ # - SD: (state_len, dim) — dim contiguous (default)
1415
+ # - DS: (dim, state_len) — TP-sharded dim on dim1,
1416
+ # consistent with SSM temporal state and HND KV cache layout.
1417
+ "VLLM_SSM_CONV_STATE_LAYOUT": env_with_choices(
1418
+ "VLLM_SSM_CONV_STATE_LAYOUT", None, ["SD", "DS"]
1419
+ ),
1420
# Enable checking whether the generated logits contain NaNs,
1421
# indicating corrupted output. Useful for debugging low level bugs
1422
# or bad hardware but it may add compute overhead.
vllm/model_executor/layers/kda.py
+11
-5
@@ -31,7 +31,11 @@ from .linear import (
31
RowParallelLinear,
32
)
33
from .mamba.abstract import MambaBase
34
-from .mamba.mamba_utils import MambaStateDtypeCalculator, MambaStateShapeCalculator
34
+from .mamba.mamba_utils import (
35
+ MambaStateDtypeCalculator,
36
+ MambaStateShapeCalculator,
37
+ is_conv_state_dim_first,
38
+)
39
from .mamba.ops.causal_conv1d import causal_conv1d_fn, causal_conv1d_update
40
from .quantization.base_config import QuantizationConfig
41
@@ -315,10 +319,12 @@ class KimiDeltaAttention(nn.Module, MambaBase):
319
beta = beta[:num_actual_tokens]
320
321
(conv_state_q, conv_state_k, conv_state_v, recurrent_state) = constant_caches
318
- # deal with strides
319
- conv_state_q = conv_state_q.transpose(-1, -2)
320
- conv_state_k = conv_state_k.transpose(-1, -2)
321
- conv_state_v = conv_state_v.transpose(-1, -2)
322
+ # conv_state must be (..., dim, width-1) for the conv kernels.
323
+ # DS layout stores it that way directly; SD layout needs a transpose.
324
+ if not is_conv_state_dim_first():
325
+ conv_state_q = conv_state_q.transpose(-1, -2)
326
+ conv_state_k = conv_state_k.transpose(-1, -2)
327
+ conv_state_v = conv_state_v.transpose(-1, -2)
328
329
q_conv_weights = self.q_conv1d.weight.view(
330
self.q_conv1d.weight.size(0), self.q_conv1d.weight.size(2)
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+15
-2
@@ -41,6 +41,7 @@ from vllm.model_executor.layers.mamba.mamba_mixer2 import mamba_v2_sharded_weigh
41
from vllm.model_executor.layers.mamba.mamba_utils import (
42
MambaStateDtypeCalculator,
43
MambaStateShapeCalculator,
44
+ is_conv_state_dim_first,
45
)
46
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
47
causal_conv1d_fn,
@@ -699,7 +700,13 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
700
spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501
701
non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
702
self_kv_cache = self.kv_cache
702
- conv_state = self_kv_cache[0].transpose(-1, -2)
703
+ # conv_state must be (..., dim, width-1) for the conv kernels.
704
+ # DS layout stores it that way directly; SD layout needs a transpose.
705
+ conv_state = (
706
+ self_kv_cache[0]
707
+ if is_conv_state_dim_first()
708
+ else self_kv_cache[0].transpose(-1, -2)
709
+ )
710
ssm_state = self_kv_cache[1]
711
num_actual_tokens = attn_metadata.num_actual_tokens
712
num_accepted_tokens = attn_metadata.num_accepted_tokens
@@ -914,7 +921,13 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
921
"""
922
non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
923
self_kv_cache = self.kv_cache
917
- conv_state = self_kv_cache[0].transpose(-1, -2)
924
+ # conv_state must be (..., dim, width-1) for the conv kernels.
925
+ # DS layout stores it that way directly; SD layout needs a transpose.
926
+ conv_state = (
927
+ self_kv_cache[0]
928
+ if is_conv_state_dim_first()
929
+ else self_kv_cache[0].transpose(-1, -2)
930
+ )
931
ssm_state = self_kv_cache[1]
932
num_actual_tokens = attn_metadata.num_actual_tokens
933
vllm/model_executor/layers/mamba/mamba_mixer.py
+7
-3
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.mamba.abstract import MambaBase
24
from vllm.model_executor.layers.mamba.mamba_utils import (
25
MambaStateDtypeCalculator,
26
MambaStateShapeCalculator,
27
+ is_conv_state_dim_first,
28
)
29
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
30
causal_conv1d_fn,
@@ -267,9 +268,12 @@ class MambaMixer(MambaBase, PluggableLayer):
268
query_start_loc_p = attn_metadata.query_start_loc_p
269
state_indices_tensor_p = attn_metadata.state_indices_tensor_p
270
state_indices_tensor_d = attn_metadata.state_indices_tensor_d
270
- self_kv_cache = self.kv_cache
271
- conv_state = self_kv_cache[0].transpose(-1, -2)
272
- ssm_state = self_kv_cache[1]
271
+ conv_state = (
272
+ self.kv_cache[0]
273
+ if is_conv_state_dim_first()
274
+ else self.kv_cache[0].transpose(-1, -2)
275
+ )
276
+ ssm_state = self.kv_cache[1]
277
has_initial_states_p = attn_metadata.has_initial_states_p
278
cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p
279
last_chunk_indices_p = attn_metadata.last_chunk_indices_p
vllm/model_executor/layers/mamba/mamba_mixer2.py
+10
-4
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.mamba.abstract import MambaBase
24
from vllm.model_executor.layers.mamba.mamba_utils import (
25
MambaStateDtypeCalculator,
26
MambaStateShapeCalculator,
27
+ is_conv_state_dim_first,
28
)
29
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
30
causal_conv1d_fn,
@@ -575,10 +576,15 @@ class MambaMixer2(MambaBase, PluggableLayer):
576
assert isinstance(attn_metadata, dict)
577
attn_metadata = attn_metadata[self.prefix]
578
assert isinstance(attn_metadata, Mamba2AttentionMetadata)
578
- self_kv_cache = self.kv_cache
579
- # conv_state = (..., dim, width-1) yet contiguous along 'dim'
580
- conv_state = self_kv_cache[0].transpose(-1, -2)
581
- ssm_state = self_kv_cache[1]
579
+ # conv_state must be (..., dim, width-1) for the conv kernels.
580
+ # DS layout stores it that way directly; SD layout needs a
581
+ # transpose (which keeps dim contiguous via stride tricks).
582
+ conv_state = (
583
+ self.kv_cache[0]
584
+ if is_conv_state_dim_first()
585
+ else self.kv_cache[0].transpose(-1, -2)
586
+ )
587
+ ssm_state = self.kv_cache[1]
588
has_initial_states_p = attn_metadata.has_initial_states_p
589
prep_initial_states = attn_metadata.prep_initial_states
590
chunk_size = attn_metadata.chunk_size
vllm/model_executor/layers/mamba/mamba_utils.py
+73
-17
@@ -1,20 +1,52 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import functools
5
from collections.abc import Callable
6
from dataclasses import dataclass
6
-from typing import TypeAlias
7
+from typing import Literal, TypeAlias
8
9
import torch
10
11
+import vllm.envs as envs
12
from vllm.config.cache import MambaDType
13
from vllm.config.model import ModelDType
14
from vllm.distributed import divide
15
+from vllm.logger import init_logger
16
from vllm.utils.torch_utils import (
17
STR_DTYPE_TO_TORCH_DTYPE,
18
get_kv_cache_torch_dtype,
19
)
20
21
+logger = init_logger(__name__)
22
+
23
+ConvStateLayoutType = Literal["SD", "DS"]
24
+
25
+
26
+@functools.lru_cache
27
+def get_conv_state_layout() -> ConvStateLayoutType:
28
+ """Return the SSM conv state layout.
29
+
30
+ SD = (state_len, dim) — dim is the innermost contiguous dimension.
31
+ DS = (dim, state_len) — TP-sharded dim is on dim-1 (like HND for KV
32
+ cache), consistent with SSM temporal state layout.
33
+ """
34
+ layout: ConvStateLayoutType | None = envs.VLLM_SSM_CONV_STATE_LAYOUT
35
+ if layout is not None:
36
+ logger.info_once(
37
+ "VLLM_SSM_CONV_STATE_LAYOUT env detected. "
38
+ "Setting SSM conv state layout to %s.",
39
+ layout,
40
+ )
41
+ return layout
42
+
43
+ return "SD"
44
+
45
+
46
+def is_conv_state_dim_first() -> bool:
47
+ """True when the conv state is stored as (dim, state_len) per block."""
48
+ return get_conv_state_layout() == "DS"
49
+
50
51
class MambaStateDtypeCalculator:
52
@classmethod
@@ -107,6 +139,13 @@ class MambaStateShapeCalculator:
139
state_shape = (num_heads // tp_size, head_dim, head_dim)
140
return (state_shape,)
141
142
+ @staticmethod
143
+ def _orient_conv_shape(dim: int, state_len: int) -> tuple[int, int]:
144
+ """Return (dim, state_len) for DS layout, (state_len, dim) for SD."""
145
+ if is_conv_state_dim_first():
146
+ return (dim, state_len)
147
+ return (state_len, dim)
148
+
149
@classmethod
150
def mamba1_state_shape(
151
cls,
@@ -115,12 +154,11 @@ class MambaStateShapeCalculator:
154
state_size: int,
155
conv_kernel: int,
156
) -> tuple[tuple[int, int], tuple[int, int]]:
118
- conv_state_shape = (divide(intermediate_size, tp_world_size), conv_kernel - 1)
157
+ conv_dim = divide(intermediate_size, tp_world_size)
158
+ conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1)
159
160
temporal_state_shape = (divide(intermediate_size, tp_world_size), state_size)
161
122
- conv_state_shape = conv_state_shape[1], conv_state_shape[0]
123
-
162
return conv_state_shape, temporal_state_shape
163
164
@classmethod
@@ -141,8 +179,9 @@ class MambaStateShapeCalculator:
179
# heads and n_groups are TP-ed
180
conv_dim = intermediate_size + 2 * n_groups * state_size
181
144
- # contiguous along 'dim' axis
145
- conv_state_shape = (conv_kernel - 1 + num_spec, divide(conv_dim, tp_world_size))
182
+ conv_state_shape = cls._orient_conv_shape(
183
+ divide(conv_dim, tp_world_size), conv_kernel - 1 + num_spec
184
+ )
185
186
# These are not TP-ed as they depend on A, dt_bias, D
187
# - they are typically small
@@ -158,7 +197,7 @@ class MambaStateShapeCalculator:
197
conv_kernel: int,
198
) -> tuple[tuple[int, int]]:
199
conv_dim = divide(intermediate_size, tp_world_size)
161
- conv_state_shape = (conv_kernel - 1, conv_dim)
200
+ conv_state_shape = cls._orient_conv_shape(conv_dim, conv_kernel - 1)
201
return (conv_state_shape,)
202
203
@classmethod
@@ -185,13 +224,11 @@ class MambaStateShapeCalculator:
224
num_spec: int = 0,
225
):
226
conv_dim = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads
188
- conv_state_shape = (
227
+ conv_state_shape = cls._orient_conv_shape(
228
divide(conv_dim, tp_world_size),
229
conv_kernel_size - 1 + num_spec,
230
)
231
193
- conv_state_shape = conv_state_shape[1], conv_state_shape[0]
194
-
232
temporal_state_shape = (
233
divide(num_v_heads, tp_world_size),
234
head_v_dim,
@@ -218,12 +255,13 @@ class MambaStateShapeCalculator:
255
proj_size = num_heads * head_dim
256
proj_k_size = num_k_heads * head_k_dim
257
221
- conv_state_shape = (divide(proj_size, tp_world_size), conv_kernel_size - 1)
222
- conv_state_k_shape = (divide(proj_k_size, tp_world_size), conv_kernel_size - 1)
258
+ conv_state_shape = cls._orient_conv_shape(
259
+ divide(proj_size, tp_world_size), conv_kernel_size - 1
260
+ )
261
+ conv_state_k_shape = cls._orient_conv_shape(
262
+ divide(proj_k_size, tp_world_size), conv_kernel_size - 1
263
+ )
264
recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim)
224
-
225
- conv_state_shape = conv_state_shape[1], conv_state_shape[0]
226
- conv_state_k_shape = conv_state_k_shape[1], conv_state_k_shape[0]
265
return (
266
conv_state_shape,
267
conv_state_k_shape,
@@ -267,9 +305,27 @@ def get_conv_copy_spec(
305
cur_block_idx: int,
306
num_accepted_tokens: int,
307
) -> MambaCopySpec:
270
- """Return a MambaCopySpec for copying a convolutional state slice."""
308
+ """Return a MambaCopySpec for copying a convolutional state slice.
309
+
310
+ Works for both SD layout ``(num_blocks, state_len, dim)`` and
311
+ DS layout ``(num_blocks, dim, state_len)``.
312
+ """
313
src_block_id = block_ids[cur_block_idx]
272
- src_state = state[src_block_id, num_accepted_tokens - 1 :]
314
+ offset = num_accepted_tokens - 1
315
+ if is_conv_state_dim_first():
316
+ # DS layout: (num_blocks, dim, state_len) — state_len is last.
317
+ if offset > 0:
318
+ # Slicing along the last dim yields a non-contiguous view
319
+ # because features (dim) are strided by state_len.
320
+ raise NotImplementedError(
321
+ "DS conv state layout does not yet support speculative "
322
+ "decoding with mamba_cache_mode='align' "
323
+ "(num_accepted_tokens > 1)."
324
+ )
325
+ src_state = state[src_block_id]
326
+ else:
327
+ # SD layout: (num_blocks, state_len, dim) — dim contiguous.
328
+ src_state = state[src_block_id, offset:]
329
return MambaCopySpec(
330
start_addr=src_state.data_ptr(), num_elements=src_state.numel()
331
)
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
-4
@@ -592,7 +592,6 @@ def causal_conv1d_fn(
592
stride_istate_seq = conv_states.stride(0)
593
stride_istate_dim = conv_states.stride(1)
594
stride_istate_token = conv_states.stride(2)
595
- assert stride_istate_dim == 1
595
if out.dim() == 2:
596
stride_o_dim = out.stride(0)
597
stride_o_token = out.stride(1)
@@ -1149,9 +1148,6 @@ def causal_conv1d_update(
1148
1149
if validate_data:
1150
assert dim == weight.size(0)
1152
- assert conv_state.stride(-2) == 1, (
1153
- f"ERROR: expect contiguous along feat-dim of conv_state (currently stride={conv_state.stride()})"
1154
- )
1151
assert state_len >= width - 1
1152
# when above happens, we don't shift-left to keep any records in conv_state
1153
assert dim == conv_state.size(1)
vllm/model_executor/layers/mamba/short_conv.py
+6
-2
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.mamba.abstract import MambaBase
17
from vllm.model_executor.layers.mamba.mamba_utils import (
18
MambaStateDtypeCalculator,
19
MambaStateShapeCalculator,
20
+ is_conv_state_dim_first,
21
)
22
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
23
causal_conv1d_fn,
@@ -117,8 +118,11 @@ class ShortConv(MambaBase, CustomOp):
118
assert isinstance(attn_metadata, dict)
119
attn_metadata = attn_metadata[self.prefix]
120
assert isinstance(attn_metadata, ShortConvAttentionMetadata)
120
- self_kv_cache = self.kv_cache
121
- conv_state = self_kv_cache[0].transpose(-1, -2)
121
+ conv_state = (
122
+ self.kv_cache[0]
123
+ if is_conv_state_dim_first()
124
+ else self.kv_cache[0].transpose(-1, -2)
125
+ )
126
state_indices_tensor_p = attn_metadata.state_indices_tensor_p
127
state_indices_tensor_d = attn_metadata.state_indices_tensor_d
128
has_initial_states_p = attn_metadata.has_initial_states_p
vllm/model_executor/models/olmo_hybrid.py
+8
-1
@@ -68,6 +68,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
68
MambaStateCopyFuncCalculator,
69
MambaStateDtypeCalculator,
70
MambaStateShapeCalculator,
71
+ is_conv_state_dim_first,
72
)
73
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
74
causal_conv1d_fn,
@@ -429,7 +430,13 @@ class OlmoHybridGatedDeltaNet(nn.Module, MambaBase):
430
spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor
431
non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor
432
self_kv_cache = self.kv_cache
432
- conv_state = self_kv_cache[0].transpose(-1, -2)
433
+ # conv_state must be (..., dim, width-1) for the conv kernels.
434
+ # DS layout stores it that way directly; SD layout needs a transpose.
435
+ conv_state = (
436
+ self_kv_cache[0]
437
+ if is_conv_state_dim_first()
438
+ else self_kv_cache[0].transpose(-1, -2)
439
+ )
440
ssm_state = self_kv_cache[1]
441
num_actual_tokens = attn_metadata.num_actual_tokens
442
num_accepted_tokens = attn_metadata.num_accepted_tokens
vllm/model_executor/models/plamo2.py
+8
-1
@@ -32,6 +32,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
32
MambaStateCopyFuncCalculator,
33
MambaStateDtypeCalculator,
34
MambaStateShapeCalculator,
35
+ is_conv_state_dim_first,
36
)
37
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
38
causal_conv1d_fn,
@@ -266,7 +267,13 @@ class Plamo2MambaMixer(MambaBase, PluggableLayer):
267
assert isinstance(attn_metadata, Mamba2AttentionMetadata)
268
self_kv_cache = self.kv_cache
269
# conv_state = (..., dim, width-1) yet contiguous along 'dim'
269
- conv_state = self_kv_cache[0].transpose(-1, -2)
270
+ # conv_state must be (..., dim, width-1) for the conv kernels.
271
+ # DS layout stores it that way directly; SD layout needs a transpose.
272
+ conv_state = (
273
+ self_kv_cache[0]
274
+ if is_conv_state_dim_first()
275
+ else self_kv_cache[0].transpose(-1, -2)
276
+ )
277
ssm_state = self_kv_cache[1]
278
state_indices_tensor_p = attn_metadata.state_indices_tensor_p
279
state_indices_tensor_d = attn_metadata.state_indices_tensor_d