| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | from torch._higher_order_ops import auto_functionalized |
| 7 | |
| 8 | import vllm.config |
| 9 | from tests.compile.backend import TestBackend |
| 10 | from tests.v1.attention.utils import ( |
| 11 | BatchSpec, |
| 12 | create_common_attn_metadata, |
| 13 | dense_kv_cache_views, |
| 14 | ) |
| 15 | from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops |
| 16 | from vllm.compilation.passes.fusion import rope_kvcache_fusion |
| 17 | from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP |
| 18 | from vllm.compilation.passes.fusion.rope_kvcache_fusion import RopeKVCacheFusionPass |
| 19 | from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass |
| 20 | from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass |
| 21 | from vllm.compilation.passes.utility.scatter_split_replace import ( |
| 22 | ScatterSplitReplacementPass, |
| 23 | ) |
| 24 | from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass |
| 25 | from vllm.config import ( |
| 26 | CacheConfig, |
| 27 | CompilationConfig, |
| 28 | CompilationMode, |
| 29 | ModelConfig, |
| 30 | PassConfig, |
| 31 | VllmConfig, |
| 32 | ) |
| 33 | from vllm.config.cache import CacheDType |
| 34 | from vllm.config.utils import Range |
| 35 | from vllm.forward_context import get_forward_context, set_forward_context |
| 36 | from vllm.model_executor.layers.attention import Attention |
| 37 | from vllm.model_executor.layers.attention import attention as attention_module |
| 38 | from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding |
| 39 | from vllm.platforms import current_platform |
| 40 | from vllm.utils.torch_utils import _encode_layer_name |
| 41 | from vllm.v1.attention.backend import ( |
| 42 | AttentionBackend, |
| 43 | CommonAttentionMetadata, |
| 44 | ) |
| 45 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 46 | from vllm.v1.kv_cache_interface import ( |
| 47 | FullAttentionSpec, |
| 48 | KVCacheLayout, |
| 49 | ) |
| 50 | |
| 51 | INDEX_SELECT_OP = torch.ops.aten.index.Tensor |
| 52 | VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update |
| 53 | FP8_DTYPE = current_platform.fp8_dtype() |
| 54 | |
| 55 | |
| 56 | def test_kv_cache_update_ops_fake_tensor_metadata(monkeypatch: pytest.MonkeyPatch): |
| 57 | query = torch.empty(1, dtype=torch.bfloat16) |
| 58 | kv_cache = torch.empty(1, dtype=torch.uint8) |
| 59 | context = (None, None, kv_cache, None) |
| 60 | monkeypatch.setattr(attention_module, "get_attention_context", lambda _: context) |
| 61 | monkeypatch.setattr(rope_kvcache_fusion, "get_attention_context", lambda _: context) |
| 62 | |
| 63 | runtime_output = attention_module.unified_kv_cache_update(query, query, "layer") |
| 64 | fake_output = attention_module.unified_kv_cache_update_fake(query, query, "layer") |
| 65 | assert runtime_output.shape == fake_output.shape |
| 66 | assert runtime_output.dtype == fake_output.dtype |
| 67 | assert runtime_output.device == fake_output.device |
| 68 | |
| 69 | args = ( |
| 70 | query, |
| 71 | query, |
| 72 | query, |
| 73 | torch.empty(1, dtype=torch.int64), |
| 74 | torch.empty(1), |
| 75 | True, |
| 76 | "layer", |
| 77 | ) |
| 78 | runtime_output = rope_kvcache_fusion.fused_rope_and_unified_kv_cache_update_impl( |
| 79 | *args |
| 80 | ) |
| 81 | fake_output = rope_kvcache_fusion.fused_rope_and_unified_kv_cache_update_fake(*args) |
| 82 | assert runtime_output.shape == fake_output.shape |
| 83 | assert runtime_output.dtype == fake_output.dtype |
| 84 | assert runtime_output.device == fake_output.device |
| 85 | |
| 86 | |
| 87 | def test_rope_kvcache_fusion_default_keeps_large_ranges_unfused(): |
| 88 | vllm_config = VllmConfig( |
| 89 | compilation_config=CompilationConfig( |
| 90 | mode=CompilationMode.VLLM_COMPILE, |
| 91 | pass_config=PassConfig(fuse_rope_kvcache=True), |
| 92 | ), |
| 93 | ) |
| 94 | fusion_pass = RopeKVCacheFusionPass(vllm_config) |
| 95 | |
| 96 | assert fusion_pass.is_applicable_for_range(Range(1, 256)) |
| 97 | assert not fusion_pass.is_applicable_for_range(Range(257, 11650)) |
| 98 | assert not fusion_pass.is_applicable_for_range(Range(11651, 16384)) |
| 99 | |
| 100 | |
| 101 | class QKRoPEKVCacheTestModel(torch.nn.Module): |
| 102 | def __init__( |
| 103 | self, |
| 104 | vllm_config: VllmConfig, |
| 105 | attn_backend: AttentionBackendEnum, |
| 106 | num_heads: int, |
| 107 | num_kv_heads: int, |
| 108 | head_size: int, |
| 109 | is_neox: bool, |
| 110 | dtype: torch.dtype, |
| 111 | device: torch.device, |
| 112 | prefix: str = "model.layers.0.self_attn.attn", |
| 113 | ): |
| 114 | super().__init__() |
| 115 | self.num_heads = num_heads |
| 116 | self.num_kv_heads = num_kv_heads |
| 117 | self.head_size = head_size |
| 118 | self.block_size = vllm_config.cache_config.block_size |
| 119 | self.q_size = num_heads * head_size |
| 120 | self.kv_size = num_kv_heads * head_size |
| 121 | self.is_neox = is_neox |
| 122 | self.dtype = dtype |
| 123 | self.device = device |
| 124 | self.layer_name = prefix |
| 125 | |
| 126 | self.rotary_emb = RotaryEmbedding( |
| 127 | head_size, |
| 128 | rotary_dim=head_size, |
| 129 | max_position_embeddings=4096, |
| 130 | base=10000, |
| 131 | is_neox_style=is_neox, |
| 132 | dtype=self.dtype, |
| 133 | ) |
| 134 | |
| 135 | # Whether to check for the RoPE custom op or component index_select |
| 136 | self.enable_rope_custom_op = self.rotary_emb.enabled() |
| 137 | |
| 138 | # Register layer metadata for the fusion pass via Attention. |
| 139 | self.attn = Attention( |
| 140 | num_heads=num_heads, |
| 141 | head_size=head_size, |
| 142 | scale=1.0 / head_size**0.5, |
| 143 | num_kv_heads=num_kv_heads, |
| 144 | cache_config=vllm_config.cache_config, |
| 145 | quant_config=vllm_config.quant_config, |
| 146 | prefix=prefix, |
| 147 | attn_backend=attn_backend.get_class(), |
| 148 | ) |
| 149 | self.attn_backend: type[AttentionBackend] = self.attn.get_attn_backend() |
| 150 | assert not self.attn_backend.forward_includes_kv_cache_update, ( |
| 151 | f"Attention backend {self.attn_backend} does not support fuse_rope_kvcache." |
| 152 | ) |
| 153 | self.attn._k_scale = self.attn._k_scale.to(device) |
| 154 | self.attn._v_scale = self.attn._v_scale.to(device) |
| 155 | |
| 156 | kv_cache_dtype_str = vllm_config.cache_config.cache_dtype |
| 157 | self.kv_cache_dtype = ( |
| 158 | FP8_DTYPE if kv_cache_dtype_str.startswith("fp8") else self.dtype |
| 159 | ) |
| 160 | |
| 161 | # Mirror the worker spec-collection loop: the backend publishes its |
| 162 | # packing (e.g. ROCM_ATTN's separate K/V head groups). |
| 163 | self.kv_cache_spec = self.attn_backend.customize_spec( |
| 164 | FullAttentionSpec( |
| 165 | block_size=self.block_size, |
| 166 | num_kv_heads=self.num_kv_heads, |
| 167 | head_size=self.head_size, |
| 168 | dtype=self.kv_cache_dtype, |
| 169 | ) |
| 170 | ) |
| 171 | |
| 172 | # Initialize attn MetadataBuilder |
| 173 | self.builder = self.attn_backend.get_builder_cls()( |
| 174 | kv_cache_spec=self.kv_cache_spec, |
| 175 | layer_names=[self.attn.layer_name], |
| 176 | vllm_config=vllm_config, |
| 177 | device=device, |
| 178 | ) |
| 179 | |
| 180 | def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata: |
| 181 | """Initialize attention metadata.""" |
| 182 | # Create common attn metadata |
| 183 | batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size) |
| 184 | common_attn_metadata = create_common_attn_metadata( |
| 185 | batch_spec, self.block_size, self.device, arange_block_indices=True |
| 186 | ) |
| 187 | |
| 188 | max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size |
| 189 | num_blocks = batch_size * max_blocks |
| 190 | |
| 191 | # A backend's preferred supported layout wins over the default, mirroring |
| 192 | # selector-time resolution (e.g. ROCM_ATTN's head groups force LHBNC). |
| 193 | layout = KVCacheLayout.LBNHC |
| 194 | supported = self.attn_backend.supported_kv_cache_layouts() |
| 195 | if supported: |
| 196 | layout = supported[0] |
| 197 | |
| 198 | raw_tensor = torch.zeros( |
| 199 | num_blocks * self.kv_cache_spec.page_size_bytes, |
| 200 | dtype=torch.int8, |
| 201 | device=self.device, |
| 202 | ) |
| 203 | self.attn.kv_cache = dense_kv_cache_views( |
| 204 | raw_tensor, |
| 205 | self.kv_cache_spec, |
| 206 | num_blocks, |
| 207 | num_layers=1, |
| 208 | layout=layout, |
| 209 | )[0] |
| 210 | |
| 211 | # Build attn metadata |
| 212 | attn_metadata = self.builder.build( |
| 213 | common_prefix_len=0, common_attn_metadata=common_attn_metadata |
| 214 | ) |
| 215 | |
| 216 | return attn_metadata |
| 217 | |
| 218 | def forward( |
| 219 | self, qkv: torch.Tensor, positions: torch.Tensor |
| 220 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 221 | # Create copy so inplace ops do not modify the original tensors |
| 222 | qkv = qkv.clone() |
| 223 | q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) |
| 224 | q, k = self.rotary_emb(positions, q, k) |
| 225 | |
| 226 | # Instead of a full forward pass, match only the KV cache update op here |
| 227 | q = q.view(-1, self.num_heads, self.head_size) |
| 228 | k = k.view(-1, self.num_kv_heads, self.head_size) |
| 229 | v = v.view(-1, self.num_kv_heads, self.head_size) |
| 230 | kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update( |
| 231 | k, v, _encode_layer_name(self.layer_name) |
| 232 | ) |
| 233 | return q, k, v, kv_cache_dummy_dep |
| 234 | |
| 235 | def ops_in_model_before(self) -> list[torch._ops.OpOverload]: |
| 236 | ops = [] |
| 237 | if self.enable_rope_custom_op: |
| 238 | if rocm_aiter_ops.is_triton_rotary_embed_enabled(): |
| 239 | ops.append(torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default) |
| 240 | else: |
| 241 | ops.append(ROTARY_OP) |
| 242 | else: |
| 243 | ops.append(INDEX_SELECT_OP) |
| 244 | ops.append(torch.ops.vllm.unified_kv_cache_update.default) |
| 245 | return ops |
| 246 | |
| 247 | def ops_in_model_after(self) -> list[torch._ops.OpOverload]: |
| 248 | return [torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default] |
| 249 | |
| 250 | |
| 251 | class QKRoPEStaticQKVCacheTestModel(QKRoPEKVCacheTestModel): |
| 252 | def __init__(self, *args, **kwargs): |
| 253 | super().__init__(*args, **kwargs) |
| 254 | self.q_scale = torch.ones((), dtype=torch.float32, device=self.device) |
| 255 | |
| 256 | def forward( |
| 257 | self, qkv: torch.Tensor, positions: torch.Tensor |
| 258 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 259 | # Create copy so inplace ops do not modify the original tensors |
| 260 | qkv = qkv.clone() |
| 261 | q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) |
| 262 | q, k = self.rotary_emb(positions, q, k) |
| 263 | |
| 264 | q_fp8 = torch.empty(q.shape, device=q.device, dtype=FP8_DTYPE) |
| 265 | _, q_fp8 = auto_functionalized( |
| 266 | torch.ops._C.static_scaled_fp8_quant.default, |
| 267 | result=q_fp8, |
| 268 | input=q, |
| 269 | scale=self.q_scale, |
| 270 | group_shape=(-1, -1), |
| 271 | ) |
| 272 | q = q_fp8.view(-1, self.num_heads, self.head_size) |
| 273 | k = k.view(-1, self.num_kv_heads, self.head_size) |
| 274 | v = v.view(-1, self.num_kv_heads, self.head_size) |
| 275 | kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update( |
| 276 | k, v, _encode_layer_name(self.layer_name) |
| 277 | ) |
| 278 | return q, k, v, kv_cache_dummy_dep |
| 279 | |
| 280 | def ops_in_model_before(self) -> list[torch._ops.OpOverload]: |
| 281 | ops = [] |
| 282 | if self.enable_rope_custom_op: |
| 283 | if rocm_aiter_ops.is_triton_rotary_embed_enabled(): |
| 284 | ops.append(torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default) |
| 285 | else: |
| 286 | ops.append(ROTARY_OP) |
| 287 | else: |
| 288 | ops.append(INDEX_SELECT_OP) |
| 289 | ops.append(torch.ops.vllm.unified_kv_cache_update.default) |
| 290 | return ops |
| 291 | |
| 292 | def ops_in_model_after(self) -> list[torch._ops.OpOverload]: |
| 293 | return [torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default] |
| 294 | |
| 295 | |
| 296 | @pytest.mark.parametrize( |
| 297 | "attn_backend", |
| 298 | [ |
| 299 | AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN, |
| 300 | AttentionBackendEnum.TRITON_ATTN, |
| 301 | AttentionBackendEnum.ROCM_ATTN, |
| 302 | AttentionBackendEnum.ROCM_AITER_FA, |
| 303 | ], |
| 304 | ) |
| 305 | @pytest.mark.parametrize("enable_rope_custom_op", [True]) # [True, False]) |
| 306 | @pytest.mark.parametrize("enable_aiter_triton_rope", [True, False]) |
| 307 | @pytest.mark.parametrize("num_heads", [64]) |
| 308 | @pytest.mark.parametrize("num_kv_heads", [8]) |
| 309 | @pytest.mark.parametrize("head_size", [64]) |
| 310 | @pytest.mark.parametrize("block_size", [16]) |
| 311 | @pytest.mark.parametrize("is_neox", [True, False]) |
| 312 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 313 | @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) |
| 314 | @pytest.mark.skipif( |
| 315 | not is_aiter_found_and_supported(), |
| 316 | reason="Only test on ROCm with AITER installed and supported", |
| 317 | ) |
| 318 | def test_rope_kvcache_fusion( |
| 319 | attn_backend: AttentionBackendEnum, |
| 320 | enable_rope_custom_op: bool, |
| 321 | enable_aiter_triton_rope: bool, |
| 322 | num_heads: int, |
| 323 | num_kv_heads: int, |
| 324 | head_size: int, |
| 325 | block_size: int, |
| 326 | is_neox: bool, |
| 327 | dtype: torch.dtype, |
| 328 | kv_cache_dtype: CacheDType, |
| 329 | monkeypatch: pytest.MonkeyPatch, |
| 330 | ): |
| 331 | torch.set_default_device("cuda") |
| 332 | torch.set_default_dtype(dtype) |
| 333 | torch.manual_seed(0) |
| 334 | |
| 335 | custom_ops: list[str] = [] |
| 336 | if enable_rope_custom_op: |
| 337 | custom_ops.append("+rotary_embedding") |
| 338 | |
| 339 | vllm_config = VllmConfig( |
| 340 | model_config=ModelConfig(dtype=dtype), |
| 341 | cache_config=CacheConfig( |
| 342 | block_size=block_size, |
| 343 | cache_dtype=kv_cache_dtype, |
| 344 | ), |
| 345 | compilation_config=CompilationConfig( |
| 346 | mode=CompilationMode.VLLM_COMPILE, |
| 347 | custom_ops=custom_ops, |
| 348 | pass_config=PassConfig( |
| 349 | fuse_rope_kvcache=True, |
| 350 | eliminate_noops=True, |
| 351 | ), |
| 352 | ), |
| 353 | ) |
| 354 | |
| 355 | with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m: |
| 356 | m.setenv("VLLM_ROCM_USE_AITER", "1") |
| 357 | m.setenv( |
| 358 | "VLLM_ROCM_USE_AITER_TRITON_ROPE", "1" if enable_aiter_triton_rope else "0" |
| 359 | ) |
| 360 | rocm_aiter_ops.refresh_env_variables() |
| 361 | |
| 362 | model = QKRoPEKVCacheTestModel( |
| 363 | vllm_config=vllm_config, |
| 364 | attn_backend=attn_backend, |
| 365 | num_heads=num_heads, |
| 366 | num_kv_heads=num_kv_heads, |
| 367 | head_size=head_size, |
| 368 | is_neox=is_neox, |
| 369 | dtype=dtype, |
| 370 | device=torch.get_default_device(), |
| 371 | ) |
| 372 | |
| 373 | fusion_pass = RopeKVCacheFusionPass(vllm_config) |
| 374 | passes = [ |
| 375 | NoOpEliminationPass(vllm_config), |
| 376 | SplitCoalescingPass(vllm_config), |
| 377 | ScatterSplitReplacementPass(vllm_config), |
| 378 | fusion_pass, |
| 379 | PostCleanupPass(vllm_config), |
| 380 | ] |
| 381 | backend = TestBackend(*passes) |
| 382 | |
| 383 | T = 5 |
| 384 | |
| 385 | qkv = torch.randn( |
| 386 | T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype |
| 387 | ) |
| 388 | pos = torch.arange(T, dtype=torch.long) |
| 389 | |
| 390 | qkv_unfused = qkv.clone() |
| 391 | pos_unfused = pos.clone() |
| 392 | |
| 393 | with set_forward_context(None, vllm_config): |
| 394 | forward_context = get_forward_context() |
| 395 | attn_metadata = model.build_attn_metadata(T) |
| 396 | forward_context.slot_mapping = { |
| 397 | model.layer_name: attn_metadata.slot_mapping |
| 398 | } |
| 399 | q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused) |
| 400 | attn_layer = forward_context.no_compile_layers[model.layer_name] |
| 401 | kv_cache_unfused = attn_layer.kv_cache |
| 402 | del dummy |
| 403 | |
| 404 | torch._dynamo.mark_dynamic(qkv, 0) |
| 405 | torch._dynamo.mark_dynamic(pos, 0) |
| 406 | with set_forward_context(None, vllm_config): |
| 407 | model_fused = torch.compile(model, backend=backend) |
| 408 | forward_context = get_forward_context() |
| 409 | attn_metadata = model_fused.build_attn_metadata(T) |
| 410 | forward_context.slot_mapping = { |
| 411 | model.layer_name: attn_metadata.slot_mapping |
| 412 | } |
| 413 | q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos) |
| 414 | attn_layer = forward_context.no_compile_layers[model.layer_name] |
| 415 | kv_cache_fused = attn_layer.kv_cache |
| 416 | del dummy |
| 417 | |
| 418 | assert fusion_pass.matched_count == 1 |
| 419 | |
| 420 | backend.check_before_ops(model.ops_in_model_before()) |
| 421 | backend.check_after_ops(model.ops_in_model_after()) |
| 422 | |
| 423 | if dtype == torch.float16: |
| 424 | ATOL, RTOL = (2e-3, 2e-3) |
| 425 | else: |
| 426 | ATOL, RTOL = (1e-2, 1e-2) |
| 427 | |
| 428 | torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL) |
| 429 | torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL) |
| 430 | torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL) |
| 431 | # Cannot compare fp8_* directly here, cast to model dtype instead |
| 432 | # TODO(charlifu): switch back to ATOL, RTOL after aiter fix is merged. |
| 433 | torch.testing.assert_close( |
| 434 | kv_cache_unfused.to(dtype), |
| 435 | kv_cache_fused.to(dtype), |
| 436 | atol=1e-1, |
| 437 | rtol=1e-1, |
| 438 | ) |
| 439 | |
| 440 | |
| 441 | @pytest.mark.parametrize( |
| 442 | "attn_backend", |
| 443 | [AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN], |
| 444 | ) |
| 445 | @pytest.mark.parametrize("enable_rope_custom_op", [True]) |
| 446 | @pytest.mark.parametrize("enable_aiter_triton_rope", [True, False]) |
| 447 | @pytest.mark.parametrize("num_heads", [64]) |
| 448 | @pytest.mark.parametrize("num_kv_heads", [8]) |
| 449 | @pytest.mark.parametrize("head_size", [64]) |
| 450 | @pytest.mark.parametrize("block_size", [16]) |
| 451 | @pytest.mark.parametrize("is_neox", [True, False]) |
| 452 | @pytest.mark.parametrize("dtype", [torch.bfloat16]) |
| 453 | @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) |
| 454 | @pytest.mark.skipif( |
| 455 | not is_aiter_found_and_supported(), |
| 456 | reason="Only test on ROCm with AITER installed and supported", |
| 457 | ) |
| 458 | @pytest.mark.skipif( |
| 459 | not hasattr(torch.ops._C, "static_scaled_fp8_quant"), |
| 460 | reason="static fp8 quant op not available on this build", |
| 461 | ) |
| 462 | def test_rope_static_qquant_kvcache_fusion( |
| 463 | attn_backend: AttentionBackendEnum, |
| 464 | enable_rope_custom_op: bool, |
| 465 | enable_aiter_triton_rope: bool, |
| 466 | num_heads: int, |
| 467 | num_kv_heads: int, |
| 468 | head_size: int, |
| 469 | block_size: int, |
| 470 | is_neox: bool, |
| 471 | dtype: torch.dtype, |
| 472 | kv_cache_dtype: CacheDType, |
| 473 | monkeypatch: pytest.MonkeyPatch, |
| 474 | ): |
| 475 | torch.set_default_device("cuda") |
| 476 | torch.set_default_dtype(dtype) |
| 477 | torch.manual_seed(0) |
| 478 | |
| 479 | custom_ops: list[str] = [] |
| 480 | if enable_rope_custom_op: |
| 481 | custom_ops.append("+rotary_embedding") |
| 482 | |
| 483 | vllm_config = VllmConfig( |
| 484 | model_config=ModelConfig(dtype=dtype), |
| 485 | cache_config=CacheConfig( |
| 486 | block_size=block_size, |
| 487 | cache_dtype=kv_cache_dtype, |
| 488 | ), |
| 489 | compilation_config=CompilationConfig( |
| 490 | mode=CompilationMode.VLLM_COMPILE, |
| 491 | custom_ops=custom_ops, |
| 492 | pass_config=PassConfig( |
| 493 | fuse_rope_kvcache=True, |
| 494 | eliminate_noops=True, |
| 495 | ), |
| 496 | ), |
| 497 | ) |
| 498 | |
| 499 | with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m: |
| 500 | m.setenv("VLLM_ROCM_USE_AITER", "1") |
| 501 | m.setenv( |
| 502 | "VLLM_ROCM_USE_AITER_TRITON_ROPE", "1" if enable_aiter_triton_rope else "0" |
| 503 | ) |
| 504 | rocm_aiter_ops.refresh_env_variables() |
| 505 | |
| 506 | model = QKRoPEStaticQKVCacheTestModel( |
| 507 | vllm_config=vllm_config, |
| 508 | attn_backend=attn_backend, |
| 509 | num_heads=num_heads, |
| 510 | num_kv_heads=num_kv_heads, |
| 511 | head_size=head_size, |
| 512 | is_neox=is_neox, |
| 513 | dtype=dtype, |
| 514 | device=torch.get_default_device(), |
| 515 | ) |
| 516 | |
| 517 | fusion_pass = RopeKVCacheFusionPass(vllm_config) |
| 518 | passes = [ |
| 519 | NoOpEliminationPass(vllm_config), |
| 520 | SplitCoalescingPass(vllm_config), |
| 521 | ScatterSplitReplacementPass(vllm_config), |
| 522 | fusion_pass, |
| 523 | PostCleanupPass(vllm_config), |
| 524 | ] |
| 525 | backend = TestBackend(*passes) |
| 526 | |
| 527 | T = 5 |
| 528 | qkv = torch.randn( |
| 529 | T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype |
| 530 | ) |
| 531 | pos = torch.arange(T, dtype=torch.long) |
| 532 | |
| 533 | qkv_unfused = qkv.clone() |
| 534 | pos_unfused = pos.clone() |
| 535 | |
| 536 | with set_forward_context(None, vllm_config): |
| 537 | forward_context = get_forward_context() |
| 538 | attn_metadata = model.build_attn_metadata(T) |
| 539 | forward_context.slot_mapping = { |
| 540 | model.layer_name: attn_metadata.slot_mapping |
| 541 | } |
| 542 | q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused) |
| 543 | attn_layer = forward_context.no_compile_layers[model.layer_name] |
| 544 | kv_cache_unfused = attn_layer.kv_cache |
| 545 | del dummy |
| 546 | |
| 547 | torch._dynamo.mark_dynamic(qkv, 0) |
| 548 | torch._dynamo.mark_dynamic(pos, 0) |
| 549 | with set_forward_context(None, vllm_config): |
| 550 | model_fused = torch.compile(model, backend=backend) |
| 551 | forward_context = get_forward_context() |
| 552 | attn_metadata = model_fused.build_attn_metadata(T) |
| 553 | forward_context.slot_mapping = { |
| 554 | model.layer_name: attn_metadata.slot_mapping |
| 555 | } |
| 556 | q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos) |
| 557 | attn_layer = forward_context.no_compile_layers[model.layer_name] |
| 558 | kv_cache_fused = attn_layer.kv_cache |
| 559 | del dummy |
| 560 | |
| 561 | assert fusion_pass.matched_count == 1 |
| 562 | backend.check_before_ops(model.ops_in_model_before()) |
| 563 | backend.check_after_ops(model.ops_in_model_after()) |
| 564 | static_quant_pre = backend.op_count( |
| 565 | torch.ops._C.static_scaled_fp8_quant.default, before=True |
| 566 | ) |
| 567 | static_quant_post = backend.op_count( |
| 568 | torch.ops._C.static_scaled_fp8_quant.default |
| 569 | ) |
| 570 | assert static_quant_pre > 0 |
| 571 | # The replacement still emits static quant, so count is expected to |
| 572 | # remain non-zero after fusion. |
| 573 | assert static_quant_post > 0 |
| 574 | |
| 575 | # Negative control: without the static-Q pattern, the generic RoPE+KV |
| 576 | # pattern cannot match this rope -> static-quant -> kv graph, so the |
| 577 | # fusion above is attributable solely to RopeStaticQQuantKVCachePattern. |
| 578 | # This is a structural property independent of the rope/dtype/neox axes, |
| 579 | # so run the (extra compile) check only once on a representative combo. |
| 580 | if is_neox and enable_aiter_triton_rope and kv_cache_dtype == "auto": |
| 581 | m.setattr( |
| 582 | rope_kvcache_fusion, |
| 583 | "_supports_static_q_fp8_quant_fusion", |
| 584 | lambda: False, |
| 585 | ) |
| 586 | generic_pass = RopeKVCacheFusionPass(vllm_config) |
| 587 | generic_backend = TestBackend( |
| 588 | NoOpEliminationPass(vllm_config), |
| 589 | SplitCoalescingPass(vllm_config), |
| 590 | ScatterSplitReplacementPass(vllm_config), |
| 591 | generic_pass, |
| 592 | PostCleanupPass(vllm_config), |
| 593 | ) |
| 594 | # Reset dynamo so the model is recompiled through generic_backend |
| 595 | # instead of reusing the cached compilation from above. |
| 596 | torch._dynamo.reset() |
| 597 | with set_forward_context(None, vllm_config): |
| 598 | model_generic = torch.compile(model, backend=generic_backend) |
| 599 | forward_context = get_forward_context() |
| 600 | attn_metadata = model_generic.build_attn_metadata(T) |
| 601 | forward_context.slot_mapping = { |
| 602 | model.layer_name: attn_metadata.slot_mapping |
| 603 | } |
| 604 | model_generic(qkv, pos) |
| 605 | # op_count reads the post-pass graph, so it also confirms the pass ran |
| 606 | # (a no-op pass would raise instead of silently passing on count 0). |
| 607 | assert generic_pass.matched_count == 0 |
| 608 | assert ( |
| 609 | generic_backend.op_count( |
| 610 | torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default |
| 611 | ) |
| 612 | == 0 |
| 613 | ) |
| 614 | |
| 615 | if dtype == torch.float16: |
| 616 | ATOL, RTOL = (2e-3, 2e-3) |
| 617 | else: |
| 618 | ATOL, RTOL = (1e-2, 1e-2) |
| 619 | |
| 620 | # TODO(charlifu): switch back to ATOL, RTOL after aiter fix is merged. |
| 621 | torch.testing.assert_close( |
| 622 | q_unfused.to(torch.float32), |
| 623 | q_fused.to(torch.float32), |
| 624 | atol=1e-1, |
| 625 | rtol=1e-1, |
| 626 | ) |
| 627 | torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL) |
| 628 | torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL) |
| 629 | # TODO(charlifu): switch back to ATOL, RTOL after aiter fix is merged. |
| 630 | torch.testing.assert_close( |
| 631 | kv_cache_unfused.to(dtype), |
| 632 | kv_cache_fused.to(dtype), |
| 633 | atol=1e-1, |
| 634 | rtol=1e-1, |
| 635 | ) |