@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b715699df4

[ROCm][Perf] Extend QK-norm/RoPE/KV-cache fusion to MRoPE (#50212)

Signed-off-by: Vorapol Assavasangthong <Vorapol.Assavasangthong@amd.com> Co-authored-by: Santosh Hiremath <Santosh.Hiremath@amd.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com>

vorapolsiloai committed Sep 23, 2026 at 19:11 UTC b715699df45f091325b307c0f5ac874010fbb23e

11 files changed +1314 -118

tests/compile/fusions_e2e/common.py

+4 -1

@@ -128,8 +128,11 @@ FUSION_LOG_PATTERNS: dict[str, re.Pattern] = {

128 ),

129 "rms_quant_fusion": re.compile(r"rms_quant_fusion.py:\d+] Replaced (\d+) patterns"),

130 "act_quant_fusion": re.compile(r"act_quant_fusion.py:\d+] Replaced (\d+) patterns"),

131 + # The combined pass can consume these sites before the standalone pass.

132 "norm_rope_fusion": re.compile(

132 - r"qk_norm_rope_fusion.py:\d+] Fused QK Norm\+RoPE on (\d+) sites"

133 + r"(?:qk_norm_rope_fusion.py:\d+] Fused QK Norm\+RoPE on|"

134 + r"qk_norm_rope_kvcache_fusion.py:\d+] "

135 + r"QK-Norm\+RoPE(?:/MRoPE)?\+KVCache fusion: replaced) (\d+)"

136 ),

137 "attn_quant_fusion": re.compile(

138 r"attn_quant_fusion.py:\d+] Fused quant onto (\d+) attention nodes"

tests/compile/fusions_e2e/conftest.py

+15 -2

@@ -182,7 +182,11 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):

182 # TODO: Remove log counting in unit tests

183 # once all matchers implement VllmFusionPatternMatcherPass

184 n_expected = tp_size * num_ranges_activated

185 - if match_name not in ("attn_quant_fusion", "act_quant_fusion"):

185 + if match_name not in (

186 + "attn_quant_fusion",

187 + "act_quant_fusion",

188 + "norm_rope_fusion",

189 + ):

190 assert len(log_matches) == n_expected, (

191 f"Could not find {n_expected} {match_name} "

192 f"(found {len(log_matches)}) in:\n {log_holder.text}"

@@ -190,7 +194,16 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):

194

195 expected_matches = getattr(matches, match_name)

196

193 - if match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check:

197 + if match_name == "norm_rope_fusion":

198 + # Opaque LayerName lets the combined pass consume short-range

199 + # sites before the standalone pass; count either owner.

200 + assert len(log_matches) >= n_expected

201 + assert all(m in (0, expected_matches) for m in log_matches)

202 + assert sum(log_matches) == expected_matches * n_expected, (

203 + f"{match_name} expected {expected_matches * n_expected} "

204 + f"sites across both passes, found: {log_matches}"

205 + )

206 + elif match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check:

207 # AR+rms+quant takes precedence over rms+quant if activated.

208 # That means we get full matching where ar+rms+quant was not

209 # activated, and less where it was (only the smallest range).

tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py

+393 -31

@@ -8,6 +8,7 @@ import pytest

8 import torch

9 from packaging.version import Version

11 +import vllm.compilation.passes.fusion.qk_norm_rope_kvcache_fusion as fusion_module

12 import vllm.config

13 from tests.compile.backend import TestBackend

14 from tests.v1.attention.utils import (

@@ -15,8 +16,12 @@ from tests.v1.attention.utils import (

16 create_common_attn_metadata,

17 dense_kv_cache_views,

18 )

18 -from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops

19 -from vllm.compilation.passes.fusion.matcher_utils import ROTARY_OP

19 +from vllm._aiter_ops import (

20 + IS_AITER_FOUND,

21 + is_aiter_found_and_supported,

22 + rocm_aiter_ops,

23 +)

24 +from vllm.compilation.passes.fusion.matcher_utils import MROPE_OP, ROTARY_OP

25 from vllm.compilation.passes.fusion.qk_norm_rope_kvcache_fusion import (

26 QkNormRopeKvCacheFusionPass,

27 )

@@ -37,8 +42,12 @@ from vllm.config import (

42 from vllm.forward_context import get_forward_context, set_forward_context

43 from vllm.model_executor.layers.attention import Attention

44 from vllm.model_executor.layers.layernorm import RMSNorm

45 +from vllm.model_executor.layers.quantization.quark.quark import QuarkConfig

46 +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

47 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding

48 +from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding

49 from vllm.platforms import current_platform

50 +from vllm.utils.torch_utils import _USE_LAYERNAME, _encode_layer_name

51 from vllm.v1.attention.backend import (

52 AttentionBackend,

53 CommonAttentionMetadata,

@@ -94,6 +103,7 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):

103 device: torch.device,

104 rotary_dim: int | None = None,

105 prefix: str = "model.layers.0.self_attn.attn",

106 + expected_query_quant_group_shape: GroupShape | None = None,

107 ):

108 super().__init__()

109 self.num_heads = num_heads

@@ -133,6 +143,9 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):

143 attn_backend=attn_backend.get_class(),

144 )

145 self.attn_backend: type[AttentionBackend] = self.attn.get_attn_backend()

146 + if expected_query_quant_group_shape is not None:

147 + assert self.attn.query_quant is not None

148 + assert self.attn.query_quant.group_shape == expected_query_quant_group_shape

149 assert not self.attn_backend.forward_includes_kv_cache_update, (

150 f"Attention backend {self.attn_backend} does not support "

151 "fuse_qk_norm_rope_kvcache."

@@ -227,18 +240,15 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):

240 self.kv_cache_dtype != self.dtype

241 and self.attn.impl.supports_quant_query_input

242 ):

230 - q_fp8 = torch.empty_like(q, dtype=FP8_DTYPE)

231 - torch.ops.vllm.rocm_aiter_per_tensor_quant(

232 - q_fp8, q, self.attn._q_scale, False

233 - )

234 - q = q_fp8

243 + assert self.attn.query_quant is not None

244 + q, _ = self.attn.query_quant(q, self.attn._q_scale)

245

246 # Final views + KV cache update

247 q = q.view(-1, self.num_heads, self.head_size)

248 k = k.view(-1, self.num_kv_heads, self.head_size)

249 v = v.view(-1, self.num_kv_heads, self.head_size)

250 kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(

241 - k, v, self.layer_name

251 + k, v, _encode_layer_name(self.layer_name)

252 )

253 return q, k, v, kv_cache_dummy_dep

254

@@ -260,6 +270,64 @@ class QKNormRoPEKVCacheTestModel(torch.nn.Module):

270 return [torch.ops.vllm.fused_qk_norm_rope_and_unified_kv_cache_update.default]

271

272

273 +class QKNormMRoPEKVCacheTestModel(QKNormRoPEKVCacheTestModel):

274 + def __init__(

275 + self,

276 + *,

277 + mrope_section: tuple[int, int, int],

278 + mrope_interleaved: bool,

279 + **kwargs,

280 + ):

281 + super().__init__(**kwargs)

282 + self.rotary_emb = MRotaryEmbedding(

283 + self.head_size,

284 + rotary_dim=self.rotary_dim,

285 + max_position_embeddings=4096,

286 + base=10000,

287 + is_neox_style=self.is_neox,

288 + dtype=self.dtype,

289 + mrope_section=list(mrope_section),

290 + mrope_interleaved=mrope_interleaved,

291 + )

292 + self.enable_rope_custom_op = self.rotary_emb.enabled()

293 +

