@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 525b7bbb3a

[Bugfix][CPU] Enable C++ causal_conv1d GDN path and float32 SSM cache on non-AMX AVX-512BF16 CPUs (#49688)

Signed-off-by: Dinesh Chitlangia <dineshchitlangia@gmail.com>

Dinesh Chitlangia committed Aug 19, 2026 at 11:50 UTC 525b7bbb3a74d811cedcee108425d9bb5c215fdd

4 files changed +290 -17

tests/kernels/mamba/cpu/test_cpu_gdn_ops.py

+270 -4

@@ -878,17 +878,18 @@ def test_causal_conv1d_update_cpu_rejects_invalid_accepted_count(

878

879

880 @pytest.mark.skipif(

881 - not torch.cpu._is_amx_tile_supported(),

882 - reason="requires AMX support",

881 + not torch.cpu._is_avx512_bf16_supported(),

882 + reason="causal_conv1d_fwd_cpu requires AVX-512BF16 (Intel Xeon or AMD EPYC)",

883 )

884 @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS)

885 @torch.inference_mode()

886 def test_causal_conv1d_fwd_cpu_two_call_split(total_tokens: int, split: int) -> None:

887 - """AMX prefill conv op must honor ``has_initial_state`` so a two-call split

887 + """C++ prefill conv op must honor ``has_initial_state`` so a two-call split

888 matches the single-call result.

889

890 Regression test for ``causal_conv1d_fwd_varlen_kernel_impl`` (``conv.cpp``)

891 - ignoring the carried conv state on continued chunks.

891 + ignoring the carried conv state on continued chunks. Runs on any

892 + AVX-512BF16 CPU since conv.cpp uses VDPBF16PS, not AMX tiles.

893 """

894 state_len = CONV_KERNEL - 1

895 x, weight, bias = _conv_inputs(total_tokens)

@@ -996,3 +997,268 @@ def test_batch_memcpy_cpu_fallback() -> None:

997

998 for src, dst in zip(srcs, dsts):

999 torch.testing.assert_close(dst, src)

1000 +

1001 +

1002 +# ---------------------------------------------------------------------------

1003 +# C++ conv (conv.cpp) uses VDPBF16PS, not AMX tiles, so it runs on any

1004 +# AVX-512BF16 CPU; weight is VNNI-packed on this same predicate at load time.

1005 +# ---------------------------------------------------------------------------

1006 +

1007 +_HAS_AVX512_BF16 = torch.cpu._is_avx512_bf16_supported()

1008 +

1009 +_STATE_LEN = CONV_KERNEL - 1

1010 +

1011 +CONV_EQUIV_SEQ_LENS = [[1], [7], [64], [65], [1, 2, 3], [63, 64, 65], [128, 129]]

1012 +

1013 +

1014 +def _conv_fp32_oracle(x, weight, bias, seq_lens, activation="silu"):

1015 + """High-precision conv reference: everything in float32, no bf16 rounding.

1016 + x: [total_tokens, dim] (no initial state). Returns [total_tokens, dim]."""

1017 + xf = x.float()

1018 + wf = weight.float().unsqueeze(1)

1019 + bf = bias.float()

1020 + out = torch.empty_like(xf)

1021 + start = 0

1022 + for n in seq_lens:

1023 + seg = xf[start : start + n].transpose(0, 1).unsqueeze(0) # [1, dim, n]

1024 + conv_in = F.pad(seg, (_STATE_LEN, 0))

1025 + seg_out = F.conv1d(conv_in, wf, bf, padding=0, groups=CONV_DIM)[..., -n:]

1026 + if activation in ("silu", "swish"):

1027 + seg_out = F.silu(seg_out)

1028 + out[start : start + n] = seg_out.squeeze(0).transpose(0, 1)

1029 + start += n

1030 + return out

1031 +

1032 +

1033 +def _run_prefill_torch(x, weight, bias, seq_lens):

1034 + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (

1035 + causal_conv1d_fn_cpu as causal_conv1d_torch,

1036 + )

1037 +

1038 + num_seqs = len(seq_lens)

1039 + conv_states = torch.zeros(num_seqs, CONV_DIM, _STATE_LEN, dtype=x.dtype)

1040 + qsl = torch.tensor(

1041 + [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32

1042 + )

1043 + out = causal_conv1d_torch(

1044 + x=x.transpose(0, 1).contiguous(),

1045 + weight=weight,

1046 + bias=bias,

1047 + conv_states=conv_states,

1048 + query_start_loc=qsl,

1049 + cache_indices=torch.arange(num_seqs, dtype=torch.int32),

1050 + has_initial_state=torch.zeros(num_seqs, dtype=torch.bool),

1051 + activation="silu",

1052 + )

1053 + return out.transpose(0, 1).contiguous(), conv_states

1054 +

1055 +

1056 +def _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False):

1057 + num_seqs = len(seq_lens)

1058 + packed_w = ops.causal_conv1d_weight_pack(weight) if is_vnni else weight

1059 + if is_vnni:

1060 + # C++-branch layout: kv-cache "SD" [slots, state_len, dim] transposed to

1061 + # [slots, dim, state_len] (a non-contiguous view).

1062 + conv_state = torch.zeros(

1063 + num_seqs, _STATE_LEN, CONV_DIM, dtype=x.dtype

1064 + ).transpose(1, 2)

1065 + else:

1066 + conv_state = torch.zeros(num_seqs, CONV_DIM, _STATE_LEN, dtype=x.dtype)

1067 + qsl = torch.tensor(

1068 + [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32

1069 + )

1070 + out = ops.causal_conv1d_fwd_cpu(

1071 + x=x.transpose(0, 1),

1072 + weight=packed_w,

1073 + bias=bias,

1074 + conv_states=conv_state,

1075 + query_start_loc=qsl,

1076 + cache_indices=torch.arange(num_seqs, dtype=torch.int32),

1077 + has_initial_state=torch.zeros(num_seqs, dtype=torch.bool),

1078 + silu_activation=True,

1079 + is_vnni=is_vnni,

1080 + )

1081 + return out.transpose(0, 1).contiguous(), conv_state

1082 +

1083 +

1084 +@pytest.mark.skipif(

1085 + not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"

1086 +)

1087 +@pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS)

1088 +@torch.inference_mode()

1089 +def test_conv_cpp_matches_torch(seq_lens):

1090 + """C++ causal_conv1d_fwd_cpu matches the torch fallback within bf16 tol."""

1091 + x, weight, bias = _conv_inputs(sum(seq_lens))

1092 + out_torch, state_torch = _run_prefill_torch(x, weight, bias, seq_lens)

1093 + out_cpp, state_cpp = _run_prefill_cpp(x, weight, bias, seq_lens)

1094 + torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2)

1095 + torch.testing.assert_close(state_cpp, state_torch, atol=1e-2, rtol=1e-2)

1096 +

1097 +

1098 +@pytest.mark.skipif(

1099 + not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"

1100 +)

1101 +@pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS)

1102 +@torch.inference_mode()

1103 +def test_conv_cpp_no_worse_than_torch_vs_fp32(seq_lens):

1104 + """Swapping torch -> C++ conv must not increase error vs an fp32 oracle."""

1105 + x, weight, bias = _conv_inputs(sum(seq_lens))

1106 + oracle = _conv_fp32_oracle(x, weight, bias, seq_lens)

1107 + out_torch, _ = _run_prefill_torch(x, weight, bias, seq_lens)

1108 + out_cpp, _ = _run_prefill_cpp(x, weight, bias, seq_lens)

1109 + err_torch = (out_torch.float() - oracle).abs().mean().item()

1110 + err_cpp = (out_cpp.float() - oracle).abs().mean().item()

1111 + assert err_cpp <= err_torch + 1e-3, (

1112 + f"C++ conv less accurate than torch: "

1113 + f"err_cpp={err_cpp:.2e} err_torch={err_torch:.2e}"

1114 + )

1115 +

1116 +

1117 +@pytest.mark.skipif(

1118 + not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"

1119 +)

1120 +@pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS)

1121 +@torch.inference_mode()

1122 +def test_conv_cpp_vnni_packed_matches_torch(seq_lens):

1123 + """The exact runtime prefill sequence (VNNI-packed weight + SD-layout

1124 + conv_state view + is_vnni=True) must match the torch fallback. Validates

1125 + the packing + layout handoff on any AVX-512BF16 CPU."""

1126 + x, weight, bias = _conv_inputs(sum(seq_lens))

1127 + out_torch, _ = _run_prefill_torch(x, weight, bias, seq_lens)

1128 + out_vnni, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=True)

1129 + torch.testing.assert_close(out_vnni, out_torch, atol=1e-2, rtol=1e-2)

1130 +

1131 +

1132 +@pytest.mark.skipif(

1133 + not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"

1134 +)

1135 +@pytest.mark.parametrize("batch", DECODE_BATCH_SIZES)

1136 +@torch.inference_mode()

1137 +def test_conv_update_cpp_matches_torch(batch):

1138 + """Decode conv: causal_conv1d_update_cpu matches causal_conv1d_update_torch,

1139 + including the in-place conv_state update (the next-step handoff)."""

1140 + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (

1141 + causal_conv1d_update_torch,

1142 + )

1143 +

1144 + x = tensor_cache(batch * CONV_DIM, torch.bfloat16).view(batch, CONV_DIM)

1145 + weight = tensor_cache(CONV_DIM * CONV_KERNEL, torch.bfloat16).view(

1146 + CONV_DIM, CONV_KERNEL

1147 + )

1148 + bias = tensor_cache(CONV_DIM, torch.bfloat16)

1149 + conv_state = tensor_cache(batch * CONV_DIM * _STATE_LEN, torch.bfloat16).view(

1150 + batch, CONV_DIM, _STATE_LEN

1151 + )

1152 +

1153 + cs_torch = conv_state.clone()

1154 + out_torch = causal_conv1d_update_torch(

1155 + x=x.unsqueeze(-1),

1156 + conv_state=cs_torch,

1157 + weight=weight,

1158 + bias=bias,

1159 + activation="silu",

1160 + ).squeeze(-1)

1161 +

1162 + cs_cpp = conv_state.clone()

1163 + out_cpp = ops.causal_conv1d_update_cpu(

1164 + x=x.contiguous(),

1165 + conv_states=cs_cpp,

1166 + weight=weight,

1167 + bias=bias,

1168 + silu_activation=True,

1169 + conv_state_indices=torch.arange(batch, dtype=torch.int32),

1170 + is_vnni=False,

1171 + )

1172 + torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2)

1173 + torch.testing.assert_close(cs_cpp, cs_torch, atol=1e-2, rtol=1e-2)

1174 +

1175 +

1176 +@pytest.mark.skipif(

1177 + not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"

1178 +)

1179 +@torch.inference_mode()

1180 +def test_conv_weight_pack_roundtrip_unpacked_matches():

1181 + """`causal_conv1d_weight_pack` repacks the (dim, width) weight, so the C++ conv op

1182 + with packed weight + is_vnni=True must equal unpacked weight + is_vnni=False.

1183 + Guards the spec-decode contract: the runtime VNNI-packs `layer.conv1d.weight`

1184 + in place and stashes the original as `_cpu_unpacked_conv_weight` for the torch

1185 + spec-decode path, so both must produce identical math.

1186 + """

1187 + seq_lens = [7, 64, 65]

1188 + x, weight, bias = _conv_inputs(sum(seq_lens))

1189 + out_unpacked, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False)

1190 + out_packed, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=True)

1191 + torch.testing.assert_close(out_packed, out_unpacked, atol=1e-2, rtol=1e-2)

1192 +

1193 +

1194 +@pytest.mark.skipif(

1195 + not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"

1196 +)

1197 +@torch.inference_mode()

1198 +def test_spec_decode_unpacked_conv_weight_stash():

1199 + """Spec-decode correctness: AVX-512BF16 CPUs VNNI-pack ``conv1d.weight`` in place

1200 + at load time and stash the original ``(dim, width)`` tensor as

1201 + ``_cpu_unpacked_conv_weight``; the spec-decode path must use that stash since

1202 + reading the packed weight directly produces garbage. Verifies the recovered

1203 + weight equals the original and that torch F.conv1d agrees with the C++ conv

1204 + fed the packed weight.

1205 + """

1206 + from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm

1207 +

1208 + torch.manual_seed(0)

1209 + # conv1d weight is stored [dim, 1, width]; bias [dim].

1210 + orig_2d = torch.rand(CONV_DIM, CONV_KERNEL, dtype=torch.bfloat16)

1211 + bias = torch.rand(CONV_DIM, dtype=torch.bfloat16)

1212 +

1213 + conv = torch.nn.Module()

1214 + conv.weight = torch.nn.Parameter(

1215 + orig_2d.view(CONV_DIM, 1, CONV_KERNEL).clone(), requires_grad=False

1216 + )

1217 +

1218 + # Load-time dispatch: packs weight in place + stashes the unpacked copy.

1219 + dispatch_cpu_unquantized_gemm(conv, remove_weight=False)

1220 +

1221 + # 1. The stash must exist and equal the original (dim, width) weight.

1222 + assert hasattr(conv, "_cpu_unpacked_conv_weight"), (

1223 + "dispatch_cpu_unquantized_gemm did not stash _cpu_unpacked_conv_weight "

1224 + "on an AVX-512BF16 CPU"

1225 + )

1226 + stash = conv._cpu_unpacked_conv_weight

1227 + torch.testing.assert_close(stash, orig_2d, atol=0, rtol=0)

1228 +

1229 + # 2. Mirror _unpacked_conv_weight()'s lookup: stash wins over conv.weight.

1230 + recovered = getattr(conv, "_cpu_unpacked_conv_weight", None)

1231 + assert recovered is not None

1232 + # conv.weight is now packed; using it directly (the bug) would differ.

1233 + packed_weight = conv.weight # [dim, 1, width], VNNI-packed contents

1234 +

1235 + # 3. Spec-path torch conv with the recovered unpacked weight must match the

1236 + # nonspec C++ conv with the packed weight (is_vnni=True).

1237 + seq_lens = [7, 65]

1238 + total = sum(seq_lens)

1239 + x = tensor_cache(total * CONV_DIM, torch.bfloat16).view(total, CONV_DIM)

1240 +

1241 + out_torch, _ = _run_prefill_torch(x, recovered, bias, seq_lens)

1242 +

1243 + num_seqs = len(seq_lens)

1244 + cs = torch.zeros(num_seqs, _STATE_LEN, CONV_DIM, dtype=x.dtype).transpose(1, 2)

1245 + qsl = torch.tensor(

1246 + [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32

1247 + )

1248 + out_cpp = (

1249 + ops.causal_conv1d_fwd_cpu(

1250 + x=x.transpose(0, 1),

1251 + weight=packed_weight.view(CONV_DIM, CONV_KERNEL),

1252 + bias=bias,

1253 + conv_states=cs,

1254 + query_start_loc=qsl,

1255 + cache_indices=torch.arange(num_seqs, dtype=torch.int32),

1256 + has_initial_state=torch.zeros(num_seqs, dtype=torch.bool),

1257 + silu_activation=True,

1258 + is_vnni=True,

1259 + )

1260 + .transpose(0, 1)

1261 + .contiguous()

1262 + )

1263 +

1264 + torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2)

vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py

+8 -6

@@ -103,13 +103,15 @@ def _cpu_gdn_attention_nonspec(

103 assert state_indices_tensor is not None

104 assert query_start_loc is not None

105

106 - is_amx = torch.cpu._is_amx_tile_supported()

106 + # C++ conv (conv.cpp) uses VDPBF16PS, not AMX tiles, so it runs on any

107 + # AVX-512BF16 CPU; weight is VNNI-packed on this same predicate at load time.

108 + use_cpp_conv = torch.cpu._is_avx512_bf16_supported()

109

110 conv_state = layer.kv_cache[0]

109 - if is_amx:

110 - # AMX causal conv requires [num_allocated_slots, kernel - 1, conv_dim].

111 + if use_cpp_conv:

112 + # C++ conv requires [num_allocated_slots, kernel - 1, conv_dim] (SD).

113 if is_conv_state_dim_first():

112 - raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")

114 + raise RuntimeError("C++ CPU GDN attention requires `SD` conv_state layout.")

115 conv_state = conv_state.transpose(1, 2)

116 else:

117 if not is_conv_state_dim_first():

@@ -138,7 +140,7 @@ def _cpu_gdn_attention_nonspec(

140 decode_b = b[:num_decode_tokens]

141 decode_a = a[:num_decode_tokens]

142 decode_state_indices = state_indices_tensor[:num_decodes]

141 - if is_amx:

143 + if use_cpp_conv:

144 decode_mixed_qkv = ops.causal_conv1d_update_cpu(

145 x=decode_mixed_qkv,

146 conv_states=conv_state,

@@ -207,7 +209,7 @@ def _cpu_gdn_attention_nonspec(

209 num_decodes : num_decodes + num_prefills

210 ]

211

210 - if is_amx:

212 + if use_cpp_conv:

213 prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu(

214 x=prefill_mixed_qkv.transpose(0, 1),

215 weight=layer.conv1d.weight,

vllm/model_executor/layers/utils.py

+5 -2

@@ -280,7 +280,10 @@ def dispatch_cpu_unquantized_gemm(

280 if layer.weight.ndim != 2:

281 # this is not a linear layer

282 # For now it should be a causal_conv1d op or MoE 3D expert weights

283 - if torch.cpu._is_amx_tile_supported() and hasattr(

283 + # The C++ causal_conv1d kernels use VDPBF16PS (no AMX tiles), so the

284 + # VNNI weight prepack applies to any AVX-512BF16 CPU, not just AMX

285 + # (e.g. AMD Zen5/Turin).

286 + if torch.cpu._is_avx512_bf16_supported() and hasattr(

287 ops, "causal_conv1d_weight_pack"

288 ):

289 # prepack conv weight

@@ -293,7 +296,7 @@ def dispatch_cpu_unquantized_gemm(

296 .clone()

297 )

298 # Stash the un-packed (dim, width) weight so the speculative-decode

296 - # GDN path (which uses torch conv, not the AMX kernel) can use it.

299 + # GDN path (which uses torch conv, not the C++ kernel) can use it.

300 layer._cpu_unpacked_conv_weight = unpacked

301 layer.weight.data = ops.causal_conv1d_weight_pack(unpacked)

302 return

vllm/platforms/cpu.py

+7 -5

@@ -181,13 +181,13 @@ class CpuPlatform(Platform):

181 "otherwise the performance is not optimized."

182 )

183

184 - # AMX GDN requires float32 state

184 + # Accelerated GDN (AMX tiles or AVX-512BF16 VDPBF16PS) requires float32 SSM state.

185 if (

186 - torch.cpu._is_amx_tile_supported()

186 + torch.cpu._is_avx512_bf16_supported()

187 and cache_config.mamba_ssm_cache_dtype != "float32"

188 ):

189 cache_config.mamba_ssm_cache_dtype = "float32"

190 - logger.warning("Reset SSM cache type to float32 for AMX mamba attention.")

190 + logger.warning("Reset SSM cache type to float32 for accelerated GDN mamba attention.")

191

192 # Lagecy setting

193 env_key = "VLLM_CPU_KVCACHE_SPACE"

@@ -290,8 +290,10 @@ class CpuPlatform(Platform):

290 # Avoid inductor generates num_thread() and breaks the thread binding

291 os.environ["TORCHINDUCTOR_CPP_DYNAMIC_THREADS"] = "1"

292

293 - # For efficient conv state memory access

294 - if torch.cpu._is_amx_tile_supported():

293 + # For efficient conv state memory access. The C++ causal_conv1d

294 + # kernels (VDPBF16PS, no AMX tiles) consume the SD layout on any

295 + # AVX-512BF16 CPU, so apply it beyond AMX (e.g. AMD Zen5/Turin).

296 + if torch.cpu._is_avx512_bf16_supported():

297 os.environ["VLLM_SSM_CONV_STATE_LAYOUT"] = "SD"

298

299 ld_preload_str = os.getenv("LD_PRELOAD", "")