@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e41a17e606

[ROCm][Quantization] Support AMD Quark per-block FP8 for fused MoE layers (#52263)

Signed-off-by: huji <huidong.ji@amd.com> Signed-off-by: jimmy-adams <41593649+jimmy-adams@users.noreply.github.com>

jimmy-adams committed Sep 8, 2026 at 18:30 UTC e41a17e606c9a8ffa639ebde5bce885667ff27ea

5 files changed +283 -56

tests/quantization/test_quark.py

+158 -1

@@ -11,7 +11,7 @@ import importlib.metadata

11 from dataclasses import dataclass

12 from importlib.util import find_spec

13 from types import SimpleNamespace

14 -from unittest.mock import MagicMock, patch

14 +from unittest.mock import MagicMock, Mock, patch

16 import huggingface_hub

17 import lm_eval

@@ -26,6 +26,7 @@ from vllm.config.cache import CacheConfig

26 from vllm.forward_context import set_forward_context

27 from vllm.model_executor.layers.attention import Attention

28 from vllm.model_executor.layers.fused_moe import (

29 + FusedMoeWeightScaleSupported,

30 RoutedExperts,

31 UnquantizedFusedMoEMethod,

32 )

@@ -48,6 +49,7 @@ from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501

49 from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E501

50 QuarkMoEMethod,

51 QuarkW4A8Fp8MoEMethod,

52 + QuarkW8A8Fp8MoEMethod,

53 QuarkW8A8Int8MoEMethod,

54 )

55 from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme

@@ -247,6 +249,26 @@ QTENSOR_CONFIGS = [

249 act_quant_key=kFp8Dynamic128Sym,

250 dispatch_cls=QuarkW8A8Fp8PerBlock,

251 ),

252 + QTensorConfig(

253 + name="fp8_w8a8_dynamic_block_fp32_moe",

254 + weight={

255 + "dtype": "fp8_e4m3",

256 + "qscheme": "per_block",

257 + "is_dynamic": False,

258 + "block_size": [128, 128],

259 + "symmetric": True,

260 + },

261 + input_tensors={

262 + "dtype": "fp8_e4m3",

263 + "qscheme": "per_group",

264 + "is_dynamic": True,

265 + "group_size": 128,

266 + "symmetric": True,

267 + },

268 + weight_quant_key=kFp8Static128BlockSym,

269 + act_quant_key=kFp8Dynamic128Sym,

270 + dispatch_cls=QuarkW8A8Fp8MoEMethod,

271 + ),

