@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 2cab53ddee

[Model][Hardware][AMD]: Part 1/2 -> Enable e2e QK Norm + RoPE + KV Cache runtime fusion for Qwen3-30B-A3B on ROCM_AITER_FA, and ROCM_AITER_UNIFIED_ATTN (#42749)

Signed-off-by: Jack Hu <Jack.Hu@amd.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>

Jack Hu committed Jul 16, 2026 at 18:39 UTC 2cab53ddeea3851a083beff8c192f068c4518b67
15 files changed +1361 -4
tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py new
+478
@@ -0,0 +1,478 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import os
5 +
6 +import pytest
7 +import torch
8 +
9 +import vllm.config
10 +from tests.compile.backend import TestBackend
11 +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
12 +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
13 +from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP
14 +from vllm.compilation.passes.fusion.qk_norm_rope_kvcache_fusion import (
15 + QkNormRopeKvCacheFusionPass,
16 +)
17 +from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
18 +from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
19 +from vllm.compilation.passes.utility.scatter_split_replace import (
20 + ScatterSplitReplacementPass,
21 +)
22 +from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
23 +from vllm.config import (
24 + CacheConfig,
25 + CompilationConfig,
26 + CompilationMode,
27 + ModelConfig,
28 + PassConfig,
29 + VllmConfig,
30 +)
31 +from vllm.forward_context import get_forward_context, set_forward_context
32 +from vllm.model_executor.layers.attention import Attention
33 +from vllm.model_executor.layers.layernorm import RMSNorm
34 +from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
35 +from vllm.platforms import current_platform
36 +from vllm.v1.attention.backend import (
37 + AttentionBackend,
38 + CommonAttentionMetadata,
39 +)
40 +from vllm.v1.attention.backends.registry import AttentionBackendEnum
41 +from vllm.v1.kv_cache_interface import AttentionSpec
42 +
43 +INDEX_SELECT_OP = torch.ops.aten.index.Tensor
44 +FP8_DTYPE = current_platform.fp8_dtype()
45 +
46 +
47 +class QKNormRoPEKVCacheTestModel(torch.nn.Module):
48 + """Minimal model that reproduces the QK-norm + RoPE + KV cache update
49 + pattern matched by QkNormRopeKvCacheFusionPass:
50 +
51 + q, k, v = split(qkv)
52 + q = rms_norm(q.view(heads, dim), q_weight).view(flat)
53 + k = rms_norm(k.view(heads, dim), k_weight).view(flat)
54 + q, k = rotary_emb(positions, q, k)
55 + q = q.view(num_heads, head_dim)
56 + k = k.view(num_kv_heads, head_dim)
57 + v = v.view(num_kv_heads, head_dim)
58 + dummy = unified_kv_cache_update(k, v, layer_name)
59 + """
60 +
61 + def __init__(
62 + self,
63 + vllm_config: VllmConfig,
64 + attn_backend: AttentionBackendEnum,
65 + num_heads: int,
66 + num_kv_heads: int,
67 + head_size: int,
68 + is_neox: bool,
69 + rms_norm_eps: float,
70 + dtype: torch.dtype,
71 + device: torch.device,
72 + rotary_dim: int | None = None,
73 + prefix: str = "model.layers.0.self_attn.attn",
74 + ):
75 + super().__init__()
76 + self.num_heads = num_heads
77 + self.num_kv_heads = num_kv_heads
78 + self.head_size = head_size
79 + self.rotary_dim = rotary_dim if rotary_dim is not None else head_size
80 + self.block_size = vllm_config.cache_config.block_size
81 + self.q_size = num_heads * head_size
82 + self.kv_size = num_kv_heads * head_size
83 + self.is_neox = is_neox
84 + self.dtype = dtype
85 + self.device = device
86 + self.layer_name = prefix
87 +
88 + self.q_norm = RMSNorm(head_size, eps=rms_norm_eps)
89 + self.k_norm = RMSNorm(head_size, eps=rms_norm_eps)
90 +
91 + self.rotary_emb = RotaryEmbedding(
92 + head_size,
93 + rotary_dim=self.rotary_dim,
94 + max_position_embeddings=4096,
95 + base=10000,
96 + is_neox_style=is_neox,
97 + dtype=self.dtype,
98 + )
99 +
100 + self.enable_rope_custom_op = self.rotary_emb.enabled()
101 +
102 + self.attn = Attention(
103 + num_heads=num_heads,
104 + head_size=head_size,
105 + scale=1.0 / head_size**0.5,
106 + num_kv_heads=num_kv_heads,
107 + cache_config=vllm_config.cache_config,
108 + quant_config=vllm_config.quant_config,
109 + prefix=prefix,
110 + attn_backend=attn_backend.get_class(),
111 + )
112 + self.attn_backend: type[AttentionBackend] = self.attn.get_attn_backend()
113 + assert not self.attn_backend.forward_includes_kv_cache_update, (
114 + f"Attention backend {self.attn_backend} does not support "
115 + "fuse_qk_norm_rope_kvcache."
116 + )
117 + kv_cache_dtype_str = vllm_config.cache_config.cache_dtype
118 + self.kv_cache_dtype = (
119 + FP8_DTYPE if kv_cache_dtype_str.startswith("fp8") else self.dtype
120 + )
121 +
122 + if self.kv_cache_dtype != self.dtype:
123 + self.attn._k_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
124 + self.attn._v_scale = torch.tensor(1.0, dtype=torch.float32, device=device)
125 + self.attn._k_scale_float = 1.0
126 + self.attn._v_scale_float = 1.0
127 + else:
128 + self.attn._k_scale = self.attn._k_scale.to(device)
129 + self.attn._v_scale = self.attn._v_scale.to(device)
130 +
131 + self.builder = self.attn.attn_backend.get_builder_cls()(
132 + kv_cache_spec=AttentionSpec(
133 + block_size=self.block_size,
134 + num_kv_heads=self.num_kv_heads,
135 + head_size=head_size,
136 + dtype=self.kv_cache_dtype,
137 + ),
138 + layer_names=[self.attn.layer_name],
139 + vllm_config=vllm_config,
140 + device=device,
141 + )
142 +
143 + def build_attn_metadata(
144 + self, batch_size: int, kv_stride_order: tuple[int, ...] | None = None
145 + ) -> CommonAttentionMetadata:
146 + batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)
147 + common_attn_metadata = create_common_attn_metadata(
148 + batch_spec, self.block_size, self.device, arange_block_indices=True
149 + )
150 +
151 + max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
152 + num_blocks = batch_size * max_blocks
153 +
154 + attn_backend = self.attn.attn_backend
155 + kv_cache_shape = attn_backend.get_kv_cache_shape(
156 + num_blocks, self.block_size, self.num_kv_heads, self.head_size
157 + )
158 + # Caller can force a physical layout; else use the backend's.
159 + if kv_stride_order is None:
160 + try:
161 + kv_stride_order = attn_backend.get_kv_cache_stride_order()
162 + except (AttributeError, NotImplementedError):
163 + kv_stride_order = tuple(range(len(kv_cache_shape)))
164 +
165 + kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_stride_order)
166 + inv_order = [kv_stride_order.index(i) for i in range(len(kv_stride_order))]
167 +
168 + raw_tensor = torch.zeros(
169 + 2 * num_blocks * self.block_size * self.num_kv_heads * self.head_size,
170 + dtype=self.kv_cache_dtype,
171 + device=self.device,
172 + )
173 + raw_tensor = raw_tensor.view(kv_cache_shape)
174 + kv_cache = raw_tensor.permute(*inv_order)
175 +
176 + # Store as a bare tensor (not wrapped in a list) to match production
177 + # `bind_kv_cache` behavior. `get_attention_context` returns this
178 + # attribute directly to the fused/unfused `do_kv_cache_update` impls,
179 + # which call `kv_cache.unbind(0)` and therefore require a tensor.
180 + self.attn.kv_cache = kv_cache
181 +
182 + attn_metadata = self.builder.build(
183 + common_prefix_len=0, common_attn_metadata=common_attn_metadata
184 + )
185 +
186 + return attn_metadata
187 +
188 + def forward(
189 + self, qkv: torch.Tensor, positions: torch.Tensor
190 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
191 + qkv = qkv.clone()
192 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
193 +
194 + # QK-norm: RMSNorm on per-head Q and K
195 + q = q.view(-1, self.num_heads, self.head_size)
196 + q = self.q_norm(q)
197 + q = q.view(-1, self.q_size)
198 +
199 + k = k.view(-1, self.num_kv_heads, self.head_size)
200 + k = self.k_norm(k)
201 + k = k.view(-1, self.kv_size)
202 +
203 + # RoPE
204 + q, k = self.rotary_emb(positions, q, k)
205 +
206 + # Mirror Attention.forward: quant-query impls consume an fp8 q.
207 + if (
208 + self.kv_cache_dtype != self.dtype
209 + and self.attn.impl.supports_quant_query_input
210 + ):
211 + q_fp8 = torch.empty_like(q, dtype=FP8_DTYPE)
212 + torch.ops.vllm.rocm_aiter_per_tensor_quant(
213 + q_fp8, q, self.attn._q_scale, False
214 + )
215 + q = q_fp8
216 +
217 + # Final views + KV cache update
218 + q = q.view(-1, self.num_heads, self.head_size)
219 + k = k.view(-1, self.num_kv_heads, self.head_size)
220 + v = v.view(-1, self.num_kv_heads, self.head_size)
221 + kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
222 + k, v, self.layer_name
223 + )
224 + return q, k, v, kv_cache_dummy_dep
225 +
226 + def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
227 + ops: list[torch._ops.OpOverload] = []
228 + # RoPE is not yet IR-migrated, so its custom op still surfaces
229 + # directly in the graph based on `enable_rope_custom_op`.
230 + if self.enable_rope_custom_op:
231 + if rocm_aiter_ops.is_triton_rotary_embed_enabled():
232 + ops.append(torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default)
233 + else:
234 + ops.append(ROTARY_OP)
235 + else:
236 + ops.append(INDEX_SELECT_OP)
237 + ops.append(torch.ops.vllm.unified_kv_cache_update.default)
238 + return ops
239 +
240 + def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
241 + return [torch.ops.vllm.fused_qk_norm_rope_and_unified_kv_cache_update.default]
242 +
243 +
244 +def _run_qk_norm_rope_kvcache_fusion_test(
245 + *,
246 + attn_backend: AttentionBackendEnum,
247 + enable_aiter_triton_rope: bool,
248 + num_tokens: int,
249 + num_heads: int,
250 + num_kv_heads: int,
251 + head_size: int,
252 + rotary_dim: int,
253 + block_size: int,
254 + is_neox: bool,
255 + use_shuffle_kv_layout: str,
256 + kv_stride_order: tuple[int, ...],
257 + dtype: torch.dtype,
258 + kv_cache_dtype: str,
259 + rms_norm_eps: float,
260 + custom_op: str,
261 + monkeypatch: pytest.MonkeyPatch,
262 +) -> None:
263 + device = os.environ.get("VLLM_TEST_CUDA_DEVICE", "cuda")
264 + torch.set_default_device(device)
265 + torch.set_default_dtype(dtype)
266 + torch.manual_seed(0)
267 +
268 + vllm_config = VllmConfig(
269 + model_config=ModelConfig(dtype=dtype),
270 + cache_config=CacheConfig(
271 + block_size=block_size,
272 + cache_dtype=kv_cache_dtype,
273 + ),
274 + compilation_config=CompilationConfig(
275 + mode=CompilationMode.VLLM_COMPILE,
276 + custom_ops=[custom_op],
277 + pass_config=PassConfig(
278 + fuse_qk_norm_rope_kvcache=True,
279 + eliminate_noops=True,
280 + ),
281 + ),
282 + )
283 +
284 + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
285 + m.setenv("VLLM_ROCM_USE_AITER", "1")
286 + m.setenv(
287 + "VLLM_ROCM_USE_AITER_TRITON_ROPE",
288 + "1" if enable_aiter_triton_rope else "0",
289 + )
290 + m.setenv("VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", use_shuffle_kv_layout)
291 + rocm_aiter_ops.refresh_env_variables()
292 +
293 + model = QKNormRoPEKVCacheTestModel(
294 + vllm_config=vllm_config,
295 + attn_backend=attn_backend,
296 + num_heads=num_heads,
297 + num_kv_heads=num_kv_heads,
298 + head_size=head_size,
299 + rotary_dim=rotary_dim,
300 + is_neox=is_neox,
301 + rms_norm_eps=rms_norm_eps,
302 + dtype=dtype,
303 + device=torch.get_default_device(),
304 + )
305 +
306 + fusion_pass = QkNormRopeKvCacheFusionPass(vllm_config)
307 + passes = [
308 + NoOpEliminationPass(vllm_config),
309 + SplitCoalescingPass(vllm_config),
310 + ScatterSplitReplacementPass(vllm_config),
311 + fusion_pass,
312 + PostCleanupPass(vllm_config),
313 + ]
314 + backend = TestBackend(*passes)
315 +
316 + qkv = torch.randn(
317 + num_tokens,
318 + num_heads * head_size + 2 * num_kv_heads * head_size,
319 + dtype=dtype,
320 + )
321 + pos = torch.arange(num_tokens, dtype=torch.long)
322 +
323 + qkv_unfused = qkv.clone()
324 + pos_unfused = pos.clone()
325 +
326 + # Run unfused (eager) forward
327 + with set_forward_context(None, vllm_config):
328 + forward_context = get_forward_context()
329 + attn_metadata = model.build_attn_metadata(num_tokens, kv_stride_order)
330 + forward_context.slot_mapping = {
331 + model.layer_name: attn_metadata.slot_mapping
332 + }
333 + q_unfused, k_unfused, v_unfused, dummy = model(qkv_unfused, pos_unfused)
334 + attn_layer = forward_context.no_compile_layers[model.layer_name]
335 + kv_cache_unfused = attn_layer.kv_cache
336 + del dummy
337 +
338 + # Run fused (compiled) forward
339 + torch._dynamo.mark_dynamic(qkv, 0)
340 + torch._dynamo.mark_dynamic(pos, 0)
341 + with set_forward_context(None, vllm_config):
342 + model_fused = torch.compile(model, backend=backend)
343 + forward_context = get_forward_context()
344 + attn_metadata = model_fused.build_attn_metadata(num_tokens, kv_stride_order)
345 + forward_context.slot_mapping = {
346 + model.layer_name: attn_metadata.slot_mapping
347 + }
348 + q_fused, k_fused, v_fused, dummy = model_fused(qkv, pos)
349 + attn_layer = forward_context.no_compile_layers[model.layer_name]
350 + kv_cache_fused = attn_layer.kv_cache
351 + del dummy
352 +
353 + assert fusion_pass.matched_count == 1
354 +
355 + backend.check_before_ops(model.ops_in_model_before())
356 + backend.check_after_ops(model.ops_in_model_after())
357 +
358 + # Sweep-backed (18.2k pts, PR #42749): native-rope ref worst 7.7e-3 -> 1e-2;
359 + # AITER-triton-rope ref is itself approximate (plateau 1.28e-2) -> 2e-2.
360 + ATOL, RTOL = (2e-2, 2e-2) if enable_aiter_triton_rope else (1e-2, 1e-2)
361 + is_fp8_cache = model.kv_cache_dtype != dtype
362 +
363 + if q_fused.dtype == FP8_DTYPE:
364 + # Quant-query path: both q are fp8; compare dequant within 1 fp8 ULP.
365 + torch.testing.assert_close(
366 + q_unfused.float(), q_fused.float(), atol=1.25e-1, rtol=1.25e-1
367 + )
368 + else:
369 + torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)
370 +
371 + if not is_fp8_cache:
372 + # The AITER PTS kernel populates k_out only for non-FP8 caches.
373 + # With FP8, the kernel writes quantized K directly to the cache
374 + # and may leave k_out uninitialised. In production this is fine
375 + # because downstream attention reads K from the cache.
376 + torch.testing.assert_close(k_unfused, k_fused, atol=ATOL, rtol=RTOL)
377 +
378 + # Should be bit exact since no processing had been done on v for both paths
379 + torch.testing.assert_close(v_unfused, v_fused, atol=0.0, rtol=0.0)
380 +
381 + # fp8 vs triton-rope ref requires loosening tolerance to 1.25e-1.
382 + if is_fp8_cache and enable_aiter_triton_rope:
383 + cache_atol = cache_rtol = 1.25e-1
384 + else:
385 + cache_atol, cache_rtol = ATOL, RTOL
386 +
387 + torch.testing.assert_close(
388 + kv_cache_unfused[0].float(),
389 + kv_cache_fused[0].float(),
390 + atol=cache_atol,
391 + rtol=cache_rtol,
392 + )
393 +
394 +
395 +_FUSION_CONFIGS = [
396 + # Full rotary, both neox styles (the original coverage).
397 + pytest.param(64, 8, 64, 64, True, id="full-neox"),
398 + pytest.param(64, 8, 64, 64, False, id="full-non_neox"),
399 + # GLM-4.5/4.6/4.7 (glm4_moe.py:275 partial_rotary_factor=0.5, neox-style)
400 + pytest.param(32, 8, 128, 64, True, id="glm4_moe"),
401 + # GLM-4 dense (glm4.py:97,124 partial_rotary_factor=0.5, non-neox)
402 + pytest.param(32, 2, 128, 64, False, id="glm4_dense"),
403 + # Moondream3-style small head (head_size=64, rotary_dim=32)
404 + pytest.param(16, 2, 64, 32, True, id="partial_small_head"),
405 +]
406 +
407 +
408 +@pytest.mark.parametrize(
409 + "num_heads, num_kv_heads, head_size, rotary_dim, is_neox",
410 + _FUSION_CONFIGS,
411 +)
412 +@pytest.mark.parametrize(
413 + "attn_backend",
414 + [
415 + AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
416 + AttentionBackendEnum.ROCM_AITER_FA,
417 + ],
418 +)
419 +@pytest.mark.parametrize("num_tokens", [5, 16, 2048])
420 +@pytest.mark.parametrize("use_shuffle_kv_layout", ["1", "0"])
421 +@pytest.mark.parametrize(
422 + "kv_stride_order",
423 + [
424 + pytest.param((0, 1, 2, 3, 4), id="block_first"),
425 + pytest.param((1, 0, 2, 3, 4), id="kv_first"),
426 + ],
427 +)
428 +@pytest.mark.parametrize("enable_aiter_triton_rope", [True, False])
429 +@pytest.mark.parametrize("block_size", [16, 32, 64])
430 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
431 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
432 +@pytest.mark.parametrize("rms_norm_eps", [1e-5, 1e-6])
433 +@pytest.mark.parametrize("custom_op", ["+rotary_embedding", "+rms_norm"])
434 +@pytest.mark.skipif(
435 + not is_aiter_found_and_supported(),
436 + reason="Only test on ROCm with AITER installed and supported",
437 +)
438 +def test_qk_norm_rope_kvcache_fusion(
439 + num_tokens: int,
440 + num_heads: int,
441 + num_kv_heads: int,
442 + head_size: int,
443 + rotary_dim: int,
444 + is_neox: bool,
445 + attn_backend: AttentionBackendEnum,
446 + enable_aiter_triton_rope: bool,
447 + use_shuffle_kv_layout: str,
448 + kv_stride_order: tuple[int, ...],
449 + block_size: int,
450 + dtype: torch.dtype,
451 + kv_cache_dtype: str,
452 + rms_norm_eps: float,
453 + custom_op: str,
454 + monkeypatch: pytest.MonkeyPatch,
455 +):
456 + if (
457 + attn_backend == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN
458 + and use_shuffle_kv_layout == "1"
459 + ):
460 + pytest.skip("ROCM_AITER_UNIFIED_ATTN is NHD-only; shuffle env is ignored")
461 + _run_qk_norm_rope_kvcache_fusion_test(
462 + attn_backend=attn_backend,
463 + enable_aiter_triton_rope=enable_aiter_triton_rope,
464 + num_tokens=num_tokens,
465 + num_heads=num_heads,
466 + num_kv_heads=num_kv_heads,
467 + head_size=head_size,
468 + rotary_dim=rotary_dim,
469 + block_size=block_size,
470 + is_neox=is_neox,
471 + use_shuffle_kv_layout=use_shuffle_kv_layout,
472 + kv_stride_order=kv_stride_order,
473 + dtype=dtype,
474 + kv_cache_dtype=kv_cache_dtype,
475 + rms_norm_eps=rms_norm_eps,
476 + custom_op=custom_op,
477 + monkeypatch=monkeypatch,
478 + )
tests/quantization/test_fp8.py
+41
@@ -17,6 +17,10 @@ from vllm.config.model import ModelConfig
17 from vllm.model_executor.kernels.linear.scaled_mm import (
18 MarlinFP8ScaledMMLinearKernel,
19 )
20 +from vllm.model_executor.layers.attention.attention import (
21 + Attention,
22 + set_default_quant_scales,
23 +)
24 from vllm.model_executor.layers.fused_moe import FusedMoE
25 from vllm.model_executor.layers.quantization.fp8 import (
26 Fp8Config,
@@ -24,6 +28,7 @@ from vllm.model_executor.layers.quantization.fp8 import (
28 Fp8LinearMethod,
29 Fp8MoEMethod,
30 )
31 +from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
32 from vllm.model_executor.layers.quantization.online.fp8 import (
33 Fp8PerTensorOnlineLinearMethod,
34 )
@@ -422,6 +427,42 @@ def test_fp8_reloading(
427 method.process_weights_after_loading(layer)
428
429
430 +@pytest.mark.parametrize("source", ["checkpoint", "runtime_calc"])
431 +def test_kv_cache_scale_sync_to_host_copies(source):
432 + """Test device-to-host sync of the k/v quantization scales, for both the
433 + checkpoint-load and runtime-calc paths that produce them.
434 + """
435 + layer = torch.nn.Module()
436 + set_default_quant_scales(layer, register_buffer=True)
437 + layer.kv_cache_dtype = "fp8"
438 +
439 + if source == "checkpoint":
440 + # Scales come from the checkpoint, so runtime calc is disabled.
441 + layer.calculate_kv_scales = False
442 + method = BaseKVCacheMethod(quant_config=None)
443 + method.create_weights(layer)
444 + # 0.3 stays != 1.0 even after the fp8_fnuz x2 rescale.
445 + checkpoint_scale = torch.tensor(0.3, dtype=torch.float32)
446 + layer.k_scale.weight_loader(layer.k_scale, checkpoint_scale)
447 + layer.v_scale.weight_loader(layer.v_scale, checkpoint_scale)
448 + method.process_weights_after_loading(layer)
449 + else:
450 + # First forward computes distinct, non-unity scales from live k/v.
451 + layer.calculate_kv_scales = True
452 + query = torch.full((4, 8), 10.0)
453 + key = torch.full((4, 8), 60.0)
454 + value = torch.full((4, 8), 50.0)
455 + Attention.calc_kv_scales(layer, query, key, value)
456 +
457 + assert layer._k_scale_float != 1.0
458 + assert layer._v_scale_float != 1.0
459 + # Host copy must mirror both the float and the device scale tensor.
460 + assert layer._k_scale_cpu.item() == pytest.approx(layer._k_scale_float)
461 + assert layer._v_scale_cpu.item() == pytest.approx(layer._v_scale_float)
462 + assert layer._k_scale_cpu.item() == pytest.approx(layer._k_scale.item())
463 + assert layer._v_scale_cpu.item() == pytest.approx(layer._v_scale.item())
464 +
465 +
466 @pytest.mark.skipif(
467 not is_quant_method_supported("fp8"),
468 reason="FP8 is not supported on this GPU type.",
vllm/_aiter_ops.py
+122
@@ -2458,6 +2458,128 @@ class rocm_aiter_ops:
2458 gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y)
2459 return y
2460
2461 + @staticmethod
2462 + def fused_qk_norm_rope_and_cache(
2463 + qkv: torch.Tensor,
2464 + q_weight: torch.Tensor,
2465 + k_weight: torch.Tensor,
2466 + cos_sin_cache: torch.Tensor,
2467 + positions: torch.Tensor,
2468 + num_heads_q: int,
2469 + num_heads_k: int,
2470 + num_heads_v: int,
2471 + head_dim: int,
2472 + is_neox: bool,
2473 + rms_norm_eps: float,
2474 + q_out: torch.Tensor,
2475 + k_cache: torch.Tensor,
2476 + v_cache: torch.Tensor,
2477 + slot_mapping: torch.Tensor,
2478 + k_scale: torch.Tensor,
2479 + v_scale: torch.Tensor,
2480 + k_out: torch.Tensor | None,
2481 + v_out: torch.Tensor | None,
2482 + return_kv: bool,
2483 + use_shuffle_layout: bool,
2484 + block_size: int,
2485 + x: int,
2486 + rotary_dim: int = 0,
2487 + ):
2488 + from aiter.ops.fused_qk_norm_rope_cache_quant import (
2489 + fused_qk_norm_rope_cache_pts_quant_shuffle,
2490 + )
2491 +
2492 + fused_qk_norm_rope_cache_pts_quant_shuffle(
2493 + qkv,
2494 + q_weight,
2495 + k_weight,
2496 + cos_sin_cache,
2497 + positions,
2498 + qkv.size(0),
2499 + num_heads_q,
2500 + num_heads_k,
2501 + num_heads_v,
2502 + head_dim,
2503 + is_neox,
2504 + rms_norm_eps,
2505 + q_out,
2506 + k_cache,
2507 + v_cache,
2508 + slot_mapping,
2509 + k_scale,
2510 + v_scale,
2511 + k_out,
2512 + v_out,
2513 + return_kv,
2514 + use_shuffle_layout,
2515 + block_size,
2516 + x,
2517 + rotary_dim,
2518 + )
2519 +
2520 + @staticmethod
2521 + def do_qk_norm_rope_kvcache_update(
2522 + qkv: torch.Tensor,
2523 + q_weight: torch.Tensor,
2524 + k_weight: torch.Tensor,
2525 + cos_sin_cache: torch.Tensor,
2526 + positions: torch.Tensor,
2527 + num_heads_q: int,
2528 + num_heads_k: int,
2529 + head_dim: int,
2530 + is_neox: bool,
2531 + rms_norm_eps: float,
2532 + q_out: torch.Tensor,
2533 + k_out: torch.Tensor,
2534 + key_cache: torch.Tensor,
2535 + value_cache: torch.Tensor,
2536 + slot_mapping: torch.Tensor,
2537 + k_scale: torch.Tensor,
2538 + v_scale: torch.Tensor,
2539 + kv_cache_dtype: str,
2540 + use_shuffle_layout: bool,
2541 + ) -> None:
2542 + """Run the fused QK-norm+RoPE+KV-cache op on already-split k/v caches.
2543 +
2544 + Shared by the AITER FA and unified-attention impls. The caller splits
2545 + kv_cache, since the unbind dim depends on the layout (e.g. the unified
2546 + encoder-decoder path is K/V-first), and passes use_shuffle_layout
2547 + (unified reads NHD and must pass False).
2548 + """
2549 + if kv_cache_dtype.startswith("fp8"):
2550 + key_cache = key_cache.view(current_platform.fp8_dtype())
2551 + value_cache = value_cache.view(current_platform.fp8_dtype())
2552 + # Partial-rotary support (e.g. GLM-4.7 applies rotary to only a prefix
2553 + # of each head's channel dim).
2554 + rotary_dim = cos_sin_cache.shape[-1]
2555 + kernel_rotary_dim = 0 if rotary_dim == head_dim else rotary_dim
2556 + rocm_aiter_ops.fused_qk_norm_rope_and_cache(
2557 + qkv=qkv,
2558 + q_weight=q_weight,
2559 + k_weight=k_weight,
2560 + cos_sin_cache=cos_sin_cache,
2561 + positions=positions,
2562 + num_heads_q=num_heads_q,
2563 + num_heads_k=num_heads_k,
2564 + num_heads_v=num_heads_k,
2565 + head_dim=head_dim,
2566 + is_neox=is_neox,
2567 + rms_norm_eps=rms_norm_eps,
2568 + q_out=q_out,
2569 + k_cache=key_cache,
2570 + v_cache=value_cache,
2571 + slot_mapping=slot_mapping,
2572 + k_scale=k_scale,
2573 + v_scale=v_scale,
2574 + k_out=k_out,
2575 + v_out=None,
2576 + return_kv=True,
2577 + use_shuffle_layout=use_shuffle_layout,
2578 + block_size=key_cache.shape[1],
2579 + x=16 // key_cache.element_size(),
2580 + rotary_dim=kernel_rotary_dim,
2581 + )
2582 +
2583 @staticmethod
2584 def triton_rope_and_cache(
2585 query: torch.Tensor,
vllm/compilation/passes/fusion/qk_norm_rope_fusion.py
+22 -1
@@ -25,6 +25,11 @@ logger = init_logger(__name__)
25
26 FUSED_QK_ROPE_OP = torch.ops._C.fused_qk_norm_rope.default
27
28 +# Head dimensions supported by csrc/fused_qknorm_rope_kernel.cu's
29 +# launchFusedQKNormRope and launchFusedQKNormRopeNTokenHeads dispatchers.
30 +# Keep in sync with the switch statements in that file.
31 +SUPPORTED_FUSED_QK_NORM_ROPE_HEAD_DIMS: tuple[int, ...] = (64, 128, 256)
32 +
33 P = ParamSpec("P")
34
35
@@ -190,7 +195,12 @@ class QkNormRopePattern:
195
196
197 class QKNormRoPEFusionPass(VllmPatternMatcherPass):
193 - """Fuse Q/K RMSNorm + RoPE into fused_qk_norm_rope when the custom op exists."""
198 + """Fuse Q/K RMSNorm + RoPE into fused_qk_norm_rope when the custom op exists.
199 +
200 + Registers patterns for both standard vLLM ops and ROCm AITER ops
201 + (when AITER is enabled), so the fusion fires regardless of which
202 + RMSNorm/RoPE implementation the graph uses.
203 + """
204
205 @enable_fake_mode
206 def __init__(self, config: VllmConfig) -> None:
@@ -216,6 +226,17 @@ class QKNormRoPEFusionPass(VllmPatternMatcherPass):
226 )
227 return
228
229 + for layer in attn_layers.values():
230 + if layer.head_size not in SUPPORTED_FUSED_QK_NORM_ROPE_HEAD_DIMS:
231 + logger.warning_once(
232 + "QK Norm+RoPE fusion not enabled: layer head_size=%d is not "
233 + "supported by fused_qk_norm_rope kernel (supported: %s). "
234 + "Falling back to unfused QK norm + RoPE path.",
235 + layer.head_size,
236 + SUPPORTED_FUSED_QK_NORM_ROPE_HEAD_DIMS,
237 + )
238 + return
239 +
240 self._attention_geometries = tuple(
241 sorted(
242 {
vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py new
+492
@@ -0,0 +1,492 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import inspect
5 +from collections.abc import Callable
6 +from typing import ParamSpec
7 +
8 +import torch
9 +import torch._inductor.pattern_matcher as pm
10 +from torch import fx
11 +from torch._higher_order_ops.auto_functionalize import auto_functionalized
12 +from torch._inductor.pattern_matcher import PatternMatcherPass
13 +
14 +import vllm.ir.ops
15 +from vllm.config import VllmConfig, get_layers_from_vllm_config
16 +from vllm.config.utils import Range
17 +from vllm.logger import init_logger
18 +from vllm.model_executor.layers.attention.attention import (
19 + Attention,
20 + get_attention_context,
21 +)
22 +from vllm.platforms import current_platform
23 +from vllm.utils.torch_utils import direct_register_custom_op
24 +
25 +from ..inductor_pass import enable_fake_mode
26 +from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
27 +from .matcher_utils import MatcherRotaryEmbedding
28 +from .rms_quant_fusion import empty_bf16, empty_fp32, empty_i64
29 +
30 +logger = init_logger(__name__)
31 +
32 +P = ParamSpec("P")
33 +
34 +# Head sizes the fused kernel fused_qk_norm_rope_cache_pts_quant_shuffle() supports
35 +# Other sizes hard-abort, so skip those layers.
36 +SUPPORTED_FUSED_QK_NORM_ROPE_KVCACHE_HEAD_DIMS: tuple[int, ...] = (64, 128, 256)
37 +
38 +
39 +# ---------------------------------------------------------------------------
40 +# Custom op: fused QK-norm + RoPE + KV cache update
41 +# ---------------------------------------------------------------------------
42 +
43 +
44 +def fused_qk_norm_rope_and_unified_kv_cache_update_impl(
45 + q_out: torch.Tensor,
46 + k_out: torch.Tensor,
47 + qkv: torch.Tensor,
48 + positions: torch.Tensor,
49 + q_weight: torch.Tensor,
50 + k_weight: torch.Tensor,
51 + rms_norm_eps: float,
52 + cos_sin_cache: torch.Tensor,
53 + is_neox: bool,
54 + layer_name: str = "",
55 +) -> torch.Tensor:
56 + _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
57 + if layer_slot_mapping is not None:
58 + attn_layer.impl.do_qk_norm_rope_kvcache_update(
59 + attn_layer,
60 + qkv,
61 + q_out,
62 + k_out,
63 + positions,
64 + q_weight,
65 + k_weight,
66 + rms_norm_eps,
67 + cos_sin_cache,
68 + is_neox,
69 + kv_cache,
70 + layer_slot_mapping,
71 + )
72 + else:
73 + # Profiling/dummy run: define q_out/k_out (consumed by attention).
74 + q_out.zero_()
75 + k_out.zero_()
76 +
77 + return torch.empty(0, device=qkv.device, dtype=qkv.dtype)
78 +
79 +
80 +def fused_qk_norm_rope_and_unified_kv_cache_update_fake(
81 + q_out: torch.Tensor,
82 + k_out: torch.Tensor,
83 + qkv: torch.Tensor,
84 + positions: torch.Tensor,
85 + q_weight: torch.Tensor,
86 + k_weight: torch.Tensor,
87 + rms_norm_eps: float,
88 + cos_sin_cache: torch.Tensor,
89 + is_neox: bool,
90 + layer_name: str = "",
91 +) -> torch.Tensor:
92 + return torch.empty(0, device=qkv.device, dtype=qkv.dtype)
93 +
94 +
95 +direct_register_custom_op(
96 + op_name="fused_qk_norm_rope_and_unified_kv_cache_update",
97 + op_func=fused_qk_norm_rope_and_unified_kv_cache_update_impl,
98 + mutates_args=["q_out", "k_out"],
99 + fake_impl=fused_qk_norm_rope_and_unified_kv_cache_update_fake,
100 +)
101 +
102 +
103 +# ---------------------------------------------------------------------------
104 +# Pattern: QK-norm + RoPE + unified_kv_cache_update
105 +# ---------------------------------------------------------------------------
106 +
107 +
108 +class QkNormRopeKvCachePattern:
109 + """
110 + Match the unfused sequence:
111 + q, k, v = split(qkv, ...)
112 + q = rms_norm(q.view(heads), q_weight).view(flat)
113 + k = rms_norm(k.view(heads), k_weight).view(flat)
114 + q, k = rotary_embedding(positions, q, k, cos_sin_cache, is_neox)
115 + q = q.view(num_heads, head_dim)
116 + k = k.view(num_kv_heads, head_dim)
117 + v = v.view(num_kv_heads, head_dim)
118 + dummy = unified_kv_cache_update(k, v, layer_name)
119 +
120 + Replace with:
121 + q_out = empty(...)
122 + k_out = empty(...)
123 + dummy = fused_qk_norm_rope_and_unified_kv_cache_update(
124 + q_out, k_out, qkv, positions, q_weight, k_weight,
125 + eps, cos_sin_cache, is_neox, layer_name)
126 + v = split(qkv, ...)[2].view(num_kv_heads, head_dim)
127 + """
128 +
129 + FUSED_OP = torch.ops.vllm.fused_qk_norm_rope_and_unified_kv_cache_update.default
130 +
131 + def __init__(
132 + self,
133 + layer: Attention,
134 + eps: float,
135 + is_neox: bool,
136 + quant_query: bool,
137 + ) -> None:
138 + self.layer_name = layer.layer_name
139 + self.num_heads = layer.num_heads
140 + self.num_kv_heads = layer.num_kv_heads
141 + self.head_size = layer.head_size
142 + self.head_size_v = layer.head_size_v
143 + self.eps = eps
144 + self.is_neox = is_neox
145 + self.quant_query = quant_query
146 +
147 + self.q_size = self.num_heads * self.head_size
148 + self.k_size = self.num_kv_heads * self.head_size
149 + self.v_size = self.num_kv_heads * self.head_size_v
150 +
151 + self.rope_matcher = MatcherRotaryEmbedding(
152 + is_neox=is_neox,
153 + head_size=self.head_size,
154 + num_heads=self.num_heads,
155 + num_kv_heads=self.num_kv_heads,
156 + )
157 +
158 + def get_inputs(self) -> list[torch.Tensor]:
159 + T = 5
160 + L = 4096
161 + qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size)
162 + positions = empty_i64(T)
163 + q_weight = empty_bf16(1, self.head_size)
164 + k_weight = empty_bf16(1, self.head_size)
165 + cos_sin_cache = empty_bf16(L, self.head_size)
166 + inputs = [qkv, positions, q_weight, k_weight, cos_sin_cache]
167 + if self.quant_query:
168 + q_scale = empty_fp32(1)
169 + inputs += [q_scale]
170 + return inputs
171 +
172 + def pattern_non_fp8_quant_query(
173 + self,
174 + qkv: torch.Tensor,
175 + positions: torch.Tensor,
176 + q_weight: torch.Tensor,
177 + k_weight: torch.Tensor,
178 + cos_sin_cache: torch.Tensor,
179 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
180 + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
181 + q_by_head = q.view(-1, self.q_size // self.head_size, self.head_size)
182 + q_normed = vllm.ir.ops.rms_norm(q_by_head, q_weight, self.eps)
183 + q_flat = q_normed.view(-1, self.q_size)
184 +
185 + k_by_head = k.view(-1, self.k_size // self.head_size, self.head_size)
186 + k_normed = vllm.ir.ops.rms_norm(k_by_head, k_weight, self.eps)
187 + k_flat = k_normed.view(-1, self.k_size)
188 +
189 + q_rope, k_rope = self.rope_matcher(positions, q_flat, k_flat, cos_sin_cache)
190 +
191 + q_rope = q_rope.view(-1, self.num_heads, self.head_size)
192 + k_rope = k_rope.view(-1, self.num_kv_heads, self.head_size)
193 + v = v.view(-1, self.num_kv_heads, self.head_size_v)
194 + dummy = torch.ops.vllm.unified_kv_cache_update(k_rope, v, self.layer_name)
195 + return dummy, q_rope, k_rope, v
196 +
197 + def replacement_non_fp8_quant_query(
198 + self,
199 + qkv: torch.Tensor,
200 + positions: torch.Tensor,
201 + q_weight: torch.Tensor,
202 + k_weight: torch.Tensor,
203 + cos_sin_cache: torch.Tensor,
204 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
205 + q_out = torch.empty(
206 + qkv.shape[0],
207 + self.num_heads,
208 + self.head_size,
209 + device=qkv.device,
210 + dtype=qkv.dtype,
211 + )
212 + k_out = torch.empty(
213 + qkv.shape[0],
214 + self.num_kv_heads,
215 + self.head_size,
216 + device=qkv.device,
217 + dtype=qkv.dtype,
218 + )
219 + _, _, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
220 + v = v.view(qkv.shape[0], self.num_kv_heads, self.head_size_v)
221 + results = auto_functionalized(
222 + self.FUSED_OP,
223 + q_out=q_out,
224 + k_out=k_out,
225 + qkv=qkv,
226 + positions=positions,
227 + q_weight=q_weight,
228 + k_weight=k_weight,
229 + rms_norm_eps=self.eps,
230 + cos_sin_cache=cos_sin_cache,
231 + is_neox=self.is_neox,
232 + layer_name=self.layer_name,
233 + )
234 + return results[0], results[1], results[2], v
235 +
236 + def pattern_fp8_quant_query(
237 + self,
238 + qkv: torch.Tensor,
239 + positions: torch.Tensor,
240 + q_weight: torch.Tensor,
241 + k_weight: torch.Tensor,
242 + cos_sin_cache: torch.Tensor,
243 + q_scale: torch.Tensor,
244 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
245 + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
246 + q_by_head = q.view(-1, self.q_size // self.head_size, self.head_size)
247 + q_normed = vllm.ir.ops.rms_norm(q_by_head, q_weight, self.eps)
248 + q_flat = q_normed.view(-1, self.q_size)
249 +
250 + k_by_head = k.view(-1, self.k_size // self.head_size, self.head_size)
251 + k_normed = vllm.ir.ops.rms_norm(k_by_head, k_weight, self.eps)
252 + k_flat = k_normed.view(-1, self.k_size)
253 +
254 + q_rope, k_rope = self.rope_matcher(positions, q_flat, k_flat, cos_sin_cache)
255 + # Match the quant-query op Attention.forward inserts (fp8 KV + UNIFIED).
256 + # Explicit auto_functionalized (out=[1]) keeps the quant node in the pattern.
257 + q_out = torch.empty_like(q_rope, dtype=current_platform.fp8_dtype())
258 + q_quant = auto_functionalized(
259 + torch.ops.vllm.rocm_aiter_per_tensor_quant.default,
260 + out=q_out,
261 + x=q_rope,
262 + scale=q_scale,
263 + is_dynamic=False,
264 + )
265 + # `scale` is mutable: its copy_ write-back to _q_scale bumps the mutation
266 + # region, so keep q flat (a reshape lands past the barrier and won't match).
267 + q_rope_fp8 = q_quant[1]
268 + q_scale_out = q_quant[2]
269 +
270 + k_rope = k_rope.view(-1, self.num_kv_heads, self.head_size)
271 + v = v.view(-1, self.num_kv_heads, self.head_size_v)
272 + dummy = torch.ops.vllm.unified_kv_cache_update(k_rope, v, self.layer_name)
273 + return dummy, q_rope_fp8, k_rope, v, q_scale_out
274 +
275 + def replacement_fp8_quant_query(
276 + self,
277 + qkv: torch.Tensor,
278 + positions: torch.Tensor,
279 + q_weight: torch.Tensor,
280 + k_weight: torch.Tensor,
281 + cos_sin_cache: torch.Tensor,
282 + q_scale: torch.Tensor,
283 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
284 + q_out = torch.empty(
285 + qkv.shape[0],
286 + self.num_heads,
287 + self.head_size,
288 + device=qkv.device,
289 + dtype=qkv.dtype,
290 + )
291 + k_out = torch.empty(
292 + qkv.shape[0],
293 + self.num_kv_heads,
294 + self.head_size,
295 + device=qkv.device,
296 + dtype=qkv.dtype,
297 + )
298 + _, _, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)
299 + v = v.view(qkv.shape[0], self.num_kv_heads, self.head_size_v)
300 + results = auto_functionalized(
301 + self.FUSED_OP,
302 + q_out=q_out,
303 + k_out=k_out,
304 + qkv=qkv,
305 + positions=positions,
306 + q_weight=q_weight,
307 + k_weight=k_weight,
308 + rms_norm_eps=self.eps,
309 + cos_sin_cache=cos_sin_cache,
310 + is_neox=self.is_neox,
311 + layer_name=self.layer_name,
312 + )
313 + # Re-apply the quant on the kernel's bf16 q_out; fused op does not quant q.
314 + # Same explicit auto_functionalized form as the pattern: [1] = quantized
315 + # q, [2] = scale (returned so the buffer-writeback use is preserved).
316 + q_fp8_flat = results[1].view(-1, self.q_size)
317 + q_fp8_out = torch.empty_like(q_fp8_flat, dtype=current_platform.fp8_dtype())
318 + q_requant = auto_functionalized(
319 + torch.ops.vllm.rocm_aiter_per_tensor_quant.default,
320 + out=q_fp8_out,
321 + x=q_fp8_flat,
322 + scale=q_scale,
323 + is_dynamic=False,
324 + )
325 + q_fp8 = q_requant[1] # flat to mirror the pattern (see note above)
326 + q_scale_out = q_requant[2]
327 + return results[0], q_fp8, results[2], v, q_scale_out
328 +
329 + @staticmethod
330 + def wrap_trace_fn(
331 + trace_fn: Callable[P, fx.GraphModule],
332 + *process_fx_fns: Callable[[fx.GraphModule], None],
333 + ) -> Callable[P, fx.GraphModule]:
334 + def wrapped(*args: P.args, **kwargs: P.kwargs) -> fx.GraphModule:
335 + gm = trace_fn(*args, **kwargs)
336 + for process_fx in process_fx_fns:
337 + process_fx(gm)
338 +
339 + return gm
340 +
341 + return wrapped
342 +
343 + @staticmethod
344 + def fx_view_to_reshape(gm: torch.fx.GraphModule) -> None:
345 + from torch._inductor.fx_passes.post_grad import view_to_reshape
346 +
347 + view_to_reshape(gm)
348 +
349 + def _register(self, pattern, replacement, pm_pass) -> None:
350 + trace_fn = QkNormRopeKvCachePattern.wrap_trace_fn(
351 + pm.fwd_only,
352 + QkNormRopeKvCachePattern.fx_view_to_reshape,
353 + )
354 +
355 + # Pre-build the search pattern with `ignore_types=(int, torch.SymInt)`
356 + # and pass it via `search_fn_pattern=` so torch skips both of its
357 + # internal `fx_to_pattern` calls and treats dynamic-shape SymInts as
358 + # wildcards.
359 + inputs = self.get_inputs()
360 + argnames = [*inspect.signature(pattern).parameters.keys()]
361 + search_gm = trace_fn(pattern, inputs)
362 + search_fn_pattern = pm.fx_to_pattern(
363 + search_gm,
364 + ignore_types=(int, torch.SymInt),
365 + argnames=argnames,
366 + )
367 +
368 + pm.register_replacement(
369 + pattern,
370 + replacement,
371 + inputs,
372 + trace_fn,
373 + pm_pass,
374 + search_fn_pattern=search_fn_pattern,
375 + )
376 +
377 + def register(self, pm_pass: PatternMatcherPass) -> None:
378 + # make_fx counts `self` in bound-method code params; wrap as plain fns.
379 + # Distinct names per branch so mypy doesn't see one name, two signatures.
380 + if self.quant_query:
381 +
382 + def pattern_q(qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale):
383 + return self.pattern_fp8_quant_query(
384 + qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
385 + )
386 +
387 + def replacement_q(
388 + qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
389 + ):
390 + return self.replacement_fp8_quant_query(
391 + qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale
392 + )
393 +
394 + self._register(pattern_q, replacement_q, pm_pass)
395 + else:
396 +
397 + def pattern_noq(qkv, positions, q_weight, k_weight, cos_sin_cache):
398 + return self.pattern_non_fp8_quant_query(
399 + qkv, positions, q_weight, k_weight, cos_sin_cache
400 + )
401 +
402 + def replacement_noq(qkv, positions, q_weight, k_weight, cos_sin_cache):
403 + return self.replacement_non_fp8_quant_query(
404 + qkv, positions, q_weight, k_weight, cos_sin_cache
405 + )
406 +
407 + self._register(pattern_noq, replacement_noq, pm_pass)
408 +
409 +
410 +# ---------------------------------------------------------------------------
411 +# Pass class
412 +# ---------------------------------------------------------------------------
413 +
414 +
415 +class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):
416 + """
417 + Fuse QK-norm + RoPE + KV cache update into a single AITER HIP kernel.
418 +
419 + Supersedes both QKNormRoPEFusionPass and RopeKVCacheFusionPass for
420 + attention layers that support the combined operation, eliminating two
421 + separate kernel launches and the intermediate memory traffic.
422 + """
423 +
424 + @enable_fake_mode
425 + def __init__(self, config: VllmConfig) -> None:
426 + super().__init__(config)
427 +
428 + self.patterns: PatternMatcherPass = PatternMatcherPass(
429 + pass_name="qk_norm_rope_kvcache_fusion_pass"
430 + )
431 +
432 + cc = config.compilation_config
433 + self.max_token_num = cc.pass_config.rope_kvcache_fusion_max_token_num
434 +
435 + dtype = config.model_config.dtype
436 + if dtype not in (torch.bfloat16, torch.float16):
437 + logger.warning_once(
438 + "QK Norm+RoPE+KVCache fusion not enabled: unsupported dtype %s", dtype
439 + )
440 + return
441 +
442 + attn_layers = get_layers_from_vllm_config(config, Attention)
443 +
444 + for _, layer in attn_layers.items():
445 + if not layer.impl.fused_qk_norm_rope_kvcache_supported():
446 + continue
447 + if layer.head_size not in SUPPORTED_FUSED_QK_NORM_ROPE_KVCACHE_HEAD_DIMS:
448 + logger.warning_once(
449 + "QK Norm+RoPE+KVCache fusion not enabled for a layer: "
450 + "head_size=%d is not supported by the "
451 + "fused_qk_norm_rope_cache_pts_quant_shuffle kernel "
452 + "(supported: %s). Falling back to the unfused path.",
453 + layer.head_size,
454 + SUPPORTED_FUSED_QK_NORM_ROPE_KVCACHE_HEAD_DIMS,
455 + )
456 + continue
457 + if layer.head_size_v != layer.head_size:
458 + # The fused kernel uses a single head_dim for q/k/v.
459 + logger.warning_once(
460 + "QK Norm+RoPE+KVCache fusion not enabled for a layer: "
461 + "head_size_v=%d differs from head_size=%d, which the fused "
462 + "kernel does not support. Falling back to the unfused path.",
463 + layer.head_size_v,
464 + layer.head_size,
465 + )
466 + continue
467 + for epsilon in [1e-5, 1e-6]:
468 + for neox in [True, False]:
469 + for quant_q in [False, True]:
470 + QkNormRopeKvCachePattern(
471 + layer=layer,
472 + eps=epsilon,
473 + is_neox=neox,
474 + quant_query=quant_q,
475 + ).register(self.patterns)
476 +
477 + self.dump_patterns(config, self.patterns)
478 +
479 + @VllmInductorPass.time_and_log
480 + def __call__(self, graph: fx.Graph) -> None:
481 + self.matched_count = self.patterns.apply(graph)
482 + logger.info(
483 + "QK-Norm+RoPE+KVCache fusion: replaced %s pattern(s) "
484 + "with AITER fused_qk_norm_rope_cache_pts_quant_shuffle",
485 + self.matched_count,
486 + )
487 +
488 + def is_applicable_for_range(self, compile_range: Range) -> bool:
489 + return compile_range.end <= self.max_token_num
490 +
491 + def uuid(self) -> str:
492 + return VllmInductorPass.hash_source(self, QkNormRopeKvCachePattern)
vllm/compilation/passes/pass_manager.py
+6
@@ -38,6 +38,7 @@ if current_platform.is_cuda_alike():
38 from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass
39 from .fusion.mla_rope_kvcache_cat_fusion import MLARoPEKVCacheCatFusionPass
40 from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
41 + from .fusion.qk_norm_rope_kvcache_fusion import QkNormRopeKvCacheFusionPass
42 from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
43 from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass
44 from .utility.scatter_split_replace import ScatterSplitReplacementPass
@@ -171,6 +172,11 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
172 if rocm_aiter_ops.is_enabled():
173 self.passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]
174
175 + if self.pass_config.fuse_qk_norm_rope_kvcache:
176 + self.passes += [SplitCoalescingPass(config)]
177 + self.passes += [ScatterSplitReplacementPass(config)]
178 + self.passes += [QkNormRopeKvCacheFusionPass(config)]
179 +
180 if (
181 self.pass_config.fuse_mla_dual_rms_norm
182 and rocm_aiter_ops.is_enabled()
vllm/config/compilation.py
+36 -2
@@ -146,10 +146,16 @@ class PassConfig:
146 """Fuse paired q/kv RMS norms in MLA attention."""
147 fuse_rope_kvcache: bool = None # type: ignore[assignment]
148 """Fuse the QK rope + KV cache ops."""
149 + fuse_qk_norm_rope_kvcache: bool = Field(default=None) # type: ignore[assignment]
150 + """Fuse QK RMSNorm + RoPE + KV cache update into a single AITER HIP
151 + kernel. Supersedes both enable_qk_norm_rope_fusion and fuse_rope_kvcache
152 + for layers that support it. Auto-enabled at O1+ on ROCm for models
153 + with QK-norm (e.g. Qwen3-MoE)."""
154
155 rope_kvcache_fusion_max_token_num: int = 256
156 """The threshold for ROCm AITER RoPE+KVCache fusion e.g. for small batch decode.
157 Larger batch sizes e.g. during prefill will use the unfused kernels.
158 + Also applies to the fused QK-Norm+RoPE+KVCache pass.
159 """
160
161 fi_allreduce_fusion_max_size_mb: float | None = None
@@ -228,6 +234,8 @@ class PassConfig:
234 "fuse_act_padding",
235 "fuse_mla_dual_rms_norm",
236 "fuse_rope_kvcache",
237 + "fuse_qk_norm_rope_kvcache",
238 + "enable_qk_norm_rope_fusion",
239 "fuse_rope_kvcache_cat_mla",
240 mode="wrap",
241 )
@@ -288,6 +296,12 @@ class PassConfig:
296 "The fusion will be disabled."
297 )
298 self.fuse_rope_kvcache = False
299 + if self.fuse_qk_norm_rope_kvcache and not current_platform.is_rocm():
300 + logger.warning_once(
301 + "QK-Norm+RoPE+KVCache fusion requires ROCm with AITER. "
302 + "The fusion will be disabled."
303 + )
304 + self.fuse_qk_norm_rope_kvcache = False
305 if self.fuse_rope_kvcache_cat_mla and not current_platform.is_cuda_alike():
306 logger.warning_once(
307 "MLA KV cache update with RoPE fusion enabled but the "
@@ -302,10 +316,13 @@ class PassConfig:
316 after all defaults are finalized.
317 TODO also log the compile ranges for which this is enabled.
318 """
319 + fusion_prefixes = ("fuse_", "enable_")
320 enabled_fusions = [
306 - f.name[len("fuse_") :]
321 + f.name[len(prefix) :]
322 for f in fields(self) # type: ignore[arg-type]
308 - if getattr(self, f.name) and f.name.startswith("fuse_")
323 + if getattr(self, f.name)
324 + for prefix in fusion_prefixes
325 + if f.name.startswith(prefix)
326 ]
327
328 if enabled_fusions:
@@ -947,6 +964,7 @@ class CompilationConfig:
964 # TODO(zhuhaoran): support rope native forward match and remove this.
965 # Linked issue: https://github.com/vllm-project/vllm/issues/28042
966 self.custom_ops.append("+rotary_embedding")
967 +
968 if (
969 self.pass_config.fuse_rope_kvcache
970 and "+rotary_embedding" not in self.custom_ops
@@ -955,6 +973,12 @@ class CompilationConfig:
973 # Linked issue: https://github.com/vllm-project/vllm/issues/28042
974 self.custom_ops.append("+rotary_embedding")
975
976 + if (
977 + self.pass_config.fuse_qk_norm_rope_kvcache
978 + and "+rotary_embedding" not in self.custom_ops
979 + ):
980 + self.custom_ops.append("+rotary_embedding")
981 +
982 if (
983 is_torch_equal_or_newer("2.9.0.dev")
984 and "combo_kernels" not in self.inductor_compile_config
@@ -1137,6 +1161,16 @@ class CompilationConfig:
1161 "to enable RoPE+KV cache fusion."
1162 )
1163 self.pass_config.fuse_rope_kvcache = False
1164 + if self.pass_config.fuse_qk_norm_rope_kvcache:
1165 + logger.warning_once(
1166 + "fuse_qk_norm_rope_kvcache is enabled, but "
1167 + "splitting_ops is None and Inductor graph partition "
1168 + "is not enabled. Disabling fuse_qk_norm_rope_kvcache. "
1169 + "Please either set splitting_ops to an empty list [] "
1170 + "or set use_inductor_graph_partition to True "
1171 + "to enable QK-Norm+RoPE+KV cache fusion."
1172 + )
1173 + self.pass_config.fuse_qk_norm_rope_kvcache = False
1174 self.splitting_ops.append("vllm::unified_kv_cache_update")
1175 self.splitting_ops.append("vllm::unified_mla_kv_cache_update")
1176
vllm/config/vllm.py
+32
@@ -193,6 +193,15 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool:
193 return rocm_aiter_ops.is_enabled() and check_aiter_fused_qk_rmsnorm()
194
195
196 +def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool:
197 + """Enable fused QK-norm + RoPE + KV cache update on ROCm with AITER."""
198 + from vllm._aiter_ops import rocm_aiter_ops
199 +
200 + if not rocm_aiter_ops.is_enabled():
201 + return False
202 + return cfg.compilation_config.is_custom_op_enabled("rotary_embedding")
203 +
204 +
205 OPTIMIZATION_LEVEL_00 = {
206 "compilation_config": {
207 "pass_config": {
@@ -205,6 +214,8 @@ OPTIMIZATION_LEVEL_00 = {
214 "fuse_act_padding": False,
215 "fuse_mla_dual_rms_norm": False,
216 "fuse_rope_kvcache": False,
217 + "fuse_qk_norm_rope_kvcache": False,
218 + "enable_qk_norm_rope_fusion": False,
219 "fuse_rope_kvcache_cat_mla": False,
220 },
221 "cudagraph_mode": CUDAGraphMode.NONE,
@@ -226,6 +237,8 @@ OPTIMIZATION_LEVEL_01 = {
237 "fuse_act_padding": enable_norm_pad_fusion,
238 "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion,
239 "fuse_rope_kvcache": False,
240 + "fuse_qk_norm_rope_kvcache": False,
241 + "enable_qk_norm_rope_fusion": False,
242 "fuse_rope_kvcache_cat_mla": False,
243 },
244 "cudagraph_mode": CUDAGraphMode.PIECEWISE,
@@ -247,6 +260,8 @@ OPTIMIZATION_LEVEL_02 = {
260 "fuse_act_padding": enable_norm_pad_fusion,
261 "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion,
262 "fuse_rope_kvcache": enable_rope_kvcache_fusion,
263 + "fuse_qk_norm_rope_kvcache": enable_qk_norm_rope_kvcache,
264 + "enable_qk_norm_rope_fusion": False,
265 "fuse_rope_kvcache_cat_mla": enable_rope_kvcache_mla_fusion,
266 },
267 "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,
@@ -268,6 +283,8 @@ OPTIMIZATION_LEVEL_03 = {
283 "fuse_act_padding": enable_norm_pad_fusion,
284 "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion,
285 "fuse_rope_kvcache": enable_rope_kvcache_fusion,
286 + "fuse_qk_norm_rope_kvcache": enable_qk_norm_rope_kvcache,
287 + "enable_qk_norm_rope_fusion": False,
288 "fuse_rope_kvcache_cat_mla": enable_rope_kvcache_mla_fusion,
289 },
290 "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,
@@ -1971,6 +1988,21 @@ class VllmConfig:
1988 compile_range_end,
1989 )
1990
1991 + if compilation_config.pass_config.fuse_qk_norm_rope_kvcache:
1992 + max_token_num = (
1993 + compilation_config.pass_config.rope_kvcache_fusion_max_token_num
1994 + )
1995 + if max_token_num is not None:
1996 + if compile_range_end is not None and max_token_num < compile_range_end:
1997 + computed_compile_ranges_endpoints.append(max_token_num)
1998 + else:
1999 + logger.debug(
2000 + "Max num batched tokens below qk_norm+rope+kvcache "
2001 + "fusion threshold, fusion enabled for "
2002 + "num_tokens <= %d.",
2003 + compile_range_end,
2004 + )
2005 +
2006 if compilation_config.compile_ranges_endpoints is not None:
2007 for x in compilation_config.compile_ranges_endpoints:
2008 assert isinstance(x, int)
vllm/model_executor/layers/attention/attention.py
+5 -1
@@ -136,10 +136,12 @@ def set_default_quant_scales(layer: nn.Module, register_buffer: bool = False) ->
136
137 # We also keep q/k/v_scale on host (cpu) memory for attention
138 # backends that require the scales to be on host instead of on device.
139 - # e.g. Flashinfer
139 + # e.g. Flashinfer & AITER
140 layer._q_scale_float = 1.0
141 layer._k_scale_float = 1.0
142 layer._v_scale_float = 1.0
143 + layer._k_scale_cpu = torch.tensor(1.0, dtype=torch.float32)
144 + layer._v_scale_cpu = torch.tensor(1.0, dtype=torch.float32)
145 layer._prob_scale_float = 1.0
146
147 # Initialize q/k/v range constants used by calc_kv_scales
@@ -586,6 +588,8 @@ class Attention(nn.Module, AttentionLayerBase):
588 self._q_scale_float = self._q_scale.item()
589 self._k_scale_float = self._k_scale.item()
590 self._v_scale_float = self._v_scale.item()
591 + self._k_scale_cpu.fill_(self._k_scale_float)
592 + self._v_scale_cpu.fill_(self._v_scale_float)
593 # We only calculate the scales once
594 self.calculate_kv_scales = False
595
vllm/model_executor/layers/attention/mla_attention.py
+2
@@ -1024,6 +1024,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1024 self._q_scale_float = self._q_scale.item()
1025 self._k_scale_float = self._k_scale.item()
1026 self._v_scale_float = self._v_scale.item()
1027 + self._k_scale_cpu.fill_(self._k_scale_float)
1028 + self._v_scale_cpu.fill_(self._v_scale_float)
1029 self.calculate_kv_scales = False
1030
1031 def get_attn_backend(self) -> type[AttentionBackend]:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+4
@@ -1200,6 +1200,10 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):
1200 layer._v_scale_float = _to_scalar(layer.v_scale)
1201 layer._q_scale_float = _to_scalar(layer.q_scale)
1202
1203 + # Sync host (cpu) scale copies read by AITER fused kernels.
1204 + layer._k_scale_cpu.fill_(layer._k_scale_float)
1205 + layer._v_scale_cpu.fill_(layer._v_scale_float)
1206 +
1207 # Discard all placeholders.
1208 del layer.k_scale
1209 del layer.v_scale
vllm/model_executor/layers/quantization/kv_cache.py
+3
@@ -144,6 +144,9 @@ class BaseKVCacheMethod(QuantizeMethodBase):
144 layer._v_scale.copy_(v_scale)
145 layer._k_scale_float = k_scale
146 layer._v_scale_float = v_scale
147 + # Sync host (cpu) scale copies read by AITER fused kernels.
148 + layer._k_scale_cpu.fill_(k_scale)
149 + layer._v_scale_cpu.fill_(v_scale)
150 if k_scale == 1.0 and v_scale == 1.0 and "e5m2" not in layer.kv_cache_dtype:
151 logger.warning_once(
152 "Using KV cache scaling factor 1.0 for fp8_e4m3. "
vllm/v1/attention/backend.py
+33
@@ -757,7 +757,9 @@ class AttentionMetadataBuilder(ABC, Generic[M]):
757 class AttentionLayer(Protocol):
758 _q_scale: torch.Tensor
759 _k_scale: torch.Tensor
760 + _k_scale_cpu: torch.Tensor
761 _v_scale: torch.Tensor
762 + _v_scale_cpu: torch.Tensor
763 _q_scale_float: float
764 _k_scale_float: float
765 _v_scale_float: float
@@ -923,6 +925,14 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):
925 """
926 return False
927
928 + def fused_qk_norm_rope_kvcache_supported(self):
929 + """
930 + Does this attention implementation support fused QKNorm+RoPE+KVCache fusion.
931 + This is used by the QkNormRopeKvCachePattern to only fuse the QKNorm ops
932 + with the RoPE ops and the KV cache update for implementations that support it.
933 + """
934 + return False
935 +
936 def fused_rope_kvcache_supported(self):
937 """
938 Does this attention implementation support RoPE+KVCache fusion.
@@ -931,6 +941,29 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):
941 """
942 return False
943
944 + def do_qk_norm_rope_kvcache_update(
945 + self,
946 + layer: AttentionLayer,
947 + qkv: torch.Tensor,
948 + q_out: torch.Tensor,
949 + k_out: torch.Tensor,
950 + positions: torch.Tensor,
951 + q_weight: torch.Tensor,
952 + k_weight: torch.Tensor,
953 + rms_norm_eps: float,
954 + cos_sin_cache: torch.Tensor,
955 + is_neox: bool,
956 + kv_cache: torch.Tensor,
957 + layer_slot_mapping: torch.Tensor,
958 + ):
959 + """
960 + If `fused_qk_norm_rope_kvcache_supported` returns True, this method
961 + will be called by the fused custom op. Applies QK-norm + RoPE and
962 + writes K/V to the KV cache. Results are written to the pre-allocated
963 + q_out and k_out tensors; V is split from QKV at the graph level.
964 + """
965 + raise NotImplementedError
966 +
967 def do_rope_and_kv_cache_update(
968 self,
969 layer: AttentionLayer,
vllm/v1/attention/backends/rocm_aiter_fa.py
+41
@@ -1449,6 +1449,47 @@ class AiterFlashAttentionImpl(AttentionImpl):
1449 and not rocm_aiter_ops.is_shuffle_kv_cache_enabled()
1450 )
1451
1452 + def fused_qk_norm_rope_kvcache_supported(self):
1453 + return rocm_aiter_ops.is_enabled()
1454 +
1455 + def do_qk_norm_rope_kvcache_update(
1456 + self,
1457 + layer: AttentionLayer,
1458 + qkv: torch.Tensor,
1459 + q_out: torch.Tensor,
1460 + k_out: torch.Tensor,
1461 + positions: torch.Tensor,
1462 + q_weight: torch.Tensor,
1463 + k_weight: torch.Tensor,
1464 + rms_norm_eps: float,
1465 + cos_sin_cache: torch.Tensor,
1466 + is_neox: bool,
1467 + kv_cache: torch.Tensor,
1468 + layer_slot_mapping: torch.Tensor,
1469 + ):
1470 + key_cache, value_cache = kv_cache.unbind(1)
1471 + rocm_aiter_ops.do_qk_norm_rope_kvcache_update(
1472 + qkv=qkv,
1473 + q_weight=q_weight,
1474 + k_weight=k_weight,
1475 + cos_sin_cache=cos_sin_cache,
1476 + positions=positions,
1477 + num_heads_q=self.num_heads,
1478 + num_heads_k=self.num_kv_heads,
1479 + head_dim=self.head_size,
1480 + is_neox=is_neox,
1481 + rms_norm_eps=rms_norm_eps,
1482 + q_out=q_out,
1483 + k_out=k_out,
1484 + key_cache=key_cache,
1485 + value_cache=value_cache,
1486 + slot_mapping=layer_slot_mapping,
1487 + k_scale=layer._k_scale_cpu,
1488 + v_scale=layer._v_scale_cpu,
1489 + kv_cache_dtype=self.kv_cache_dtype,
1490 + use_shuffle_layout=rocm_aiter_ops.is_shuffle_kv_cache_enabled(),
1491 + )
1492 +
1493 def do_rope_and_kv_cache_update(
1494 self,
1495 layer: AttentionLayer,
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+44
@@ -311,6 +311,50 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
311 def fused_rope_kvcache_supported(self):
312 return rocm_aiter_ops.is_enabled()
313
314 + def fused_qk_norm_rope_kvcache_supported(self):
315 + return rocm_aiter_ops.is_enabled()
316 +
317 + def do_qk_norm_rope_kvcache_update(
318 + self,
319 + layer: AttentionLayer,
320 + qkv: torch.Tensor,
321 + q_out: torch.Tensor,
322 + k_out: torch.Tensor,
323 + positions: torch.Tensor,
324 + q_weight: torch.Tensor,
325 + k_weight: torch.Tensor,
326 + rms_norm_eps: float,
327 + cos_sin_cache: torch.Tensor,
328 + is_neox: bool,
329 + kv_cache: torch.Tensor,
330 + layer_slot_mapping: torch.Tensor,
331 + ):
332 + # _split_kv_cache picks the unbind dim per layout (incl. the K/V-first
333 + # encoder-decoder path). unified reads NHD, so never write the shuffle
334 + # layout here.
335 + key_cache, value_cache = self._split_kv_cache(kv_cache)
336 + rocm_aiter_ops.do_qk_norm_rope_kvcache_update(
337 + qkv=qkv,
338 + q_weight=q_weight,
339 + k_weight=k_weight,
340 + cos_sin_cache=cos_sin_cache,
341 + positions=positions,
342 + num_heads_q=self.num_heads,
343 + num_heads_k=self.num_kv_heads,
344 + head_dim=self.head_size,
345 + is_neox=is_neox,
346 + rms_norm_eps=rms_norm_eps,
347 + q_out=q_out,
348 + k_out=k_out,
349 + key_cache=key_cache,
350 + value_cache=value_cache,
351 + slot_mapping=layer_slot_mapping,
352 + k_scale=layer._k_scale_cpu,
353 + v_scale=layer._v_scale_cpu,
354 + kv_cache_dtype=self.kv_cache_dtype,
355 + use_shuffle_layout=False,
356 + )
357 +
358 def do_rope_and_kv_cache_update(
359 self,
360 layer: AttentionLayer,