| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Utility functions for attention-related v1 tests.""" |
| 4 | |
| 5 | from dataclasses import dataclass |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.config import ( |
| 11 | CacheConfig, |
| 12 | CompilationConfig, |
| 13 | DeviceConfig, |
| 14 | LoadConfig, |
| 15 | ModelConfig, |
| 16 | ParallelConfig, |
| 17 | SchedulerConfig, |
| 18 | VllmConfig, |
| 19 | ) |
| 20 | from vllm.config.model import ModelDType |
| 21 | from vllm.v1.attention.backend import ( |
| 22 | AttentionImpl, |
| 23 | AttentionMetadataBuilder, |
| 24 | AttentionType, |
| 25 | CommonAttentionMetadata, |
| 26 | ) |
| 27 | from vllm.v1.attention.backends.mamba_attn import ( |
| 28 | BaseMambaAttentionMetadata, |
| 29 | BaseMambaAttentionMetadataBuilder, |
| 30 | ) |
| 31 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 32 | from vllm.v1.kv_cache_interface import ( |
| 33 | EncoderOnlyAttentionSpec, |
| 34 | FullAttentionSpec, |
| 35 | KVCacheLayout, |
| 36 | KVCacheSpec, |
| 37 | KVCacheTensor, |
| 38 | MambaSpec, |
| 39 | compute_layout_strides, |
| 40 | create_kv_cache_views, |
| 41 | get_kv_quant_mode, |
| 42 | ) |
| 43 | |
| 44 | |
| 45 | @dataclass |
| 46 | class BatchSpec: |
| 47 | """Specification for a batch configuration (workload shape only).""" |
| 48 | |
| 49 | seq_lens: list[int] |
| 50 | query_lens: list[int] |
| 51 | |
| 52 | name: str = "unnamed" |
| 53 | |
| 54 | @property |
| 55 | def batch_size(self): |
| 56 | return len(self.seq_lens) |
| 57 | |
| 58 | def __post_init__(self): |
| 59 | assert len(self.seq_lens) == len(self.query_lens) |
| 60 | |
| 61 | def compute_num_tokens(self): |
| 62 | return sum(self.query_lens) |
| 63 | |
| 64 | |
| 65 | def create_common_attn_metadata( |
| 66 | batch_spec: BatchSpec, |
| 67 | block_size: int, |
| 68 | device: torch.device, |
| 69 | max_block_idx: int = 1000, |
| 70 | arange_block_indices: bool = False, |
| 71 | ) -> CommonAttentionMetadata: |
| 72 | """Create CommonAttentionMetadata from a BatchSpec and ModelParams.""" |
| 73 | # Create query start locations |
| 74 | query_start_loc = torch.zeros( |
| 75 | batch_spec.batch_size + 1, dtype=torch.int32, device=device |
| 76 | ) |
| 77 | query_start_loc[1:] = torch.tensor( |
| 78 | batch_spec.query_lens, dtype=torch.int32, device=device |
| 79 | ).cumsum(0) |
| 80 | query_start_loc_cpu = query_start_loc.cpu() |
| 81 | num_tokens = batch_spec.compute_num_tokens() |
| 82 | |
| 83 | # Create sequence lengths |
| 84 | seq_lens = torch.tensor(batch_spec.seq_lens, dtype=torch.int32, device=device) |
| 85 | seq_lens_cpu = seq_lens.cpu() |
| 86 | max_seq_len = int(seq_lens_cpu.max()) |
| 87 | |
| 88 | # Create block table and slot mapping |
| 89 | max_blocks = (max(batch_spec.seq_lens) + block_size - 1) // block_size |
| 90 | if arange_block_indices: |
| 91 | num_blocks = batch_spec.batch_size * max_blocks |
| 92 | block_table_tensor = torch.arange( |
| 93 | num_blocks, dtype=torch.int32, device=device |
| 94 | ).view(batch_spec.batch_size, max_blocks) |
| 95 | slot_mapping = torch.arange(num_tokens, dtype=torch.int64, device=device).view( |
| 96 | num_tokens |
| 97 | ) |
| 98 | else: |
| 99 | block_table_tensor = torch.randint( |
| 100 | 0, |
| 101 | max_block_idx, |
| 102 | (batch_spec.batch_size, max_blocks), |
| 103 | dtype=torch.int32, |
| 104 | device=device, |
| 105 | ) |
| 106 | slot_mapping = torch.randint( |
| 107 | 0, max_block_idx, (num_tokens,), dtype=torch.int64, device=device |
| 108 | ) |
| 109 | |
| 110 | # Calculate max query length |
| 111 | max_query_len = max(batch_spec.query_lens) |
| 112 | |
| 113 | return CommonAttentionMetadata( |
| 114 | query_start_loc=query_start_loc, |
| 115 | query_start_loc_cpu=query_start_loc_cpu, |
| 116 | seq_lens=seq_lens, |
| 117 | seq_lens_cpu_upper_bound=seq_lens_cpu, |
| 118 | num_reqs=batch_spec.batch_size, |
| 119 | num_actual_tokens=num_tokens, |
| 120 | max_query_len=max_query_len, |
| 121 | max_seq_len=max_seq_len, |
| 122 | block_table_tensor=block_table_tensor, |
| 123 | slot_mapping=slot_mapping, |
| 124 | causal=True, |
| 125 | ) |
| 126 | |
| 127 | |
| 128 | def try_get_attention_backend( |
| 129 | backend: AttentionBackendEnum, |
| 130 | ) -> tuple[type[AttentionMetadataBuilder], type[AttentionImpl]]: |
| 131 | """Try to get the attention backend class, skipping test if not found.""" |
| 132 | try: |
| 133 | backend_class = backend.get_class() |
| 134 | return backend_class.get_builder_cls(), backend_class.get_impl_cls() |
| 135 | except ImportError as e: |
| 136 | pytest.skip(f"{backend.name} not available: {e}") |
| 137 | raise AssertionError("unreachable") from None |
| 138 | |
| 139 | |
| 140 | def try_backend_includes_kv_cache_update( |
| 141 | backend: AttentionBackendEnum, |
| 142 | ) -> bool: |
| 143 | """Try to get the attention backend class, skipping test if not found.""" |
| 144 | try: |
| 145 | backend_class = backend.get_class() |
| 146 | return backend_class.forward_includes_kv_cache_update |
| 147 | except ImportError as e: |
| 148 | pytest.skip(f"{backend.name} not available: {e}") |
| 149 | raise AssertionError("unreachable") from None |
| 150 | |
| 151 | |
| 152 | def create_standard_kv_cache_spec( |
| 153 | vllm_config: VllmConfig, |
| 154 | attn_type: AttentionType = AttentionType.DECODER, |
| 155 | ) -> FullAttentionSpec | EncoderOnlyAttentionSpec: |
| 156 | """Create an AttentionSpec from VllmConfig. |
| 157 | |
| 158 | Returns an EncoderOnlyAttentionSpec for encoder-only attention (no KV |
| 159 | cache), and a FullAttentionSpec otherwise. |
| 160 | """ |
| 161 | if attn_type == AttentionType.ENCODER_ONLY: |
| 162 | return EncoderOnlyAttentionSpec( |
| 163 | block_size=vllm_config.cache_config.block_size, |
| 164 | num_kv_heads=vllm_config.model_config.get_num_kv_heads( |
| 165 | vllm_config.parallel_config |
| 166 | ), |
| 167 | head_size=vllm_config.model_config.get_head_size(), |
| 168 | dtype=vllm_config.model_config.dtype, |
| 169 | ) |
| 170 | return FullAttentionSpec( |
| 171 | block_size=vllm_config.cache_config.block_size, |
| 172 | num_kv_heads=vllm_config.model_config.get_num_kv_heads( |
| 173 | vllm_config.parallel_config |
| 174 | ), |
| 175 | head_size=vllm_config.model_config.get_head_size(), |
| 176 | dtype=vllm_config.model_config.dtype, |
| 177 | sliding_window=vllm_config.model_config.get_sliding_window(), |
| 178 | kv_quant_mode=get_kv_quant_mode(vllm_config.cache_config.cache_dtype), |
| 179 | ) |
| 180 | |
| 181 | |
| 182 | def create_vllm_config( |
| 183 | model_name: str = "meta-llama/Meta-Llama-3-8B", |
| 184 | tensor_parallel_size: int = 1, |
| 185 | max_model_len: int = 1024, |
| 186 | dtype: ModelDType | torch.dtype = "auto", |
| 187 | num_gpu_blocks: int = 1000, |
| 188 | block_size: int = 16, |
| 189 | max_num_seqs: int = 256, |
| 190 | max_num_batched_tokens: int = 8192, |
| 191 | enable_chunked_prefill: bool = True, |
| 192 | add_mock_model_methods: bool = True, |
| 193 | hf_config_override: dict | None = None, |
| 194 | ) -> VllmConfig: |
| 195 | """Create a VllmConfig for testing with reasonable defaults.""" |
| 196 | model_config = ModelConfig( |
| 197 | model=model_name, |
| 198 | tokenizer=model_name, |
| 199 | trust_remote_code=False, |
| 200 | dtype=dtype, |
| 201 | seed=0, |
| 202 | max_model_len=max_model_len, |
| 203 | ) |
| 204 | |
| 205 | cache_config = CacheConfig( |
| 206 | block_size=block_size, |
| 207 | cache_dtype="auto", |
| 208 | ) |
| 209 | # Set cache blocks for testing |
| 210 | # (these may be set during initialization normally) |
| 211 | cache_config.num_gpu_blocks = num_gpu_blocks |
| 212 | cache_config.num_cpu_blocks = 0 |
| 213 | # Builders read the resolved layout from the config; tests that exercise a |
| 214 | # different layout overwrite this field. |
| 215 | cache_config.kv_cache_layout = "LBNHC" |
| 216 | |
| 217 | parallel_config = ParallelConfig( |
| 218 | tensor_parallel_size=tensor_parallel_size, |
| 219 | ) |
| 220 | |
| 221 | scheduler_config = SchedulerConfig( |
| 222 | max_num_seqs=max_num_seqs, |
| 223 | max_num_batched_tokens=max_num_batched_tokens, |
| 224 | enable_chunked_prefill=enable_chunked_prefill, |
| 225 | max_model_len=model_config.max_model_len, |
| 226 | is_encoder_decoder=model_config.is_encoder_decoder, |
| 227 | ) |
| 228 | |
| 229 | device_config = DeviceConfig() |
| 230 | load_config = LoadConfig() |
| 231 | compilation_config = CompilationConfig() |
| 232 | |
| 233 | if add_mock_model_methods: |
| 234 | # Add mock methods to satisfy backends that need them |
| 235 | # This is a workaround because tests don't build full, real models, |
| 236 | # but some backends expect to query the model for layer-specific |
| 237 | # parameters |
| 238 | import types |
| 239 | |
| 240 | model_config.get_num_layers = types.MethodType(lambda self: 1, model_config) |
| 241 | model_config.get_sliding_window_for_layer = types.MethodType( |
| 242 | lambda self, i: None, model_config |
| 243 | ) |
| 244 | model_config.get_logits_soft_cap_for_layer = types.MethodType( |
| 245 | lambda self, i: 0.0, model_config |
| 246 | ) |
| 247 | model_config.get_sm_scale_for_layer = types.MethodType( |
| 248 | lambda self, i: 1.0 / model_config.get_head_size() ** 0.5, model_config |
| 249 | ) |
| 250 | |
| 251 | if hf_config_override: |
| 252 | model_config.hf_config.update(hf_config_override) |
| 253 | |
| 254 | return VllmConfig( |
| 255 | model_config=model_config, |
| 256 | cache_config=cache_config, |
| 257 | parallel_config=parallel_config, |
| 258 | scheduler_config=scheduler_config, |
| 259 | device_config=device_config, |
| 260 | load_config=load_config, |
| 261 | compilation_config=compilation_config, |
| 262 | ) |
| 263 | |
| 264 | |
| 265 | @dataclass |
| 266 | class BackendConfig: |
| 267 | name: str |
| 268 | attention_config: dict |
| 269 | comp_config: dict |
| 270 | specific_gpu_arch: tuple | None = None |
| 271 | |
| 272 | |
| 273 | # Define all backend configurations of full cudagraph to be tested |
| 274 | full_cg_backend_configs = { |
| 275 | # FA3 on Hopper |
| 276 | "FA3": BackendConfig( |
| 277 | name="FA3", |
| 278 | attention_config={ |
| 279 | "backend": "FLASH_ATTN", |
| 280 | "flash_attn_version": 3, |
| 281 | "flash_attn_max_num_splits_for_cuda_graph": 16, |
| 282 | }, |
| 283 | comp_config={ |
| 284 | "cudagraph_mode": "FULL", |
| 285 | }, |
| 286 | specific_gpu_arch=(9, 0), |
| 287 | ), |
| 288 | # FlashMLA on Hopper |
| 289 | "FlashMLA": BackendConfig( |
| 290 | name="FlashMLA", |
| 291 | attention_config={"backend": "FLASHMLA"}, |
| 292 | comp_config={ |
| 293 | "cudagraph_mode": "FULL_AND_PIECEWISE", |
| 294 | }, |
| 295 | specific_gpu_arch=(9, 0), |
| 296 | ), |
| 297 | # Cutlass MLA on Blackwell |
| 298 | "CutlassMLA": BackendConfig( |
| 299 | name="CutlassMLA", |
| 300 | attention_config={"backend": "CUTLASS_MLA"}, |
| 301 | comp_config={ |
| 302 | "cudagraph_mode": "FULL_AND_PIECEWISE", |
| 303 | }, |
| 304 | specific_gpu_arch=(10, 0), |
| 305 | ), |
| 306 | # FlashInfer MLA on Blackwell |
| 307 | "FlashInferMLA": BackendConfig( |
| 308 | name="FlashInferMLA", |
| 309 | attention_config={"backend": "FLASHINFER_MLA"}, |
| 310 | comp_config={ |
| 311 | "cudagraph_mode": "FULL_AND_PIECEWISE", |
| 312 | }, |
| 313 | specific_gpu_arch=(10, 0), |
| 314 | ), |
| 315 | # FlashAttention MLA on Hopper |
| 316 | "FlashAttentionMLA": BackendConfig( |
| 317 | name="FlashAttentionMLA", |
| 318 | attention_config={ |
| 319 | "backend": "FLASH_ATTN_MLA", |
| 320 | "flash_attn_max_num_splits_for_cuda_graph": 16, |
| 321 | }, |
| 322 | comp_config={ |
| 323 | "cudagraph_mode": "FULL_DECODE_ONLY", |
| 324 | }, |
| 325 | specific_gpu_arch=(9, 0), |
| 326 | ), |
| 327 | # FA2 |
| 328 | "FA2": BackendConfig( |
| 329 | name="FA2", |
| 330 | attention_config={ |
| 331 | "backend": "FLASH_ATTN", |
| 332 | "flash_attn_version": 2, |
| 333 | "flash_attn_max_num_splits_for_cuda_graph": 16, |
| 334 | }, |
| 335 | comp_config={ |
| 336 | "cudagraph_mode": "FULL_AND_PIECEWISE", |
| 337 | }, |
| 338 | ), |
| 339 | # Triton Attention |
| 340 | "TritonAttn": BackendConfig( |
| 341 | name="TritonAttn", |
| 342 | attention_config={"backend": "TRITON_ATTN"}, |
| 343 | comp_config={ |
| 344 | "cudagraph_mode": "FULL_AND_PIECEWISE", |
| 345 | }, |
| 346 | ), |
| 347 | # FlashInfer |
| 348 | "FlashInfer": BackendConfig( |
| 349 | name="FlashInfer", |
| 350 | attention_config={"backend": "FLASHINFER"}, |
| 351 | comp_config={ |
| 352 | "cudagraph_mode": "FULL_AND_PIECEWISE", |
| 353 | }, |
| 354 | ), |
| 355 | "RocmAttn": BackendConfig( |
| 356 | name="RocmAttn", |
| 357 | attention_config={ |
| 358 | "backend": "ROCM_ATTN", |
| 359 | }, |
| 360 | comp_config={ |
| 361 | "cudagraph_mode": "FULL", |
| 362 | }, |
| 363 | ), |
| 364 | } |
| 365 | |
| 366 | |
| 367 | class MockMambaBuilder(BaseMambaAttentionMetadataBuilder[BaseMambaAttentionMetadata]): |
| 368 | """Minimal concrete subclass for testing (base class is ABC).""" |
| 369 | |
| 370 | metadata_cls = BaseMambaAttentionMetadata |
| 371 | |
| 372 | @classmethod |
| 373 | def build_mamba_metadata( |
| 374 | cls, |
| 375 | vllm_config: VllmConfig, |
| 376 | seq_lens: list[int], |
| 377 | query_lens: list[int], |
| 378 | is_prefilling: list[bool], |
| 379 | *, |
| 380 | device: torch.device | None = None, |
| 381 | ) -> BaseMambaAttentionMetadata: |
| 382 | block_size = vllm_config.cache_config.block_size |
| 383 | device = device or torch.device("cpu") |
| 384 | mamba_spec = MambaSpec( |
| 385 | block_size=block_size, shapes=((1,), (1,)), dtypes=(torch.float32,) |
| 386 | ) |
| 387 | builder = cls(mamba_spec, ["layer0"], vllm_config, device) |
| 388 | batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=query_lens) |
| 389 | common_metadata = create_common_attn_metadata( |
| 390 | batch_spec, block_size=block_size, device=device, arange_block_indices=True |
| 391 | ) |
| 392 | common_metadata = common_metadata.replace( |
| 393 | is_prefilling=torch.tensor(is_prefilling, dtype=torch.bool) |
| 394 | ) |
| 395 | return builder.build(0, common_metadata) |
| 396 | |
| 397 | |
| 398 | def dense_kv_cache_tensor( |
| 399 | raw: torch.Tensor, |
| 400 | spec: KVCacheSpec, |
| 401 | num_blocks: int, |
| 402 | num_layers: int, |
| 403 | layout: KVCacheLayout, |
| 404 | layer_names: list[str] | None = None, |
| 405 | ) -> KVCacheTensor: |
| 406 | """The ``KVCacheTensor`` for a dense allocation of ``num_layers`` layers.""" |
| 407 | layer_stride, block_stride, _, _, _ = compute_layout_strides( |
| 408 | spec, num_blocks, num_layers, layout |
| 409 | ) |
| 410 | return KVCacheTensor( |
| 411 | size=raw.numel(), |
| 412 | layers=layer_names or [str(i) for i in range(num_layers)], |
| 413 | layer_stride=layer_stride, |
| 414 | block_stride=block_stride, |
| 415 | ) |
| 416 | |
| 417 | |
| 418 | def dense_kv_cache_views( |
| 419 | raw: torch.Tensor, |
| 420 | spec: KVCacheSpec, |
| 421 | num_blocks: int, |
| 422 | num_layers: int, |
| 423 | layout: KVCacheLayout, |
| 424 | kernel_block_size: int | None = None, |
| 425 | ) -> list[torch.Tensor]: |
| 426 | """``create_kv_cache_views`` for a dense allocation of ``num_layers`` layers.""" |
| 427 | tensor = dense_kv_cache_tensor(raw, spec, num_blocks, num_layers, layout) |
| 428 | return create_kv_cache_views( |
| 429 | raw, |
| 430 | spec, |
| 431 | num_blocks, |
| 432 | layout, |
| 433 | tensor, |
| 434 | kernel_block_size=kernel_block_size, |
| 435 | ) |