@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 09663abde0

[ROCm][MLA] Fuse MLA q/kv RMSNorm + FP8 per-token quant in the FP8 attention path (#44977)

Signed-off-by: Xavier Aguilar <xavier.aguilarfruto@amd.com> Signed-off-by: Xavier Aguilar <Xavier.AguilarFruto@amd.com>

xaguilar-amd committed Jul 2, 2026 at 07:00 UTC 09663abde0f50944a8d5ea30120666024b503faa

4 files changed +371 -5

docs/design/fusions.md

+25 -4

@@ -30,7 +30,7 @@ or just on the low or high end.

35 ## Support Matrix

@@ -381,11 +381,32 @@ q_normed, kv_normed = fused_mla_dual_rms_norm(

381 Requires: AMD ROCm with AITER enabled. Enabled by default at optimization level O1 and above

382 when AITER is available.

383

384 +**FP8 attention variant (per-token quant).** With a per-token FP8 `q_b_proj`,

385 +only the *q* latent is FP8-quantized while *kv* stays bf16.

386 +`RocmAiterRMSNormQuantFusionPass` first folds the q side into

387 +`rocm_aiter_rmsnorm_fused_dynamic_quant`, leaving kv a plain

388 +`rms_norm` — breaking the symmetric pattern above. The same pass then matches

389 +this asymmetric pair and lowers it to `fused_mla_dual_rms_norm_per_token_quant`.

390 +

391 +```text

392 +# Unfused (q norm+quant fused; kv still plain rms_norm):

393 +q_c, kv_lora = split(projected, [q_dim, kv_dim])

394 +kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])

395 +q_fp8, q_scale = rocm_aiter_rmsnorm_fused_dynamic_quant(q_c, q_weight, eps, fp8)

396 +kv_normed = rms_norm(kv_c, kv_weight, eps) # bf16

397 +

398 +# Fused:

399 +q_c, kv_lora = split(projected, [q_dim, kv_dim])

400 +kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])

401 +q_fp8, q_scale, kv_normed = fused_mla_dual_rms_norm_per_token_quant(

402 + q_c, q_weight, kv_c, kv_weight, eps1, eps2)

403 +```

404 +

405 **Code locations.**

406

386 -- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`)

387 -- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`)

388 -- AITER kernel: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442)

407 +- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`, `MLADualRMSPerTokenQuantPattern`)

408 +- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`, `fused_mla_dual_rms_norm_per_token_quant`)

409 +- AITER kernels: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442), `fused_qk_rmsnorm_per_token_quant`

410

411 ## See Also

412

tests/compile/passes/test_fuse_mla_dual_rms_norm.py

+143 -1

@@ -12,7 +12,10 @@ import torch

13 import vllm.config

14 from tests.compile.backend import TestBackend

15 -from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops

15 +from vllm._aiter_ops import (

16 + is_aiter_found_and_supported,

17 + rocm_aiter_ops,

18 +)

19 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass

20 from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass

21 from vllm.config import (

@@ -23,6 +26,7 @@ from vllm.config import (

26 VllmConfig,

27 )

28 from vllm.model_executor.layers.layernorm import RMSNorm

29 +from vllm.platforms import current_platform

31 # MLA attention geometry for DeepSeek-V3 / Kimi-K2

32 Q_DIM = 1536

@@ -30,6 +34,8 @@ KV_C_DIM = 512

34 K_PE_DIM = 64

35 EPS = 1e-6

37 +FP8_DTYPE = current_platform.fp8_dtype()

38 +

40 class MLADualRMSNormTestModel(torch.nn.Module):

41 """

@@ -146,3 +152,139 @@ def test_fuse_mla_dual_rms_norm(

152

153 backend.check_before_ops(model.ops_in_model_before())

154 backend.check_after_ops(model.ops_in_model_after())

155 +

156 +

157 +class MLADualRMSNormFp8PerTokenTestModel(torch.nn.Module):

158 + """

159 + Minimal model reproducing the FP8 MLA attention path with *per-token* quant:

160 + linear -> split([q_dim, kv_dim])

161 + +-- q_c (getitem 0) -> rocm_aiter_rmsnorm_fused_dynamic_quant -> dequant

162 + +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim])

