[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>

tests/kernels/mamba/cpu/test_cpu_gdn_ops.py

@@ -878,17 +878,18 @@ def test_causal_conv1d_update_cpu_rejects_invalid_accepted_count(
 
 
 @pytest.mark.skipif(
-    not torch.cpu._is_amx_tile_supported(),
-    reason="requires AMX support",
+    not torch.cpu._is_avx512_bf16_supported(),
+    reason="causal_conv1d_fwd_cpu requires AVX-512BF16 (Intel Xeon or AMD EPYC)",
 )
 @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS)
 @torch.inference_mode()
 def test_causal_conv1d_fwd_cpu_two_call_split(total_tokens: int, split: int) -> None:
-    """AMX prefill conv op must honor ``has_initial_state`` so a two-call split
+    """C++ prefill conv op must honor ``has_initial_state`` so a two-call split
     matches the single-call result.
 
     Regression test for ``causal_conv1d_fwd_varlen_kernel_impl`` (``conv.cpp``)
-    ignoring the carried conv state on continued chunks.
+    ignoring the carried conv state on continued chunks. Runs on any
+    AVX-512BF16 CPU since conv.cpp uses VDPBF16PS, not AMX tiles.
     """
     state_len = CONV_KERNEL - 1
     x, weight, bias = _conv_inputs(total_tokens)
@@ -996,3 +997,268 @@ def test_batch_memcpy_cpu_fallback() -> None:
 
     for src, dst in zip(srcs, dsts):
         torch.testing.assert_close(dst, src)
+
+
+# ---------------------------------------------------------------------------
+# C++ conv (conv.cpp) uses VDPBF16PS, not AMX tiles, so it runs on any
+# AVX-512BF16 CPU; weight is VNNI-packed on this same predicate at load time.
+# ---------------------------------------------------------------------------
+
+_HAS_AVX512_BF16 = torch.cpu._is_avx512_bf16_supported()
+
+_STATE_LEN = CONV_KERNEL - 1
+
+CONV_EQUIV_SEQ_LENS = [[1], [7], [64], [65], [1, 2, 3], [63, 64, 65], [128, 129]]
+
+
+def _conv_fp32_oracle(x, weight, bias, seq_lens, activation="silu"):
+    """High-precision conv reference: everything in float32, no bf16 rounding.
+    x: [total_tokens, dim] (no initial state). Returns [total_tokens, dim]."""
+    xf = x.float()
+    wf = weight.float().unsqueeze(1)
+    bf = bias.float()
+    out = torch.empty_like(xf)
+    start = 0
+    for n in seq_lens:
+        seg = xf[start : start + n].transpose(0, 1).unsqueeze(0)  # [1, dim, n]
+        conv_in = F.pad(seg, (_STATE_LEN, 0))
+        seg_out = F.conv1d(conv_in, wf, bf, padding=0, groups=CONV_DIM)[..., -n:]
+        if activation in ("silu", "swish"):
+            seg_out = F.silu(seg_out)
+        out[start : start + n] = seg_out.squeeze(0).transpose(0, 1)
+        start += n
+    return out
+
+
+def _run_prefill_torch(x, weight, bias, seq_lens):
+    from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
+        causal_conv1d_fn_cpu as causal_conv1d_torch,
+    )
+
+    num_seqs = len(seq_lens)
+    conv_states = torch.zeros(num_seqs, CONV_DIM, _STATE_LEN, dtype=x.dtype)
+    qsl = torch.tensor(
+        [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32
+    )
+    out = causal_conv1d_torch(
+        x=x.transpose(0, 1).contiguous(),
+        weight=weight,
+        bias=bias,
+        conv_states=conv_states,
+        query_start_loc=qsl,
+        cache_indices=torch.arange(num_seqs, dtype=torch.int32),
+        has_initial_state=torch.zeros(num_seqs, dtype=torch.bool),
+        activation="silu",
+    )
+    return out.transpose(0, 1).contiguous(), conv_states
+
+
+def _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False):
+    num_seqs = len(seq_lens)
+    packed_w = ops.causal_conv1d_weight_pack(weight) if is_vnni else weight
+    if is_vnni:
+        # C++-branch layout: kv-cache "SD" [slots, state_len, dim] transposed to
+        # [slots, dim, state_len] (a non-contiguous view).
+        conv_state = torch.zeros(
+            num_seqs, _STATE_LEN, CONV_DIM, dtype=x.dtype
+        ).transpose(1, 2)
+    else:
+        conv_state = torch.zeros(num_seqs, CONV_DIM, _STATE_LEN, dtype=x.dtype)
+    qsl = torch.tensor(
+        [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32
+    )
+    out = ops.causal_conv1d_fwd_cpu(
+        x=x.transpose(0, 1),
+        weight=packed_w,
+        bias=bias,
+        conv_states=conv_state,
+        query_start_loc=qsl,
+        cache_indices=torch.arange(num_seqs, dtype=torch.int32),
+        has_initial_state=torch.zeros(num_seqs, dtype=torch.bool),
+        silu_activation=True,
+        is_vnni=is_vnni,
+    )
+    return out.transpose(0, 1).contiguous(), conv_state
+
+
+@pytest.mark.skipif(
+    not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
+)
+@pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS)
+@torch.inference_mode()
+def test_conv_cpp_matches_torch(seq_lens):
+    """C++ causal_conv1d_fwd_cpu matches the torch fallback within bf16 tol."""
+    x, weight, bias = _conv_inputs(sum(seq_lens))
+    out_torch, state_torch = _run_prefill_torch(x, weight, bias, seq_lens)
+    out_cpp, state_cpp = _run_prefill_cpp(x, weight, bias, seq_lens)
+    torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(state_cpp, state_torch, atol=1e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(
+    not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
+)
+@pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS)
+@torch.inference_mode()
+def test_conv_cpp_no_worse_than_torch_vs_fp32(seq_lens):
+    """Swapping torch -> C++ conv must not increase error vs an fp32 oracle."""
+    x, weight, bias = _conv_inputs(sum(seq_lens))
+    oracle = _conv_fp32_oracle(x, weight, bias, seq_lens)
+    out_torch, _ = _run_prefill_torch(x, weight, bias, seq_lens)
+    out_cpp, _ = _run_prefill_cpp(x, weight, bias, seq_lens)
+    err_torch = (out_torch.float() - oracle).abs().mean().item()
+    err_cpp = (out_cpp.float() - oracle).abs().mean().item()
+    assert err_cpp <= err_torch + 1e-3, (
+        f"C++ conv less accurate than torch: "
+        f"err_cpp={err_cpp:.2e} err_torch={err_torch:.2e}"
+    )
+
+
+@pytest.mark.skipif(
+    not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
+)
+@pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS)
+@torch.inference_mode()
+def test_conv_cpp_vnni_packed_matches_torch(seq_lens):
+    """The exact runtime prefill sequence (VNNI-packed weight + SD-layout
+    conv_state view + is_vnni=True) must match the torch fallback. Validates
+    the packing + layout handoff on any AVX-512BF16 CPU."""
+    x, weight, bias = _conv_inputs(sum(seq_lens))
+    out_torch, _ = _run_prefill_torch(x, weight, bias, seq_lens)
+    out_vnni, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=True)
+    torch.testing.assert_close(out_vnni, out_torch, atol=1e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(
+    not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
+)
+@pytest.mark.parametrize("batch", DECODE_BATCH_SIZES)
+@torch.inference_mode()
+def test_conv_update_cpp_matches_torch(batch):
+    """Decode conv: causal_conv1d_update_cpu matches causal_conv1d_update_torch,
+    including the in-place conv_state update (the next-step handoff)."""
+    from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
+        causal_conv1d_update_torch,
+    )
+
+    x = tensor_cache(batch * CONV_DIM, torch.bfloat16).view(batch, CONV_DIM)
+    weight = tensor_cache(CONV_DIM * CONV_KERNEL, torch.bfloat16).view(
+        CONV_DIM, CONV_KERNEL
+    )
+    bias = tensor_cache(CONV_DIM, torch.bfloat16)
+    conv_state = tensor_cache(batch * CONV_DIM * _STATE_LEN, torch.bfloat16).view(
+        batch, CONV_DIM, _STATE_LEN
+    )
+
+    cs_torch = conv_state.clone()
+    out_torch = causal_conv1d_update_torch(
+        x=x.unsqueeze(-1),
+        conv_state=cs_torch,
+        weight=weight,
+        bias=bias,
+        activation="silu",
+    ).squeeze(-1)
+
+    cs_cpp = conv_state.clone()
+    out_cpp = ops.causal_conv1d_update_cpu(
+        x=x.contiguous(),
+        conv_states=cs_cpp,
+        weight=weight,
+        bias=bias,
+        silu_activation=True,
+        conv_state_indices=torch.arange(batch, dtype=torch.int32),
+        is_vnni=False,
+    )
+    torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(cs_cpp, cs_torch, atol=1e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(
+    not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
+)
+@torch.inference_mode()
+def test_conv_weight_pack_roundtrip_unpacked_matches():
+    """`causal_conv1d_weight_pack` repacks the (dim, width) weight, so the C++ conv op
+    with packed weight + is_vnni=True must equal unpacked weight + is_vnni=False.
+    Guards the spec-decode contract: the runtime VNNI-packs `layer.conv1d.weight`
+    in place and stashes the original as `_cpu_unpacked_conv_weight` for the torch
+    spec-decode path, so both must produce identical math.
+    """
+    seq_lens = [7, 64, 65]
+    x, weight, bias = _conv_inputs(sum(seq_lens))
+    out_unpacked, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False)
+    out_packed, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=True)
+    torch.testing.assert_close(out_packed, out_unpacked, atol=1e-2, rtol=1e-2)
+
+
+@pytest.mark.skipif(
+    not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
+)
+@torch.inference_mode()
+def test_spec_decode_unpacked_conv_weight_stash():
+    """Spec-decode correctness: AVX-512BF16 CPUs VNNI-pack ``conv1d.weight`` in place
+    at load time and stash the original ``(dim, width)`` tensor as
+    ``_cpu_unpacked_conv_weight``; the spec-decode path must use that stash since
+    reading the packed weight directly produces garbage. Verifies the recovered
+    weight equals the original and that torch F.conv1d agrees with the C++ conv
+    fed the packed weight.
+    """
+    from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm
+
+    torch.manual_seed(0)
+    # conv1d weight is stored [dim, 1, width]; bias [dim].
+    orig_2d = torch.rand(CONV_DIM, CONV_KERNEL, dtype=torch.bfloat16)
+    bias = torch.rand(CONV_DIM, dtype=torch.bfloat16)
+
+    conv = torch.nn.Module()
+    conv.weight = torch.nn.Parameter(
+        orig_2d.view(CONV_DIM, 1, CONV_KERNEL).clone(), requires_grad=False
+    )
+
+    # Load-time dispatch: packs weight in place + stashes the unpacked copy.
+    dispatch_cpu_unquantized_gemm(conv, remove_weight=False)
+
+    # 1. The stash must exist and equal the original (dim, width) weight.
+    assert hasattr(conv, "_cpu_unpacked_conv_weight"), (
+        "dispatch_cpu_unquantized_gemm did not stash _cpu_unpacked_conv_weight "
+        "on an AVX-512BF16 CPU"
+    )
+    stash = conv._cpu_unpacked_conv_weight
+    torch.testing.assert_close(stash, orig_2d, atol=0, rtol=0)
+
+    # 2. Mirror _unpacked_conv_weight()'s lookup: stash wins over conv.weight.
+    recovered = getattr(conv, "_cpu_unpacked_conv_weight", None)
+    assert recovered is not None
+    # conv.weight is now packed; using it directly (the bug) would differ.
+    packed_weight = conv.weight  # [dim, 1, width], VNNI-packed contents
+
+    # 3. Spec-path torch conv with the recovered unpacked weight must match the
+    #    nonspec C++ conv with the packed weight (is_vnni=True).
+    seq_lens = [7, 65]
+    total = sum(seq_lens)
+    x = tensor_cache(total * CONV_DIM, torch.bfloat16).view(total, CONV_DIM)
+
+    out_torch, _ = _run_prefill_torch(x, recovered, bias, seq_lens)
+
+    num_seqs = len(seq_lens)
+    cs = torch.zeros(num_seqs, _STATE_LEN, CONV_DIM, dtype=x.dtype).transpose(1, 2)
+    qsl = torch.tensor(
+        [0, *torch.tensor(seq_lens).cumsum(0).tolist()], dtype=torch.int32
+    )
+    out_cpp = (
+        ops.causal_conv1d_fwd_cpu(
+            x=x.transpose(0, 1),
+            weight=packed_weight.view(CONV_DIM, CONV_KERNEL),
+            bias=bias,
+            conv_states=cs,
+            query_start_loc=qsl,
+            cache_indices=torch.arange(num_seqs, dtype=torch.int32),
+            has_initial_state=torch.zeros(num_seqs, dtype=torch.bool),
+            silu_activation=True,
+            is_vnni=True,
+        )
+        .transpose(0, 1)
+        .contiguous()
+    )
+
+    torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2)

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

@@ -103,13 +103,15 @@ def _cpu_gdn_attention_nonspec(
     assert state_indices_tensor is not None
     assert query_start_loc is not None
 
-    is_amx = torch.cpu._is_amx_tile_supported()
+    # C++ conv (conv.cpp) uses VDPBF16PS, not AMX tiles, so it runs on any
+    # AVX-512BF16 CPU; weight is VNNI-packed on this same predicate at load time.
+    use_cpp_conv = torch.cpu._is_avx512_bf16_supported()
 
     conv_state = layer.kv_cache[0]
-    if is_amx:
-        # AMX causal conv requires [num_allocated_slots, kernel - 1, conv_dim].
+    if use_cpp_conv:
+        # C++ conv requires [num_allocated_slots, kernel - 1, conv_dim] (SD).
         if is_conv_state_dim_first():
-            raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")
+            raise RuntimeError("C++ CPU GDN attention requires `SD` conv_state layout.")
         conv_state = conv_state.transpose(1, 2)
     else:
         if not is_conv_state_dim_first():
@@ -138,7 +140,7 @@ def _cpu_gdn_attention_nonspec(
         decode_b = b[:num_decode_tokens]
         decode_a = a[:num_decode_tokens]
         decode_state_indices = state_indices_tensor[:num_decodes]
-        if is_amx:
+        if use_cpp_conv:
             decode_mixed_qkv = ops.causal_conv1d_update_cpu(
                 x=decode_mixed_qkv,
                 conv_states=conv_state,
@@ -207,7 +209,7 @@ def _cpu_gdn_attention_nonspec(
             num_decodes : num_decodes + num_prefills
         ]
 
-        if is_amx:
+        if use_cpp_conv:
             prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu(
                 x=prefill_mixed_qkv.transpose(0, 1),
                 weight=layer.conv1d.weight,

vllm/model_executor/layers/utils.py

@@ -280,7 +280,10 @@ def dispatch_cpu_unquantized_gemm(
     if layer.weight.ndim != 2:
         # this is not a linear layer
         # For now it should be a causal_conv1d op or MoE 3D expert weights
-        if torch.cpu._is_amx_tile_supported() and hasattr(
+        # The C++ causal_conv1d kernels use VDPBF16PS (no AMX tiles), so the
+        # VNNI weight prepack applies to any AVX-512BF16 CPU, not just AMX
+        # (e.g. AMD Zen5/Turin).
+        if torch.cpu._is_avx512_bf16_supported() and hasattr(
             ops, "causal_conv1d_weight_pack"
         ):
             # prepack conv weight
@@ -293,7 +296,7 @@ def dispatch_cpu_unquantized_gemm(
                 .clone()
             )
             # Stash the un-packed (dim, width) weight so the speculative-decode
-            # GDN path (which uses torch conv, not the AMX kernel) can use it.
+            # GDN path (which uses torch conv, not the C++ kernel) can use it.
             layer._cpu_unpacked_conv_weight = unpacked
             layer.weight.data = ops.causal_conv1d_weight_pack(unpacked)
         return

vllm/platforms/cpu.py

@@ -181,13 +181,13 @@ class CpuPlatform(Platform):
                 "otherwise the performance is not optimized."
             )
 
-        # AMX GDN requires float32 state
+        # Accelerated GDN (AMX tiles or AVX-512BF16 VDPBF16PS) requires float32 SSM state.
         if (
-            torch.cpu._is_amx_tile_supported()
+            torch.cpu._is_avx512_bf16_supported()
             and cache_config.mamba_ssm_cache_dtype != "float32"
         ):
             cache_config.mamba_ssm_cache_dtype = "float32"
-            logger.warning("Reset SSM cache type to float32 for AMX mamba attention.")
+            logger.warning("Reset SSM cache type to float32 for accelerated GDN mamba attention.")
 
         # Lagecy setting
         env_key = "VLLM_CPU_KVCACHE_SPACE"
@@ -290,8 +290,10 @@ class CpuPlatform(Platform):
         # Avoid inductor generates num_thread() and breaks the thread binding
         os.environ["TORCHINDUCTOR_CPP_DYNAMIC_THREADS"] = "1"
 
-        # For efficient conv state memory access
-        if torch.cpu._is_amx_tile_supported():
+        # For efficient conv state memory access. The C++ causal_conv1d
+        # kernels (VDPBF16PS, no AMX tiles) consume the SD layout on any
+        # AVX-512BF16 CPU, so apply it beyond AMX (e.g. AMD Zen5/Turin).
+        if torch.cpu._is_avx512_bf16_supported():
             os.environ["VLLM_SSM_CONV_STATE_LAYOUT"] = "SD"
 
         ld_preload_str = os.getenv("LD_PRELOAD", "")