@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8c11001ba2

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 02:45 UTC

[ROCM] DSfp4 mla projection gemms weight dynamic quantization (#32238)

Signed-off-by: Aleksandr Malyshev <maleksan@amd.com> Co-authored-by: Aleksandr Malyshev <maleksan@amd.com>

Aleksandr Malyshev committed Jan 15, 2026 at 12:13 UTC 8c11001ba2110134df3c3aecc63f2559ad1f5996

4 files changed +97 -80

vllm/_aiter_ops.py

+30

@@ -837,6 +837,7 @@ class rocm_aiter_ops:

837 _TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION

838 # TODO: Consolidate under _LINEAR_ENABLED

839 _FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM

840 + _FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM

841 # TODO: Consolidate under _LINEAR_ENABLED

842 _FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM

843 # TODO: Consolidate under VLLM_ROCM_USE_AITER_ROPE

@@ -863,6 +864,7 @@ class rocm_aiter_ops:

864 cls._SHUFFLE_KV_CACHE_ENABLED = envs.VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT

865 cls._TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION

866 cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM

867 + cls._FP4BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP4BMM

868 cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM

869 cls._TRITON_ROTARY_EMBED = envs.VLLM_ROCM_USE_AITER_TRITON_ROPE

870 cls._MOE_SHARED_EXPERTS_ENABLED = envs.VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS

@@ -923,6 +925,11 @@ class rocm_aiter_ops:

925 def is_fp8bmm_enabled(cls) -> bool:

926 return cls._AITER_ENABLED and cls._FP8BMM_ENABLED

927

928 + @classmethod

929 + @if_aiter_supported

930 + def is_fp4bmm_enabled(cls) -> bool:

931 + return cls._AITER_ENABLED and cls._FP4BMM_ENABLED

932 +

933 @classmethod

934 @if_aiter_supported

935 def is_asm_fp4_gemm_dynamic_quant_enabled(cls) -> bool:

@@ -1403,6 +1410,29 @@ class rocm_aiter_ops:

1410 query = query.view(query_shape)

1411 key = key.view(key_shape)

1412

1413 + @staticmethod

1414 + def batched_gemm_a16wfp4(

1415 + X: torch.Tensor,

1416 + W: torch.Tensor,

1417 + w_scale: torch.Tensor,

1418 + Y: torch.Tensor,

1419 + transpose_bm: bool | None = False,

1420 + prequant: bool | None = False,

1421 + y_scale: torch.Tensor | None = None,

1422 + ) -> torch.Tensor:

1423 + # ruff: noqa: E501 # isort: skip

1424 + from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4

1425 +

1426 + return batched_gemm_a16wfp4(

1427 + X,

1428 + W,

1429 + w_scale,

1430 + y=Y,

1431 + transpose_bm=transpose_bm,

1432 + prequant=prequant,

1433 + y_scale=y_scale,

1434 + )

1435 +

1436 @staticmethod

1437 def triton_fp8_bmm(

1438 X: torch.Tensor,

vllm/envs.py

@@ -121,6 +121,7 @@ if TYPE_CHECKING:

121 VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False

122 VLLM_ROCM_USE_AITER_TRITON_ROPE: bool = False

123 VLLM_ROCM_USE_AITER_FP8BMM: bool = True

124 + VLLM_ROCM_USE_AITER_FP4BMM: bool = True

125 VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False

126 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False

127 VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True

@@ -991,6 +992,11 @@ environment_variables: dict[str, Callable[[], Any]] = {

992 "VLLM_ROCM_USE_AITER_FP8BMM": lambda: (

993 os.getenv("VLLM_ROCM_USE_AITER_FP8BMM", "True").lower() in ("true", "1")

994 ),

995 + # Whether to use aiter triton fp4 bmm kernel

996 + # By default is enabled.

997 + "VLLM_ROCM_USE_AITER_FP4BMM": lambda: (

998 + os.getenv("VLLM_ROCM_USE_AITER_FP4BMM", "True").lower() in ("true", "1")

999 + ),

1000 # Use AITER triton unified attention for V1 attention

1001 "VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION": lambda: (

1002 os.getenv("VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION", "False").lower()

vllm/model_executor/layers/attention/mla_attention.py

+47 -80

@@ -1183,6 +1183,12 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):

1183 self.q_pad_num_heads = q_pad_num_heads

1184 self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()

1185

1186 + # If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported

1187 + self.is_aiter_triton_fp4_bmm_enabled = (

1188 + rocm_aiter_ops.is_fp4bmm_enabled()

1189 + and self.kv_b_proj.weight.dtype == torch.bfloat16

1190 + )

1191 +

1192 def process_weights_after_loading(self, act_dtype: torch.dtype):

1193 # we currently do not have quantized bmm's which are needed for

1194 # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform

@@ -1211,7 +1217,21 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):

1217 [self.qk_nope_head_dim, self.v_head_dim], dim=-1

1218 )

1219

1214 - if self.is_aiter_triton_fp8_bmm_enabled:

1220 + # If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported

1221 + if self.is_aiter_triton_fp4_bmm_enabled:

1222 + from vllm.model_executor.layers.quantization.quark.utils import (

1223 + quark_quantize_weight_to_mxfp4,

1224 + )

1225 +

1226 + self.W_K, self.W_K_scale = quark_quantize_weight_to_mxfp4(W_UK)

1227 + # Convert from (L, N, P) to (N, L, P)

1228 + self.W_K = self.W_K.transpose(0, 1)

1229 + self.W_K_scale = self.W_K_scale.transpose(0, 1)

1230 +

1231 + self.W_V, self.W_V_scale = quark_quantize_weight_to_mxfp4(

1232 + W_UV.permute(1, 2, 0)

1233 + )

1234 + elif self.is_aiter_triton_fp8_bmm_enabled:

1235 W_K = W_UK.transpose(0, 1) # 16 512 128

1236 W_V = W_UV.permute(1, 2, 0) # 16 128 512

1237 self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(

@@ -1261,16 +1281,26 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):

1281 def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):

