@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 32ad1400d7

[ROCm][Perf] Fuse SWA q/kv RMSNorm and q FP8 group quant for DeepSeek-V4 (#53540)

Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>

Shanshan Shen committed Aug 28, 2026 at 00:37 UTC 32ad1400d7fa08872430de3f1d0f981dc8db5cb6
4 files changed +333 -12
tests/kernels/core/test_rocm_aiter_ops.py
+76 -1
@@ -24,7 +24,7 @@ import pytest
24 import torch
25
26 import vllm._aiter_ops # noqa: F401 - ensure ops are registered
27 -from tests.kernels.utils import bf16_ulp_distance
27 +from tests.kernels.utils import bf16_ulp_distance, fp8_ulp_distance
28 from vllm.platforms import current_platform
29 from vllm.utils.torch_utils import set_random_seed
30
@@ -531,6 +531,81 @@ def test_rocm_aiter_rmsnorm_with_add_fp8_group_quant_residual_accuracy():
531 )
532
533
534 +@pytest.mark.parametrize("m", [1, 32])
535 +def test_rocm_aiter_mla_dual_rmsnorm_group_quant_vs_sequential(m):
536 + """Fused q/kv RMSNorm + q group-quant matches rms_norm -> group_fp8_quant.
537 +
538 + M=1 covers the single-token decode case. M=0 is guarded by the caller
539 + (the ROCm DeepSeek-V4 attention checks num_tokens > 0) and the op has
540 + no zero-row guard itself, so it is intentionally not tested here.
541 + """
542 + require_aiter()
543 + require_fp8()
544 + from vllm._aiter_ops import rocm_aiter_ops
545 +
546 + nq, nk, group_size = 1536, 576, 128
547 + q = torch.randn(m, nq, dtype=torch.bfloat16)
548 + kv = torch.randn(m, nk, dtype=torch.bfloat16)
549 + q_weight = torch.ones(nq, dtype=torch.bfloat16)
550 + kv_weight = torch.ones(nk, dtype=torch.bfloat16)
551 + eps = 1e-5
552 + fp8_dtype = current_platform.fp8_dtype()
553 +
554 + fused_q, fused_scales, fused_kv = rocm_aiter_ops.fused_qk_rmsnorm_group_quant(
555 + q, q_weight, eps, kv, kv_weight, eps, group_size=group_size
556 + )
557 + _assert_quant_shapes(fused_q, fused_scales, m, nq, fp8_dtype, nq // group_size)
558 + assert fused_kv.shape == (m, nk)
559 + assert fused_kv.dtype == kv.dtype
560 +
561 + # q: dequantized comparison against the sequential composition.
562 + normed_q = rocm_aiter_ops.rms_norm(q, q_weight, eps)
563 + ref_q, ref_scales = rocm_aiter_ops.group_fp8_quant(normed_q, group_size)
564 + _assert_rel_error_quality(
565 + _dequantize_grouped(fused_q, fused_scales, group_size),
566 + _dequantize_grouped(ref_q, ref_scales, group_size),
567 + label="mla_dual_rmsnorm_group_quant",
568 + mean_limit=0.05,
569 + preferred_rel=0.05,
570 + max_rel=0.5,
571 + max_fail_rate=0.01,
572 + )
573 +
574 + # kv: both sides are aiter RMSNorm kernels; allow one bf16 ULP.
575 + ref_kv = rocm_aiter_ops.rms_norm(kv, kv_weight, eps)
576 + max_ulp = int(bf16_ulp_distance(fused_kv, ref_kv).max().item())
577 + assert max_ulp <= 1, f"fused kv: max ULP {max_ulp} > 1"
578 +
579 +
580 +def test_rocm_aiter_mla_dual_rmsnorm_group_quant_strided_split_view():
581 + """Split-view inputs, as produced by the DeepSeek-V4 q/kv split, match
582 + contiguous inputs."""
583 + require_aiter()
584 + require_fp8()
585 + from vllm._aiter_ops import rocm_aiter_ops
586 +
587 + m, nq, nk, group_size = 16, 1536, 576, 128
588 + qr_kv = torch.randn(m, nq + nk, dtype=torch.bfloat16)
589 + q_view, kv_view = qr_kv.split([nq, nk], dim=-1)
590 + assert not q_view.is_contiguous()
591 + q_weight, kv_weight = torch.ones(nq + nk, dtype=torch.bfloat16).split([nq, nk])
592 + eps = 1e-5
593 +
594 + def run(q, kv):
595 + return rocm_aiter_ops.fused_qk_rmsnorm_group_quant(
596 + q, q_weight, eps, kv, kv_weight, eps, group_size=group_size
597 + )
598 +
599 + out_strided = run(q_view, kv_view)
600 + out_contig = run(q_view.contiguous(), kv_view.contiguous())
601 +
602 + q_ulp = int(fp8_ulp_distance(out_strided[0], out_contig[0]).max().item())
603 + kv_ulp = int(bf16_ulp_distance(out_strided[2], out_contig[2]).max().item())
604 + assert q_ulp <= 1, f"strided q: max fp8 ULP {q_ulp} > 1"
605 + assert kv_ulp <= 1, f"strided kv: max bf16 ULP {kv_ulp} > 1"
606 + torch.testing.assert_close(out_strided[1], out_contig[1])
607 +
608 +
609 # -- End-to-end inference chain test ---------------------------------------
610
611
vllm/_aiter_ops.py
+94
@@ -1455,6 +1455,70 @@ def _fused_mla_dual_rms_norm_per_token_quant_fake(
1455 return q_out, q_scale, kv_normed
1456
1457
1458 +def _fused_mla_dual_rms_norm_group_quant_impl(
1459 + q: torch.Tensor,
1460 + q_weight: torch.Tensor,
1461 + kv: torch.Tensor,
1462 + kv_weight: torch.Tensor,
1463 + q_epsilon: float,
1464 + kv_epsilon: float,
1465 + group_size: int,
1466 + transpose_scale: bool,
1467 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1468 + """Fused MLA q/kv RMSNorm (+ FP8 per-1x128 group quant on q) via AITER.
1469 +
1470 + Backs the ``fused_mla_dual_rms_norm_group_quant`` custom op used by the
1471 + DeepSeek-V4 ROCm attention path when the q latent is group-quantized
1472 + (``(M, N/128)`` scales) for the block-scaled FP8 wq_b GEMMs. The q latent
1473 + is quantized directly from the fp32 norm accumulator, skipping the bf16
1474 + intermediate the standalone quant path would re-read; the kv latent is
1475 + RMS-normed to bf16 and consumed by the fused insert kernel unchanged.
1476 + """
1477 + from aiter.ops.fused_qk_rmsnorm_group_quant import (
1478 + fused_qk_rmsnorm_group_quant,
1479 + )
1480 +
1481 + mq, nq = q.shape
1482 + q_out = torch.empty((mq, nq), dtype=FP8_DTYPE, device=q.device)
1483 + q_scale = torch.empty((mq, nq // group_size), dtype=torch.float32, device=q.device)
1484 + kv_normed = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
1485 +
1486 + # q -> RMSNorm + FP8 group quant (q slot); kv -> RMSNorm only (k slot).
1487 + # `split` views are accepted directly (unit inner stride); the kernel
1488 + # handles strided inputs, matching the aiter op-test usage.
1489 + fused_qk_rmsnorm_group_quant(
1490 + q_out_quantized=q_out,
1491 + q_out_scale=q_scale,
1492 + q=q,
1493 + q_weight=q_weight,
1494 + q_epsilon=q_epsilon,
1495 + k_out=kv_normed,
1496 + k=kv,
1497 + k_weight=kv_weight,
1498 + k_epsilon=kv_epsilon,
1499 + group_size=group_size,
1500 + transpose_scale=transpose_scale,
1501 + )
1502 + return q_out, q_scale, kv_normed
1503 +
1504 +
1505 +def _fused_mla_dual_rms_norm_group_quant_fake(
1506 + q: torch.Tensor,
1507 + q_weight: torch.Tensor,
1508 + kv: torch.Tensor,
1509 + kv_weight: torch.Tensor,
1510 + q_epsilon: float,
1511 + kv_epsilon: float,
1512 + group_size: int,
1513 + transpose_scale: bool,
1514 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1515 + mq, nq = q.shape
1516 + q_out = torch.empty((mq, nq), dtype=FP8_DTYPE, device=q.device)
1517 + q_scale = torch.empty((mq, nq // group_size), dtype=torch.float32, device=q.device)
1518 + kv_normed = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)
1519 + return q_out, q_scale, kv_normed
1520 +
1521 +
1522 def _rocm_aiter_gemm_a8wfp4_impl(
1523 x: torch.Tensor,
1524 w: torch.Tensor,
@@ -2277,6 +2341,13 @@ class rocm_aiter_ops:
2341 fake_impl=_fused_mla_dual_rms_norm_per_token_quant_fake,
2342 )
2343
2344 + direct_register_custom_op(
2345 + op_name="fused_mla_dual_rms_norm_group_quant",
2346 + op_func=_fused_mla_dual_rms_norm_group_quant_impl,
2347 + mutates_args=[],
2348 + fake_impl=_fused_mla_dual_rms_norm_group_quant_fake,
2349 + )
2350 +
2351 _OPS_REGISTERED = True
2352
2353 @staticmethod
@@ -2340,6 +2411,29 @@ class rocm_aiter_ops:
2411 def get_fused_mla_dual_rms_norm_per_token_quant_op() -> OpOverload:
2412 return torch.ops.vllm.fused_mla_dual_rms_norm_per_token_quant.default
2413
2414 + @staticmethod
2415 + def fused_qk_rmsnorm_group_quant(
2416 + q: torch.Tensor,
2417 + q_weight: torch.Tensor,
2418 + q_epsilon: float,
2419 + kv: torch.Tensor,
2420 + kv_weight: torch.Tensor,
2421 + kv_epsilon: float,
2422 + group_size: int = 128,
2423 + transpose_scale: bool = False,
2424 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
2425 + """Fused q/kv RMSNorm + per-1x128 FP8 group quant on q via AITER."""
2426 + return torch.ops.vllm.fused_mla_dual_rms_norm_group_quant(
2427 + q,
2428 + q_weight,
2429 + kv,
2430 + kv_weight,
2431 + q_epsilon,
2432 + kv_epsilon,
2433 + group_size,
2434 + transpose_scale,
2435 + )
2436 +
2437 @staticmethod
2438 def rms_norm(
2439 x: torch.Tensor,
vllm/models/deepseek_v4/amd/rocm.py
+96
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import functools
5 from dataclasses import dataclass
6 from typing import cast
7
@@ -60,6 +61,36 @@ def _build_indptr_from_lengths(lengths: torch.Tensor) -> torch.Tensor:
61 return indptr
62
63
64 +def apply_pre_quantized_block_scaled_mm(
65 + linear: torch.nn.Module,
66 + x_fp8: torch.Tensor,
67 + x_scale: torch.Tensor,
68 +) -> torch.Tensor:
69 + """Block-scaled fp8 GEMM on pre-quantized activations.
70 +
71 + The fused q/kv norm kernel writes fp8 qr + per-1x128 scales; this
72 + drives the linear's block-scaled GEMM directly with them, bypassing
73 + apply_weights which would re-quantize the fp8 input. Only valid for
74 + the wq_b-style column/replicated linears: their output is the local
75 + TP shard, so no all-reduce is needed.
76 + """
77 + from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel import (
78 + FP8BlockParams,
79 + )
80 +
81 + params = FP8BlockParams.from_layer(linear)
82 + weight_scale = (
83 + params.weight_scale
84 + if params.weight_scale_inv is None
85 + else params.weight_scale_inv
86 + )
87 + kernel = linear.quant_method.fp8_linear
88 + out = kernel.apply_block_scaled_mm(
89 + A=x_fp8, B=params.weight, As=x_scale, Bs=weight_scale
90 + )
91 + return out.to(dtype=kernel.config.out_dtype)
92 +
93 +
94 # ROCm sparse prefill keeps this dense combine local so AMD-specific SWA changes
95 # do not touch the shared DeepSeek V4 cache utilities.
96 _SPARSE_PREFILL_TOPK_ALIGNMENT = 128
@@ -621,6 +652,71 @@ class DeepseekV4ROCMAiterMLAAttention(DeepseekV4Attention):
652 indexer_weights, _ = indexer.weights_proj(hidden_states)
653 return qr_kv, kv_score, indexer_kv_score, indexer_weights
654
655 + @functools.cached_property
656 + def _wq_b_uses_aiter_block_scaled(self) -> bool:
657 + """True when both wq_b GEMMs run the aiter block-scaled fp8 kernel.
658 +
659 + Cached: the linear kernels and the aiter env gates are fixed once
660 + the model is built, so this is evaluated at the first forward
661 + only.
662 +
663 + The fused norm+quant path is only valid if the quant and GEMM it
664 + replaces are exactly the aiter ones; otherwise fall back to the
665 + shared path.
666 + """
667 + from vllm._aiter_ops import rocm_aiter_ops
668 + from vllm.model_executor.kernels.linear.scaled_mm import (
669 + Fp8BlockScaledMMLinearKernel,
670 + )
671 +
672 + if not rocm_aiter_ops.is_linear_fp8_enabled():
673 + return False
674 +
675 + linears = [self.wq_b]
676 + if self.indexer is not None:
677 + linears.append(self.indexer.wq_b)
678 + for linear in linears:
679 + kernel = getattr(getattr(linear, "quant_method", None), "fp8_linear", None)
680 + if not isinstance(kernel, Fp8BlockScaledMMLinearKernel):
681 + return False
682 + return True
683 +
684 + def _split_qkv_and_norm(
685 + self, qr_kv: torch.Tensor
686 + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
687 + """Fuse q/kv RMSNorm + per-1x128 fp8 q quant into one aiter kernel.
688 +
689 + The shared path norms q and kv in one triton kernel and the wq_b
690 + linears then re-read the bf16 qr to quantize it. The aiter kernel
691 + computes both RMSNorms (fp32 accumulate) and the fp8 group quant
692 + in a single pass, writing fp8 qr + group scales directly; both
693 + wq_b GEMMs (attention and indexer) then consume that pair and
694 + skip their own input quant. kv stays bf16: the fused insert
695 + kernel RoPE/quantizes it itself. Falls back to the shared path
696 + when the aiter linear path is not active.
697 + """
698 + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
699 + if not (
700 + qr.dim() == 2
701 + and qr.shape[0] > 0
702 + and self.q_lora_rank % 128 == 0
703 + and self._wq_b_uses_aiter_block_scaled
704 + ):
705 + return super()._split_qkv_and_norm(qr_kv)
706 +
707 + from vllm._aiter_ops import rocm_aiter_ops
708 +
709 + return rocm_aiter_ops.fused_qk_rmsnorm_group_quant(
710 + q=qr,
711 + q_weight=self.q_norm.weight.data,
712 + q_epsilon=self.eps,
713 + kv=kv,
714 + kv_weight=self.kv_norm.weight.data,
715 + kv_epsilon=self.eps,
716 + group_size=128,
717 + transpose_scale=False,
718 + )
719 +
720 def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
721 # ROCm BF16 reference wo_a path (inverse RoPE + einsum) + wo_b.
722 z = rocm_inv_rope_einsum(
vllm/models/deepseek_v4/attention.py
+67 -11
@@ -373,19 +373,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
373 qr_kv, kv_score, indexer_kv_score, indexer_weights = (
374 self._run_parallel_input_projections(hidden_states)
375 )
376 - qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
377 - qr, kv = fused_q_kv_rmsnorm(
378 - qr,
379 - kv,
380 - self.q_norm.weight.data,
381 - self.kv_norm.weight.data,
382 - self.eps,
383 - )
376 + qr, qr_scale, kv = self._split_qkv_and_norm(qr_kv)
377
378 self._prepare_and_attn_fn(
379 hidden_states,
380 qr,
381 kv,
382 + qr_scale,
383 kv_score,
384 indexer_kv_score,
385 indexer_weights,
@@ -397,12 +391,32 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
391 # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
392 return self._o_proj(o, positions)
393
394 + def _split_qkv_and_norm(
395 + self, qr_kv: torch.Tensor
396 + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
397 + """Split the fused q-lora / kv projection and RMSNorm both halves.
398 +
399 + ``qr_scale`` is None on the shared path. The ROCm subclass returns
400 + a pre-quantized fp8 ``qr`` together with its per-1x128 fp32 scales,
401 + which the downstream wq_b projections consume directly.
402 + """
403 + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1)
404 + qr, kv = fused_q_kv_rmsnorm(
405 + qr,
406 + kv,
407 + self.q_norm.weight.data,
408 + self.kv_norm.weight.data,
409 + self.eps,
410 + )
411 + return qr, None, kv
412 +
413 @eager_break_during_capture
414 def _prepare_and_attn_eager(
415 self,
416 hidden_states: torch.Tensor,
417 qr: torch.Tensor,
418 kv: torch.Tensor,
419 + qr_scale: torch.Tensor | None,
420 kv_score: torch.Tensor,
421 indexer_kv_score: torch.Tensor,
422 indexer_weights: torch.Tensor,
@@ -418,6 +432,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
432 hidden_states,
433 qr,
434 kv,
435 + qr_scale,
436 kv_score,
437 indexer_kv_score,
438 indexer_weights,
@@ -430,6 +445,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
445 hidden_states: torch.Tensor,
446 qr: torch.Tensor,
447 kv: torch.Tensor,
448 + qr_scale: torch.Tensor | None,
449 kv_score: torch.Tensor,
450 indexer_kv_score: torch.Tensor,
451 indexer_weights: torch.Tensor,
@@ -446,7 +462,9 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
462 aux_streams = self.aux_stream_list
463
464 def project_query_and_cache_kv() -> torch.Tensor:
449 - q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
465 + q = self._wq_b_proj(qr, qr_scale).view(
466 + -1, self.n_local_heads, self.head_dim
467 + )
468 return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata)
469
470 index_q: torch.Tensor | None = None
@@ -468,6 +486,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
486 indexer_weights,
487 positions,
488 self.indexer_rotary_emb,
489 + qr_scale,
490 ),
491 lambda: compressor(kv_score, positions, self.rotary_emb),
492 ],
@@ -507,6 +526,24 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
526 qr_kv, _ = self.fused_wqa_wkv(hidden_states)
527 return qr_kv
528
529 + def _wq_b_proj(
530 + self, qr: torch.Tensor, qr_scale: torch.Tensor | None = None
531 + ) -> torch.Tensor:
532 + """Project the q-lora input to the query heads.
533 +
534 + ``qr_scale`` is only set by the ROCm fused path, where ``qr`` is
535 + already fp8-quantized with per-1x128 scales and must bypass the
536 + linear's internal re-quantization (apply_weights would quantize
537 + the fp8 input again).
538 + """
539 + if qr_scale is None:
540 + return self.wq_b(qr)
541 + from vllm.models.deepseek_v4.amd.rocm import (
542 + apply_pre_quantized_block_scaled_mm,
543 + )
544 +
545 + return apply_pre_quantized_block_scaled_mm(self.wq_b, qr, qr_scale)
546 +
547 def _run_parallel_input_projections(
548 self, hidden_states: torch.Tensor
549 ) -> tuple[
@@ -881,6 +918,7 @@ class DeepseekV4Indexer(nn.Module):
918 indexer_weights: torch.Tensor,
919 positions: torch.Tensor,
920 rotary_emb: nn.Module,
921 + qr_scale: torch.Tensor | None = None,
922 ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
923 compressor = self.compressor
924
@@ -911,8 +949,7 @@ class DeepseekV4Indexer(nn.Module):
949 return None, None, None
950
951 def wq_b_and_q_quant():
914 - # ReplicatedLinear returns (output, bias); bias is None.
915 - q, _ = self.wq_b(qr)
952 + q = self._wq_b_proj(qr, qr_scale)
953 q = q.view(-1, self.n_head, self.head_dim)
954 return fused_indexer_q_rope_quant(
955 positions,
@@ -938,3 +975,22 @@ class DeepseekV4Indexer(nn.Module):
975 else:
976 q, q_scale = q_quant, None
977 return q, q_scale, weights
978 +
979 + def _wq_b_proj(
980 + self, qr: torch.Tensor, qr_scale: torch.Tensor | None = None
981 + ) -> torch.Tensor:
982 + """Project the q-lora input to the indexer query heads.
983 +
984 + Same bypass as ``DeepseekV4Attention._wq_b_proj``: on ROCm the fused
985 + norm path hands over pre-quantized fp8 ``qr`` with per-1x128 scales,
986 + so the linear's internal re-quantization must be skipped.
987 + """
988 + if qr_scale is None:
989 + # ReplicatedLinear returns (output, bias); bias is None.
990 + q, _ = self.wq_b(qr)
991 + return q
992 + from vllm.models.deepseek_v4.amd.rocm import (
993 + apply_pre_quantized_block_scaled_mm,
994 + )
995 +
996 + return apply_pre_quantized_block_scaled_mm(self.wq_b, qr, qr_scale)