@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 04d1dd7f4a

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 14, 05:30 UTC

[ROCm][Aiter] Add triton fp8 bmm kernel for mla (#23264)

Signed-off-by: Divakar Verma <divakar.verma@amd.com> Co-authored-by: ShaoChunLee <Shao-Chun.Lee@amd.com>

Divakar Verma committed Aug 28, 2025 at 13:18 UTC 04d1dd7f4a444a61ae4b01ea0271490082dbd605

2 files changed +104 -12

vllm/envs.py

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

99 VLLM_ROCM_USE_AITER_RMSNORM: bool = True

100 VLLM_ROCM_USE_AITER_MLA: bool = True

101 VLLM_ROCM_USE_AITER_MHA: bool = True

102 + VLLM_ROCM_USE_AITER_FP8BMM: bool = True

103 VLLM_ROCM_USE_SKINNY_GEMM: bool = True

104 VLLM_ROCM_FP8_PADDING: bool = True

105 VLLM_ROCM_MOE_PADDING: bool = True

@@ -774,6 +775,12 @@ environment_variables: dict[str, Callable[[], Any]] = {

775 lambda: (os.getenv("VLLM_ROCM_USE_AITER_MHA", "True").lower() in

776 ("true", "1")),

777

778 + # Whether to use aiter triton fp8 bmm kernel

779 + # By default is enabled.

780 + "VLLM_ROCM_USE_AITER_FP8BMM":

781 + lambda: (os.getenv("VLLM_ROCM_USE_AITER_FP8BMM", "True").lower() in

782 + ("true", "1")),

783 +

784 # use rocm skinny gemms

785 "VLLM_ROCM_USE_SKINNY_GEMM":

786 lambda: (os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in

@@ -1272,6 +1279,7 @@ def compute_hash() -> str:

1279 "VLLM_ROCM_USE_AITER_RMSNORM",

1280 "VLLM_ROCM_USE_AITER_MLA",

1281 "VLLM_ROCM_USE_AITER_MHA",

1282 + "VLLM_ROCM_USE_AITER_FP8BMM",

1283 "VLLM_ROCM_USE_SKINNY_GEMM",

1284 "VLLM_ROCM_FP8_PADDING",

1285 "VLLM_ROCM_MOE_PADDING",

vllm/v1/attention/backends/mla/common.py

+96 -12

@@ -193,6 +193,7 @@ from dataclasses import dataclass, field

193 from typing import ClassVar, Generic, Optional, TypeVar, Union

194

195 import torch

196 +from tqdm import tqdm

197

198 import vllm.envs as envs

199 from vllm import _custom_ops as ops

@@ -203,6 +204,7 @@ from vllm.attention.backends.utils import get_mla_dims

204 from vllm.attention.ops.merge_attn_states import merge_attn_states

205 from vllm.attention.utils.fa_utils import get_flash_attn_version

206 from vllm.config import VllmConfig

207 +from vllm.distributed.parallel_state import is_global_first_rank

208 from vllm.logger import init_logger

209 from vllm.model_executor.layers.linear import (ColumnParallelLinear,

210 LinearBase,

@@ -234,6 +236,28 @@ try:

236 except ImportError:

237 flashinfer_available = False

238

239 +

240 +def is_rocm_aiter_fp8bmm_enabled() -> bool:

241 + return current_platform.is_rocm() \

242 + and envs.VLLM_ROCM_USE_AITER_FP8BMM \

243 + and envs.VLLM_ROCM_USE_AITER

244 +

245 +

246 +if is_rocm_aiter_fp8bmm_enabled():

247 + from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import ( # noqa: E501 # isort: skip

248 + batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant

249 + as aiter_triton_fp8_bmm)

250 +

251 + def dynamic_per_batched_tensor_quant(

252 + x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn):

253 + DTYPE_MAX = torch.finfo(dtype).max

254 + min_val, max_val = x.aminmax()

255 + amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-10)

256 + scale = DTYPE_MAX / amax

257 + x_scl_sat = (x * scale).clamp(min=-DTYPE_MAX, max=DTYPE_MAX)

258 + return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal()

259 +

260 +

261 logger = init_logger(__name__)

262

263 CUDNN_WORKSPACE_SIZE = 12800

@@ -945,10 +969,21 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

969 def _v_up_proj(self, x):

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

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

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

949 - x = torch.bmm(x, self.W_UV)

950 - # Convert from (N, B, V) to (B, N * V)

951 - return x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)

972 + if is_rocm_aiter_fp8bmm_enabled():

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

974 + x = aiter_triton_fp8_bmm(x,

975 + self.W_V,

976 + self.W_V_scale,

977 + group_size=128,

978 + transpose_bm=True)

979 + # Convert from (B, N, V) to (B, N * V)

980 + x = x.reshape(-1, self.num_heads * self.v_head_dim)

981 + else:

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

983 + x = torch.bmm(x, self.W_UV)

984 + # Convert from (N, B, V) to (B, N * V)

985 + x = x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)

986 + return x

987

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

989

@@ -996,10 +1031,50 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

1031 W_UK, W_UV = kv_b_proj_weight.split(

1032 [self.qk_nope_head_dim, self.v_head_dim], dim=-1)

1033

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

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

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

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

1034 + if is_rocm_aiter_fp8bmm_enabled():

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

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

1037 + self.W_K, self.W_K_scale = dynamic_per_batched_tensor_quant(

1038 + W_K, dtype=current_platform.fp8_dtype())

1039 + self.W_V, self.W_V_scale = dynamic_per_batched_tensor_quant(

1040 + W_V, dtype=current_platform.fp8_dtype())

1041 +

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

1043 + # triton kernel to avoid runtime compilation for unseen batch sizes

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

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

1046 + max_batch_size = 1024 # [ToDo] Find the optimal upper limit

1047 + pre_compilation_list = list(range(1, max_batch_size + 1))

1048 + if is_global_first_rank():

1049 + pre_compilation_list = tqdm(

1050 + pre_compilation_list,

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

1052 + total=max_batch_size,

1053 + )

1054 +

1055 + for m in pre_compilation_list:

1056 + x = torch.empty((self.W_K.shape[0], m, self.W_K.shape[2]),

1057 + dtype=torch.bfloat16,

1058 + device=self.W_K.device)

1059 + aiter_triton_fp8_bmm(x,

1060 + self.W_K,

1061 + self.W_K_scale,

1062 + group_size=128,

1063 + transpose_bm=True)

1064 +

1065 + x = torch.empty((self.W_V.shape[0], m, self.W_V.shape[2]),

1066 + dtype=torch.bfloat16,

1067 + device=self.W_V.device)

1068 + aiter_triton_fp8_bmm(x,

1069 + self.W_V,

1070 + self.W_V_scale,

1071 + group_size=128,

1072 + transpose_bm=True)

1073 + else:

1074 + # Convert from (L, N, V) to (N, L, V)

1075 + self.W_UV = W_UV.transpose(0, 1)

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

1077 + self.W_UK_T = W_UK.permute(1, 2, 0)

1078

1079 def _compute_prefill_context(

1080 self,

@@ -1203,10 +1278,19 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

1278 [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

1279 # Convert from (B, N, P) to (N, B, P)

1280 decode_q_nope = decode_q_nope.transpose(0, 1)

1206 - # Multiply (N, B, P) x (N, P, L) -> (N, B, L)

1207 - decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T)

1208 - # Convert from (N, B, L) to (B, N, L)

1209 - decode_ql_nope = decode_ql_nope.transpose(0, 1)

1281 +

1282 + if is_rocm_aiter_fp8bmm_enabled():

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

1284 + decode_ql_nope = aiter_triton_fp8_bmm(decode_q_nope,

1285 + self.W_K,

1286 + self.W_K_scale,

1287 + group_size=128,

1288 + transpose_bm=True)

1289 + else:

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

1291 + decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T)

1292 + # Convert from (N, B, L) to (B, N, L)

1293 + decode_ql_nope = decode_ql_nope.transpose(0, 1)

1294

1295 if fp8_attention:

1296 ql_nope_shape = decode_ql_nope.shape