1282 # Convert from (B, N, L) to (N, B, L)

1283 x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)

1264 -

1265 - if self.is_aiter_triton_fp8_bmm_enabled:

1266 - out = out.view(-1, self.num_heads, self.v_head_dim)

1284 + out = out.view(-1, self.num_heads, self.v_head_dim)

1285 + if self.is_aiter_triton_fp4_bmm_enabled:

1286 + out = rocm_aiter_ops.batched_gemm_a16wfp4(

1287 + x,

1288 + self.W_V,

1289 + self.W_V_scale,

1290 + out,

1291 + transpose_bm=True,

1292 + prequant=True,

1293 + y_scale=None,

1294 + )

1295 + x = out.view(-1, self.num_heads * self.v_head_dim)

1296 + elif self.is_aiter_triton_fp8_bmm_enabled:

1297 # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)

1298 x = rocm_aiter_ops.triton_fp8_bmm(

1299 x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True, YQ=out

1300 )

1301 else:

1302 # Convert from (B, N * V) to (N, B, V)

1273 - out = out.view(-1, self.num_heads, self.v_head_dim).transpose(0, 1)

1303 + out = out.transpose(0, 1)

1304

1305 # Multiply (N, B, L) x (N, L, V) -> (N, B, V)

1306 torch.bmm(x, self.W_UV, out=out) # Reuse "out" to make it "hot"

@@ -1579,80 +1609,6 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1609 # Convert from (q_len, num_heads) to (num_heads, q_len)

1610 return attn_out, lse.transpose(0, 1).contiguous()

1611

1582 - def process_weights_after_loading(self, act_dtype: torch.dtype):

1583 - # we currently do not have quantized bmm's which are needed for

1584 - # `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform

1585 - # the bmm's in 16-bit, the extra memory overhead of this is fairly low

1586 - kv_b_proj_weight = get_and_maybe_dequant_weights(

1587 - self.kv_b_proj, out_dtype=act_dtype

1588 - ).T

1589 - assert kv_b_proj_weight.shape == (

1590 - self.kv_lora_rank,

1591 - self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),

1592 - ), (

1593 - f"{kv_b_proj_weight.shape=}, "

1594 - f"{self.kv_lora_rank=}, "

1595 - f"{self.num_heads=}, "

1596 - f"{self.qk_nope_head_dim=}, "

1597 - f"{self.v_head_dim=}"

1598 - )

1599 - kv_b_proj_weight = kv_b_proj_weight.view(

1600 - self.kv_lora_rank,

1601 - self.num_heads,

1602 - self.qk_nope_head_dim + self.v_head_dim,

1603 - )

1604 -

1605 - W_UK, W_UV = kv_b_proj_weight.split(

1606 - [self.qk_nope_head_dim, self.v_head_dim], dim=-1

1607 - )

1608 -

1609 - if self.is_aiter_triton_fp8_bmm_enabled:

1610 - W_K = W_UK.transpose(0, 1) # 16 512 128

1611 - W_V = W_UV.permute(1, 2, 0) # 16 128 512

1612 - self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(

1613 - W_K, dtype=current_platform.fp8_dtype()

1614 - )

1615 - self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant(

1616 - W_V, dtype=current_platform.fp8_dtype()

1617 - )

1618 -

1619 - # The kernel operates on non-padded inputs. Hence, pre-compiling

1620 - # triton kernel to avoid runtime compilation for unseen batch sizes

1621 - # Pre-compile for batch sizes 1 to 1024 to cover most use-cases.

1622 - # On DS-R1, this step adds roughly 50s to the model loading time.

1623 - max_batch_size = 1024 # [ToDo] Find the optimal upper limit

1624 - pre_compilation_list = list(range(1, max_batch_size + 1))

1625 - if is_global_first_rank():

1626 - pre_compilation_list = tqdm(

1627 - pre_compilation_list,

1628 - desc="[Aiter Triton] Pre-compiling fp8 BMM kernel",

1629 - total=max_batch_size,

1630 - )

1631 -

1632 - for m in pre_compilation_list:

1633 - x = torch.empty(

1634 - (self.W_K.shape[0], m, self.W_K.shape[2]),

1635 - dtype=torch.bfloat16,

1636 - device=self.W_K.device,

1637 - )

1638 - rocm_aiter_ops.triton_fp8_bmm(

1639 - x, self.W_K, self.W_K_scale, group_size=128, transpose_bm=True

1640 - )

1641 -

1642 - x = torch.empty(

1643 - (self.W_V.shape[0], m, self.W_V.shape[2]),

1644 - dtype=torch.bfloat16,

1645 - device=self.W_V.device,

1646 - )

1647 - rocm_aiter_ops.triton_fp8_bmm(

1648 - x, self.W_V, self.W_V_scale, group_size=128, transpose_bm=True

1649 - )

1650 - else:

1651 - # Convert from (L, N, V) to (N, L, V)

1652 - self.W_UV = W_UV.transpose(0, 1)

1653 - # Convert from (L, N, P) to (N, P, L)

1654 - self.W_UK_T = W_UK.permute(1, 2, 0)

1655 -

1612 def _concat_k_nope_k_pe(

1613 self, k_nope: torch.Tensor, k_pe: torch.Tensor

1614 ) -> torch.Tensor:

@@ -2033,7 +1989,18 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

1989 decode_pe_padded.copy_(decode_q_pe)

1990 decode_q_pe = decode_pe_padded

1991

2036 - if self.is_aiter_triton_fp8_bmm_enabled:

1992 + if self.is_aiter_triton_fp4_bmm_enabled:

1993 + from aiter.ops.triton.batched_gemm_a16wfp4 import batched_gemm_a16wfp4

1994 +

1995 + decode_ql_nope = batched_gemm_a16wfp4(

1996 + decode_q_nope,

1997 + self.W_K,

1998 + self.W_K_scale,

1999 + transpose_bm=True,

2000 + prequant=True,

2001 + y_scale=layer._q_scale if fp8_attention else None,

2002 + )

2003 + elif self.is_aiter_triton_fp8_bmm_enabled:

2004 # Multiply+Transpose (N, B, P)x(N, P, L)->(N, B, L)->(B, N, L)

2005 decode_ql_nope = rocm_aiter_ops.triton_fp8_bmm(

2006 decode_q_nope,

vllm/model_executor/layers/quantization/quark/utils.py

+14

@@ -6,6 +6,7 @@ from types import MappingProxyType

6 from typing import Any

8 import regex as re

9 +import torch

12 def deep_compare(dict1: Any, dict2: Any) -> bool:

@@ -103,3 +104,16 @@ def _is_equal_or_regex_match(

104 elif target == value:

105 return True

106 return False

107 +

108 +

109 +# utility for tensor dims > 2 cases

110 +def quark_quantize_weight_to_mxfp4(w: torch.Tensor):

111 + assert w.dtype == torch.bfloat16, (

112 + "Quark dynamic quantization is supported only for fp16 weights and only to MXF4"

113 + )

114 +

115 + from aiter.ops.triton.quant import dynamic_mxfp4_quant

116 +

117 + *dims, d = w.shape

118 + w, w_scales = dynamic_mxfp4_quant(w.reshape(-1, d))

119 + return w.view(*dims, d // 2), w_scales.view(*dims, d // 32)