294 + def forward(

295 + self, qkv: torch.Tensor, positions: torch.Tensor

296 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

297 + q, k, v, kv_cache_dummy_dep = super().forward(qkv, positions)

298 + output = torch.empty(

299 + q.shape[0],

300 + self.num_heads * self.head_size,

301 + device=q.device,

302 + dtype=self.dtype,

303 + )

304 + torch.ops.vllm.unified_attention_with_output(

305 + q,

306 + k,

307 + v,

308 + output,

309 + _encode_layer_name(self.layer_name),

310 + kv_cache_dummy_dep=kv_cache_dummy_dep,

311 + )

312 + # Production decoder attention consumes K from the cache rather than

313 + # its direct K argument. Returning `output` keeps that production user

314 + # in the graph without exposing K to an unsupported consumer.

315 + return q, output, v, kv_cache_dummy_dep

316 +

317 + def ops_in_model_before(self) -> list[torch._ops.OpOverload]:

318 + return [MROPE_OP, torch.ops.vllm.unified_kv_cache_update.default]

319 +

320 + def ops_in_model_after(self) -> list[torch._ops.OpOverload]:

321 + return [torch.ops.vllm.fused_qk_norm_mrope_and_unified_kv_cache_update.default]

322 +

323 +

324 +class QKNormMRoPEObservableKTestModel(QKNormMRoPEKVCacheTestModel):

325 + def forward(

326 + self, qkv: torch.Tensor, positions: torch.Tensor

327 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

328 + return QKNormRoPEKVCacheTestModel.forward(self, qkv, positions)

329 +

330 +

331 def _run_qk_norm_rope_kvcache_fusion_test(

332 *,

333 attn_backend: AttentionBackendEnum,

@@ -278,18 +346,33 @@ def _run_qk_norm_rope_kvcache_fusion_test(

346 rms_norm_eps: float,

347 custom_op: str,

348 monkeypatch: pytest.MonkeyPatch,

349 + mrope_section: tuple[int, int, int] | None = None,

350 + mrope_interleaved: bool = False,

351 + observe_mrope_k: bool = False,

352 + expect_fusion: bool = True,

353 + use_quark_scalar_query_scale: bool = False,

354 ) -> None:

355 device = os.environ.get("VLLM_TEST_CUDA_DEVICE", "cuda")

356 torch.set_default_device(device)

357 torch.set_default_dtype(dtype)

358 torch.manual_seed(0)

359

360 + # Set AITER toggles before VllmConfig initialization enables env caching.

361 + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")

362 + monkeypatch.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1")

363 + rocm_aiter_ops.refresh_env_variables()

364 +

365 vllm_config = VllmConfig(

366 model_config=ModelConfig(dtype=dtype),

367 cache_config=CacheConfig(

368 block_size=block_size,

369 cache_dtype=kv_cache_dtype,

370 ),

371 + quant_config=(

372 + QuarkConfig(quant_config={"exclude": []})

373 + if use_quark_scalar_query_scale

374 + else None

375 + ),

376 compilation_config=CompilationConfig(

377 mode=CompilationMode.VLLM_COMPILE,

378 custom_ops=[custom_op],

@@ -302,6 +385,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(

385

386 with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:

387 m.setenv("VLLM_ROCM_USE_AITER", "1")

388 + m.setenv("VLLM_ROCM_USE_AITER_LINEAR", "1")

389 m.setenv(

390 "VLLM_ROCM_USE_AITER_TRITON_ROPE",

391 "1" if enable_aiter_triton_rope else "0",

@@ -309,18 +393,49 @@ def _run_qk_norm_rope_kvcache_fusion_test(

393 m.setenv("VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", use_shuffle_kv_layout)

394 rocm_aiter_ops.refresh_env_variables()

395

312 - model = QKNormRoPEKVCacheTestModel(

313 - vllm_config=vllm_config,

314 - attn_backend=attn_backend,

315 - num_heads=num_heads,

316 - num_kv_heads=num_kv_heads,

317 - head_size=head_size,

318 - rotary_dim=rotary_dim,

319 - is_neox=is_neox,

320 - rms_norm_eps=rms_norm_eps,

321 - dtype=dtype,

322 - device=torch.get_default_device(),

323 - )

396 + expected_query_quant_group_shape = None

397 + if use_quark_scalar_query_scale:

398 + expected_query_quant_group_shape = (

399 + GroupShape(-1, num_heads * head_size)

400 + if num_kv_heads == 1

401 + else GroupShape.PER_TENSOR

402 + )

403 +

404 + model_kwargs = {

405 + "vllm_config": vllm_config,

406 + "attn_backend": attn_backend,

407 + "num_heads": num_heads,

408 + "num_kv_heads": num_kv_heads,

409 + "head_size": head_size,

410 + "rotary_dim": rotary_dim,

411 + "is_neox": is_neox,

412 + "rms_norm_eps": rms_norm_eps,

413 + "dtype": dtype,

414 + "device": torch.get_default_device(),

415 + "expected_query_quant_group_shape": expected_query_quant_group_shape,

416 + }

417 + if mrope_section is None:

418 + model = QKNormRoPEKVCacheTestModel(**model_kwargs)

419 + else:

420 + alternate_section = (8, 8, 8)

421 + m.setattr(

422 + fusion_module,

423 + "_discover_mrope_configs",

424 + lambda _: (

425 + (mrope_section, mrope_interleaved),

426 + (alternate_section, not mrope_interleaved),

427 + ),

428 + )

429 + model_cls = (

430 + QKNormMRoPEObservableKTestModel

431 + if observe_mrope_k

432 + else QKNormMRoPEKVCacheTestModel

433 + )

434 + model = model_cls(

435 + **model_kwargs,

436 + mrope_section=mrope_section,

437 + mrope_interleaved=mrope_interleaved,

438 + )

439

440 fusion_pass = QkNormRopeKvCacheFusionPass(vllm_config)

441 passes = [

@@ -337,7 +452,12 @@ def _run_qk_norm_rope_kvcache_fusion_test(

452 num_heads * head_size + 2 * num_kv_heads * head_size,

453 dtype=dtype,

454 )

340 - pos = torch.arange(num_tokens, dtype=torch.long)

455 + if mrope_section is None:

456 + pos = torch.arange(num_tokens, dtype=torch.long)

457 + else:

458 + # Real Qwen-VL position IDs can be a strided [3, N] view.

459 + pos = torch.randint(0, 4096, (3, 2 * num_tokens), dtype=torch.long)[:, ::2]

460 + assert not pos.is_contiguous()

461

462 qkv_unfused = qkv.clone()

463 pos_unfused = pos.clone()

@@ -356,7 +476,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(

476

477 # Run fused (compiled) forward

478 torch._dynamo.mark_dynamic(qkv, 0)

359 - torch._dynamo.mark_dynamic(pos, 0)

479 + torch._dynamo.mark_dynamic(pos, 0 if mrope_section is None else 1)

480 with set_forward_context(None, vllm_config):

481 model_fused = torch.compile(model, backend=backend)

482 forward_context = get_forward_context()

@@ -369,14 +489,33 @@ def _run_qk_norm_rope_kvcache_fusion_test(

489 kv_cache_fused = attn_layer.kv_cache

490 del dummy

491

372 - assert fusion_pass.matched_count == 1

492 + assert fusion_pass.matched_count == int(expect_fusion)

493

374 - backend.check_before_ops(model.ops_in_model_before())

375 - backend.check_after_ops(model.ops_in_model_after())

494 + if expect_fusion:

495 + backend.check_before_ops(model.ops_in_model_before())

496 + backend.check_after_ops(model.ops_in_model_after())

497 + else:

498 + for op in model.ops_in_model_before():

499 + assert backend.op_count(op, before=True) > 0

500 + assert backend.op_count(op) > 0

501 +

502 + if use_quark_scalar_query_scale:

503 + expected_quant_op = (

504 + torch.ops.vllm.rocm_aiter_per_tensor_quant.default

505 + if num_kv_heads != 1

506 + else torch.ops._C.static_scaled_fp8_quant.default

507 + )

508 + assert backend.op_count(expected_quant_op, before=True) > 0

509 + # Query quantization remains separate after the prologue fusion.

510 + assert backend.op_count(expected_quant_op) > 0

511

512 # Sweep-backed (18.2k pts, PR #42749): native-rope ref worst 7.7e-3 -> 1e-2;

513 # AITER-triton-rope ref is itself approximate (plateau 1.28e-2) -> 2e-2.

379 - ATOL, RTOL = (2e-2, 2e-2) if enable_aiter_triton_rope else (1e-2, 1e-2)

514 + ATOL, RTOL = (

515 + (2e-2, 2e-2)

516 + if enable_aiter_triton_rope or mrope_section is not None

517 + else (1e-2, 1e-2)

518 + )

519 is_fp8_cache = model.kv_cache_dtype != dtype

520

521 if q_fused.dtype == FP8_DTYPE:

@@ -387,7 +526,7 @@ def _run_qk_norm_rope_kvcache_fusion_test(

526 else:

527 torch.testing.assert_close(q_unfused, q_fused, atol=ATOL, rtol=RTOL)

528

390 - if not is_fp8_cache:

529 + if not expect_fusion or (not is_fp8_cache and mrope_section is None):

530 # The AITER PTS kernel populates k_out only for non-FP8 caches.

531 # With FP8, the kernel writes quantized K directly to the cache

532 # and may leave k_out uninitialised. In production this is fine

@@ -397,15 +536,16 @@ def _run_qk_norm_rope_kvcache_fusion_test(

536 # Should be bit exact since no processing had been done on v for both paths

537 torch.testing.assert_close(v_unfused, v_fused, atol=0.0, rtol=0.0)

538

400 - # fp8 vs triton-rope ref requires loosening tolerance to 1.25e-1.

401 - if is_fp8_cache and enable_aiter_triton_rope:

539 + # Fused and unfused arithmetic can straddle an FP8 rounding boundary.

540 + # Allow one E4M3 quantization step while still comparing every block.

541 + if is_fp8_cache:

542 cache_atol = cache_rtol = 1.25e-1

543 else:

544 cache_atol, cache_rtol = ATOL, RTOL

545

546 torch.testing.assert_close(

407 - kv_cache_unfused[0].float(),

408 - kv_cache_fused[0].float(),

547 + kv_cache_unfused.float(),

548 + kv_cache_fused.float(),

549 atol=cache_atol,

550 rtol=cache_rtol,

551 )

@@ -494,3 +634,225 @@ def test_qk_norm_rope_kvcache_fusion(

634 custom_op=custom_op,

635 monkeypatch=monkeypatch,

636 )

637 +

638 +

639 +@pytest.mark.parametrize(

640 + "mrope_section",

641 + [

642 + pytest.param((16, 24, 24), id="full_rotary"),

643 + pytest.param((4, 6, 6), id="partial_rotary"),

644 + ],

645 +)

646 +@pytest.mark.parametrize("mrope_interleaved", [False, True])

647 +@pytest.mark.parametrize("is_neox", [False, True])

648 +@pytest.mark.parametrize(

649 + "kv_layout",

650 + [

651 + pytest.param(KVCacheLayout.LBHNC, id="head_major"),

652 + pytest.param(KVCacheLayout.LBNHC, id="token_major"),

653 + pytest.param(KVCacheLayout.LHBNC, id="block_major"),

654 + ],

655 +)

656 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])

657 +@pytest.mark.skipif(

658 + not IS_AITER_FOUND,

659 + reason="Requires AITER",

660 +)

661 +def test_qk_norm_mrope_kvcache_fusion(

662 + mrope_section: tuple[int, int, int],

663 + mrope_interleaved: bool,

664 + is_neox: bool,

665 + kv_layout: KVCacheLayout,

666 + kv_cache_dtype: str,

667 + monkeypatch: pytest.MonkeyPatch,

668 +):

669 + _run_qk_norm_rope_kvcache_fusion_test(

670 + attn_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,

671 + enable_aiter_triton_rope=False,

672 + num_tokens=5,

673 + num_heads=8,

674 + num_kv_heads=2,

675 + head_size=128,

676 + rotary_dim=2 * sum(mrope_section),

677 + block_size=16,

678 + is_neox=is_neox,

679 + use_shuffle_kv_layout="0",

680 + kv_layout=kv_layout,

681 + dtype=torch.bfloat16,

682 + kv_cache_dtype=kv_cache_dtype,

683 + rms_norm_eps=1e-6,

684 + custom_op="+rotary_embedding",

685 + monkeypatch=monkeypatch,

686 + mrope_section=mrope_section,

687 + mrope_interleaved=mrope_interleaved,

688 + )

689 +

690 +

691 +@pytest.mark.parametrize(

692 + ("num_heads", "num_kv_heads"),

693 + [(32, 2), (16, 1), (8, 1)],

694 + ids=["tp2", "tp4", "tp8"],

695 +)

696 +def test_qk_norm_mrope_kvcache_fusion_quark_scalar_query_scale(

697 + num_heads: int,

698 + num_kv_heads: int,

699 + monkeypatch: pytest.MonkeyPatch,

700 +):

701 + _run_qk_norm_rope_kvcache_fusion_test(

702 + attn_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,

703 + enable_aiter_triton_rope=False,

704 + num_tokens=5,

705 + num_heads=num_heads,

706 + num_kv_heads=num_kv_heads,

707 + head_size=128,

708 + rotary_dim=128,

709 + block_size=16,

710 + is_neox=True,

711 + use_shuffle_kv_layout="0",

712 + kv_layout=KVCacheLayout.LBHNC,

713 + dtype=torch.bfloat16,

714 + kv_cache_dtype="fp8",

715 + rms_norm_eps=1e-6,

716 + custom_op="+rotary_embedding",

717 + monkeypatch=monkeypatch,

718 + mrope_section=(24, 20, 20),

719 + mrope_interleaved=True,

720 + use_quark_scalar_query_scale=True,

721 + )

722 +

723 +

724 +def test_mrope_fusion_rejects_observable_k(monkeypatch: pytest.MonkeyPatch):

725 + _run_qk_norm_rope_kvcache_fusion_test(

726 + attn_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,

727 + enable_aiter_triton_rope=False,

728 + num_tokens=5,

729 + num_heads=8,

730 + num_kv_heads=2,

731 + head_size=128,

732 + rotary_dim=128,

733 + block_size=16,

734 + is_neox=True,

735 + use_shuffle_kv_layout="0",

736 + kv_layout=KVCacheLayout.LBHNC,

737 + dtype=torch.bfloat16,

738 + kv_cache_dtype="auto",

739 + rms_norm_eps=1e-6,

740 + custom_op="+rotary_embedding",

741 + monkeypatch=monkeypatch,

742 + mrope_section=(16, 24, 24),

743 + observe_mrope_k=True,

744 + expect_fusion=False,

745 + )

746 +

747 +

748 +def test_mrope_fusion_rejects_unaligned_rotary_dim(

749 + monkeypatch: pytest.MonkeyPatch,

750 +):

751 + _run_qk_norm_rope_kvcache_fusion_test(

752 + attn_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,

753 + enable_aiter_triton_rope=False,

754 + num_tokens=5,

755 + num_heads=8,

756 + num_kv_heads=2,

757 + head_size=128,

758 + rotary_dim=6,

759 + block_size=16,

760 + is_neox=True,

761 + use_shuffle_kv_layout="0",

762 + kv_layout=KVCacheLayout.LBHNC,

763 + dtype=torch.bfloat16,

764 + kv_cache_dtype="auto",

765 + rms_norm_eps=1e-6,

766 + custom_op="+rotary_embedding",

767 + monkeypatch=monkeypatch,

768 + mrope_section=(1, 1, 1),

769 + expect_fusion=False,

770 + )

771 +

772 +

773 +@pytest.mark.skipif(not _USE_LAYERNAME, reason="requires opaque LayerName")

774 +def test_opaque_layer_name_patterns_register_once(monkeypatch: pytest.MonkeyPatch):

775 + device = os.environ.get("VLLM_TEST_CUDA_DEVICE", "cuda")

776 + vllm_config = VllmConfig(

777 + model_config=ModelConfig(dtype=torch.bfloat16),

778 + cache_config=CacheConfig(block_size=16, cache_dtype="fp8_e4m3"),

779 + compilation_config=CompilationConfig(

780 + mode=CompilationMode.VLLM_COMPILE,

781 + custom_ops=["+rotary_embedding"],

782 + pass_config=PassConfig(fuse_qk_norm_rope_kvcache=True),

783 + ),

784 + )

785 +

786 + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:

787 + m.setenv("VLLM_ROCM_USE_AITER", "1")

788 + rocm_aiter_ops.refresh_env_variables()

789 + model_kwargs = {

790 + "vllm_config": vllm_config,

791 + "attn_backend": AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,

792 + "num_heads": 8,

793 + "num_kv_heads": 2,

794 + "head_size": 128,

795 + "is_neox": True,

796 + "rms_norm_eps": 1e-6,

797 + "dtype": torch.bfloat16,

798 + "device": torch.device(device),

799 + }

800 + models = [

801 + QKNormRoPEKVCacheTestModel(

802 + **model_kwargs,

803 + prefix=f"model.layers.{index}.self_attn.attn",

804 + )

805 + for index in range(2)

806 + ]

807 +

808 + # Opaque layer names are wildcards in the pattern graph. Registering

809 + # once per physical layer raises "Duplicate pattern" on torch >= 2.11.

810 + fusion_pass = QkNormRopeKvCacheFusionPass(vllm_config)

811 +

812 + assert models

813 + assert fusion_pass.patterns

814 +

815 +

816 +def _call_fused_mrope_impl(q_out: torch.Tensor) -> None:

817 + device = q_out.device

818 + fusion_module.fused_qk_norm_mrope_and_unified_kv_cache_update_impl(

819 + q_out=q_out,

820 + qkv=torch.zeros(2, 16, device=device),

821 + positions=torch.arange(2, device=device),

822 + q_weight=torch.ones(4, device=device),

823 + k_weight=torch.ones(4, device=device),

824 + rms_norm_eps=1e-6,

825 + cos_sin_cache=torch.zeros(16, 4, device=device),

826 + is_neox=True,

827 + mrope_section_t=1,

828 + mrope_section_h=1,

829 + mrope_section_w=0,

830 + is_interleaved=False,

831 + rotary_dim=4,

832 + layer_name="model.layers.0.self_attn.attn",

833 + )

834 +

835 +

836 +def test_mrope_profiling_pass_initializes_query(monkeypatch: pytest.MonkeyPatch):

837 + monkeypatch.setattr(

838 + fusion_module,

839 + "get_attention_context",

840 + lambda _: (None, None, None, None),

841 + )

842 + q_out = torch.full((2, 8), torch.nan, device="cpu")

843 +

844 + _call_fused_mrope_impl(q_out)

845 +

846 + torch.testing.assert_close(q_out, torch.zeros_like(q_out))

847 +

848 +

849 +def test_mrope_attention_context_failure_propagates(

850 + monkeypatch: pytest.MonkeyPatch,

851 +):

852 + def fail_context(_):

853 + raise RuntimeError("missing attention context")

854 +

855 + monkeypatch.setattr(fusion_module, "get_attention_context", fail_context)

856 +

857 + with pytest.raises(RuntimeError, match="missing attention context"):

858 + _call_fused_mrope_impl(torch.empty(2, 8, device="cpu"))

vllm/_aiter_ops.py

+109

@@ -3362,6 +3362,115 @@ class rocm_aiter_ops:

3362 rotary_dim=kernel_rotary_dim,

3363 )

3364

3365 + @staticmethod

3366 + def fused_qk_norm_mrope_and_cache(

3367 + qkv: torch.Tensor,

3368 + q_weight: torch.Tensor,

3369 + k_weight: torch.Tensor,

3370 + cos_sin_cache: torch.Tensor,

3371 + positions: torch.Tensor,

3372 + num_heads_q: int,

3373 + num_heads_k: int,

3374 + head_dim: int,

3375 + is_neox: bool,

3376 + mrope_section: list[int],

3377 + is_interleaved: bool,

3378 + rms_norm_eps: float,

3379 + q_out: torch.Tensor,

3380 + k_cache: torch.Tensor,

3381 + v_cache: torch.Tensor,

3382 + slot_mapping: torch.Tensor,

3383 + k_scale: torch.Tensor,

3384 + v_scale: torch.Tensor,

3385 + block_size: int,

3386 + x: int,

3387 + rotary_dim: int = 0,

3388 + ) -> None:

3389 + from aiter.ops.fused_qk_norm_mrope_cache_quant import (

3390 + fused_qk_norm_mrope_3d_cache_pts_quant_shuffle,

3391 + )

3392 +

3393 + fused_qk_norm_mrope_3d_cache_pts_quant_shuffle(

3394 + qkv,

3395 + q_weight,

3396 + k_weight,

3397 + cos_sin_cache,

3398 + positions,

3399 + qkv.size(0),

3400 + num_heads_q,

3401 + num_heads_k,

3402 + num_heads_k,

3403 + head_dim,

3404 + is_neox,

3405 + mrope_section,

3406 + is_interleaved,

3407 + rms_norm_eps,

3408 + q_out,

3409 + k_cache,

3410 + v_cache,

3411 + slot_mapping,

3412 + k_scale,

3413 + v_scale,

3414 + None,

3415 + None,

3416 + False,

3417 + False,

3418 + block_size,

3419 + x,

3420 + rotary_dim,

3421 + )

3422 +

3423 + @staticmethod

3424 + def do_qk_norm_mrope_kvcache_update(

3425 + qkv: torch.Tensor,

3426 + q_weight: torch.Tensor,

3427 + k_weight: torch.Tensor,

3428 + cos_sin_cache: torch.Tensor,

3429 + positions: torch.Tensor,

3430 + num_heads_q: int,

3431 + num_heads_k: int,

3432 + head_dim: int,

3433 + is_neox: bool,

3434 + mrope_section: list[int],

3435 + is_interleaved: bool,

3436 + rms_norm_eps: float,

3437 + q_out: torch.Tensor,

3438 + key_cache: torch.Tensor,

3439 + value_cache: torch.Tensor,

3440 + slot_mapping: torch.Tensor,

3441 + k_scale: torch.Tensor,

3442 + v_scale: torch.Tensor,

3443 + kv_cache_dtype: str,

3444 + rotary_dim: int = 0,

3445 + ) -> None:

3446 + if kv_cache_dtype.startswith("fp8"):

3447 + key_cache = key_cache.view(current_platform.fp8_dtype())

3448 + value_cache = value_cache.view(current_platform.fp8_dtype())

3449 +

3450 + rocm_aiter_ops.fused_qk_norm_mrope_and_cache(

3451 + qkv=qkv,

3452 + q_weight=q_weight,

3453 + k_weight=k_weight,

3454 + cos_sin_cache=cos_sin_cache,

3455 + positions=positions,

3456 + num_heads_q=num_heads_q,

3457 + num_heads_k=num_heads_k,

3458 + head_dim=head_dim,

3459 + is_neox=is_neox,

3460 + mrope_section=mrope_section,

3461 + is_interleaved=is_interleaved,

3462 + rms_norm_eps=rms_norm_eps,

3463 + q_out=q_out,

3464 + k_cache=key_cache,

3465 + v_cache=value_cache,

3466 + slot_mapping=slot_mapping,

3467 + k_scale=k_scale,

3468 + v_scale=v_scale,

3469 + block_size=key_cache.shape[1],

3470 + x=16 // key_cache.element_size(),

3471 + rotary_dim=rotary_dim,

3472 + )

3473 +

3474 @staticmethod

3475 def triton_rope_and_cache(

3476 query: torch.Tensor,

vllm/compilation/passes/fusion/matcher_utils.py

+69

@@ -31,6 +31,7 @@ from vllm.platforms import current_platform

32 ROTARY_OP = torch.ops._C.rotary_embedding.default

33 FLASHINFER_ROTARY_OP = torch.ops.vllm.flashinfer_rotary_embedding.default

34 +MROPE_OP = torch.ops.vllm.mrope.default

36 QUANT_OPS: dict[QuantKey, OpOverload] = {

37 kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default, # noqa: E501

@@ -162,6 +163,74 @@ class MatcherRotaryEmbedding(MatcherCustomOp):

163 return result

164

165

166 +class MatcherMRotaryEmbedding(MatcherCustomOp):

167 + def __init__(

168 + self,

169 + is_neox: bool,

170 + head_size: int,

171 + rotary_dim: int,

172 + num_heads: int,

173 + num_kv_heads: int,

174 + mrope_section: tuple[int, int, int],

175 + mrope_interleaved: bool,

176 + enabled: bool | None = None,

177 + ) -> None:

178 + if enabled is None:

179 + enabled = RotaryEmbedding.enabled()

180 +

181 + super().__init__(enabled)

182 + self.is_neox = is_neox

183 + self.head_size = head_size

184 + self.rotary_dim = rotary_dim

185 + self.num_heads = num_heads

186 + self.num_kv_heads = num_kv_heads

187 + self.q_size = self.num_heads * self.head_size

188 + self.kv_size = self.num_kv_heads * self.head_size

189 + self.mrope_section = mrope_section

190 + self.mrope_interleaved = mrope_interleaved

191 +

192 + def inputs(self) -> list[torch.Tensor]:

193 + positions = self.empty_int64(3, 5)

194 + query = self.empty(5, self.q_size)

195 + key = self.empty(5, self.kv_size)

196 + cos_sin_cache = self.empty(4096, self.rotary_dim)

197 + return [positions, query, key, cos_sin_cache]

198 +

199 + def forward_custom(

200 + self,

201 + positions: torch.Tensor,

202 + query: torch.Tensor,

203 + key: torch.Tensor,

204 + cos_sin_cache: torch.Tensor,

205 + ) -> tuple[torch.Tensor, torch.Tensor]:

206 + result = auto_functionalized(

207 + MROPE_OP,

208 + positions=positions,

209 + query=query,

210 + key=key,

211 + cos_sin_cache=cos_sin_cache,

212 + head_size=self.head_size,

213 + rotary_dim=self.rotary_dim,

214 + mrope_section_t=self.mrope_section[0],

215 + mrope_section_h=self.mrope_section[1],

216 + mrope_section_w=self.mrope_section[2],

217 + mrope_interleaved=self.mrope_interleaved,

218 + is_neox_style=self.is_neox,

219 + )

220 + return result[1], result[2]

221 +

222 + def forward_native(

223 + self,

224 + positions: torch.Tensor,

225 + query: torch.Tensor,

226 + key: torch.Tensor,

227 + cos_sin_cache: torch.Tensor,

228 + ) -> tuple[torch.Tensor, torch.Tensor]:

229 + raise NotImplementedError(

230 + "MRoPE fusion requires the rotary_embedding custom op"

231 + )

232 +

233 +

234 class MatcherRMSNormGated(MatcherCustomOp):

235 """Matches RMSNormGated with norm_before_gate=True and group_size=None."""

236

vllm/compilation/passes/fusion/qk_norm_rope_kvcache_fusion.py

+564 -78

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import inspect

5 -from collections.abc import Callable

5 +from collections.abc import Callable, Mapping

6 from typing import ParamSpec

8 import torch

@@ -19,12 +19,19 @@ from vllm.model_executor.layers.attention.attention import (

19 Attention,

20 get_attention_context,

21 )

22 +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

23 from vllm.platforms import current_platform

23 -from vllm.utils.torch_utils import direct_register_custom_op

24 +from vllm.utils.torch_utils import (

25 + _USE_LAYERNAME,

26 + LayerNameType,

27 + _encode_layer_name,

28 + _resolve_layer_name,

29 + direct_register_custom_op,

30 +)

32 from ..inductor_pass import enable_fake_mode

33 from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass

27 -from .matcher_utils import MatcherRotaryEmbedding

34 +from .matcher_utils import MROPE_OP, MatcherMRotaryEmbedding, MatcherRotaryEmbedding

35 from .rms_quant_fusion import empty_bf16, empty_fp32, empty_i64

37 logger = init_logger(__name__)

@@ -51,8 +58,9 @@ def fused_qk_norm_rope_and_unified_kv_cache_update_impl(

58 rms_norm_eps: float,

59 cos_sin_cache: torch.Tensor,

60 is_neox: bool,

54 - layer_name: str = "",

61 + layer_name: LayerNameType,

62 ) -> torch.Tensor:

63 + layer_name = _resolve_layer_name(layer_name)

64 _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)

65 if layer_slot_mapping is not None:

66 attn_layer.impl.do_qk_norm_rope_kvcache_update(

@@ -87,7 +95,7 @@ def fused_qk_norm_rope_and_unified_kv_cache_update_fake(

95 rms_norm_eps: float,

96 cos_sin_cache: torch.Tensor,

97 is_neox: bool,

90 - layer_name: str = "",

98 + layer_name: LayerNameType,

99 ) -> torch.Tensor:

100 return torch.empty(0, device=qkv.device, dtype=qkv.dtype)

101

@@ -100,6 +108,76 @@ direct_register_custom_op(

108 )

109

110

111 +def fused_qk_norm_mrope_and_unified_kv_cache_update_impl(

112 + q_out: torch.Tensor,

113 + qkv: torch.Tensor,

114 + positions: torch.Tensor,

115 + q_weight: torch.Tensor,

116 + k_weight: torch.Tensor,

117 + rms_norm_eps: float,

118 + cos_sin_cache: torch.Tensor,

119 + is_neox: bool,

120 + mrope_section_t: int,

121 + mrope_section_h: int,

122 + mrope_section_w: int,

123 + is_interleaved: bool,

124 + rotary_dim: int,

125 + layer_name: LayerNameType,

126 +) -> torch.Tensor:

127 + layer_name = _resolve_layer_name(layer_name)

128 + _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)

129 + if layer_slot_mapping is not None:

130 + if positions.ndim == 1:

131 + positions = positions.unsqueeze(0).expand(3, -1).contiguous()

132 + attn_layer.impl.do_qk_norm_mrope_kvcache_update(

133 + attn_layer,

134 + qkv,

135 + q_out,

136 + positions,

137 + q_weight,

138 + k_weight,

139 + rms_norm_eps,

140 + cos_sin_cache,

141 + is_neox,

142 + (mrope_section_t, mrope_section_h, mrope_section_w),

143 + is_interleaved,

144 + rotary_dim,

145 + kv_cache,

146 + layer_slot_mapping,

147 + )

148 + else:

149 + q_out.zero_()

150 +

151 + return torch.empty(0, device=qkv.device, dtype=qkv.dtype)

152 +

153 +

154 +def fused_qk_norm_mrope_and_unified_kv_cache_update_fake(

155 + q_out: torch.Tensor,

156 + qkv: torch.Tensor,

157 + positions: torch.Tensor,

158 + q_weight: torch.Tensor,

159 + k_weight: torch.Tensor,

160 + rms_norm_eps: float,

161 + cos_sin_cache: torch.Tensor,

162 + is_neox: bool,

163 + mrope_section_t: int,

164 + mrope_section_h: int,

165 + mrope_section_w: int,

166 + is_interleaved: bool,

167 + rotary_dim: int,

168 + layer_name: LayerNameType,

169 +) -> torch.Tensor:

170 + return torch.empty(0, device=qkv.device, dtype=qkv.dtype)

171 +

172 +

173 +direct_register_custom_op(

174 + op_name="fused_qk_norm_mrope_and_unified_kv_cache_update",

175 + op_func=fused_qk_norm_mrope_and_unified_kv_cache_update_impl,

176 + mutates_args=["q_out"],

177 + fake_impl=fused_qk_norm_mrope_and_unified_kv_cache_update_fake,

178 +)

179 +

180 +

181 # ---------------------------------------------------------------------------

182 # Pattern: QK-norm + RoPE + unified_kv_cache_update

183 # ---------------------------------------------------------------------------

@@ -142,19 +220,27 @@ class QkNormRopeKvCachePattern:

220 self.eps = eps

221 self.is_neox = is_neox

222 self.quant_query = quant_query

223 + self.encoded_layer_name = _encode_layer_name(self.layer_name)

224 + self.query_quant_group_shape = (

225 + layer.query_quant.group_shape

226 + if layer.query_quant is not None

227 + else GroupShape.PER_TENSOR

228 + )

229

230 self.q_size = self.num_heads * self.head_size

231 self.k_size = self.num_kv_heads * self.head_size

232 self.v_size = self.num_kv_heads * self.head_size_v

233

150 - self.rope_matcher = MatcherRotaryEmbedding(

151 - is_neox=is_neox,

152 - head_size=self.head_size,

153 - num_heads=self.num_heads,

154 - num_kv_heads=self.num_kv_heads,

234 + self.rope_matcher: MatcherRotaryEmbedding | MatcherMRotaryEmbedding = (

235 + MatcherRotaryEmbedding(

236 + is_neox=is_neox,

237 + head_size=self.head_size,

238 + num_heads=self.num_heads,

239 + num_kv_heads=self.num_kv_heads,

240 + )

241 )

242

157 - def get_inputs(self) -> list[torch.Tensor]:

243 + def get_inputs(self) -> list:

244 T = 5

245 L = 4096

246 qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size)

@@ -166,8 +252,13 @@ class QkNormRopeKvCachePattern:

252 if self.quant_query:

253 q_scale = empty_fp32(1)

254 inputs += [q_scale]

255 + if _USE_LAYERNAME:

256 + inputs.append(self.encoded_layer_name)

257 return inputs

258

259 + def _get_layer_name(self, layer_name: LayerNameType | None) -> LayerNameType:

260 + return self.encoded_layer_name if layer_name is None else layer_name

261 +

262 def pattern_non_fp8_quant_query(

263 self,

264 qkv: torch.Tensor,

@@ -175,6 +266,7 @@ class QkNormRopeKvCachePattern:

266 q_weight: torch.Tensor,

267 k_weight: torch.Tensor,

268 cos_sin_cache: torch.Tensor,

269 + layer_name: LayerNameType | None = None,

270 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

271 q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)

272 q_by_head = q.view(-1, self.q_size // self.head_size, self.head_size)

@@ -190,7 +282,9 @@ class QkNormRopeKvCachePattern:

282 q_rope = q_rope.view(-1, self.num_heads, self.head_size)

283 k_rope = k_rope.view(-1, self.num_kv_heads, self.head_size)

284 v = v.view(-1, self.num_kv_heads, self.head_size_v)

193 - dummy = torch.ops.vllm.unified_kv_cache_update(k_rope, v, self.layer_name)

285 + dummy = torch.ops.vllm.unified_kv_cache_update(

286 + k_rope, v, self._get_layer_name(layer_name)

287 + )

288 return dummy, q_rope, k_rope, v

289

290 def replacement_non_fp8_quant_query(

@@ -200,6 +294,7 @@ class QkNormRopeKvCachePattern:

294 q_weight: torch.Tensor,

295 k_weight: torch.Tensor,

296 cos_sin_cache: torch.Tensor,

297 + layer_name: LayerNameType | None = None,

298 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

299 q_out = torch.empty(

300 qkv.shape[0],

@@ -228,10 +323,49 @@ class QkNormRopeKvCachePattern:

323 rms_norm_eps=self.eps,

324 cos_sin_cache=cos_sin_cache,

325 is_neox=self.is_neox,

231 - layer_name=self.layer_name,

326 + layer_name=self._get_layer_name(layer_name),

327 )

328 return results[0], results[1], results[2], v

329

330 + def _quantize_query(

331 + self,

332 + query: torch.Tensor,

333 + q_scale: torch.Tensor,

334 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

335 + if self.query_quant_group_shape.is_per_tensor():

336 + q_out = torch.empty_like(query, dtype=current_platform.fp8_dtype())

337 + q_quant = auto_functionalized(

338 + torch.ops.vllm.rocm_aiter_per_tensor_quant.default,

339 + out=q_out,

340 + x=query,

341 + scale=q_scale,

342 + is_dynamic=False,

343 + )

344 + # The AITER op's schema marks scale mutable, so preserve its

345 + # functionalized write-back as an output of the matched region.

346 + return q_quant[1], q_quant[2]

347 +

348 + # QuantFP8.forward_hip falls back to the generic static op for the

349 + # historical per-head shape (-1, block_size). This also covers the

350 + # degenerate one-local-KV-head case where a scalar scale is represented

351 + # as (-1, q_size), rather than changing Attention's classification.

352 + q_out = torch.empty(

353 + query.shape,

354 + device=query.device,

355 + dtype=current_platform.fp8_dtype(),

356 + )

357 + q_quant = auto_functionalized(

358 + torch.ops._C.static_scaled_fp8_quant.default,

359 + result=q_out,

360 + input=query,

361 + scale=q_scale,

362 + group_shape=[

363 + self.query_quant_group_shape.row,

364 + self.query_quant_group_shape.col,

365 + ],

366 + )

367 + return q_quant[1], None

368 +

369 def pattern_fp8_quant_query(

370 self,

371 qkv: torch.Tensor,

@@ -240,7 +374,8 @@ class QkNormRopeKvCachePattern:

374 k_weight: torch.Tensor,

375 cos_sin_cache: torch.Tensor,

376 q_scale: torch.Tensor,

243 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

377 + layer_name: LayerNameType | None = None,

378 + ) -> tuple[torch.Tensor, ...]:

379 q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)

380 q_by_head = q.view(-1, self.q_size // self.head_size, self.head_size)

381 q_normed = vllm.ir.ops.rms_norm(q_by_head, q_weight, self.eps)

@@ -251,24 +386,15 @@ class QkNormRopeKvCachePattern:

386 k_flat = k_normed.view(-1, self.k_size)

387

388 q_rope, k_rope = self.rope_matcher(positions, q_flat, k_flat, cos_sin_cache)

254 - # Match the quant-query op Attention.forward inserts (fp8 KV + UNIFIED).

255 - # Explicit auto_functionalized (out=[1]) keeps the quant node in the pattern.

256 - q_out = torch.empty_like(q_rope, dtype=current_platform.fp8_dtype())

257 - q_quant = auto_functionalized(

258 - torch.ops.vllm.rocm_aiter_per_tensor_quant.default,

259 - out=q_out,

260 - x=q_rope,

261 - scale=q_scale,

262 - is_dynamic=False,

263 - )

264 - # `scale` is mutable: its copy_ write-back to _q_scale bumps the mutation

265 - # region, so keep q flat (a reshape lands past the barrier and won't match).

266 - q_rope_fp8 = q_quant[1]

267 - q_scale_out = q_quant[2]

389 + q_rope_fp8, q_scale_out = self._quantize_query(q_rope, q_scale)

390

391 k_rope = k_rope.view(-1, self.num_kv_heads, self.head_size)

392 v = v.view(-1, self.num_kv_heads, self.head_size_v)

271 - dummy = torch.ops.vllm.unified_kv_cache_update(k_rope, v, self.layer_name)

393 + dummy = torch.ops.vllm.unified_kv_cache_update(

394 + k_rope, v, self._get_layer_name(layer_name)

395 + )

396 + if q_scale_out is None:

397 + return dummy, q_rope_fp8, k_rope, v

398 return dummy, q_rope_fp8, k_rope, v, q_scale_out

399

400 def replacement_fp8_quant_query(

@@ -279,7 +405,8 @@ class QkNormRopeKvCachePattern:

405 k_weight: torch.Tensor,

406 cos_sin_cache: torch.Tensor,

407 q_scale: torch.Tensor,

282 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

408 + layer_name: LayerNameType | None = None,

409 + ) -> tuple[torch.Tensor, ...]:

410 q_out = torch.empty(

411 qkv.shape[0],

412 self.num_heads,

@@ -307,22 +434,14 @@ class QkNormRopeKvCachePattern:

434 rms_norm_eps=self.eps,

435 cos_sin_cache=cos_sin_cache,

436 is_neox=self.is_neox,

310 - layer_name=self.layer_name,

437 + layer_name=self._get_layer_name(layer_name),

438 )

312 - # Re-apply the quant on the kernel's bf16 q_out; fused op does not quant q.

313 - # Same explicit auto_functionalized form as the pattern: [1] = quantized

314 - # q, [2] = scale (returned so the buffer-writeback use is preserved).

439 + # Re-apply the same quant form on the kernel's bf16 q_out; the fused

440 + # operation does not quantize Q.

441 q_fp8_flat = results[1].view(-1, self.q_size)

316 - q_fp8_out = torch.empty_like(q_fp8_flat, dtype=current_platform.fp8_dtype())

317 - q_requant = auto_functionalized(

318 - torch.ops.vllm.rocm_aiter_per_tensor_quant.default,

319 - out=q_fp8_out,

320 - x=q_fp8_flat,

321 - scale=q_scale,

322 - is_dynamic=False,

323 - )

324 - q_fp8 = q_requant[1] # flat to mirror the pattern (see note above)

325 - q_scale_out = q_requant[2]

442 + q_fp8, q_scale_out = self._quantize_query(q_fp8_flat, q_scale)

443 + if q_scale_out is None:

444 + return results[0], q_fp8, results[2], v

445 return results[0], q_fp8, results[2], v, q_scale_out

446

447 @staticmethod

@@ -345,6 +464,9 @@ class QkNormRopeKvCachePattern:

464

465 view_to_reshape(gm)

466

467 + def _extra_check(self, match: pm.Match) -> bool:

468 + return True

469 +

470 def _register(self, pattern, replacement, pm_pass) -> None:

471 trace_fn = QkNormRopeKvCachePattern.wrap_trace_fn(

472 pm.fwd_only,

@@ -370,40 +492,338 @@ class QkNormRopeKvCachePattern:

492 inputs,

493 trace_fn,

494 pm_pass,

495 + extra_check=self._extra_check,

496 search_fn_pattern=search_fn_pattern,

497 )

498

499 def register(self, pm_pass: PatternMatcherPass) -> None:

500 # make_fx counts `self` in bound-method code params; wrap as plain fns.

378 - # Distinct names per branch so mypy doesn't see one name, two signatures.

501 if self.quant_query:

380 -

381 - def pattern_q(qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale):

382 - return self.pattern_fp8_quant_query(

502 + if _USE_LAYERNAME:

503 +

504 + def pattern_q_with_layer(

505 + qkv,

506 + positions,

507 + q_weight,

508 + k_weight,

509 + cos_sin_cache,

510 + q_scale,

511 + layer_name,

512 + ):

513 + return self.pattern_fp8_quant_query(

514 + qkv,

515 + positions,

516 + q_weight,

517 + k_weight,

518 + cos_sin_cache,

519 + q_scale,

520 + layer_name,

521 + )

522 +

523 + def replacement_q_with_layer(

524 + qkv,

525 + positions,

526 + q_weight,

527 + k_weight,

528 + cos_sin_cache,

529 + q_scale,

530 + layer_name,

531 + ):

532 + return self.replacement_fp8_quant_query(

533 + qkv,

534 + positions,

535 + q_weight,

536 + k_weight,

537 + cos_sin_cache,

538 + q_scale,

539 + layer_name,

540 + )

541 +

542 + self._register(pattern_q_with_layer, replacement_q_with_layer, pm_pass)

543 + else:

544 +

545 + def pattern_q_without_layer(

546 qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale

384 - )

547 + ):

548 + return self.pattern_fp8_quant_query(

549 + qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale

550 + )

551

386 - def replacement_q(

387 - qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale

388 - ):

389 - return self.replacement_fp8_quant_query(

552 + def replacement_q_without_layer(

553 qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale

391 - )

554 + ):

555 + return self.replacement_fp8_quant_query(

556 + qkv, positions, q_weight, k_weight, cos_sin_cache, q_scale

557 + )

558

393 - self._register(pattern_q, replacement_q, pm_pass)

559 + self._register(

560 + pattern_q_without_layer, replacement_q_without_layer, pm_pass

561 + )

562 else:

563 + if _USE_LAYERNAME:

564 +

565 + def pattern_noq_with_layer(

566 + qkv,

567 + positions,

568 + q_weight,

569 + k_weight,

570 + cos_sin_cache,

571 + layer_name,

572 + ):

573 + return self.pattern_non_fp8_quant_query(

574 + qkv,

575 + positions,

576 + q_weight,

577 + k_weight,

578 + cos_sin_cache,

579 + layer_name,

580 + )

581 +

582 + def replacement_noq_with_layer(

583 + qkv,

584 + positions,

585 + q_weight,

586 + k_weight,

587 + cos_sin_cache,

588 + layer_name,

589 + ):

590 + return self.replacement_non_fp8_quant_query(

591 + qkv,

592 + positions,

593 + q_weight,

594 + k_weight,

595 + cos_sin_cache,

596 + layer_name,

597 + )

598 +

599 + self._register(

600 + pattern_noq_with_layer, replacement_noq_with_layer, pm_pass

601 + )

602 + else:

603

396 - def pattern_noq(qkv, positions, q_weight, k_weight, cos_sin_cache):

397 - return self.pattern_non_fp8_quant_query(

604 + def pattern_noq_without_layer(

605 qkv, positions, q_weight, k_weight, cos_sin_cache

399 - )

606 + ):

607 + return self.pattern_non_fp8_quant_query(

608 + qkv, positions, q_weight, k_weight, cos_sin_cache

609 + )

610

401 - def replacement_noq(qkv, positions, q_weight, k_weight, cos_sin_cache):

402 - return self.replacement_non_fp8_quant_query(

611 + def replacement_noq_without_layer(

612 qkv, positions, q_weight, k_weight, cos_sin_cache

613 + ):

614 + return self.replacement_non_fp8_quant_query(

615 + qkv, positions, q_weight, k_weight, cos_sin_cache

616 + )

617 +

618 + self._register(

619 + pattern_noq_without_layer,

620 + replacement_noq_without_layer,

621 + pm_pass,

622 )

623

406 - self._register(pattern_noq, replacement_noq, pm_pass)

624 +

625 +class QkNormMRopeKvCachePattern(QkNormRopeKvCachePattern):

626 + FUSED_OP = torch.ops.vllm.fused_qk_norm_mrope_and_unified_kv_cache_update.default

627 +

628 + def __init__(

629 + self,

630 + layer: Attention,

631 + eps: float,

632 + is_neox: bool,

633 + quant_query: bool,

634 + rotary_dim: int,

635 + mrope_section: tuple[int, int, int],

636 + is_interleaved: bool,

637 + ) -> None:

638 + super().__init__(

639 + layer=layer,

640 + eps=eps,

641 + is_neox=is_neox,

642 + quant_query=quant_query,

643 + )

644 + self.rotary_dim = rotary_dim

645 + self.mrope_section = mrope_section

646 + self.is_interleaved = is_interleaved

647 + self.rope_matcher = MatcherMRotaryEmbedding(

648 + is_neox=is_neox,

649 + head_size=self.head_size,

650 + rotary_dim=rotary_dim,

651 + num_heads=self.num_heads,

652 + num_kv_heads=self.num_kv_heads,

653 + mrope_section=mrope_section,

654 + mrope_interleaved=is_interleaved,

655 + )

656 +

657 + def get_inputs(self) -> list[torch.Tensor]:

658 + T = 5

659 + L = 4096

660 + qkv = empty_bf16(T, self.q_size + self.k_size + self.v_size)

661 + positions = empty_i64(3, T)

662 + q_weight = empty_bf16(1, self.head_size)

663 + k_weight = empty_bf16(1, self.head_size)

664 + cos_sin_cache = empty_bf16(L, self.rotary_dim)

665 + inputs = [qkv, positions, q_weight, k_weight, cos_sin_cache]

666 + if self.quant_query:

667 + inputs.append(empty_fp32(1))

668 + if _USE_LAYERNAME:

669 + inputs.append(self.encoded_layer_name)

670 + return inputs

671 +

672 + def _unfused_kv_views(self, qkv: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

673 + # The supported backend reads K/V from the cache after the fused write.

674 + # Keep view-compatible placeholders for the attention custom-op ABI.

675 + _, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)

676 + k = k.view(qkv.shape[0], self.num_kv_heads, self.head_size)

677 + v = v.view(qkv.shape[0], self.num_kv_heads, self.head_size_v)

678 + return k, v

679 +

680 + def _fused_results(

681 + self,

682 + qkv: torch.Tensor,

683 + positions: torch.Tensor,

684 + q_weight: torch.Tensor,

685 + k_weight: torch.Tensor,

686 + cos_sin_cache: torch.Tensor,

687 + layer_name: LayerNameType | None,

688 + ) -> tuple[torch.Tensor, ...]:

689 + q_out = torch.empty(

690 + qkv.shape[0],

691 + self.num_heads,

692 + self.head_size,

693 + device=qkv.device,

694 + dtype=qkv.dtype,

695 + )

696 + return auto_functionalized(

697 + self.FUSED_OP,

698 + q_out=q_out,

699 + qkv=qkv,

700 + positions=positions,

701 + q_weight=q_weight,

702 + k_weight=k_weight,

703 + rms_norm_eps=self.eps,

704 + cos_sin_cache=cos_sin_cache,

705 + is_neox=self.is_neox,

706 + mrope_section_t=self.mrope_section[0],

707 + mrope_section_h=self.mrope_section[1],

708 + mrope_section_w=self.mrope_section[2],

709 + is_interleaved=self.is_interleaved,

710 + rotary_dim=self.rotary_dim,

711 + layer_name=self._get_layer_name(layer_name),

712 + )

713 +

714 + def replacement_non_fp8_quant_query(

715 + self,

716 + qkv: torch.Tensor,

717 + positions: torch.Tensor,

718 + q_weight: torch.Tensor,

719 + k_weight: torch.Tensor,

720 + cos_sin_cache: torch.Tensor,

721 + layer_name: LayerNameType | None = None,

722 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

723 + k, v = self._unfused_kv_views(qkv)

724 + results = self._fused_results(

725 + qkv, positions, q_weight, k_weight, cos_sin_cache, layer_name

726 + )

727 + return results[0], results[1], k, v

728 +

729 + def replacement_fp8_quant_query(

730 + self,

731 + qkv: torch.Tensor,

732 + positions: torch.Tensor,

733 + q_weight: torch.Tensor,

734 + k_weight: torch.Tensor,

735 + cos_sin_cache: torch.Tensor,

736 + q_scale: torch.Tensor,

737 + layer_name: LayerNameType | None = None,

738 + ) -> tuple[torch.Tensor, ...]:

739 + k, v = self._unfused_kv_views(qkv)

740 + results = self._fused_results(

741 + qkv, positions, q_weight, k_weight, cos_sin_cache, layer_name

742 + )

743 + q_flat = results[1].view(-1, self.q_size)

744 + q_fp8, q_scale_out = self._quantize_query(q_flat, q_scale)

745 + if q_scale_out is None:

746 + return results[0], q_fp8, k, v

747 + return results[0], q_fp8, k, v, q_scale_out

748 +

749 + def _extra_check(self, match: pm.Match) -> bool:

750 + expected = {

751 + "head_size": self.head_size,

752 + "rotary_dim": self.rotary_dim,

753 + "mrope_section_t": self.mrope_section[0],

754 + "mrope_section_h": self.mrope_section[1],

755 + "mrope_section_w": self.mrope_section[2],

756 + "mrope_interleaved": self.is_interleaved,

757 + "is_neox_style": self.is_neox,

758 + }

759 + mrope_node = next(

760 + (node for node in match.nodes if node.args and node.args[0] == MROPE_OP),

761 + None,

762 + )

763 + if mrope_node is None or any(

764 + mrope_node.kwargs.get(name) != value for name, value in expected.items()

765 + ):

766 + return False

767 +

768 + # The fused AITER MRoPE entrypoint writes rotated K directly to the

769 + # cache and does not expose a BF16 K result. Returning the raw K view is

770 + # valid only for decoder unified attention, whose backend consumes K

771 + # from that cache. Reject graphs that observe K in any other way.

772 + output_nodes = match.output_nodes()

773 + if len(output_nodes) < 3 or output_nodes[2] is None:

774 + return False

775 + matched_nodes = set(match.nodes)

776 + attention_op = torch.ops.vllm.unified_attention_with_output.default

777 +

778 + def is_unified_attention_user(user: fx.Node) -> bool:

779 + return user.op == "call_function" and (

780 + user.target == attention_op

781 + or (user.args and user.args[0] == attention_op)

782 + )

783 +

784 + return all(

785 + user in matched_nodes or is_unified_attention_user(user)

786 + for user in output_nodes[2].users

787 + )

788 +

789 +

790 +def _discover_mrope_configs(

791 + config: VllmConfig,

792 +) -> tuple[tuple[tuple[int, int, int], bool], ...]:

793 + model_config = config.model_config

794 + hf_configs = (

795 + getattr(model_config, "hf_text_config", None),

796 + getattr(model_config, "hf_config", None),

797 + )

798 + configs: set[tuple[tuple[int, int, int], bool]] = set()

799 + for hf_config in hf_configs:

800 + if hf_config is None:

801 + continue

802 + rope_parameters = getattr(hf_config, "rope_parameters", None)

803 + if rope_parameters is None:

804 + rope_parameters = getattr(hf_config, "rope_scaling", None)

805 + if not isinstance(rope_parameters, Mapping):

806 + continue

807 + section = rope_parameters.get("mrope_section")

808 + if not isinstance(section, (list, tuple)) or len(section) != 3:

809 + continue

810 + try:

811 + parsed_section = (

812 + int(section[0]),

813 + int(section[1]),

814 + int(section[2]),

815 + )

816 + except (TypeError, ValueError):

817 + continue

818 + if any(value <= 0 for value in parsed_section):

819 + continue

820 + configs.add(

821 + (

822 + parsed_section,

823 + bool(rope_parameters.get("mrope_interleaved", False)),

824 + )

825 + )

826 + return tuple(sorted(configs))

827

828

829 # ---------------------------------------------------------------------------

@@ -412,7 +832,7 @@ class QkNormRopeKvCachePattern:

832

833

834 class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):

415 - """Fuse QK-norm + RoPE + KV cache update into a single AITER HIP kernel.

835 + """Fuse QK-norm + RoPE/MRoPE + KV cache update into an AITER HIP kernel.

836

837 Supersedes both QKNormRoPEFusionPass and RopeKVCacheFusionPass for

838 attention layers that support the combined operation, eliminating two

@@ -429,22 +849,33 @@ class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):

849

850 cc = config.compilation_config

851 self.max_token_num = cc.pass_config.rope_kvcache_fusion_max_token_num

852 + self.mrope_configs = (

853 + _discover_mrope_configs(config)

854 + if cc.is_custom_op_enabled("rotary_embedding")

855 + else ()

856 + )

857

858 dtype = config.model_config.dtype

859 if dtype not in (torch.bfloat16, torch.float16):

860 logger.warning_once(

436 - "QK Norm+RoPE+KVCache fusion not enabled: unsupported dtype %s", dtype

861 + "QK Norm+RoPE/MRoPE+KVCache fusion not enabled: unsupported dtype %s",

862 + dtype,

863 )

864 return

865

866 attn_layers = get_layers_from_vllm_config(config, Attention)

867

868 for _, layer in attn_layers.items():

443 - if not layer.impl.fused_qk_norm_rope_kvcache_supported():

869 + supports_rope = layer.impl.fused_qk_norm_rope_kvcache_supported()

870 + supports_mrope = (

871 + bool(self.mrope_configs)

872 + and layer.impl.fused_qk_norm_mrope_kvcache_supported()

873 + )

874 + if not supports_rope and not supports_mrope:

875 continue

876 if layer.head_size not in SUPPORTED_FUSED_QK_NORM_ROPE_KVCACHE_HEAD_DIMS:

877 logger.warning_once(

447 - "QK Norm+RoPE+KVCache fusion not enabled for a layer: "

878 + "QK Norm+RoPE/MRoPE+KVCache fusion not enabled for a layer: "

879 "head_size=%d is not supported by the "

880 "fused_qk_norm_rope_cache_pts_quant_shuffle kernel "

881 "(supported: %s). Falling back to the unfused path.",

@@ -455,22 +886,71 @@ class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):

886 if layer.head_size_v != layer.head_size:

887 # The fused kernel uses a single head_dim for q/k/v.

888 logger.warning_once(

458 - "QK Norm+RoPE+KVCache fusion not enabled for a layer: "

889 + "QK Norm+RoPE/MRoPE+KVCache fusion not enabled for a layer: "

890 "head_size_v=%d differs from head_size=%d, which the fused "

891 "kernel does not support. Falling back to the unfused path.",

892 layer.head_size_v,

893 layer.head_size,

894 )

895 continue

465 - for epsilon in [1e-5, 1e-6]:

466 - for neox in [True, False]:

467 - for quant_q in [False, True]:

468 - QkNormRopeKvCachePattern(

469 - layer=layer,

470 - eps=epsilon,

471 - is_neox=neox,

472 - quant_query=quant_q,

473 - ).register(self.patterns)

896 + if supports_rope:

897 + for epsilon in [1e-5, 1e-6]:

898 + for neox in [True, False]:

899 + for quant_q in [False, True]:

900 + QkNormRopeKvCachePattern(

901 + layer=layer,

902 + eps=epsilon,

903 + is_neox=neox,

904 + quant_query=quant_q,

905 + ).register(self.patterns)

906 +

907 + if supports_mrope:

908 + for section, is_interleaved in self.mrope_configs:

909 + rotary_dim = 2 * sum(section)

910 + if rotary_dim > layer.head_size:

911 + logger.warning_once(

912 + "QK Norm+MRoPE+KVCache fusion not enabled for a "

913 + "layer: rotary_dim=%d exceeds head_size=%d.",

914 + rotary_dim,

915 + layer.head_size,

916 + )

917 + continue

918 + vector_width = layer.head_size // 32

919 + for epsilon in [1e-5, 1e-6]:

920 + for neox in [True, False]:

921 + # Each warp lane processes `vector_width` adjacent

922 + # values. NeoX additionally exchanges vectors

923 + # across the half-rotary boundary.

924 + rotary_alignment = vector_width * (2 if neox else 1)

925 + if rotary_dim % rotary_alignment != 0:

926 + logger.warning_once(

927 + "QK Norm+MRoPE+KVCache fusion not enabled "

928 + "for a layer/style: rotary_dim=%d must be "

929 + "divisible by %d for head_size=%d and "

930 + "is_neox=%s.",

931 + rotary_dim,

932 + rotary_alignment,

933 + layer.head_size,

934 + neox,

935 + )

936 + continue

937 + for quant_q in [False, True]:

938 + QkNormMRopeKvCachePattern(

939 + layer=layer,

940 + eps=epsilon,

941 + is_neox=neox,

942 + quant_query=quant_q,

943 + rotary_dim=rotary_dim,

944 + mrope_section=section,

945 + is_interleaved=is_interleaved,

946 + ).register(self.patterns)

947 +

948 + # Opaque LayerName is a pattern wildcard on torch >= 2.11, so

949 + # homogeneous attention layers produce duplicate search patterns.

950 + # As in the other LayerName-aware fusion passes, one supported

951 + # layer registers all shape/style variants for the model.

952 + if _USE_LAYERNAME:

953 + break

954

955 self.dump_patterns(config, self.patterns)

956

@@ -478,8 +958,8 @@ class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):

958 def __call__(self, graph: fx.Graph) -> None:

959 self.matched_count = self.patterns.apply(graph)

960 logger.info(

481 - "QK-Norm+RoPE+KVCache fusion: replaced %s pattern(s) "

482 - "with AITER fused_qk_norm_rope_cache_pts_quant_shuffle",

961 + "QK-Norm+RoPE/MRoPE+KVCache fusion: replaced %s pattern(s) "

962 + "with an AITER fused attention-prologue kernel",

963 self.matched_count,

964 )

965

@@ -487,4 +967,10 @@ class QkNormRopeKvCacheFusionPass(VllmPatternMatcherPass):

967 return compile_range.end <= self.max_token_num

968

969 def uuid(self) -> str:

490 - return VllmInductorPass.hash_source(self, QkNormRopeKvCachePattern)

970 + return VllmInductorPass.hash_source(

971 + self,

972 + QkNormRopeKvCachePattern,

973 + QkNormMRopeKvCachePattern,

974 + _discover_mrope_configs,

975 + repr(self.mrope_configs),

976 + )

vllm/config/compilation.py

+4 -4

@@ -154,15 +154,15 @@ class PassConfig:

154 fuse_rope_kvcache: bool = None # type: ignore[assignment]

155 """Fuse the QK rope + KV cache ops."""

156 fuse_qk_norm_rope_kvcache: bool = Field(default=None) # type: ignore[assignment]

157 - """Fuse QK RMSNorm + RoPE + KV cache update into a single AITER HIP

157 + """Fuse QK RMSNorm + RoPE/MRoPE + KV cache update into an AITER HIP

158 kernel. Supersedes both enable_qk_norm_rope_fusion and fuse_rope_kvcache

159 - for layers that support it. Auto-enabled at O1+ on ROCm for models

160 - with QK-norm (e.g. Qwen3-MoE)."""

159 + for layers that support it. Auto-enabled at O2+ on ROCm for models

160 + with QK-norm (e.g. Qwen3-MoE and Qwen3-VL-class architectures)."""

161

162 rope_kvcache_fusion_max_token_num: int = 256

163 """The threshold for ROCm AITER RoPE+KVCache fusion e.g. for small batch decode.

164 Larger batch sizes e.g. during prefill will use the unfused kernels.

165 - Also applies to the fused QK-Norm+RoPE+KVCache pass.

165 + Also applies to the fused QK-Norm+RoPE/MRoPE+KVCache pass.

166 """

167

168 fi_allreduce_fusion_max_size_mb: float | None = None

vllm/config/vllm.py

+1 -1

@@ -242,7 +242,7 @@ def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool:

242

243

244 def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool:

245 - """Enable fused QK-norm + RoPE + KV cache update on ROCm with AITER."""

245 + """Enable fused QK-norm + RoPE/MRoPE + KV cache update with AITER."""

246 from vllm._aiter_ops import rocm_aiter_ops

247

248 if not rocm_aiter_ops.is_enabled():

vllm/model_executor/layers/rotary_embedding/mrope.py

+86

@@ -7,6 +7,7 @@ import torch

8 from vllm.platforms import current_platform

9 from vllm.triton_utils import tl, triton

10 +from vllm.utils.torch_utils import direct_register_custom_op

12 from .base import RotaryEmbeddingBase

13 from .yarn_scaling_rope import YaRNScalingRotaryEmbedding, yarn_get_mscale

@@ -240,6 +241,62 @@ def apply_interleaved_rope(x: torch.Tensor, mrope_section: list[int]) -> torch.T

241 return torch.where(is_width, x[2], result)

242

243

244 +def _mrope_apply(

245 + positions: torch.Tensor,

246 + query: torch.Tensor,

247 + key: torch.Tensor,

248 + cos_sin_cache: torch.Tensor,

249 + head_size: int,

250 + rotary_dim: int,

251 + mrope_section_t: int,

252 + mrope_section_h: int,

253 + mrope_section_w: int,

254 + mrope_interleaved: bool,

255 + is_neox_style: bool,

256 +) -> None:

257 + cos_sin = cos_sin_cache[positions]

258 + cos, sin = cos_sin.chunk(2, dim=-1)

259 + query_out, key_out = triton_mrope(

260 + query,

261 + key,

262 + cos,

263 + sin,

264 + [mrope_section_t, mrope_section_h, mrope_section_w],

265 + head_size,

266 + rotary_dim,

267 + mrope_interleaved,

268 + is_neox_style,

269 + )

270 + if query_out.data_ptr() != query.data_ptr():

271 + query.copy_(query_out.view_as(query))

272 + if key_out.data_ptr() != key.data_ptr():

273 + key.copy_(key_out.view_as(key))

274 +

275 +

276 +def _mrope_apply_fake(

277 + positions: torch.Tensor,

278 + query: torch.Tensor,

279 + key: torch.Tensor,

280 + cos_sin_cache: torch.Tensor,

281 + head_size: int,

282 + rotary_dim: int,

283 + mrope_section_t: int,

284 + mrope_section_h: int,

285 + mrope_section_w: int,

286 + mrope_interleaved: bool,

287 + is_neox_style: bool,

288 +) -> None:

289 + return

290 +

291 +

292 +direct_register_custom_op(

293 + op_name="mrope",

294 + op_func=_mrope_apply,

295 + mutates_args=["query", "key"],

296 + fake_impl=_mrope_apply_fake,

297 +)

298 +

299 +

300 class MRotaryEmbedding(RotaryEmbeddingBase):

301 """Rotary Embedding with Multimodal Sections."""

302

@@ -426,6 +483,35 @@ class MRotaryEmbedding(RotaryEmbeddingBase):

483 key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)

484 return query, key

485

486 + def forward_hip(

487 + self,

488 + positions: torch.Tensor,

489 + query: torch.Tensor,

490 + key: torch.Tensor | None = None,

491 + offsets: torch.Tensor | None = None,

492 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

493 + if positions.ndim == 2:

494 + assert key is not None

495 + assert self.mrope_section

496 + query_shape = query.shape

497 + key_shape = key.shape

498 + cos_sin_cache = self._match_cos_sin_cache_dtype(query)

499 + torch.ops.vllm.mrope(

500 + positions,

501 + query,

502 + key,

503 + cos_sin_cache,

504 + self.head_size,

505 + self.rotary_dim,

506 + self.mrope_section[0],

507 + self.mrope_section[1],

508 + self.mrope_section[2],

509 + self.mrope_interleaved,

510 + self.is_neox_style,

511 + )

512 + return query.reshape(query_shape), key.reshape(key_shape)

513 + return super().forward_hip(positions, query, key)

514 +

515 def forward_cpu(

516 self,

517 positions: torch.Tensor,

vllm/v1/attention/backend.py

+24

@@ -926,6 +926,10 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):

926 """

927 return False

928

929 + def fused_qk_norm_mrope_kvcache_supported(self):

930 + """Whether this implementation supports fused QKNorm+MRoPE+KVCache."""

931 + return False

932 +

933 def fused_rope_kvcache_supported(self):

934 """Does this attention implementation support RoPE+KVCache fusion.

935 This is used by the RopeKVCacheFusionPass to only fuse the RoPE ops

@@ -955,6 +959,26 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):

959 """

960 raise NotImplementedError

961

962 + def do_qk_norm_mrope_kvcache_update(

963 + self,

964 + layer: AttentionLayer,

965 + qkv: torch.Tensor,

966 + q_out: torch.Tensor,

967 + positions: torch.Tensor,

968 + q_weight: torch.Tensor,

969 + k_weight: torch.Tensor,

970 + rms_norm_eps: float,

971 + cos_sin_cache: torch.Tensor,

972 + is_neox: bool,

973 + mrope_section: tuple[int, int, int],

974 + is_interleaved: bool,

975 + rotary_dim: int,

976 + kv_cache: torch.Tensor,

977 + layer_slot_mapping: torch.Tensor,

978 + ):

979 + """Apply QK-norm and MRoPE, then write K/V to the cache."""

980 + raise NotImplementedError

981 +

982 def do_rope_and_kv_cache_update(

983 self,

984 layer: AttentionLayer,

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

+45 -1

@@ -7,7 +7,7 @@ from typing import ClassVar

7 import torch

9 from vllm import _custom_ops as ops

10 -from vllm._aiter_ops import rocm_aiter_ops

10 +from vllm._aiter_ops import IS_AITER_FOUND, rocm_aiter_ops

11 from vllm.config import get_current_vllm_config_or_none

12 from vllm.config.cache import CacheDType

13 from vllm.logger import init_logger

@@ -359,6 +359,9 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

359 def fused_qk_norm_rope_kvcache_supported(self):

360 return rocm_aiter_ops.is_enabled()

361

362 + def fused_qk_norm_mrope_kvcache_supported(self) -> bool:

363 + return IS_AITER_FOUND

364 +

365 def do_qk_norm_rope_kvcache_update(

366 self,

367 layer: AttentionLayer,

@@ -397,6 +400,47 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

400 use_shuffle_layout=False,

401 )

402

403 + def do_qk_norm_mrope_kvcache_update(

404 + self,

405 + layer: AttentionLayer,

406 + qkv: torch.Tensor,

407 + q_out: torch.Tensor,

408 + positions: torch.Tensor,

409 + q_weight: torch.Tensor,

410 + k_weight: torch.Tensor,

411 + rms_norm_eps: float,

412 + cos_sin_cache: torch.Tensor,

413 + is_neox: bool,

414 + mrope_section: tuple[int, int, int],

415 + is_interleaved: bool,

416 + rotary_dim: int,

417 + kv_cache: torch.Tensor,

418 + layer_slot_mapping: torch.Tensor,

419 + ):

420 + key_cache, value_cache = self._split_kv_cache(kv_cache)

421 + rocm_aiter_ops.do_qk_norm_mrope_kvcache_update(

422 + qkv=qkv,

423 + q_weight=q_weight,

424 + k_weight=k_weight,

425 + cos_sin_cache=cos_sin_cache,

426 + positions=positions,

427 + num_heads_q=self.num_heads,

428 + num_heads_k=self.num_kv_heads,

429 + head_dim=self.head_size,

430 + is_neox=is_neox,

431 + mrope_section=list(mrope_section),

432 + is_interleaved=is_interleaved,

433 + rms_norm_eps=rms_norm_eps,

434 + q_out=q_out,

435 + key_cache=key_cache,

436 + value_cache=value_cache,

437 + slot_mapping=layer_slot_mapping,

438 + k_scale=layer._k_scale_cpu,

439 + v_scale=layer._v_scale_cpu,

440 + kv_cache_dtype=self.kv_cache_dtype,

441 + rotary_dim=rotary_dim,

442 + )

443 +

444 def do_rope_and_kv_cache_update(

445 self,

446 layer: AttentionLayer,