272 QTensorConfig(

273 name="fp8_w8a8_block_static_input",

274 weight={

@@ -968,12 +990,111 @@ def test_quant_method_dispatch_instantiation(case, monkeypatch, default_vllm_con

990 lambda: True,

991 )

992

993 + # default_vllm_config carries no model, but the FP8 and OCP MX methods

994 + # read the model type off the HF config.

995 + monkeypatch.setattr(

996 + "vllm.model_executor.layers.quantization.quark.quark_moe."

997 + "get_current_vllm_config",

998 + lambda: SimpleNamespace(

999 + model_config=SimpleNamespace(hf_config=SimpleNamespace())

1000 + ),

1001 + )

1002 +

1003 layer = TestRoutedExperts()

1004 method = config.get_quant_method(layer, "experts")

1005

1006 assert isinstance(method, case.dispatch_cls)

1007

1008

1009 +QUARK_MOE_MODULE = "vllm.model_executor.layers.quantization.quark.quark_moe"

1010 +# validate_fp8_block_shape_moe imports this from vllm.distributed when called,

1011 +# so the patch has to target the source module rather than a local binding.

1012 +TP_WORLD_SIZE = "vllm.distributed.get_tensor_model_parallel_world_size"

1013 +

1014 +

1015 +def _make_per_block_fp8_moe_method(

1016 + activation_quant_key: QuantKey = kFp8Dynamic128Sym,

1017 +) -> QuarkW8A8Fp8MoEMethod:

1018 + return QuarkW8A8Fp8MoEMethod(

1019 + _make_test_moe_config(),

1020 + kFp8Static128BlockSym,

1021 + activation_quant_key,

1022 + )

1023 +

1024 +

1025 +def test_quark_w8a8_fp8_moe_per_block_requires_dynamic_group_input():

1026 + with (

1027 + patch(

1028 + f"{QUARK_MOE_MODULE}.select_fp8_moe_backend",

1029 + return_value=(Mock(), Mock()),

1030 + ),

1031 + patch(f"{QUARK_MOE_MODULE}.get_current_vllm_config"),

1032 + pytest.raises(ValueError, match="per-block scales"),

1033 + ):

1034 + _make_per_block_fp8_moe_method(kFp8StaticTensorSym)

1035 +

1036 +

1037 +def test_quark_w8a8_fp8_moe_per_block_weight_shapes():

1038 + with (

1039 + patch(

1040 + f"{QUARK_MOE_MODULE}.select_fp8_moe_backend",

1041 + return_value=(Mock(), Mock()),

1042 + ),

1043 + patch(f"{QUARK_MOE_MODULE}.get_current_vllm_config"),

1044 + patch(TP_WORLD_SIZE, return_value=1),

1045 + ):

1046 + method = _make_per_block_fp8_moe_method()

1047 + assert method.block_quant

1048 + assert method.weight_block_size == [128, 128]

1049 +

1050 + layer = torch.nn.Module()

1051 + method.create_weights(

1052 + layer,

1053 + num_experts=4,

1054 + hidden_size=512,

1055 + intermediate_size_per_partition=256,

1056 + params_dtype=torch.bfloat16,

1057 + )

1058 +

1059 + w13_num_shards = method.moe.w13_num_shards

1060 + assert layer.weight_block_size == [128, 128]

1061 + assert layer.w13_weight.shape == (4, w13_num_shards * 256, 512)

1062 + assert layer.w2_weight.shape == (4, 512, 256)

1063 + # Quark exports block scales as `weight_scale`, like the per-tensor and

1064 + # per-channel schemes, so all schemes register the same parameter name.

1065 + assert not hasattr(layer, "w13_weight_scale_inv")

1066 + assert not hasattr(layer, "w2_weight_scale_inv")

1067 + # One scale per 128x128 tile of each expert's weight.

1068 + assert layer.w13_weight_scale.shape == (

1069 + 4,

1070 + w13_num_shards * (256 // 128),

1071 + 512 // 128,

1072 + )

1073 + assert layer.w2_weight_scale.shape == (4, 512 // 128, 256 // 128)

1074 + # The loader shards block scales on the block grid, not per row.

1075 + for scale in (layer.w13_weight_scale, layer.w2_weight_scale):

1076 + assert scale.quant_method == FusedMoeWeightScaleSupported.BLOCK.value

1077 +

1078 +

1079 +def test_quark_w8a8_fp8_moe_per_block_rejects_misaligned_partition():

1080 + with (

1081 + patch(

1082 + f"{QUARK_MOE_MODULE}.select_fp8_moe_backend",

1083 + return_value=(Mock(), Mock()),

1084 + ),

1085 + patch(f"{QUARK_MOE_MODULE}.get_current_vllm_config"),

1086 + patch(TP_WORLD_SIZE, return_value=1),

1087 + pytest.raises(ValueError, match="not divisible by"),

1088 + ):

1089 + _make_per_block_fp8_moe_method().create_weights(

1090 + torch.nn.Module(),

1091 + num_experts=4,

1092 + hidden_size=512,

1093 + intermediate_size_per_partition=192,

1094 + params_dtype=torch.bfloat16,

1095 + )

1096 +

1097 +

1098 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])

1099 def test_quark_fp8_w_per_tensor_a_per_tensor(

1100 kv_cache_dtype: str, monkeypatch, dist_init, workspace_init

@@ -1111,6 +1232,42 @@ def test_quark_int8_w8a8_moe(vllm_runner, tp):

1232 assert output

1233

1234

1235 +@pytest.mark.parametrize("tp", [1])

1236 +def test_quark_fp8_w8a8_per_block_moe(vllm_runner, tp):

1237 + """Test per-block (128x128) FP8 MoE quantization with a tiny Qwen3 MoE model."""

1238 + model_path = "Adamji/tiny-qwen3-moe-fp8-per-block"

1239 + with vllm_runner(

1240 + model_path,

1241 + enforce_eager=True,

1242 + tensor_parallel_size=tp,

1243 + gpu_memory_utilization=0.1,

1244 + ) as llm:

1245 +

1246 + def check_model(model):

1247 + experts = model.model.layers[0].mlp.experts

1248 + method = experts._quant_method

1249 + assert isinstance(method, QuarkW8A8Fp8MoEMethod), (

1250 + f"Expected QuarkW8A8Fp8MoEMethod, got {type(method)}"

1251 + )

1252 + assert method.weight_qscheme == "per_block"

1253 + assert method.weight_block_size == [128, 128]

1254 +

1255 + # hidden_size=128, moe_intermediate_size=256 and 4 experts, so one

1256 + # scale per 128x128 tile, with w13 stacking gate on top of up.

1257 + routed_experts = experts.routed_experts

1258 + assert routed_experts.w13_weight_scale.shape == (4, 4, 1)

1259 + assert routed_experts.w2_weight_scale.shape == (4, 1, 2)

1260 + # Quark exports the block scales under the same name as the

1261 + # per-tensor and per-channel schemes.

1262 + assert not hasattr(routed_experts, "w13_weight_scale_inv")

1263 + assert not hasattr(routed_experts, "w2_weight_scale_inv")

1264 +

1265 + llm.apply_model(check_model)

1266 +

1267 + output = llm.generate_greedy("Hello", max_tokens=4)

1268 + assert output

1269 +

1270 +

1271 @pytest.mark.skipif(

1272 not (on_gfx950() or on_gfx942()),

1273 reason="Quark W4A8 (INT4-FP8) MoE requires the AITER kernel on gfx942/gfx950",

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py

+5 -19

@@ -8,7 +8,6 @@ from compressed_tensors.quantization import (

8 QuantizationStrategy,

9 )

11 -from vllm.distributed import get_tensor_model_parallel_world_size

11 from vllm.logger import init_logger

12 from vllm.model_executor.layers.fused_moe import (

13 FusedMoeWeightScaleSupported,

@@ -31,6 +30,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso

30 from vllm.model_executor.layers.quantization.utils.fp8_utils import (

31 process_fp8_input_tensor_strategy_moe,

32 process_fp8_weight_tensor_strategy_moe,

33 + validate_fp8_block_shape_moe,

34 )

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

36 kFp8Dynamic128Sym,

@@ -129,28 +129,14 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

129 if self.block_quant:

130 assert self.weight_block_size is not None

131 layer.weight_block_size = self.weight_block_size

132 - tp_size = get_tensor_model_parallel_world_size()

132 block_n, block_k = (

133 self.weight_block_size[0],

134 self.weight_block_size[1],

135 )

137 - # NOTE: To ensure proper alignment of the block-wise quantization

138 - # scales, the output_size of the weights for both the gate and up

139 - # layers must be divisible by block_n.

140 - # Required by column parallel or enabling merged weights

141 - if intermediate_size_per_partition % block_n != 0:

142 - raise ValueError(

143 - f"The output_size of gate's and up's weight = "

144 - f"{intermediate_size_per_partition} is not divisible by "

145 - f"weight quantization block_n = {block_n}."

146 - )

147 - if tp_size > 1 and intermediate_size_per_partition % block_k != 0:

148 - # Required by row parallel

149 - raise ValueError(

150 - f"The input_size of down's weight = "

151 - f"{intermediate_size_per_partition} is not divisible by "

152 - f"weight quantization block_k = {block_k}."

153 - )

136 + validate_fp8_block_shape_moe(

137 + intermediate_size_per_partition,

138 + self.weight_block_size,

139 + )

140

141 # WEIGHTS

142 w13_weight = torch.nn.Parameter(

vllm/model_executor/layers/quantization/fp8.py

+12 -22

@@ -53,6 +53,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (

53 process_fp8_weight_tensor_strategy,

54 process_fp8_weight_tensor_strategy_moe,

55 validate_fp8_block_shape,

56 + validate_fp8_block_shape_moe,

57 )

58 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

59 get_marlin_input_dtype,

@@ -575,34 +576,23 @@ class Fp8MoEMethod(FusedMoEMethodBase):

576 assert self.moe_block_shape is not None

577 moe_block_shape = self.moe_block_shape

578 layer.weight_block_size = self.weight_block_size

578 - tp_size = get_tensor_model_parallel_world_size()

579 block_n, block_k = (

580 self.weight_block_size[0],

581 self.weight_block_size[1],

582 )

583 - # NOTE: To ensure proper alignment of the block-wise quantization

584 - # scales, the output_size of the weights for both the gate and up

585 - # layers must be divisible by block_n.

586 - # Required by column parallel or enabling merged weights

587 - if intermediate_size_per_partition % block_n != 0:

588 - if self.weight_scale_refine is None:

589 - raise ValueError(

590 - f"The output_size of gate's and up's weight = "

591 - f"{intermediate_size_per_partition} is not divisible by "

592 - f"weight quantization block_n = {block_n}."

593 - )

583 + if self.weight_scale_refine is None:

584 + validate_fp8_block_shape_moe(

585 + intermediate_size_per_partition,

586 + self.weight_block_size,

587 + )

588 + else:

589 # Use the refined block grid for the scale parameters; the

590 # loader upsamples the checkpoint scales accordingly.

596 - block_n, block_k = moe_block_shape

597 - if tp_size > 1 and intermediate_size_per_partition % block_k != 0:

598 - # Required by row parallel

599 - if self.weight_scale_refine is None:

600 - raise ValueError(

601 - f"The input_size of down's weight = "

602 - f"{intermediate_size_per_partition} is not divisible by "

603 - f"weight quantization block_k = {block_k}."

604 - )

605 - block_n, block_k = moe_block_shape

591 + tp_size = get_tensor_model_parallel_world_size()

592 + if intermediate_size_per_partition % block_n != 0:

593 + block_n, block_k = moe_block_shape

594 + if tp_size > 1 and intermediate_size_per_partition % block_k != 0:

595 + block_n, block_k = moe_block_shape

596

597 # WEIGHTS

598 w13_weight = torch.nn.Parameter(

vllm/model_executor/layers/quantization/quark/quark_moe.py

+79 -14

@@ -57,6 +57,9 @@ from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (

57 select_nvfp4_moe_backend,

58 )

59 from vllm.model_executor.layers.quantization.quark.utils import QuarkQTensorHint

60 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (

61 + validate_fp8_block_shape_moe,

62 +)

63 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (

64 _ACTIVATION_QUANT_KEY_MAP,

65 _WEIGHT_QUANT_KEY_MAP,

@@ -66,8 +69,10 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (

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

70 GroupShape,

71 QuantKey,

72 + kFp8Dynamic128Sym,

73 kFp8DynamicTensorSym,

74 kFp8DynamicTokenSym,

75 + kFp8Static128BlockSym,

76 kFp8StaticChannelSym,

77 kFp8StaticTensorSym,

78 kInt4W4A8StaticChannelSym,

@@ -230,10 +235,12 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

235 kFp8StaticTensorSym,

236 kFp8DynamicTensorSym,

237 kFp8DynamicTokenSym,

238 + kFp8Dynamic128Sym,

239 ]

240 supported_weight_quant_keys = [

241 kFp8StaticChannelSym,

242 kFp8StaticTensorSym,

243 + kFp8Static128BlockSym,

244 ]

245

246 def __init__(

@@ -244,14 +251,18 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

251 ):

252 super().__init__(moe, weight_quant_key, activation_quant_key)

253 self.weight_dtype = "fp8"

247 - self.weight_qscheme = (

248 - "per_channel" if weight_quant_key == kFp8StaticChannelSym else "per_tensor"

249 - )

250 - self.input_qscheme = (

251 - "per_channel"

252 - if activation_quant_key == kFp8DynamicTokenSym

253 - else "per_tensor"

254 - )

254 + if weight_quant_key == kFp8StaticChannelSym:

255 + self.weight_qscheme = "per_channel"

256 + elif weight_quant_key == kFp8Static128BlockSym:

257 + self.weight_qscheme = "per_block"

258 + else:

259 + self.weight_qscheme = "per_tensor"

260 + if activation_quant_key == kFp8DynamicTokenSym:

261 + self.input_qscheme = "per_channel"

262 + elif activation_quant_key == kFp8Dynamic128Sym:

263 + self.input_qscheme = "per_group"

264 + else:

265 + self.input_qscheme = "per_tensor"

266 self.static_input_scales = activation_quant_key == kFp8StaticTensorSym

267 per_tensor = (

268 self.weight_qscheme == "per_tensor" and self.input_qscheme == "per_tensor"

@@ -259,16 +270,31 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

270 per_channel = (

271 self.weight_qscheme == "per_channel" and self.input_qscheme == "per_channel"

272 )

262 - self.act_quant_group_shape = (

263 - GroupShape.PER_TOKEN if per_channel else GroupShape.PER_TENSOR

273 + per_block = (

274 + self.weight_qscheme == "per_block" and self.input_qscheme == "per_group"

275 )

265 - if not (per_tensor or per_channel):

276 + if not (per_tensor or per_channel or per_block):

277 raise ValueError(

267 - "For FP8 Fused MoE layers, only per-tensor and per-channel "

268 - "scales for weights and activations are supported. Found "

269 - f"{self.weight_qscheme}, {self.input_qscheme}"

278 + "For FP8 Fused MoE layers, only per-tensor, per-channel and "

279 + "per-block scales for weights and activations are supported. "

280 + f"Found {self.weight_qscheme}, {self.input_qscheme}"

281 ) # noqa E501

282

283 + # One scale per (block_n, block_k) tile of the weight, so the block

284 + # shape is carried by the weight key's group shape.

285 + self.weight_block_size: list[int] | None = (

286 + list(weight_quant_key.scale.group_shape) if per_block else None

287 + )

288 + self.block_quant = self.weight_block_size is not None

289 +

290 + if per_channel:

291 + self.act_quant_group_shape = GroupShape.PER_TOKEN

292 + elif per_block:

293 + assert self.weight_block_size is not None

294 + self.act_quant_group_shape = GroupShape(1, self.weight_block_size[1])

295 + else:

296 + self.act_quant_group_shape = GroupShape.PER_TENSOR

297 +

298 self.fp8_backend, self.experts_cls = select_fp8_moe_backend(

299 config=moe,

300 weight_key=self.weight_quant_key,

@@ -293,6 +319,15 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

319 layer.weight_block_size = None

320 params_dtype = torch.float8_e4m3fn

321

322 + if self.block_quant:

323 + assert self.weight_block_size is not None

324 + layer.weight_block_size = self.weight_block_size

325 + block_n, block_k = self.weight_block_size[0], self.weight_block_size[1]

326 + validate_fp8_block_shape_moe(

327 + intermediate_size_per_partition,

328 + self.weight_block_size,

329 + )

330 +

331 # WEIGHTS

332 w13_weight = torch.nn.Parameter(

333 torch.zeros(

@@ -368,6 +403,35 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

403 )

404 set_weight_attrs(w13_weight_scale, extra_weight_attrs)

405 set_weight_attrs(w2_weight_scale, extra_weight_attrs)

406 + elif self.weight_qscheme == "per_block":

407 + # One scale per (block_n, block_k) tile of each expert's weight.

408 + w13_weight_scale = torch.nn.Parameter(

409 + torch.ones(

410 + num_experts,

411 + self.moe.w13_num_shards

412 + * ((intermediate_size_per_partition + block_n - 1) // block_n),

413 + (hidden_size + block_k - 1) // block_k,

414 + dtype=torch.float32,

415 + ),

416 + requires_grad=False,

417 + )

418 + layer.register_parameter("w13_weight_scale", w13_weight_scale)

419 + w2_weight_scale = torch.nn.Parameter(

420 + torch.ones(

421 + num_experts,

422 + (hidden_size + block_n - 1) // block_n,

423 + (intermediate_size_per_partition + block_k - 1) // block_k,

424 + dtype=torch.float32,

425 + ),

426 + requires_grad=False,

427 + )

428 + layer.register_parameter("w2_weight_scale", w2_weight_scale)

429 + # Add PER-BLOCK quantization for RoutedExperts.weight_loader.

430 + extra_weight_attrs.update(

431 + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}

432 + )

433 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)

434 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)

435

436 # INPUT_SCALES

437 if self.static_input_scales:

@@ -548,6 +612,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

612 a2_scale=layer.w2_input_scale,

613 w1_bias=getattr(layer, "w13_bias", None),

614 w2_bias=getattr(layer, "w2_bias", None),

615 + block_shape=self.weight_block_size,

616 per_act_token_quant=self.input_qscheme == "per_channel",

617 per_out_ch_quant=self.weight_qscheme == "per_channel",

618 swiglu_limit=getattr(layer, "swiglu_limit", None),

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+29

@@ -1265,6 +1265,35 @@ def validate_fp8_block_shape(

1265 )

1266

1267

1268 +def validate_fp8_block_shape_moe(

1269 + intermediate_size_per_partition: int,

1270 + block_size: list[int],

1271 +) -> None:

1272 + """Validate fused MoE block quantization shapes for tensor parallelism."""

1273 + from vllm.distributed import get_tensor_model_parallel_world_size

1274 +

1275 + tp_size = get_tensor_model_parallel_world_size()

1276 + block_n, block_k = block_size[0], block_size[1]

1277 +

1278 + # NOTE: To ensure proper alignment of the block-wise quantization

1279 + # scales, the output_size of the weights for both the gate and up

1280 + # layers must be divisible by block_n.

1281 + # Required by column parallel or enabling merged weights

1282 + if intermediate_size_per_partition % block_n != 0:

1283 + raise ValueError(

1284 + f"The output_size of gate's and up's weight = "

1285 + f"{intermediate_size_per_partition} is not divisible by "

1286 + f"weight quantization block_n = {block_n}."

1287 + )

1288 + if tp_size > 1 and intermediate_size_per_partition % block_k != 0:

1289 + # Required by row parallel

1290 + raise ValueError(

1291 + f"The input_size of down's weight = "

1292 + f"{intermediate_size_per_partition} is not divisible by "

1293 + f"weight quantization block_k = {block_k}."

1294 + )

1295 +

1296 +

1297 def create_fp8_weight_parameter(

1298 output_size_per_partition: int,

1299 input_size_per_partition: int,