[CPU] Conv1d optimised kernel for aarch64 (#54093)

Signed-off-by: Anna Mayne <anna.mayne@arm.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>

cmake/cpu_extension.cmake

@@ -524,6 +524,7 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)
         "csrc/cpu/cpu_tanhf_neon.hpp"
         ${VLLM_EXT_SRC})
     if (ARM_BF16_FOUND)
+        set(VLLM_EXT_SRC "csrc/cpu/sgl-kernels/conv.cpp" ${VLLM_EXT_SRC})
         if (ARM_I8MM_FOUND)
             set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe_int8.cpp" ${VLLM_EXT_SRC})
         endif()

csrc/cpu/sgl-kernels/conv.cpp

@@ -66,6 +66,10 @@ struct tinygemm_kernel {
   }
 };
 
+#if defined(ARM_BF16_SUPPORT)
+#include "conv_arm.h"
+#endif
+
 #if defined(CPU_CAPABILITY_AVX512)
 template <int K, int BLOCK_N, bool has_bias, bool has_silu>
 struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> {
@@ -543,16 +547,24 @@ at::Tensor causal_conv1d_weight_pack(const at::Tensor& weight) {
 
   int64_t dim = weight.size(0);
   int64_t width = weight.size(1);
-  constexpr int64_t BLOCK_N = block_size_n();
+  constexpr int64_t BLOCK_N =
+#if defined(ARM_BF16_SUPPORT)
+      8;// 8 channel blocks for microtile
+#else
+      block_size_n();
+#endif
   TORCH_CHECK(width == 4, "causal_conv1d_weight_pack: support only width of 4");
-  TORCH_CHECK(dim % BLOCK_N == 0, "causal_conv1d_weight_pack: invalid dim size ", dim);
+  TORCH_CHECK(
+      dim % block_size_n() == 0,
+      "causal_conv1d_weight_pack: invalid dim size ",
+      dim);
 
   const int64_t N = dim, K2 = width >> 1;
   const int64_t NB = div_up(N, BLOCK_N);
 
   auto packed_weight = at::empty_like(weight);
   AT_DISPATCH_REDUCED_FLOATING_TYPES(weight.scalar_type(), "causal_conv1d_fwd_kernel_impl", [&] {
-    // cast to float32 as vnni size is 2
+    // Copy each pair as an opaque 32-bit word. This is not a bf16->fp32 conversion.
     const float* w_data = reinterpret_cast<float*>(weight.data_ptr<scalar_t>());
     float* packed_data = reinterpret_cast<float*>(packed_weight.data_ptr<scalar_t>());
 
@@ -642,9 +654,9 @@ at::Tensor causal_conv1d_fwd_cpu(
     const std::optional<at::Tensor>& has_initial_state,
     bool silu_activation,
     int64_t pad_slot_id,
-    bool is_vnni) {
+    bool is_weight_packed) {
   CHECK_CONTIGUOUS(weight);
-  auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
+  auto packed_w = is_weight_packed ? weight : causal_conv1d_weight_pack(weight);
 
   const bool is_var_seqlen = query_start_loc.has_value();
   const int64_t input_ndim = is_var_seqlen ? 2 : 3;
@@ -755,10 +767,10 @@ at::Tensor causal_conv1d_update_cpu(
     const std::optional<at::Tensor>& num_accepted_tokens,
     const std::optional<at::Tensor>& conv_state_indices,
     int64_t pad_slot_id,
-    bool is_vnni) {
+    bool is_weight_packed) {
   CHECK_CONTIGUOUS(x);
   CHECK_CONTIGUOUS(weight);
-  auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
+  auto packed_w = is_weight_packed ? weight : causal_conv1d_weight_pack(weight);
 
   TORCH_CHECK(
       x.dim() == 2 || x.dim() == 3,

csrc/cpu/sgl-kernels/conv_arm.h

@@ -0,0 +1,126 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+#pragma once
+
+#if defined(ARM_BF16_SUPPORT)
+
+template <int K, int BLOCK_N, bool has_bias, bool has_silu>
+struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> {
+  static inline bfloat16x8_t load_bf16x8(const at::BFloat16* ptr) {
+    return vld1q_bf16(reinterpret_cast<const bfloat16_t*>(ptr));
+  }
+
+  // Input layout: A[m][channel]
+  // Weight layout: B[N/MICRO_N, K/2, MICRO_N, 2]
+  // The packed weights are reused across all M positions.
+  static inline void tinygemm_neon_8(
+      const at::BFloat16* __restrict__ A, const at::BFloat16* __restrict__ B,
+      at::BFloat16* __restrict__ C, const at::BFloat16* __restrict__ bias,
+      const at::BFloat16* __restrict__ conv_states, bool has_initial_state,
+      int64_t M, int64_t lda, bool is_first_token) {
+    bfloat16x8_t va[4];
+    // Load 4 pairs of kernel weights into 4 registers
+    const bfloat16x8_t vb0 = load_bf16x8(B);       // channels 0-3, taps 0-1
+    const bfloat16x8_t vb1 = load_bf16x8(B + 8);   // channels 4-7, taps 0-1
+    const bfloat16x8_t vb2 = load_bf16x8(B + 16);  // channels 0-3, taps 2-3
+    const bfloat16x8_t vb3 = load_bf16x8(B + 24);  // channels 4-7, taps 2-3
+
+    // Load the input data. A single register holds the values for a
+    // single token across all 8 channels for this microtile.
+    auto load_a = [&](int64_t m) {
+      if (m < 0 && is_first_token) {
+        if (!has_initial_state) {
+          return vreinterpretq_bf16_u16(vdupq_n_u16(0));
+        }
+        return load_bf16x8(conv_states + (m + K - 1) * lda);
+      }
+      return load_bf16x8(A + m * lda);
+    };
+
+    // Load the previous inputs from the cached state if available.
+    va[1] = load_a(-3);
+    va[2] = load_a(-2);
+    va[3] = load_a(-1);
+
+    // Iterate over the sequence block, updating the previous conv inputs
+    // and reading the next token's input data (8 channels).
+    for (int64_t m = 0; m < M; ++m) {
+      va[0] = va[1];
+      va[1] = va[2];
+      va[2] = va[3];
+      va[3] = load_a(m);
+
+      float32x4_t vc0;
+      float32x4_t vc1;
+      if constexpr (has_bias) {
+        const bfloat16x8_t bias_vec = load_bf16x8(bias);
+        vc0 = vcvtq_low_f32_bf16(bias_vec);
+        vc1 = vcvtq_high_f32_bf16(bias_vec);
+      } else {
+        vc0 = vdupq_n_f32(0.0f);
+        vc1 = vdupq_n_f32(0.0f);
+      }
+
+      // Use the zip operation to match the input layout to the weights layout
+      bfloat16x8_t va_pair =
+          vreinterpretq_bf16_u16(  // channels 0-3, inputs for taps 0-1
+              vzip1q_u16(vreinterpretq_u16_bf16(va[0]),
+                         vreinterpretq_u16_bf16(va[1])));
+      vc0 = vbfdotq_f32(vc0, va_pair,
+                        vb0);  // 4 independent length-two dot products
+      va_pair = vreinterpretq_bf16_u16(  // channels 4-7, inputs for taps 0-1
+          vzip2q_u16(vreinterpretq_u16_bf16(va[0]),
+                     vreinterpretq_u16_bf16(va[1])));
+      vc1 = vbfdotq_f32(vc1, va_pair, vb1);
+      va_pair = vreinterpretq_bf16_u16(  // channels 0-3, inputs for taps 2-3
+          vzip1q_u16(vreinterpretq_u16_bf16(va[2]),
+                     vreinterpretq_u16_bf16(va[3])));
+      // Add to the previous result for the same channels
+      vc0 = vbfdotq_f32(vc0, va_pair, vb2);
+      va_pair = vreinterpretq_bf16_u16(  // channels 4-7, inputs for taps 2-3
+          vzip2q_u16(vreinterpretq_u16_bf16(va[2]),
+                     vreinterpretq_u16_bf16(va[3])));
+      vc1 = vbfdotq_f32(vc1, va_pair, vb3);
+
+      using fVec = at::vec::Vectorized<float>;
+      fVec out0(vc0);
+      fVec out1(vc1);
+      if constexpr (has_silu) {
+        const fVec one(1.0f);
+        out0 = out0 / (one + out0.neg().fexp_u20());
+        out1 = out1 / (one + out1.neg().fexp_u20());
+      }
+
+      // Convert the 2 x 4 fp32 results to 1 x 8 bf16
+      const bfloat16x8_t out =
+          vcombine_bf16(vcvt_bf16_f32(static_cast<float32x4_t>(out0)),
+                        vcvt_bf16_f32(static_cast<float32x4_t>(out1)));
+      vst1q_bf16(reinterpret_cast<bfloat16_t*>(C + m * lda), out);
+    }
+  }
+
+  static inline void apply(const at::BFloat16* __restrict__ A,
+                           const at::BFloat16* __restrict__ B,
+                           at::BFloat16* __restrict__ C,
+                           const at::BFloat16* __restrict__ bias,
+                           const at::BFloat16* __restrict__ conv_states,
+                           bool has_initial_state, int64_t M, int64_t lda,
+                           bool is_first_token) {
+    static_assert(K == 4);
+    constexpr int64_t MICRO_N = 8;
+    static_assert(BLOCK_N % MICRO_N == 0);
+
+    // BLOCK_N is always 32 or 64. M is the block sequence length.
+    // Input layout: A[m][channel]
+    // Weight layout: B[N/MICRO_N, K/2, MICRO_N, 2]
+    // Use an 8 channel microtile and iterate over the 8 channel blocks.
+    for (int64_t n = 0; n < BLOCK_N; n += MICRO_N) {
+      tinygemm_neon_8(A + n, B + n * K, C + n, has_bias ? bias + n : nullptr,
+                      conv_states ? conv_states + n : nullptr,
+                      has_initial_state, M, lda, is_first_token);
+    }
+  }
+};
+
+#endif

csrc/cpu/torch_bindings.cpp

@@ -281,14 +281,14 @@ at::Tensor causal_conv1d_fwd_cpu(
     const std::optional<at::Tensor>& query_start_loc,
     const std::optional<at::Tensor>& cache_indices,
     const std::optional<at::Tensor>& has_initial_state, bool silu_activation,
-    int64_t pad_slot_id, bool is_vnni);
+    int64_t pad_slot_id, bool is_weight_packed);
 
 at::Tensor causal_conv1d_update_cpu(
     const at::Tensor& x, const at::Tensor& conv_states,
     const at::Tensor& weight, const std::optional<at::Tensor>& bias,
     bool silu_activation, const std::optional<at::Tensor>& num_accepted_tokens,
     const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id,
-    bool is_vnni);
+    bool is_weight_packed);
 
 void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,
                          const std::string& activation);
@@ -517,6 +517,33 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 
 #endif  // (defined(__aarch64__) && !defined(__APPLE__))
 
+#if (defined(__aarch64__) && defined(ARM_BF16_SUPPORT) && \
+     !defined(__APPLE__)) ||                              \
+    (defined(__AVX512BF16__) && defined(__AVX512F__) &&   \
+     defined(__AVX512VNNI__))
+
+  // Adapted from sglang: casual_conv1d kernels
+  ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
+  ops.impl("causal_conv1d_weight_pack", torch::kCPU,
+           &causal_conv1d_weight_pack);
+
+  ops.def(
+      "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
+      "conv_states, Tensor? query_start_loc,"
+      "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
+      "int pad_slot_id, bool is_weight_packed) -> "
+      "Tensor");
+  ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
+
+  ops.def(
+      "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
+      "weight, Tensor? bias, bool silu_activation,"
+      "Tensor? num_accepted_tokens, Tensor? conv_state_indices, int "
+      "pad_slot_id, "
+      "bool is_weight_packed) -> Tensor");
+  ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
+#endif
+
   // Layernorm
   // Apply Root Mean Square (RMS) Normalization to the input tensor.
   ops.def(
@@ -665,24 +692,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
       "-> (Tensor, Tensor)");
   ops.impl("quantize_fp8e4m3_vec", torch::kCPU, &quantize_fp8e4m3_vec);
 
-  ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");
-  ops.impl("causal_conv1d_weight_pack", torch::kCPU,
-           &causal_conv1d_weight_pack);
-  ops.def(
-      "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "
-      "conv_states, Tensor? query_start_loc,"
-      "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "
-      "int pad_slot_id, bool is_vnni) -> "
-      "Tensor");
-  ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);
-  ops.def(
-      "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "
-      "weight, Tensor? bias, bool silu_activation,"
-      "Tensor? num_accepted_tokens, Tensor? conv_state_indices, int "
-      "pad_slot_id, "
-      "bool is_vnni) -> Tensor");
-  ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
-
   // Adapted from sglang: MLA CPU kernels (AMX-only, DeepSeek V2/V3/R1)
   ops.def(
       "decode_attention_cpu(Tensor query, Tensor k_buffer, Tensor v_buffer, "

tests/kernels/mamba/cpu/test_cpu_gdn_ops.py

@@ -494,8 +494,10 @@ def _sd_conv_states(
     return storage.transpose(1, 2)
 
 
-def _maybe_pack_conv_weight(weight: torch.Tensor, is_vnni: bool) -> torch.Tensor:
-    return ops.causal_conv1d_weight_pack(weight) if is_vnni else weight
+def _maybe_pack_conv_weight(
+    weight: torch.Tensor, is_weight_packed: bool
+) -> torch.Tensor:
+    return ops.causal_conv1d_weight_pack(weight) if is_weight_packed else weight
 
 
 @torch.inference_mode()
@@ -523,7 +525,7 @@ def test_causal_conv1d_update_cpu_preserves_sd_state(
         bias=None,
         silu_activation=False,
         conv_state_indices=state_indices,
-        is_vnni=False,
+        is_weight_packed=False,
     )
 
     assert out is x
@@ -666,7 +668,10 @@ def test_spec_aware_nonspec_materializes_state_indices(
 
 
 @torch.inference_mode()
+@pytest.mark.parametrize(("has_amx", "has_arm_bf16"), [(True, False), (False, True)])
 def test_spec_forward_prepares_native_conv_metadata(
+    has_amx: bool,
+    has_arm_bf16: bool,
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     block_table = torch.tensor([[0, 1], [4, 5]], dtype=torch.int32)
@@ -687,14 +692,17 @@ def test_spec_forward_prepares_native_conv_metadata(
     )
     forwarded_indices = None
     forwarded_counts = None
+    forwarded_is_weight_packed = None
 
     def causal_conv1d_update_cpu(**kwargs):
-        nonlocal forwarded_counts, forwarded_indices
+        nonlocal forwarded_counts, forwarded_indices, forwarded_is_weight_packed
         forwarded_indices = kwargs["conv_state_indices"]
         forwarded_counts = kwargs["num_accepted_tokens"]
+        forwarded_is_weight_packed = kwargs["is_weight_packed"]
         return kwargs["x"]
 
-    monkeypatch.setattr(torch.cpu, "_is_amx_tile_supported", lambda: True)
+    monkeypatch.setattr(torch.cpu, "_is_amx_tile_supported", lambda: has_amx)
+    monkeypatch.setattr(gdn_attention, "is_arm_bf16", lambda: has_arm_bf16)
     monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False)
     monkeypatch.setattr(
         gdn_attention.ops, "causal_conv1d_update_cpu", causal_conv1d_update_cpu
@@ -732,6 +740,7 @@ def test_spec_forward_prepares_native_conv_metadata(
         assert actual is not None
         assert actual.is_contiguous()
         assert actual.dtype == torch.int32
+        assert forwarded_is_weight_packed is True
         torch.testing.assert_close(actual, reference)
 
 
@@ -792,15 +801,15 @@ def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> No
 
 
 @pytest.mark.skipif(
-    not torch.cpu._is_amx_tile_supported(),
-    reason="requires AMX support",
+    not (torch.cpu._is_amx_tile_supported() or gdn_attention.is_arm_bf16()),
+    reason="native causal conv1d requires AMX or Arm BF16 support",
 )
 @torch.inference_mode()
 def test_causal_conv1d_update_cpu_accepts_wide_state() -> None:
     state_len = CONV_KERNEL - 1
     wide_state_len = state_len + 5
     batch_size = 3
-    is_vnni = True
+    is_weight_packed = True
     x, weight, bias = _conv_inputs(batch_size)
     conv_state_indices = torch.tensor([2, 0, 1], dtype=torch.int32)
 
@@ -813,7 +822,7 @@ def test_causal_conv1d_update_cpu_accepts_wide_state() -> None:
     wide_state[:, :, state_len:].fill_(7)
     wide_tail = wide_state[:, :, state_len:].clone()
 
-    conv_weight = _maybe_pack_conv_weight(weight, is_vnni)
+    conv_weight = _maybe_pack_conv_weight(weight, is_weight_packed)
     out_narrow = ops.causal_conv1d_update_cpu(
         x=x,
         conv_states=narrow_state,
@@ -821,7 +830,7 @@ def test_causal_conv1d_update_cpu_accepts_wide_state() -> None:
         bias=bias,
         silu_activation=True,
         conv_state_indices=conv_state_indices,
-        is_vnni=is_vnni,
+        is_weight_packed=is_weight_packed,
     )
     out_wide = ops.causal_conv1d_update_cpu(
         x=x,
@@ -830,7 +839,7 @@ def test_causal_conv1d_update_cpu_accepts_wide_state() -> None:
         bias=bias,
         silu_activation=True,
         conv_state_indices=conv_state_indices,
-        is_vnni=is_vnni,
+        is_weight_packed=is_weight_packed,
     )
 
     torch.testing.assert_close(out_wide, out_narrow, atol=1e-2, rtol=1e-2)
@@ -872,11 +881,18 @@ def _ref_causal_conv1d_update_cpu_multi(
 
 
 @pytest.mark.skipif(
-    not torch.cpu._is_amx_tile_supported(),
-    reason="requires AMX support",
+    not (torch.cpu._is_amx_tile_supported() or gdn_attention.is_arm_bf16()),
+    reason="native causal conv1d requires AMX or Arm BF16 support",
 )
 @pytest.mark.parametrize(
-    ("batch_size, seq_len, accepted_counts, has_bias, silu_activation, is_vnni"),
+    (
+        "batch_size, "
+        "seq_len, "
+        "accepted_counts, "
+        "has_bias, "
+        "silu_activation, "
+        "is_weight_packed"
+    ),
     [
         (1, 1, [1], False, False, False),
         (1, 1, [1], True, True, True),
@@ -893,7 +909,7 @@ def test_causal_conv1d_update_cpu_multi_token_matches_python(
     accepted_counts: list[int],
     has_bias: bool,
     silu_activation: bool,
-    is_vnni: bool,
+    is_weight_packed: bool,
 ) -> None:
     dim = 96
     state_len = seq_len + 2
@@ -911,7 +927,7 @@ def test_causal_conv1d_update_cpu_multi_token_matches_python(
     )
     conv_states = conv_states_ref.clone()
 
-    conv_weight = _maybe_pack_conv_weight(weight, is_vnni)
+    conv_weight = _maybe_pack_conv_weight(weight, is_weight_packed)
     out = ops.causal_conv1d_update_cpu(
         x=x,
         conv_states=conv_states,
@@ -919,7 +935,7 @@ def test_causal_conv1d_update_cpu_multi_token_matches_python(
         bias=bias,
         silu_activation=silu_activation,
         conv_state_indices=conv_state_indices,
-        is_vnni=is_vnni,
+        is_weight_packed=is_weight_packed,
         num_accepted_tokens=num_accepted_tokens,
     )
     ref_out = _ref_causal_conv1d_update_cpu_multi(
@@ -937,8 +953,8 @@ def test_causal_conv1d_update_cpu_multi_token_matches_python(
 
 
 @pytest.mark.skipif(
-    not torch.cpu._is_amx_tile_supported(),
-    reason="requires AMX support",
+    not (torch.cpu._is_amx_tile_supported() or gdn_attention.is_arm_bf16()),
+    reason="native causal conv1d requires AMX or Arm BF16 support",
 )
 @pytest.mark.parametrize("num_accepted", [0, 17])
 @torch.inference_mode()
@@ -961,14 +977,15 @@ def test_causal_conv1d_update_cpu_rejects_invalid_accepted_count(
             bias=None,
             silu_activation=True,
             conv_state_indices=torch.tensor([0], dtype=torch.int32),
-            is_vnni=False,
+            is_weight_packed=False,
             num_accepted_tokens=torch.tensor([num_accepted], dtype=torch.int32),
         )
 
 
 @pytest.mark.skipif(
-    not torch.cpu._is_avx512_bf16_supported(),
-    reason="causal_conv1d_fwd_cpu requires AVX-512BF16 (Intel Xeon or AMD EPYC)",
+    not (torch.cpu._is_avx512_bf16_supported() or gdn_attention.is_arm_bf16()),
+    reason="causal_conv1d_fwd_cpu requires AVX-512BF16"
+    " (Intel Xeon or AMD EPYC) or Arm BF16 support",
 )
 @pytest.mark.parametrize("total_tokens, split", TWO_CALL_SPLITS)
 @torch.inference_mode()
@@ -977,13 +994,13 @@ def test_causal_conv1d_fwd_cpu_two_call_split(total_tokens: int, split: int) ->
     matches the single-call result.
 
     Regression test for ``causal_conv1d_fwd_varlen_kernel_impl`` (``conv.cpp``)
-    ignoring the carried conv state on continued chunks. Runs on any
-    AVX-512BF16 CPU since conv.cpp uses VDPBF16PS, not AMX tiles.
+    ignoring the carried conv state on continued chunks. Runs on aarch64 cpus with
+    bf16 support or any AVX-512BF16 CPU since conv.cpp uses VDPBF16PS, not AMX tiles.
     """
     state_len = CONV_KERNEL - 1
     x, weight, bias = _conv_inputs(total_tokens)
 
-    def amx(x_seg, conv_states, has_init):
+    def native(x_seg, conv_states, has_init):
         seq = x_seg.shape[0]
         return ops.causal_conv1d_fwd_cpu(
             x=x_seg.transpose(0, 1),  # [dim, seq]; stride(-2)==1 (view of [seq,dim])
@@ -994,30 +1011,30 @@ def test_causal_conv1d_fwd_cpu_two_call_split(total_tokens: int, split: int) ->
             cache_indices=torch.tensor([0], dtype=torch.int32),
             has_initial_state=torch.tensor([has_init]),
             silu_activation=True,
-            is_vnni=False,
+            is_weight_packed=False,
         ).contiguous()
 
     # conv_state layout passed by the AMX branch: [num_slots, dim, state_len].
     cs_full = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype)
-    out_full = amx(x, cs_full, False)
+    out_full = native(x, cs_full, False)
 
     cs_split = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype)
-    out1 = amx(x[:split], cs_split, False)
-    out2 = amx(x[split:], cs_split, True)
+    out1 = native(x[:split], cs_split, False)
+    out2 = native(x[split:], cs_split, True)
     out_split = torch.cat([out1, out2], dim=1)
 
     torch.testing.assert_close(out_split, out_full, atol=1e-2, rtol=1e-2)
 
 
 @pytest.mark.skipif(
-    not torch.cpu._is_amx_tile_supported(),
-    reason="requires AMX support",
+    not (torch.cpu._is_amx_tile_supported() or gdn_attention.is_arm_bf16()),
+    reason="native causal conv1d requires AMX or Arm BF16 support",
 )
 @torch.inference_mode()
 def test_causal_conv1d_fwd_cpu_accepts_wide_state() -> None:
     state_len = CONV_KERNEL - 1
     wide_state_len = state_len + 5
-    is_vnni = True
+    is_weight_packed = True
     seq_lens = [CHUNK_SIZE - 1, CHUNK_SIZE + 5]
     total_tokens = sum(seq_lens)
     x, weight, bias = _conv_inputs(total_tokens)
@@ -1034,7 +1051,7 @@ def test_causal_conv1d_fwd_cpu_accepts_wide_state() -> None:
     wide_state[:, :, state_len:].fill_(7)
     wide_tail = wide_state[:, :, state_len:].clone()
 
-    conv_weight = _maybe_pack_conv_weight(weight, is_vnni)
+    conv_weight = _maybe_pack_conv_weight(weight, is_weight_packed)
     out_narrow = ops.causal_conv1d_fwd_cpu(
         x=x.transpose(0, 1),
         weight=conv_weight,
@@ -1044,7 +1061,7 @@ def test_causal_conv1d_fwd_cpu_accepts_wide_state() -> None:
         cache_indices=cache_indices,
         has_initial_state=has_initial_state,
         silu_activation=True,
-        is_vnni=is_vnni,
+        is_weight_packed=is_weight_packed,
     )
     out_wide = ops.causal_conv1d_fwd_cpu(
         x=x.transpose(0, 1),
@@ -1055,7 +1072,7 @@ def test_causal_conv1d_fwd_cpu_accepts_wide_state() -> None:
         cache_indices=cache_indices,
         has_initial_state=has_initial_state,
         silu_activation=True,
-        is_vnni=is_vnni,
+        is_weight_packed=is_weight_packed,
     )
 
     torch.testing.assert_close(out_wide, out_narrow, atol=1e-2, rtol=1e-2)
@@ -1100,15 +1117,23 @@ def test_batch_memcpy_cpu_fallback() -> None:
 
 
 # ---------------------------------------------------------------------------
-# 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.
+# C++ conv (conv.cpp) uses AArch64BF16 or VDPBF16PS, not AMX tiles, so it can
+# run on any AVX-512BF16 CPU; weight is 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]]
+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"):
@@ -1153,10 +1178,10 @@ def _run_prefill_torch(x, weight, bias, seq_lens):
     return out.transpose(0, 1).contiguous(), conv_states
 
 
-def _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False):
+def _run_prefill_cpp(x, weight, bias, seq_lens, is_weight_packed=False):
     num_seqs = len(seq_lens)
-    packed_w = ops.causal_conv1d_weight_pack(weight) if is_vnni else weight
-    if is_vnni:
+    packed_w = ops.causal_conv1d_weight_pack(weight) if is_weight_packed else weight
+    if is_weight_packed:
         # C++-branch layout: kv-cache "SD" [slots, state_len, dim] transposed to
         # [slots, dim, state_len] (a non-contiguous view).
         conv_state = torch.zeros(
@@ -1176,7 +1201,7 @@ def _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False):
         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,
+        is_weight_packed=is_weight_packed,
     )
     return out.transpose(0, 1).contiguous(), conv_state
 
@@ -1245,7 +1270,7 @@ def test_causal_conv1d_fwd_cpu_ds_adapter_varlen_mixed_initial_state(
         cache_indices=state_indices,
         has_initial_state=has_initial_state,
         silu_activation=True,
-        is_vnni=False,
+        is_weight_packed=False,
     )
     monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True)
     out_ds = gdn_attention._causal_conv1d_fwd_cpu(
@@ -1257,7 +1282,7 @@ def test_causal_conv1d_fwd_cpu_ds_adapter_varlen_mixed_initial_state(
         cache_indices=state_indices,
         has_initial_state=has_initial_state,
         silu_activation=True,
-        is_vnni=False,
+        is_weight_packed=False,
     )
 
     torch.testing.assert_close(out_ds, out_sd)
@@ -1302,7 +1327,7 @@ def test_causal_conv1d_update_cpu_ds_adapter_wide_multi_token_copyback(
         bias=None,
         silu_activation=False,
         conv_state_indices=state_indices,
-        is_vnni=False,
+        is_weight_packed=False,
         num_accepted_tokens=num_accepted_tokens,
     )
 
@@ -1321,7 +1346,8 @@ def test_causal_conv1d_update_cpu_ds_adapter_wide_multi_token_copyback(
 
 
 @pytest.mark.skipif(
-    not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
+    not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()),
+    reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support",
 )
 @pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS)
 @torch.inference_mode()
@@ -1335,7 +1361,8 @@ def test_conv_cpp_matches_torch(seq_lens):
 
 
 @pytest.mark.skipif(
-    not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
+    not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()),
+    reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support",
 )
 @pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS)
 @torch.inference_mode()
@@ -1354,89 +1381,114 @@ def test_conv_cpp_no_worse_than_torch_vs_fp32(seq_lens):
 
 
 @pytest.mark.skipif(
-    not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
+    not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()),
+    reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support",
 )
 @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."""
+def test_conv_cpp_weight_packed_matches_torch(seq_lens):
+    """The exact runtime prefill sequence (packed weight + SD-layout
+    conv_state view + is_weight_packed=True) must match the torch fallback. Validates
+    the native packing + layout handoff."""
     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)
+    out_weight_packed, _ = _run_prefill_cpp(
+        x, weight, bias, seq_lens, is_weight_packed=True
+    )
+    torch.testing.assert_close(out_weight_packed, out_torch, atol=1e-2, rtol=1e-2)
 
 
 @pytest.mark.skipif(
-    not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
+    not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()),
+    reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support",
 )
 @pytest.mark.parametrize("batch", DECODE_BATCH_SIZES)
+@pytest.mark.parametrize(
+    ("dim", "has_bias", "silu_activation", "is_weight_packed"),
+    [
+        (32, False, False, False),
+        (32, True, True, True),
+        (96, False, True, True),
+        (96, True, False, False),
+        (8192, False, True, False),
+        (8192, True, True, True),
+    ],
+)
 @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)."""
+def test_conv_update_cpp_matches_torch(
+    batch, dim, has_bias, silu_activation, is_weight_packed
+):
+    """Decode preserves indexed, padded state across successive signed inputs."""
     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
-    )
+    torch.manual_seed(0)
+    weight = torch.randn(dim, CONV_KERNEL, dtype=torch.bfloat16)
+    bias = torch.randn(dim, dtype=torch.bfloat16) if has_bias else None
+    conv_weight = _maybe_pack_conv_weight(weight, is_weight_packed)
+    indices = torch.arange(batch, 0, -1, dtype=torch.int32)
+    storage_ref = torch.randn(batch + 2, _STATE_LEN + 2, dim, dtype=torch.bfloat16)
+    storage_cpp = storage_ref.clone()
+    cs_cpp = storage_cpp[:, :_STATE_LEN].transpose(1, 2)
 
-    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)
+    for _ in range(3):
+        x = torch.randn(batch, dim, dtype=torch.bfloat16)
+        cs_torch = storage_ref[indices, :_STATE_LEN].transpose(1, 2).contiguous()
+        out_torch = causal_conv1d_update_torch(
+            x=x.unsqueeze(-1),
+            conv_state=cs_torch,
+            weight=weight,
+            bias=bias,
+            activation="silu" if silu_activation else None,
+        ).squeeze(-1)
+        storage_ref[indices, :_STATE_LEN] = cs_torch.transpose(1, 2)
+        out_cpp = ops.causal_conv1d_update_cpu(
+            x=x,
+            conv_states=cs_cpp,
+            weight=conv_weight,
+            bias=bias,
+            silu_activation=silu_activation,
+            conv_state_indices=indices,
+            is_weight_packed=is_weight_packed,
+        )
+        torch.testing.assert_close(out_cpp, out_torch, atol=1e-2, rtol=1e-2)
+        torch.testing.assert_close(storage_cpp, storage_ref, atol=0, rtol=0)
 
 
 @pytest.mark.skipif(
-    not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
+    not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()),
+    reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support",
 )
 @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.
+    with packed weight + is_weight_packed=True must equal unpacked weight
+    + is_weight_packed=False. Guards the spec-decode contract: the runtime 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)
+    out_unpacked, state_unpacked = _run_prefill_cpp(
+        x, weight, bias, seq_lens, is_weight_packed=False
+    )
+    out_packed, state_packed = _run_prefill_cpp(
+        x, weight, bias, seq_lens, is_weight_packed=True
+    )
+    torch.testing.assert_close(out_packed, out_unpacked, atol=0, rtol=0)
+    torch.testing.assert_close(state_packed, state_unpacked, atol=0, rtol=0)
 
 
 @pytest.mark.skipif(
-    not _HAS_AVX512_BF16, reason="C++ causal_conv1d requires AVX-512BF16"
+    not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()),
+    reason="C++ causal_conv1d requires AVX-512BF16 or AArch64 BF16",
 )
 @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
+    """Spec-decode correctness: AVX-512BF16 and AArch64 BF16 CPUs 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
@@ -1460,7 +1512,7 @@ def test_spec_decode_unpacked_conv_weight_stash():
     # 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"
+        "on an AArch64-BF16 or AVX-512BF16 CPU"
     )
     stash = conv._cpu_unpacked_conv_weight
     torch.testing.assert_close(stash, orig_2d, atol=0, rtol=0)
@@ -1469,10 +1521,10 @@ def test_spec_decode_unpacked_conv_weight_stash():
     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
+    packed_weight = conv.weight  # [dim, 1, width], 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).
+    #    nonspec C++ conv with the packed weight (is_weight_packed=True).
     seq_lens = [7, 65]
     total = sum(seq_lens)
     x = tensor_cache(total * CONV_DIM, torch.bfloat16).view(total, CONV_DIM)
@@ -1494,7 +1546,7 @@ def test_spec_decode_unpacked_conv_weight_stash():
             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,
+            is_weight_packed=True,
         )
         .transpose(0, 1)
         .contiguous()

tests/model_executor/test_cpu_unquantized_gemm_dispatch.py

@@ -6,10 +6,61 @@ import pytest
 import torch
 
 from vllm.model_executor.layers import utils
-from vllm.platforms import current_platform
+from vllm.platforms import CpuArchEnum, current_platform
 from vllm.utils.torch_utils import set_default_torch_dtype
 
 
+def test_dispatch_prepacks_arm_bf16_causal_conv(monkeypatch):
+    monkeypatch.setattr(
+        current_platform, "get_cpu_architecture", lambda: CpuArchEnum.ARM
+    )
+    monkeypatch.setattr(torch.cpu, "get_capabilities", lambda: {"bf16": True})
+    monkeypatch.setattr(torch.cpu, "_is_avx512_bf16_supported", lambda: False)
+
+    packed = torch.randn(32, 4, dtype=torch.bfloat16)
+    pack_inputs = []
+
+    def pack(weight):
+        pack_inputs.append(weight.clone())
+        return packed
+
+    monkeypatch.setattr(utils.ops, "causal_conv1d_weight_pack", pack)
+    original = torch.randn(32, 1, 4, dtype=torch.bfloat16)
+    layer = torch.nn.Module()
+    layer.weight = torch.nn.Parameter(original.clone(), requires_grad=False)
+
+    utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
+
+    assert len(pack_inputs) == 1
+    # Check causal_conv1d_weight_pack received the expected weights
+    torch.testing.assert_close(pack_inputs[0], original.view(32, 4))
+    # Check we have correctly stashed the unpacked weights
+    torch.testing.assert_close(layer._cpu_unpacked_conv_weight, original.view(32, 4))
+    # Check we have correctly stored the packed weights
+    assert layer.weight.data_ptr() == packed.data_ptr()
+
+
+def test_dispatch_does_not_pack_3d_expert_weight(monkeypatch):
+    monkeypatch.setattr(
+        current_platform, "get_cpu_architecture", lambda: CpuArchEnum.ARM
+    )
+    monkeypatch.setattr(torch.cpu, "get_capabilities", lambda: {"bf16": True})
+    monkeypatch.setattr(torch.cpu, "_is_avx512_bf16_supported", lambda: False)
+    pack_calls = []
+    monkeypatch.setattr(
+        utils.ops,
+        "causal_conv1d_weight_pack",
+        lambda weight: pack_calls.append(weight),
+    )
+
+    layer = torch.nn.Module()
+    layer.weight = torch.nn.Parameter(
+        torch.randn(8, 32, 4, dtype=torch.bfloat16), requires_grad=False
+    )
+    utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
+    assert pack_calls == []
+
+
 @pytest.fixture(scope="module")
 def _mock_zentorch_linear_unary():
     """Register a mock zentorch_linear_unary op when zentorch is not installed.

vllm/_custom_ops.py

@@ -3666,7 +3666,7 @@ def causal_conv1d_fwd_cpu(
     cache_indices: torch.Tensor | None,
     has_initial_state: torch.Tensor | None,
     silu_activation: bool,
-    is_vnni: bool,
+    is_weight_packed: bool,
 ) -> torch.Tensor:
     return torch.ops._C.causal_conv1d_fwd_cpu(
         x,
@@ -3678,7 +3678,7 @@ def causal_conv1d_fwd_cpu(
         has_initial_state,
         silu_activation,
         -1,
-        is_vnni,
+        is_weight_packed,
     )
 
 
@@ -3689,7 +3689,7 @@ def causal_conv1d_update_cpu(
     bias: torch.Tensor | None,
     silu_activation: bool,
     conv_state_indices: torch.Tensor | None,
-    is_vnni: bool,
+    is_weight_packed: bool,
     num_accepted_tokens: torch.Tensor | None = None,
 ) -> torch.Tensor:
     return torch.ops._C.causal_conv1d_update_cpu(
@@ -3701,7 +3701,7 @@ def causal_conv1d_update_cpu(
         num_accepted_tokens,
         conv_state_indices,
         -1,
-        is_vnni,
+        is_weight_packed,
     )
 
 

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

@@ -14,7 +14,6 @@ from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
 )
 from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (
     causal_conv1d_update_cpu,
-    causal_conv1d_update_torch,
 )
 from vllm.platforms import CpuArchEnum, current_platform
 from vllm.utils.torch_utils import (
@@ -56,7 +55,7 @@ def _causal_conv1d_fwd_cpu(
     cache_indices: torch.Tensor,
     has_initial_state: torch.Tensor,
     silu_activation: bool,
-    is_vnni: bool,
+    is_weight_packed: bool,
 ) -> torch.Tensor:
     if not is_conv_state_dim_first():
         return ops.causal_conv1d_fwd_cpu(
@@ -68,7 +67,7 @@ def _causal_conv1d_fwd_cpu(
             cache_indices=cache_indices,
             has_initial_state=has_initial_state,
             silu_activation=silu_activation,
-            is_vnni=is_vnni,
+            is_weight_packed=is_weight_packed,
         )
 
     scratch, identity_indices = _stage_ds_conv_states(conv_states, cache_indices)
@@ -81,7 +80,7 @@ def _causal_conv1d_fwd_cpu(
         cache_indices=identity_indices,
         has_initial_state=has_initial_state,
         silu_activation=silu_activation,
-        is_vnni=is_vnni,
+        is_weight_packed=is_weight_packed,
     )
     _scatter_ds_conv_states(conv_states, cache_indices, scratch)
     return out
@@ -95,7 +94,7 @@ def _causal_conv1d_update_cpu(
     bias: torch.Tensor | None,
     silu_activation: bool,
     conv_state_indices: torch.Tensor,
-    is_vnni: bool,
+    is_weight_packed: bool,
     num_accepted_tokens: torch.Tensor | None = None,
 ) -> torch.Tensor:
     if not is_conv_state_dim_first():
@@ -106,7 +105,7 @@ def _causal_conv1d_update_cpu(
             bias=bias,
             silu_activation=silu_activation,
             conv_state_indices=conv_state_indices,
-            is_vnni=is_vnni,
+            is_weight_packed=is_weight_packed,
             num_accepted_tokens=num_accepted_tokens,
         )
 
@@ -118,13 +117,20 @@ def _causal_conv1d_update_cpu(
         bias=bias,
         silu_activation=silu_activation,
         conv_state_indices=identity_indices,
-        is_vnni=is_vnni,
+        is_weight_packed=is_weight_packed,
         num_accepted_tokens=num_accepted_tokens,
     )
     _scatter_ds_conv_states(conv_states, conv_state_indices, scratch)
     return out
 
 
+def is_arm_bf16() -> bool:
+    return (
+        current_platform.get_cpu_architecture() == CpuArchEnum.ARM
+        and torch.cpu.get_capabilities().get("bf16", False)
+    )
+
+
 def cpu_gdn_attention_core(
     mixed_qkv: torch.Tensor,
     b: torch.Tensor,
@@ -201,9 +207,10 @@ def _cpu_gdn_attention_nonspec(
     assert state_indices_tensor is not None
     assert query_start_loc is not None
 
-    # 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()
+    # C++ conv (conv.cpp) runs on aarch64 with bf16 support and on x86 uses VDPBF16PS,
+    # not AMX tiles, so it runs on any AVX-512BF16 CPU. The weights are packed on this
+    # same predicate at load time.
+    use_cpp_conv = torch.cpu._is_avx512_bf16_supported() or is_arm_bf16()
 
     conv_state = layer.kv_cache[0]
     if use_cpp_conv:
@@ -244,28 +251,17 @@ def _cpu_gdn_attention_nonspec(
                 bias=layer.conv1d.bias,
                 silu_activation=(layer.activation == "silu"),
                 conv_state_indices=decode_state_indices,
-                is_vnni=True,
+                is_weight_packed=True,
             )
         else:
-            if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
-                decode_conv_state = conv_state[decode_state_indices].contiguous()
-                decode_mixed_qkv = causal_conv1d_update_torch(
-                    x=decode_mixed_qkv.unsqueeze(-1),
-                    conv_state=decode_conv_state,
-                    weight=conv_weights,
-                    bias=layer.conv1d.bias,
-                    activation=layer.activation,
-                ).squeeze(-1)
-                conv_state[decode_state_indices] = decode_conv_state
-            else:
-                decode_mixed_qkv = causal_conv1d_update_cpu(
-                    x=decode_mixed_qkv,
-                    conv_state=conv_state,
-                    weight=conv_weights,
-                    bias=layer.conv1d.bias,
-                    activation=layer.activation,
-                    conv_state_indices=decode_state_indices,
-                )
+            decode_mixed_qkv = causal_conv1d_update_cpu(
+                x=decode_mixed_qkv,
+                conv_state=conv_state,
+                weight=conv_weights,
+                bias=layer.conv1d.bias,
+                activation=layer.activation,
+                conv_state_indices=decode_state_indices,
+            )
 
         query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)
 
@@ -315,7 +311,7 @@ def _cpu_gdn_attention_nonspec(
                 cache_indices=prefill_state_indices,
                 has_initial_state=prefill_has_initial_state,
                 silu_activation=layer.activation == "silu",
-                is_vnni=True,
+                is_weight_packed=True,
             ).transpose(0, 1)
         else:
             prefill_mixed_qkv = causal_conv1d_torch(
@@ -383,9 +379,9 @@ def _ssm_state_view(layer) -> torch.Tensor:
 def _unpacked_conv_weight(layer) -> torch.Tensor:
     """Return the plain (dim, width) conv weight.
 
-    On AMX the conv1d weight is VNNI-packed in place at load time (only usable
-    by the AMX C++ kernel), so the torch spec-decode path relies on the
-    un-packed copy stashed by ``dispatch_cpu_unquantized_gemm``.
+    On AMX and aarch64+bf16 the conv1d weight is packed in place at load time,
+    so the torch spec-decode path relies on the un-packed copy stashed by
+    ``dispatch_cpu_unquantized_gemm``.
     """
     w = getattr(layer.conv1d, "_cpu_unpacked_conv_weight", None)
     if w is not None:
@@ -505,13 +501,15 @@ def _spec_forward(
     seq_starts = spec_qsl_cpu[:-1]
     seq_lens = spec_qsl_cpu[1:] - spec_qsl_cpu[:-1]
 
+    is_amx = torch.cpu._is_amx_tile_supported()
+
     # ---- 1. Convolution (per-sequence rolling buffer) ----
     dim = mixed_qkv_spec.size(-1)
     bias = layer.conv1d.bias
     silu = layer.activation == "silu"
 
     can_use_native_conv = (
-        torch.cpu._is_amx_tile_supported()
+        (is_amx or is_arm_bf16())
         and width == 4
         and num_spec_decodes > 0
         and bool(torch.all(seq_lens == seq_lens[0]).item())
@@ -528,7 +526,7 @@ def _spec_forward(
             conv_state_indices=spec_state_indices[:num_spec_decodes, 0]
             .to("cpu", torch.int32)
             .contiguous(),
-            is_vnni=True,
+            is_weight_packed=True,
             num_accepted_tokens=num_accepted[:num_spec_decodes].to("cpu", torch.int32),
         ).view_as(mixed_qkv_spec)
     else:
@@ -603,9 +601,9 @@ def _spec_aware_nonspec(
     assert query_start_loc is not None
     state_indices_tensor = state_indices_tensor.contiguous()
 
-    is_amx = torch.cpu._is_amx_tile_supported()
+    can_use_native_conv = torch.cpu._is_amx_tile_supported() or is_arm_bf16()
 
-    if not is_amx:
+    if not can_use_native_conv:
         conv_weights = _unpacked_conv_weight(layer)
 
     num_decodes = attn_metadata_i.num_decodes
@@ -618,7 +616,8 @@ def _spec_aware_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 can_use_native_conv:
             decode_mixed_qkv = _causal_conv1d_update_cpu(
                 x=decode_mixed_qkv,
                 conv_states=conv_buf,
@@ -626,30 +625,19 @@ def _spec_aware_nonspec(
                 bias=layer.conv1d.bias,
                 silu_activation=layer.activation == "silu",
                 conv_state_indices=decode_state_indices,
-                is_vnni=True,
+                is_weight_packed=True,
             )
         else:
             # Only the first ``width-1`` columns hold the real conv state.
             conv_state_view = conv_buf[:, :, : width - 1]
-            if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
-                decode_conv_state = conv_state_view[decode_state_indices].contiguous()
-                decode_mixed_qkv = causal_conv1d_update_torch(
-                    x=decode_mixed_qkv.unsqueeze(-1),
-                    conv_state=decode_conv_state,
-                    weight=conv_weights,
-                    bias=layer.conv1d.bias,
-                    activation=layer.activation,
-                ).squeeze(-1)
-                conv_state_view[decode_state_indices] = decode_conv_state
-            else:
-                decode_mixed_qkv = causal_conv1d_update_cpu(
-                    x=decode_mixed_qkv,
-                    conv_state=conv_state_view,
-                    weight=conv_weights,
-                    bias=layer.conv1d.bias,
-                    activation=layer.activation,
-                    conv_state_indices=decode_state_indices,
-                )
+            decode_mixed_qkv = causal_conv1d_update_cpu(
+                x=decode_mixed_qkv,
+                conv_state=conv_state_view,
+                weight=conv_weights,
+                bias=layer.conv1d.bias,
+                activation=layer.activation,
+                conv_state_indices=decode_state_indices,
+            )
 
         query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)
         attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu(
@@ -685,7 +673,8 @@ def _spec_aware_nonspec(
         prefill_has_initial_state = has_initial_state[
             num_decodes : num_decodes + num_prefills
         ]
-        if is_amx:
+
+        if can_use_native_conv:
             prefill_mixed_qkv = _causal_conv1d_fwd_cpu(
                 x=prefill_mixed_qkv.transpose(0, 1),
                 weight=layer.conv1d.weight,
@@ -695,7 +684,7 @@ def _spec_aware_nonspec(
                 cache_indices=prefill_state_indices,
                 has_initial_state=prefill_has_initial_state,
                 silu_activation=layer.activation == "silu",
-                is_vnni=True,
+                is_weight_packed=True,
             ).transpose(0, 1)
         else:
             prefill_mixed_qkv = causal_conv1d_torch(
@@ -757,8 +746,9 @@ def _spec_aware_nonspec_subset(
     prefill_state_indices = prefill_state_indices.contiguous()
 
     is_amx = torch.cpu._is_amx_tile_supported()
+    can_use_native_conv = is_amx or is_arm_bf16()
 
-    if is_amx:
+    if can_use_native_conv:
         conv_out = _causal_conv1d_fwd_cpu(
             x=mixed_qkv.transpose(0, 1),
             weight=layer.conv1d.weight,
@@ -768,7 +758,7 @@ def _spec_aware_nonspec_subset(
             cache_indices=prefill_state_indices,
             has_initial_state=has_initial_state,
             silu_activation=layer.activation == "silu",
-            is_vnni=True,
+            is_weight_packed=True,
         ).transpose(0, 1)
     else:
         conv_weights = _unpacked_conv_weight(layer)

vllm/model_executor/layers/utils.py

@@ -475,13 +475,27 @@ def dispatch_cpu_unquantized_gemm(
     # Skip CPU GEMM dispatch for non-2D weights (e.g. MoE 3D expert weights).
     # These layers are handled by their own specialized methods.
     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
-        # 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"
+        # This is not a linear layer.
+        # For now it should be a causal_conv1d op or MoE 3D expert weights.
+        # Causal-conv weights use [dim, 1, width].
+        # The C++ causal_conv1d kernels support aarch64 with bf16 and
+        # x86 CPUs that use VDPBF16PS (no AMX tiles). So the
+        # weight prepack applies to aarch64 with bf16 and any AVX-512BF16 CPU,
+        # not just AMX (e.g. AMD Zen5/Turin).
+        is_causal_conv1d_weight = (
+            layer.weight.ndim == 3
+            and layer.weight.size(1) == 1
+            and layer.weight.size(2) == 4
+            and layer.weight.dtype == torch.bfloat16
+        )
+        is_arm_bf16 = (
+            current_platform.get_cpu_architecture() == CpuArchEnum.ARM
+            and torch.cpu.get_capabilities().get("bf16", False)
+        )
+        if (
+            (is_arm_bf16 or torch.cpu._is_avx512_bf16_supported())
+            and is_causal_conv1d_weight
+            and hasattr(ops, "causal_conv1d_weight_pack")
         ):
             # prepack conv weight
             unpacked = (
@@ -493,7 +507,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 C++ kernel) can use it.
+            # and unsupported-layout fallbacks can still use ordinary weights.
             layer._cpu_unpacked_conv_weight = unpacked
             layer.weight.data = ops.causal_conv1d_weight_pack(unpacked)
         return