163 + +-- kv_c (getitem 0) -> rms_norm (bf16)

164 + +-- k_pe

165 + """

166 +

167 + def __init__(

168 + self,

169 + hidden_size: int,

170 + q_dim: int = Q_DIM,

171 + kv_c_dim: int = KV_C_DIM,

172 + k_pe_dim: int = K_PE_DIM,

173 + eps: float = EPS,

174 + ):

175 + super().__init__()

176 + self.q_dim = q_dim

177 + self.kv_dim = kv_c_dim + k_pe_dim

178 + self.kv_c_dim = kv_c_dim

179 + self.k_pe_dim = k_pe_dim

180 + self.eps = eps

181 +

182 + self.proj = torch.nn.Linear(hidden_size, q_dim + self.kv_dim, bias=False)

183 + self.q_weight = torch.nn.Parameter(torch.ones(q_dim))

184 + self.kv_norm = RMSNorm(kv_c_dim, eps=eps)

185 +

186 + def _dequant(self, x_fp8: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:

187 + # Per-token: a single (M, 1) scale broadcast across the row.

188 + return (x_fp8.to(torch.float32) * scale).to(torch.bfloat16)

189 +

190 + def forward(self, x: torch.Tensor):

191 + # Avoid graph input being a direct arg to a matched pattern node

192 + x = torch.relu(x)

193 +

194 + projected = self.proj(x)

195 +

196 + q_c, kv_lora = projected.split([self.q_dim, self.kv_dim], dim=-1)

197 + kv_c, k_pe = kv_lora.split([self.kv_c_dim, self.k_pe_dim], dim=-1)

198 +

199 + q_fp8, q_scale = torch.ops.vllm.rocm_aiter_rmsnorm_fused_dynamic_quant(

200 + q_c, self.q_weight, self.eps, FP8_DTYPE

201 + )

202 + kv_normed = self.kv_norm(kv_c)

203 +

204 + return self._dequant(q_fp8, q_scale), kv_normed, k_pe

205 +

206 + def ops_in_model_before(self):

207 + return [

208 + torch.ops.vllm.rocm_aiter_rmsnorm_fused_dynamic_quant.default,

209 + torch.ops.vllm_ir.rms_norm.default,

210 + ]

211 +

212 + def ops_in_model_after(self):

213 + return [torch.ops.vllm.fused_mla_dual_rms_norm_per_token_quant.default]

214 +

215 +

216 +@pytest.mark.parametrize("dtype", [torch.bfloat16])

217 +@pytest.mark.parametrize("hidden_size", [7168])

218 +@pytest.mark.skipif(

219 + not is_aiter_found_and_supported(),

220 + reason="Only test on ROCm with AITER installed and supported",

221 +)

222 +def test_fuse_mla_dual_rms_norm_fp8_per_token(

223 + dtype: torch.dtype,

224 + hidden_size: int,

225 + monkeypatch: pytest.MonkeyPatch,

226 +):

227 + torch._dynamo.reset()

228 +

229 + vllm_config = VllmConfig(

230 + model_config=ModelConfig(dtype=dtype),

231 + compilation_config=CompilationConfig(

232 + mode=CompilationMode.VLLM_COMPILE,

233 + custom_ops=["+rms_norm"],

234 + pass_config=PassConfig(

235 + fuse_mla_dual_rms_norm=True,

236 + eliminate_noops=True,

237 + ),

238 + ),

239 + )

240 +

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

242 + from vllm.compilation.passes.fusion.rocm_aiter_fusion import (

243 + MLADualRMSNormFusionPass,

244 + )

245 +

246 + torch.set_default_device("cuda")

247 + torch.set_default_dtype(dtype)

248 + torch.manual_seed(42)

249 +

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

251 + rocm_aiter_ops.refresh_env_variables()

252 +

253 + fusion_pass = MLADualRMSNormFusionPass(vllm_config)

254 + passes = [

255 + NoOpEliminationPass(vllm_config),

256 + fusion_pass,

257 + PostCleanupPass(vllm_config),

258 + ]

259 + backend = TestBackend(*passes)

260 + model = MLADualRMSNormFp8PerTokenTestModel(hidden_size)

261 +

262 + x = torch.randn(4, hidden_size)

263 + torch._dynamo.mark_dynamic(x, 0)

264 +

265 + with torch.inference_mode():

266 + outputs_unfused = model(x)

267 +

268 + model_fused = torch.compile(model, backend=backend)

269 + outputs_fused = model_fused(x)

270 +

271 + q_deq_u, kv_normed_u, k_pe_u = outputs_unfused

272 + q_deq_f, kv_normed_f, k_pe_f = outputs_fused

273 +

274 + torch.testing.assert_close(k_pe_u, k_pe_f, atol=0, rtol=0)

275 +

276 + torch.testing.assert_close(kv_normed_u, kv_normed_f, atol=1e-2, rtol=1e-2)

277 +

278 + E4M3_STEP = 0.125

279 + exact_frac = (q_deq_u == q_deq_f).float().mean().item()

280 + assert exact_frac > 0.99, (

281 + f"q: only {exact_frac:.4f} of elements bit-exact; scales likely differ"

282 + )

283 + torch.testing.assert_close(q_deq_u, q_deq_f, atol=1e-2, rtol=E4M3_STEP)

284 +

285 + assert fusion_pass.matched_count == 1, (

286 + f"Expected 1 fused pair, got {fusion_pass.matched_count}"

287 + )

288 +

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

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

vllm/_aiter_ops.py

+68

@@ -1355,6 +1355,63 @@ def _fused_mla_dual_rms_norm_fake(

1355 return (torch.empty_like(x1), torch.empty_like(x2))

1356

1357

1358 +def _fused_mla_dual_rms_norm_per_token_quant_impl(

1359 + q: torch.Tensor,

1360 + q_weight: torch.Tensor,

1361 + kv: torch.Tensor,

1362 + kv_weight: torch.Tensor,

1363 + q_epsilon: float,

1364 + kv_epsilon: float,

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

1366 + """Fused MLA q/kv RMSNorm (+ FP8 per-token quant on q) via AITER.

