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] [quantization] Fix quark ptpc moe and add test case (#24649)
Signed-off-by: Haoyang Li <lihaoyang0109@gmail.com> Co-authored-by: Haoyang Li <haoyang.li@amd.com>
haoyangli-amd committed
Sep 17, 2025 at 13:15 UTC
ca2d1925ef5ad309061c2d5dd9a1e409c5ca28ee
2 files changed
+196
-52
tests/quantization/test_quark.py
+25
@@ -77,6 +77,31 @@ def test_quark_fp8_w_per_tensor_a_per_tensor(vllm_runner, kv_cache_dtype, tp):
77
assert output
78
79
80
+@pytest.mark.parametrize('tp', [1])
81
+def test_quark_fp8_w_per_channel_a_per_token(vllm_runner, tp):
82
+ model_path = "amd/Qwen2.5-1.5B-Instruct-ptpc-Quark-ts"
83
+ with vllm_runner(model_path, tensor_parallel_size=tp) as llm:
84
+
85
+ def check_model(model):
86
+ layer = model.model.layers[0]
87
+
88
+ qkv_proj = layer.self_attn.qkv_proj
89
+
90
+ assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
91
+ assert isinstance(qkv_proj.scheme, QuarkW8A8Fp8)
92
+
93
+ if isinstance(qkv_proj.scheme, QuarkW8A8Fp8):
94
+ assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
95
+ assert qkv_proj.weight_scale.shape[0] == qkv_proj.weight.shape[
96
+ 1]
97
+ assert qkv_proj.weight_scale.shape[1] == 1
98
+
99
+ llm.apply_model(check_model)
100
+
101
+ output = llm.generate_greedy("Hello my name is", max_tokens=20)
102
+ assert output
103
+
104
+
105
@pytest.mark.parametrize('tp', [1])
106
def test_quark_int8_w_per_tensor_a_per_tensor(vllm_runner, tp):
107
model_path = "amd/Llama-3.1-8B-Instruct-w-int8-a-int8-sym-test"
vllm/model_executor/layers/quantization/quark/quark_moe.py
+171
-52
@@ -5,17 +5,25 @@ from typing import Any, Callable, Optional, Union
5
6
import torch
7
8
+import vllm.envs as envs
9
from vllm import _custom_ops as ops
10
from vllm.logger import init_logger
11
from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
12
FusedMoEMethodBase,
13
FusedMoeWeightScaleSupported)
14
+from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
15
+ is_rocm_aiter_moe_enabled)
16
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
17
+ prepare_moe_fp8_layer_for_marlin)
18
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
19
OCP_MX_BLOCK_SIZE)
20
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
21
+ GroupShape)
22
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
23
all_close_1d, normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize)
24
from vllm.model_executor.utils import set_weight_attrs
25
from vllm.platforms import current_platform
26
+from vllm.scalar_type import scalar_types
27
28
logger = init_logger(__name__)
29
@@ -67,21 +75,45 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
75
self.weight_quant = weight_config
76
self.input_quant = input_config
77
70
- weight_qscheme = self.weight_quant.get("qscheme")
71
- input_qscheme = self.input_quant.get("qscheme")
72
- if not (weight_qscheme == "per_tensor"
73
- and input_qscheme == "per_tensor"):
78
+ self.weight_qscheme = self.weight_quant.get("qscheme")
79
+ self.input_qscheme = self.input_quant.get("qscheme")
80
+ per_tensor = (self.weight_qscheme == "per_tensor"
81
+ and self.input_qscheme == "per_tensor")
82
+ per_channel = (self.weight_qscheme == "per_channel"
83
+ and self.input_qscheme == "per_channel")
84
+ self.act_quant_group_shape = GroupShape.PER_TOKEN \
85
+ if per_channel else GroupShape.PER_TENSOR
86
+ if not (per_tensor or per_channel):
87
raise ValueError(
75
- "For FP8 Fused MoE layers, only per-tensor scales "
76
- "for weights and activations are supported. Found "
77
- f"{weight_qscheme}, {input_qscheme}") # noqa E501
88
+ "For FP8 Fused MoE layers, only per-tensor and per-channel "
89
+ "scales for weights and activations are supported. Found "
90
+ f"{self.weight_qscheme}, {self.input_qscheme}") # noqa E501
91
92
self.static_input_scales = not self.input_quant.get("is_dynamic")
93
+ if self.static_input_scales and per_channel:
94
+ raise ValueError(
95
+ "For FP8 Fused MoE layer, we require either per tensor or "
96
+ "channelwise, dynamic per token quantization.")
97
+
98
+ # For GPUs that lack FP8 hardware support, we can leverage the Marlin
99
+ # kernel for fast weight-only FP8 quantization
100
+ self.use_marlin = (not current_platform.has_device_capability(89)
101
+ or envs.VLLM_TEST_FORCE_FP8_MARLIN)
102
+ # Disable marlin for rocm
103
+ if current_platform.is_rocm():
104
+ self.use_marlin = False
105
+
106
+ self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
107
108
def create_weights(self, layer: torch.nn.Module, num_experts: int,
109
hidden_size: int, intermediate_size_per_partition: int,
110
params_dtype: torch.dtype, **extra_weight_attrs):
111
112
+ layer.intermediate_size_per_partition = intermediate_size_per_partition
113
+ layer.hidden_size = hidden_size
114
+ layer.num_experts = num_experts
115
+ layer.orig_dtype = params_dtype
116
+ layer.weight_block_size = None
117
params_dtype = torch.float8_e4m3fn
118
119
# WEIGHTS
@@ -104,24 +136,39 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
136
set_weight_attrs(w2_weight, extra_weight_attrs)
137
138
# WEIGHT_SCALES
107
- # Allocate 2 scales for w1 and w3 respectively.
108
- # They will be combined to a single scale after weight loading.
109
- w13_weight_scale = torch.nn.Parameter(torch.ones(num_experts,
110
- 2,
111
- dtype=torch.float32),
112
- requires_grad=False)
113
- layer.register_parameter("w13_weight_scale", w13_weight_scale)
114
-
115
- w2_weight_scale = torch.nn.Parameter(torch.ones(num_experts,
116
- dtype=torch.float32),
117
- requires_grad=False)
118
- layer.register_parameter("w2_weight_scale", w2_weight_scale)
119
- # Add the quantization method used (per tensor/grouped/channel)
120
- # to ensure the weight scales are loaded in properly
121
- extra_weight_attrs.update(
122
- {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value})
123
- set_weight_attrs(w13_weight_scale, extra_weight_attrs)
124
- set_weight_attrs(w2_weight_scale, extra_weight_attrs)
139
+ if self.weight_qscheme == "per_tensor":
140
+ # Allocate 2 scales for w1 and w3 respectively.
141
+ # They are combined to a single scale after weight loading.
142
+ w13_weight_scale = torch.nn.Parameter(torch.ones(
143
+ num_experts, 2, dtype=torch.float32),
144
+ requires_grad=False)
145
+ layer.register_parameter("w13_weight_scale", w13_weight_scale)
146
+ w2_weight_scale = torch.nn.Parameter(torch.ones(
147
+ num_experts, dtype=torch.float32),
148
+ requires_grad=False)
149
+ layer.register_parameter("w2_weight_scale", w2_weight_scale)
150
+ # Add PER-TENSOR quantization for FusedMoE.weight_loader.
151
+ extra_weight_attrs.update(
152
+ {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value})
153
+ set_weight_attrs(w13_weight_scale, extra_weight_attrs)
154
+ set_weight_attrs(w2_weight_scale, extra_weight_attrs)
155
+ elif self.weight_qscheme == "per_channel":
156
+ # quark's scale is 1 dim.
157
+ w13_weight_scale = torch.nn.Parameter(torch.ones(
158
+ num_experts,
159
+ 2 * intermediate_size_per_partition,
160
+ dtype=torch.float32),
161
+ requires_grad=False)
162
+ layer.register_parameter("w13_weight_scale", w13_weight_scale)
163
+ w2_weight_scale = torch.nn.Parameter(torch.ones(
164
+ num_experts, hidden_size, dtype=torch.float32),
165
+ requires_grad=False)
166
+ layer.register_parameter("w2_weight_scale", w2_weight_scale)
167
+ # Add PER-CHANNEL quantization for FusedMoE.weight_loader.
168
+ extra_weight_attrs.update(
169
+ {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value})
170
+ set_weight_attrs(w13_weight_scale, extra_weight_attrs)
171
+ set_weight_attrs(w2_weight_scale, extra_weight_attrs)
172
173
# INPUT_SCALES
174
if self.static_input_scales:
@@ -185,24 +232,60 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
232
layer.w2_input_scale = torch.nn.Parameter(w2_input_scale,
233
requires_grad=False)
234
188
- # Fp8 moe kernel needs single weight scale for w13 per expert.
189
- # We take the max then dequant and requant each expert.
190
- assert layer.w13_weight_scale is not None
191
- shard_size = layer.intermediate_size_per_partition
192
- max_w13_scales = layer.w13_weight_scale.max(dim=1).values
193
- for expert_id in range(layer.local_num_experts):
194
- start = 0
195
- for shard_id in range(2):
196
- dq_weight = per_tensor_dequantize(
197
- layer.w13_weight[expert_id][start:start + shard_size, :],
198
- layer.w13_weight_scale[expert_id][shard_id])
199
- layer.w13_weight[expert_id][
200
- start:start + shard_size, :], _ = ops.scaled_fp8_quant(
201
- dq_weight, max_w13_scales[expert_id])
202
- start += shard_size
203
-
204
- layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
205
- requires_grad=False)
235
+ # For per-tensor case, Fp8 moe kernel needs single weight scale
236
+ # for w13 per expert. Use max then dequant and requant each expert.
237
+ if self.weight_qscheme == "per_tensor":
238
+ assert layer.w13_weight_scale is not None
239
+ shard_size = layer.intermediate_size_per_partition
240
+ max_w13_scales = layer.w13_weight_scale.max(dim=1).values
241
+ for expert_id in range(layer.local_num_experts):
242
+ start = 0
243
+ for shard_id in range(2):
244
+ dq_weight = per_tensor_dequantize(
245
+ layer.w13_weight[expert_id][start:start +
246
+ shard_size, :],
247
+ layer.w13_weight_scale[expert_id][shard_id])
248
+ layer.w13_weight[expert_id][
249
+ start:start + shard_size, :], _ = ops.scaled_fp8_quant(
250
+ dq_weight, max_w13_scales[expert_id])
251
+ start += shard_size
252
+
253
+ layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
254
+ requires_grad=False)
255
+ # quark's scale is 1 dim.
256
+ elif self.weight_qscheme == "per_channel":
257
+ if self.act_quant_group_shape == GroupShape.PER_TOKEN:
258
+ w13_weight_scale = layer.w13_weight_scale.unsqueeze(-1)
259
+ layer.w13_weight_scale = torch.nn.Parameter(
260
+ w13_weight_scale, requires_grad=False)
261
+ w2_weight_scale = layer.w2_weight_scale.unsqueeze(-1)
262
+ layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale,
263
+ requires_grad=False)
264
+ # Property to determine if AITER is used
265
+ if self.rocm_aiter_moe_enabled:
266
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa E501
267
+ rocm_aiter_fused_experts, shuffle_weights)
268
+
269
+ # reshaping weights is required for aiter moe kernel.
270
+ shuffled_w13, shuffled_w2 = shuffle_weights(
271
+ layer.w13_weight.data, layer.w2_weight.data)
272
+
273
+ layer.w13_weight = torch.nn.Parameter(shuffled_w13,
274
+ requires_grad=False)
275
+ layer.w2_weight = torch.nn.Parameter(shuffled_w2,
276
+ requires_grad=False)
277
+
278
+ self.rocm_aiter_fused_experts_func = rocm_aiter_fused_experts
279
+ elif self.use_marlin:
280
+
281
+ prepare_moe_fp8_layer_for_marlin(layer, False)
282
+ # Activations not quantized for marlin.
283
+ del layer.w13_input_scale
284
+ del layer.w2_input_scale
285
+ self.fused_experts_func = None
286
+ else:
287
+ from vllm.model_executor.layers.fused_moe import fused_experts
288
+ self.fused_experts_func = fused_experts
289
290
def apply(
291
self,
@@ -233,8 +316,6 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
316
raise NotImplementedError(
317
"EPLB not supported for `QuarkW8A8Fp8MoEMethod` yet.")
318
236
- from vllm.model_executor.layers.fused_moe import fused_experts
237
-
319
topk_weights, topk_ids = FusedMoE.select_experts(
320
hidden_states=x,
321
router_logits=router_logits,
@@ -249,22 +330,60 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
330
e_score_correction_bias=e_score_correction_bias,
331
indices_type=self.topk_indices_dtype)
332
252
- return fused_experts(
253
- x,
254
- layer.w13_weight,
255
- layer.w2_weight,
333
+ if self.rocm_aiter_moe_enabled:
334
+ return self.rocm_aiter_fused_experts_func(
335
+ hidden_states=x,
336
+ w1=layer.w13_weight,
337
+ w2=layer.w2_weight,
338
+ topk_weights=topk_weights,
339
+ topk_ids=topk_ids,
340
+ activation=activation,
341
+ apply_router_weight_on_input=apply_router_weight_on_input,
342
+ use_fp8_w8a8=True,
343
+ per_channel_quant=self.weight_qscheme == "per_channel",
344
+ w1_scale=layer.w13_weight_scale,
345
+ w2_scale=layer.w2_weight_scale,
346
+ a1_scale=layer.w13_input_scale,
347
+ a2_scale=layer.w2_input_scale,
348
+ expert_map=expert_map)
349
+ if self.use_marlin:
350
+ assert activation == "silu", (
351
+ f"{activation} not supported for Marlin MoE.")
352
+ return torch.ops.vllm.fused_marlin_moe(
353
+ x,
354
+ layer.w13_weight,
355
+ layer.w2_weight,
356
+ None,
357
+ None,
358
+ layer.w13_weight_scale,
359
+ layer.w2_weight_scale,
360
+ router_logits,
361
+ topk_weights,
362
+ topk_ids,
363
+ quant_type_id=scalar_types.float8_e4m3fn.id,
364
+ apply_router_weight_on_input=apply_router_weight_on_input,
365
+ global_num_experts=global_num_experts,
366
+ expert_map=expert_map)
367
+
368
+ assert self.fused_experts_func is not None
369
+
370
+ return self.fused_experts_func(
371
+ hidden_states=x,
372
+ w1=layer.w13_weight,
373
+ w2=layer.w2_weight,
374
topk_weights=topk_weights,
375
topk_ids=topk_ids,
376
inplace=True,
377
+ activation=activation,
378
+ apply_router_weight_on_input=apply_router_weight_on_input,
379
use_fp8_w8a8=True,
380
+ per_channel_quant=self.weight_qscheme == "per_channel",
381
global_num_experts=global_num_experts,
261
- apply_router_weight_on_input=apply_router_weight_on_input,
382
expert_map=expert_map,
383
w1_scale=layer.w13_weight_scale,
384
w2_scale=layer.w2_weight_scale,
385
a1_scale=layer.w13_input_scale,
266
- a2_scale=layer.w2_input_scale,
267
- activation=activation)
386
+ a2_scale=layer.w2_input_scale)
387
388
389
class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):