@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 93dd80f636

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

almayne committed Sep 30, 2026 at 13:23 UTC 93dd80f63639453591e68891614729fb0bd44855

9 files changed +461 -206

cmake/cpu_extension.cmake

@@ -524,6 +524,7 @@ if (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND)

524 "csrc/cpu/cpu_tanhf_neon.hpp"

525 ${VLLM_EXT_SRC})

526 if (ARM_BF16_FOUND)

527 + set(VLLM_EXT_SRC "csrc/cpu/sgl-kernels/conv.cpp" ${VLLM_EXT_SRC})

528 if (ARM_I8MM_FOUND)

529 set(VLLM_EXT_SRC "csrc/cpu/cpu_fused_moe_int8.cpp" ${VLLM_EXT_SRC})

530 endif()

csrc/cpu/sgl-kernels/conv.cpp

+19 -7

@@ -66,6 +66,10 @@ struct tinygemm_kernel {

66 }

67 };

69 +#if defined(ARM_BF16_SUPPORT)

70 +#include "conv_arm.h"

71 +#endif

72 +

73 #if defined(CPU_CAPABILITY_AVX512)

74 template <int K, int BLOCK_N, bool has_bias, bool has_silu>

75 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) {

547

548 int64_t dim = weight.size(0);

549 int64_t width = weight.size(1);

546 - constexpr int64_t BLOCK_N = block_size_n();

550 + constexpr int64_t BLOCK_N =

551 +#if defined(ARM_BF16_SUPPORT)

552 + 8;// 8 channel blocks for microtile

553 +#else

554 + block_size_n();

555 +#endif

556 TORCH_CHECK(width == 4, "causal_conv1d_weight_pack: support only width of 4");

548 - TORCH_CHECK(dim % BLOCK_N == 0, "causal_conv1d_weight_pack: invalid dim size ", dim);

557 + TORCH_CHECK(

558 + dim % block_size_n() == 0,

559 + "causal_conv1d_weight_pack: invalid dim size ",

560 + dim);

561

562 const int64_t N = dim, K2 = width >> 1;

563 const int64_t NB = div_up(N, BLOCK_N);

564

565 auto packed_weight = at::empty_like(weight);

566 AT_DISPATCH_REDUCED_FLOATING_TYPES(weight.scalar_type(), "causal_conv1d_fwd_kernel_impl", [&] {

555 - // cast to float32 as vnni size is 2

567 + // Copy each pair as an opaque 32-bit word. This is not a bf16->fp32 conversion.

568 const float* w_data = reinterpret_cast<float*>(weight.data_ptr<scalar_t>());

569 float* packed_data = reinterpret_cast<float*>(packed_weight.data_ptr<scalar_t>());

570

@@ -642,9 +654,9 @@ at::Tensor causal_conv1d_fwd_cpu(

654 const std::optional<at::Tensor>& has_initial_state,

655 bool silu_activation,

656 int64_t pad_slot_id,

645 - bool is_vnni) {

657 + bool is_weight_packed) {

658 CHECK_CONTIGUOUS(weight);

647 - auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);

659 + auto packed_w = is_weight_packed ? weight : causal_conv1d_weight_pack(weight);

660

661 const bool is_var_seqlen = query_start_loc.has_value();

662 const int64_t input_ndim = is_var_seqlen ? 2 : 3;

@@ -755,10 +767,10 @@ at::Tensor causal_conv1d_update_cpu(

767 const std::optional<at::Tensor>& num_accepted_tokens,

768 const std::optional<at::Tensor>& conv_state_indices,

769 int64_t pad_slot_id,

758 - bool is_vnni) {

770 + bool is_weight_packed) {

771 CHECK_CONTIGUOUS(x);

772 CHECK_CONTIGUOUS(weight);

761 - auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);

773 + auto packed_w = is_weight_packed ? weight : causal_conv1d_weight_pack(weight);

774

775 TORCH_CHECK(

776 x.dim() == 2 || x.dim() == 3,

csrc/cpu/sgl-kernels/conv_arm.h new

+126

@@ -0,0 +1,126 @@

1 +// SPDX-License-Identifier: Apache-2.0

2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +#pragma once

5 +

6 +#if defined(ARM_BF16_SUPPORT)

7 +

8 +template <int K, int BLOCK_N, bool has_bias, bool has_silu>

9 +struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> {

10 + static inline bfloat16x8_t load_bf16x8(const at::BFloat16* ptr) {

11 + return vld1q_bf16(reinterpret_cast<const bfloat16_t*>(ptr));

12 + }

13 +

14 + // Input layout: A[m][channel]

15 + // Weight layout: B[N/MICRO_N, K/2, MICRO_N, 2]

16 + // The packed weights are reused across all M positions.

17 + static inline void tinygemm_neon_8(

18 + const at::BFloat16* __restrict__ A, const at::BFloat16* __restrict__ B,

19 + at::BFloat16* __restrict__ C, const at::BFloat16* __restrict__ bias,

20 + const at::BFloat16* __restrict__ conv_states, bool has_initial_state,

21 + int64_t M, int64_t lda, bool is_first_token) {

22 + bfloat16x8_t va[4];

23 + // Load 4 pairs of kernel weights into 4 registers

24 + const bfloat16x8_t vb0 = load_bf16x8(B); // channels 0-3, taps 0-1

25 + const bfloat16x8_t vb1 = load_bf16x8(B + 8); // channels 4-7, taps 0-1

26 + const bfloat16x8_t vb2 = load_bf16x8(B + 16); // channels 0-3, taps 2-3

27 + const bfloat16x8_t vb3 = load_bf16x8(B + 24); // channels 4-7, taps 2-3

28 +

29 + // Load the input data. A single register holds the values for a

30 + // single token across all 8 channels for this microtile.

31 + auto load_a = [&](int64_t m) {

32 + if (m < 0 && is_first_token) {

33 + if (!has_initial_state) {

34 + return vreinterpretq_bf16_u16(vdupq_n_u16(0));

35 + }

36 + return load_bf16x8(conv_states + (m + K - 1) * lda);

37 + }

38 + return load_bf16x8(A + m * lda);

39 + };

40 +

41 + // Load the previous inputs from the cached state if available.

42 + va[1] = load_a(-3);

43 + va[2] = load_a(-2);

44 + va[3] = load_a(-1);

45 +

46 + // Iterate over the sequence block, updating the previous conv inputs

47 + // and reading the next token's input data (8 channels).

48 + for (int64_t m = 0; m < M; ++m) {

49 + va[0] = va[1];

50 + va[1] = va[2];

51 + va[2] = va[3];

52 + va[3] = load_a(m);

53 +

54 + float32x4_t vc0;

55 + float32x4_t vc1;

56 + if constexpr (has_bias) {

57 + const bfloat16x8_t bias_vec = load_bf16x8(bias);

58 + vc0 = vcvtq_low_f32_bf16(bias_vec);

59 + vc1 = vcvtq_high_f32_bf16(bias_vec);

60 + } else {

61 + vc0 = vdupq_n_f32(0.0f);

62 + vc1 = vdupq_n_f32(0.0f);

63 + }

64 +

65 + // Use the zip operation to match the input layout to the weights layout

66 + bfloat16x8_t va_pair =

67 + vreinterpretq_bf16_u16( // channels 0-3, inputs for taps 0-1

68 + vzip1q_u16(vreinterpretq_u16_bf16(va[0]),

69 + vreinterpretq_u16_bf16(va[1])));

70 + vc0 = vbfdotq_f32(vc0, va_pair,

71 + vb0); // 4 independent length-two dot products

72 + va_pair = vreinterpretq_bf16_u16( // channels 4-7, inputs for taps 0-1

73 + vzip2q_u16(vreinterpretq_u16_bf16(va[0]),

74 + vreinterpretq_u16_bf16(va[1])));

75 + vc1 = vbfdotq_f32(vc1, va_pair, vb1);

76 + va_pair = vreinterpretq_bf16_u16( // channels 0-3, inputs for taps 2-3

77 + vzip1q_u16(vreinterpretq_u16_bf16(va[2]),

78 + vreinterpretq_u16_bf16(va[3])));

79 + // Add to the previous result for the same channels

80 + vc0 = vbfdotq_f32(vc0, va_pair, vb2);

81 + va_pair = vreinterpretq_bf16_u16( // channels 4-7, inputs for taps 2-3

82 + vzip2q_u16(vreinterpretq_u16_bf16(va[2]),

83 + vreinterpretq_u16_bf16(va[3])));

84 + vc1 = vbfdotq_f32(vc1, va_pair, vb3);

85 +

86 + using fVec = at::vec::Vectorized<float>;

87 + fVec out0(vc0);

88 + fVec out1(vc1);

89 + if constexpr (has_silu) {

90 + const fVec one(1.0f);

91 + out0 = out0 / (one + out0.neg().fexp_u20());

92 + out1 = out1 / (one + out1.neg().fexp_u20());

93 + }

94 +

95 + // Convert the 2 x 4 fp32 results to 1 x 8 bf16

96 + const bfloat16x8_t out =

97 + vcombine_bf16(vcvt_bf16_f32(static_cast<float32x4_t>(out0)),

98 + vcvt_bf16_f32(static_cast<float32x4_t>(out1)));

99 + vst1q_bf16(reinterpret_cast<bfloat16_t*>(C + m * lda), out);

100 + }

101 + }

102 +

103 + static inline void apply(const at::BFloat16* __restrict__ A,

104 + const at::BFloat16* __restrict__ B,

105 + at::BFloat16* __restrict__ C,

106 + const at::BFloat16* __restrict__ bias,

107 + const at::BFloat16* __restrict__ conv_states,

108 + bool has_initial_state, int64_t M, int64_t lda,

109 + bool is_first_token) {

110 + static_assert(K == 4);

111 + constexpr int64_t MICRO_N = 8;

112 + static_assert(BLOCK_N % MICRO_N == 0);

113 +

114 + // BLOCK_N is always 32 or 64. M is the block sequence length.

115 + // Input layout: A[m][channel]

116 + // Weight layout: B[N/MICRO_N, K/2, MICRO_N, 2]

117 + // Use an 8 channel microtile and iterate over the 8 channel blocks.

118 + for (int64_t n = 0; n < BLOCK_N; n += MICRO_N) {

119 + tinygemm_neon_8(A + n, B + n * K, C + n, has_bias ? bias + n : nullptr,

120 + conv_states ? conv_states + n : nullptr,

121 + has_initial_state, M, lda, is_first_token);

122 + }

123 + }

124 +};

125 +

126 +#endif

csrc/cpu/torch_bindings.cpp

+29 -20

@@ -281,14 +281,14 @@ at::Tensor causal_conv1d_fwd_cpu(

281 const std::optional<at::Tensor>& query_start_loc,

282 const std::optional<at::Tensor>& cache_indices,

283 const std::optional<at::Tensor>& has_initial_state, bool silu_activation,

284 - int64_t pad_slot_id, bool is_vnni);

284 + int64_t pad_slot_id, bool is_weight_packed);

285

286 at::Tensor causal_conv1d_update_cpu(

287 const at::Tensor& x, const at::Tensor& conv_states,

288 const at::Tensor& weight, const std::optional<at::Tensor>& bias,

289 bool silu_activation, const std::optional<at::Tensor>& num_accepted_tokens,

290 const std::optional<at::Tensor>& conv_state_indices, int64_t pad_slot_id,

291 - bool is_vnni);

291 + bool is_weight_packed);

292

293 void activation_lut_bf16(torch::Tensor& out, torch::Tensor& input,

294 const std::string& activation);

@@ -517,6 +517,33 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

517

518 #endif // (defined(__aarch64__) && !defined(__APPLE__))

519

520 +#if (defined(__aarch64__) && defined(ARM_BF16_SUPPORT) && \

521 + !defined(__APPLE__)) || \

522 + (defined(__AVX512BF16__) && defined(__AVX512F__) && \

523 + defined(__AVX512VNNI__))

524 +

525 + // Adapted from sglang: casual_conv1d kernels

526 + ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");

527 + ops.impl("causal_conv1d_weight_pack", torch::kCPU,

528 + &causal_conv1d_weight_pack);

529 +

530 + ops.def(

531 + "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "

532 + "conv_states, Tensor? query_start_loc,"

533 + "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "

534 + "int pad_slot_id, bool is_weight_packed) -> "

535 + "Tensor");

536 + ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);