1367 +

1368 + Backs the ``fused_mla_dual_rms_norm_per_token_quant`` custom op used by the

1369 + MLA FP8 attention fusion when the q latent is quantized *per token* (a single

1370 + ``(M, 1)`` scale). Only the *q* latent is FP8 quantized (it feeds the

1371 + FP8 ``q_b_proj`` GEMM); the *kv* latent is RMS-normed and consumed by attention as bf16.

1372 + """

1373 + from aiter.ops.fused_qk_rmsnorm_group_quant import (

1374 + fused_qk_rmsnorm_per_token_quant,

1375 + )

1376 +

1377 + mq, nq = q.shape

1378 + q_out = torch.empty((mq, nq), dtype=FP8_DTYPE, device=q.device)

1379 + q_scale = torch.empty((mq, 1), dtype=torch.float32, device=q.device)

1380 + kv_normed = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)

1381 +

1382 + # q -> RMSNorm + FP8 per-token quant (q slot); kv -> RMSNorm only (k slot).

1383 + # `split` views are accepted directly (unit inner stride); the kernel

1384 + # handles strided inputs, matching the aiter op-test usage.

1385 + fused_qk_rmsnorm_per_token_quant(

1386 + q_out_quantized=q_out,

1387 + q_out_scale=q_scale,

1388 + q=q,

1389 + q_weight=q_weight,

1390 + q_epsilon=q_epsilon,

1391 + k_out=kv_normed,

1392 + k=kv,

1393 + k_weight=kv_weight,

1394 + k_epsilon=kv_epsilon,

1395 + gemma_norm=False,

1396 + )

1397 + return q_out, q_scale, kv_normed

1398 +

1399 +

1400 +def _fused_mla_dual_rms_norm_per_token_quant_fake(

1401 + q: torch.Tensor,

1402 + q_weight: torch.Tensor,

1403 + kv: torch.Tensor,

1404 + kv_weight: torch.Tensor,

1405 + q_epsilon: float,

1406 + kv_epsilon: float,

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

1408 + mq, nq = q.shape

1409 + q_out = torch.empty((mq, nq), dtype=FP8_DTYPE, device=q.device)

1410 + q_scale = torch.empty((mq, 1), dtype=torch.float32, device=q.device)

1411 + kv_normed = torch.empty(kv.shape, dtype=kv.dtype, device=kv.device)

1412 + return q_out, q_scale, kv_normed

1413 +

1414 +

1415 def _rocm_aiter_gemm_a8wfp4_impl(

1416 x: torch.Tensor,

1417 w: torch.Tensor,

@@ -2046,6 +2103,13 @@ class rocm_aiter_ops:

2103 fake_impl=_fused_mla_dual_rms_norm_fake,

2104 )

2105

2106 + direct_register_custom_op(

2107 + op_name="fused_mla_dual_rms_norm_per_token_quant",

2108 + op_func=_fused_mla_dual_rms_norm_per_token_quant_impl,

2109 + mutates_args=[],

2110 + fake_impl=_fused_mla_dual_rms_norm_per_token_quant_fake,

2111 + )

2112 +

2113 _OPS_REGISTERED = True

2114

2115 @staticmethod

@@ -2120,6 +2184,10 @@ class rocm_aiter_ops:

2184 def get_fused_mla_dual_rms_norm_op() -> OpOverload:

2185 return torch.ops.vllm.fused_mla_dual_rms_norm.default

2186

2187 + @staticmethod

2188 + def get_fused_mla_dual_rms_norm_per_token_quant_op() -> OpOverload:

2189 + return torch.ops.vllm.fused_mla_dual_rms_norm_per_token_quant.default

2190 +

2191 @staticmethod

2192 def w8a8_gemm(

2193 A: torch.Tensor,

vllm/compilation/passes/fusion/rocm_aiter_fusion.py

+135

@@ -922,11 +922,145 @@ class MLADualRMSNormPattern(

922 return _replacement

923

924

925 +class MLADualRMSPerTokenQuantPattern(

926 + VllmPatternReplacement[

927 + ...,

928 + tuple[

929 + torch.Tensor,

930 + torch.Tensor,

931 + torch.Tensor,

932 + torch.Tensor,

933 + ],

934 + ]

935 +):

936 + """

