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 11, 18:45 UTC
[ROCm] Aiter Quant Kernels (#25552)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
vllmellm committed
Dec 9, 2025 at 22:27 UTC
ee14644ba9a3696c83ede2c948b73ebc3e1ffb33
3 files changed
+123
-2
vllm/_aiter_ops.py
+87
@@ -9,6 +9,8 @@ import vllm.envs as envs
9
from vllm.platforms import current_platform
10
from vllm.utils.torch_utils import direct_register_custom_op, is_torch_equal_or_newer
11
12
+_FP8_DTYPE = current_platform.fp8_dtype()
13
+
14
15
def is_aiter_found() -> bool:
16
from importlib.util import find_spec
@@ -467,6 +469,59 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_fake(
469
return torch.empty_like(x), torch.empty_like(residual)
470
471
472
+def _rocm_aiter_per_tensor_quant_impl(
473
+ x: torch.Tensor,
474
+ quant_dtype: torch.dtype,
475
+ scale: torch.Tensor | None = None,
476
+) -> tuple[torch.Tensor, torch.Tensor]:
477
+ from aiter.ops.quant import per_tensor_quant_hip
478
+
479
+ return per_tensor_quant_hip(x, scale, quant_dtype)
480
+
481
+
482
+def _rocm_aiter_per_tensor_quant_fake(
483
+ x: torch.Tensor,
484
+ quant_dtype: torch.dtype,
485
+ scale: torch.Tensor | None = None,
486
+) -> tuple[torch.Tensor, torch.Tensor]:
487
+ return torch.empty_like(x, dtype=quant_dtype), torch.empty(
488
+ 1, dtype=torch.float32, device=x.device
489
+ )
490
+
491
+
492
+def _rocm_aiter_per_token_quant_impl(
493
+ x: torch.Tensor, quant_dtype: torch.dtype, scale: torch.Tensor | None = None
494
+) -> tuple[torch.Tensor, torch.Tensor]:
495
+ from aiter.ops.quant import dynamic_per_token_scaled_quant
496
+
497
+ assert quant_dtype in [torch.int8, _FP8_DTYPE]
498
+
499
+ out_shape = x.shape
500
+ out = torch.empty(x.shape, dtype=_FP8_DTYPE, device=x.device)
501
+ if scale is None:
502
+ scale = torch.empty((*out_shape[:-1], 1), dtype=torch.float32, device=x.device)
503
+ dynamic_per_token_scaled_quant(
504
+ out,
505
+ x,
506
+ scale,
507
+ scale_ub=None,
508
+ shuffle_scale=False,
509
+ num_rows=None,
510
+ num_rows_factor=1,
511
+ )
512
+ return out, scale
513
+
514
+
515
+def _rocm_aiter_per_token_quant_fake(
516
+ x: torch.Tensor, quant_dtype: torch.dtype, scale: torch.Tensor | None = None
517
+) -> tuple[torch.Tensor, torch.Tensor]:
518
+ out_shape = x.shape
519
+ return (
520
+ torch.empty(x.shape, dtype=_FP8_DTYPE, device=x.device),
521
+ torch.empty((*out_shape[:-1], 1), dtype=torch.float32, device=x.device),
522
+ )
523
+
524
+
525
# Global flag to ensure ops are registered only once
526
_OPS_REGISTERED = False
527
@@ -665,6 +720,22 @@ class rocm_aiter_ops:
720
dispatch_key=current_platform.dispatch_key,
721
)
722
723
+ direct_register_custom_op(
724
+ op_name="rocm_aiter_per_tensor_quant",
725
+ op_func=_rocm_aiter_per_tensor_quant_impl,
726
+ mutates_args=[],
727
+ fake_impl=_rocm_aiter_per_tensor_quant_fake,
728
+ dispatch_key=current_platform.dispatch_key,
729
+ )
730
+
731
+ direct_register_custom_op(
732
+ op_name="rocm_aiter_per_token_quant",
733
+ op_func=_rocm_aiter_per_token_quant_impl,
734
+ mutates_args=["scale"],
735
+ fake_impl=_rocm_aiter_per_token_quant_fake,
736
+ dispatch_key=current_platform.dispatch_key,
737
+ )
738
+
739
_OPS_REGISTERED = True
740
741
@staticmethod
@@ -859,6 +930,22 @@ class rocm_aiter_ops:
930
kv_scale=kv_scale,
931
)
932
933
+ @staticmethod
934
+ def per_tensor_quant(
935
+ x: torch.Tensor,
936
+ quant_dtype: torch.dtype,
937
+ scale: torch.Tensor | None = None,
938
+ ) -> tuple[torch.Tensor, torch.Tensor]:
939
+ return torch.ops.vllm.rocm_aiter_per_tensor_quant(x, quant_dtype, scale)
940
+
941
+ @staticmethod
942
+ def per_token_quant(
943
+ x: torch.Tensor,
944
+ quant_dtype: torch.dtype,
945
+ scale: torch.Tensor | None = None,
946
+ ) -> tuple[torch.Tensor, torch.Tensor]:
947
+ return torch.ops.vllm.rocm_aiter_per_token_quant(x, quant_dtype, scale)
948
+
949
@staticmethod
950
def triton_fp4_gemm_dynamic_qaunt(
951
x: torch.Tensor,
vllm/model_executor/layers/quantization/input_quant_fp8.py
+31
@@ -5,6 +5,7 @@ import torch
5
import torch.nn.functional as F
6
7
from vllm import _custom_ops as ops
8
+from vllm._aiter_ops import rocm_aiter_ops
9
from vllm.model_executor.custom_op import CustomOp
10
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
11
from vllm.platforms import current_platform
@@ -45,10 +46,13 @@ class QuantFP8(CustomOp):
46
super().__init__()
47
self.static = static
48
self.group_shape = group_shape
49
+ self.use_per_token_if_dynamic = group_shape == GroupShape.PER_TOKEN
50
self.num_token_padding = num_token_padding
51
self.column_major_scales = column_major_scales
52
self.use_ue8m0 = use_ue8m0
53
54
+ self.use_aiter = rocm_aiter_ops.is_linear_fp8_enaled()
55
+
56
self.is_group_quant = group_shape.is_per_group()
57
if self.is_group_quant:
58
assert not static, "Group quantization only supports dynamic mode"
@@ -92,6 +96,33 @@ class QuantFP8(CustomOp):
96
use_per_token_if_dynamic=self.use_per_token_if_dynamic,
97
)
98
99
+ def forward_hip(
100
+ self,
101
+ x: torch.Tensor,
102
+ scale: torch.Tensor | None = None,
103
+ scale_ub: torch.Tensor | None = None,
104
+ ) -> tuple[torch.Tensor, torch.Tensor]:
105
+ use_aiter_quant = (
106
+ not self.is_group_quant
107
+ and self.use_aiter
108
+ and scale_ub is None
109
+ and x.is_contiguous()
110
+ )
111
+ use_aiter_per_tensor_quant = (
112
+ use_aiter_quant and self.group_shape == GroupShape.PER_TENSOR
113
+ )
114
+ use_aiter_per_token_quant = (
115
+ use_aiter_quant and self.group_shape == GroupShape.PER_TOKEN
116
+ )
117
+
118
+ if use_aiter_per_tensor_quant:
119
+ return rocm_aiter_ops.per_tensor_quant(x, _FP8_DTYPE, scale)
120
+ if use_aiter_per_token_quant:
121
+ return rocm_aiter_ops.per_token_quant(x, _FP8_DTYPE, scale)
122
+
123
+ # Fallback to CUDA implementation
124
+ return self.forward_cuda(x, scale, scale_ub)
125
+
126
def forward_native(
127
self,
128
x: torch.Tensor,
vllm/platforms/rocm.py
+5
-2
@@ -381,6 +381,8 @@ class RocmPlatform(Platform):
381
compilation_config = vllm_config.compilation_config
382
parallel_config = vllm_config.parallel_config
383
is_eager_execution = compilation_config == CUDAGraphMode.NONE
384
+ use_aiter_rms_norm = rocm_aiter_ops.is_rmsnorm_enabled()
385
+ use_aiter_fp8_linear = rocm_aiter_ops.is_linear_fp8_enaled()
386
387
if compilation_config.cudagraph_mode.has_full_cudagraphs():
388
# decode context parallel does not support full cudagraphs
@@ -400,8 +402,6 @@ class RocmPlatform(Platform):
402
)
403
compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
404
403
- use_aiter_rms_norm = rocm_aiter_ops.is_rmsnorm_enabled()
404
-
405
if cache_config and cache_config.block_size is None:
406
cache_config.block_size = 16
407
@@ -415,6 +415,9 @@ class RocmPlatform(Platform):
415
):
416
compilation_config.custom_ops.append("+rms_norm")
417
418
+ if use_aiter_fp8_linear and "-quant_fp8" not in compilation_config.custom_ops:
419
+ compilation_config.custom_ops.append("+quant_fp8")
420
+
421
@classmethod
422
def verify_model_arch(cls, model_arch: str) -> None:
423
if model_arch in _ROCM_UNSUPPORTED_MODELS: