[CPU][GDN] Support NIXL DS convolution-state layout (#53300)
Signed-off-by: Li, Tianmu <tianmu.li@intel.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Tianmu Li committed
Sep 23, 2026 at 23:55 UTC
67ddac3ea8d5b36978187f92908f555ff71190e4
4 files changed
+383
-21
tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
+186
@@ -473,6 +473,20 @@ def _conv_inputs(total_tokens: int):
473
return x, weight, bias
474
475
476
+def _conv_states(
477
+ num_slots: int,
478
+ state_len: int,
479
+ dim: int = CONV_DIM,
480
+ layout: str = "SD",
481
+) -> torch.Tensor:
482
+ if layout == "SD":
483
+ return torch.zeros(num_slots, state_len, dim, dtype=torch.bfloat16).transpose(
484
+ 1, 2
485
+ )
486
+ assert layout == "DS"
487
+ return torch.zeros(num_slots, dim, state_len, dtype=torch.bfloat16)
488
+
489
+
490
def _sd_conv_states(
491
num_slots: int, state_len: int, dim: int = CONV_DIM
492
) -> torch.Tensor:
@@ -484,6 +498,39 @@ def _maybe_pack_conv_weight(weight: torch.Tensor, is_vnni: bool) -> torch.Tensor
498
return ops.causal_conv1d_weight_pack(weight) if is_vnni else weight
499
500
501
+@torch.inference_mode()
502
+def test_causal_conv1d_update_cpu_preserves_sd_state(
503
+ monkeypatch: pytest.MonkeyPatch,
504
+) -> None:
505
+ """SD state and indices reach the native operator unchanged."""
506
+ num_slots = 2
507
+ state_len = CONV_KERNEL - 1
508
+ conv_state = _sd_conv_states(num_slots, state_len)
509
+ state_indices = torch.tensor([1], dtype=torch.int32)
510
+ x = torch.zeros(1, CONV_DIM, dtype=torch.bfloat16)
511
+ recorded = {}
512
+
513
+ def record_update(**kwargs):
514
+ recorded.update(kwargs)
515
+ return x
516
+
517
+ monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False)
518
+ monkeypatch.setattr(gdn_attention.ops, "causal_conv1d_update_cpu", record_update)
519
+ out = gdn_attention._causal_conv1d_update_cpu(
520
+ x=x,
521
+ conv_states=conv_state,
522
+ weight=torch.zeros(CONV_DIM, CONV_KERNEL, dtype=torch.bfloat16),
523
+ bias=None,
524
+ silu_activation=False,
525
+ conv_state_indices=state_indices,
526
+ is_vnni=False,
527
+ )
528
+
529
+ assert out is x
530
+ assert recorded["conv_states"] is conv_state
531
+ assert recorded["conv_state_indices"] is state_indices
532
+
533
+
534
@torch.inference_mode()
535
def test_spec_aware_mixed_routing_preserves_token_order(
536
monkeypatch: pytest.MonkeyPatch,
@@ -1123,6 +1170,145 @@ def _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False):
1170
return out.transpose(0, 1).contiguous(), conv_state
1171
1172
1173
+def _python_native_fwd(**kwargs):
1174
+ from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
1175
+ causal_conv1d_fn_cpu,
1176
+ )
1177
+
1178
+ return causal_conv1d_fn_cpu(
1179
+ x=kwargs["x"],
1180
+ weight=kwargs["weight"],
1181
+ bias=kwargs["bias"],
1182
+ conv_states=kwargs["conv_states"],
1183
+ query_start_loc=kwargs["query_start_loc"],
1184
+ cache_indices=kwargs["cache_indices"],
1185
+ has_initial_state=kwargs["has_initial_state"],
1186
+ activation=kwargs["silu_activation"],
1187
+ )
1188
+
1189
+
1190
+@torch.inference_mode()
1191
+def test_causal_conv1d_fwd_cpu_ds_adapter_varlen_mixed_initial_state(
1192
+ monkeypatch: pytest.MonkeyPatch,
1193
+) -> None:
1194
+ dim = 8
1195
+ seq_lens = [2, 3]
1196
+ state_len = _STATE_LEN
1197
+ state_indices = torch.tensor([2, 0], dtype=torch.int32)
1198
+ query_start_loc = torch.tensor([0, 2, 5], dtype=torch.int32)
1199
+ has_initial_state = torch.tensor([False, True])
1200
+ x_tokens, weight, bias = _conv_inputs(sum(seq_lens))
1201
+ x_tokens = x_tokens[:, :dim]
1202
+ weight = weight[:dim]
1203
+ bias = bias[:dim]
1204
+ x = x_tokens.transpose(0, 1).contiguous()
1205
+ logical_state = (
1206
+ torch.arange(3 * dim * state_len, dtype=torch.float32)
1207
+ .view(3, dim, state_len)
1208
+ .to(torch.bfloat16)
1209
+ )
1210
+ state_sd = _conv_states(3, state_len, dim, "SD")
1211
+ state_ds = _conv_states(3, state_len, dim, "DS")
1212
+ state_sd.copy_(logical_state)
1213
+ state_ds.copy_(logical_state)
1214
+ recorded = []
1215
+
1216
+ def record_fwd(**kwargs):
1217
+ recorded.append(
1218
+ (
1219
+ kwargs["conv_states"].shape,
1220
+ kwargs["conv_states"].stride(),
1221
+ kwargs["cache_indices"].clone(),
1222
+ )
1223
+ )
1224
+ return _python_native_fwd(**kwargs)
1225
+
1226
+ monkeypatch.setattr(gdn_attention.ops, "causal_conv1d_fwd_cpu", record_fwd)
1227
+ monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False)
1228
+ out_sd = gdn_attention._causal_conv1d_fwd_cpu(
1229
+ x=x,
1230
+ weight=weight,
1231
+ bias=bias,
1232
+ conv_states=state_sd,
1233
+ query_start_loc=query_start_loc,
1234
+ cache_indices=state_indices,
1235
+ has_initial_state=has_initial_state,
1236
+ silu_activation=True,
1237
+ is_vnni=False,
1238
+ )
1239
+ monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True)
1240
+ out_ds = gdn_attention._causal_conv1d_fwd_cpu(
1241
+ x=x,
1242
+ weight=weight,
1243
+ bias=bias,
1244
+ conv_states=state_ds,
1245
+ query_start_loc=query_start_loc,
1246
+ cache_indices=state_indices,
1247
+ has_initial_state=has_initial_state,
1248
+ silu_activation=True,
1249
+ is_vnni=False,
1250
+ )
1251
+
1252
+ torch.testing.assert_close(out_ds, out_sd)
1253
+ torch.testing.assert_close(state_ds, state_sd)
1254
+ assert recorded[0][2].equal(state_indices)
1255
+ assert recorded[1][0] == (2, dim, state_len)
1256
+ assert recorded[1][1] == (dim * state_len, 1, dim)
1257
+ assert recorded[1][2].equal(torch.arange(2, dtype=torch.int32))
1258
+
1259
+
1260
+@torch.inference_mode()
1261
+def test_causal_conv1d_update_cpu_ds_adapter_wide_multi_token_copyback(
1262
+ monkeypatch: pytest.MonkeyPatch,
1263
+) -> None:
1264
+ dim = 8
1265
+ seq_len = 4
1266
+ state_len = seq_len + 5
1267
+ state = _conv_states(4, state_len, dim, "DS")
1268
+ state.copy_(
1269
+ torch.arange(state.numel(), dtype=torch.float32).view_as(state).to(state.dtype)
1270
+ )
1271
+ original_state = state.clone()
1272
+ state_indices = torch.tensor([3, 1, 0], dtype=torch.int32)
1273
+ num_accepted_tokens = torch.tensor([1, 4, 2], dtype=torch.int32)
1274
+ mutated_scratch = torch.arange(
1275
+ state_indices.numel() * dim * state_len, dtype=torch.float32
1276
+ ).view(state_indices.numel(), dim, state_len)
1277
+ mutated_scratch = (mutated_scratch + 200).to(state.dtype)
1278
+ recorded = {}
1279
+
1280
+ def mutate_update(**kwargs):
1281
+ recorded.update(kwargs)
1282
+ kwargs["conv_states"].copy_(mutated_scratch)
1283
+ return torch.zeros_like(kwargs["x"])
1284
+
1285
+ monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True)
1286
+ monkeypatch.setattr(gdn_attention.ops, "causal_conv1d_update_cpu", mutate_update)
1287
+ gdn_attention._causal_conv1d_update_cpu(
1288
+ x=torch.zeros(state_indices.numel(), seq_len, dim, dtype=state.dtype),
1289
+ conv_states=state,
1290
+ weight=torch.zeros(dim, CONV_KERNEL, dtype=state.dtype),
1291
+ bias=None,
1292
+ silu_activation=False,
1293
+ conv_state_indices=state_indices,
1294
+ is_vnni=False,
1295
+ num_accepted_tokens=num_accepted_tokens,
1296
+ )
1297
+
1298
+ expected_state = original_state.clone()
1299
+ expected_state.index_copy_(0, state_indices.to(torch.long), mutated_scratch)
1300
+ torch.testing.assert_close(state, expected_state)
1301
+ torch.testing.assert_close(
1302
+ state.index_select(0, state_indices)[:, :, _STATE_LEN:],
1303
+ mutated_scratch[:, :, _STATE_LEN:],
1304
+ )
1305
+ torch.testing.assert_close(
1306
+ recorded["conv_state_indices"],
1307
+ torch.arange(state_indices.numel(), dtype=torch.int32),
1308
+ )
1309
+ torch.testing.assert_close(recorded["num_accepted_tokens"], num_accepted_tokens)
1310
+
1311
+
1312
@pytest.mark.skipif(
1313
not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
1314
)
tests/platforms/test_cpu.py
+74
-1
@@ -1,12 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import os
5
from types import SimpleNamespace
6
7
import pytest
8
import torch
9
9
-from vllm.config import CacheConfig, VllmConfig
10
+from vllm.config import CacheConfig, KVTransferConfig, VllmConfig
11
from vllm.platforms.cpu import CpuPlatform
12
13
@@ -44,6 +45,7 @@ def _cpu_config(
45
architecture=architecture or model_type,
46
has_inner_state=True,
47
use_mla=False,
48
+ hf_config=SimpleNamespace(model_type=model_type),
49
hf_text_config=SimpleNamespace(
50
model_type=model_type,
51
layer_types=layer_types,
@@ -150,3 +152,74 @@ def test_cpu_accelerated_gdn_dtype_policy(
152
153
CpuPlatform.check_and_update_config(config)
154
assert cache_config.mamba_ssm_cache_dtype == expected_dtype
155
+
156
+
157
+@pytest.mark.parametrize(
158
+ (
159
+ "connector",
160
+ "extra_config",
161
+ "explicit_layout",
162
+ "avx512_bf16_supported",
163
+ "expected_layout",
164
+ ),
165
+ [
166
+ pytest.param("NixlConnector", None, None, True, "DS", id="nixl"),
167
+ pytest.param("NixlPullConnector", None, None, True, "DS", id="nixl-pull"),
168
+ pytest.param("NixlPushConnector", None, None, True, "DS", id="nixl-push"),
169
+ pytest.param(
170
+ "MultiConnector",
171
+ {
172
+ "connectors": [
173
+ {
174
+ "kv_connector": "ExampleConnector",
175
+ "kv_connector_extra_config": {},
176
+ },
177
+ {
178
+ "kv_connector": "NixlConnector",
179
+ "kv_connector_extra_config": {},
180
+ },
181
+ ]
182
+ },
183
+ None,
184
+ True,
185
+ "DS",
186
+ id="multi-nixl",
187
+ ),
188
+ pytest.param("NixlConnector", None, "SD", True, "SD", id="explicit-override"),
189
+ pytest.param("OffloadingConnector", None, None, True, "SD", id="non-nixl"),
190
+ pytest.param(
191
+ "OffloadingConnector", None, None, False, None, id="non-nixl-no-avx"
192
+ ),
193
+ ],
194
+)
195
+def test_cpu_conv_state_layout_selection(
196
+ monkeypatch: pytest.MonkeyPatch,
197
+ connector: str,
198
+ extra_config: dict | None,
199
+ explicit_layout: str | None,
200
+ avx512_bf16_supported: bool,
201
+ expected_layout: str | None,
202
+) -> None:
203
+ layout_env = "VLLM_SSM_CONV_STATE_LAYOUT"
204
+ monkeypatch.setattr(
205
+ "torch.cpu._is_avx512_bf16_supported",
206
+ lambda: avx512_bf16_supported,
207
+ )
208
+
209
+ kv_transfer_config = KVTransferConfig(
210
+ kv_connector=connector,
211
+ kv_connector_extra_config=extra_config or {},
212
+ kv_role="kv_both",
213
+ )
214
+ config = _cpu_config(
215
+ CacheConfig(mamba_ssm_cache_dtype="float32"),
216
+ model_type="qwen3_5",
217
+ resolved_dtype="float32",
218
+ )
219
+ config.kv_transfer_config = kv_transfer_config
220
+ monkeypatch.delenv(layout_env, raising=False)
221
+ if explicit_layout is not None:
222
+ monkeypatch.setenv(layout_env, explicit_layout)
223
+
224
+ CpuPlatform.check_and_update_config(config)
225
+ assert os.environ.get(layout_env) == expected_layout
vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py
+106
-15
@@ -27,6 +27,104 @@ from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
27
_CPU_GDN_ATTENTION_OPS_REGISTERED = False
28
29
30
+def _stage_ds_conv_states(
31
+ conv_states: torch.Tensor, state_indices: torch.Tensor
32
+) -> tuple[torch.Tensor, torch.Tensor]:
33
+ storage = conv_states.transpose(1, 2).index_select(0, state_indices)
34
+ scratch = storage.transpose(1, 2)
35
+ identity_indices = torch.arange(
36
+ state_indices.numel(), dtype=torch.int32, device=state_indices.device
37
+ )
38
+ return scratch, identity_indices
39
+
40
+
41
+def _scatter_ds_conv_states(
42
+ conv_states: torch.Tensor,
43
+ state_indices: torch.Tensor,
44
+ scratch: torch.Tensor,
45
+) -> None:
46
+ conv_states.index_copy_(0, state_indices.to(torch.int64), scratch)
47
+
48
+
49
+def _causal_conv1d_fwd_cpu(
50
+ *,
51
+ x: torch.Tensor,
52
+ weight: torch.Tensor,
53
+ bias: torch.Tensor | None,
54
+ conv_states: torch.Tensor,
55
+ query_start_loc: torch.Tensor,
56
+ cache_indices: torch.Tensor,
57
+ has_initial_state: torch.Tensor,
58
+ silu_activation: bool,
59
+ is_vnni: bool,
60
+) -> torch.Tensor:
61
+ if not is_conv_state_dim_first():
62
+ return ops.causal_conv1d_fwd_cpu(
63
+ x=x,
64
+ weight=weight,
65
+ bias=bias,
66
+ conv_states=conv_states,
67
+ query_start_loc=query_start_loc,
68
+ cache_indices=cache_indices,
69
+ has_initial_state=has_initial_state,
70
+ silu_activation=silu_activation,
71
+ is_vnni=is_vnni,
72
+ )
73
+
74
+ scratch, identity_indices = _stage_ds_conv_states(conv_states, cache_indices)
75
+ out = ops.causal_conv1d_fwd_cpu(
76
+ x=x,
77
+ weight=weight,
78
+ bias=bias,
79
+ conv_states=scratch,
80
+ query_start_loc=query_start_loc,
81
+ cache_indices=identity_indices,
82
+ has_initial_state=has_initial_state,
83
+ silu_activation=silu_activation,
84
+ is_vnni=is_vnni,
85
+ )
86
+ _scatter_ds_conv_states(conv_states, cache_indices, scratch)
87
+ return out
88
+
89
+
90
+def _causal_conv1d_update_cpu(
91
+ *,
92
+ x: torch.Tensor,
93
+ conv_states: torch.Tensor,
94
+ weight: torch.Tensor,
95
+ bias: torch.Tensor | None,
96
+ silu_activation: bool,
97
+ conv_state_indices: torch.Tensor,
98
+ is_vnni: bool,
99
+ num_accepted_tokens: torch.Tensor | None = None,
100
+) -> torch.Tensor:
101
+ if not is_conv_state_dim_first():
102
+ return ops.causal_conv1d_update_cpu(
103
+ x=x,
104
+ conv_states=conv_states,
105
+ weight=weight,
106
+ bias=bias,
107
+ silu_activation=silu_activation,
108
+ conv_state_indices=conv_state_indices,
109
+ is_vnni=is_vnni,
110
+ num_accepted_tokens=num_accepted_tokens,
111
+ )
112
+
113
+ scratch, identity_indices = _stage_ds_conv_states(conv_states, conv_state_indices)
114
+ out = ops.causal_conv1d_update_cpu(
115
+ x=x,
116
+ conv_states=scratch,
117
+ weight=weight,
118
+ bias=bias,
119
+ silu_activation=silu_activation,
120
+ conv_state_indices=identity_indices,
121
+ is_vnni=is_vnni,
122
+ num_accepted_tokens=num_accepted_tokens,
123
+ )
124
+ _scatter_ds_conv_states(conv_states, conv_state_indices, scratch)
125
+ return out
126
+
127
+
128
def cpu_gdn_attention_core(
129
mixed_qkv: torch.Tensor,
130
b: torch.Tensor,
@@ -109,10 +207,8 @@ def _cpu_gdn_attention_nonspec(
207
208
conv_state = layer.kv_cache[0]
209
if use_cpp_conv:
112
- # C++ conv requires [num_allocated_slots, kernel - 1, conv_dim] (SD).
113
- if is_conv_state_dim_first():
114
- raise RuntimeError("C++ CPU GDN attention requires `SD` conv_state layout.")
115
- conv_state = conv_state.transpose(1, 2)
210
+ if not is_conv_state_dim_first():
211
+ conv_state = conv_state.transpose(1, 2)
212
else:
213
if not is_conv_state_dim_first():
214
conv_state = conv_state.transpose(-1, -2)
@@ -141,7 +237,7 @@ def _cpu_gdn_attention_nonspec(
237
decode_a = a[:num_decode_tokens]
238
decode_state_indices = state_indices_tensor[:num_decodes]
239
if use_cpp_conv:
144
- decode_mixed_qkv = ops.causal_conv1d_update_cpu(
240
+ decode_mixed_qkv = _causal_conv1d_update_cpu(
241
x=decode_mixed_qkv,
242
conv_states=conv_state,
243
weight=layer.conv1d.weight,
@@ -210,7 +306,7 @@ def _cpu_gdn_attention_nonspec(
306
]
307
308
if use_cpp_conv:
213
- prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu(
309
+ prefill_mixed_qkv = _causal_conv1d_fwd_cpu(
310
x=prefill_mixed_qkv.transpose(0, 1),
311
weight=layer.conv1d.weight,
312
bias=layer.conv1d.bias,
@@ -416,7 +512,6 @@ def _spec_forward(
512
513
can_use_native_conv = (
514
torch.cpu._is_amx_tile_supported()
419
- and not is_conv_state_dim_first()
515
and width == 4
516
and num_spec_decodes > 0
517
and bool(torch.all(seq_lens == seq_lens[0]).item())
@@ -424,7 +519,7 @@ def _spec_forward(
519
)
520
if can_use_native_conv:
521
q_i = int(seq_lens[0].item())
427
- conv_out = ops.causal_conv1d_update_cpu(
522
+ conv_out = _causal_conv1d_update_cpu(
523
x=mixed_qkv_spec.view(num_spec_decodes, q_i, dim),
524
conv_states=conv_buf,
525
weight=layer.conv1d.weight,
@@ -509,8 +604,6 @@ def _spec_aware_nonspec(
604
state_indices_tensor = state_indices_tensor.contiguous()
605
606
is_amx = torch.cpu._is_amx_tile_supported()
512
- if is_amx and is_conv_state_dim_first():
513
- raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")
607
608
if not is_amx:
609
conv_weights = _unpacked_conv_weight(layer)
@@ -526,7 +619,7 @@ def _spec_aware_nonspec(
619
decode_a = a[:num_decode_tokens]
620
decode_state_indices = state_indices_tensor[:num_decodes]
621
if is_amx:
529
- decode_mixed_qkv = ops.causal_conv1d_update_cpu(
622
+ decode_mixed_qkv = _causal_conv1d_update_cpu(
623
x=decode_mixed_qkv,
624
conv_states=conv_buf,
625
weight=layer.conv1d.weight,
@@ -593,7 +686,7 @@ def _spec_aware_nonspec(
686
num_decodes : num_decodes + num_prefills
687
]
688
if is_amx:
596
- prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu(
689
+ prefill_mixed_qkv = _causal_conv1d_fwd_cpu(
690
x=prefill_mixed_qkv.transpose(0, 1),
691
weight=layer.conv1d.weight,
692
bias=layer.conv1d.bias,
@@ -664,11 +757,9 @@ def _spec_aware_nonspec_subset(
757
prefill_state_indices = prefill_state_indices.contiguous()
758
759
is_amx = torch.cpu._is_amx_tile_supported()
667
- if is_amx and is_conv_state_dim_first():
668
- raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")
760
761
if is_amx:
671
- conv_out = ops.causal_conv1d_fwd_cpu(
762
+ conv_out = _causal_conv1d_fwd_cpu(
763
x=mixed_qkv.transpose(0, 1),
764
weight=layer.conv1d.weight,
765
bias=layer.conv1d.bias,
vllm/platforms/cpu.py
+17
-5
@@ -395,11 +395,23 @@ class CpuPlatform(Platform):
395
# Avoid inductor generates num_thread() and breaks the thread binding
396
os.environ["TORCHINDUCTOR_CPP_DYNAMIC_THREADS"] = "1"
397
398
- # For efficient conv state memory access. The C++ causal_conv1d
399
- # kernels (VDPBF16PS, no AMX tiles) consume the SD layout on any
400
- # AVX-512BF16 CPU, so apply it beyond AMX (e.g. AMD Zen5/Turin).
401
- if torch.cpu._is_avx512_bf16_supported():
402
- os.environ["VLLM_SSM_CONV_STATE_LAYOUT"] = "SD"
398
+ # NIXL's Mamba descriptors require DS conv state storage. Select it
399
+ # before cache shapes are created, while preserving an explicit layout.
400
+ conv_state_layout_env = "VLLM_SSM_CONV_STATE_LAYOUT"
401
+ if conv_state_layout_env not in os.environ:
402
+ kv_transfer_config = vllm_config.kv_transfer_config
403
+ uses_nixl = kv_transfer_config is not None and any(
404
+ kv_transfer_config.has_connector(name)
405
+ for name in (
406
+ "NixlConnector",
407
+ "NixlPullConnector",
408
+ "NixlPushConnector",
409
+ )
410
+ )
411
+ if uses_nixl:
412
+ os.environ[conv_state_layout_env] = "DS"
413
+ elif torch.cpu._is_avx512_bf16_supported():
414
+ os.environ[conv_state_layout_env] = "SD"
415
416
ld_preload_str = os.getenv("LD_PRELOAD", "")
417
cpu_architecture = Platform.get_cpu_architecture()