@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2ff4e51152
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 21:30 UTC

[ROCm] AITER fused RoPE+KVCache (#33443)

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com> Signed-off-by: charlifu <charlifu@amd.com> Signed-off-by: Rohan Potdar <66227218+Rohan138@users.noreply.github.com> Co-authored-by: charlifu <charlifu@amd.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>

Rohan Potdar committed Feb 23, 2026 at 21:06 UTC 2ff4e51152d8975303529731069033ba313e1c46
19 files changed +1211 -83
tests/compile/passes/test_functionalization.py
+84 -12
@@ -1,10 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import copy
5 +
6 import pytest
7 import torch
8
7 -import vllm.envs as envs
9 from tests.compile.backend import TestBackend
10 from tests.utils import TestFP8Layer
11 from vllm.compilation.passes.fusion.act_quant_fusion import (
@@ -31,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
32 )
33 from vllm.model_executor.layers.rotary_embedding import get_rope
34 from vllm.platforms import current_platform
35 +from vllm.utils.torch_utils import direct_register_custom_op
36
37 TEST_FP8 = current_platform.supports_fp8()
38 FP8_DTYPE = current_platform.fp8_dtype()
@@ -198,23 +200,82 @@ class TestRotaryEmbeddingSliceScatter(torch.nn.Module):
200 return [torch.ops.aten.slice_scatter.default]
201
202
201 -MODELS = [
202 - TestSiluMul,
203 - TestFusedAddRMSNorm,
204 - TestRotaryEmbedding,
205 - TestRotaryEmbeddingSliceScatter,
206 -]
203 +class TestFunctionWithMutatedArgsAndReturn(torch.nn.Module):
204 + OP_REGISTERED = False
205 +
206 + def __init__(self):
207 + super().__init__()
208 + self.register_test_custom_op()
209 +
210 + @classmethod
211 + def register_test_custom_op(cls):
212 + if not cls.OP_REGISTERED:
213 +
214 + def function_with_mutated_args_and_return_impl(
215 + x: torch.Tensor,
216 + ) -> torch.Tensor:
217 + ret = x + 1
218 + x.add_(2)
219 + return ret
220 +
221 + def function_with_mutated_args_and_return_fake(
222 + x: torch.Tensor,
223 + ) -> torch.Tensor:
224 + return torch.empty_like(x)
225 +
226 + direct_register_custom_op(
227 + op_name="function_with_mutated_args_and_return",
228 + op_func=function_with_mutated_args_and_return_impl,
229 + mutates_args=["x"],
230 + fake_impl=function_with_mutated_args_and_return_fake,
231 + )
232 +
233 + cls.OP_REGISTERED = True
234 +
235 + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
236 + # Clone x to avoid mutating the original tensor
237 + ret = torch.ops.vllm.function_with_mutated_args_and_return(x)
238 + return x, ret
239 +
240 + def example_inputs(self, num_tokens=32):
241 + hidden_states = torch.randn(num_tokens)
242 + return (hidden_states,)
243 +
244 + def ops_in_model(self, do_fusion):
245 + return [torch.ops.vllm.function_with_mutated_args_and_return.default]
246 +
247 + def ops_not_in_model(self):
248 + return []
249 +
250 +
251 +MODELS_AND_DO_FUSION = {
252 + TestSiluMul: [True, False],
253 + TestFusedAddRMSNorm: [True, False],
254 + TestRotaryEmbedding: [False],
255 + TestRotaryEmbeddingSliceScatter: [False],
256 + TestFunctionWithMutatedArgsAndReturn: [False],
257 +}
258
259
260 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
210 -@pytest.mark.parametrize("model_class", MODELS)
211 -@pytest.mark.parametrize("do_fusion", [True, False])
212 -@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda", reason="Only test on CUDA")
261 +@pytest.mark.parametrize(
262 + "model_class, do_fusion",
263 + [
264 + (model_class, do_fusion)
265 + for model_class, fusions in MODELS_AND_DO_FUSION.items()
266 + for do_fusion in fusions
267 + ],
268 +)
269 +@pytest.mark.skipif(
270 + not current_platform.is_cuda_alike(),
271 + reason="Only test on cuda and rocm platform",
272 +)
273 def test_fix_functionalization(
274 model_class: torch.nn.Module, do_fusion: bool, dtype: torch.dtype
275 ):
276 torch.set_default_device("cuda")
277 torch.set_default_dtype(dtype)
278 + torch.manual_seed(0)
279
280 vllm_config = VllmConfig(
281 model_config=ModelConfig(dtype=dtype),
@@ -246,8 +307,14 @@ def test_fix_functionalization(
307 backend_no_func = TestBackend(*passes)
308
309 model = model_class()
249 - torch.compile(model, backend=backend_func)(*model.example_inputs())
250 - torch.compile(model, backend=backend_no_func)(*model.example_inputs())
310 + inputs_func = model.example_inputs()
311 + inputs_no_func = copy.deepcopy(inputs_func)
312 + model_func = model_class()
313 + model_no_func = copy.deepcopy(model_func)
314 + model_func = torch.compile(model_func, backend=backend_func)
315 + model_no_func = torch.compile(model_no_func, backend=backend_no_func)
316 + model_func(*inputs_func)
317 + model_no_func(*inputs_no_func)
318
319 # check if the functionalization pass is applied
320 for op in model.ops_in_model(do_fusion):
@@ -265,3 +332,8 @@ def test_fix_functionalization(
332 found[op] = True
333 assert all(found[op] for op in model.ops_in_model(do_fusion))
334 assert all(not found.get(op) for op in model.ops_not_in_model())
335 +
336 + # TODO (Rohan138): compare the outputs from model_func and model_no_func
337 + # currently runs into errors while comparing `TestFusedAddRMSNorm`
338 + # Linked issue: https://github.com/vllm-project/vllm/issues/34996
339 + # torch.testing.assert_close(outputs_func, outputs_no_func)
tests/compile/passes/test_rope_kvcache_fusion.py new
+325
@@ -0,0 +1,325 @@
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.matcher_utils import ROTARY_OP
12 +from vllm.compilation.passes.fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
13 +from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
14 +from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
15 +from vllm.compilation.passes.utility.scatter_split_replace import (
16 + ScatterSplitReplacementPass,
17 +)
18 +from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
19 +from vllm.config import (
20 + CacheConfig,
21 + CompilationConfig,
22 + CompilationMode,
23 + ModelConfig,
24 + PassConfig,
25 + VllmConfig,
26 +)
27 +from vllm.forward_context import get_forward_context, set_forward_context
28 +from vllm.model_executor.layers.attention import Attention
29 +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
30 +from vllm.platforms import current_platform
31 +from vllm.v1.attention.backend import (
32 + AttentionBackend,
33 + CommonAttentionMetadata,
34 +)
35 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
36 +from vllm.v1.kv_cache_interface import AttentionSpec
37 +
38 +INDEX_SELECT_OP = torch.ops.aten.index.Tensor
39 +VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
40 +FP8_DTYPE = current_platform.fp8_dtype()
41 +
42 +
43 +class QKRoPEKVCacheTestModel(torch.nn.Module):
44 + def __init__(
45 + self,
46 + vllm_config: VllmConfig,
47 + attn_backend: AttentionBackendEnum,
48 + num_heads: int,
49 + num_kv_heads: int,
50 + head_size: int,
51 + is_neox: bool,
52 + dtype: torch.dtype,
53 + device: torch.device,
54 + prefix: str = "model.layers.0.self_attn.attn",
55 + ):
56 + super().__init__()
57 + self.num_heads = num_heads
58 + self.num_kv_heads = num_kv_heads
59 + self.head_size = head_size
60 + self.block_size = vllm_config.cache_config.block_size
61 + self.q_size = num_heads * head_size
62 + self.kv_size = num_kv_heads * head_size
63 + self.is_neox = is_neox
64 + self.dtype = dtype
65 + self.device = device
66 + self.layer_name = prefix
67 +
68 + self.rotary_emb = RotaryEmbedding(
69 + head_size,
70 + rotary_dim=head_size,
71 + max_position_embeddings=4096,
72 + base=10000,
73 + is_neox_style=is_neox,
74 + dtype=self.dtype,
75 + )
76 +
77 + # Whether to check for the RoPE custom op or component index_select
78 + self.enable_rope_custom_op = self.rotary_emb.enabled()
79 +
80 + # Register layer metadata for the fusion pass via Attention.
81 + self.attn = Attention(
82 + num_heads=num_heads,
83 + head_size=head_size,
84 + scale=1.0 / head_size**0.5,
85 + num_kv_heads=num_kv_heads,
86 + cache_config=vllm_config.cache_config,
87 + quant_config=vllm_config.quant_config,
88 + prefix=prefix,
89 + attn_backend=attn_backend.get_class(),
90 + )
91 + self.attn_backend: type[AttentionBackend] = self.attn.get_attn_backend()
92 + assert not self.attn_backend.forward_includes_kv_cache_update, (
93 + f"Attention backend {self.attn_backend} does not support fuse_rope_kvcache."
94 + )
95 + self.attn._k_scale = self.attn._k_scale.to(device)
96 + self.attn._v_scale = self.attn._v_scale.to(device)
97 +
98 + kv_cache_dtype_str = vllm_config.cache_config.cache_dtype
99 + self.kv_cache_dtype = (
100 + FP8_DTYPE if kv_cache_dtype_str.startswith("fp8") else self.dtype
101 + )
102 +
103 + # Initialize attn MetadataBuilder
104 + self.builder = self.attn.attn_backend.get_builder_cls()(
105 + kv_cache_spec=AttentionSpec(
106 + block_size=self.block_size,
107 + num_kv_heads=self.num_kv_heads,
108 + head_size=head_size,
109 + dtype=self.kv_cache_dtype,
110 + ),
111 + layer_names=[self.attn.layer_name],
112 + vllm_config=vllm_config,
113 + device=device,
114 + )
115 +
116 + def build_attn_metadata(self, batch_size: int) -> CommonAttentionMetadata:
117 + """Initialize attention metadata."""
118 + # Create common attn metadata
119 + batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
120 + common_attn_metadata = create_common_attn_metadata(
121 + batch_spec, self.block_size, self.device, arange_block_indices=True
122 + )
123 +
124 + max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
125 + num_blocks = batch_size * max_blocks
126 +
127 + # Fetch the attention backend and kv cache shape and stride order
128 + attn_backend = self.attn.attn_backend
129 + kv_cache_shape = attn_backend.get_kv_cache_shape(
130 + num_blocks, self.block_size, self.num_kv_heads, self.head_size
131 + )
132 + try:
133 + kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
134 + except (AttributeError, NotImplementedError):
135 + kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
136 +
137 + kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
138 + inv_order = [
139 + kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
140 + ]
141 +
142 + # Create dummy KV cache
143 + raw_tensor = torch.zeros(
144 + 2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
145 + dtype=self.kv_cache_dtype,
146 + device=self.device,
147 + )
148 + raw_tensor = raw_tensor.view(kv_cache_shape)
149 + kv_cache = raw_tensor.permute(*inv_order)
150 +
151 + self.attn.kv_cache = [kv_cache]
152 +
153 + # Build attn metadata
154 + attn_metadata = self.builder.build(
155 + common_prefix_len=0, common_attn_metadata=common_attn_metadata
156 + )
157 +
158 + return attn_metadata
159 +
160 + def forward(
161 + self, qkv: torch.Tensor, positions: torch.Tensor
162 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
163 + # Create copy so inplace ops do not modify the original tensors
164 + qkv = qkv.clone()
165 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
166 + q, k = self.rotary_emb(positions, q, k)
167 +
168 + # Instead of a full forward pass, match only the KV cache update op here
169 + q = q.view(-1, self.num_heads, self.head_size)
170 + k = k.view(-1, self.num_kv_heads, self.head_size)
171 + v = v.view(-1, self.num_kv_heads, self.head_size)
172 + kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
173 + k, v, self.layer_name
174 + )
175 + return q, k, v, kv_cache_dummy_dep
176 +
177 + def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
178 + ops = []
179 + if self.enable_rope_custom_op:
180 + ops.append(ROTARY_OP)
181 + else:
182 + ops.append(INDEX_SELECT_OP)
183 + ops.append(torch.ops.vllm.unified_kv_cache_update.default)
184 + return ops
185 +
186 + def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
187 + return [torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default]
188 +
189 +
190 +@pytest.mark.parametrize(
191 + "attn_backend",
192 + [
193 + AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
194 + AttentionBackendEnum.TRITON_ATTN,
195 + AttentionBackendEnum.ROCM_ATTN,
196 + ],
197 +)
198 +@pytest.mark.parametrize("enable_rope_custom_op", [True]) # [True, False])
199 +@pytest.mark.parametrize("num_heads", [64])
200 +@pytest.mark.parametrize("num_kv_heads", [8])
201 +@pytest.mark.parametrize("head_size", [64])
202 +@pytest.mark.parametrize("block_size", [16])
203 +@pytest.mark.parametrize("is_neox", [True, False])
204 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
205 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
206 +@pytest.mark.skipif(
207 + not is_aiter_found_and_supported(),
208 + reason="Only test on ROCm with AITER installed and supported",
209 +)
210 +def test_rope_kvcache_fusion(
211 + attn_backend: AttentionBackendEnum,
212 + enable_rope_custom_op: bool,
213 + num_heads: int,
214 + num_kv_heads: int,
215 + head_size: int,
216 + block_size: int,
217 + is_neox: bool,
218 + dtype: torch.dtype,
219 + kv_cache_dtype: str,
220 + monkeypatch: pytest.MonkeyPatch,
221 +):
222 + torch.set_default_device("cuda")
223 + torch.set_default_dtype(dtype)
224 + torch.manual_seed(0)
225 +
226 + custom_ops: list[str] = []
227 + if enable_rope_custom_op:
228 + custom_ops.append("+rotary_embedding")
229 +
230 + vllm_config = VllmConfig(
231 + model_config=ModelConfig(dtype=dtype),
232 + cache_config=CacheConfig(
233 + block_size=block_size,
234 + cache_dtype=kv_cache_dtype,
235 + ),
236 + compilation_config=CompilationConfig(
237 + mode=CompilationMode.VLLM_COMPILE,
238 + custom_ops=custom_ops,
239 + pass_config=PassConfig(
240 + fuse_rope_kvcache=True,
241 + eliminate_noops=True,
242 + ),
243 + ),
244 + )
245 +
246 + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
247 + m.setenv("VLLM_ROCM_USE_AITER", "1")
248 + rocm_aiter_ops.refresh_env_variables()
249 +
250 + model = QKRoPEKVCacheTestModel(
251 + vllm_config=vllm_config,
252 + attn_backend=attn_backend,
253 + num_heads=num_heads,
254 + num_kv_heads=num_kv_heads,
255 + head_size=head_size,
256 + is_neox=is_neox,
257 + dtype=dtype,
258 + device=torch.get_default_device(),
259 + )
260 +
261 + fusion_pass = RopeKVCacheFusionPass(vllm_config)
262 + passes = [
263 + NoOpEliminationPass(vllm_config),
264 + SplitCoalescingPass(vllm_config),
265 + ScatterSplitReplacementPass(vllm_config),
266 + fusion_pass,
267 + PostCleanupPass(vllm_config),
268 + ]
269 + backend = TestBackend(*passes)
270 +
271 + T = 5
272 +
273 + qkv = torch.randn(
274 + T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
275 + )
276 + pos = torch.arange(T, dtype=torch.long)
277 +
278 + qkv_unfused = qkv.clone()
279 + pos_unfused = pos.clone()
280 +
281 + with set_forward_context(None, vllm_config):
282 + forward_context = get_forward_context()
283 + attn_metadata = model.build_attn_metadata(T)
284 + forward_context.slot_mapping = {
285 + model.layer_name: attn_metadata.slot_mapping
286 + }
287 + q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)
288 + attn_layer = forward_context.no_compile_layers[model.layer_name]
289 + kv_cache_unfused = attn_layer.kv_cache[forward_context.virtual_engine]
290 + del dummy
291 +
292 + torch._dynamo.mark_dynamic(qkv, 0)
293 + torch._dynamo.mark_dynamic(pos, 0)
294 + with set_forward_context(None, vllm_config):
295 + model_fused = torch.compile(model, backend=backend)
296 + forward_context = get_forward_context()
297 + attn_metadata = model_fused.build_attn_metadata(T)
298 + forward_context.slot_mapping = {
299 + model.layer_name: attn_metadata.slot_mapping
300 + }
301 + q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)
302 + attn_layer = forward_context.no_compile_layers[model.layer_name]
303 + kv_cache_fused = attn_layer.kv_cache[forward_context.virtual_engine]
304 + del dummy
305 +
306 + assert fusion_pass.matched_count == 1
307 +
308 + backend.check_before_ops(model.ops_in_model_before())
309 + backend.check_after_ops(model.ops_in_model_after())
310 +
311 + if dtype == torch.float16:
312 + ATOL, RTOL = (2e-3, 2e-3)
313 + else:
314 + ATOL, RTOL = (1e-2, 1e-2)
315 +
316 + torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
317 + torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
318 + torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
319 + # Cannot compare fp8_* directly here, cast to model dtype instead
320 + torch.testing.assert_close(
321 + kv_cache_unfused.view(dtype),
322 + kv_cache_fused.view(dtype),
323 + atol=ATOL,
324 + rtol=RTOL,
325 + )
tests/compile/passes/test_scatter_split_replace.py new
+107
@@ -0,0 +1,107 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +import torch.nn as nn
7 +
8 +import vllm
9 +from tests.compile.backend import TestBackend
10 +from vllm.compilation.passes.utility.scatter_split_replace import (
11 + ScatterSplitReplacementPass,
12 +)
13 +from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
14 +from vllm.config import CompilationConfig, CompilationMode, VllmConfig
15 +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
16 +
17 +
18 +class ScatterSplitReplacementModel(nn.Module):
19 + """Model with a rope+getitem+slice_scatter+split_with_sizes sequence."""
20 +
21 + def __init__(
22 + self,
23 + num_heads: int,
24 + num_kv_heads: int,
25 + head_size: int,
26 + dtype: torch.dtype,
27 + ):
28 + super().__init__()
29 + self.q_size = num_heads * head_size
30 + self.kv_size = num_kv_heads * head_size
31 +
32 + self.rotary_emb = RotaryEmbedding(
33 + head_size,
34 + rotary_dim=head_size,
35 + max_position_embeddings=4096,
36 + base=10000,
37 + is_neox_style=True,
38 + dtype=dtype,
39 + )
40 +
41 + def forward(self, qkv: torch.Tensor, positions: torch.Tensor):
42 + # Create copy so inplace ops do not modify the original tensors
43 + qkv = qkv.clone()
44 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
45 + q, k = self.rotary_emb(positions, q, k)
46 + q = q + 1
47 + k = k + 2
48 + v = v + 3
49 + return q, k, v
50 +
51 + def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
52 + return [
53 + torch.ops.aten.slice_scatter.default,
54 + torch.ops.aten.split_with_sizes.default,
55 + torch.ops.aten.getitem.default,
56 + ]
57 +
58 + def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
59 + return [torch.ops.aten.getitem.default]
60 +
61 +
62 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
63 +def test_scatter_split_replace(dtype):
64 + torch.set_default_device("cuda")
65 + torch.set_default_dtype(dtype)
66 + torch.manual_seed(0)
67 +
68 + num_heads = 8
69 + num_kv_heads = 4
70 + head_size = 64
71 +
72 + vllm_config = VllmConfig(
73 + compilation_config=CompilationConfig(
74 + mode=CompilationMode.VLLM_COMPILE,
75 + custom_ops=["+rotary_embedding"],
76 + ),
77 + )
78 + with vllm.config.set_current_vllm_config(vllm_config):
79 + # ScatterSplitReplacementPass requires SplitCoalescingPass to be run before it
80 + coalesce_pass = SplitCoalescingPass(vllm_config)
81 + replace_pass = ScatterSplitReplacementPass(vllm_config)
82 + passes = [coalesce_pass, replace_pass]
83 + backend = TestBackend(*passes)
84 +
85 + model = ScatterSplitReplacementModel(num_heads, num_kv_heads, head_size, dtype)
86 +
87 + T = 5
88 + qkv = torch.randn(
89 + T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
90 + )
91 + pos = torch.arange(T, dtype=torch.long)
92 +
93 + qkv_eager = qkv.clone()
94 + pos_eager = pos.clone()
95 + result_eager = model(qkv_eager, pos_eager)
96 +
97 + torch._dynamo.mark_dynamic(qkv, 0)
98 + torch._dynamo.mark_dynamic(pos, 0)
99 +
100 + model_compiled = torch.compile(model, backend=backend)
101 + result_compiled = model_compiled(qkv, pos)
102 +
103 + for eager, compiled in zip(result_eager, result_compiled):
104 + torch.testing.assert_close(eager, compiled)
105 +
106 + assert backend.op_count(torch.ops.aten.slice_scatter.default) == 0
107 + assert backend.op_count(torch.ops.aten.split_with_sizes.default) == 1
tests/v1/attention/test_attention_backends.py
+1 -1
@@ -179,7 +179,7 @@ def create_and_prepopulate_kv_cache(
179 block_table[i, :num_blocks_for_seq] = inv_perm[start:end]
180 start_block_idx += num_blocks_for_seq
181
182 - # Create a realistic slot mapping that corresponds to the block table
182 + # Create a realistic slot mapping that corresponds to the block table
183 for i in range(batch_size):
184 token_offsets = torch.arange(int(query_lens[i])) + int(context_lens[i])
185 block_indices = token_offsets // block_size
vllm/_aiter_ops.py
+39
@@ -1518,6 +1518,45 @@ class rocm_aiter_ops:
1518 query = query.view(query_shape)
1519 key = key.view(key_shape)
1520
1521 + @staticmethod
1522 + def triton_rope_and_cache(
1523 + query: torch.Tensor,
1524 + key: torch.Tensor,
1525 + value: torch.Tensor,
1526 + positions: torch.Tensor,
1527 + cos_sin_cache: torch.Tensor,
1528 + is_neox: bool,
1529 + key_cache: torch.Tensor,
1530 + value_cache: torch.Tensor,
1531 + layer_slot_mapping: torch.Tensor,
1532 + k_scale: torch.Tensor,
1533 + v_scale: torch.Tensor,
1534 + flash_layout: bool,
1535 + apply_scale: bool,
1536 + ):
1537 + from aiter.ops.triton.fused_kv_cache import fused_qk_rope_reshape_and_cache
1538 +
1539 + cos, sin = cos_sin_cache.chunk(2, dim=-1)
1540 + fused_qk_rope_reshape_and_cache(
1541 + query,
1542 + key,
1543 + value,
1544 + key_cache,
1545 + value_cache,
1546 + layer_slot_mapping,
1547 + positions,
1548 + cos,
1549 + sin,
1550 + k_scale,
1551 + v_scale,
1552 + is_neox,
1553 + flash_layout=flash_layout,
1554 + apply_scale=apply_scale,
1555 + q_out=query,
1556 + k_out=key,
1557 + output_zeros=False,
1558 + )
1559 +
1560 @staticmethod
1561 def batched_gemm_a16wfp4(
1562 X: torch.Tensor,
vllm/compilation/passes/fusion/rope_kvcache_fusion.py new
+230
@@ -0,0 +1,230 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +import torch._inductor.pattern_matcher as pm
6 +from torch import fx
7 +from torch._higher_order_ops import auto_functionalized
8 +from torch._inductor.fx_passes.post_grad import view_to_reshape
9 +from torch._inductor.pattern_matcher import PatternMatcherPass
10 +
11 +from vllm.config import VllmConfig, get_layers_from_vllm_config
12 +from vllm.config.utils import Range
13 +from vllm.logger import init_logger
14 +from vllm.model_executor.layers.attention.attention import (
15 + Attention,
16 + get_attention_context,
17 +)
18 +from vllm.utils.torch_utils import direct_register_custom_op
19 +
20 +from ..inductor_pass import enable_fake_mode
21 +from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
22 +from .matcher_utils import (
23 + MatcherRotaryEmbedding,
24 +)
25 +from .rms_quant_fusion import (
26 + empty_bf16,
27 + empty_i64,
28 +)
29 +
30 +logger = init_logger(__name__)
31 +
32 +
33 +def fused_rope_and_unified_kv_cache_update_impl(
34 + query: torch.Tensor,
35 + key: torch.Tensor,
36 + value: torch.Tensor,
37 + positions: torch.Tensor,
38 + cos_sin_cache: torch.Tensor,
39 + is_neox: bool,
40 + layer_name: str = "",
41 +) -> torch.Tensor:
42 + """
43 + This impl fetches the KV cache and slot mapping from the forward context,
44 + then calls the layer impl's `AttentionImpl.do_rope_and_kv_cache_update` method.
45 + It also returns a dummy tensor, similar to `Attention.unified_kv_cache_update`,
46 + that is passed to unified_attention to signal a side effect and
47 + the data dependency between them to ensure torch.compile preserves ordering.
48 + """
49 + _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
50 + if layer_slot_mapping is not None:
51 + attn_layer.impl.do_rope_and_kv_cache_update(
52 + attn_layer,
53 + query,
54 + key,
55 + value,
56 + positions,
57 + cos_sin_cache,
58 + is_neox,
59 + kv_cache,
60 + layer_slot_mapping,
61 + )
62 +
63 + return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype)
64 +
65 +
66 +def fused_rope_and_unified_kv_cache_update_fake(
67 + query: torch.Tensor,
68 + key: torch.Tensor,
69 + value: torch.Tensor,
70 + positions: torch.Tensor,
71 + cos_sin_cache: torch.Tensor,
72 + is_neox: bool,
73 + layer_name: str = "",
74 +) -> torch.Tensor:
75 + return torch.empty(0, device=query.device, dtype=query.dtype)
76 +
77 +
78 +direct_register_custom_op(
79 + op_name="fused_rope_and_unified_kv_cache_update",
80 + op_func=fused_rope_and_unified_kv_cache_update_impl,
81 + mutates_args=["query", "key"],
82 + fake_impl=fused_rope_and_unified_kv_cache_update_fake,
83 +)
84 +
85 +
86 +class RopeReshapeKVCachePattern:
87 + """
88 + This pattern matches the following unfused inplace ops:
89 + q, k = rotary_embedding(positions, q, k, head_size, cos_sin_cache, is_neox)
90 + kv_cache_dummy = unified_kv_cache_update(k, v, layer_name)
91 +
92 + and replaces it with the fused inplace op:
93 + kv_cache_dummy = fused_rope_and_unified_kv_cache_update(
94 + q, k, v, positions, cos_sin_cache, is_neox, layer_name
95 + )
96 + """
97 +
98 + FUSED_OP = torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default
99 +
100 + def __init__(
101 + self,
102 + layer: Attention,
103 + is_neox: bool,
104 + ) -> None:
105 + self.layer_name = layer.layer_name
106 + self.num_heads = layer.num_heads
107 + self.num_kv_heads = layer.num_kv_heads
108 + self.head_size = layer.head_size
109 + self.head_size_v = layer.head_size_v
110 + self.is_neox = is_neox
111 +
112 + self.q_size = self.num_heads * self.head_size
113 + self.k_size = self.num_kv_heads * self.head_size
114 + self.v_size = self.num_kv_heads * self.head_size_v
115 +
116 + self.rope_matcher = MatcherRotaryEmbedding(
117 + is_neox=self.is_neox,
118 + head_size=self.head_size,
119 + num_heads=self.num_heads,
120 + num_kv_heads=self.num_kv_heads,
121 + )
122 +
123 + def get_inputs(self) -> list[torch.Tensor]:
124 + # Sample inputs to help pattern tracing
125 + T = 5
126 + L = 4096
127 + qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size)
128 + positions = empty_i64(T)
129 + cos_sin_cache = empty_bf16(L, self.head_size)
130 + return [
131 + qkv,
132 + positions,
133 + cos_sin_cache,
134 + ]
135 +
136 + def register(self, pm_pass: PatternMatcherPass) -> None:
137 + def pattern(
138 + qkv: torch.Tensor,
139 + positions: torch.Tensor,
140 + cos_sin_cache: torch.Tensor,
141 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
142 + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
143 + q, k = self.rope_matcher(positions, q, k, cos_sin_cache)
144 + q = q.view(-1, self.num_heads, self.head_size)
145 + k = k.view(-1, self.num_kv_heads, self.head_size)
146 + v = v.view(-1, self.num_kv_heads, self.head_size_v)
147 + dummy = torch.ops.vllm.unified_kv_cache_update(k, v, self.layer_name)
148 + return dummy, q, k, v
149 +
150 + def replacement(
151 + qkv: torch.Tensor,
152 + positions: torch.Tensor,
153 + cos_sin_cache: torch.Tensor,
154 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
155 + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
156 + q = q.view(-1, self.num_heads, self.head_size)
157 + k = k.view(-1, self.num_kv_heads, self.head_size)
158 + v = v.view(-1, self.num_kv_heads, self.head_size_v)
159 + results = auto_functionalized(
160 + self.FUSED_OP,
161 + query=q,
162 + key=k,
163 + value=v,
164 + positions=positions,
165 + cos_sin_cache=cos_sin_cache,
166 + is_neox=self.is_neox,
167 + layer_name=self.layer_name,
168 + )
169 + return results[0], results[1], results[2], v
170 +
171 + # NOTE: use view_to_reshape to unify view/reshape to simplify
172 + # pattern and increase matching opportunities
173 + def fwd_and_view_to_reshape(*args, **kwargs) -> fx.GraphModule:
174 + gm = pm.fwd_only(*args, **kwargs)
175 + view_to_reshape(gm)
176 + return gm
177 +
178 + pm.register_replacement(
179 + pattern, replacement, self.get_inputs(), fwd_and_view_to_reshape, pm_pass
180 + )
181 +
182 +
183 +class RopeKVCacheFusionPass(VllmPatternMatcherPass):
184 + """
185 + This pass fuses the rotary embedding and KV cache update operations
186 + into a single fused kernel if available.
187 +
188 + It uses the pattern matcher and matches each layer manually, as strings
189 + cannot be wildcarded. This also lets us check support on attention layers
190 + upon registration instead of during pattern matching.
191 +
192 + This fusion eliminates the need for separate kernel launches and
193 + intermediate memory operations between the RoPE and cache update steps.
194 + """
195 +
196 + @enable_fake_mode
197 + def __init__(self, config: VllmConfig) -> None:
198 + super().__init__(config)
199 +
200 + self.patterns: PatternMatcherPass = PatternMatcherPass(
201 + pass_name="rope_kv_cache_fusion_pass"
202 + )
203 +
204 + cc = config.compilation_config
205 + self.max_token_num = cc.pass_config.rope_kvcache_fusion_max_token_num
206 +
207 + attn_layers = get_layers_from_vllm_config(config, Attention)
208 + for _, layer in attn_layers.items():
209 + if layer.impl.fused_rope_kvcache_supported():
210 + for is_neox in [True, False]:
211 + RopeReshapeKVCachePattern(
212 + layer=layer,
213 + is_neox=is_neox,
214 + ).register(self.patterns)
215 +
216 + self.dump_patterns(config, self.patterns)
217 +
218 + @VllmInductorPass.time_and_log
219 + def __call__(self, graph: fx.Graph) -> None:
220 + self.matched_count = self.patterns.apply(graph)
221 + logger.debug("Replaced %s patterns", self.matched_count)
222 +
223 + def is_applicable_for_range(self, compile_range: Range) -> bool:
224 + # This pass works best for the small-batch decode setting.
225 + # For large-batch e.g. prefill, it is better to use two separate kernels
226 + # since they are compute bound and the fused kernels require further tuning.
227 + return compile_range.end <= self.max_token_num
228 +
229 + def uuid(self) -> str:
230 + return VllmInductorPass.hash_source(self, RopeReshapeKVCachePattern)
vllm/compilation/passes/pass_manager.py
+7
@@ -28,7 +28,9 @@ if current_platform.is_cuda_alike():
28 from .fusion.attn_quant_fusion import AttnFusionPass
29 from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
30 from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
31 + from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
32 from .fusion.sequence_parallelism import SequenceParallelismPass
33 + from .utility.scatter_split_replace import ScatterSplitReplacementPass
34 from .utility.split_coalescing import SplitCoalescingPass
35
36 if current_platform.is_cuda():
@@ -136,6 +138,11 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
138 if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
139 self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
140
141 + if self.pass_config.fuse_rope_kvcache:
142 + self.passes += [SplitCoalescingPass(config)]
143 + self.passes += [ScatterSplitReplacementPass(config)]
144 + self.passes += [RopeKVCacheFusionPass(config)]
145 +
146 if self.pass_config.fuse_attn_quant:
147 self.passes += [AttnFusionPass(config)]
148
vllm/compilation/passes/utility/fix_functionalization.py
+38 -3
@@ -162,6 +162,24 @@ class FixFunctionalizationPass(VllmInductorPass):
162 "position_ids",
163 )
164 self.defunctionalize(graph, node, mutated_args=mutated_args, args=args)
165 + elif (
166 + hasattr(torch.ops.vllm, "fused_rope_and_unified_kv_cache_update")
167 + and at_target
168 + == torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default
169 + ):
170 + mutated_args = {
171 + 1: "query",
172 + 2: "key",
173 + }
174 + self.defunctionalize(graph, node, mutated_args=mutated_args)
175 + # only used for test_functionalization::TestFunctionWithMutatedArgsAndReturn
176 + elif (
177 + hasattr(torch.ops.vllm, "function_with_mutated_args_and_return")
178 + and at_target
179 + == torch.ops.vllm.function_with_mutated_args_and_return.default
180 + ):
181 + mutated_args = {1: "x"}
182 + self.defunctionalize(graph, node, mutated_args=mutated_args)
183 else:
184 continue # skip the count
185
@@ -208,13 +226,20 @@ class FixFunctionalizationPass(VllmInductorPass):
226 self, node: torch.fx.Node, mutated_args: dict[int, torch.fx.Node | str]
227 ) -> None:
228 """
211 - Replace all getitem users of the auto-functionalized node with the
229 + Replace mutated getitem users of the auto-functionalized node with the
230 mutated arguments.
231 :param node: The auto-functionalized node
232 :param mutated_args: The mutated arguments, indexed by getitem index.
233 If the value of an arg is a string, `node.kwargs[arg]` is used.
234 """
235 for idx, user in self.getitem_users(node).items():
236 + # Some functionalized nodes may return both a result at getitem[0]
237 + # as well as mutated args at getitem[1:...]
238 + if idx == 0:
239 + assert idx not in mutated_args, (
240 + f"result at getitem[0] should not be in mutated_args for {node}"
241 + )
242 + continue
243 arg = mutated_args[idx]
244 arg = node.kwargs[arg] if isinstance(arg, str) else arg
245 user.replace_all_uses_with(arg)
@@ -257,10 +282,20 @@ class FixFunctionalizationPass(VllmInductorPass):
282 with graph.inserting_before(node):
283 function = node.args[0]
284 if args is None:
260 - graph.call_function(function, kwargs=node.kwargs)
285 + fn_node = graph.call_function(function, kwargs=node.kwargs)
286 else:
287 # Args passed as strings refer to items in node.kwargs
288 args = tuple(
289 node.kwargs[arg] if isinstance(arg, str) else arg for arg in args
290 )
266 - graph.call_function(function, args=args)
291 + fn_node = graph.call_function(function, args=args)
292 +
293 + # If the function returns a value as well as mutating args inplace,
294 + # the functionalized node will have a getitem[0] user that holds this value
295 + # Replace getitem[0] user of the auto-functionalized node
296 + # with the new defunctionalized node directly if it exists
297 + users = self.getitem_users(node)
298 + if 0 in users:
299 + user = users[0]
300 + user.replace_all_uses_with(fn_node)
301 + self._remove(user)
vllm/compilation/passes/utility/scatter_split_replace.py new
+134
@@ -0,0 +1,134 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Replace ``slice_scatter`` and ``split_with_sizes`` nodes with a single
5 +assignment if there are no users for the inplace tensor written to by
6 +the slice_scatter call.
7 +
8 +The inplace rotary_embedding custom op takes in mutable query and key inputs
9 +that are split+getitem outputs of a single qkv tensor.
10 +When functionalized, we fetch the rotated query and key from the functionalized op
11 +using `getitem` calls. However, we also write to the qkv tensor inplace using a
12 +`slice_scatter`, then split the inplace tensor to get the output tensors again.
13 +Instead, if the inplace tensor has no subsequent users, we can just replace the
14 +`slice_scatter` and `split_with_sizes` nodes with the `getitem` calls.
15 +
16 +This is already done in fix_functionalization::FixFunctionalizationPass, but
17 +writing a custom pass for it before defunctionalization allows matching against the
18 +qkv split+rotary_embedding subpattern as part of e.g. the RoPE+KVCache fusion pass.
19 +"""
20 +
21 +import operator
22 +
23 +import torch
24 +from torch import fx
25 +from torch._higher_order_ops.auto_functionalize import auto_functionalized
26 +
27 +from vllm.logger import init_logger
28 +
29 +from ..fx_utils import is_func
30 +from ..vllm_inductor_pass import VllmInductorPass
31 +
32 +logger = init_logger(__name__)
33 +
34 +
35 +class ScatterSplitReplacementPass(VllmInductorPass):
36 + """Replace getitem+slice_scatter+split nodes with a single getitem when
37 + the inplace subtensor written to by the slice_scatter has no other users.
38 +
39 + Here's an example graph with q_size = 512, kv_size = 64:
40 + split_with_sizes_1 = torch.ops.aten.split_with_sizes.default(qkv, (512, 64, 64), -1)
41 + at = auto_functionalized(torch.ops._C.rotary_embedding.default(positions, q, k))
42 + q = operator.getitem(at, 1)
43 + k = operator.getitem(at, 2)
44 + torch.ops.aten.slice_scatter.default(qkv, q, [0, 512], -1)
45 + torch.ops.aten.slice_scatter.default(qkv, k, [512, 512 + 64], -1)
46 + split_with_sizes_2 = torch.ops.aten.split_with_sizes.default(qkv, (512, 64, 64), -1)
47 + q = operator.getitem(split_with_sizes_2, 0)
48 + k = operator.getitem(split_with_sizes_2, 1)
49 + v = operator.getitem(split_with_sizes_2, 2)
50 +
51 + After this pass, this sequence of nodes is replaced with:
52 + split_with_sizes_1 = torch.ops.aten.split_with_sizes.default(qkv, (512, 64, 64), -1)
53 + at = auto_functionalized(torch.ops._C.rotary_embedding.default(positions, q, k))
54 + q = operator.getitem(at, 1)
55 + k = operator.getitem(at, 2)
56 + v = operator.getitem(split_with_sizes_1, 2)
57 + """
58 +
59 + @VllmInductorPass.time_and_log
60 + def __call__(self, graph: fx.Graph) -> None:
61 + count = 0
62 +
63 + for node in graph.nodes:
64 + if not is_func(node, auto_functionalized):
65 + continue
66 +
67 + kwargs = node.kwargs
68 + at_target = node.args[0]
69 +
70 + if at_target == torch.ops._C.rotary_embedding.default:
71 + query = kwargs["query"]
72 + key = kwargs["key"]
73 + getitem_nodes = {}
74 + for user in node.users:
75 + if is_func(user, operator.getitem):
76 + getitem_nodes[user.args[1]] = user
77 +
78 + if (
79 + is_func(query, operator.getitem)
80 + and is_func(key, operator.getitem)
81 + and query.args[0] == key.args[0]
82 + and is_func(query.args[0], torch.ops.aten.split_with_sizes.default)
83 + and all(
84 + is_func(user, torch.ops.aten.slice_scatter.default)
85 + for getitem_node in getitem_nodes.values()
86 + for user in getitem_node.users
87 + )
88 + ):
89 + # Pattern where query and key are slices of a qkv tensor.
90 + # While functionalized, results at [1] and [2] are scattered
91 + # back into qkv, then split again to get query and key.
92 + # If the inplace tensor has no other users, we can replace
93 + # the slice_scatter+split nodes with the original results.
94 + for user in getitem_nodes[1].users:
95 + slice_scatter_1_node = user
96 + if not is_func(
97 + slice_scatter_1_node, torch.ops.aten.slice_scatter.default
98 + ):
99 + continue
100 +
101 + for user in getitem_nodes[2].users:
102 + slice_scatter_2_node = user
103 + if not is_func(
104 + slice_scatter_2_node, torch.ops.aten.slice_scatter.default
105 + ):
106 + continue
107 +
108 + for user in slice_scatter_2_node.users:
109 + split_node = user
110 + if not is_func(split_node, torch.ops.aten.split_with_sizes.default):
111 + continue
112 +
113 + split_getitem_users = {}
114 + for user in split_node.users:
115 + if is_func(user, operator.getitem):
116 + split_getitem_users[user.args[1]] = user
117 +
118 + # Replace query node
119 + split_getitem_users[0].replace_all_uses_with(getitem_nodes[1])
120 + graph.erase_node(split_getitem_users[0])
121 + # Replace key node
122 + split_getitem_users[1].replace_all_uses_with(getitem_nodes[2])
123 + graph.erase_node(split_getitem_users[1])
124 + # Redirect value node to original qkv tensor
125 + split_getitem_users[2].replace_input_with(split_node, query.args[0])
126 +
127 + # Erase unused nodes
128 + graph.erase_node(split_node)
129 + graph.erase_node(slice_scatter_2_node)
130 + graph.erase_node(slice_scatter_1_node)
131 +
132 + count += 1
133 +
134 + logger.debug("Eliminated %d slice_scatter+split nodes", count)
vllm/config/compilation.py
+27
@@ -127,6 +127,13 @@ class PassConfig:
127 # ROCm/AITER specific fusions
128 fuse_act_padding: bool = Field(default=None)
129 """Fuse the custom RMSNorm + padding ops."""
130 + fuse_rope_kvcache: bool = Field(default=None)
131 + """Fuse the QK rope + KV cache ops."""
132 +
133 + rope_kvcache_fusion_max_token_num: int = 256
134 + """The threshold for ROCm AITER RoPE+KVCache fusion e.g. for small batch decode.
135 + Larger batch sizes e.g. during prefill will use the unfused kernels.
136 + """
137
138 fi_allreduce_fusion_max_size_mb: float | None = None
139 """The threshold of the communicated tensor sizes under which
@@ -198,6 +205,7 @@ class PassConfig:
205 "fuse_gemm_comms",
206 "fuse_allreduce_rms",
207 "fuse_act_padding",
208 + "fuse_rope_kvcache",
209 mode="wrap",
210 )
211 @classmethod
@@ -243,6 +251,12 @@ class PassConfig:
251 "The fusion will be disabled."
252 )
253 self.fuse_act_padding = False
254 + if self.fuse_rope_kvcache and not current_platform.is_rocm():
255 + logger.warning_once(
256 + "KV cache fusion currently only enabled on ROCm. "
257 + "The fusion will be disabled."
258 + )
259 + self.fuse_rope_kvcache = False
260
261
262 class DynamicShapesType(str, enum.Enum):
@@ -824,6 +838,19 @@ class CompilationConfig:
838 # TODO(zhuhaoran): support rope native forward match and remove this.
839 # Linked issue: https://github.com/vllm-project/vllm/issues/28042
840 self.custom_ops.append("+rotary_embedding")
841 + if self.pass_config.fuse_rope_kvcache:
842 + from vllm._aiter_ops import rocm_aiter_ops
843 +
844 + if rocm_aiter_ops.is_triton_rotary_embed_enabled():
845 + logger.warning(
846 + "Cannot use VLLM_ROCM_USE_AITER_TRITON_ROPE with "
847 + "fuse_rope_kvcache. Disabling fuse_rope_kvcache."
848 + )
849 + self.pass_config.fuse_rope_kvcache = False
850 + else:
851 + # TODO(Rohan138): support rope native forward match and remove this.
852 + # Linked issue: https://github.com/vllm-project/vllm/issues/28042
853 + self.custom_ops.append("+rotary_embedding")
854
855 if (
856 is_torch_equal_or_newer("2.9.0.dev")
vllm/config/vllm.py
+14
@@ -1401,6 +1401,20 @@ class VllmConfig:
1401 "allreduce-rms fusion will be enabled for all num_tokens."
1402 )
1403
1404 + if compilation_config.pass_config.fuse_rope_kvcache:
1405 + max_token_num = (
1406 + compilation_config.pass_config.rope_kvcache_fusion_max_token_num
1407 + )
1408 + if max_token_num is not None:
1409 + if compile_range_end is not None and max_token_num < compile_range_end:
1410 + computed_compile_ranges_split_points.append(max_token_num)
1411 + else:
1412 + logger.debug(
1413 + "Max num batched tokens below rope+kvcache fusion threshold, "
1414 + "rope+kvcache fusion enabled for num_tokens <= %d.",
1415 + compile_range_end,
1416 + )
1417 +
1418 if compilation_config.compile_ranges_split_points is not None:
1419 for x in compilation_config.compile_ranges_split_points:
1420 assert isinstance(x, int)
vllm/model_executor/layers/attention/attention.py
+13 -15
@@ -570,11 +570,11 @@ direct_register_custom_op(
570
571 def get_attention_context(
572 layer_name: str,
573 -) -> tuple[Any, "Attention | MLAAttention", torch.Tensor]:
573 +) -> tuple[Any, "Attention | MLAAttention", torch.Tensor, torch.Tensor]:
574 """Extract attention context for a given layer.
575
576 This helper function extracts the attention metadata, attention layer
577 - instance, and KV cache tensor for a specific layer.
577 + instance, KV cache tensor, and slot mapping for a specific layer.
578
579 Args:
580 layer_name: The name/identifier of the attention layer.
@@ -585,6 +585,7 @@ def get_attention_context(
585 no metadata available
586 - attn_layer: The attention layer instance (Attention or MLAAttention)
587 - kv_cache: The KV cache tensor for current virtual engine
588 + - slot_mapping: The slot mapping for this specific layer
589
590 Note: attn_metadata may be None, but attn_layer and kv_cache are always
591 extracted from the forward context.
@@ -593,9 +594,14 @@ def get_attention_context(
594 attn_metadata = forward_context.attn_metadata
595 if isinstance(attn_metadata, dict):
596 attn_metadata = attn_metadata[layer_name]
596 - attn_layer = forward_context.no_compile_layers[layer_name]
597 + attn_layer: Attention | MLAAttention = forward_context.no_compile_layers[layer_name]
598 kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
598 - return attn_metadata, attn_layer, kv_cache
599 + slot_mapping = forward_context.slot_mapping
600 + assert isinstance(slot_mapping, dict), (
601 + f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
602 + )
603 + layer_slot_mapping = slot_mapping.get(layer_name)
604 + return attn_metadata, attn_layer, kv_cache, layer_slot_mapping
605
606
607 @maybe_transfer_kv_layer
@@ -605,7 +611,7 @@ def unified_attention(
611 value: torch.Tensor,
612 layer_name: str,
613 ) -> torch.Tensor:
608 - attn_metadata, self, kv_cache = get_attention_context(layer_name)
614 + attn_metadata, self, kv_cache, _ = get_attention_context(layer_name)
615 output = self.impl.forward(self, query, key, value, kv_cache, attn_metadata)
616
617 return output
@@ -636,15 +642,7 @@ def unified_kv_cache_update(
642 Returns a dummy that is passed to unified_attention to signal a side effect and
643 the data dependency between them to ensure torch.compile preserves ordering.
644 """
639 - forward_context = get_forward_context()
640 - attn_layer = forward_context.no_compile_layers[layer_name]
641 - kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
642 -
643 - slot_mapping = forward_context.slot_mapping
644 - assert isinstance(slot_mapping, dict), (
645 - f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
646 - )
647 - layer_slot_mapping = slot_mapping.get(layer_name)
645 + _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
646 if layer_slot_mapping is not None:
647 assert hasattr(attn_layer.impl, "do_kv_cache_update"), (
648 f"{attn_layer.impl.__class__.__name__} does not support kv cache update"
@@ -691,7 +689,7 @@ def unified_attention_with_output(
689 # that ensures torch.compile preserves ordering between KV cache update and
690 # attention forward.
691 del kv_cache_dummy_dep
694 - attn_metadata, self, kv_cache = get_attention_context(layer_name)
692 + attn_metadata, self, kv_cache, _ = get_attention_context(layer_name)
693
694 self.impl.forward(
695 self,
vllm/model_executor/layers/attention/kv_transfer_utils.py
+2 -2
@@ -40,8 +40,8 @@ def maybe_transfer_kv_layer(func: Callable) -> Callable:
40
41 layer_name: str = args[layer_name_index]
42
43 - # Extract attention context (layer-specific metadata, layer, and kv_cache)
44 - attn_metadata, attn_layer, kv_cache = get_attention_context(layer_name)
43 + # Extract attention context (metadata, layer, kv_cache, layer_slot_mapping)
44 + attn_metadata, _, kv_cache, _ = get_attention_context(layer_name)
45 connector = get_kv_transfer_group()
46 if attn_metadata is None or not connector.has_connector_metadata():
47 return func(*args, **kwargs)
vllm/model_executor/layers/attention/mla_attention.py
+2 -2
@@ -828,7 +828,7 @@ def unified_mla_attention(
828 k_pe: torch.Tensor,
829 layer_name: str,
830 ) -> torch.Tensor:
831 - attn_metadata, layer, kv_cache = get_attention_context(layer_name)
831 + attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
832 output = layer.forward_impl(q, kv_c_normed, k_pe, kv_cache, attn_metadata)
833
834 return output
@@ -862,7 +862,7 @@ def unified_mla_attention_with_output(
862 output_scale: torch.Tensor | None = None,
863 output_block_scale: torch.Tensor | None = None,
864 ) -> None:
865 - attn_metadata, layer, kv_cache = get_attention_context(layer_name)
865 + attn_metadata, layer, kv_cache, _ = get_attention_context(layer_name)
866 layer.forward_impl(
867 q,
868 kv_c_normed,
vllm/v1/attention/backend.py
+27
@@ -723,6 +723,33 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):
723 """
724 return False
725
726 + def fused_rope_kvcache_supported(self):
727 + """
728 + Does this attention implementation support RoPE+KVCache fusion.
729 + This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops
730 + with the KV cache update for implementations that support it.
731 + """
732 + return False
733 +
734 + def do_rope_and_kv_cache_update(
735 + self,
736 + layer: AttentionLayer,
737 + query: torch.Tensor,
738 + key: torch.Tensor,
739 + value: torch.Tensor,
740 + positions: torch.Tensor,
741 + cos_sin_cache: torch.Tensor,
742 + is_neox: bool,
743 + kv_cache: torch.Tensor,
744 + layer_slot_mapping: torch.Tensor,
745 + ):
746 + """
747 + If `fused_rope_kvcache_supported` returns True, this method will be called
748 + by torch.ops.vllm.fused_rope_and_unified_kv_cache_update
749 + to perform the inplace RoPE and KV cache update.
750 + """
751 + raise NotImplementedError
752 +
753
754 class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
755 """MLA attention implementation with forward_mqa and forward_mha methods."""
vllm/v1/attention/backends/rocm_aiter_fa.py
+37 -48
@@ -11,7 +11,6 @@ from vllm._aiter_ops import rocm_aiter_ops
11 from vllm.config import VllmConfig, get_layers_from_vllm_config
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.attention import Attention
14 -from vllm.model_executor.layers.attention.attention import get_attention_context
14 from vllm.platforms import current_platform
15 from vllm.utils.math_utils import cdiv
16 from vllm.utils.platform_utils import get_cu_count
@@ -1290,11 +1289,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
1289 kv_cache: torch.Tensor,
1290 slot_mapping: torch.Tensor,
1291 ):
1293 - attn_metadata, _, _ = get_attention_context(layer.layer_name)
1294 - if attn_metadata is None:
1295 - # Profiling run.
1296 - return
1297 -
1292 key_cache, value_cache = kv_cache.unbind(0)
1293
1294 # key and value may be None in the case of cross attention. They are
@@ -1303,45 +1297,40 @@ class AiterFlashAttentionImpl(AttentionImpl):
1297 if self.kv_cache_dtype.startswith("fp8"):
1298 key_cache = key_cache.view(current_platform.fp8_dtype())
1299 value_cache = value_cache.view(current_platform.fp8_dtype())
1306 - if (
1307 - self.kv_sharing_target_layer_name is None
1308 - and key is not None
1309 - and value is not None
1310 - ):
1311 - # Reshape the input keys and values and store them in the cache.
1312 - # Skip this if sharing KV cache with an earlier attention layer.
1313 - # NOTE(woosuk): Here, key and value are padded while slot_mapping
1314 - # is not padded. However, we don't need to do
1315 - # key[:num_actual_tokens] and value[:num_actual_tokens] because
1316 - # the reshape_and_cache_flash op uses the slot_mapping's shape
1317 - # to determine the number of actual tokens.
1318 - if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
1319 - # We may calculate per token quant scale in
1320 - # reshape_and_cache_shuffle_triton which might differ from
1321 - # vllm's style when shuffle layout is used.
1322 - k_scale = attn_metadata.k_scale
1323 - v_scale = attn_metadata.v_scale
1324 - assert k_scale is not None and v_scale is not None, (
1325 - "k_scale and v_scale are required for shuffled update"
1326 - )
1327 - reshape_and_cache_shuffle_triton(
1328 - key,
1329 - value,
1330 - key_cache,
1331 - value_cache,
1332 - slot_mapping,
1333 - self.kv_cache_dtype,
1334 - k_scale,
1335 - v_scale,
1336 - )
1337 - else:
1338 - torch.ops._C_cache_ops.reshape_and_cache_flash(
1339 - key,
1340 - value,
1341 - key_cache,
1342 - value_cache,
1343 - slot_mapping,
1344 - self.kv_cache_dtype,
1345 - layer._k_scale,
1346 - layer._v_scale,
1347 - )
1300 + # Reshape the input keys and values and store them in the cache.
1301 + # Skip this if sharing KV cache with an earlier attention layer.
1302 + # NOTE(woosuk): Here, key and value are padded while slot_mapping
1303 + # is not padded. However, we don't need to do
1304 + # key[:num_actual_tokens] and value[:num_actual_tokens] because
1305 + # the reshape_and_cache_flash op uses the slot_mapping's shape
1306 + # to determine the number of actual tokens.
1307 + if rocm_aiter_ops.is_shuffle_kv_cache_enabled():
1308 + # We may calculate per token quant scale in
1309 + # reshape_and_cache_shuffle_triton which might differ from
1310 + # vllm's style when shuffle layout is used.
1311 + k_scale = layer._k_scale
1312 + v_scale = layer._v_scale
1313 + assert k_scale is not None and v_scale is not None, (
1314 + "k_scale and v_scale are required for shuffled update"
1315 + )
1316 + reshape_and_cache_shuffle_triton(
1317 + key,
1318 + value,
1319 + key_cache,
1320 + value_cache,
1321 + slot_mapping,
1322 + self.kv_cache_dtype,
1323 + k_scale,
1324 + v_scale,
1325 + )
1326 + else:
1327 + torch.ops._C_cache_ops.reshape_and_cache_flash(
1328 + key,
1329 + value,
1330 + key_cache,
1331 + value_cache,
1332 + slot_mapping,
1333 + self.kv_cache_dtype,
1334 + layer._k_scale,
1335 + layer._v_scale,
1336 + )
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+40
@@ -5,6 +5,7 @@
5 import torch
6
7 from vllm import _custom_ops as ops
8 +from vllm._aiter_ops import rocm_aiter_ops
9 from vllm.logger import init_logger
10 from vllm.model_executor.layers.quantization.utils.quant_utils import (
11 QuantKey,
@@ -207,3 +208,42 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
208 layer._k_scale,
209 layer._v_scale,
210 )
211 +
212 + def fused_rope_kvcache_supported(self):
213 + return rocm_aiter_ops.is_enabled()
214 +
215 + def do_rope_and_kv_cache_update(
216 + self,
217 + layer: AttentionLayer,
218 + query: torch.Tensor,
219 + key: torch.Tensor,
220 + value: torch.Tensor,
221 + positions: torch.Tensor,
222 + cos_sin_cache: torch.Tensor,
223 + is_neox: bool,
224 + kv_cache: torch.Tensor,
225 + layer_slot_mapping: torch.Tensor,
226 + ):
227 + key_cache, value_cache = kv_cache.unbind(0)
228 + flash_layout = True
229 +
230 + is_fp8_kv_cache = self.kv_cache_dtype.startswith("fp8")
231 + if is_fp8_kv_cache:
232 + key_cache = key_cache.view(self.fp8_dtype)
233 + value_cache = value_cache.view(self.fp8_dtype)
234 +
235 + rocm_aiter_ops.triton_rope_and_cache(
236 + query,
237 + key,
238 + value,
239 + positions,
240 + cos_sin_cache,
241 + is_neox,
242 + key_cache,
243 + value_cache,
244 + layer_slot_mapping,
245 + layer._k_scale,
246 + layer._v_scale,
247 + flash_layout,
248 + is_fp8_kv_cache,
249 + )
vllm/v1/attention/backends/rocm_attn.py
+44
@@ -7,6 +7,7 @@ from typing import ClassVar
7
8 import torch
9
10 +from vllm._aiter_ops import rocm_aiter_ops
11 from vllm.config import VllmConfig
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -415,3 +416,46 @@ class RocmAttentionImpl(AttentionImpl):
416 layer._k_scale,
417 layer._v_scale,
418 )
419 +
420 + def fused_rope_kvcache_supported(self):
421 + return rocm_aiter_ops.is_enabled()
422 +
423 + def do_rope_and_kv_cache_update(
424 + self,
425 + layer: AttentionLayer,
426 + query: torch.Tensor,
427 + key: torch.Tensor,
428 + value: torch.Tensor,
429 + positions: torch.Tensor,
430 + cos_sin_cache: torch.Tensor,
431 + is_neox: bool,
432 + kv_cache: torch.Tensor,
433 + layer_slot_mapping: torch.Tensor,
434 + ):
435 + key_cache, value_cache = PagedAttention.split_kv_cache(
436 + kv_cache,
437 + layer.num_kv_heads, # type: ignore[attr-defined]
438 + layer.head_size, # type: ignore[attr-defined]
439 + )
440 + flash_layout = False
441 +
442 + is_fp8_kv_cache = self.kv_cache_dtype.startswith("fp8")
443 + if is_fp8_kv_cache:
444 + key_cache = key_cache.view(self.fp8_dtype)
445 + value_cache = value_cache.view(self.fp8_dtype)
446 +
447 + rocm_aiter_ops.triton_rope_and_cache(
448 + query,
449 + key,
450 + value,
451 + positions,
452 + cos_sin_cache,
453 + is_neox,
454 + key_cache,
455 + value_cache,
456 + layer_slot_mapping,
457 + layer._k_scale,
458 + layer._v_scale,
459 + flash_layout,
460 + is_fp8_kv_cache,
461 + )
vllm/v1/attention/backends/triton_attn.py
+40
@@ -7,6 +7,7 @@ from typing import ClassVar
7
8 import torch
9
10 +from vllm._aiter_ops import rocm_aiter_ops
11 from vllm.config import CUDAGraphMode, VllmConfig
12 from vllm.config.cache import CacheDType
13 from vllm.logger import init_logger
@@ -596,3 +597,42 @@ class TritonAttentionImpl(AttentionImpl):
597 layer._k_scale,
598 layer._v_scale,
599 )
600 +
601 + def fused_rope_kvcache_supported(self):
602 + return rocm_aiter_ops.is_enabled()
603 +
604 + def do_rope_and_kv_cache_update(
605 + self,
606 + layer: AttentionLayer,
607 + query: torch.Tensor,
608 + key: torch.Tensor,
609 + value: torch.Tensor,
610 + positions: torch.Tensor,
611 + cos_sin_cache: torch.Tensor,
612 + is_neox: bool,
613 + kv_cache: torch.Tensor,
614 + layer_slot_mapping: torch.Tensor,
615 + ):
616 + key_cache, value_cache = kv_cache.unbind(1)
617 + flash_layout = True
618 +
619 + is_fp8_kv_cache = self.kv_cache_dtype.startswith("fp8")
620 + if is_fp8_kv_cache:
621 + key_cache = key_cache.view(self.fp8_dtype)
622 + value_cache = value_cache.view(self.fp8_dtype)
623 +
624 + rocm_aiter_ops.triton_rope_and_cache(
625 + query,
626 + key,
627 + value,
628 + positions,
629 + cos_sin_cache,
630 + is_neox,
631 + key_cache,
632 + value_cache,
633 + layer_slot_mapping,
634 + layer._k_scale,
635 + layer._v_scale,
636 + flash_layout,
637 + is_fp8_kv_cache,
638 + )