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 1, 20:00 UTC
[ROCm][GPT-OSS] Fuse RoPE + static Q FP8 quant on fused RoPE+KV path (#42832)
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
akii96 committed
Jun 6, 2026 at 00:22 UTC
4200f621472925f8039191b9ff3b72951134cdf5
2 files changed
+451
-1
tests/compile/passes/test_rope_kvcache_fusion.py
+257
@@ -3,11 +3,13 @@
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 BatchSpec, create_common_attn_metadata
11
from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
12
+from vllm.compilation.passes.fusion import rope_kvcache_fusion
13
from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
14
from vllm.compilation.passes.fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
15
from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
@@ -24,6 +26,7 @@ from vllm.config import (
26
PassConfig,
27
VllmConfig,
28
)
29
+from vllm.config.utils import Range
30
from vllm.forward_context import get_forward_context, set_forward_context
31
from vllm.model_executor.layers.attention import Attention
32
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
@@ -40,6 +43,20 @@ VLLM_UNIFIED_KV_CACHE_UPDATE_OP = torch.ops.vllm.unified_kv_cache_update
43
FP8_DTYPE = current_platform.fp8_dtype()
44
45
46
+def test_rope_kvcache_fusion_default_keeps_large_ranges_unfused():
47
+ vllm_config = VllmConfig(
48
+ compilation_config=CompilationConfig(
49
+ mode=CompilationMode.VLLM_COMPILE,
50
+ pass_config=PassConfig(fuse_rope_kvcache=True),
51
+ ),
52
+ )
53
+ fusion_pass = RopeKVCacheFusionPass(vllm_config)
54
+
55
+ assert fusion_pass.is_applicable_for_range(Range(1, 256))
56
+ assert not fusion_pass.is_applicable_for_range(Range(257, 11650))
57
+ assert not fusion_pass.is_applicable_for_range(Range(11651, 16384))
58
+
59
+
60
class QKRoPEKVCacheTestModel(torch.nn.Module):
61
def __init__(
62
self,
@@ -184,6 +201,51 @@ class QKRoPEKVCacheTestModel(torch.nn.Module):
201
return [torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default]
202
203
204
+class QKRoPEStaticQKVCacheTestModel(QKRoPEKVCacheTestModel):
205
+ def __init__(self, *args, **kwargs):
206
+ super().__init__(*args, **kwargs)
207
+ self.q_scale = torch.ones((), dtype=torch.float32, device=self.device)
208
+
209
+ def forward(
210
+ self, qkv: torch.Tensor, positions: torch.Tensor
211
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
212
+ # Create copy so inplace ops do not modify the original tensors
213
+ qkv = qkv.clone()
214
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
215
+ q, k = self.rotary_emb(positions, q, k)
216
+
217
+ q_fp8 = torch.empty(q.shape, device=q.device, dtype=FP8_DTYPE)
218
+ _, q_fp8 = auto_functionalized(
219
+ torch.ops._C.static_scaled_fp8_quant.default,
220
+ result=q_fp8,
221
+ input=q,
222
+ scale=self.q_scale,
223
+ group_shape=(-1, -1),
224
+ )
225
+ q = q_fp8.view(-1, self.num_heads, self.head_size)
226
+ k = k.view(-1, self.num_kv_heads, self.head_size)
227
+ v = v.view(-1, self.num_kv_heads, self.head_size)
228
+ kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
229
+ k, v, _encode_layer_name(self.layer_name)
230
+ )
231
+ return q, k, v, kv_cache_dummy_dep
232
+
233
+ def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
234
+ ops = []
235
+ if self.enable_rope_custom_op:
236
+ if rocm_aiter_ops.is_triton_rotary_embed_enabled():
237
+ ops.append(torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default)
238
+ else:
239
+ ops.append(ROTARY_OP)
240
+ else:
241
+ ops.append(INDEX_SELECT_OP)
242
+ ops.append(torch.ops.vllm.unified_kv_cache_update.default)
243
+ return ops
244
+
245
+ def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
246
+ return [torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default]
247
+
248
+
249
@pytest.mark.parametrize(
250
"attn_backend",
251
[
@@ -326,3 +388,198 @@ def test_rope_kvcache_fusion(
388
atol=ATOL,
389
rtol=RTOL,
390
)
391
+
392
+
393
+@pytest.mark.parametrize(
394
+ "attn_backend",
395
+ [AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN],
396
+)
397
+@pytest.mark.parametrize("enable_rope_custom_op", [True])
398
+@pytest.mark.parametrize("enable_aiter_triton_rope", [True, False])
399
+@pytest.mark.parametrize("num_heads", [64])
400
+@pytest.mark.parametrize("num_kv_heads", [8])
401
+@pytest.mark.parametrize("head_size", [64])
402
+@pytest.mark.parametrize("block_size", [16])
403
+@pytest.mark.parametrize("is_neox", [True, False])
404
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
405
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
406
+@pytest.mark.skipif(
407
+ not is_aiter_found_and_supported(),
408
+ reason="Only test on ROCm with AITER installed and supported",
409
+)
410
+@pytest.mark.skipif(
411
+ not hasattr(torch.ops._C, "static_scaled_fp8_quant"),
412
+ reason="static fp8 quant op not available on this build",
413
+)
414
+def test_rope_static_qquant_kvcache_fusion(
415
+ attn_backend: AttentionBackendEnum,
416
+ enable_rope_custom_op: bool,
417
+ enable_aiter_triton_rope: bool,
418
+ num_heads: int,
419
+ num_kv_heads: int,
420
+ head_size: int,
421
+ block_size: int,
422
+ is_neox: bool,
423
+ dtype: torch.dtype,
424
+ kv_cache_dtype: str,
425
+ monkeypatch: pytest.MonkeyPatch,
426
+):
427
+ torch.set_default_device("cuda")
428
+ torch.set_default_dtype(dtype)
429
+ torch.manual_seed(0)
430
+
431
+ custom_ops: list[str] = []
432
+ if enable_rope_custom_op:
433
+ custom_ops.append("+rotary_embedding")
434
+
435
+ vllm_config = VllmConfig(
436
+ model_config=ModelConfig(dtype=dtype),
437
+ cache_config=CacheConfig(
438
+ block_size=block_size,
439
+ cache_dtype=kv_cache_dtype,
440
+ ),
441
+ compilation_config=CompilationConfig(
442
+ mode=CompilationMode.VLLM_COMPILE,
443
+ custom_ops=custom_ops,
444
+ pass_config=PassConfig(
445
+ fuse_rope_kvcache=True,
446
+ eliminate_noops=True,
447
+ ),
448
+ ),
449
+ )
450
+
451
+ with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
452
+ m.setenv("VLLM_ROCM_USE_AITER", "1")
453
+ m.setenv(
454
+ "VLLM_ROCM_USE_AITER_TRITON_ROPE", "1" if enable_aiter_triton_rope else "0"
455
+ )
456
+ rocm_aiter_ops.refresh_env_variables()
457
+
458
+ model = QKRoPEStaticQKVCacheTestModel(
459
+ vllm_config=vllm_config,
460
+ attn_backend=attn_backend,
461
+ num_heads=num_heads,
462
+ num_kv_heads=num_kv_heads,
463
+ head_size=head_size,
464
+ is_neox=is_neox,
465
+ dtype=dtype,
466
+ device=torch.get_default_device(),
467
+ )
468
+
469
+ fusion_pass = RopeKVCacheFusionPass(vllm_config)
470
+ passes = [
471
+ NoOpEliminationPass(vllm_config),
472
+ SplitCoalescingPass(vllm_config),
473
+ ScatterSplitReplacementPass(vllm_config),
474
+ fusion_pass,
475
+ PostCleanupPass(vllm_config),
476
+ ]
477
+ backend = TestBackend(*passes)
478
+
479
+ T = 5
480
+ qkv = torch.randn(
481
+ T, num_heads * head_size + 2 * num_kv_heads * head_size, dtype=dtype
482
+ )
483
+ pos = torch.arange(T, dtype=torch.long)
484
+
485
+ qkv_unfused = qkv.clone()
486
+ pos_unfused = pos.clone()
487
+
488
+ with set_forward_context(None, vllm_config):
489
+ forward_context = get_forward_context()
490
+ attn_metadata = model.build_attn_metadata(T)
491
+ forward_context.slot_mapping = {
492
+ model.layer_name: attn_metadata.slot_mapping
493
+ }
494
+ q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)
495
+ attn_layer = forward_context.no_compile_layers[model.layer_name]
496
+ kv_cache_unfused = attn_layer.kv_cache
497
+ del dummy
498
+
499
+ torch._dynamo.mark_dynamic(qkv, 0)
500
+ torch._dynamo.mark_dynamic(pos, 0)
501
+ with set_forward_context(None, vllm_config):
502
+ model_fused = torch.compile(model, backend=backend)
503
+ forward_context = get_forward_context()
504
+ attn_metadata = model_fused.build_attn_metadata(T)
505
+ forward_context.slot_mapping = {
506
+ model.layer_name: attn_metadata.slot_mapping
507
+ }
508
+ q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)
509
+ attn_layer = forward_context.no_compile_layers[model.layer_name]
510
+ kv_cache_fused = attn_layer.kv_cache
511
+ del dummy
512
+
513
+ assert fusion_pass.matched_count == 1
514
+ backend.check_before_ops(model.ops_in_model_before())
515
+ backend.check_after_ops(model.ops_in_model_after())
516
+ static_quant_pre = backend.op_count(
517
+ torch.ops._C.static_scaled_fp8_quant.default, before=True
518
+ )
519
+ static_quant_post = backend.op_count(
520
+ torch.ops._C.static_scaled_fp8_quant.default
521
+ )
522
+ assert static_quant_pre > 0
523
+ # The replacement still emits static quant, so count is expected to
524
+ # remain non-zero after fusion.
525
+ assert static_quant_post > 0
526
+
527
+ # Negative control: without the static-Q pattern, the generic RoPE+KV
528
+ # pattern cannot match this rope -> static-quant -> kv graph, so the
529
+ # fusion above is attributable solely to RopeStaticQQuantKVCachePattern.
530
+ # This is a structural property independent of the rope/dtype/neox axes,
531
+ # so run the (extra compile) check only once on a representative combo.
532
+ if is_neox and enable_aiter_triton_rope and kv_cache_dtype == "auto":
533
+ m.setattr(
534
+ rope_kvcache_fusion,
535
+ "_supports_static_q_fp8_quant_fusion",
536
+ lambda: False,
537
+ )
538
+ generic_pass = RopeKVCacheFusionPass(vllm_config)
539
+ generic_backend = TestBackend(
540
+ NoOpEliminationPass(vllm_config),
541
+ SplitCoalescingPass(vllm_config),
542
+ ScatterSplitReplacementPass(vllm_config),
543
+ generic_pass,
544
+ PostCleanupPass(vllm_config),
545
+ )
546
+ # Reset dynamo so the model is recompiled through generic_backend
547
+ # instead of reusing the cached compilation from above.
548
+ torch._dynamo.reset()
549
+ with set_forward_context(None, vllm_config):
550
+ model_generic = torch.compile(model, backend=generic_backend)
551
+ forward_context = get_forward_context()
552
+ attn_metadata = model_generic.build_attn_metadata(T)
553
+ forward_context.slot_mapping = {
554
+ model.layer_name: attn_metadata.slot_mapping
555
+ }
556
+ model_generic(qkv, pos)
557
+ # op_count reads the post-pass graph, so it also confirms the pass ran
558
+ # (a no-op pass would raise instead of silently passing on count 0).
559
+ assert generic_pass.matched_count == 0
560
+ assert (
561
+ generic_backend.op_count(
562
+ torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default
563
+ )
564
+ == 0
565
+ )
566
+
567
+ if dtype == torch.float16:
568
+ ATOL, RTOL = (2e-3, 2e-3)
569
+ else:
570
+ ATOL, RTOL = (1e-2, 1e-2)
571
+
572
+ torch.testing.assert_close(
573
+ q_unfused.to(torch.float32),
574
+ q_fused.to(torch.float32),
575
+ atol=ATOL,
576
+ rtol=RTOL,
577
+ )
578
+ torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
579
+ torch.testing.assert_close(v_unfused, v_fused, atol=ATOL, rtol=RTOL)
580
+ torch.testing.assert_close(
581
+ kv_cache_unfused.view(dtype),
582
+ kv_cache_fused.view(dtype),
583
+ atol=ATOL,
584
+ rtol=RTOL,
585
+ )
vllm/compilation/passes/fusion/rope_kvcache_fusion.py
+194
-1
@@ -15,6 +15,7 @@ from vllm.model_executor.layers.attention.attention import (
15
Attention,
16
get_attention_context,
17
)
18
+from vllm.platforms import current_platform
19
from vllm.utils.torch_utils import (
20
_USE_LAYERNAME,
21
LayerNameType,
@@ -34,6 +35,12 @@ from .rms_quant_fusion import (
35
)
36
37
logger = init_logger(__name__)
38
+FP8_DTYPE = current_platform.fp8_dtype()
39
+STATIC_FP8_QUANT_OP = getattr(torch.ops._C, "static_scaled_fp8_quant", None)
40
+
41
+
42
+def _supports_static_q_fp8_quant_fusion() -> bool:
43
+ return STATIC_FP8_QUANT_OP is not None and FP8_DTYPE is not None
44
45
46
def fused_rope_and_unified_kv_cache_update_impl(
@@ -90,6 +97,185 @@ direct_register_custom_op(
97
)
98
99
100
+class RopeStaticQQuantKVCachePattern:
101
+ """
102
+ Fuse rope + static Q fp8 quant while preserving explicit KV-cache update
103
+ dependency ordering.
104
+ """
105
+
106
+ FUSED_ROPE_KV_OP = torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default
107
+
108
+ def __init__(
109
+ self,
110
+ layer: Attention,
111
+ is_neox: bool,
112
+ ) -> None:
113
+ self.layer_name = layer.layer_name
114
+ self.num_heads = layer.num_heads
115
+ self.num_kv_heads = layer.num_kv_heads
116
+ self.head_size = layer.head_size
117
+ self.head_size_v = layer.head_size_v
118
+ self.is_neox = is_neox
119
+
120
+ self.q_size = self.num_heads * self.head_size
121
+ self.k_size = self.num_kv_heads * self.head_size
122
+ self.v_size = self.num_kv_heads * self.head_size_v
123
+
124
+ self.rope_matcher = MatcherRotaryEmbedding(
125
+ is_neox=self.is_neox,
126
+ head_size=self.head_size,
127
+ num_heads=self.num_heads,
128
+ num_kv_heads=self.num_kv_heads,
129
+ )
130
+
131
+ def get_inputs(self) -> list:
132
+ T = 5
133
+ L = 4096
134
+ qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size)
135
+ positions = empty_i64(T)
136
+ cos_sin_cache = empty_bf16(L, self.head_size)
137
+ q_scale = torch.empty((), dtype=torch.float32, device=qkv.device)
138
+ inputs: list = [qkv, positions, cos_sin_cache, q_scale]
139
+ if _USE_LAYERNAME:
140
+ inputs.append(_encode_layer_name(self.layer_name))
141
+ return inputs
142
+
143
+ def _mk_pattern_with_layer_name_input(self, _ln):
144
+ def pattern(qkv, positions, cos_sin_cache, q_scale, layer_name):
145
+ q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
146
+ q, k = self.rope_matcher(positions, q, k, cos_sin_cache)
147
+ q_fp8 = torch.empty(q.shape, device=q.device, dtype=FP8_DTYPE)
148
+ _, q_fp8 = auto_functionalized(
149
+ torch.ops._C.static_scaled_fp8_quant.default,
150
+ result=q_fp8,
151
+ input=q,
152
+ scale=q_scale,
153
+ group_shape=(-1, -1),
154
+ )
155
+ q_view = q_fp8.view(-1, self.num_heads, self.head_size)
156
+ k_view = k.view(-1, self.num_kv_heads, self.head_size)
157
+ v_view = v.view(-1, self.num_kv_heads, self.head_size_v)
158
+ kv_cache_dummy = torch.ops.vllm.unified_kv_cache_update(
159
+ k_view, v_view, layer_name
160
+ )
161
+ return kv_cache_dummy, q_view, k_view, v_view
162
+
163
+ def replacement(qkv, positions, cos_sin_cache, q_scale, layer_name):
164
+ q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
165
+ q_view = q.view(-1, self.num_heads, self.head_size)
166
+ k_view = k.view(-1, self.num_kv_heads, self.head_size)
167
+ v_view = v.view(-1, self.num_kv_heads, self.head_size_v)
168
+ rope_kv_results = auto_functionalized(
169
+ self.FUSED_ROPE_KV_OP,
170
+ query=q_view,
171
+ key=k_view,
172
+ value=v_view,
173
+ positions=positions,
174
+ cos_sin_cache=cos_sin_cache,
175
+ is_neox=self.is_neox,
176
+ layer_name=layer_name,
177
+ )
178
+ kv_cache_dummy = rope_kv_results[0]
179
+ q_after = rope_kv_results[1]
180
+ k_after = rope_kv_results[2]
181
+ q_after_flat = q_after.view(-1, self.q_size)
182
+ q_fp8 = torch.empty(
183
+ q_after_flat.shape, device=q_after_flat.device, dtype=FP8_DTYPE
184
+ )
185
+ _, q_fp8 = auto_functionalized(
186
+ torch.ops._C.static_scaled_fp8_quant.default,
187
+ result=q_fp8,
188
+ input=q_after_flat,
189
+ scale=q_scale,
190
+ group_shape=(-1, -1),
191
+ )
192
+ return (
193
+ kv_cache_dummy,
194
+ q_fp8.view(-1, self.num_heads, self.head_size),
195
+ k_after,
196
+ v_view,
197
+ )
198
+
199
+ return pattern, replacement
200
+
201
+ def _mk_pattern_with_layer_name_closure(self, _ln):
202
+ def pattern(qkv, positions, cos_sin_cache, q_scale):
203
+ q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
204
+ q, k = self.rope_matcher(positions, q, k, cos_sin_cache)
205
+ q_fp8 = torch.empty(q.shape, device=q.device, dtype=FP8_DTYPE)
206
+ _, q_fp8 = auto_functionalized(
207
+ torch.ops._C.static_scaled_fp8_quant.default,
208
+ result=q_fp8,
209
+ input=q,
210
+ scale=q_scale,
211
+ group_shape=(-1, -1),
212
+ )
213
+ q_view = q_fp8.view(-1, self.num_heads, self.head_size)
214
+ k_view = k.view(-1, self.num_kv_heads, self.head_size)
215
+ v_view = v.view(-1, self.num_kv_heads, self.head_size_v)
216
+ kv_cache_dummy = torch.ops.vllm.unified_kv_cache_update(k_view, v_view, _ln)
217
+ return kv_cache_dummy, q_view, k_view, v_view
218
+
219
+ def replacement(qkv, positions, cos_sin_cache, q_scale):
220
+ q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
221
+ q_view = q.view(-1, self.num_heads, self.head_size)
222
+ k_view = k.view(-1, self.num_kv_heads, self.head_size)
223
+ v_view = v.view(-1, self.num_kv_heads, self.head_size_v)
224
+ rope_kv_results = auto_functionalized(
225
+ self.FUSED_ROPE_KV_OP,
226
+ query=q_view,
227
+ key=k_view,
228
+ value=v_view,
229
+ positions=positions,
230
+ cos_sin_cache=cos_sin_cache,
231
+ is_neox=self.is_neox,
232
+ layer_name=_ln,
233
+ )
234
+ kv_cache_dummy = rope_kv_results[0]
235
+ q_after = rope_kv_results[1]
236
+ k_after = rope_kv_results[2]
237
+ q_after_flat = q_after.view(-1, self.q_size)
238
+ q_fp8 = torch.empty(
239
+ q_after_flat.shape, device=q_after_flat.device, dtype=FP8_DTYPE
240
+ )
241
+ _, q_fp8 = auto_functionalized(
242
+ torch.ops._C.static_scaled_fp8_quant.default,
243
+ result=q_fp8,
244
+ input=q_after_flat,
245
+ scale=q_scale,
246
+ group_shape=(-1, -1),
247
+ )
248
+ return (
249
+ kv_cache_dummy,
250
+ q_fp8.view(-1, self.num_heads, self.head_size),
251
+ k_after,
252
+ v_view,
253
+ )
254
+
255
+ return pattern, replacement
256
+
257
+ def register(self, pm_pass: PatternMatcherPass) -> None:
258
+ _ln = _encode_layer_name(self.layer_name)
259
+
260
+ if _USE_LAYERNAME:
261
+ pattern, replacement = self._mk_pattern_with_layer_name_input(_ln)
262
+ else:
263
+ pattern, replacement = self._mk_pattern_with_layer_name_closure(_ln)
264
+
265
+ def fwd_and_view_to_reshape(*args, **kwargs) -> fx.GraphModule:
266
+ gm = pm.fwd_only(*args, **kwargs)
267
+ view_to_reshape(gm)
268
+ return gm
269
+
270
+ pm.register_replacement(
271
+ pattern,
272
+ replacement,
273
+ self.get_inputs(),
274
+ fwd_and_view_to_reshape,
275
+ pm_pass,
276
+ )
277
+
278
+
279
class RopeReshapeKVCachePattern:
280
"""
281
This pattern matches the following unfused inplace ops:
@@ -253,6 +439,11 @@ class RopeKVCacheFusionPass(VllmPatternMatcherPass):
439
for _, layer in attn_layers.items():
440
if layer.impl.fused_rope_kvcache_supported():
441
for is_neox in [True, False]:
442
+ if _supports_static_q_fp8_quant_fusion():
443
+ RopeStaticQQuantKVCachePattern(
444
+ layer=layer,
445
+ is_neox=is_neox,
446
+ ).register(self.patterns)
447
RopeReshapeKVCachePattern(
448
layer=layer,
449
is_neox=is_neox,
@@ -274,4 +465,6 @@ class RopeKVCacheFusionPass(VllmPatternMatcherPass):
465
return compile_range.end <= self.max_token_num
466
467
def uuid(self) -> str:
277
- return VllmInductorPass.hash_source(self, RopeReshapeKVCachePattern)
468
+ return VllmInductorPass.hash_source(
469
+ self, RopeStaticQQuantKVCachePattern, RopeReshapeKVCachePattern
470
+ )