537 +

538 + ops.def(

539 + "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "

540 + "weight, Tensor? bias, bool silu_activation,"

541 + "Tensor? num_accepted_tokens, Tensor? conv_state_indices, int "

542 + "pad_slot_id, "

543 + "bool is_weight_packed) -> Tensor");

544 + ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);

545 +#endif

546 +

547 // Layernorm

548 // Apply Root Mean Square (RMS) Normalization to the input tensor.

549 ops.def(

@@ -665,24 +692,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

692 "-> (Tensor, Tensor)");

693 ops.impl("quantize_fp8e4m3_vec", torch::kCPU, &quantize_fp8e4m3_vec);

694

668 - ops.def("causal_conv1d_weight_pack(Tensor weight) -> Tensor");

669 - ops.impl("causal_conv1d_weight_pack", torch::kCPU,

670 - &causal_conv1d_weight_pack);

671 - ops.def(

672 - "causal_conv1d_fwd_cpu(Tensor x, Tensor weight, Tensor? bias, Tensor? "

673 - "conv_states, Tensor? query_start_loc,"

674 - "Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation, "

675 - "int pad_slot_id, bool is_vnni) -> "

676 - "Tensor");

677 - ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu);

678 - ops.def(

679 - "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "

680 - "weight, Tensor? bias, bool silu_activation,"

681 - "Tensor? num_accepted_tokens, Tensor? conv_state_indices, int "

682 - "pad_slot_id, "

683 - "bool is_vnni) -> Tensor");

684 - ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);

685 -

695 // Adapted from sglang: MLA CPU kernels (AMX-only, DeepSeek V2/V3/R1)

