@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 5840d95284

[Bugfix][ROCm] AMD-Quark mixed-precision DeepSeek-V4.1 support (#57071)

Signed-off-by: Xiao Yu <xiao.yu.dc@outlook.com> Co-authored-by: Cursor <cursoragent@cursor.com>

xiao-llm committed Sep 25, 2026 at 21:51 UTC 5840d95284fe3ae9f56f2366074f58d277d69ad8
11 files changed +647 -102
tests/quantization/test_fp8.py
+79
@@ -228,6 +228,85 @@ def test_deepseek_v41_vl_mapper_routes_linear_scales(
228 ]
229
230
231 +def test_deepseek_v41_vl_exposes_quant_mappings_on_class():
232 + """``configure_quant_config`` reads both mappings off the class, before
233 + ``__init__`` builds the instance mapper. Without them a Quark config's
234 + per-layer keys never match this wrapper's ``language_model.``-rooted
235 + prefixes and every attention shard falls back to the global spec."""
236 + from vllm.models.deepseek_v41.amd.vl_model import DeepseekV41ForCausalLM
237 +
238 + assert "fused_wqa_wkv" in DeepseekV41ForCausalLM.packed_modules_mapping
239 + mapper = DeepseekV41ForCausalLM.hf_to_vllm_mapper.get_rename_mapper()
240 + assert mapper.apply_list(["layers.0.attn.wq_a"]) == [
241 + "language_model.model.layers.0.attn.wq_a"
242 + ]
243 +
244 +
245 +def test_deepseek_v41_engram_scale_accepts_quark_name():
246 + """Quark exports name the engram scale ``embed.weight_scale``; the
247 + ``\\.scale$`` rules only match a literal ``.scale`` suffix, so without an
248 + explicit rule the tensor is never routed and loading fails."""
249 + from vllm.models.deepseek_v41.amd import model as model_module
250 +
251 + mapper = model_module._make_deepseek_v4_weights_mapper("fp4", "weight_scale")
252 + weight = torch.empty(0)
253 + mapped = [
254 + name
255 + for name, _ in mapper.apply(
256 + [
257 + ("layers.1.engram.embed.weight_scale", weight),
258 + ("layers.1.engram.embed.scale", weight),
259 + ]
260 + )
261 + ]
262 + assert mapped == [
263 + "model.layers.1.engram.embed_tokens.weight_scale_inv",
264 + "model.layers.1.engram.embed_tokens.weight_scale_inv",
265 + ]
266 +
267 +
268 +def test_deepseek_v41_declines_quark_configs():
269 + """``from_config`` rewrites a Quark config into a single global FP8 scheme,
270 + which would discard the per-layer specs of a mixed-precision export (MXFP4
271 + experts + 2-D block MXFP8 attention in DeepSeek-V4.1-Flash). QuarkConfig
272 + knows how to dispatch each scheme, so it must handle every Quark export."""
273 + from vllm.models.deepseek_v41.quant_config import DeepseekV4FP8Config
274 +
275 + hf_config = SimpleNamespace(model_type="deepseek_v41")
276 + mxfp4_global = {
277 + "global_quant_config": {
278 + "weight": {"dtype": "fp4", "qscheme": "per_group", "group_size": 32}
279 + },
280 + "quant_method": "quark",
281 + }
282 +
283 + assert (
284 + DeepseekV4FP8Config.override_quantization_method(
285 + mxfp4_global, None, hf_config=hf_config
286 + )
287 + is None
288 + )
289 +
290 + mixed = {
291 + **mxfp4_global,
292 + "layer_quant_config": {
293 + "layers.0.attn.wkv": {
294 + "weight": {
295 + "dtype": "fp8_e4m3",
296 + "qscheme": "per_block",
297 + "block_size": [32, 32],
298 + }
299 + }
300 + },
301 + }
302 + assert (
303 + DeepseekV4FP8Config.override_quantization_method(
304 + mixed, None, hf_config=hf_config
305 + )
306 + is None
307 + )
308 +
309 +
310 @pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA")
311 @pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu])
312 @pytest.mark.parametrize(
tests/quantization/test_quark.py
+282
@@ -909,6 +909,184 @@ def test_quant_method_dispatch_target(case):
909 assert method_cls is (QuarkLinearMethod if is_linear else case.dispatch_cls)
910
911
912 +def test_quant_method_dispatch_mxfp8_2d_block(default_vllm_config):
913 + """Shape B (DeepSeek-V4.1): 32x32 per-block e8m0 routes to QuarkOCP_MX.
914 +
915 + Verifies that get_quant_method and get_quant_method_target agree (the
916 + short-circuit that previously caused them to diverge has been removed).
917 + """
918 + from vllm.model_executor.layers.quantization.quark.schemes import QuarkOCP_MX
919 + from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Static
920 +
921 + default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
922 + mxfp8_spec = {
923 + "weight": {
924 + "dtype": "fp8_e4m3",
925 + "qscheme": "per_block",
926 + "block_size": [32, 32],
927 + "symmetric": True,
928 + "is_dynamic": False,
929 + "scale_type": "float8_e8m0fnu",
930 + },
931 + "input_tensors": {
932 + "dtype": "fp8_e4m3",
933 + "qscheme": "per_group",
934 + "group_size": 32,
935 + "symmetric": True,
936 + "is_dynamic": True,
937 + },
938 + }
939 + config = QuarkConfig(
940 + {
941 + "global_quant_config": {
942 + "weight": {
943 + "dtype": "fp4",
944 + "qscheme": "per_group",
945 + "group_size": 32,
946 + "scale_format": "e8m0",
947 + "is_dynamic": False,
948 + },
949 + "input_tensors": {
950 + "dtype": "fp4",
951 + "qscheme": "per_group",
952 + "group_size": 32,
953 + "scale_format": "e8m0",
954 + "is_dynamic": True,
955 + },
956 + },
957 + "layer_type_quant_config": {},
958 + "layer_quant_config": {"layers.0.attn.wkv": mxfp8_spec},
959 + "exclude": [],
960 + }
961 + )
962 +
963 + class TestLinear(LinearBase):
964 + def __init__(self):
965 + torch.nn.Module.__init__(self)
966 +
967 + # get_quant_method_target now routes MXFP8 through the matcher chain.
968 + wk, ak, mcls = config.get_quant_method_target("layers.0.attn.wkv", LinearBase)
969 + assert wk == kMxfp8Static
970 + assert mcls is QuarkLinearMethod
971 +
972 + linear = TestLinear()
973 + method = config.get_quant_method(linear, "layers.0.attn.wkv")
974 + assert isinstance(method, QuarkLinearMethod)
975 + assert isinstance(linear.scheme, QuarkOCP_MX)
976 + assert linear.scheme.weight_quant_key == kMxfp8Static
977 + assert linear.scheme.scale_block_rows == 32
978 +
979 + # Experts still fall through to the global MXFP4 spec.
980 + assert (
981 + config.get_quant_method_target("layers.0.ffn.experts", RoutedExperts)[2]
982 + is not QuarkLinearMethod
983 + )
984 +
985 + # 128x128 per-block FP8 (DeepSeek V4) keeps its existing Quark scheme.
986 + v4_spec = {
987 + "weight": {**mxfp8_spec["weight"], "block_size": [128, 128]},
988 + "input_tensors": {**mxfp8_spec["input_tensors"], "group_size": 128},
989 + }
990 + v4_config = _make_qtensor_config(v4_spec["weight"], v4_spec["input_tensors"])
991 + assert v4_config.get_quant_method_target("linear", LinearBase)[2] is (
992 + QuarkLinearMethod
993 + )
994 + linear = TestLinear()
995 + assert isinstance(v4_config.get_quant_method(linear, "linear"), QuarkLinearMethod)
996 + assert isinstance(linear.scheme, QuarkW8A8Fp8PerBlock)
997 +
998 +
999 +def test_quant_method_dispatch_mxfp8_canonical(default_vllm_config):
1000 + """Shape A (canonical MX): per_group/group_size=32/e8m0 FP8 routes
1001 + to QuarkOCP_MX with scale_block_rows == 1.
1002 +
1003 + This covers the canonical MXFP8 spelling (no 2-D block). Validated by
1004 + synthetic dispatch test only — no canonical-MXFP8 checkpoint on hand.
1005 + """
1006 + from vllm.model_executor.layers.quantization.quark.schemes import QuarkOCP_MX
1007 + from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Static
1008 +
1009 + default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
1010 + mxfp8_canonical_spec = {
1011 + "weight": {
1012 + "dtype": "fp8_e4m3",
1013 + "qscheme": "per_group",
1014 + "group_size": 32,
1015 + "scale_format": "e8m0",
1016 + "is_dynamic": False,
1017 + },
1018 + "input_tensors": {
1019 + "dtype": "fp8_e4m3",
1020 + "qscheme": "per_group",
1021 + "group_size": 32,
1022 + "scale_format": "e8m0",
1023 + "is_dynamic": True,
1024 + },
1025 + }
1026 + config = QuarkConfig(
1027 + {
1028 + "global_quant_config": mxfp8_canonical_spec,
1029 + "layer_type_quant_config": {},
1030 + "layer_quant_config": {},
1031 + "exclude": [],
1032 + }
1033 + )
1034 +
1035 + wk, ak, mcls = config.get_quant_method_target("linear", LinearBase)
1036 + assert wk == kMxfp8Static
1037 + assert mcls is QuarkLinearMethod
1038 +
1039 + class TestLinear(LinearBase):
1040 + def __init__(self):
1041 + torch.nn.Module.__init__(self)
1042 +
1043 + linear = TestLinear()
1044 + method = config.get_quant_method(linear, "linear")
1045 + assert isinstance(method, QuarkLinearMethod)
1046 + assert isinstance(linear.scheme, QuarkOCP_MX)
1047 + assert linear.scheme.weight_quant_key == kMxfp8Static
1048 + assert linear.scheme.scale_block_rows == 1
1049 +
1050 +
1051 +def test_quant_method_dispatch_mxfp8_moe_raises(default_vllm_config):
1052 + """MXFP8 in a MoE config raises ValueError — experts are unsupported."""
1053 + from vllm.model_executor.layers.quantization.quark.quark_moe import (
1054 + QuarkOCP_MX_MoEMethod,
1055 + )
1056 +
1057 + default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
1058 + mxfp8_spec = {
1059 + "weight": {
1060 + "dtype": "fp8_e4m3",
1061 + "qscheme": "per_group",
1062 + "group_size": 32,
1063 + "scale_format": "e8m0",
1064 + "is_dynamic": False,
1065 + },
1066 + "input_tensors": {
1067 + "dtype": "fp8_e4m3",
1068 + "qscheme": "per_group",
1069 + "group_size": 32,
1070 + "scale_format": "e8m0",
1071 + "is_dynamic": True,
1072 + },
1073 + }
1074 + config = QuarkConfig(
1075 + {
1076 + "global_quant_config": mxfp8_spec,
1077 + "layer_type_quant_config": {},
1078 + "layer_quant_config": {},
1079 + "exclude": [],
1080 + }
1081 + )
1082 + wk, ak, mcls = config.get_quant_method_target("experts", RoutedExperts)
1083 + assert mcls is QuarkOCP_MX_MoEMethod
1084 + # The OCP MX MoE constructor should fail loudly for MXFP8.
1085 + fake_moe_config = MagicMock()
1086 + with pytest.raises(ValueError, match="MXFP8 experts are not supported"):
1087 + QuarkOCP_MX_MoEMethod(fake_moe_config, wk, ak)
1088 +
1089 +
1090 @pytest.mark.parametrize(
1091 ("weight", "input_tensors"),
1092 [
@@ -2313,3 +2491,107 @@ def test_quark_int4_canonicalizes_pack_for_kernel_layout(pack_method, symmetric)
2491 )
2492
2493 assert torch.equal(actual, expected)
2494 +
2495 +
2496 +# ---------------------------------------------------------------------------
2497 +# override_quantization_method: Quark always delegates to QuarkConfig
2498 +# ---------------------------------------------------------------------------
2499 +
2500 +_QUARK_MXFP4_CFG = {
2501 + "quant_method": "quark",
2502 + "global_quant_config": {
2503 + "weight": {"dtype": "fp4", "qscheme": "per_group", "group_size": 32},
2504 + },
2505 +}
2506 +
2507 +_QUARK_MXFP8_CFG = {
2508 + "quant_method": "quark",
2509 + "global_quant_config": {
2510 + "weight": {
2511 + "dtype": "fp8_e4m3",
2512 + "qscheme": "per_block",
2513 + "block_size": [32, 32],
2514 + "scale_type": "float8_e8m0fnu",
2515 + "symmetric": True,
2516 + },
2517 + },
2518 +}
2519 +
2520 +_QUARK_MIXED_CFG = {
2521 + "quant_method": "quark",
2522 + "global_quant_config": {
2523 + "weight": {"dtype": "fp4", "qscheme": "per_group", "group_size": 32},
2524 + },
2525 + "layer_quant_config": {"layers.0.attn": {"weight": {"dtype": "fp8_e4m3"}}},
2526 +}
2527 +
2528 +_FP8_CFG = {"quant_method": "fp8"}
2529 +
2530 +_DSV4_FP8_CFG = {"quant_method": "deepseek_v4_fp8"}
2531 +
2532 +
2533 +@pytest.mark.parametrize(
2534 + "quant_config_cls",
2535 + [
2536 + "deepseek_v41",
2537 + "deepseek_v4",
2538 + ],
2539 +)
2540 +@pytest.mark.parametrize(
2541 + "hf_quant_cfg, expect_none",
2542 + [
2543 + (_QUARK_MXFP4_CFG, True),
2544 + (_QUARK_MXFP8_CFG, True),
2545 + (_QUARK_MIXED_CFG, True),
2546 + (_FP8_CFG, False),
2547 + (_DSV4_FP8_CFG, False),
2548 + ],
2549 + ids=[
2550 + "quark-mxfp4",
2551 + "quark-mxfp8",
2552 + "quark-mixed",
2553 + "native-fp8",
2554 + "deepseek-v4-fp8",
2555 + ],
2556 +)
2557 +def test_quark_override_delegates_to_quark_config(
2558 + quant_config_cls, hf_quant_cfg, expect_none
2559 +):
2560 + """Any quant_method=='quark' must return None (delegate to QuarkConfig).
2561 +
2562 + Native FP8 configs should still be claimed by the model-specific config.
2563 + """
2564 + import importlib
2565 +
2566 + if quant_config_cls == "deepseek_v41":
2567 + # Module was renamed from deepseek_v4_1 to deepseek_v41 upstream.
2568 + for mod_name in (
2569 + "vllm.models.deepseek_v41.quant_config",
2570 + "vllm.models.deepseek_v4_1.quant_config",
2571 + ):
2572 + try:
2573 + mod = importlib.import_module(mod_name)
2574 + break
2575 + except ModuleNotFoundError:
2576 + continue
2577 + else:
2578 + pytest.skip("deepseek_v41 quant_config not found")
2579 + model_type = "deepseek_v41"
2580 + else:
2581 + mod = importlib.import_module("vllm.models.deepseek_v4.quant_config")
2582 + model_type = "deepseek_v4"
2583 +
2584 + config_cls = mod.DeepseekV4FP8Config
2585 + hf_config = SimpleNamespace(model_type=model_type)
2586 + result = config_cls.override_quantization_method(
2587 + hf_quant_cfg, None, hf_config=hf_config
2588 + )
2589 +
2590 + if expect_none:
2591 + assert result is None, (
2592 + f"Quark config should delegate to QuarkConfig (return None), got {result!r}"
2593 + )
2594 + else:
2595 + assert result == "deepseek_v4_fp8", (
2596 + f"Native FP8 should be claimed by DeepseekV4FP8Config, got {result!r}"
2597 + )
vllm/model_executor/layers/quantization/quark/quark.py
+41 -42
@@ -45,6 +45,8 @@ from vllm.model_executor.layers.quantization.quark.utils import (
45 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
46 _ACTIVATION_QUANT_KEY_MAP,
47 _WEIGHT_QUANT_KEY_MAP,
48 + OCP_MX_BLOCK_SIZE,
49 + ocp_mx_weight_dtype_and_rows,
50 )
51 from vllm.model_executor.layers.quantization.utils.quant_utils import (
52 QuantKey,
@@ -158,6 +160,20 @@ class QuarkConfig(QuantizationConfig):
160
161 self.quant_config = quant_config_with_hf_to_vllm_mapper
162
163 + def _scale_block_rows(self, prefix: str, layer_type: type[torch.nn.Module]) -> int:
164 + """How many weight rows one checkpoint scale row covers.
165 +
166 + 1 for the canonical 1-D per-group MX spelling; ``block_size[0]`` for
167 + MXFP8 exports that spell the scale 2-D (DeepSeek-V4.1 uses 32). The
168 + scheme cannot work this out itself: it has no handle on the layer name,
169 + and tensor-parallel sharding narrows the loaded scale so the ratio is
170 + not recoverable from shapes either.
171 + """
172 + config = self._find_matched_config(prefix, layer_type)
173 + weight = self._unwrap_single_quant_config(config.get("weight"))
174 + normalized = ocp_mx_weight_dtype_and_rows(weight)
175 + return normalized[1] if normalized is not None else 1
176 +
177 def get_quant_method(
178 self, layer: torch.nn.Module, prefix: str
179 ) -> "QuantizeMethodBase | None":
@@ -175,6 +191,11 @@ class QuarkConfig(QuantizationConfig):
191 scheme_cls,
192 weight_quant_key=weight_quant_key,
193 activation_quant_key=activation_quant_key,
194 + scale_block_rows=(
195 + self._scale_block_rows(prefix, type(layer))
196 + if scheme_cls is QuarkOCP_MX
197 + else 1
198 + ),
199 weight_config=self._find_matched_config(prefix, type(layer)).get(
200 "weight"
201 ),
@@ -640,52 +661,18 @@ class QuarkConfig(QuantizationConfig):
661 The rationale for checking only the weight type is that
662 the model loading concept and process primarily concerns the weights themselves.
663 """
643 - # Confirm weights quantized.
644 - if not isinstance(weight_quant, dict):
645 - logger.debug(
646 - "Quark model's weight quantization is incompatible with OCP_MX format: "
647 - "weight_quant is not a dictionary."
648 - )
649 - return QuantKeyMatch(False, None, None)
650 -
651 - # Input and weight qscheme needs to be per group.
652 - if weight_quant.get("qscheme") != "per_group":
653 - logger.debug(
654 - "Quark model's weight quantization is incompatible with OCP MX format: "
655 - "weight is not per_group."
656 - )
657 - return QuantKeyMatch(False, None, None)
658 -
659 - # Input and weight group size needs to be 32.
660 - if weight_quant.get("group_size") != 32:
661 - logger.debug(
662 - "Quark model's weight quantization is incompatible with OCP MX format: "
663 - "group_size of weight is not 32."
664 - )
665 - return QuantKeyMatch(False, None, None)
666 -
667 - # Activations and weight scales need to be in e8m0 format.
668 - if weight_quant.get("scale_format") != "e8m0":
664 + # Confirm weights quantized, and normalize the two OCP MX config
665 + # dialects (1-D per_group, and MXFP8's 2-D per_block) to a dtype.
666 + normalized = ocp_mx_weight_dtype_and_rows(weight_quant)
667 + if normalized is None:
668 logger.debug(
670 - "Quark model's weight quantization is incompatible with OCP MX format: "
671 - "scale_format of weight is not e8m0."
669 + "Quark model's weight quantization is incompatible with OCP MX "
670 + "format: %s",
671 + weight_quant,
672 )
673 return QuantKeyMatch(False, None, None)
674
675 - # Input and weight dtypes need to be any of fp4,
676 - # fp6_e3m2 or fp6_e3m2, possibly mixed.
677 - if weight_quant.get("dtype") not in {
678 - "fp4",
679 - "fp6_e3m2",
680 - "fp6_e2m3",
681 - }:
682 - logger.debug(
683 - "Quark model's weight quantization is incompatible with OCP MX format: "
684 - "dtype is not in {fp4, fp6_e3m2, fp6_e2m3}."
685 - )
686 - return QuantKeyMatch(False, None, None)
687 -
688 - weight_dtype = weight_quant["dtype"].replace("fp", "mxfp")
675 + weight_dtype, _ = normalized
676 weight_quant_key = _WEIGHT_QUANT_KEY_MAP[weight_dtype]
677 if input_quant is None:
678 activation_quant_key = None
@@ -703,6 +690,14 @@ class QuarkConfig(QuantizationConfig):
690 "input scales."
691 )
692 return QuantKeyMatch(False, None, None)
693 + elif (
694 + input_quant["dtype"] == "fp8_e4m3"
695 + and input_quant.get("qscheme") == "per_group"
696 + and input_quant.get("group_size") == OCP_MX_BLOCK_SIZE
697 + ):
698 + # MXFP8 activations. Must be checked before the per-tensor FP8
699 + # branch below, which would otherwise claim them.
700 + activation_quant_key = _ACTIVATION_QUANT_KEY_MAP["mxfp8_e4m3"]
701 elif input_quant["dtype"] == "fp8_e4m3":
702 activation_quant_key = kFp8DynamicTensorSym
703 else:
@@ -903,6 +898,7 @@ class QuarkConfig(QuantizationConfig):
898 scheme_cls: type["QuarkScheme"],
899 weight_quant_key: QuantKey | None,
900 activation_quant_key: QuantKey | None,
901 + scale_block_rows: int = 1,
902 weight_config: dict[str, Any] | None = None,
903 ) -> "QuarkScheme":
904 """Construct a Quark scheme selected by get_scheme_cls."""
@@ -942,6 +938,9 @@ class QuarkConfig(QuantizationConfig):
938 ):
939 kwargs["weight_quant_key"] = weight_quant_key
940
941 + if scheme_cls is QuarkOCP_MX:
942 + kwargs["scale_block_rows"] = scale_block_rows
943 +
944 scheme = scheme_cls(**kwargs)
945
946 # Raise error if device does not support the scheme
vllm/model_executor/layers/quantization/quark/quark_moe.py
+10 -2
@@ -1587,7 +1587,9 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1587 raise ValueError(
1588 f"Unsupported OCP MX dtype combination for MoE: "
1589 f"input_dtype={self.input_dtype}, weight_dtype={self.weight_dtype}. "
1590 - f"Please check that the combination is supported in OCP_MX_Scheme."
1590 + f"MXFP8 experts are not supported through the MoE path; only "
1591 + f"the linear path handles MXFP8. Supported MoE weight dtypes "
1592 + f"are mxfp4, mxfp6_e3m2, and mxfp6_e2m3."
1593 )
1594
1595 # TODO(bowenbao): refactor and introduce backends for other OCP MX schemes,
@@ -1676,10 +1678,16 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1678 if quant_dtype == "mxfp4":
1679 assert dim % 2 == 0
1680 return dim // 2
1679 - else:
1681 + elif quant_dtype in {"mxfp6_e3m2", "mxfp6_e2m3"}:
1682 # FP6 packs 4 * 6 = 24 bits on 3 bytes.
1683 assert (dim * 3) % 4 == 0
1684 return (dim * 3) // 4
1685 + else:
1686 + raise NotImplementedError(
1687 + f"Unsupported quant_dtype in QuarkOCP_MX_MoEMethod."
1688 + f"get_packed_dim: {quant_dtype}. MXFP8 is not supported for "
1689 + f"MoE experts."
1690 + )
1691
1692 def create_weights(
1693 self,
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+119 -35
@@ -10,22 +10,31 @@ from vllm.logger import init_logger
10 from vllm.model_executor.kernels.linear import (
11 MxFp4LinearKernel,
12 MxFp6LinearKernel,
13 + Mxfp8LinearKernel,
14 init_mxfp4_linear_kernel,
15 init_mxfp6_linear_kernel,
16 + init_mxfp8_linear_kernel,
17 +)
18 +from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
19 + MXFP8_SCALE_DTYPE,
20 + MXFP8_VALUE_DTYPE,
21 )
22 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
23 _ACTIVATION_QUANT_KEY_MAP,
24 _WEIGHT_QUANT_KEY_MAP,
25 OCP_MX_BLOCK_SIZE,
26 + OCP_MX_UNPACKED_DTYPES,
27 )
28 from vllm.model_executor.layers.quantization.utils.quant_utils import (
29 QuantKey,
30 kMxfp4Static,
31 kMxfp6E2M3Static,
32 kMxfp6E3M2Static,
33 + kMxfp8Static,
34 )
35 from vllm.model_executor.parameter import (
36 GroupQuantScaleParameter,
37 + ModelWeightParameter,
38 PackedvLLMParameter,
39 )
40 from vllm.platforms import current_platform
@@ -36,7 +45,7 @@ logger = init_logger(__name__)
45
46
47 class QuarkOCP_MX(QuarkScheme):
39 - ocp_mx_linear: MxFp6LinearKernel | MxFp4LinearKernel
48 + ocp_mx_linear: MxFp6LinearKernel | MxFp4LinearKernel | Mxfp8LinearKernel
49 supported_activation_quant_keys = [*_ACTIVATION_QUANT_KEY_MAP.values(), None]
50 supported_weight_quant_keys = [*_WEIGHT_QUANT_KEY_MAP.values()]
51
@@ -44,9 +53,10 @@ class QuarkOCP_MX(QuarkScheme):
53 self,
54 weight_quant_key: QuantKey,
55 activation_quant_key: QuantKey | None,
56 + scale_block_rows: int = 1,
57 ):
58 super().__init__(weight_quant_key, activation_quant_key)
49 -
59 + self.scale_block_rows = scale_block_rows
60 self.weight_dtype = next(
61 dtype
62 for dtype, quant_key in _WEIGHT_QUANT_KEY_MAP.items()
@@ -62,11 +72,21 @@ class QuarkOCP_MX(QuarkScheme):
72 else None
73 )
74
65 - if self.weight_dtype == "mxfp4":
66 - self.packed_factor: int | Fraction = 2
75 + self.is_unpacked = self.weight_dtype in OCP_MX_UNPACKED_DTYPES
76 +
77 + if self.is_unpacked:
78 + self.packed_factor: int | Fraction = 1
79 + elif self.weight_dtype == "mxfp4":
80 + self.packed_factor = 2
81 else:
82 self.packed_factor = Fraction(numerator=8, denominator=6)
83
84 + if self.is_unpacked:
85 + # MXFP8 runs on the shared kernel abstraction, which picks a native
86 + # backend where one exists and emulates otherwise. The MXFP4/MXFP6
87 + # warnings below do not apply.
88 + return
89 +
90 if not current_platform.supports_mx():
91 logger.warning_once(
92 "The current platform does not support native MXFP4/MXFP6 "
@@ -88,7 +108,9 @@ class QuarkOCP_MX(QuarkScheme):
108 )
109
110 def get_packed_dim(self, dim: int, quant_dtype: str):
91 - if quant_dtype == "mxfp4":
111 + if quant_dtype in OCP_MX_UNPACKED_DTYPES:
112 + return dim
113 + elif quant_dtype == "mxfp4":
114 assert dim % 2 == 0
115 return dim // 2
116 elif quant_dtype in {"mxfp6_e3m2", "mxfp6_e2m3"}:
@@ -102,6 +124,26 @@ class QuarkOCP_MX(QuarkScheme):
124 "open an issue."
125 )
126
127 + def _expanding_scale_loader(self, weight_loader: Callable) -> Callable:
128 + """Expand a 2-D block scale to one row per weight row.
129 +
130 + A ``block_size=[R, 32]`` export carries one e8m0 value per R weight
131 + rows; the kernels want one per row. At ``R == 1`` the layouts already
132 + coincide, so the loader is passed through untouched.
133 + """
134 + if self.scale_block_rows == 1:
135 + return weight_loader
136 +
137 + rows = self.scale_block_rows
138 +
139 + def expanding_loader(param, loaded_weight, *args, **kwargs):
140 + loaded_weight = loaded_weight.view(torch.uint8).repeat_interleave(
141 + rows, dim=0
142 + )
143 + return weight_loader(param, loaded_weight, *args, **kwargs)
144 +
145 + return expanding_loader
146 +
147 @classmethod
148 def get_min_capability(cls) -> int:
149 return 70
@@ -130,36 +172,68 @@ class QuarkOCP_MX(QuarkScheme):
172 "tensor-parallel sharding for this layer."
173 )
174
133 - output_size_per_partition = sum(output_partition_sizes)
134 - layer.logical_widths = output_partition_sizes
135 -
136 - # WEIGHT
137 - weight = PackedvLLMParameter(
138 - data=torch.empty(
139 - output_size_per_partition,
140 - self.get_packed_dim(input_size_per_partition, self.weight_dtype),
141 - dtype=torch.uint8,
142 - ),
143 - input_dim=1,
144 - output_dim=0,
145 - packed_dim=1,
146 - packed_factor=self.packed_factor,
147 - weight_loader=weight_loader,
148 - )
149 - layer.register_parameter("weight", weight)
150 -
151 - # WEIGHT SCALE
152 - weight_scale = GroupQuantScaleParameter(
153 - data=torch.empty(
154 - output_size_per_partition,
155 - input_size_per_partition // OCP_MX_BLOCK_SIZE,
156 - dtype=torch.uint8,
157 - ),
158 - input_dim=1,
159 - output_dim=0,
160 - weight_loader=weight_loader,
161 - )
162 - layer.register_parameter("weight_scale", weight_scale)
175 + if self.is_unpacked:
176 + output_size_per_partition = sum(output_partition_sizes)
177 + layer.logical_widths = output_partition_sizes
178 + layer.input_size_per_partition = input_size_per_partition
179 + layer.output_size_per_partition = output_size_per_partition
180 + layer.params_dtype = params_dtype
181 +
182 + weight = ModelWeightParameter(
183 + data=torch.empty(
184 + output_size_per_partition,
185 + input_size_per_partition,
186 + dtype=MXFP8_VALUE_DTYPE,
187 + ),
188 + input_dim=1,
189 + output_dim=0,
190 + weight_loader=weight_loader,
191 + )
192 + layer.register_parameter("weight", weight)
193 +
194 + weight_scale = GroupQuantScaleParameter(
195 + data=torch.empty(
196 + output_size_per_partition,
197 + input_size_per_partition // OCP_MX_BLOCK_SIZE,
198 + dtype=MXFP8_SCALE_DTYPE,
199 + ),
200 + input_dim=1,
201 + output_dim=0,
202 + weight_loader=self._expanding_scale_loader(weight_loader),
203 + )
204 + layer.register_parameter("weight_scale", weight_scale)
205 + layer.weight_block_size = [1, OCP_MX_BLOCK_SIZE]
206 + else:
207 + output_size_per_partition = sum(output_partition_sizes)
208 + layer.logical_widths = output_partition_sizes
209 +
210 + # WEIGHT
211 + weight = PackedvLLMParameter(
212 + data=torch.empty(
213 + output_size_per_partition,
214 + self.get_packed_dim(input_size_per_partition, self.weight_dtype),
215 + dtype=torch.uint8,
216 + ),
217 + input_dim=1,
218 + output_dim=0,
219 + packed_dim=1,
220 + packed_factor=self.packed_factor,
221 + weight_loader=weight_loader,
222 + )
223 + layer.register_parameter("weight", weight)
224 +
225 + # WEIGHT SCALE
226 + weight_scale = GroupQuantScaleParameter(
227 + data=torch.empty(
228 + output_size_per_partition,
229 + input_size_per_partition // OCP_MX_BLOCK_SIZE,
230 + dtype=torch.uint8,
231 + ),
232 + input_dim=1,
233 + output_dim=0,
234 + weight_loader=weight_loader,
235 + )
236 + layer.register_parameter("weight_scale", weight_scale)
237
238 if self.weight_quant_key == kMxfp4Static:
239 self.ocp_mx_linear = init_mxfp4_linear_kernel(
@@ -170,6 +244,16 @@ class QuarkOCP_MX(QuarkScheme):
244 weight_quant_key=self.weight_quant_key,
245 activation_quant_key=self.activation_quant_key,
246 )
247 + elif self.weight_quant_key == kMxfp8Static:
248 + self.ocp_mx_linear = init_mxfp8_linear_kernel(
249 + weight_shape=layer.weight.shape
250 + )
251 + else:
252 + raise NotImplementedError(
253 + "No OCP MX linear kernel for weight_quant_key="
254 + f"{self.weight_quant_key}. Something is wrong, please open an "
255 + "issue."
256 + )
257
258 def apply_weights(
259 self,
vllm/model_executor/layers/quantization/utils/config_utils.py
+6 -1
@@ -44,7 +44,12 @@ def get_quark_ocp_mx_group_size(
44 return None
45
46 assert weight_quant is not None
47 - return int(weight_quant["group_size"])
47 + group_size = weight_quant.get("group_size")
48 + if group_size is None:
49 + block_size = weight_quant.get("block_size")
50 + if isinstance(block_size, list | tuple) and len(block_size) >= 2:
51 + group_size = block_size[1]
52 + return int(group_size) if group_size is not None else None
53
54
55 def is_shared_expert_quant_fse_compatible(
vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py
+58
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from enum import Enum
4 +from typing import Any
5
6 from vllm.logger import init_logger
7 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -11,6 +12,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
12 kMxfp6E2M3Static,
13 kMxfp6E3M2Dynamic,
14 kMxfp6E3M2Static,
15 + kMxfp8Dynamic,
16 + kMxfp8Static,
17 )
18
19 logger = init_logger(__name__)
@@ -21,14 +24,19 @@ _WEIGHT_QUANT_KEY_MAP: dict[str, QuantKey] = {
24 "mxfp4": kMxfp4Static,
25 "mxfp6_e3m2": kMxfp6E3M2Static,
26 "mxfp6_e2m3": kMxfp6E2M3Static,
27 + "mxfp8_e4m3": kMxfp8Static,
28 }
29
30 _ACTIVATION_QUANT_KEY_MAP: dict[str, QuantKey] = {
31 "mxfp4": kMxfp4Dynamic,
32 "mxfp6_e3m2": kMxfp6E3M2Dynamic,
33 "mxfp6_e2m3": kMxfp6E2M3Dynamic,
34 + "mxfp8_e4m3": kMxfp8Dynamic,
35 }
36
37 +# Unpacked OCP MX dtypes: one byte per element, no `packed_factor`.
38 +OCP_MX_UNPACKED_DTYPES = {"mxfp8_e4m3"}
39 +
40 OCP_MX_DTYPES = {
41 "mxfp4",
42 "mxfp6_e3m2",
@@ -39,6 +47,56 @@ OCP_MX_DTYPES = {
47 }
48
49
50 +def ocp_mx_weight_dtype_and_rows(
51 + weight_quant: dict[str, Any] | None,
52 +) -> tuple[str, int] | None:
53 + """Normalize a Quark weight quant config to ``(mx_dtype, scale_block_rows)``.
54 +
55 + Quark spells OCP MX weights two ways. The canonical one is 1-D per-group:
56 + ``qscheme="per_group"``, ``group_size=32``, ``scale_format="e8m0"``. MXFP8
57 + checkpoints may instead use a 2-D per-block spelling
58 + (``qscheme="per_block"``, ``block_size=[R, 32]``,
59 + ``scale_type="float8_e8m0fnu"``), which carries one scale per ``R`` weight
60 + rows rather than one per row. ``scale_block_rows`` is that ``R``; it is 1
61 + for the canonical spelling, where the two layouts coincide.
62 +
63 + Returns ``None`` when the config is not an OCP MX weight quantization.
64 + """
65 + if not isinstance(weight_quant, dict):
66 + return None
67 +
68 + dtype = weight_quant.get("dtype")
69 + if not isinstance(dtype, str):
70 + return None
71 + mx_dtype = dtype.replace("fp", "mxfp")
72 + if mx_dtype not in _WEIGHT_QUANT_KEY_MAP:
73 + return None
74 +
75 + qscheme = weight_quant.get("qscheme")
76 + if qscheme == "per_group":
77 + if weight_quant.get("group_size") != OCP_MX_BLOCK_SIZE:
78 + return None
79 + if weight_quant.get("scale_format") != "e8m0":
80 + return None
81 + return mx_dtype, 1
82 +
83 + if qscheme == "per_block":
84 + # Only the unpacked dtypes are accepted here: expanding a 2-D scale
85 + # over a sub-byte packed weight has no checkpoint to validate against.
86 + if mx_dtype not in OCP_MX_UNPACKED_DTYPES:
87 + return None
88 + block_size = list(weight_quant.get("block_size") or [])
89 + if len(block_size) != 2 or block_size[1] != OCP_MX_BLOCK_SIZE:
90 + return None
91 + if weight_quant.get("scale_type") != "float8_e8m0fnu":
92 + return None
93 + if weight_quant.get("symmetric") is not True or weight_quant.get("is_dynamic"):
94 + return None
95 + return mx_dtype, block_size[0]
96 +
97 + return None
98 +
99 +
100 class OCP_MX_Scheme(str, Enum):
101 w_mxfp4 = "w_mxfp4"
102 w_mxfp4_a_mxfp4 = "w_mxfp4_a_mxfp4"
vllm/models/deepseek_v4/quant_config.py
+7 -7
@@ -135,15 +135,15 @@ class DeepseekV4FP8Config(Fp8Config):
135 def override_quantization_method(
136 cls, hf_quant_cfg, user_quant, hf_config=None
137 ) -> QuantizationMethods | None:
138 + # Quark checkpoints are always handled by QuarkConfig.
139 + if (
140 + isinstance(hf_quant_cfg, dict)
141 + and hf_quant_cfg.get("quant_method") == "quark"
142 + ):
143 + return None
144 if not (
145 isinstance(hf_quant_cfg, dict)
140 - and (
141 - hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8")
142 - or (
143 - hf_quant_cfg.get("quant_method") == "quark"
144 - and cls._is_quark_mxfp4_ocp(hf_quant_cfg)
145 - )
146 - )
146 + and hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8")
147 ):
148 return None
149 model_type = getattr(hf_config, "model_type", None)
vllm/models/deepseek_v41/amd/model.py
+10
@@ -895,6 +895,11 @@ def _make_deepseek_v4_weights_mapper(
895 # renames the engram fp8 table but not its scale; route the
896 # scale explicitly to the same module.
897 re.compile(r"(engram\.embed)\.scale$"): r"\1_tokens.weight_scale_inv",
898 + # Quark exports spell the same tensor ``embed.weight_scale``,
899 + # which the ``\.scale$`` rules never match.
900 + re.compile(
901 + r"(engram\.embed)\.weight_scale$"
902 + ): r"\1_tokens.weight_scale_inv",
903 re.compile(r"\.scale$"): f".{linear_scale_name}",
904 }
905 else:
@@ -908,6 +913,11 @@ def _make_deepseek_v4_weights_mapper(
913 ): r"\1.weight_scale_inv",
914 # Same engram reroute as the fp4 branch above.
915 re.compile(r"(engram\.embed)\.scale$"): r"\1_tokens.weight_scale_inv",
916 + # Quark exports spell the same tensor ``embed.weight_scale``,
917 + # which the ``\.scale$`` rules never match.
918 + re.compile(
919 + r"(engram\.embed)\.weight_scale$"
920 + ): r"\1_tokens.weight_scale_inv",
921 re.compile(r"\.scale$"): f".{linear_scale_name}",
922 }
923 return WeightsMapper(
vllm/models/deepseek_v41/amd/vl_model.py
+26 -8
@@ -79,20 +79,22 @@ class DeepseekV4VLImagePixelInputs(TensorSchema):
79 types: Annotated[torch.Tensor, TensorShape("ns", dynamic_dims={"ns"})]
80
81
82 +_VL_PREFIX_MAPPING: dict[str, str | None] = {
83 + "layers.": "language_model.model.layers.",
84 + "embed.": "language_model.model.embed.",
85 + "norm.": "language_model.model.norm.",
86 + "hc_head": "language_model.model.hc_head",
87 + "mtp.": "language_model.model.mtp.",
88 +}
89 +
90 +
91 def _make_deepseek_v4_vl_weights_mapper(
92 expert_dtype: str, linear_scale_name: str
93 ) -> WeightsMapper:
94 """Text-checkpoint mapping rules re-rooted under ``language_model.``."""
95 base = _make_deepseek_v4_weights_mapper(expert_dtype, linear_scale_name)
87 - orig_to_new_prefix: dict[str, str | None] = {
88 - "layers.": "language_model.model.layers.",
89 - "embed.": "language_model.model.embed.",
90 - "norm.": "language_model.model.norm.",
91 - "hc_head": "language_model.model.hc_head",
92 - "mtp.": "language_model.model.mtp.",
93 - }
96 return WeightsMapper(
95 - orig_to_new_prefix=orig_to_new_prefix,
97 + orig_to_new_prefix=dict(_VL_PREFIX_MAPPING),
98 orig_to_new_regex=base.orig_to_new_regex,
99 orig_to_new_suffix={
100 "head.weight": "language_model.lm_head.weight",
@@ -131,6 +133,22 @@ class DeepseekV41ForCausalLM(
133
134 supports_encoder_tp_data = True
135
136 + # Both of these are read off the *class* by
137 + # ``configure_quant_config``/``SupportsQuant``, before ``__init__`` runs,
138 + # to translate a quantization config's checkpoint-style module names into
139 + # runtime prefixes. ``__init__`` later overwrites the instance
140 + # ``hf_to_vllm_mapper`` with the full weight-loading mapper; only the
141 + # prefix rules matter here. Without the class attributes, per-layer specs
142 + # naming ``layers.N.attn.wq_a`` never match this wrapper's
143 + # ``language_model.model.layers.N.attn.wq_a`` and silently fall back to
144 + # the global spec.
145 + packed_modules_mapping = DeepseekV41LLMForCausalLM.packed_modules_mapping
146 +
147 + hf_to_vllm_mapper = WeightsMapper(
148 + orig_to_new_prefix=dict(_VL_PREFIX_MAPPING),
149 + orig_to_new_substr={".shared_experts.w2": ".shared_experts.down_proj"},
150 + )
151 +
152 # The MoE router needs raw token ids to detect image-span tokens
153 # (all carrying image_token_id, see common/mm_preprocess.py) and apply
154 # bias_vl.
vllm/models/deepseek_v41/quant_config.py
+9 -7
@@ -139,15 +139,17 @@ class DeepseekV4FP8Config(Fp8Config):
139 def override_quantization_method(
140 cls, hf_quant_cfg, user_quant, hf_config=None
141 ) -> QuantizationMethods | None:
142 + # Quark checkpoints are always handled by QuarkConfig, which knows
143 + # how to dispatch each per-layer scheme (MXFP4, MXFP8, etc.).
144 + # ``from_config`` below rewrites the config into a single global FP8
145 + # scheme, which would discard per-layer specs — so never claim Quark.
146 + if isinstance(hf_quant_cfg, dict) and (
147 + hf_quant_cfg.get("quant_method") == "quark"
148 + ):
149 + return None
150 if not (
151 isinstance(hf_quant_cfg, dict)
144 - and (
145 - hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8")
146 - or (
147 - hf_quant_cfg.get("quant_method") == "quark"
148 - and cls._is_quark_mxfp4_ocp(hf_quant_cfg)
149 - )
150 - )
152 + and hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8")
153 ):
154 return None
155 model_type = getattr(hf_config, "model_type", None)