937 + Fuse the MLA FP8 attention path -- q-latent RMSNorm + FP8 *per-token* quant

938 + plus kv-latent RMSNorm -- into AITER's ``fused_qk_rmsnorm_per_token_quant``.

939 +

940 + With a per-token FP8 ``q_b_proj`` (Quark / ModelOpt), the earlier

941 + ``RocmAiterRMSNormQuantFusionPass`` folds the q side into

942 + ``rocm_aiter_rmsnorm_fused_dynamic_quant`` and leaves the kv side a plain

943 + ``vllm_ir.rms_norm``. This pattern matches that asymmetric pair::

944 +

945 + gemm -> split_with_sizes([q_dim, kv_dim])

946 + +-- q_c -> rocm_aiter_rmsnorm_fused_dynamic_quant -> (q_fp8, q_scale)

947 + +-- kv_lora -> split_with_sizes([kv_c_dim, k_pe_dim])

948 + +-- kv_c -> vllm_ir.rms_norm -> kv_normed (bf16)

949 + +-- k_pe

950 + """

951 +

952 + DYNAMIC_QUANT_OP = rocm_aiter_ops.get_rmsnorm_fused_dynamic_quant_op()

953 + FUSED_OP = rocm_aiter_ops.get_fused_mla_dual_rms_norm_per_token_quant_op()

954 +

955 + def __init__(self, epsilon: float) -> None:

956 + self._epsilon = epsilon

957 +

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

959 + q_dim, kv_c_dim, k_pe_dim = 256, 128, 64

960 + return [

961 + self.empty_bf16(5, q_dim + kv_c_dim + k_pe_dim),

962 + self.empty_bf16(q_dim),

963 + self.empty_bf16(kv_c_dim),

964 + ]

965 +

966 + @property

967 + def pattern(

968 + self,

969 + ) -> Callable[

970 + ...,

971 + tuple[

972 + torch.Tensor,

973 + torch.Tensor,

974 + torch.Tensor,

975 + torch.Tensor,

976 + ],

977 + ]:

978 + eps = self._epsilon

979 + dynamic_quant_op = self.DYNAMIC_QUANT_OP

980 +

981 + def _pattern(

982 + projected: torch.Tensor,

983 + q_weight: torch.Tensor,

984 + kv_weight: torch.Tensor,

985 + ) -> tuple[

986 + torch.Tensor,

987 + torch.Tensor,

988 + torch.Tensor,

989 + torch.Tensor,

990 + ]:

991 + q_dim = q_weight.shape[0]

992 + kv_dim = projected.shape[-1] - q_dim

993 + kv_c_dim = kv_weight.shape[0]

994 + k_pe_dim = kv_dim - kv_c_dim

995 + q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1)

996 + kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1)

997 + q_quant = dynamic_quant_op(

998 + x=q_c,

999 + weight=q_weight,

1000 + epsilon=eps,

1001 + quant_dtype=FP8_DTYPE,

1002 + )

1003 + kv_normed = vllm.ir.ops.rms_norm(kv_c, kv_weight, eps)

1004 + return q_quant[0], q_quant[1], kv_normed, k_pe

1005 +

1006 + return _pattern

1007 +

1008 + @property

1009 + def replacement(

1010 + self,

1011 + ) -> Callable[

1012 + ...,

1013 + tuple[

1014 + torch.Tensor,

1015 + torch.Tensor,

1016 + torch.Tensor,

1017 + torch.Tensor,

1018 + ],

1019 + ]:

1020 + eps = self._epsilon

1021 + fused_op = self.FUSED_OP

1022 +

1023 + def _replacement(

1024 + projected: torch.Tensor,

1025 + q_weight: torch.Tensor,

1026 + kv_weight: torch.Tensor,

1027 + ) -> tuple[

1028 + torch.Tensor,

1029 + torch.Tensor,

1030 + torch.Tensor,

1031 + torch.Tensor,

1032 + ]:

1033 + q_dim = q_weight.shape[0]

1034 + kv_dim = projected.shape[-1] - q_dim

1035 + kv_c_dim = kv_weight.shape[0]

1036 + k_pe_dim = kv_dim - kv_c_dim

1037 + q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1)

1038 + kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1)

1039 + at = fused_op(

1040 + q_c,

1041 + q_weight,

1042 + kv_c,

1043 + kv_weight,

1044 + eps,

1045 + eps,

1046 + )

1047 + # q_fp8, q_scale, kv_normed, k_pe

1048 + return at[0], at[1], at[2], k_pe

1049 +

1050 + return _replacement

1051 +

1052 +

1053 class MLADualRMSNormFusionPass(VllmFusionPatternMatcherPass):

1054 """

1055 Post-grad PatternMatcher pass that fuses paired q / kv RMS norms in

1056 MLA attention into ``fused_mla_dual_rms_norm`` backed by aiter's

1057 ``fused_qk_rmsnorm`` HIP kernel.

1058 +

1059 + The FP8 attention path is also handled via

1060 + :class:`MLADualRMSPerTokenQuantPattern`, which fuses the q-latent RMSNorm +

1061 + FP8 per-token quant together with the kv-latent RMSNorm into

1062 + ``fused_mla_dual_rms_norm_per_token_quant`` backed by aiter's

1063 + ``fused_qk_rmsnorm_per_token_quant`` HIP kernel.

1064 """

1065

1066 def __init__(self, config: VllmConfig) -> None:

@@ -934,3 +1068,4 @@ class MLADualRMSNormFusionPass(VllmFusionPatternMatcherPass):

1068

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

1070 self.register(MLADualRMSNormPattern(epsilon))

1071 + self.register(MLADualRMSPerTokenQuantPattern(epsilon))