696 ops.def(

697 "decode_attention_cpu(Tensor query, Tensor k_buffer, Tensor v_buffer, "

tests/kernels/mamba/cpu/test_cpu_gdn_ops.py

+155 -103

@@ -494,8 +494,10 @@ def _sd_conv_states(

494 return storage.transpose(1, 2)

495

496

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

497 +def _maybe_pack_conv_weight(

498 + weight: torch.Tensor, is_weight_packed: bool

499 +) -> torch.Tensor:

500 + return ops.causal_conv1d_weight_pack(weight) if is_weight_packed else weight

501

502

503 @torch.inference_mode()

@@ -523,7 +525,7 @@ def test_causal_conv1d_update_cpu_preserves_sd_state(

525 bias=None,

526 silu_activation=False,

527 conv_state_indices=state_indices,

526 - is_vnni=False,

528 + is_weight_packed=False,

529 )

530

531 assert out is x

@@ -666,7 +668,10 @@ def test_spec_aware_nonspec_materializes_state_indices(

668

669

670 @torch.inference_mode()

671 +@pytest.mark.parametrize(("has_amx", "has_arm_bf16"), [(True, False), (False, True)])

672 def test_spec_forward_prepares_native_conv_metadata(

673 + has_amx: bool,

674 + has_arm_bf16: bool,

675 monkeypatch: pytest.MonkeyPatch,

676 ) -> None:

677 block_table = torch.tensor([[0, 1], [4, 5]], dtype=torch.int32)

@@ -687,14 +692,17 @@ def test_spec_forward_prepares_native_conv_metadata(

692 )

693 forwarded_indices = None

694 forwarded_counts = None

695 + forwarded_is_weight_packed = None

696

697 def causal_conv1d_update_cpu(**kwargs):

692 - nonlocal forwarded_counts, forwarded_indices

698 + nonlocal forwarded_counts, forwarded_indices, forwarded_is_weight_packed

699 forwarded_indices = kwargs["conv_state_indices"]

700 forwarded_counts = kwargs["num_accepted_tokens"]

701 + forwarded_is_weight_packed = kwargs["is_weight_packed"]

702 return kwargs["x"]

703

697 - monkeypatch.setattr(torch.cpu, "_is_amx_tile_supported", lambda: True)

704 + monkeypatch.setattr(torch.cpu, "_is_amx_tile_supported", lambda: has_amx)

705 + monkeypatch.setattr(gdn_attention, "is_arm_bf16", lambda: has_arm_bf16)

706 monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: False)

707 monkeypatch.setattr(

708 gdn_attention.ops, "causal_conv1d_update_cpu", causal_conv1d_update_cpu

@@ -732,6 +740,7 @@ def test_spec_forward_prepares_native_conv_metadata(

740 assert actual is not None

741 assert actual.is_contiguous()

742 assert actual.dtype == torch.int32

743 + assert forwarded_is_weight_packed is True

744 torch.testing.assert_close(actual, reference)

745

746

@@ -792,15 +801,15 @@ def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> No

801

802

803 @pytest.mark.skipif(

795 - not torch.cpu._is_amx_tile_supported(),

796 - reason="requires AMX support",

804 + not (torch.cpu._is_amx_tile_supported() or gdn_attention.is_arm_bf16()),

805 + reason="native causal conv1d requires AMX or Arm BF16 support",

806 )

807 @torch.inference_mode()

808 def test_causal_conv1d_update_cpu_accepts_wide_state() -> None:

809 state_len = CONV_KERNEL - 1

810 wide_state_len = state_len + 5

811 batch_size = 3

803 - is_vnni = True

812 + is_weight_packed = True

813 x, weight, bias = _conv_inputs(batch_size)

814 conv_state_indices = torch.tensor([2, 0, 1], dtype=torch.int32)

815

@@ -813,7 +822,7 @@ def test_causal_conv1d_update_cpu_accepts_wide_state() -> None:

822 wide_state[:, :, state_len:].fill_(7)

823 wide_tail = wide_state[:, :, state_len:].clone()

824

816 - conv_weight = _maybe_pack_conv_weight(weight, is_vnni)

825 + conv_weight = _maybe_pack_conv_weight(weight, is_weight_packed)

826 out_narrow = ops.causal_conv1d_update_cpu(

827 x=x,

828 conv_states=narrow_state,

@@ -821,7 +830,7 @@ def test_causal_conv1d_update_cpu_accepts_wide_state() -> None:

830 bias=bias,

831 silu_activation=True,

832 conv_state_indices=conv_state_indices,

824 - is_vnni=is_vnni,

833 + is_weight_packed=is_weight_packed,

834 )

835 out_wide = ops.causal_conv1d_update_cpu(

836 x=x,

@@ -830,7 +839,7 @@ def test_causal_conv1d_update_cpu_accepts_wide_state() -> None:

839 bias=bias,

840 silu_activation=True,

841 conv_state_indices=conv_state_indices,

833 - is_vnni=is_vnni,

842 + is_weight_packed=is_weight_packed,

843 )

844

845 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(

881

882

883 @pytest.mark.skipif(

875 - not torch.cpu._is_amx_tile_supported(),

876 - reason="requires AMX support",

884 + not (torch.cpu._is_amx_tile_supported() or gdn_attention.is_arm_bf16()),

885 + reason="native causal conv1d requires AMX or Arm BF16 support",

886 )

887 @pytest.mark.parametrize(

879 - ("batch_size, seq_len, accepted_counts, has_bias, silu_activation, is_vnni"),

888 + (

889 + "batch_size, "

890 + "seq_len, "

891 + "accepted_counts, "

892 + "has_bias, "

893 + "silu_activation, "

894 + "is_weight_packed"

895 + ),

896 [

897 (1, 1, [1], False, False, False),

898 (1, 1, [1], True, True, True),

@@ -893,7 +909,7 @@ def test_causal_conv1d_update_cpu_multi_token_matches_python(

909 accepted_counts: list[int],

910 has_bias: bool,

911 silu_activation: bool,

896 - is_vnni: bool,

912 + is_weight_packed: bool,

913 ) -> None:

914 dim = 96

915 state_len = seq_len + 2

@@ -911,7 +927,7 @@ def test_causal_conv1d_update_cpu_multi_token_matches_python(

927 )

928 conv_states = conv_states_ref.clone()

929

914 - conv_weight = _maybe_pack_conv_weight(weight, is_vnni)

930 + conv_weight = _maybe_pack_conv_weight(weight, is_weight_packed)

931 out = ops.causal_conv1d_update_cpu(

932 x=x,

933 conv_states=conv_states,

@@ -919,7 +935,7 @@ def test_causal_conv1d_update_cpu_multi_token_matches_python(

935 bias=bias,

936 silu_activation=silu_activation,

937 conv_state_indices=conv_state_indices,

922 - is_vnni=is_vnni,

938 + is_weight_packed=is_weight_packed,

939 num_accepted_tokens=num_accepted_tokens,

940 )

941 ref_out = _ref_causal_conv1d_update_cpu_multi(

@@ -937,8 +953,8 @@ def test_causal_conv1d_update_cpu_multi_token_matches_python(

953

954

955 @pytest.mark.skipif(

940 - not torch.cpu._is_amx_tile_supported(),

941 - reason="requires AMX support",

956 + not (torch.cpu._is_amx_tile_supported() or gdn_attention.is_arm_bf16()),

957 + reason="native causal conv1d requires AMX or Arm BF16 support",

958 )

959 @pytest.mark.parametrize("num_accepted", [0, 17])

960 @torch.inference_mode()

@@ -961,14 +977,15 @@ def test_causal_conv1d_update_cpu_rejects_invalid_accepted_count(

977 bias=None,

978 silu_activation=True,

979 conv_state_indices=torch.tensor([0], dtype=torch.int32),

964 - is_vnni=False,

980 + is_weight_packed=False,

981 num_accepted_tokens=torch.tensor([num_accepted], dtype=torch.int32),

982 )

983

984

985 @pytest.mark.skipif(

970 - not torch.cpu._is_avx512_bf16_supported(),

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

986 + not (torch.cpu._is_avx512_bf16_supported() or gdn_attention.is_arm_bf16()),

987 + reason="causal_conv1d_fwd_cpu requires AVX-512BF16"

988 + " (Intel Xeon or AMD EPYC) or Arm BF16 support",

989 )

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

991 @torch.inference_mode()

@@ -977,13 +994,13 @@ def test_causal_conv1d_fwd_cpu_two_call_split(total_tokens: int, split: int) ->

994 matches the single-call result.

995

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

980 - ignoring the carried conv state on continued chunks. Runs on any

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

997 + ignoring the carried conv state on continued chunks. Runs on aarch64 cpus with

998 + bf16 support or any AVX-512BF16 CPU since conv.cpp uses VDPBF16PS, not AMX tiles.

999 """

1000 state_len = CONV_KERNEL - 1

1001 x, weight, bias = _conv_inputs(total_tokens)

1002

986 - def amx(x_seg, conv_states, has_init):

1003 + def native(x_seg, conv_states, has_init):

1004 seq = x_seg.shape[0]

1005 return ops.causal_conv1d_fwd_cpu(

1006 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) ->

1011 cache_indices=torch.tensor([0], dtype=torch.int32),

1012 has_initial_state=torch.tensor([has_init]),

1013 silu_activation=True,

997 - is_vnni=False,

1014 + is_weight_packed=False,

1015 ).contiguous()

1016

1017 # conv_state layout passed by the AMX branch: [num_slots, dim, state_len].

1018 cs_full = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype)

1002 - out_full = amx(x, cs_full, False)

1019 + out_full = native(x, cs_full, False)

1020

1021 cs_split = torch.zeros(1, CONV_DIM, state_len, dtype=x.dtype)

1005 - out1 = amx(x[:split], cs_split, False)

1006 - out2 = amx(x[split:], cs_split, True)

1022 + out1 = native(x[:split], cs_split, False)

1023 + out2 = native(x[split:], cs_split, True)

1024 out_split = torch.cat([out1, out2], dim=1)

1025

1026 torch.testing.assert_close(out_split, out_full, atol=1e-2, rtol=1e-2)

1027

1028

1029 @pytest.mark.skipif(

1013 - not torch.cpu._is_amx_tile_supported(),

1014 - reason="requires AMX support",

1030 + not (torch.cpu._is_amx_tile_supported() or gdn_attention.is_arm_bf16()),

1031 + reason="native causal conv1d requires AMX or Arm BF16 support",

1032 )

1033 @torch.inference_mode()

1034 def test_causal_conv1d_fwd_cpu_accepts_wide_state() -> None:

1035 state_len = CONV_KERNEL - 1

1036 wide_state_len = state_len + 5

1020 - is_vnni = True

1037 + is_weight_packed = True

1038 seq_lens = [CHUNK_SIZE - 1, CHUNK_SIZE + 5]

1039 total_tokens = sum(seq_lens)

1040 x, weight, bias = _conv_inputs(total_tokens)

@@ -1034,7 +1051,7 @@ def test_causal_conv1d_fwd_cpu_accepts_wide_state() -> None:

1051 wide_state[:, :, state_len:].fill_(7)

1052 wide_tail = wide_state[:, :, state_len:].clone()

1053

1037 - conv_weight = _maybe_pack_conv_weight(weight, is_vnni)

1054 + conv_weight = _maybe_pack_conv_weight(weight, is_weight_packed)

1055 out_narrow = ops.causal_conv1d_fwd_cpu(

1056 x=x.transpose(0, 1),

1057 weight=conv_weight,

@@ -1044,7 +1061,7 @@ def test_causal_conv1d_fwd_cpu_accepts_wide_state() -> None:

1061 cache_indices=cache_indices,

1062 has_initial_state=has_initial_state,

1063 silu_activation=True,

1047 - is_vnni=is_vnni,

1064 + is_weight_packed=is_weight_packed,

1065 )

1066 out_wide = ops.causal_conv1d_fwd_cpu(

1067 x=x.transpose(0, 1),

@@ -1055,7 +1072,7 @@ def test_causal_conv1d_fwd_cpu_accepts_wide_state() -> None:

1072 cache_indices=cache_indices,

1073 has_initial_state=has_initial_state,

1074 silu_activation=True,

1058 - is_vnni=is_vnni,

1075 + is_weight_packed=is_weight_packed,

1076 )

1077

1078 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:

1117

1118

1119 # ---------------------------------------------------------------------------

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

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

1120 +# C++ conv (conv.cpp) uses AArch64BF16 or VDPBF16PS, not AMX tiles, so it can

1121 +# run on any AVX-512BF16 CPU; weight is packed on this same predicate at load time.

1122 # ---------------------------------------------------------------------------

1123

1124 _HAS_AVX512_BF16 = torch.cpu._is_avx512_bf16_supported()

1125

1126 _STATE_LEN = CONV_KERNEL - 1

1127

1111 -CONV_EQUIV_SEQ_LENS = [[1], [7], [64], [65], [1, 2, 3], [63, 64, 65], [128, 129]]

1128 +CONV_EQUIV_SEQ_LENS = [

1129 + [1],

1130 + [7],

1131 + [64],

1132 + [65],

1133 + [1, 2, 3],

1134 + [63, 64, 65],

1135 + [128, 129],

1136 +]

1137

1138

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

@@ -1153,10 +1178,10 @@ def _run_prefill_torch(x, weight, bias, seq_lens):

1178 return out.transpose(0, 1).contiguous(), conv_states

1179

1180

1156 -def _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False):

1181 +def _run_prefill_cpp(x, weight, bias, seq_lens, is_weight_packed=False):

1182 num_seqs = len(seq_lens)

1158 - packed_w = ops.causal_conv1d_weight_pack(weight) if is_vnni else weight

1159 - if is_vnni:

1183 + packed_w = ops.causal_conv1d_weight_pack(weight) if is_weight_packed else weight

1184 + if is_weight_packed:

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

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

1187 conv_state = torch.zeros(

@@ -1176,7 +1201,7 @@ def _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False):

1201 cache_indices=torch.arange(num_seqs, dtype=torch.int32),

1202 has_initial_state=torch.zeros(num_seqs, dtype=torch.bool),

1203 silu_activation=True,

1179 - is_vnni=is_vnni,

1204 + is_weight_packed=is_weight_packed,

1205 )

1206 return out.transpose(0, 1).contiguous(), conv_state

1207

@@ -1245,7 +1270,7 @@ def test_causal_conv1d_fwd_cpu_ds_adapter_varlen_mixed_initial_state(

1270 cache_indices=state_indices,

1271 has_initial_state=has_initial_state,

1272 silu_activation=True,

1248 - is_vnni=False,

1273 + is_weight_packed=False,

1274 )

1275 monkeypatch.setattr(gdn_attention, "is_conv_state_dim_first", lambda: True)

1276 out_ds = gdn_attention._causal_conv1d_fwd_cpu(

@@ -1257,7 +1282,7 @@ def test_causal_conv1d_fwd_cpu_ds_adapter_varlen_mixed_initial_state(

1282 cache_indices=state_indices,

1283 has_initial_state=has_initial_state,

1284 silu_activation=True,

1260 - is_vnni=False,

1285 + is_weight_packed=False,

1286 )

1287

1288 torch.testing.assert_close(out_ds, out_sd)

@@ -1302,7 +1327,7 @@ def test_causal_conv1d_update_cpu_ds_adapter_wide_multi_token_copyback(

1327 bias=None,

1328 silu_activation=False,

1329 conv_state_indices=state_indices,

1305 - is_vnni=False,

1330 + is_weight_packed=False,

1331 num_accepted_tokens=num_accepted_tokens,

1332 )

1333

@@ -1321,7 +1346,8 @@ def test_causal_conv1d_update_cpu_ds_adapter_wide_multi_token_copyback(

1346

1347

1348 @pytest.mark.skipif(

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

1349 + not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()),

1350 + reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support",

1351 )

1352 @pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS)

1353 @torch.inference_mode()

@@ -1335,7 +1361,8 @@ def test_conv_cpp_matches_torch(seq_lens):

1361

1362

1363 @pytest.mark.skipif(

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

1364 + not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()),

1365 + reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support",

1366 )

1367 @pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS)

1368 @torch.inference_mode()

@@ -1354,89 +1381,114 @@ def test_conv_cpp_no_worse_than_torch_vs_fp32(seq_lens):

1381

1382

1383 @pytest.mark.skipif(

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

1384 + not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()),

1385 + reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support",

1386 )

1387 @pytest.mark.parametrize("seq_lens", CONV_EQUIV_SEQ_LENS)

1388 @torch.inference_mode()

1361 -def test_conv_cpp_vnni_packed_matches_torch(seq_lens):

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

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

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

1389 +def test_conv_cpp_weight_packed_matches_torch(seq_lens):

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

1391 + conv_state view + is_weight_packed=True) must match the torch fallback. Validates

1392 + the native packing + layout handoff."""

1393 x, weight, bias = _conv_inputs(sum(seq_lens))

1394 out_torch, _ = _run_prefill_torch(x, weight, bias, seq_lens)

1367 - out_vnni, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=True)

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

1395 + out_weight_packed, _ = _run_prefill_cpp(

1396 + x, weight, bias, seq_lens, is_weight_packed=True

1397 + )

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

1399

1400

1401 @pytest.mark.skipif(

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

1402 + not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()),

1403 + reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support",

1404 )

1405 @pytest.mark.parametrize("batch", DECODE_BATCH_SIZES)

1406 +@pytest.mark.parametrize(

1407 + ("dim", "has_bias", "silu_activation", "is_weight_packed"),

1408 + [

1409 + (32, False, False, False),

1410 + (32, True, True, True),

1411 + (96, False, True, True),

1412 + (96, True, False, False),

1413 + (8192, False, True, False),

1414 + (8192, True, True, True),

1415 + ],

1416 +)

1417 @torch.inference_mode()

1376 -def test_conv_update_cpp_matches_torch(batch):

1377 - """Decode conv: causal_conv1d_update_cpu matches causal_conv1d_update_torch,

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

1418 +def test_conv_update_cpp_matches_torch(

1419 + batch, dim, has_bias, silu_activation, is_weight_packed

1420 +):

1421 + """Decode preserves indexed, padded state across successive signed inputs."""

1422 from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (

1423 causal_conv1d_update_torch,

1424 )

1425

1383 - x = tensor_cache(batch * CONV_DIM, torch.bfloat16).view(batch, CONV_DIM)

1384 - weight = tensor_cache(CONV_DIM * CONV_KERNEL, torch.bfloat16).view(

1385 - CONV_DIM, CONV_KERNEL

1386 - )

1387 - bias = tensor_cache(CONV_DIM, torch.bfloat16)

1388 - conv_state = tensor_cache(batch * CONV_DIM * _STATE_LEN, torch.bfloat16).view(

1389 - batch, CONV_DIM, _STATE_LEN

1390 - )

1391 -

1392 - cs_torch = conv_state.clone()

1393 - out_torch = causal_conv1d_update_torch(

1394 - x=x.unsqueeze(-1),

1395 - conv_state=cs_torch,

1396 - weight=weight,

1397 - bias=bias,

1398 - activation="silu",

1399 - ).squeeze(-1)

1400 -

1401 - cs_cpp = conv_state.clone()

1402 - out_cpp = ops.causal_conv1d_update_cpu(

1403 - x=x.contiguous(),

1404 - conv_states=cs_cpp,

1405 - weight=weight,

1406 - bias=bias,

1407 - silu_activation=True,

1408 - conv_state_indices=torch.arange(batch, dtype=torch.int32),

1409 - is_vnni=False,

1410 - )

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

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

1426 + torch.manual_seed(0)

1427 + weight = torch.randn(dim, CONV_KERNEL, dtype=torch.bfloat16)

1428 + bias = torch.randn(dim, dtype=torch.bfloat16) if has_bias else None

1429 + conv_weight = _maybe_pack_conv_weight(weight, is_weight_packed)

1430 + indices = torch.arange(batch, 0, -1, dtype=torch.int32)

1431 + storage_ref = torch.randn(batch + 2, _STATE_LEN + 2, dim, dtype=torch.bfloat16)

1432 + storage_cpp = storage_ref.clone()

1433 + cs_cpp = storage_cpp[:, :_STATE_LEN].transpose(1, 2)

1434 +

1435 + for _ in range(3):

1436 + x = torch.randn(batch, dim, dtype=torch.bfloat16)

1437 + cs_torch = storage_ref[indices, :_STATE_LEN].transpose(1, 2).contiguous()

1438 + out_torch = causal_conv1d_update_torch(

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

1440 + conv_state=cs_torch,

1441 + weight=weight,

1442 + bias=bias,

1443 + activation="silu" if silu_activation else None,

1444 + ).squeeze(-1)

1445 + storage_ref[indices, :_STATE_LEN] = cs_torch.transpose(1, 2)

1446 + out_cpp = ops.causal_conv1d_update_cpu(

1447 + x=x,

1448 + conv_states=cs_cpp,

1449 + weight=conv_weight,

1450 + bias=bias,

1451 + silu_activation=silu_activation,

1452 + conv_state_indices=indices,

1453 + is_weight_packed=is_weight_packed,

1454 + )

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

1456 + torch.testing.assert_close(storage_cpp, storage_ref, atol=0, rtol=0)

1457

1458

1459 @pytest.mark.skipif(

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

1460 + not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()),

1461 + reason="C++ causal_conv1d requires AVX-512BF16 or Arm BF16 support",

1462 )

1463 @torch.inference_mode()

1464 def test_conv_weight_pack_roundtrip_unpacked_matches():

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

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

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

1423 - in place and stashes the original as `_cpu_unpacked_conv_weight` for the torch

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

1466 + with packed weight + is_weight_packed=True must equal unpacked weight

1467 + + is_weight_packed=False. Guards the spec-decode contract: the runtime packs

1468 + `layer.conv1d.weight` in place and stashes the original as

1469 + `_cpu_unpacked_conv_weight` for the torch spec-decode path, so both must produce

1470 + identical math.

1471 """

1472 seq_lens = [7, 64, 65]

1473 x, weight, bias = _conv_inputs(sum(seq_lens))

1428 - out_unpacked, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=False)

1429 - out_packed, _ = _run_prefill_cpp(x, weight, bias, seq_lens, is_vnni=True)

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

1474 + out_unpacked, state_unpacked = _run_prefill_cpp(

1475 + x, weight, bias, seq_lens, is_weight_packed=False

1476 + )

1477 + out_packed, state_packed = _run_prefill_cpp(

1478 + x, weight, bias, seq_lens, is_weight_packed=True

1479 + )

1480 + torch.testing.assert_close(out_packed, out_unpacked, atol=0, rtol=0)

1481 + torch.testing.assert_close(state_packed, state_unpacked, atol=0, rtol=0)

1482

1483

1484 @pytest.mark.skipif(

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

1485 + not (_HAS_AVX512_BF16 or gdn_attention.is_arm_bf16()),

1486 + reason="C++ causal_conv1d requires AVX-512BF16 or AArch64 BF16",

1487 )

1488 @torch.inference_mode()

1489 def test_spec_decode_unpacked_conv_weight_stash():

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

1439 - at load time and stash the original ``(dim, width)`` tensor as

1490 + """Spec-decode correctness: AVX-512BF16 and AArch64 BF16 CPUs pack ``conv1d.weight``

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

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

1493 reading the packed weight directly produces garbage. Verifies the recovered

1494 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():

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

1513 assert hasattr(conv, "_cpu_unpacked_conv_weight"), (

1514 "dispatch_cpu_unquantized_gemm did not stash _cpu_unpacked_conv_weight "

1463 - "on an AVX-512BF16 CPU"

1515 + "on an AArch64-BF16 or AVX-512BF16 CPU"

1516 )

1517 stash = conv._cpu_unpacked_conv_weight

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

@@ -1469,10 +1521,10 @@ def test_spec_decode_unpacked_conv_weight_stash():

1521 recovered = getattr(conv, "_cpu_unpacked_conv_weight", None)

1522 assert recovered is not None

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

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

1524 + packed_weight = conv.weight # [dim, 1, width], packed contents

1525

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

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

1527 + # nonspec C++ conv with the packed weight (is_weight_packed=True).

1528 seq_lens = [7, 65]

1529 total = sum(seq_lens)

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

@@ -1494,7 +1546,7 @@ def test_spec_decode_unpacked_conv_weight_stash():

1546 cache_indices=torch.arange(num_seqs, dtype=torch.int32),

1547 has_initial_state=torch.zeros(num_seqs, dtype=torch.bool),

1548 silu_activation=True,

1497 - is_vnni=True,

1549 + is_weight_packed=True,

1550 )

1551 .transpose(0, 1)

1552 .contiguous()

tests/model_executor/test_cpu_unquantized_gemm_dispatch.py

+52 -1

@@ -6,10 +6,61 @@ import pytest

6 import torch

8 from vllm.model_executor.layers import utils

9 -from vllm.platforms import current_platform

9 +from vllm.platforms import CpuArchEnum, current_platform

10 from vllm.utils.torch_utils import set_default_torch_dtype

13 +def test_dispatch_prepacks_arm_bf16_causal_conv(monkeypatch):

14 + monkeypatch.setattr(

15 + current_platform, "get_cpu_architecture", lambda: CpuArchEnum.ARM

16 + )

17 + monkeypatch.setattr(torch.cpu, "get_capabilities", lambda: {"bf16": True})

18 + monkeypatch.setattr(torch.cpu, "_is_avx512_bf16_supported", lambda: False)

19 +

20 + packed = torch.randn(32, 4, dtype=torch.bfloat16)

21 + pack_inputs = []

22 +

23 + def pack(weight):

24 + pack_inputs.append(weight.clone())

25 + return packed

26 +

27 + monkeypatch.setattr(utils.ops, "causal_conv1d_weight_pack", pack)

28 + original = torch.randn(32, 1, 4, dtype=torch.bfloat16)

29 + layer = torch.nn.Module()

30 + layer.weight = torch.nn.Parameter(original.clone(), requires_grad=False)

31 +

32 + utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)

33 +

34 + assert len(pack_inputs) == 1

35 + # Check causal_conv1d_weight_pack received the expected weights

36 + torch.testing.assert_close(pack_inputs[0], original.view(32, 4))

37 + # Check we have correctly stashed the unpacked weights

38 + torch.testing.assert_close(layer._cpu_unpacked_conv_weight, original.view(32, 4))

39 + # Check we have correctly stored the packed weights

40 + assert layer.weight.data_ptr() == packed.data_ptr()

41 +

42 +

43 +def test_dispatch_does_not_pack_3d_expert_weight(monkeypatch):

44 + monkeypatch.setattr(

45 + current_platform, "get_cpu_architecture", lambda: CpuArchEnum.ARM

46 + )

47 + monkeypatch.setattr(torch.cpu, "get_capabilities", lambda: {"bf16": True})

48 + monkeypatch.setattr(torch.cpu, "_is_avx512_bf16_supported", lambda: False)

49 + pack_calls = []

50 + monkeypatch.setattr(

51 + utils.ops,

52 + "causal_conv1d_weight_pack",

53 + lambda weight: pack_calls.append(weight),

54 + )

55 +

56 + layer = torch.nn.Module()

57 + layer.weight = torch.nn.Parameter(

58 + torch.randn(8, 32, 4, dtype=torch.bfloat16), requires_grad=False

59 + )

60 + utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)

61 + assert pack_calls == []

62 +

63 +

64 @pytest.fixture(scope="module")

65 def _mock_zentorch_linear_unary():

66 """Register a mock zentorch_linear_unary op when zentorch is not installed.

vllm/_custom_ops.py

+4 -4

@@ -3666,7 +3666,7 @@ def causal_conv1d_fwd_cpu(

3666 cache_indices: torch.Tensor | None,

3667 has_initial_state: torch.Tensor | None,

3668 silu_activation: bool,

3669 - is_vnni: bool,

3669 + is_weight_packed: bool,

3670 ) -> torch.Tensor:

3671 return torch.ops._C.causal_conv1d_fwd_cpu(

3672 x,

@@ -3678,7 +3678,7 @@ def causal_conv1d_fwd_cpu(

3678 has_initial_state,

3679 silu_activation,

3680 -1,

3681 - is_vnni,

3681 + is_weight_packed,

3682 )

3683

3684

@@ -3689,7 +3689,7 @@ def causal_conv1d_update_cpu(

3689 bias: torch.Tensor | None,

3690 silu_activation: bool,

3691 conv_state_indices: torch.Tensor | None,

3692 - is_vnni: bool,

3692 + is_weight_packed: bool,

3693 num_accepted_tokens: torch.Tensor | None = None,

3694 ) -> torch.Tensor:

3695 return torch.ops._C.causal_conv1d_update_cpu(

@@ -3701,7 +3701,7 @@ def causal_conv1d_update_cpu(

3701 num_accepted_tokens,

3702 conv_state_indices,

3703 -1,

3704 - is_vnni,

3704 + is_weight_packed,

3705 )

3706

3707

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

+53 -63

@@ -14,7 +14,6 @@ from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (

14 )

15 from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import (

16 causal_conv1d_update_cpu,

17 - causal_conv1d_update_torch,

17 )

18 from vllm.platforms import CpuArchEnum, current_platform

19 from vllm.utils.torch_utils import (

@@ -56,7 +55,7 @@ def _causal_conv1d_fwd_cpu(

55 cache_indices: torch.Tensor,

56 has_initial_state: torch.Tensor,

57 silu_activation: bool,

59 - is_vnni: bool,

58 + is_weight_packed: bool,

59 ) -> torch.Tensor:

60 if not is_conv_state_dim_first():

61 return ops.causal_conv1d_fwd_cpu(

@@ -68,7 +67,7 @@ def _causal_conv1d_fwd_cpu(

67 cache_indices=cache_indices,

68 has_initial_state=has_initial_state,

69 silu_activation=silu_activation,

71 - is_vnni=is_vnni,

70 + is_weight_packed=is_weight_packed,

71 )

73 scratch, identity_indices = _stage_ds_conv_states(conv_states, cache_indices)

@@ -81,7 +80,7 @@ def _causal_conv1d_fwd_cpu(

80 cache_indices=identity_indices,

81 has_initial_state=has_initial_state,

82 silu_activation=silu_activation,

84 - is_vnni=is_vnni,

83 + is_weight_packed=is_weight_packed,

84 )

85 _scatter_ds_conv_states(conv_states, cache_indices, scratch)

86 return out

@@ -95,7 +94,7 @@ def _causal_conv1d_update_cpu(

94 bias: torch.Tensor | None,

95 silu_activation: bool,

96 conv_state_indices: torch.Tensor,

98 - is_vnni: bool,

97 + is_weight_packed: bool,

98 num_accepted_tokens: torch.Tensor | None = None,

99 ) -> torch.Tensor:

100 if not is_conv_state_dim_first():

@@ -106,7 +105,7 @@ def _causal_conv1d_update_cpu(

105 bias=bias,

106 silu_activation=silu_activation,

107 conv_state_indices=conv_state_indices,

109 - is_vnni=is_vnni,

108 + is_weight_packed=is_weight_packed,

109 num_accepted_tokens=num_accepted_tokens,

110 )

111

@@ -118,13 +117,20 @@ def _causal_conv1d_update_cpu(

117 bias=bias,

118 silu_activation=silu_activation,

119 conv_state_indices=identity_indices,

121 - is_vnni=is_vnni,

120 + is_weight_packed=is_weight_packed,

121 num_accepted_tokens=num_accepted_tokens,

122 )

123 _scatter_ds_conv_states(conv_states, conv_state_indices, scratch)

124 return out

125

126

127 +def is_arm_bf16() -> bool:

128 + return (

129 + current_platform.get_cpu_architecture() == CpuArchEnum.ARM

130 + and torch.cpu.get_capabilities().get("bf16", False)

131 + )

132 +

133 +

134 def cpu_gdn_attention_core(

135 mixed_qkv: torch.Tensor,

136 b: torch.Tensor,

@@ -201,9 +207,10 @@ def _cpu_gdn_attention_nonspec(

207 assert state_indices_tensor is not None

208 assert query_start_loc is not None

209

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

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

206 - use_cpp_conv = torch.cpu._is_avx512_bf16_supported()

210 + # C++ conv (conv.cpp) runs on aarch64 with bf16 support and on x86 uses VDPBF16PS,

211 + # not AMX tiles, so it runs on any AVX-512BF16 CPU. The weights are packed on this

212 + # same predicate at load time.

213 + use_cpp_conv = torch.cpu._is_avx512_bf16_supported() or is_arm_bf16()

214

215 conv_state = layer.kv_cache[0]

216 if use_cpp_conv:

@@ -244,28 +251,17 @@ def _cpu_gdn_attention_nonspec(

251 bias=layer.conv1d.bias,

252 silu_activation=(layer.activation == "silu"),

253 conv_state_indices=decode_state_indices,

247 - is_vnni=True,

254 + is_weight_packed=True,

255 )

256 else:

250 - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:

251 - decode_conv_state = conv_state[decode_state_indices].contiguous()

252 - decode_mixed_qkv = causal_conv1d_update_torch(

253 - x=decode_mixed_qkv.unsqueeze(-1),

254 - conv_state=decode_conv_state,

255 - weight=conv_weights,

256 - bias=layer.conv1d.bias,

257 - activation=layer.activation,

258 - ).squeeze(-1)

259 - conv_state[decode_state_indices] = decode_conv_state

260 - else:

261 - decode_mixed_qkv = causal_conv1d_update_cpu(

262 - x=decode_mixed_qkv,

263 - conv_state=conv_state,

264 - weight=conv_weights,

265 - bias=layer.conv1d.bias,

266 - activation=layer.activation,

267 - conv_state_indices=decode_state_indices,

268 - )

257 + decode_mixed_qkv = causal_conv1d_update_cpu(

258 + x=decode_mixed_qkv,

259 + conv_state=conv_state,

260 + weight=conv_weights,

261 + bias=layer.conv1d.bias,

262 + activation=layer.activation,

263 + conv_state_indices=decode_state_indices,

264 + )

265

266 query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)

267

@@ -315,7 +311,7 @@ def _cpu_gdn_attention_nonspec(

311 cache_indices=prefill_state_indices,

312 has_initial_state=prefill_has_initial_state,

313 silu_activation=layer.activation == "silu",

318 - is_vnni=True,

314 + is_weight_packed=True,

315 ).transpose(0, 1)

316 else:

317 prefill_mixed_qkv = causal_conv1d_torch(

@@ -383,9 +379,9 @@ def _ssm_state_view(layer) -> torch.Tensor:

379 def _unpacked_conv_weight(layer) -> torch.Tensor:

380 """Return the plain (dim, width) conv weight.

381

386 - On AMX the conv1d weight is VNNI-packed in place at load time (only usable

387 - by the AMX C++ kernel), so the torch spec-decode path relies on the

388 - un-packed copy stashed by ``dispatch_cpu_unquantized_gemm``.

382 + On AMX and aarch64+bf16 the conv1d weight is packed in place at load time,

383 + so the torch spec-decode path relies on the un-packed copy stashed by

384 + ``dispatch_cpu_unquantized_gemm``.

385 """

386 w = getattr(layer.conv1d, "_cpu_unpacked_conv_weight", None)

387 if w is not None:

@@ -505,13 +501,15 @@ def _spec_forward(

501 seq_starts = spec_qsl_cpu[:-1]

502 seq_lens = spec_qsl_cpu[1:] - spec_qsl_cpu[:-1]

503

504 + is_amx = torch.cpu._is_amx_tile_supported()

505 +

506 # ---- 1. Convolution (per-sequence rolling buffer) ----

507 dim = mixed_qkv_spec.size(-1)

508 bias = layer.conv1d.bias

509 silu = layer.activation == "silu"

510

511 can_use_native_conv = (

514 - torch.cpu._is_amx_tile_supported()

512 + (is_amx or is_arm_bf16())

513 and width == 4

514 and num_spec_decodes > 0

515 and bool(torch.all(seq_lens == seq_lens[0]).item())

@@ -528,7 +526,7 @@ def _spec_forward(

526 conv_state_indices=spec_state_indices[:num_spec_decodes, 0]

527 .to("cpu", torch.int32)

528 .contiguous(),

531 - is_vnni=True,

529 + is_weight_packed=True,

530 num_accepted_tokens=num_accepted[:num_spec_decodes].to("cpu", torch.int32),

531 ).view_as(mixed_qkv_spec)

532 else:

@@ -603,9 +601,9 @@ def _spec_aware_nonspec(

601 assert query_start_loc is not None

602 state_indices_tensor = state_indices_tensor.contiguous()

603

606 - is_amx = torch.cpu._is_amx_tile_supported()

604 + can_use_native_conv = torch.cpu._is_amx_tile_supported() or is_arm_bf16()

605

608 - if not is_amx:

606 + if not can_use_native_conv:

607 conv_weights = _unpacked_conv_weight(layer)

608

609 num_decodes = attn_metadata_i.num_decodes

@@ -618,7 +616,8 @@ def _spec_aware_nonspec(

616 decode_b = b[:num_decode_tokens]

617 decode_a = a[:num_decode_tokens]

618 decode_state_indices = state_indices_tensor[:num_decodes]

621 - if is_amx:

619 +

620 + if can_use_native_conv:

621 decode_mixed_qkv = _causal_conv1d_update_cpu(

622 x=decode_mixed_qkv,

623 conv_states=conv_buf,

@@ -626,30 +625,19 @@ def _spec_aware_nonspec(

625 bias=layer.conv1d.bias,

626 silu_activation=layer.activation == "silu",

627 conv_state_indices=decode_state_indices,

629 - is_vnni=True,

628 + is_weight_packed=True,

629 )

630 else:

631 # Only the first ``width-1`` columns hold the real conv state.

632 conv_state_view = conv_buf[:, :, : width - 1]

634 - if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:

635 - decode_conv_state = conv_state_view[decode_state_indices].contiguous()

636 - decode_mixed_qkv = causal_conv1d_update_torch(

637 - x=decode_mixed_qkv.unsqueeze(-1),

638 - conv_state=decode_conv_state,

639 - weight=conv_weights,

640 - bias=layer.conv1d.bias,

641 - activation=layer.activation,

642 - ).squeeze(-1)

643 - conv_state_view[decode_state_indices] = decode_conv_state

644 - else:

645 - decode_mixed_qkv = causal_conv1d_update_cpu(

646 - x=decode_mixed_qkv,

647 - conv_state=conv_state_view,

648 - weight=conv_weights,

649 - bias=layer.conv1d.bias,

650 - activation=layer.activation,

651 - conv_state_indices=decode_state_indices,

652 - )

633 + decode_mixed_qkv = causal_conv1d_update_cpu(

634 + x=decode_mixed_qkv,

635 + conv_state=conv_state_view,

636 + weight=conv_weights,

637 + bias=layer.conv1d.bias,

638 + activation=layer.activation,

639 + conv_state_indices=decode_state_indices,

640 + )

641

642 query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv)

643 attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu(

@@ -685,7 +673,8 @@ def _spec_aware_nonspec(

673 prefill_has_initial_state = has_initial_state[

674 num_decodes : num_decodes + num_prefills

675 ]

688 - if is_amx:

676 +

677 + if can_use_native_conv:

678 prefill_mixed_qkv = _causal_conv1d_fwd_cpu(

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

680 weight=layer.conv1d.weight,

@@ -695,7 +684,7 @@ def _spec_aware_nonspec(

684 cache_indices=prefill_state_indices,

685 has_initial_state=prefill_has_initial_state,

686 silu_activation=layer.activation == "silu",

698 - is_vnni=True,

687 + is_weight_packed=True,

688 ).transpose(0, 1)

689 else:

690 prefill_mixed_qkv = causal_conv1d_torch(

@@ -757,8 +746,9 @@ def _spec_aware_nonspec_subset(

746 prefill_state_indices = prefill_state_indices.contiguous()

747

748 is_amx = torch.cpu._is_amx_tile_supported()

749 + can_use_native_conv = is_amx or is_arm_bf16()

750

761 - if is_amx:

751 + if can_use_native_conv:

752 conv_out = _causal_conv1d_fwd_cpu(

753 x=mixed_qkv.transpose(0, 1),

754 weight=layer.conv1d.weight,

@@ -768,7 +758,7 @@ def _spec_aware_nonspec_subset(

758 cache_indices=prefill_state_indices,

759 has_initial_state=has_initial_state,

760 silu_activation=layer.activation == "silu",

771 - is_vnni=True,

761 + is_weight_packed=True,

762 ).transpose(0, 1)

763 else:

764 conv_weights = _unpacked_conv_weight(layer)

vllm/model_executor/layers/utils.py

+22 -8

@@ -475,13 +475,27 @@ def dispatch_cpu_unquantized_gemm(

475 # Skip CPU GEMM dispatch for non-2D weights (e.g. MoE 3D expert weights).

476 # These layers are handled by their own specialized methods.

477 if layer.weight.ndim != 2:

478 - # this is not a linear layer

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

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

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

482 - # (e.g. AMD Zen5/Turin).

483 - if torch.cpu._is_avx512_bf16_supported() and hasattr(

484 - ops, "causal_conv1d_weight_pack"

478 + # This is not a linear layer.

479 + # For now it should be a causal_conv1d op or MoE 3D expert weights.

480 + # Causal-conv weights use [dim, 1, width].

481 + # The C++ causal_conv1d kernels support aarch64 with bf16 and

482 + # x86 CPUs that use VDPBF16PS (no AMX tiles). So the

483 + # weight prepack applies to aarch64 with bf16 and any AVX-512BF16 CPU,

484 + # not just AMX (e.g. AMD Zen5/Turin).

485 + is_causal_conv1d_weight = (

486 + layer.weight.ndim == 3

487 + and layer.weight.size(1) == 1

488 + and layer.weight.size(2) == 4

489 + and layer.weight.dtype == torch.bfloat16

490 + )

491 + is_arm_bf16 = (

492 + current_platform.get_cpu_architecture() == CpuArchEnum.ARM

493 + and torch.cpu.get_capabilities().get("bf16", False)

494 + )

495 + if (

496 + (is_arm_bf16 or torch.cpu._is_avx512_bf16_supported())

497 + and is_causal_conv1d_weight

498 + and hasattr(ops, "causal_conv1d_weight_pack")

499 ):

500 # prepack conv weight

501 unpacked = (

@@ -493,7 +507,7 @@ def dispatch_cpu_unquantized_gemm(

507 .clone()

508 )

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

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

510 + # and unsupported-layout fallbacks can still use ordinary weights.

511 layer._cpu_unpacked_conv_weight = unpacked

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

513 return