vllm-project--vllm/tests/kernels/mamba/test_cpu_short_conv.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest
import torch

from vllm.config import CompilationConfig, VllmConfig
from vllm.forward_context import set_forward_context
from vllm.model_executor.layers.mamba.short_conv import ShortConv
from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm
from vllm.platforms import current_platform
from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadata

if not current_platform.is_cpu():
    pytest.skip("skipping CPU-only tests", allow_module_level=True)


@pytest.fixture(autouse=True)
def mock_dist():
    with (
        patch(
            "vllm.model_executor.layers.linear.get_tensor_model_parallel_rank",
            return_value=0,
        ),
        patch(
            "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
            return_value=1,
        ),
        patch(
            "vllm.distributed.parallel_state.model_parallel_is_initialized",
            return_value=True,
        ),
        patch(
            "vllm.distributed.parallel_state.get_tp_group",
            return_value=MagicMock(rank_in_group=0),
        ),
    ):
        yield


@pytest.fixture
def vllm_config():
    # ShortConv only needs compilation_config from the current vLLM config, so a
    # minimal config (model_config=None) avoids mocking ModelConfig and the
    # associated VllmConfig validation churn.
    return VllmConfig(compilation_config=CompilationConfig())


def test_short_conv_forward_native_prefill(vllm_config):
    prefix = "test_layer"
    config = SimpleNamespace(conv_L_cache=4, conv_bias=True)
    dim = 16

    from vllm.config import set_current_vllm_config

    with set_current_vllm_config(vllm_config):
        layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix)

    layer.to("cpu")
    # vLLM Linear layers allocate weights with torch.empty (uninitialized).
    # On ARM these come back as zero-filled pages, so in_proj output is zero and
    # the prefill state stays zero. Seed + init to make the test platform-safe.
    torch.manual_seed(0)
    for p in layer.parameters():
        torch.nn.init.normal_(p)
    dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False)
    dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False)

    # Mock AttentionMetadata
    num_prefills = 1
    num_prefill_tokens = 5
    query_start_loc_p = torch.tensor([0, 5], dtype=torch.int32)
    state_indices_tensor_p = torch.tensor([0], dtype=torch.int32)

    # ShortConvAttentionMetadata
    attn_metadata = ShortConvAttentionMetadata(
        num_prefills=num_prefills,
        num_prefill_tokens=num_prefill_tokens,
        num_decodes=0,
        num_decode_tokens=0,
        num_reqs=1,
        query_start_loc_p=query_start_loc_p,
        has_initial_states_p=torch.tensor([False]),
        state_indices_tensor_p=state_indices_tensor_p,
        state_indices_tensor_d=torch.empty((0, 1), dtype=torch.int32),
        num_accepted_tokens=None,
        query_start_loc_d=None,
        block_idx_last_scheduled_token=None,
        block_idx_first_scheduled_token_p=None,
        block_idx_last_computed_token=None,
        block_idx_last_scheduled_token_prev_step=None,
        num_computed_tokens_p=None,
        seq_lens=torch.tensor([5]),
    )

    # Mock KV cache
    # conv_state shape (num_blocks, L_cache - 1, dim)
    conv_state = torch.zeros((1, config.conv_L_cache - 1, dim))
    layer.kv_cache = (conv_state,)

    hidden_states = torch.randn((num_prefill_tokens, dim))
    output = torch.zeros_like(hidden_states)

    attn_metadata_dict = {prefix: attn_metadata}
    with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config):
        layer.forward_native(hidden_states, output)

    # Check if KV cache was updated
    assert not torch.allclose(conv_state, torch.zeros_like(conv_state))


def test_short_conv_forward_native_decode(vllm_config):
    prefix = "test_layer_decode"
    config = SimpleNamespace(conv_L_cache=4, conv_bias=True)
    dim = 16

    from vllm.config import set_current_vllm_config

    with set_current_vllm_config(vllm_config):
        layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix)

    layer.to("cpu")
    torch.manual_seed(0)
    for p in layer.parameters():
        torch.nn.init.normal_(p)
    dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False)
    dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False)

    # Mock AttentionMetadata for 2 decode requests
    num_decodes = 2
    state_indices_tensor_d = torch.tensor([0, 1], dtype=torch.int32)

    attn_metadata = ShortConvAttentionMetadata(
        num_prefills=0,
        num_prefill_tokens=0,
        num_decodes=num_decodes,
        num_decode_tokens=num_decodes,
        num_reqs=num_decodes,
        query_start_loc_p=None,
        has_initial_states_p=None,
        state_indices_tensor_p=torch.empty((0,), dtype=torch.int32),
        state_indices_tensor_d=state_indices_tensor_d,
        num_accepted_tokens=None,
        query_start_loc_d=torch.tensor([0, 1, 2], dtype=torch.int32),
        block_idx_last_scheduled_token=None,
        block_idx_first_scheduled_token_p=None,
        block_idx_last_computed_token=None,
        block_idx_last_scheduled_token_prev_step=None,
        num_computed_tokens_p=None,
        seq_lens=torch.tensor([1, 1]),
    )

    # Mock KV cache (2 blocks for 2 requests)
    conv_state = torch.randn((2, config.conv_L_cache - 1, dim))
    layer.kv_cache = (conv_state,)

    hidden_states = torch.randn((num_decodes, dim))
    output = torch.zeros_like(hidden_states)

    old_conv_state = conv_state.clone()

    attn_metadata_dict = {prefix: attn_metadata}
    with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config):
        layer.forward_native(hidden_states, output)

    # Check if KV cache was updated
    assert not torch.allclose(conv_state, old_conv_state)


def test_dispatch_cpu_unquantized_gemm_conv_layer():
    # Convolution layers have >2D weights; dispatch should skip them gracefully.
    # Shape/dtype are AMX-pack safe (bf16, width==4, dim % block_size == 0) so
    # the AMX prepack branch does not raise on AMX-capable CPUs.
    class MockConvLayer(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.weight = torch.nn.Parameter(
                torch.randn(32, 1, 4, dtype=torch.bfloat16)
            )
            self.bias = torch.nn.Parameter(torch.randn(32, dtype=torch.bfloat16))

    layer = MockConvLayer()
    # The ndim != 2 guard returns early without raising.
    dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
    # No cpu_linear set — conv layers are handled elsewhere.
    assert not hasattr(layer, "cpu_linear")