| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | import vllm.config |
| 8 | from tests.compile.backend import TestBackend |
| 9 | from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata |
| 10 | from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops |
| 11 | from vllm.compilation.passes.fusion.mla_rope_kvcache_cat_fusion import ( |
| 12 | MLARoPEKVCacheCatFusionPass, |
| 13 | ) |
| 14 | from vllm.compilation.passes.utility.fix_functionalization import ( |
| 15 | FixFunctionalizationPass, |
| 16 | ) |
| 17 | from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass |
| 18 | from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass |
| 19 | from vllm.config import ( |
| 20 | CacheConfig, |
| 21 | CompilationConfig, |
| 22 | CompilationMode, |
| 23 | ModelConfig, |
| 24 | PassConfig, |
| 25 | VllmConfig, |
| 26 | ) |
| 27 | from vllm.config.cache import CacheDType |
| 28 | from vllm.forward_context import get_forward_context, set_forward_context |
| 29 | from vllm.model_executor.layers.attention import MLAAttention |
| 30 | from vllm.model_executor.layers.linear import ColumnParallelLinear |
| 31 | from vllm.model_executor.layers.rotary_embedding import ( |
| 32 | DeepseekScalingRotaryEmbedding, |
| 33 | RotaryEmbedding, |
| 34 | ) |
| 35 | from vllm.platforms import current_platform |
| 36 | from vllm.utils.torch_utils import _encode_layer_name |
| 37 | from vllm.v1.attention.backend import ( |
| 38 | AttentionBackend, |
| 39 | CommonAttentionMetadata, |
| 40 | ) |
| 41 | from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla |
| 42 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 43 | |
| 44 | INDEX_SELECT_OP = torch.ops.aten.index.Tensor |
| 45 | VLLM_UNIFIED_MLA_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_mla_kv_cache_update |
| 46 | FP8_DTYPE = current_platform.fp8_dtype() |
| 47 | |
| 48 | |
| 49 | class MLARoPEKVCacheCatTestModel(torch.nn.Module): |
| 50 | def __init__( |
| 51 | self, |
| 52 | vllm_config: VllmConfig, |
| 53 | attn_backend: AttentionBackendEnum, |
| 54 | use_deepseek_scaling_rope: bool, |
| 55 | num_heads: int, |
| 56 | qk_nope_head_dim: int, |
| 57 | qk_rope_head_dim: int, |
| 58 | v_head_dim: int, |
| 59 | q_lora_rank: int, |
| 60 | kv_lora_rank: int, |
| 61 | is_neox: bool, |
| 62 | dtype: torch.dtype, |
| 63 | device: torch.device, |
| 64 | prefix: str = "model.layers.0.self_attn.attn", |
| 65 | ): |
| 66 | super().__init__() |
| 67 | self.num_heads = num_heads |
| 68 | self.qk_nope_head_dim = qk_nope_head_dim |
| 69 | self.qk_rope_head_dim = qk_rope_head_dim |
| 70 | self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim |
| 71 | self.v_head_dim = v_head_dim |
| 72 | self.q_lora_rank = q_lora_rank |
| 73 | self.kv_lora_rank = kv_lora_rank |
| 74 | self.dtype = dtype |
| 75 | self.device = device |
| 76 | self.layer_name = prefix |
| 77 | |
| 78 | self.num_kv_heads = 1 |
| 79 | self.head_size = kv_lora_rank + qk_rope_head_dim |
| 80 | self.block_size = vllm_config.cache_config.block_size |
| 81 | self.scale = self.qk_head_dim**-0.5 |
| 82 | |
| 83 | if use_deepseek_scaling_rope: |
| 84 | self.rotary_emb = DeepseekScalingRotaryEmbedding( |
| 85 | head_size=qk_rope_head_dim, |
| 86 | rotary_dim=qk_rope_head_dim, |
| 87 | max_position_embeddings=4096, |
| 88 | base=10000, |
| 89 | is_neox_style=is_neox, |
| 90 | scaling_factor=1.0, |
| 91 | dtype=dtype, |
| 92 | ) |
| 93 | else: |
| 94 | self.rotary_emb = RotaryEmbedding( |
| 95 | head_size=qk_rope_head_dim, |
| 96 | rotary_dim=qk_rope_head_dim, |
| 97 | max_position_embeddings=4096, |
| 98 | base=10000, |
| 99 | is_neox_style=is_neox, |
| 100 | dtype=dtype, |
| 101 | ) |
| 102 | |
| 103 | # Initialize intermediate mm layers for unit test |
| 104 | self.q_b_proj = ColumnParallelLinear( |
| 105 | self.q_lora_rank, |
| 106 | self.num_heads * self.qk_head_dim, |
| 107 | bias=False, |
| 108 | prefix=f"{prefix}.q_b_proj", |
| 109 | ).to(device) |
| 110 | self.kv_b_proj = ColumnParallelLinear( |
| 111 | self.kv_lora_rank, |
| 112 | self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), |
| 113 | bias=False, |
| 114 | prefix=f"{prefix}.kv_b_proj", |
| 115 | ).to(device) |
| 116 | |
| 117 | # ColumnParallelLinear default init in bf16 with seed 0 produces |
| 118 | # near-zero weights (7/4.7M nonzero), making the GEMM output almost |
| 119 | # entirely zero and masking correctness bugs. Reinitialize to get |
| 120 | # dense outputs. |
| 121 | with torch.no_grad(): |
| 122 | torch.nn.init.normal_(self.q_b_proj.weight, std=0.02) |
| 123 | torch.nn.init.normal_(self.kv_b_proj.weight, std=0.02) |
| 124 | |
| 125 | # Register layer metadata for the fusion pass via MLAAttention |
| 126 | self.mla_attn = MLAAttention( |
| 127 | num_heads=self.num_heads, |
| 128 | scale=self.scale, |
| 129 | qk_nope_head_dim=self.qk_nope_head_dim, |
| 130 | qk_rope_head_dim=self.qk_rope_head_dim, |
| 131 | v_head_dim=self.v_head_dim, |
| 132 | q_lora_rank=self.q_lora_rank, |
| 133 | kv_lora_rank=self.kv_lora_rank, |
| 134 | kv_b_proj=self.kv_b_proj, |
| 135 | cache_config=vllm_config.cache_config, |
| 136 | quant_config=vllm_config.quant_config, |
| 137 | prefix=prefix, |
| 138 | attn_backend=attn_backend.get_class(), |
| 139 | ) |
| 140 | self.attn_backend: type[AttentionBackend] = self.mla_attn.get_attn_backend() |
| 141 | self.mla_attn._k_scale = self.mla_attn._k_scale.to(device) |
| 142 | self.mla_attn._v_scale = self.mla_attn._v_scale.to(device) |
| 143 | |
| 144 | # Keep both the string dtype (for ops) and torch dtype (for tensors) |
| 145 | self.kv_cache_dtype_str = vllm_config.cache_config.cache_dtype |
| 146 | self.kv_cache_dtype = ( |
| 147 | FP8_DTYPE if self.kv_cache_dtype_str.startswith("fp8") else self.dtype |
| 148 | ) |
| 149 | |
| 150 | # Initialize attn MetadataBuilder |
| 151 | self.builder = self.attn_backend.get_builder_cls()( |
| 152 | kv_cache_spec=self.mla_attn.get_kv_cache_spec(vllm_config), |
| 153 | layer_names=[self.mla_attn.layer_name], |
| 154 | vllm_config=vllm_config, |
| 155 | device=device, |
| 156 | ) |
| 157 | |
| 158 | def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata: |
| 159 | """Initialize attention metadata.""" |
| 160 | # Create common attn metadata |
| 161 | batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size) |
| 162 | common_attn_metadata = create_common_attn_metadata( |
| 163 | batch_spec, self.block_size, self.device, arange_block_indices=True |
| 164 | ) |
| 165 | |
| 166 | max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size |
| 167 | num_blocks = batch_size * max_blocks |
| 168 | |
| 169 | # MLA KV cache is 4D: (num_blocks, num_heads=1, block_size, head_size) |
| 170 | kv_cache = torch.zeros( |
| 171 | (num_blocks, 1, self.block_size, self.head_size), |
| 172 | dtype=self.kv_cache_dtype, |
| 173 | device=self.device, |
| 174 | ) |
| 175 | |
| 176 | self.mla_attn.bind_kv_cache(kv_cache) |
| 177 | |
| 178 | # Build attn metadata |
| 179 | attn_metadata = self.builder.build( |
| 180 | common_prefix_len=0, common_attn_metadata=common_attn_metadata |
| 181 | ) |
| 182 | |
| 183 | return attn_metadata |
| 184 | |
| 185 | def forward( |
| 186 | self, qkv_lora: torch.Tensor, positions: torch.Tensor |
| 187 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 188 | qkv_lora = qkv_lora.clone() |
| 189 | q_c, kv_lora = qkv_lora.split( |
| 190 | [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], |
| 191 | dim=-1, |
| 192 | ) |
| 193 | q = self.q_b_proj(q_c)[0] |
| 194 | kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) |
| 195 | |
| 196 | q = q.view(-1, self.num_heads, self.qk_head_dim) |
| 197 | k_pe = k_pe.unsqueeze(1) |
| 198 | |
| 199 | q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb( |
| 200 | positions, q[..., self.qk_nope_head_dim :], k_pe |
| 201 | ) |
| 202 | |
| 203 | dummy = torch.ops.vllm.unified_mla_kv_cache_update( |
| 204 | kv_c, |
| 205 | k_pe, |
| 206 | _encode_layer_name(self.layer_name), |
| 207 | self.kv_cache_dtype_str, |
| 208 | self.mla_attn._k_scale, |
| 209 | ) |
| 210 | return q, kv_c, k_pe, dummy |
| 211 | |
| 212 | def ops_in_model_before(self) -> list[torch._ops.OpOverload]: |
| 213 | ops = [ |
| 214 | INDEX_SELECT_OP, |
| 215 | torch.ops.vllm.unified_mla_kv_cache_update.default, |
| 216 | ] |
| 217 | return ops |
| 218 | |
| 219 | def ops_in_model_after(self) -> list[torch._ops.OpOverload]: |
| 220 | return [torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default] |
| 221 | |
| 222 | |
| 223 | MLA_BACKENDS = [AttentionBackendEnum.TRITON_MLA] |
| 224 | if flash_attn_supports_mla(): |
| 225 | MLA_BACKENDS += [AttentionBackendEnum.FLASH_ATTN_MLA] |
| 226 | if is_aiter_found_and_supported(): |
| 227 | MLA_BACKENDS += [AttentionBackendEnum.ROCM_AITER_MLA] |
| 228 | |
| 229 | |
| 230 | @pytest.mark.parametrize("attn_backend", MLA_BACKENDS) |
| 231 | @pytest.mark.parametrize("use_deepseek_scaling_rope", [True]) |
| 232 | @pytest.mark.parametrize("num_heads", [16]) |
| 233 | @pytest.mark.parametrize("qk_nope_head_dim", [128]) |
| 234 | @pytest.mark.parametrize("qk_rope_head_dim", [64]) |
| 235 | @pytest.mark.parametrize("v_head_dim", [128]) |
| 236 | @pytest.mark.parametrize("q_lora_rank", [1536]) |
| 237 | @pytest.mark.parametrize("kv_lora_rank", [512]) |
| 238 | @pytest.mark.parametrize("block_size", [16]) |
| 239 | @pytest.mark.parametrize("is_neox", [True, False]) |
| 240 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 241 | @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) |
| 242 | @pytest.mark.skipif( |
| 243 | not current_platform.is_cuda_alike(), |
| 244 | reason="MLA RoPE+KVCache+Cat fusion is only supported on CUDA and ROCm.", |
| 245 | ) |
| 246 | def test_mla_rope_kvcache_cat_fusion( |
| 247 | attn_backend: AttentionBackendEnum, |
| 248 | use_deepseek_scaling_rope: bool, |
| 249 | num_heads: int, |
| 250 | qk_nope_head_dim: int, |
| 251 | qk_rope_head_dim: int, |
| 252 | v_head_dim: int, |
| 253 | q_lora_rank: int, |
| 254 | kv_lora_rank: int, |
| 255 | block_size: int, |
| 256 | is_neox: bool, |
| 257 | dtype: torch.dtype, |
| 258 | kv_cache_dtype: CacheDType, |
| 259 | monkeypatch: pytest.MonkeyPatch, |
| 260 | ): |
| 261 | torch.set_default_device("cuda") |
| 262 | torch.set_default_dtype(dtype) |
| 263 | torch.manual_seed(0) |
| 264 | |
| 265 | vllm_config = VllmConfig( |
| 266 | model_config=ModelConfig( |
| 267 | model="deepseek-ai/DeepSeek-V2-Lite", |
| 268 | dtype=dtype, |
| 269 | ), |
| 270 | cache_config=CacheConfig( |
| 271 | block_size=block_size, |
| 272 | cache_dtype=kv_cache_dtype, |
| 273 | ), |
| 274 | compilation_config=CompilationConfig( |
| 275 | mode=CompilationMode.VLLM_COMPILE, |
| 276 | pass_config=PassConfig( |
| 277 | fuse_rope_kvcache_cat_mla=True, |
| 278 | eliminate_noops=True, |
| 279 | ), |
| 280 | ), |
| 281 | ) |
| 282 | |
| 283 | with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m: |
| 284 | if not torch.distributed.is_initialized(): |
| 285 | from vllm.distributed.parallel_state import ( |
| 286 | init_distributed_environment, |
| 287 | initialize_model_parallel, |
| 288 | ) |
| 289 | from vllm.utils.system_utils import update_environment_variables |
| 290 | |
| 291 | update_environment_variables( |
| 292 | { |
| 293 | "RANK": "0", |
| 294 | "LOCAL_RANK": "0", |
| 295 | "WORLD_SIZE": "1", |
| 296 | "MASTER_ADDR": "localhost", |
| 297 | "MASTER_PORT": "54321", |
| 298 | } |
| 299 | ) |
| 300 | init_distributed_environment() |
| 301 | initialize_model_parallel() |
| 302 | |
| 303 | if attn_backend == AttentionBackendEnum.ROCM_AITER_MLA: |
| 304 | m.setenv("VLLM_ROCM_USE_AITER", "1") |
| 305 | rocm_aiter_ops.refresh_env_variables() |
| 306 | |
| 307 | model = MLARoPEKVCacheCatTestModel( |
| 308 | vllm_config=vllm_config, |
| 309 | attn_backend=attn_backend, |
| 310 | use_deepseek_scaling_rope=use_deepseek_scaling_rope, |
| 311 | num_heads=num_heads, |
| 312 | qk_nope_head_dim=qk_nope_head_dim, |
| 313 | qk_rope_head_dim=qk_rope_head_dim, |
| 314 | v_head_dim=v_head_dim, |
| 315 | q_lora_rank=q_lora_rank, |
| 316 | kv_lora_rank=kv_lora_rank, |
| 317 | is_neox=is_neox, |
| 318 | dtype=dtype, |
| 319 | device=torch.get_default_device(), |
| 320 | ) |
| 321 | |
| 322 | fusion_pass = MLARoPEKVCacheCatFusionPass(vllm_config) |
| 323 | # note: FixFunctionalizationPass is required to correctly lower |
| 324 | # the fused op to its inplace version with auto-functionalization v1. |
| 325 | # Without it, decompose_auto_functionalized calls clone_preserve_strides |
| 326 | # on the non-contiguous q_pe slice directly, and inductor's lowering |
| 327 | # of the resulting as_strided chain incorrectly drops the storage offset. |
| 328 | # auto-functionalization v2 avoids this: it clones the contiguous base |
| 329 | # tensor (_all_bases) and reconstructs the slice as a view, so the |
| 330 | # offset is never passed through as_strided lowering. |
| 331 | passes = [ |
| 332 | NoOpEliminationPass(vllm_config), |
| 333 | fusion_pass, |
| 334 | PostCleanupPass(vllm_config), |
| 335 | FixFunctionalizationPass(vllm_config), |
| 336 | ] |
| 337 | backend = TestBackend(*passes) |
| 338 | |
| 339 | T = 5 |
| 340 | |
| 341 | qkv_lora = torch.randn( |
| 342 | T, |
| 343 | q_lora_rank + kv_lora_rank + qk_rope_head_dim, |
| 344 | dtype=dtype, |
| 345 | ) |
| 346 | pos = torch.arange(T, dtype=torch.long) |
| 347 | |
| 348 | qkv_unfused = qkv_lora.clone() |
| 349 | pos_unfused = pos.clone() |
| 350 | |
| 351 | # Run unfused version |
| 352 | with set_forward_context(None, vllm_config): |
| 353 | forward_context = get_forward_context() |
| 354 | attn_metadata = model.build_attn_metadata(T) |
| 355 | forward_context.slot_mapping = { |
| 356 | model.layer_name: attn_metadata.slot_mapping |
| 357 | } |
| 358 | q_unfused, kv_c_unfused, k_pe_unfused, dummy = model( |
| 359 | qkv_unfused, pos_unfused |
| 360 | ) |
| 361 | attn_layer = forward_context.no_compile_layers[model.layer_name] |
| 362 | kv_cache_unfused = attn_layer.kv_cache.clone() |
| 363 | del dummy |
| 364 | |
| 365 | # Run fused version (compiled) |
| 366 | torch._dynamo.mark_dynamic(qkv_lora, 0) |
| 367 | torch._dynamo.mark_dynamic(pos, 0) |
| 368 | with set_forward_context(None, vllm_config): |
| 369 | model_fused = torch.compile(model, backend=backend) |
| 370 | forward_context = get_forward_context() |
| 371 | attn_metadata = model.build_attn_metadata(T) |
| 372 | forward_context.slot_mapping = { |
| 373 | model.layer_name: attn_metadata.slot_mapping |
| 374 | } |
| 375 | q_fused, kv_c_fused, k_pe_fused, dummy = model_fused(qkv_lora, pos) |
| 376 | attn_layer = forward_context.no_compile_layers[model.layer_name] |
| 377 | kv_cache_fused = attn_layer.kv_cache |
| 378 | del dummy |
| 379 | |
| 380 | assert fusion_pass.matched_count == 1 |
| 381 | |
| 382 | backend.check_before_ops(model.ops_in_model_before()) |
| 383 | backend.check_after_ops(model.ops_in_model_after()) |
| 384 | |
| 385 | if dtype == torch.float16: |
| 386 | ATOL, RTOL = (2e-3, 2e-3) |
| 387 | else: |
| 388 | ATOL, RTOL = (1e-2, 1e-2) |
| 389 | |
| 390 | torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL) |
| 391 | torch.testing.assert_close(kv_c_unfused, kv_c_fused, atol=ATOL, rtol=RTOL) |
| 392 | torch.testing.assert_close(k_pe_unfused, k_pe_fused, atol=ATOL, rtol=RTOL) |
| 393 | # Cannot compare fp8_* directly here, cast to model dtype instead |
| 394 | torch.testing.assert_close( |
| 395 | kv_cache_unfused.view(dtype), |
| 396 | kv_cache_fused.view(dtype), |
| 397 | atol=ATOL, |
| 398 | rtol=RTOL, |
| 399 | ) |