| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | from unittest.mock import MagicMock, patch |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import CompilationConfig, VllmConfig |
| 11 | from vllm.forward_context import set_forward_context |
| 12 | from vllm.model_executor.layers.mamba.short_conv import ShortConv |
| 13 | from vllm.model_executor.layers.utils import dispatch_cpu_unquantized_gemm |
| 14 | from vllm.platforms import current_platform |
| 15 | from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadata |
| 16 | |
| 17 | if not current_platform.is_cpu(): |
| 18 | pytest.skip("skipping CPU-only tests", allow_module_level=True) |
| 19 | |
| 20 | |
| 21 | @pytest.fixture(autouse=True) |
| 22 | def mock_dist(): |
| 23 | with ( |
| 24 | patch( |
| 25 | "vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", |
| 26 | return_value=0, |
| 27 | ), |
| 28 | patch( |
| 29 | "vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", |
| 30 | return_value=1, |
| 31 | ), |
| 32 | patch( |
| 33 | "vllm.distributed.parallel_state.model_parallel_is_initialized", |
| 34 | return_value=True, |
| 35 | ), |
| 36 | patch( |
| 37 | "vllm.distributed.parallel_state.get_tp_group", |
| 38 | return_value=MagicMock(rank_in_group=0), |
| 39 | ), |
| 40 | ): |
| 41 | yield |
| 42 | |
| 43 | |
| 44 | @pytest.fixture |
| 45 | def vllm_config(): |
| 46 | # ShortConv only needs compilation_config from the current vLLM config, so a |
| 47 | # minimal config (model_config=None) avoids mocking ModelConfig and the |
| 48 | # associated VllmConfig validation churn. |
| 49 | return VllmConfig(compilation_config=CompilationConfig()) |
| 50 | |
| 51 | |
| 52 | def test_short_conv_forward_native_prefill(vllm_config): |
| 53 | prefix = "test_layer" |
| 54 | config = SimpleNamespace(conv_L_cache=4, conv_bias=True) |
| 55 | dim = 16 |
| 56 | |
| 57 | from vllm.config import set_current_vllm_config |
| 58 | |
| 59 | with set_current_vllm_config(vllm_config): |
| 60 | layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix) |
| 61 | |
| 62 | layer.to("cpu") |
| 63 | # vLLM Linear layers allocate weights with torch.empty (uninitialized). |
| 64 | # On ARM these come back as zero-filled pages, so in_proj output is zero and |
| 65 | # the prefill state stays zero. Seed + init to make the test platform-safe. |
| 66 | torch.manual_seed(0) |
| 67 | for p in layer.parameters(): |
| 68 | torch.nn.init.normal_(p) |
| 69 | dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False) |
| 70 | dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False) |
| 71 | |
| 72 | # Mock AttentionMetadata |
| 73 | num_prefills = 1 |
| 74 | num_prefill_tokens = 5 |
| 75 | query_start_loc_p = torch.tensor([0, 5], dtype=torch.int32) |
| 76 | state_indices_tensor_p = torch.tensor([0], dtype=torch.int32) |
| 77 | |
| 78 | # ShortConvAttentionMetadata |
| 79 | attn_metadata = ShortConvAttentionMetadata( |
| 80 | num_prefills=num_prefills, |
| 81 | num_prefill_tokens=num_prefill_tokens, |
| 82 | num_decodes=0, |
| 83 | num_decode_tokens=0, |
| 84 | num_reqs=1, |
| 85 | query_start_loc_p=query_start_loc_p, |
| 86 | has_initial_states_p=torch.tensor([False]), |
| 87 | state_indices_tensor_p=state_indices_tensor_p, |
| 88 | state_indices_tensor_d=torch.empty((0, 1), dtype=torch.int32), |
| 89 | num_accepted_tokens=None, |
| 90 | query_start_loc_d=None, |
| 91 | seq_lens=torch.tensor([5]), |
| 92 | ) |
| 93 | |
| 94 | # Mock KV cache |
| 95 | # conv_state shape (num_blocks, L_cache - 1, dim) |
| 96 | conv_state = torch.zeros((1, config.conv_L_cache - 1, dim)) |
| 97 | layer.kv_cache = (conv_state,) |
| 98 | |
| 99 | hidden_states = torch.randn((num_prefill_tokens, dim)) |
| 100 | output = torch.zeros_like(hidden_states) |
| 101 | |
| 102 | attn_metadata_dict = {prefix: attn_metadata} |
| 103 | with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config): |
| 104 | layer.forward_native(hidden_states, output) |
| 105 | |
| 106 | # Check if KV cache was updated |
| 107 | assert not torch.allclose(conv_state, torch.zeros_like(conv_state)) |
| 108 | |
| 109 | |
| 110 | def test_short_conv_forward_native_decode(vllm_config): |
| 111 | prefix = "test_layer_decode" |
| 112 | config = SimpleNamespace(conv_L_cache=4, conv_bias=True) |
| 113 | dim = 16 |
| 114 | |
| 115 | from vllm.config import set_current_vllm_config |
| 116 | |
| 117 | with set_current_vllm_config(vllm_config): |
| 118 | layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix) |
| 119 | |
| 120 | layer.to("cpu") |
| 121 | torch.manual_seed(0) |
| 122 | for p in layer.parameters(): |
| 123 | torch.nn.init.normal_(p) |
| 124 | dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False) |
| 125 | dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False) |
| 126 | |
| 127 | # Mock AttentionMetadata for 2 decode requests |
| 128 | num_decodes = 2 |
| 129 | state_indices_tensor_d = torch.tensor([0, 1], dtype=torch.int32) |
| 130 | |
| 131 | attn_metadata = ShortConvAttentionMetadata( |
| 132 | num_prefills=0, |
| 133 | num_prefill_tokens=0, |
| 134 | num_decodes=num_decodes, |
| 135 | num_decode_tokens=num_decodes, |
| 136 | num_reqs=num_decodes, |
| 137 | query_start_loc_p=None, |
| 138 | has_initial_states_p=None, |
| 139 | state_indices_tensor_p=torch.empty((0,), dtype=torch.int32), |
| 140 | state_indices_tensor_d=state_indices_tensor_d, |
| 141 | num_accepted_tokens=None, |
| 142 | query_start_loc_d=torch.tensor([0, 1, 2], dtype=torch.int32), |
| 143 | seq_lens=torch.tensor([1, 1]), |
| 144 | ) |
| 145 | |
| 146 | # Mock KV cache (2 blocks for 2 requests) |
| 147 | conv_state = torch.randn((2, config.conv_L_cache - 1, dim)) |
| 148 | layer.kv_cache = (conv_state,) |
| 149 | |
| 150 | hidden_states = torch.randn((num_decodes, dim)) |
| 151 | output = torch.zeros_like(hidden_states) |
| 152 | |
| 153 | old_conv_state = conv_state.clone() |
| 154 | |
| 155 | attn_metadata_dict = {prefix: attn_metadata} |
| 156 | with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config): |
| 157 | layer.forward_native(hidden_states, output) |
| 158 | |
| 159 | # Check if KV cache was updated |
| 160 | assert not torch.allclose(conv_state, old_conv_state) |
| 161 | |
| 162 | |
| 163 | def test_dispatch_cpu_unquantized_gemm_conv_layer(): |
| 164 | # Convolution layers have >2D weights; dispatch should skip them gracefully. |
| 165 | # Shape/dtype are AMX-pack safe (bf16, width==4, dim % block_size == 0) so |
| 166 | # the AMX prepack branch does not raise on AMX-capable CPUs. |
| 167 | class MockConvLayer(torch.nn.Module): |
| 168 | def __init__(self): |
| 169 | super().__init__() |
| 170 | self.weight = torch.nn.Parameter( |
| 171 | torch.randn(32, 1, 4, dtype=torch.bfloat16) |
| 172 | ) |
| 173 | self.bias = torch.nn.Parameter(torch.randn(32, dtype=torch.bfloat16)) |
| 174 | |
| 175 | layer = MockConvLayer() |
| 176 | # The ndim != 2 guard returns early without raising. |
| 177 | dispatch_cpu_unquantized_gemm(layer, remove_weight=False) |
| 178 | # No cpu_linear set — conv layers are handled elsewhere. |
| 179 | assert not hasattr(layer, "cpu_linear") |