| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Attention metadata geometry must come from the builder's own group.""" |
| 4 | |
| 5 | from types import SimpleNamespace |
| 6 | from unittest.mock import MagicMock, patch |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.v1.attention.backends.cpu_attn import ( |
| 13 | CPUAttentionBackendImpl, |
| 14 | CPUAttentionMetadataBuilder, |
| 15 | ) |
| 16 | from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadataBuilder |
| 17 | |
| 18 | requires_cpu = pytest.mark.skipif( |
| 19 | not current_platform.is_cpu(), reason="CPU attention backend" |
| 20 | ) |
| 21 | |
| 22 | # Laguna's shape: 48 query heads model-wide, 64 on its sliding layers, both |
| 23 | # against 8 KV heads. |
| 24 | MODEL_WIDE_NUM_HEADS = 48 |
| 25 | NUM_KV_HEADS = 8 |
| 26 | |
| 27 | |
| 28 | def _layers(layer_num_heads: list[int]): |
| 29 | """Stand-in attention layers, one per head count, as one attention group.""" |
| 30 | return { |
| 31 | f"layer_{i}": SimpleNamespace( |
| 32 | impl=MagicMock( |
| 33 | spec=CPUAttentionBackendImpl, |
| 34 | num_heads=num_heads, |
| 35 | sliding_window=None, |
| 36 | ) |
| 37 | ) |
| 38 | for i, num_heads in enumerate(layer_num_heads) |
| 39 | } |
| 40 | |
| 41 | |
| 42 | def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder: |
| 43 | layers = _layers(layer_num_heads) |
| 44 | vllm_config = MagicMock() |
| 45 | vllm_config.model_config.dtype = torch.bfloat16 |
| 46 | vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS |
| 47 | vllm_config.cache_config.cache_dtype = "auto" |
| 48 | kv_cache_spec = SimpleNamespace( |
| 49 | num_kv_heads=NUM_KV_HEADS, head_size=64, block_size=16 |
| 50 | ) |
| 51 | |
| 52 | with ( |
| 53 | patch( |
| 54 | "vllm.v1.attention.backends.utils.get_layers_from_vllm_config", |
| 55 | return_value=layers, |
| 56 | ), |
| 57 | patch( |
| 58 | "vllm.v1.attention.backends.cpu_attn.get_layers_from_vllm_config", |
| 59 | return_value=layers, |
| 60 | ), |
| 61 | ): |
| 62 | return CPUAttentionMetadataBuilder( |
| 63 | kv_cache_spec=kv_cache_spec, |
| 64 | layer_names=list(layers), |
| 65 | vllm_config=vllm_config, |
| 66 | device=torch.device("cpu"), |
| 67 | ) |
| 68 | |
| 69 | |
| 70 | @requires_cpu |
| 71 | @pytest.mark.parametrize("group_num_heads", [MODEL_WIDE_NUM_HEADS, 64, 16]) |
| 72 | def test_num_heads_comes_from_the_group(group_num_heads): |
| 73 | """The group's own count wins, even when it is not the model-wide one.""" |
| 74 | builder = _build([group_num_heads, group_num_heads]) |
| 75 | assert builder.num_heads == group_num_heads |
| 76 | |
| 77 | |
| 78 | @requires_cpu |
| 79 | def test_mixed_head_counts_in_one_group_are_rejected(): |
| 80 | """Grouping guarantees uniformity; a mixed group means that broke.""" |
| 81 | with pytest.raises(AssertionError, match="share num_heads"): |
| 82 | _build([MODEL_WIDE_NUM_HEADS, 64]) |
| 83 | |
| 84 | |
| 85 | def test_flash_attention_geometry_comes_from_the_group(): |
| 86 | """FA3 must use group geometry for layers without an ``impl`` wrapper.""" |
| 87 | layers = {f"layer_{i}": SimpleNamespace(num_heads=16) for i in range(2)} |
| 88 | vllm_config = MagicMock() |
| 89 | vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS |
| 90 | vllm_config.model_config.get_num_kv_heads.return_value = NUM_KV_HEADS |
| 91 | vllm_config.model_config.get_head_size.return_value = 128 |
| 92 | vllm_config.model_config.rswa_window = None |
| 93 | vllm_config.model_config.is_mm_prefix_lm = False |
| 94 | vllm_config.parallel_config.cp_kv_cache_interleave_size = 1 |
| 95 | vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs.return_value = ( |
| 96 | False |
| 97 | ) |
| 98 | vllm_config.compilation_config.max_cudagraph_capture_size = None |
| 99 | kv_cache_spec = SimpleNamespace( |
| 100 | block_size=16, |
| 101 | num_kv_heads=2, |
| 102 | head_size=64, |
| 103 | dtype=torch.bfloat16, |
| 104 | dcp_sharded=True, |
| 105 | ) |
| 106 | |
| 107 | with ( |
| 108 | patch( |
| 109 | "vllm.v1.attention.backends.utils.get_layers_from_vllm_config", |
| 110 | return_value=layers, |
| 111 | ), |
| 112 | patch( |
| 113 | "vllm.distributed.parallel_state.get_dcp_group", |
| 114 | side_effect=AssertionError, |
| 115 | ), |
| 116 | patch( |
| 117 | "vllm.v1.attention.backends.flash_attn.get_flash_attn_version", |
| 118 | return_value=3, |
| 119 | ), |
| 120 | ): |
| 121 | builder = FlashAttentionMetadataBuilder( |
| 122 | kv_cache_spec=kv_cache_spec, |
| 123 | layer_names=list(layers), |
| 124 | vllm_config=vllm_config, |
| 125 | device=torch.device("cpu"), |
| 126 | ) |
| 127 | |
| 128 | assert builder.aot_schedule |
| 129 | assert builder.num_heads_q == 16 |
| 130 | assert builder.num_heads_kv == 2 |
| 131 | assert builder.headdim == 64 |