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 6, 19:15 UTC
[ROCm][Quantization] add quark w4a8 mxfp4_fp8 for LinearLayer (#35316)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Mar 13, 2026 at 15:44 UTC
6341d43043517a431d49b9c571fc5fa8afe182cb
4 files changed
+311
-1
vllm/_aiter_ops.py
+53
@@ -861,6 +861,39 @@ def _rocm_aiter_triton_add_rmsnorm_pad_fake(
861
return out, residual_out
862
863
864
+def _rocm_aiter_gemm_a8wfp4_impl(
865
+ x: torch.Tensor,
866
+ w: torch.Tensor,
867
+ x_scales: torch.Tensor,
868
+ w_scales: torch.Tensor,
869
+ out_dtype: torch.dtype,
870
+) -> torch.Tensor:
871
+ from aiter.ops.triton.gemm_a8wfp4 import gemm_a8wfp4
872
+
873
+ M, N = x.shape[0], w.shape[0]
874
+ y = torch.empty(M, N, dtype=out_dtype, device=x.device)
875
+ gemm_a8wfp4(
876
+ x=x,
877
+ w=w,
878
+ y=y,
879
+ x_scales=x_scales,
880
+ w_scales=w_scales,
881
+ dtype=out_dtype,
882
+ config=None,
883
+ )
884
+ return y
885
+
886
+
887
+def _rocm_aiter_gemm_a8wfp4_fake(
888
+ x: torch.Tensor,
889
+ w: torch.Tensor,
890
+ x_scales: torch.Tensor,
891
+ w_scales: torch.Tensor,
892
+ out_dtype: torch.dtype,
893
+) -> torch.Tensor:
894
+ return torch.empty(x.shape[0], w.shape[0], dtype=out_dtype, device=x.device)
895
+
896
+
897
def _triton_rotary_embedding_impl(
898
positions: torch.Tensor,
899
query: torch.Tensor,
@@ -1337,6 +1370,14 @@ class rocm_aiter_ops:
1370
dispatch_key=current_platform.dispatch_key,
1371
)
1372
1373
+ direct_register_custom_op(
1374
+ op_name="rocm_aiter_gemm_a8wfp4",
1375
+ op_func=_rocm_aiter_gemm_a8wfp4_impl,
1376
+ mutates_args=[],
1377
+ fake_impl=_rocm_aiter_gemm_a8wfp4_fake,
1378
+ dispatch_key=current_platform.dispatch_key,
1379
+ )
1380
+
1381
# Register rocm aiter rotary embedding custom op
1382
direct_register_custom_op(
1383
op_name="rocm_aiter_triton_rotary_embedding",
@@ -1646,6 +1687,18 @@ class rocm_aiter_ops:
1687
) -> tuple[torch.Tensor, torch.Tensor]:
1688
return torch.ops.vllm.rocm_aiter_per_token_quant(x, quant_dtype, scale)
1689
1690
+ @staticmethod
1691
+ def gemm_a8wfp4(
1692
+ x: torch.Tensor,
1693
+ w: torch.Tensor,
1694
+ x_scales: torch.Tensor,
1695
+ w_scales: torch.Tensor,
1696
+ out_dtype: torch.dtype,
1697
+ ) -> torch.Tensor:
1698
+ return torch.ops.vllm.rocm_aiter_gemm_a8wfp4(
1699
+ x, w, x_scales, w_scales, out_dtype
1700
+ )
1701
+
1702
@staticmethod
1703
def triton_fp4_gemm_dynamic_qaunt(
1704
x: torch.Tensor,
vllm/model_executor/layers/quantization/quark/quark.py
+32
@@ -26,6 +26,7 @@ from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E
26
from vllm.model_executor.layers.quantization.quark.schemes import (
27
QuarkOCP_MX,
28
QuarkScheme,
29
+ QuarkW4A8_MXFP4_FP8,
30
QuarkW8A8Fp8,
31
QuarkW8A8Int8,
32
)
@@ -350,6 +351,31 @@ class QuarkConfig(QuantizationConfig):
351
# Only symmetric weight quantization supported.
352
return is_int8_dtype and is_tensor and is_weight_symmetric and is_static
353
354
+ def _is_w4a8_mxfp4_fp8(
355
+ self,
356
+ weight_quant: dict[str, Any] | None,
357
+ input_quant: dict[str, Any] | None,
358
+ ) -> bool:
359
+ if weight_quant is None or input_quant is None:
360
+ return False
361
+
362
+ is_weight_mxfp4 = (
363
+ weight_quant.get("dtype") == "fp4"
364
+ and weight_quant.get("qscheme") == "per_group"
365
+ and weight_quant.get("group_size") == 32
366
+ and weight_quant.get("scale_format") == "e8m0"
367
+ and not weight_quant.get("is_dynamic")
368
+ )
369
+
370
+ is_input_fp8 = (
371
+ input_quant.get("dtype") == "fp8_e4m3"
372
+ and input_quant.get("qscheme") == "per_tensor"
373
+ and not input_quant.get("is_dynamic") # Static per-tensor
374
+ and input_quant.get("symmetric") is True # Symmetric quantization
375
+ )
376
+
377
+ return is_weight_mxfp4 and is_input_fp8
378
+
379
def _is_w_ocp_mx_a_x(
380
self, weight_quant: dict[str, Any] | None, input_quant: dict[str, Any] | None
381
) -> bool:
@@ -504,6 +530,12 @@ class QuarkConfig(QuantizationConfig):
530
is_static_input_scheme=True,
531
input_symmetric=input_config.get("symmetric"),
532
)
533
+ elif self._is_w4a8_mxfp4_fp8(weight_config, input_config):
534
+ is_w4a8_supported = self._check_scheme_supported(
535
+ QuarkW4A8_MXFP4_FP8.get_min_capability(), error=False
536
+ )
537
+ if is_w4a8_supported:
538
+ return QuarkW4A8_MXFP4_FP8(weight_config, input_config)
539
elif self._is_w_ocp_mx_a_x(weight_config, input_config):
540
return QuarkOCP_MX(
541
weight_config, input_config, dynamic_mxfp4_quant=dynamic_mxfp4_quant
vllm/model_executor/layers/quantization/quark/schemes/__init__.py
+8
-1
@@ -3,7 +3,14 @@
3
4
from .quark_ocp_mx import QuarkOCP_MX
5
from .quark_scheme import QuarkScheme
6
+from .quark_w4a8_mxfp4_fp8 import QuarkW4A8_MXFP4_FP8
7
from .quark_w8a8_fp8 import QuarkW8A8Fp8
8
from .quark_w8a8_int8 import QuarkW8A8Int8
9
9
-__all__ = ["QuarkScheme", "QuarkW8A8Fp8", "QuarkW8A8Int8", "QuarkOCP_MX"]
10
+__all__ = [
11
+ "QuarkScheme",
12
+ "QuarkW8A8Fp8",
13
+ "QuarkW8A8Int8",
14
+ "QuarkOCP_MX",
15
+ "QuarkW4A8_MXFP4_FP8",
16
+]
vllm/model_executor/layers/quantization/quark/schemes/quark_w4a8_mxfp4_fp8.py
new
+218
@@ -0,0 +1,218 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Callable
5
+from fractions import Fraction
6
+from typing import Any
7
+
8
+import torch
9
+import torch.nn.functional as F
10
+
11
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
12
+from vllm.logger import init_logger
13
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
14
+ get_fp8_min_max,
15
+)
16
+from vllm.model_executor.parameter import (
17
+ GroupQuantScaleParameter,
18
+ PackedvLLMParameter,
19
+ PerTensorScaleParameter,
20
+)
21
+from vllm.platforms import current_platform
22
+
23
+from .quark_scheme import QuarkScheme
24
+
25
+logger = init_logger(__name__)
26
+
27
+
28
+__all__ = ["QuarkW4A8_MXFP4_FP8"]
29
+
30
+OCP_MX_BLOCK_SIZE = 32
31
+
32
+
33
+class QuarkW4A8_MXFP4_FP8(QuarkScheme):
34
+ """
35
+ - Weights: MXFP4 with E8M0 scales per block of 32
36
+ - Activations: FP8 E4M3 (static per-tensor quantization)
37
+
38
+ Uses the AITER Triton kernel and falls back to emulation if AITER not available.
39
+ """
40
+
41
+ def __init__(
42
+ self,
43
+ weight_quant_spec: dict[str, Any],
44
+ input_quant_spec: dict[str, Any],
45
+ ):
46
+ self.out_dtype = None
47
+
48
+ self.weight_dtype = "mxfp4"
49
+ self.packed_factor: Fraction = Fraction(2, 1) # 2 FP4 values per byte
50
+ self.weight_block_size = OCP_MX_BLOCK_SIZE
51
+
52
+ self.is_static_input_scheme = not input_quant_spec.get("is_dynamic")
53
+ self.input_qscheme = input_quant_spec.get("qscheme") # "per_tensor"
54
+
55
+ self.fp8_min, self.fp8_max = get_fp8_min_max()
56
+ self.fp8_dtype = current_platform.fp8_dtype()
57
+
58
+ if not self.is_static_input_scheme:
59
+ raise NotImplementedError(
60
+ "Dynamic FP8 activation quantization is not yet supported "
61
+ "for W4A8. The current implementation expects static per-tensor "
62
+ "FP8 scales stored in the checkpoint."
63
+ )
64
+
65
+ kernel_supported_gpu = False
66
+ if current_platform.is_rocm():
67
+ from vllm.platforms.rocm import on_gfx950
68
+
69
+ kernel_supported_gpu = on_gfx950()
70
+
71
+ self.use_aiter_kernel = (
72
+ is_aiter_found_and_supported()
73
+ and self.is_static_input_scheme
74
+ and kernel_supported_gpu
75
+ )
76
+
77
+ if not self.use_aiter_kernel:
78
+ logger.warning_once(
79
+ "[W4A8 MXFP4+FP8] Aiter Triton kernel not found. Using emulation mode."
80
+ )
81
+
82
+ @classmethod
83
+ def get_min_capability(cls) -> int:
84
+ return 70
85
+
86
+ def get_packed_dim(self, dim: int) -> int:
87
+ assert dim % 2 == 0, f"Dimension {dim} must be even for MXFP4 packing"
88
+ return dim // 2
89
+
90
+ def create_weights(
91
+ self,
92
+ layer: torch.nn.Module,
93
+ output_partition_sizes: list[int],
94
+ input_size_per_partition: int,
95
+ params_dtype: torch.dtype,
96
+ weight_loader: Callable,
97
+ **kwargs,
98
+ ):
99
+ output_size_per_partition = sum(output_partition_sizes)
100
+ layer.logical_widths = output_partition_sizes
101
+ layer.input_size_per_partition = input_size_per_partition
102
+ layer.output_size_per_partition = output_size_per_partition
103
+
104
+ # MXFP4 WEIGHT (packed, 2 values per byte)
105
+ weight = PackedvLLMParameter(
106
+ data=torch.empty(
107
+ output_size_per_partition,
108
+ self.get_packed_dim(input_size_per_partition),
109
+ dtype=torch.uint8,
110
+ ),
111
+ input_dim=1,
112
+ output_dim=0,
113
+ packed_dim=1,
114
+ packed_factor=self.packed_factor,
115
+ weight_loader=weight_loader,
116
+ )
117
+ layer.register_parameter("weight", weight)
118
+
119
+ # WEIGHT SCALE (E8M0 format, per block of 32)
120
+ weight_scale = GroupQuantScaleParameter(
121
+ data=torch.empty(
122
+ output_size_per_partition,
123
+ input_size_per_partition // self.weight_block_size,
124
+ dtype=torch.uint8,
125
+ ),
126
+ input_dim=1,
127
+ output_dim=0,
128
+ weight_loader=weight_loader,
129
+ )
130
+ layer.register_parameter("weight_scale", weight_scale)
131
+
132
+ # INPUT SCALE (FP8 per-tensor static scale)
133
+ if self.is_static_input_scheme:
134
+ input_scale = PerTensorScaleParameter(
135
+ data=torch.empty(
136
+ len(output_partition_sizes),
137
+ dtype=torch.float32,
138
+ ),
139
+ weight_loader=weight_loader,
140
+ )
141
+ # Initialize to avoid NaN
142
+ input_scale[:] = torch.finfo(torch.float32).min
143
+ layer.register_parameter("input_scale", input_scale)
144
+
145
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
146
+ # Ensuring weights & scales are non-trainable
147
+ layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False)
148
+ layer.weight_scale = torch.nn.Parameter(
149
+ layer.weight_scale.data, requires_grad=False
150
+ )
151
+
152
+ if self.is_static_input_scheme:
153
+ input_scale = layer.input_scale.data
154
+ # For fused modules (QKV), take the max scale
155
+ if input_scale.numel() != 1:
156
+ input_scale = input_scale.max()
157
+
158
+ layer.input_scale = torch.nn.Parameter(
159
+ torch.tensor(input_scale, dtype=torch.float32),
160
+ requires_grad=False,
161
+ )
162
+
163
+ def apply_weights(
164
+ self,
165
+ layer: torch.nn.Module,
166
+ x: torch.Tensor,
167
+ bias: torch.Tensor | None = None,
168
+ ) -> torch.Tensor:
169
+ if self.use_aiter_kernel:
170
+ return self._apply_aiter_kernel(layer, x, bias)
171
+ else:
172
+ return self._apply_emulation(layer, x, bias)
173
+
174
+ def _apply_aiter_kernel(
175
+ self,
176
+ layer: torch.nn.Module,
177
+ x: torch.Tensor,
178
+ bias: torch.Tensor | None = None,
179
+ ) -> torch.Tensor:
180
+ M = x.shape[0]
181
+ out_dtype = x.dtype if self.out_dtype is None else self.out_dtype
182
+
183
+ input_scale = layer.input_scale
184
+ x_fp8 = (x / input_scale).clamp(self.fp8_min, self.fp8_max).to(self.fp8_dtype)
185
+
186
+ # Broadcast per-tensor scale to per-row (M, 1) for Aiter kernel
187
+ x_scales = input_scale.expand(M, 1).to(dtype=torch.float32, device=x.device)
188
+
189
+ y = rocm_aiter_ops.gemm_a8wfp4(
190
+ x_fp8, layer.weight, x_scales, layer.weight_scale, out_dtype
191
+ )
192
+
193
+ if bias is not None:
194
+ y = y + bias
195
+
196
+ return y
197
+
198
+ def _apply_emulation(
199
+ self,
200
+ layer: torch.nn.Module,
201
+ x: torch.Tensor,
202
+ bias: torch.Tensor | None = None,
203
+ ) -> torch.Tensor:
204
+ from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
205
+ dequant_mxfp4,
206
+ )
207
+
208
+ weight_dq = dequant_mxfp4(
209
+ layer.weight,
210
+ layer.weight_scale,
211
+ x.dtype,
212
+ )
213
+
214
+ input_scale = layer.input_scale
215
+ x_fp8 = (x / input_scale).clamp(self.fp8_min, self.fp8_max).to(self.fp8_dtype)
216
+ x_dq = (x_fp8.to(x.dtype) * input_scale).to(x.dtype)
217
+
218
+ return F.linear(x_dq, weight_dq, bias)