@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 544cb724c8

[CPU][Spec Decode] Optimize GDN conv path for speculative decoding (#48577)

Signed-off-by: Li, Tianmu <tianmu.li@intel.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>

Tianmu Li committed Jul 26, 2026 at 23:20 UTC 544cb724c888dda1eff06a02cb0baca783c3dd86

5 files changed +767 -122

csrc/cpu/sgl-kernels/conv.cpp

+173 -12

@@ -451,6 +451,90 @@ void causal_conv1d_update_kernel_impl(

451 });

452 }

453

454 +template <typename scalar_t>

455 +void causal_conv1d_update_multi_kernel_impl(

456 + scalar_t* __restrict__ out,

457 + const scalar_t* __restrict__ input,

458 + scalar_t* __restrict__ conv_states,

459 + const scalar_t* __restrict__ weight,

460 + const scalar_t* __restrict__ bias,

461 + const int32_t* __restrict__ num_accepted_tokens,

462 + const int32_t* __restrict__ conv_indices,

463 + bool silu_activation,

464 + int64_t batch,

465 + int64_t dim,

466 + int64_t seqlen,

467 + int64_t width,

468 + int64_t state_len,

469 + int64_t conv_state_slot_stride) {

470 + constexpr int64_t BLOCK_N = block_size_n() * 2;

471 + const int64_t NB = div_up(dim, BLOCK_N);

472 +

473 + AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {

474 + at::parallel_for(0, batch * NB, 0, [&](int64_t begin, int64_t end) {

475 + int64_t bs{0}, nb{0};

476 + data_index_init(begin, bs, batch, nb, NB);

477 +

478 + for (int64_t i = begin; i < end; ++i) {

479 + const int64_t nb_start = nb * BLOCK_N;

480 + const int64_t nb_size = std::min(dim - nb_start, BLOCK_N);

481 + const int32_t conv_state_index = conv_indices[bs];

482 + const int32_t history_offset = num_accepted_tokens[bs] - 1;

483 +

484 + switch (width << 4 | nb_size >> 4) {

485 + case 0x42:

486 + tinygemm_kernel<scalar_t, 4, 32, has_bias, has_silu>::apply(

487 + input + bs * seqlen * dim + nb_start,

488 + weight + nb_start * width,

489 + out + bs * seqlen * dim + nb_start,

490 + has_bias ? bias + nb_start : nullptr,

491 + conv_states + conv_state_index * conv_state_slot_stride +

492 + history_offset * dim + nb_start,

493 + true,

494 + seqlen,

495 + dim,

496 + true);

497 + break;

498 + case 0x44:

499 + tinygemm_kernel<scalar_t, 4, 64, has_bias, has_silu>::apply(

500 + input + bs * seqlen * dim + nb_start,

501 + weight + nb_start * width,

502 + out + bs * seqlen * dim + nb_start,

503 + has_bias ? bias + nb_start : nullptr,

504 + conv_states + conv_state_index * conv_state_slot_stride +

505 + history_offset * dim + nb_start,

506 + true,

507 + seqlen,

508 + dim,

509 + true);

510 + break;

511 + default:

512 + TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);

513 + }

514 +

515 + data_index_step(bs, batch, nb, NB);

516 + }

517 + });

518 + });

519 +

520 + at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {

521 + for (int64_t bs = begin; bs < end; ++bs) {

522 + const int32_t conv_state_index = conv_indices[bs];

523 + const int32_t num_accepted = num_accepted_tokens[bs];

524 + scalar_t* state = conv_states + conv_state_index * conv_state_slot_stride;

525 +

526 + std::memmove(

527 + state,

528 + state + num_accepted * dim,

529 + (state_len - seqlen) * dim * sizeof(scalar_t));

530 + std::memcpy(

531 + state + (state_len - seqlen) * dim,

532 + input + bs * seqlen * dim,

533 + seqlen * dim * sizeof(scalar_t));

534 + }

535 + });

536 +}

537 +

538 } // anonymous namespace

539

540 // from [dim, width] or [N, K]

@@ -545,7 +629,7 @@ at::Tensor get_block_indices(const std::optional<at::Tensor>& offsets, int64_t n

629 // query_start_loc: (batch + 1) int32

630 // cache_indices: (batch) int32

631 // has_initial_state: (batch) bool

548 -// conv_states: (..., dim, width - 1) itype

632 +// conv_states: (..., dim, state_len) itype, where state_len >= width - 1

633 // activation: either None or "silu" or "swish"

634 // pad_slot_id: int

635 //

@@ -586,11 +670,14 @@ at::Tensor causal_conv1d_fwd_cpu(

670 CHECK_EQ(conv_states_val.scalar_type(), scalar_type);

671 CHECK_GE(padded_batch, batch);

672 CHECK_EQ(conv_states_val.size(1), dim);

589 - CHECK_EQ(conv_states_val.size(2), width - 1);

673 + const int64_t state_len = conv_states_val.size(2);

674 + CHECK_GE(state_len, width - 1);

675

676 // adjust `conv_states` to be contiguous on `dim`

677 // should happen only once

678 if (conv_states_val.stride(-2) != 1) {

679 + TORCH_CHECK(state_len == width - 1,

680 + "causal_conv1d_fwd_cpu: wide conv_states must be contiguous on dim.");

681 auto conv_states_copy = conv_states_val.clone();

682 conv_states_val.as_strided_({padded_batch, dim, width - 1}, {(width - 1) * dim, 1, dim});

683 conv_states_val.copy_(conv_states_copy);

@@ -651,14 +738,14 @@ at::Tensor causal_conv1d_fwd_cpu(

738

739 // API aligned with GPUs

740 //

654 -// x: (batch, dim) or (batch, dim, seqlen)

741 +// x: (batch, dim) or (batch, seqlen, dim)

742 // conv_state: (..., dim, state_len), where state_len >= width - 1

743 // weight: (dim, width)

744 // bias: (dim,)

658 -// cache_seqlens: (batch,), dtype int32.

745 +// num_accepted_tokens: (batch,), dtype int32.

746 // conv_state_indices: (batch,), dtype int32

747 // pad_slot_id: int

661 -// out: (batch, dim) or (batch, dim, seqlen)

748 +// out: (batch, dim) or (batch, seqlen, dim)

749 //

750 at::Tensor causal_conv1d_update_cpu(

751 const at::Tensor& x,

@@ -666,7 +753,7 @@ at::Tensor causal_conv1d_update_cpu(

753 const at::Tensor& weight,

754 const std::optional<at::Tensor>& bias,

755 bool silu_activation,

669 - const std::optional<at::Tensor>& cache_seqlens,

756 + const std::optional<at::Tensor>& num_accepted_tokens,

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

758 int64_t pad_slot_id,

759 bool is_vnni) {

@@ -674,13 +761,13 @@ at::Tensor causal_conv1d_update_cpu(

761 CHECK_CONTIGUOUS(weight);

762 auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);

763

677 - // TODO: add multi-token prediction support

678 - TORCH_CHECK(x.dim() == 2, "causal_conv1d_update_cpu: expect x to be 2D tensor.");

679 - TORCH_CHECK(!cache_seqlens.has_value(), "causal_conv1d_update_cpu: don't support cache_seqlens.");

764 + TORCH_CHECK(

765 + x.dim() == 2 || x.dim() == 3,

766 + "causal_conv1d_update_cpu: expect x to be 2D or 3D tensor.");

767

768 int64_t batch = x.size(0);

682 - int64_t dim = x.size(1);

683 - int64_t seqlen = 1;

769 + int64_t dim = x.dim() == 2 ? x.size(1) : x.size(2);

770 + int64_t seqlen = x.dim() == 2 ? 1 : x.size(1);

771 int64_t width = weight.size(-1);

772

773 const auto scalar_type = x.scalar_type();

@@ -690,10 +777,84 @@ at::Tensor causal_conv1d_update_cpu(

777

778 CHECK_EQ(conv_states.scalar_type(), scalar_type);

779 CHECK_EQ(conv_states.size(1), dim);

693 - CHECK_EQ(conv_states.size(2), width - 1);

780 + const int64_t state_len = conv_states.size(2);

781 + CHECK_GE(state_len, width - 1);

782 +

783 + if (x.dim() == 3) {

784 + TORCH_CHECK(

785 + num_accepted_tokens.has_value(),

786 + "causal_conv1d_update_cpu: num_accepted_tokens is required for 3D x.");

787 + TORCH_CHECK(

788 + conv_state_indices.has_value(),

789 + "causal_conv1d_update_cpu: conv_state_indices is required for 3D x.");

790 + CHECK_OPTIONAL_SHAPE_DTYPE(num_accepted_tokens, batch, at::kInt);

791 + TORCH_CHECK(

792 + width == 4,

793 + "causal_conv1d_update_cpu: support only width of 4 for 3D x.");

794 + TORCH_CHECK(

795 + seqlen > 0,

796 + "causal_conv1d_update_cpu: expect non-empty sequence for 3D x.");

797 + TORCH_CHECK(

798 + state_len >= seqlen,

799 + "causal_conv1d_update_cpu: state_len must be >= seqlen for 3D x.");

800 + TORCH_CHECK(

801 + conv_states.stride(-2) == 1 && conv_states.stride(-1) == dim,

802 + "causal_conv1d_update_cpu: 3D x requires SD conv_states layout.");

803 +

804 + const int32_t* accepted_counts =

805 + num_accepted_tokens.value().data_ptr<int32_t>();

806 + const int32_t* indices = conv_state_indices.value().data_ptr<int32_t>();

807 + const int64_t num_slots = conv_states.size(0);

808 + for (int64_t bs = 0; bs < batch; ++bs) {

809 + const int32_t num_accepted = accepted_counts[bs];

810 + const int32_t conv_state_index = indices[bs];

811 + TORCH_CHECK(

812 + conv_state_index != pad_slot_id,

813 + "causal_conv1d_update_cpu: 3D x does not support pad slots.");

814 + TORCH_CHECK(

815 + conv_state_index >= 0 && conv_state_index < num_slots,

816 + "causal_conv1d_update_cpu: conv_state_indices out of range.");

817 + TORCH_CHECK(

818 + num_accepted >= 1 && num_accepted <= seqlen,

819 + "causal_conv1d_update_cpu: num_accepted_tokens must be in [1, "

820 + "seqlen].");

821 + TORCH_CHECK(

822 + num_accepted - 1 + width - 1 <= state_len,

823 + "causal_conv1d_update_cpu: history window exceeds conv_states.");

824 + }

825 +

826 + int64_t conv_state_slot_stride = conv_states.stride(0);

827 + at::Tensor out = at::empty_like(x);

828 + AT_DISPATCH_REDUCED_FLOATING_TYPES(

829 + scalar_type, "causal_conv1d_update_multi_kernel_impl", [&] {

830 + causal_conv1d_update_multi_kernel_impl<scalar_t>(

831 + out.data_ptr<scalar_t>(),

832 + x.data_ptr<scalar_t>(),

833 + conv_states.data_ptr<scalar_t>(),

834 + packed_w.data_ptr<scalar_t>(),

835 + conditional_data_ptr<scalar_t>(bias),

836 + accepted_counts,

837 + indices,

838 + silu_activation,

839 + batch,

840 + dim,

841 + seqlen,

842 + width,

843 + state_len,

844 + conv_state_slot_stride);

845 + });

846 + return out;

847 + }

848 +

849 + TORCH_CHECK(

850 + !num_accepted_tokens.has_value(),

851 + "causal_conv1d_update_cpu: num_accepted_tokens is only supported for 3D "

852 + "x.");

853

854 // adjust `conv_states` to be contiguous on `dim`

855 if (conv_states.stride(-2) != 1) {

856 + TORCH_CHECK(state_len == width - 1,

857 + "causal_conv1d_update_cpu: wide conv_states must be contiguous on dim.");

858 int64_t num_cache_lines = conv_states.size(0);

859 auto conv_states_copy = conv_states.clone();

860 conv_states.as_strided_({num_cache_lines, dim, width - 1}, {(width - 1) * dim, 1, dim});

csrc/cpu/torch_bindings.cpp

+3 -2

@@ -147,7 +147,7 @@ at::Tensor causal_conv1d_fwd_cpu(

147 at::Tensor causal_conv1d_update_cpu(

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

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

150 - bool silu_activation, const std::optional<at::Tensor>& cache_seqlens,

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

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

152 bool is_vnni);

153

@@ -516,7 +516,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

516 ops.def(

517 "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor "

518 "weight, Tensor? bias, bool silu_activation,"

519 - "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, "

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

520 + "pad_slot_id, "

521 "bool is_vnni) -> Tensor");

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

523 #endif

tests/kernels/mamba/cpu/test_cpu_gdn_ops.py

+435 -1

@@ -2,12 +2,14 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import functools

5 +import types

7 import pytest

8 import torch

9 import torch.nn.functional as F

11 import vllm._custom_ops as ops

12 +from vllm.model_executor.layers.mamba.ops.cpu import gdn_attention

13 from vllm.platforms import current_platform

14 from vllm.utils.torch_utils import set_random_seed

@@ -417,6 +419,207 @@ def _conv_inputs(total_tokens: int):

419 return x, weight, bias

420

421

422 +def _sd_conv_states(

423 + num_slots: int, state_len: int, dim: int = CONV_DIM

424 +) -> torch.Tensor:

425 + storage = torch.zeros(num_slots, state_len, dim, dtype=torch.bfloat16)

426 + return storage.transpose(1, 2)

427 +

428 +

429 +def _maybe_pack_conv_weight(weight: torch.Tensor, is_vnni: bool) -> torch.Tensor:

430 + return ops.causal_conv1d_weight_pack(weight) if is_vnni else weight

431 +

432 +

433 +@torch.inference_mode()

434 +def test_spec_aware_mixed_routing_preserves_token_order(

435 + monkeypatch: pytest.MonkeyPatch,

436 +) -> None:

437 + num_tokens = 4

438 + projection = torch.arange(num_tokens * 16, dtype=torch.float32).view(num_tokens, 16)

439 + mixed_qkv, b, a = projection[:, :4], projection[:, 4:6], projection[:, 6:8]

440 + assert all(not tensor.is_contiguous() for tensor in (mixed_qkv, b, a))

441 +

442 + spec_indices = torch.tensor([0, 2])

443 + nonspec_indices = torch.tensor([1, 3])

444 + metadata = types.SimpleNamespace(

445 + spec_sequence_masks=torch.ones(1, dtype=torch.bool),

446 + spec_token_indx=spec_indices,

447 + non_spec_token_indx=nonspec_indices,

448 + num_prefills=1,

449 + num_decodes=0,

450 + )

451 + routed = []

452 +

453 + def record(*args):

454 + routed.append(args[2:5])

455 + return args[2]

456 +

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

458 + monkeypatch.setattr(gdn_attention, "_spec_forward", record)

459 + monkeypatch.setattr(gdn_attention, "_spec_aware_nonspec_subset", record)

460 +

461 + layer = types.SimpleNamespace(

462 + kv_cache=[torch.empty(1, 4, 6), torch.empty(1, 1, 1, 1)]

463 + )

464 + core_attn_out = torch.empty_like(mixed_qkv)

465 + gdn_attention._cpu_gdn_attention_spec_aware(

466 + layer=layer,

467 + attn_metadata_i=metadata,

468 + mixed_qkv=mixed_qkv,

469 + b=b,

470 + a=a,

471 + core_attn_out=core_attn_out,

472 + width=CONV_KERNEL,

473 + state_len=6,

474 + )

475 +

476 + expected_inputs = (mixed_qkv, b, a)

477 + assert len(routed) == 2

478 + for actual_inputs, indices in zip(routed, (spec_indices, nonspec_indices)):

479 + for actual, expected in zip(actual_inputs, expected_inputs):

480 + assert actual.is_contiguous()

481 + torch.testing.assert_close(actual, expected.index_select(0, indices))

482 + torch.testing.assert_close(core_attn_out, mixed_qkv)

483 +

484 +

485 +@torch.inference_mode()

486 +def test_spec_aware_nonspec_materializes_state_indices(

487 + monkeypatch: pytest.MonkeyPatch,

488 +) -> None:

489 + block_table = torch.arange(8, dtype=torch.int32).view(2, 4)

490 + state_indices = block_table[:, 0]

491 + assert not state_indices.is_contiguous()

492 +

493 + metadata = types.SimpleNamespace(

494 + non_spec_state_indices_tensor=state_indices,

495 + non_spec_query_start_loc=torch.tensor([0, 2, 4], dtype=torch.int32),

496 + num_decodes=0,

497 + num_decode_tokens=0,

498 + num_prefills=2,

499 + num_prefill_tokens=4,

500 + has_initial_state=torch.tensor([False, False]),

501 + )

502 +

503 + recorded_indices = None

504 +

505 + def causal_conv1d_fwd_cpu(**kwargs):

506 + nonlocal recorded_indices

507 + recorded_indices = kwargs["cache_indices"]

508 + return kwargs["x"]

509 +

510 + def fused_gdn_gating_cpu(**kwargs):

511 + return kwargs["a"], kwargs["b"]

512 +

513 + def chunk_gated_delta_rule_cpu(**kwargs):

514 + out = torch.zeros(1, 4, 1, 1)

515 + return out, kwargs["initial_state"]

516 +

517 + monkeypatch.setattr(torch.cpu, "_is_amx_tile_supported", lambda: True)

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

519 + monkeypatch.setattr(

520 + gdn_attention.ops, "causal_conv1d_fwd_cpu", causal_conv1d_fwd_cpu

521 + )

522 + monkeypatch.setattr(gdn_attention.ops, "fused_gdn_gating_cpu", fused_gdn_gating_cpu)

523 + monkeypatch.setattr(

524 + gdn_attention.ops,

525 + "chunk_gated_delta_rule_cpu",

526 + chunk_gated_delta_rule_cpu,

527 + )

528 +

529 + layer = types.SimpleNamespace(

530 + activation="silu",

531 + conv1d=types.SimpleNamespace(weight=torch.empty(0), bias=None),

532 + A_log=torch.empty(0),

533 + dt_bias=torch.empty(0),

534 + rearrange_mixed_qkv=lambda x: (

535 + x[:, :1].view(1, 4, 1, 1),

536 + x[:, :1].view(1, 4, 1, 1),

537 + x[:, :1].view(1, 4, 1, 1),

538 + ),

539 + )

540 + gdn_attention._spec_aware_nonspec(

541 + layer=layer,

542 + attn_metadata_i=metadata,

543 + mixed_qkv=torch.zeros(4, 4),

544 + b=torch.zeros(4, 1),

545 + a=torch.zeros(4, 1),

546 + core_attn_out=torch.zeros(4, 1, 1),

547 + conv_buf=torch.zeros(8, 1, 6),

548 + ssm_state=torch.zeros(8, 1, 1, 1),

549 + width=4,

550 + )

551 +

552 + assert recorded_indices is not None

553 + assert recorded_indices.is_contiguous()

554 + torch.testing.assert_close(

555 + recorded_indices, torch.tensor([0, 4], dtype=torch.int32)

556 + )

557 +

558 +

559 +@torch.inference_mode()

560 +def test_spec_forward_prepares_native_conv_metadata(

561 + monkeypatch: pytest.MonkeyPatch,

562 +) -> None:

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

564 + state_indices = block_table[:, 0]

565 + accepted_counts = torch.tensor([1, 4], dtype=torch.int32)

566 + assert not state_indices.is_contiguous()

567 + metadata = types.SimpleNamespace(

568 + num_spec_decodes=2,

569 + spec_state_indices_tensor=block_table,

570 + spec_query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32),

571 + num_accepted_tokens=accepted_counts,

572 + )

573 + forwarded_indices = None

574 + forwarded_counts = None

575 +

576 + def causal_conv1d_update_cpu(**kwargs):

577 + nonlocal forwarded_counts, forwarded_indices

578 + forwarded_indices = kwargs["conv_state_indices"]

579 + forwarded_counts = kwargs["num_accepted_tokens"]

580 + return kwargs["x"]

581 +

582 + monkeypatch.setattr(torch.cpu, "_is_amx_tile_supported", lambda: True)

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

584 + monkeypatch.setattr(

585 + gdn_attention.ops, "causal_conv1d_update_cpu", causal_conv1d_update_cpu

586 + )

587 + monkeypatch.setattr(

588 + gdn_attention.ops,

589 + "fused_sigmoid_gating_delta_rule_update_spec_cpu",

590 + lambda **kwargs: kwargs["q"],

591 + )

592 +

593 + layer = types.SimpleNamespace(

594 + activation="silu",

595 + conv1d=types.SimpleNamespace(weight=torch.empty(1, CONV_KERNEL), bias=None),

596 + A_log=None,

597 + dt_bias=None,

598 + rearrange_mixed_qkv=lambda x: (x.unsqueeze(0),) * 3,

599 + )

600 + gdn_attention._spec_forward(

601 + layer=layer,

602 + attn_metadata_i=metadata,

603 + mixed_qkv_spec=torch.zeros(8, 1, dtype=torch.bfloat16),

604 + b_spec=torch.empty(0),

605 + a_spec=torch.empty(0),

606 + conv_buf=torch.empty(0),

607 + ssm_state=torch.empty(0),

608 + width=CONV_KERNEL,

609 + state_len=0,

610 + )

611 +

612 + expected = (

613 + (forwarded_indices, torch.tensor([0, 4], dtype=torch.int32)),

614 + (forwarded_counts, torch.tensor([1, 4], dtype=torch.int32)),

615 + )

616 + for actual, reference in expected:

617 + assert actual is not None

618 + assert actual.is_contiguous()

619 + assert actual.dtype == torch.int32

620 + torch.testing.assert_close(actual, reference)

621 +

622 +

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

624 @torch.inference_mode()

625 def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> None:

@@ -475,7 +678,182 @@ def test_causal_conv1d_torch_two_call_split(total_tokens: int, split: int) -> No

678

679 @pytest.mark.skipif(

680 not torch.cpu._is_amx_tile_supported(),

478 - reason="causal_conv1d_fwd_cpu requires AMX/AVX512",

681 + reason="requires AMX support",

682 +)

683 +@torch.inference_mode()

684 +def test_causal_conv1d_update_cpu_accepts_wide_state() -> None:

685 + state_len = CONV_KERNEL - 1

686 + wide_state_len = state_len + 5

687 + batch_size = 3

688 + is_vnni = True

689 + x, weight, bias = _conv_inputs(batch_size)

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

691 +

692 + narrow_state = _sd_conv_states(batch_size, state_len)

693 + narrow_state.copy_(

694 + tensor_cache(narrow_state.numel(), torch.bfloat16).view_as(narrow_state)

695 + )

696 + wide_state = _sd_conv_states(batch_size, wide_state_len)

697 + wide_state[:, :, :state_len].copy_(narrow_state)

698 + wide_state[:, :, state_len:].fill_(7)

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

700 +

701 + conv_weight = _maybe_pack_conv_weight(weight, is_vnni)

702 + out_narrow = ops.causal_conv1d_update_cpu(

703 + x=x,

704 + conv_states=narrow_state,

705 + weight=conv_weight,

706 + bias=bias,

707 + silu_activation=True,

708 + conv_state_indices=conv_state_indices,

709 + is_vnni=is_vnni,

710 + )

711 + out_wide = ops.causal_conv1d_update_cpu(

712 + x=x,

713 + conv_states=wide_state,

714 + weight=conv_weight,

715 + bias=bias,

716 + silu_activation=True,

717 + conv_state_indices=conv_state_indices,

718 + is_vnni=is_vnni,

719 + )

720 +

721 + torch.testing.assert_close(out_wide, out_narrow, atol=1e-2, rtol=1e-2)

722 + torch.testing.assert_close(

723 + wide_state[:, :, :state_len], narrow_state, atol=0, rtol=0

724 + )

725 + torch.testing.assert_close(wide_state[:, :, state_len:], wide_tail, atol=0, rtol=0)

726 +

727 +

728 +def _ref_causal_conv1d_update_cpu_multi(

729 + x: torch.Tensor,

730 + conv_states: torch.Tensor,

731 + weight: torch.Tensor,

732 + bias: torch.Tensor | None,

733 + silu_activation: bool,

734 + conv_state_indices: torch.Tensor,

735 + num_accepted_tokens: torch.Tensor,

736 +) -> torch.Tensor:

737 + batch_size, seq_len, dim = x.shape

738 + state_len = conv_states.size(2)

739 + conv_out = torch.empty_like(x)

740 + conv_weight = weight.unsqueeze(1)

741 +

742 + for i in range(batch_size):

743 + slot = int(conv_state_indices[i].item())

744 + offset = int(num_accepted_tokens[i].item()) - 1

745 + state = conv_states[slot]

746 + x_seq = x[i].transpose(0, 1).to(state.dtype)

747 + prior = state[:, offset : offset + CONV_KERNEL - 1]

748 + conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0)

749 + out = F.conv1d(conv_in, conv_weight, bias, groups=dim)[0]

750 + if silu_activation:

751 + out = F.silu(out)

752 + conv_out[i] = out.transpose(0, 1).to(conv_out.dtype)

753 + keep = state[:, offset + 1 : offset + 1 + (state_len - seq_len)]

754 + state.copy_(torch.cat([keep, x_seq], dim=-1))

755 +

756 + return conv_out

757 +

758 +

759 +@pytest.mark.skipif(

760 + not torch.cpu._is_amx_tile_supported(),

761 + reason="requires AMX support",

762 +)

763 +@pytest.mark.parametrize(

764 + ("batch_size, seq_len, accepted_counts, has_bias, silu_activation, is_vnni"),

765 + [

766 + (1, 1, [1], False, False, False),

767 + (1, 1, [1], True, True, True),

768 + (4, 4, [1, 2, 3, 4], False, True, False),

769 + (4, 4, [4, 3, 2, 1], True, False, True),

770 + (4, 16, [1, 5, 10, 16], False, False, True),

771 + (4, 16, [16, 10, 5, 1], True, True, False),

772 + ],

773 +)

774 +@torch.inference_mode()

775 +def test_causal_conv1d_update_cpu_multi_token_matches_python(

776 + batch_size: int,

777 + seq_len: int,

778 + accepted_counts: list[int],

779 + has_bias: bool,

780 + silu_activation: bool,

781 + is_vnni: bool,

782 +) -> None:

783 + dim = 96

784 + state_len = seq_len + 2

785 + x = tensor_cache(batch_size * seq_len * dim, torch.bfloat16).view(

786 + batch_size, seq_len, dim

787 + )

788 + weight = tensor_cache(dim * CONV_KERNEL, torch.bfloat16).view(dim, CONV_KERNEL)

789 + bias = tensor_cache(dim, torch.bfloat16) if has_bias else None

790 + conv_state_indices = torch.arange(batch_size - 1, -1, -1, dtype=torch.int32)

791 + num_accepted_tokens = torch.tensor(accepted_counts, dtype=torch.int32)

792 +

793 + conv_states_ref = _sd_conv_states(batch_size, state_len, dim)

794 + conv_states_ref.copy_(

795 + tensor_cache(conv_states_ref.numel(), torch.bfloat16).view_as(conv_states_ref)

796 + )

797 + conv_states = conv_states_ref.clone()

798 +

799 + conv_weight = _maybe_pack_conv_weight(weight, is_vnni)

800 + out = ops.causal_conv1d_update_cpu(

801 + x=x,

802 + conv_states=conv_states,

803 + weight=conv_weight,

804 + bias=bias,

805 + silu_activation=silu_activation,

806 + conv_state_indices=conv_state_indices,

807 + is_vnni=is_vnni,

808 + num_accepted_tokens=num_accepted_tokens,

809 + )

810 + ref_out = _ref_causal_conv1d_update_cpu_multi(

811 + x=x,

812 + conv_states=conv_states_ref,

813 + weight=weight,

814 + bias=bias,

815 + silu_activation=silu_activation,

816 + conv_state_indices=conv_state_indices,

817 + num_accepted_tokens=num_accepted_tokens,

818 + )

819 +

820 + torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)

821 + torch.testing.assert_close(conv_states, conv_states_ref, atol=0, rtol=0)

822 +

823 +

824 +@pytest.mark.skipif(

825 + not torch.cpu._is_amx_tile_supported(),

826 + reason="requires AMX support",

827 +)

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

829 +@torch.inference_mode()

830 +def test_causal_conv1d_update_cpu_rejects_invalid_accepted_count(

831 + num_accepted: int,

832 +) -> None:

833 + batch_size = 1

834 + seq_len = 16

835 + dim = 96

836 + state_len = seq_len + 2

837 + x = torch.zeros(batch_size, seq_len, dim, dtype=torch.bfloat16)

838 + weight = torch.zeros(dim, CONV_KERNEL, dtype=torch.bfloat16)

839 + conv_states = _sd_conv_states(batch_size, state_len, dim)

840 +

841 + with pytest.raises(RuntimeError, match="num_accepted_tokens must be in.*seqlen"):

842 + ops.causal_conv1d_update_cpu(

843 + x=x,

844 + conv_states=conv_states,

845 + weight=weight,

846 + bias=None,

847 + silu_activation=True,

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

849 + is_vnni=False,

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

851 + )

852 +

853 +

854 +@pytest.mark.skipif(

855 + not torch.cpu._is_amx_tile_supported(),

856 + reason="requires AMX support",

857 )

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

859 @torch.inference_mode()

@@ -515,6 +893,62 @@ def test_causal_conv1d_fwd_cpu_two_call_split(total_tokens: int, split: int) ->

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

894

895

896 +@pytest.mark.skipif(

897 + not torch.cpu._is_amx_tile_supported(),

898 + reason="requires AMX support",

899 +)

900 +@torch.inference_mode()

901 +def test_causal_conv1d_fwd_cpu_accepts_wide_state() -> None:

902 + state_len = CONV_KERNEL - 1

903 + wide_state_len = state_len + 5

904 + is_vnni = True

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

906 + total_tokens = sum(seq_lens)

907 + x, weight, bias = _conv_inputs(total_tokens)

908 + query_start_loc = torch.tensor([0, seq_lens[0], total_tokens], dtype=torch.int32)

909 + cache_indices = torch.tensor([2, 0], dtype=torch.int32)

910 + has_initial_state = torch.tensor([True, False])

911 +

912 + narrow_state = _sd_conv_states(3, state_len)

913 + narrow_state.copy_(

914 + tensor_cache(narrow_state.numel(), torch.bfloat16).view_as(narrow_state)

915 + )

916 + wide_state = _sd_conv_states(3, wide_state_len)

917 + wide_state[:, :, :state_len].copy_(narrow_state)

918 + wide_state[:, :, state_len:].fill_(7)

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

920 +

921 + conv_weight = _maybe_pack_conv_weight(weight, is_vnni)

922 + out_narrow = ops.causal_conv1d_fwd_cpu(

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

924 + weight=conv_weight,

925 + bias=bias,

926 + conv_states=narrow_state,

927 + query_start_loc=query_start_loc,

928 + cache_indices=cache_indices,

929 + has_initial_state=has_initial_state,

930 + silu_activation=True,

931 + is_vnni=is_vnni,

932 + )

933 + out_wide = ops.causal_conv1d_fwd_cpu(

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

935 + weight=conv_weight,

936 + bias=bias,

937 + conv_states=wide_state,

938 + query_start_loc=query_start_loc,

939 + cache_indices=cache_indices,

940 + has_initial_state=has_initial_state,

941 + silu_activation=True,

942 + is_vnni=is_vnni,

943 + )

944 +

945 + torch.testing.assert_close(out_wide, out_narrow, atol=1e-2, rtol=1e-2)

946 + torch.testing.assert_close(

947 + wide_state[:, :, :state_len], narrow_state, atol=0, rtol=0

948 + )

949 + torch.testing.assert_close(wide_state[:, :, state_len:], wide_tail, atol=0, rtol=0)

950 +

951 +

952 @torch.inference_mode()

953 def test_batch_memcpy_cpu_fallback() -> None:

954 """The ctypes batch_memcpy fallback (used when triton-cpu is absent) must

vllm/_custom_ops.py

+2 -1

@@ -3505,6 +3505,7 @@ def causal_conv1d_update_cpu(

3505 silu_activation: bool,

3506 conv_state_indices: torch.Tensor | None,

3507 is_vnni: bool,

3508 + num_accepted_tokens: torch.Tensor | None = None,

3509 ) -> torch.Tensor:

3510 return torch.ops._C.causal_conv1d_update_cpu(

3511 x,

@@ -3512,7 +3513,7 @@ def causal_conv1d_update_cpu(

3513 weight,

3514 bias,

3515 silu_activation,

3515 - None,

3516 + num_accepted_tokens,

3517 conv_state_indices,

3518 -1,

3519 is_vnni,

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

+154 -106

@@ -306,25 +306,21 @@ def _cpu_gdn_attention_spec_aware(

306 width: int,

307 state_len: int,

308 ) -> None:

309 - mixed_qkv = mixed_qkv.contiguous()

310 - a = a.contiguous()

311 - b = b.contiguous()

312 -

309 spec_sequence_masks = attn_metadata_i.spec_sequence_masks

310 conv_buf = _conv_buffer_view(layer) # (num_slots, dim, state_len)

311 ssm_state = _ssm_state_view(layer)

312

313 if spec_sequence_masks is None:

314 # No spec sequences in this batch (e.g. the prompt prefill step while

319 - # speculative decoding is configured). Process as prefill/decode using

320 - # torch conv (which only touches the first ``width-1`` columns of the

321 - # wide buffer, leaving the rolling history untouched).

315 + # speculative decoding is configured). Process as ordinary

316 + # prefill/decode while touching only the first ``width-1`` columns of

317 + # the wide buffer, leaving the rolling history untouched.

318 _spec_aware_nonspec(

319 layer,

320 attn_metadata_i,

325 - mixed_qkv,

326 - b,

327 - a,

321 + mixed_qkv.contiguous(),

322 + b.contiguous(),

323 + a.contiguous(),

324 core_attn_out,

325 conv_buf,

326 ssm_state,

@@ -339,9 +335,9 @@ def _cpu_gdn_attention_spec_aware(

335 num_decodes = attn_metadata_i.num_decodes

336

337 if num_prefills == 0 and num_decodes == 0:

342 - mixed_qkv_spec = mixed_qkv

343 - b_spec = b

344 - a_spec = a

338 + mixed_qkv_spec = mixed_qkv.contiguous()

339 + b_spec = b.contiguous()

340 + a_spec = a.contiguous()

341 spec_out_indx = None

342 else:

343 assert spec_token_indx is not None

@@ -405,50 +401,71 @@ def _spec_forward(

401 assert spec_qsl is not None

402 assert num_accepted is not None

403

408 - spec_qsl_cpu = spec_qsl[: num_spec_decodes + 1].to("cpu", torch.int64)

409 - num_acc_cpu = num_accepted[:num_spec_decodes].to("cpu", torch.int64)

404 + spec_qsl_cpu = spec_qsl[: num_spec_decodes + 1]

405 seq_starts = spec_qsl_cpu[:-1]

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

407

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

414 - w2d = _unpacked_conv_weight(layer) # (dim, width)

415 - dim = w2d.size(0)

416 - w = w2d.unsqueeze(1) # (dim, 1, width) for F.conv1d depthwise

409 + dim = mixed_qkv_spec.size(-1)

410 bias = layer.conv1d.bias

411 silu = layer.activation == "silu"

412

420 - conv_out = torch.empty_like(mixed_qkv_spec)

421 - col0 = spec_state_indices[:, 0].to("cpu", torch.int64)

422 - for i in range(num_spec_decodes):

423 - q_i = int(seq_lens[i].item())

424 - if q_i == 0:

425 - continue

426 - start = int(seq_starts[i].item())

427 - slot0 = int(col0[i].item())

428 - a_prev = int(num_acc_cpu[i].item())

429 - offset = a_prev - 1

430 - B = conv_buf[slot0] # (dim, state_len)

431 - x_seq = mixed_qkv_spec[start : start + q_i].transpose(0, 1).to(B.dtype)

432 - prior = B[:, offset : offset + (width - 1)]

433 - conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0) # (1, dim, w-1+q)

434 - out = F.conv1d(conv_in, w, bias, groups=dim)[0] # (dim, q_i)

435 - if silu:

436 - out = F.silu(out)

437 - conv_out[start : start + q_i] = out.transpose(0, 1).to(conv_out.dtype)

438 - # Roll the buffer: drop ``a_prev`` from the front, append the new

439 - # draft tokens, keep total length == state_len.

440 - keep = B[:, offset + 1 : offset + 1 + (state_len - q_i)]

441 - new_B = torch.cat([keep, x_seq], dim=-1)

442 - B.copy_(new_B)

413 + can_use_native_conv = (

414 + torch.cpu._is_amx_tile_supported()

415 + and not is_conv_state_dim_first()

416 + and width == 4

417 + and num_spec_decodes > 0

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

419 + and int(seq_lens[0].item()) > 0

420 + )

421 + if can_use_native_conv:

422 + q_i = int(seq_lens[0].item())

423 + conv_out = ops.causal_conv1d_update_cpu(

424 + x=mixed_qkv_spec.view(num_spec_decodes, q_i, dim),

425 + conv_states=conv_buf,

426 + weight=layer.conv1d.weight,

427 + bias=bias,

428 + silu_activation=silu,

429 + conv_state_indices=spec_state_indices[:num_spec_decodes, 0]

430 + .to("cpu", torch.int32)

431 + .contiguous(),

432 + is_vnni=True,

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

434 + ).view_as(mixed_qkv_spec)

435 + else:

436 + w = _unpacked_conv_weight(layer).unsqueeze(1)

437 + col0 = spec_state_indices[:num_spec_decodes, 0]

438 + num_acc_cpu = num_accepted[:num_spec_decodes]

439 + conv_out = torch.empty_like(mixed_qkv_spec)

440 + for i in range(num_spec_decodes):

441 + q_i = int(seq_lens[i].item())

442 + if q_i == 0:

443 + continue

444 + start = int(seq_starts[i].item())

445 + slot0 = int(col0[i].item())

446 + offset = int(num_acc_cpu[i].item()) - 1

447 + B = conv_buf[slot0] # (dim, state_len)

448 + x_seq = mixed_qkv_spec[start : start + q_i].transpose(0, 1).to(B.dtype)

449 + prior = B[:, offset : offset + (width - 1)]

450 + conv_in = torch.cat([prior, x_seq], dim=-1).unsqueeze(0)

451 + out = F.conv1d(conv_in, w, bias, groups=dim)[0] # (dim, q_i)

452 + if silu:

453 + out = F.silu(out)

454 + conv_out[start : start + q_i] = out.transpose(0, 1).to(conv_out.dtype)

455 + # Roll the buffer: drop the accepted history from the front, append

456 + # the new draft tokens, keep total length == state_len.

457 + keep = B[:, offset + 1 : offset + 1 + (state_len - q_i)]

458 + new_B = torch.cat([keep, x_seq], dim=-1)

459 + B.copy_(new_B)

460

461 # ---- 2. Recurrent (multi-slot SSM state) ----

462 # Single fused kernel call: it runs the recurrence over each sequence's

463 # draft tokens internally, resumes from slot ``num_accepted-1`` and stores

464 # the state after token ``t`` into slot ``t`` (rollback for the next step).

465 query, key, value = layer.rearrange_mixed_qkv(conv_out)

449 - query = query.squeeze(0).contiguous()

450 - key = key.squeeze(0).contiguous()

451 - value = value.squeeze(0).contiguous()

466 + query = query.squeeze(0)

467 + key = key.squeeze(0)

468 + value = value.squeeze(0)

469 spec_idx = spec_state_indices[:num_spec_decodes].to(torch.int32).contiguous()

470 num_acc = num_accepted[:num_spec_decodes].to(torch.int32).contiguous()

471 cu = spec_qsl[: num_spec_decodes + 1].to(torch.int32).contiguous()

@@ -458,8 +475,8 @@ def _spec_forward(

475 q=query,

476 k=key,

477 v=value,

461 - a=a_spec.contiguous(),

462 - b=b_spec.contiguous(),

478 + a=a_spec,

479 + b=b_spec,

480 initial_state_source=ssm_state,

481 spec_state_indices=spec_idx,

482 num_accepted_tokens=num_acc,

@@ -469,15 +486,6 @@ def _spec_forward(

486 return out_spec

487

488

472 -def core_attn_out_like(layer, mixed_qkv_spec: torch.Tensor) -> torch.Tensor:

473 - num_tokens = mixed_qkv_spec.size(0)

474 - return torch.zeros(

475 - (num_tokens, layer.num_v_heads // layer.tp_size, layer.head_v_dim),

476 - dtype=mixed_qkv_spec.dtype,

477 - device=mixed_qkv_spec.device,

478 - )

479 -

480 -

489 def _spec_aware_nonspec(

490 layer,

491 attn_metadata_i: GDNAttentionMetadata,

@@ -489,13 +497,19 @@ def _spec_aware_nonspec(

497 ssm_state: torch.Tensor,

498 width: int,

499 ) -> None:

492 - """Non-spec prefill/decode with a wide conv buffer (torch path)."""

500 + """Non-spec prefill/decode with a wide conv buffer."""

501 state_indices_tensor = attn_metadata_i.non_spec_state_indices_tensor

502 query_start_loc = attn_metadata_i.non_spec_query_start_loc

503 assert state_indices_tensor is not None

504 assert query_start_loc is not None

505 + state_indices_tensor = state_indices_tensor.contiguous()

506 +

507 + is_amx = torch.cpu._is_amx_tile_supported()

508 + if is_amx and is_conv_state_dim_first():

509 + raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")

510

498 - conv_weights = _unpacked_conv_weight(layer)

511 + if not is_amx:

512 + conv_weights = _unpacked_conv_weight(layer)

513

514 num_decodes = attn_metadata_i.num_decodes

515 num_decode_tokens = attn_metadata_i.num_decode_tokens

@@ -507,42 +521,48 @@ def _spec_aware_nonspec(

521 decode_b = b[:num_decode_tokens]

522 decode_a = a[:num_decode_tokens]

523 decode_state_indices = state_indices_tensor[:num_decodes]

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

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

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

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

514 - decode_mixed_qkv = causal_conv1d_update_torch(

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

516 - conv_state=decode_conv_state,

517 - weight=conv_weights,

518 - bias=layer.conv1d.bias,

519 - activation=layer.activation,

520 - ).squeeze(-1)

521 - conv_state_view[decode_state_indices] = decode_conv_state

522 - else:

523 - decode_mixed_qkv = causal_conv1d_update_cpu(

524 + if is_amx:

525 + decode_mixed_qkv = ops.causal_conv1d_update_cpu(

526 x=decode_mixed_qkv,

525 - conv_state=conv_buf[:, :, : width - 1],

526 - weight=conv_weights,

527 + conv_states=conv_buf,

528 + weight=layer.conv1d.weight,

529 bias=layer.conv1d.bias,

528 - activation=layer.activation,

530 + silu_activation=layer.activation == "silu",

531 conv_state_indices=decode_state_indices,

532 + is_vnni=True,

533 )

534 + else:

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

536 + conv_state_view = conv_buf[:, :, : width - 1]

537 + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:

538 + decode_conv_state = conv_state_view[decode_state_indices].contiguous()

539 + decode_mixed_qkv = causal_conv1d_update_torch(

540 + x=decode_mixed_qkv.unsqueeze(-1),

541 + conv_state=decode_conv_state,

542 + weight=conv_weights,

543 + bias=layer.conv1d.bias,

544 + activation=layer.activation,

545 + ).squeeze(-1)

546 + conv_state_view[decode_state_indices] = decode_conv_state

547 + else:

548 + decode_mixed_qkv = causal_conv1d_update_cpu(

549 + x=decode_mixed_qkv,

550 + conv_state=conv_state_view,

551 + weight=conv_weights,

552 + bias=layer.conv1d.bias,

553 + activation=layer.activation,

554 + conv_state_indices=decode_state_indices,

555 + )

556

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

533 - # rearrange_mixed_qkv can return views whose last dim is not

534 - # contiguous; the fused CPU kernel requires a contiguous last dim.

535 - query = query.contiguous()

536 - key = key.contiguous()

537 - value = value.contiguous()

558 attn_out = ops.fused_sigmoid_gating_delta_rule_update_cpu(

559 A_log=layer.A_log,

560 dt_bias=layer.dt_bias,

561 q=query,

562 k=key,

563 v=value,

544 - a=decode_a.contiguous(),

545 - b=decode_b.contiguous(),

564 + a=decode_a,

565 + b=decode_b,

566 initial_state_source=ssm_state,

567 initial_state_indices=decode_state_indices,

568 cu_seqlens=query_start_loc[: num_decodes + 1],

@@ -568,17 +588,29 @@ def _spec_aware_nonspec(

588 prefill_has_initial_state = has_initial_state[

589 num_decodes : num_decodes + num_prefills

590 ]

571 - # ``causal_conv1d_torch`` only touches columns [:width-1] of the buffer.

572 - prefill_mixed_qkv = causal_conv1d_torch(

573 - x=prefill_mixed_qkv.transpose(0, 1),

574 - weight=conv_weights,

575 - bias=layer.conv1d.bias,

576 - conv_states=conv_buf,

577 - query_start_loc=prefill_query_start_loc,

578 - cache_indices=prefill_state_indices,

579 - has_initial_state=prefill_has_initial_state,

580 - activation=layer.activation,

581 - ).transpose(0, 1)

591 + if is_amx:

592 + prefill_mixed_qkv = ops.causal_conv1d_fwd_cpu(

593 + x=prefill_mixed_qkv.transpose(0, 1),

594 + weight=layer.conv1d.weight,

595 + bias=layer.conv1d.bias,

596 + conv_states=conv_buf,

597 + query_start_loc=prefill_query_start_loc,

598 + cache_indices=prefill_state_indices,

599 + has_initial_state=prefill_has_initial_state,

600 + silu_activation=layer.activation == "silu",

601 + is_vnni=True,

602 + ).transpose(0, 1)

603 + else:

604 + prefill_mixed_qkv = causal_conv1d_torch(

605 + x=prefill_mixed_qkv.transpose(0, 1),

606 + weight=conv_weights,

607 + bias=layer.conv1d.bias,

608 + conv_states=conv_buf,

609 + query_start_loc=prefill_query_start_loc,

610 + cache_indices=prefill_state_indices,

611 + has_initial_state=prefill_has_initial_state,

612 + activation=layer.activation,

613 + ).transpose(0, 1)

614

615 query, key, value = layer.rearrange_mixed_qkv(prefill_mixed_qkv)

616 g, beta = ops.fused_gdn_gating_cpu(

@@ -618,24 +650,41 @@ def _spec_aware_nonspec_subset(

650

651 Returns outputs ordered like ``non_spec_token_indx``.

652 """

621 - out = core_attn_out_like(layer, mixed_qkv)

653 has_initial_state = attn_metadata_i.has_initial_state

654 prefill_state_indices = attn_metadata_i.prefill_state_indices

655 prefill_qsl = attn_metadata_i.prefill_query_start_loc

656 assert prefill_state_indices is not None and prefill_qsl is not None

657 assert has_initial_state is not None

658 + prefill_state_indices = prefill_state_indices.contiguous()

659 +

660 + is_amx = torch.cpu._is_amx_tile_supported()

661 + if is_amx and is_conv_state_dim_first():

662 + raise RuntimeError("AMX GDN attention requires `SD` conv_state layout.")

663

628 - conv_weights = _unpacked_conv_weight(layer)

629 - conv_out = causal_conv1d_torch(

630 - x=mixed_qkv.transpose(0, 1),

631 - weight=conv_weights,

632 - bias=layer.conv1d.bias,

633 - conv_states=conv_buf,

634 - query_start_loc=prefill_qsl,

635 - cache_indices=prefill_state_indices,

636 - has_initial_state=has_initial_state,

637 - activation=layer.activation,

638 - ).transpose(0, 1)

664 + if is_amx:

665 + conv_out = ops.causal_conv1d_fwd_cpu(

666 + x=mixed_qkv.transpose(0, 1),

667 + weight=layer.conv1d.weight,

668 + bias=layer.conv1d.bias,

669 + conv_states=conv_buf,

670 + query_start_loc=prefill_qsl,

671 + cache_indices=prefill_state_indices,

672 + has_initial_state=has_initial_state,

673 + silu_activation=layer.activation == "silu",

674 + is_vnni=True,

675 + ).transpose(0, 1)

676 + else:

677 + conv_weights = _unpacked_conv_weight(layer)

678 + conv_out = causal_conv1d_torch(

679 + x=mixed_qkv.transpose(0, 1),

680 + weight=conv_weights,

681 + bias=layer.conv1d.bias,

682 + conv_states=conv_buf,

683 + query_start_loc=prefill_qsl,

684 + cache_indices=prefill_state_indices,

685 + has_initial_state=has_initial_state,

686 + activation=layer.activation,

687 + ).transpose(0, 1)

688

689 query, key, value = layer.rearrange_mixed_qkv(conv_out)

690 g, beta = ops.fused_gdn_gating_cpu(

@@ -658,8 +707,7 @@ def _spec_aware_nonspec_subset(

707 ssm_state[prefill_state_indices] = last_recurrent_state.to(

708 ssm_state.dtype, copy=False

709 )

661 - out[:] = attn_out.squeeze(0)

662 - return out

710 + return attn_out.squeeze(0)

711

712

713 def cpu_gdn_attention_core_fake(