@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5f7fab881a

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 5, 02:15 UTC

[ROCm][FEAT] Integrate aiter gemm w8a8 ptpc (#33773)

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

vllmellm committed Apr 16, 2026 at 09:55 UTC 5f7fab881a13df2926f50cb8d9f67b0bbc2ee41f

9 files changed +279 -23

tests/utils.py

@@ -1856,6 +1856,7 @@ class TestFP8Layer(torch.nn.Module):

1856 out_dtype=out_dtype,

1857 force_kernel=force_kernel,

1858 )

1859 + self.kernel.process_weights_after_loading(self)

1860

1861 def is_quant_fp8_enabled(self) -> bool:

1862 return self.kernel.quant_fp8.enabled()

vllm/_aiter_ops.py

+97 -9

@@ -3,6 +3,7 @@

3 import functools

4 from collections.abc import Callable

6 +import pandas as pd

7 import torch

8 from torch._ops import OpOverload

@@ -56,6 +57,29 @@ def is_aiter_found_and_supported() -> bool:

57 return False

60 +@functools.cache

61 +def _load_gemm_tuned_configs(

62 + q_dtype_w: torch.dtype, csv_path: str

63 +) -> set[tuple[int, int, int]]:

64 + try:

65 + df = pd.read_csv(csv_path).drop_duplicates()

66 + df = df[df["q_dtype_w"] == str(q_dtype_w)]

67 + return set(zip(df["N"].astype(int), df["K"].astype(int), df["M"].astype(int)))

68 + except Exception:

69 + return set()

70 +

71 +

72 +def _check_kernel_tuned(N: int, K: int, q_dtype_w: torch.dtype, csv_path: str) -> bool:

73 + configs = _load_gemm_tuned_configs(q_dtype_w, csv_path)

74 + l_m = (

75 + [1, 2, 4]

76 + + list(range(8, 513, 8))

77 + + [1024, 1536]

78 + + [2**i for i in range(11, 19)]

79 + )

80 + return any((N, K, M) in configs for M in l_m)

81 +

82 +

83 def if_aiter_supported(func: Callable) -> Callable:

84 """Decorator that only executes the function if

85 ROCm AITER package is supported and enabled on gfx9 archs.

@@ -468,7 +492,7 @@ def _rocm_aiter_mla_decode_fwd_fake(

492 pass

493

494

471 -def _rocm_aiter_gemm_a8w8_impl(

495 +def _rocm_aiter_w8a8_gemm_impl(

496 A: torch.Tensor,

497 B: torch.Tensor,

498 As: torch.Tensor,

@@ -485,7 +509,7 @@ def _rocm_aiter_gemm_a8w8_impl(

509 return gemm_a8w8_CK(A, B, As, Bs, bias, output_dtype)

510

511

488 -def _rocm_aiter_gemm_a8w8_fake(

512 +def _rocm_aiter_w8a8_gemm_fake(

513 A: torch.Tensor,

514 B: torch.Tensor,

515 As: torch.Tensor,

@@ -499,6 +523,35 @@ def _rocm_aiter_gemm_a8w8_fake(

523 return Y

524

525

526 +def _rocm_aiter_preshuffled_per_token_w8a8_gemm_impl(

527 + A: torch.Tensor,

528 + B: torch.Tensor,

529 + As: torch.Tensor,

530 + Bs: torch.Tensor,

531 + bias: torch.Tensor | None = None,

532 + output_dtype: torch.dtype = torch.float16,

533 +) -> torch.Tensor:

534 + from aiter import gemm_a8w8_bpreshuffle

535 +

536 + output = gemm_a8w8_bpreshuffle(A, B, As, Bs, None, output_dtype)

537 + if bias is not None:

538 + output.add_(bias)

539 + return output

540 +

541 +

542 +def _rocm_aiter_preshuffled_per_token_w8a8_gemm_fake(

543 + A: torch.Tensor,

544 + B: torch.Tensor,

545 + As: torch.Tensor,

546 + Bs: torch.Tensor,

547 + bias: torch.Tensor | None = None,

548 + output_dtype: torch.dtype = torch.float16,

549 +) -> torch.Tensor:

550 + m = A.shape[0]

551 + n = B.shape[0]

552 + return torch.empty(m, n, dtype=output_dtype, device=A.device)

553 +

554 +

555 def _rocm_aiter_triton_gemm_a8w8_blockscale_impl(

556 A: torch.Tensor,

557 B: torch.Tensor,

@@ -1313,11 +1366,15 @@ class rocm_aiter_ops:

1366 )

1367

1368 direct_register_custom_op(

1316 - op_name="rocm_aiter_gemm_a8w8",

1317 - op_func=_rocm_aiter_gemm_a8w8_impl,

1318 - mutates_args=[],

1319 - fake_impl=_rocm_aiter_gemm_a8w8_fake,

1320 - dispatch_key=current_platform.dispatch_key,

1369 + op_name="rocm_aiter_w8a8_gemm",

1370 + op_func=_rocm_aiter_w8a8_gemm_impl,

1371 + fake_impl=_rocm_aiter_w8a8_gemm_fake,

1372 + )

1373 +

1374 + direct_register_custom_op(

1375 + op_name="_rocm_aiter_preshuffled_per_token_w8a8_gemm",

1376 + op_func=_rocm_aiter_preshuffled_per_token_w8a8_gemm_impl,

1377 + fake_impl=_rocm_aiter_preshuffled_per_token_w8a8_gemm_fake,

1378 )

1379

1380 direct_register_custom_op(

@@ -1493,7 +1550,7 @@ class rocm_aiter_ops:

1550 )

1551

1552 @staticmethod

1496 - def gemm_a8w8(

1553 + def w8a8_gemm(

1554 A: torch.Tensor,

1555 B: torch.Tensor,

1556 As: torch.Tensor,

@@ -1501,7 +1558,20 @@ class rocm_aiter_ops:

1558 bias: torch.Tensor | None = None,

1559 output_dtype: torch.dtype = torch.float16,

1560 ) -> torch.Tensor:

1504 - return torch.ops.vllm.rocm_aiter_gemm_a8w8(A, B, As, Bs, bias, output_dtype)

1561 + return torch.ops.vllm.rocm_aiter_w8a8_gemm(A, B, As, Bs, bias, output_dtype)

1562 +

1563 + @staticmethod

1564 + def preshuffled_per_token_w8a8_gemm(

1565 + A: torch.Tensor,

1566 + B: torch.Tensor,

1567 + As: torch.Tensor,

1568 + Bs: torch.Tensor,

1569 + bias: torch.Tensor | None = None,

1570 + output_dtype: torch.dtype = torch.float16,

1571 + ) -> torch.Tensor:

1572 + return torch.ops.vllm._rocm_aiter_preshuffled_per_token_w8a8_gemm(

1573 + A, B, As, Bs, bias, output_dtype

1574 + )

1575

1576 @staticmethod

1577 def triton_gemm_a8w8_blockscale(

@@ -1920,6 +1990,24 @@ class rocm_aiter_ops:

1990 (8192, 3584),

1991 ]

1992

1993 + @staticmethod

1994 + def is_shuffled_per_token_w8a8_gemm_tuned(

1995 + N: int, K: int, q_dtype_w: torch.dtype

1996 + ) -> bool:

1997 + import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops

1998 +

1999 + csv_path = (

2000 + aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_BPRESHUFFLE_FILE

2001 + )

2002 + return _check_kernel_tuned(N, K, q_dtype_w, csv_path)

2003 +

2004 + @staticmethod

2005 + def is_per_token_w8a8_gemm_tuned(N: int, K: int, q_dtype_w: torch.dtype) -> bool:

2006 + import aiter.ops.gemm_op_a8w8 as aiter_gemm_a8w8_ops

2007 +

2008 + csv_path = aiter_gemm_a8w8_ops.AITER_CONFIGS.AITER_CONFIG_GEMM_A8W8_FILE

2009 + return _check_kernel_tuned(N, K, q_dtype_w, csv_path)

2010 +

2011 @staticmethod

2012 def shuffle_weight(

2013 tensor: torch.Tensor, layout: tuple[int, int] = (16, 16)

vllm/model_executor/kernels/linear/__init__.py

+9 -3

@@ -106,6 +106,8 @@ from vllm.model_executor.kernels.linear.scaled_mm import (

106 from vllm.model_executor.kernels.linear.scaled_mm.aiter import (

107 AiterFp8BlockScaledMMKernel,

108 AiterInt8ScaledMMLinearKernel,

109 + AiterPerTokenFp8ScaledMMLinearKernel,

110 + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,

111 )

112 from vllm.model_executor.kernels.linear.scaled_mm.cpu import (

113 CPUInt8ScaledMMLinearKernel,

@@ -165,6 +167,8 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =

167 ChannelWiseTorchFP8ScaledMMLinearKernel,

168 ],

169 PlatformEnum.ROCM: [

170 + AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,

171 + AiterPerTokenFp8ScaledMMLinearKernel,

172 ROCmFP8ScaledMMLinearKernel,

173 PerTensorTorchFP8ScaledMMLinearKernel,

174 RowWiseTorchFP8ScaledMMLinearKernel,

@@ -360,18 +364,18 @@ def choose_scaled_mm_linear_kernel(

364 def init_fp8_linear_kernel(

365 activation_quant_key: QuantKey,

366 weight_quant_key: QuantKey,

363 - weight_shape: tuple[int, int],

367 input_dtype: torch.dtype,

368 out_dtype: torch.dtype,

366 - force_kernel: type[_KernelT] | None = None,

369 + weight_shape: tuple[int, int],

370 + force_kernel: type[FP8ScaledMMLinearKernel] | None = None,

371 module_name: str | None = None,

372 ) -> FP8ScaledMMLinearKernel | Fp8BlockScaledMMLinearKernel:

373 scaled_mm_linear_kernel_config = FP8ScaledMMLinearLayerConfig(

374 weight_quant_key=weight_quant_key,

375 activation_quant_key=activation_quant_key,

372 - weight_shape=weight_shape,

376 input_dtype=input_dtype,

377 out_dtype=out_dtype,

378 + weight_shape=weight_shape,

379 )

380

381 if activation_quant_key.scale.group_shape.is_per_group():

@@ -725,6 +729,8 @@ __all__ = [

729 "FP8ScaledMMLinearLayerConfig",

730 "Int8ScaledMMLinearLayerConfig",

731 "ScaledMMLinearLayerConfig",

732 + "AiterPreshuffledPerTokenFp8ScaledMMLinearKernel",

733 + "AiterPerTokenFp8ScaledMMLinearKernel",

734 "NvFp4LinearKernel",

735 "NvFp4LinearLayerConfig",

736 "AiterInt8ScaledMMLinearKernel",

vllm/model_executor/kernels/linear/scaled_mm/aiter.py

+160 -4

@@ -5,18 +5,27 @@

5 import torch

7 from vllm import _custom_ops as ops

8 -from vllm._aiter_ops import rocm_aiter_ops

8 +from vllm._aiter_ops import (

9 + rocm_aiter_ops,

10 +)

11 +from vllm.logger import init_logger

12 from vllm.model_executor.layers.quantization.utils.quant_utils import (

13 GroupShape,

14 )

15 +from vllm.model_executor.utils import replace_parameter

16 from vllm.platforms import current_platform

18 from .BlockScaledMMLinearKernel import (

19 Fp8BlockScaledMMLinearKernel,

16 - FP8ScaledMMLinearLayerConfig,

20 )

21 from .cutlass import CutlassInt8ScaledMMLinearKernel

19 -from .ScaledMMLinearKernel import Int8ScaledMMLinearLayerConfig

22 +from .ScaledMMLinearKernel import (

23 + FP8ScaledMMLinearKernel,

24 + FP8ScaledMMLinearLayerConfig,

25 + Int8ScaledMMLinearLayerConfig,

26 +)

27 +

28 +logger = init_logger(__name__)

31 class AiterInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):

@@ -113,7 +122,154 @@ class AiterInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):

122 # a to be [M, K]

123 # b to be [N, K]

124 # CutlassInt8ScaledMMLinearKernel prepare weight `w_q` in [K, N] format

116 - return rocm_aiter_ops.gemm_a8w8(x_q, w_q.t(), x_s, w_s, bias, out_dtype)

125 + return rocm_aiter_ops.w8a8_gemm(x_q, w_q.t(), x_s, w_s, bias, out_dtype)

126 +

127 +

128 +class AiterPreshuffledPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):

129 + @classmethod

130 + def is_supported(

131 + cls, compute_capability: int | None = None

132 + ) -> tuple[bool, str | None]:

133 + if not current_platform.is_rocm():

134 + return False, "requires ROCm."

135 + if not rocm_aiter_ops.is_linear_fp8_enabled():

136 + return (

137 + False,

138 + "requires setting `VLLM_ROCM_USE_AITER=1` "

139 + "and `VLLM_ROCM_USE_AITER_LINEAR=1`. "

140 + "`VLLM_ROCM_USE_AITER_LINEAR` default is True.",

141 + )

142 + try:

143 + import aiter # noqa: F401

144 + except Exception:

145 + return False, "requires aiter library to be installed."

146 + return True, None

147 +

148 + @classmethod

149 + def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:

150 + is_ptpc = (

151 + c.activation_quant_key.scale.group_shape.is_per_token()

152 + and c.weight_quant_key.scale.group_shape.is_per_channel()

153 + )

154 + if c.weight_shape is None:

155 + return False, "weight_shape is required for Aiter kernels"

156 + N, K = c.weight_shape

157 + fp8_dtype = current_platform.fp8_dtype()

158 +

159 + if c.out_dtype is not torch.bfloat16:

160 + return False, "requires bfloat16 output dtype."

161 +

162 + if not is_ptpc:

163 + return (

164 + False,

165 + "requires per token activation scales and per channel weight scales.",

166 + )

167 +

168 + if not (N % 16 == 0 and K % 16 == 0):

169 + return (

170 + False,

171 + f"requires N and K dimensions divisible by 16, received "

172 + f"N={N} and K={K}.",

173 + )

174 +

175 + # Aiter's shuffled per-token Gemm performs better than torch only when its

176 + # tuned.

177 + if not rocm_aiter_ops.is_shuffled_per_token_w8a8_gemm_tuned(N, K, fp8_dtype):

178 + return (

179 + False,

180 + f"requires a tuned configuration for N: {N} and K: {K} "

181 + f"and fp8 dtype {fp8_dtype}.",

182 + )

183 +

184 + return True, None

185 +

186 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

187 + w_name, *_ = self.layer_param_names

188 + w, *_ = self._get_layer_params(layer)

189 +

190 + replace_parameter(

191 + layer,

192 + w_name,

193 + torch.nn.Parameter(

194 + rocm_aiter_ops.shuffle_weight(w.t().contiguous()).data,

195 + requires_grad=False,

196 + ),

197 + )

198 +

199 + def apply_scaled_mm(

200 + self,

201 + *,

202 + A: torch.Tensor,

203 + B: torch.Tensor,

204 + out_dtype: torch.dtype,

205 + As: torch.Tensor,

206 + Bs: torch.Tensor,

207 + bias: torch.Tensor | None,

208 + output_shape: list,

209 + ) -> torch.Tensor:

210 + return rocm_aiter_ops.preshuffled_per_token_w8a8_gemm(

211 + A, B, As, Bs, bias, out_dtype

212 + )

213 +

214 +

215 +class AiterPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):

216 + @classmethod

217 + def is_supported(

218 + cls, compute_capability: int | None = None

219 + ) -> tuple[bool, str | None]:

220 + return AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.is_supported(

221 + compute_capability

222 + )

223 +

224 + @classmethod

225 + def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:

226 + is_ptpc = (

227 + c.activation_quant_key.scale.group_shape.is_per_token()

228 + and c.weight_quant_key.scale.group_shape.is_per_channel()

229 + )

230 + if c.weight_shape is None:

231 + return False, "weight_shape is required for Aiter kernels"

232 + N, K = c.weight_shape

233 + fp8_dtype = current_platform.fp8_dtype()

234 +

235 + if not is_ptpc:

236 + return (

237 + False,

238 + "requires per token activation scales and per channel weight scales.",

239 + )

240 +

241 + # Aiter's per-token Gemm performs better than torch oonly when its

242 + # tuned.

243 + if not rocm_aiter_ops.is_per_token_w8a8_gemm_tuned(N, K, fp8_dtype):

244 + return (

245 + False,

246 + f"requires a tuned configuration for N: {N} and K: {K} "

247 + f"and fp8 dtype {fp8_dtype}.",

248 + )

249 + return True, None

250 +

251 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

252 + w_name, *_ = self.layer_param_names

253 + w, *_ = self._get_layer_params(layer)

254 +

255 + replace_parameter(

256 + layer,

257 + w_name,

258 + torch.nn.Parameter(w.t(), requires_grad=False),

259 + )

260 +

261 + def apply_scaled_mm(

262 + self,

263 + *,

264 + A: torch.Tensor,

265 + B: torch.Tensor,

266 + out_dtype: torch.dtype,

267 + As: torch.Tensor,

268 + Bs: torch.Tensor,

269 + bias: torch.Tensor | None,

270 + output_shape: list,

271 + ) -> torch.Tensor:

272 + return rocm_aiter_ops.w8a8_gemm(A, B, As, Bs, bias, out_dtype)

273

274

275 class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py

+4 -5

@@ -31,8 +31,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

31 GroupShape,

32 create_fp8_quant_key,

33 kFp8DynamicTokenSym,

34 + kFp8StaticChannelSym,

35 kFp8StaticTensorSym,

35 - kFp8StaticTokenSym,

36 )

37 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

38 cutlass_block_fp8_supported,

@@ -47,7 +47,7 @@ activation_quant_key_mapping = {

47 DYNAMIC_QUANT: kFp8DynamicTokenSym,

48 }

49 weight_quant_key_mapping = {

50 - QuantizationStrategy.CHANNEL: kFp8StaticTokenSym,

50 + QuantizationStrategy.CHANNEL: kFp8StaticChannelSym,

51 QuantizationStrategy.TENSOR: kFp8StaticTensorSym,

52 }

53 logger = init_logger(__name__)

@@ -67,7 +67,6 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):

67 self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enabled()

68 assert not self.is_static_input_scheme

69 self.act_q_group_shape = GroupShape(1, self.weight_block_size[0])

70 -

70 self.weight_quant_key = create_fp8_quant_key(

71 static=True, group_shape=GroupShape(*self.weight_block_size)

72 )

@@ -76,7 +75,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):

75 )

76 else:

77 self.activation_quant_key = activation_quant_key_mapping[

79 - is_static_input_scheme

78 + self.is_static_input_scheme

79 ]

80 self.weight_quant_key = weight_quant_key_mapping[self.strategy]

@@ -138,9 +137,9 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):

137 self.fp8_linear = init_fp8_linear_kernel(

138 activation_quant_key=self.activation_quant_key,

139 weight_quant_key=self.weight_quant_key,

141 - weight_shape=layer.weight.shape,

140 input_dtype=self.input_dtype,

141 out_dtype=self.out_dtype,

142 + weight_shape=(output_size_per_partition, input_size_per_partition),

143 module_name=self.__class__.__name__,

144 )

145

vllm/model_executor/layers/quantization/fbgemm_fp8.py

@@ -175,6 +175,8 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):

175 # Activations not quantized for marlin.

176 del layer.input_scale_ub

177

178 + self.fp8_linear.process_weights_after_loading(layer)

179 +

180 def apply(

181 self,

182 layer: torch.nn.Module,

vllm/model_executor/layers/quantization/fp8.py

+2 -2

@@ -397,8 +397,6 @@ class Fp8LinearMethod(LinearMethodBase):

397 if self.block_quant:

398 assert not self.act_q_static

399

400 - self.fp8_linear.process_weights_after_loading(layer)

401 -

400 # If checkpoint not serialized fp8, quantize the weights.

401 else:

402 # If checkpoint is fp8 per-tensor, handle that there are N scales for N

@@ -428,6 +426,8 @@ class Fp8LinearMethod(LinearMethodBase):

426 else:

427 layer.input_scale = None

428

429 + self.fp8_linear.process_weights_after_loading(layer)

430 +

431 def apply(

432 self,

433 layer: torch.nn.Module,

vllm/model_executor/layers/quantization/modelopt.py

@@ -517,6 +517,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase):

517 layer.weight = Parameter(weight.t(), requires_grad=False)

518 layer.weight_scale = Parameter(max_w_scale, requires_grad=False)

519 layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)

520 + self.fp8_linear.process_weights_after_loading(layer)

521

522 def apply(

523 self,

@@ -597,6 +598,7 @@ class ModelOptFp8PcPtLinearMethod(LinearMethodBase):

598 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

599 layer.weight = Parameter(layer.weight.t(), requires_grad=False)

600 layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)

601 + self.fp8_linear.process_weights_after_loading(layer)

602

603 def apply(

604 self,

vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py

@@ -120,6 +120,8 @@ class QuarkW8A8Fp8(QuarkScheme):

120 if self.is_static_input_scheme:

121 layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)

122

123 + self.fp8_linear.process_weights_after_loading(layer)

124 +

125 def create_weights(

126 self,

127 layer: torch.nn.Module,