[Bugfix][ROCm] AMD-Quark mixed-precision DeepSeek-V4.1 support with Quark OCP MXFP8

tests/quantization/test_fp8.py

@@ -237,6 +237,86 @@ def test_deepseek_v41_vl_mapper_routes_linear_scales(
     ]
 
 
+def test_deepseek_v41_vl_exposes_quant_mappings_on_class():
+    """``configure_quant_config`` reads both mappings off the class, before
+    ``__init__`` builds the instance mapper. Without them a Quark config's
+    per-layer keys never match this wrapper's ``language_model.``-rooted
+    prefixes and every attention shard falls back to the global spec."""
+    from vllm.models.deepseek_v41.amd.vl_model import DeepseekV41ForCausalLM
+
+    assert "fused_wqa_wkv" in DeepseekV41ForCausalLM.packed_modules_mapping
+    mapper = DeepseekV41ForCausalLM.hf_to_vllm_mapper.get_rename_mapper()
+    assert mapper.apply_list(["layers.0.attn.wq_a"]) == [
+        "language_model.model.layers.0.attn.wq_a"
+    ]
+
+
+def test_deepseek_v41_engram_scale_accepts_quark_name():
+    """Quark exports name the engram scale ``embed.weight_scale``; the
+    ``\\.scale$`` rules only match a literal ``.scale`` suffix, so without an
+    explicit rule the tensor is never routed and loading fails."""
+    from vllm.models.deepseek_v41.amd import model as model_module
+
+    mapper = model_module._make_deepseek_v4_weights_mapper("fp4", "weight_scale")
+    weight = torch.empty(0)
+    mapped = [
+        name
+        for name, _ in mapper.apply(
+            [
+                ("layers.1.engram.embed.weight_scale", weight),
+                ("layers.1.engram.embed.scale", weight),
+            ]
+        )
+    ]
+    assert mapped == [
+        "model.layers.1.engram.embed_tokens.weight_scale_inv",
+        "model.layers.1.engram.embed_tokens.weight_scale_inv",
+    ]
+
+
+def test_deepseek_v41_declines_mixed_precision_quark_config():
+    """A Quark export with per-layer specs is mixed precision (MXFP4 experts +
+    2-D block MXFP8 attention in DeepSeek-V4.1-Flash). ``from_config`` rewrites
+    the config into a single global FP8 scheme, so claiming such a checkpoint
+    would discard the per-layer specs; QuarkConfig must handle it instead."""
+    from vllm.models.deepseek_v41.quant_config import DeepseekV4FP8Config
+
+    hf_config = SimpleNamespace(model_type="deepseek_v41")
+    mxfp4_global = {
+        "global_quant_config": {
+            "weight": {"dtype": "fp4", "qscheme": "per_group", "group_size": 32}
+        },
+        "quant_method": "quark",
+    }
+
+    # Single-scheme Quark MXFP4 exports (DeepSeek V4) are still claimed.
+    assert (
+        DeepseekV4FP8Config.override_quantization_method(
+            mxfp4_global, None, hf_config=hf_config
+        )
+        == "deepseek_v4_fp8"
+    )
+
+    mixed = {
+        **mxfp4_global,
+        "layer_quant_config": {
+            "layers.0.attn.wkv": {
+                "weight": {
+                    "dtype": "fp8_e4m3",
+                    "qscheme": "per_block",
+                    "block_size": [32, 32],
+                }
+            }
+        },
+    }
+    assert (
+        DeepseekV4FP8Config.override_quantization_method(
+            mixed, None, hf_config=hf_config
+        )
+        is None
+    )
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="DeepGEMM requires CUDA")
 @pytest.mark.parametrize("scale_dtype", [torch.uint8, torch.float8_e8m0fnu])
 @pytest.mark.parametrize(

tests/quantization/test_quark.py

@@ -889,6 +889,186 @@ def test_quant_method_dispatch_target(case):
     assert method_cls is (QuarkLinearMethod if is_linear else case.dispatch_cls)
 
 
+def test_quant_method_dispatch_mxfp8_2d_block(default_vllm_config):
+    """Shape B (DeepSeek-V4.1): 32x32 per-block e8m0 routes to QuarkOCP_MX.
+
+    Verifies that get_quant_method and get_quant_method_target agree (the
+    short-circuit that previously caused them to diverge has been removed).
+    """
+    from vllm.model_executor.layers.quantization.quark.schemes import QuarkOCP_MX
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Static
+
+    default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
+    mxfp8_spec = {
+        "weight": {
+            "dtype": "fp8_e4m3",
+            "qscheme": "per_block",
+            "block_size": [32, 32],
+            "symmetric": True,
+            "is_dynamic": False,
+            "scale_type": "float8_e8m0fnu",
+        },
+        "input_tensors": {
+            "dtype": "fp8_e4m3",
+            "qscheme": "per_group",
+            "group_size": 32,
+            "symmetric": True,
+            "is_dynamic": True,
+        },
+    }
+    config = QuarkConfig(
+        {
+            "global_quant_config": {
+                "weight": {
+                    "dtype": "fp4",
+                    "qscheme": "per_group",
+                    "group_size": 32,
+                    "scale_format": "e8m0",
+                    "is_dynamic": False,
+                },
+                "input_tensors": {
+                    "dtype": "fp4",
+                    "qscheme": "per_group",
+                    "group_size": 32,
+                    "scale_format": "e8m0",
+                    "is_dynamic": True,
+                },
+            },
+            "layer_type_quant_config": {},
+            "layer_quant_config": {"layers.0.attn.wkv": mxfp8_spec},
+            "exclude": [],
+        }
+    )
+
+    class TestLinear(LinearBase):
+        def __init__(self):
+            torch.nn.Module.__init__(self)
+
+    # get_quant_method_target now routes MXFP8 through the matcher chain.
+    wk, ak, mcls = config.get_quant_method_target(
+        "layers.0.attn.wkv", LinearBase
+    )
+    assert wk == kMxfp8Static
+    assert mcls is QuarkLinearMethod
+
+    linear = TestLinear()
+    method = config.get_quant_method(linear, "layers.0.attn.wkv")
+    assert isinstance(method, QuarkLinearMethod)
+    assert isinstance(linear.scheme, QuarkOCP_MX)
+    assert linear.scheme.weight_quant_key == kMxfp8Static
+    assert linear.scheme.scale_block_rows == 32
+
+    # Experts still fall through to the global MXFP4 spec.
+    assert (
+        config.get_quant_method_target("layers.0.ffn.experts", RoutedExperts)[2]
+        is not QuarkLinearMethod
+    )
+
+    # 128x128 per-block FP8 (DeepSeek V4) keeps its existing Quark scheme.
+    v4_spec = {
+        "weight": {**mxfp8_spec["weight"], "block_size": [128, 128]},
+        "input_tensors": {**mxfp8_spec["input_tensors"], "group_size": 128},
+    }
+    v4_config = _make_qtensor_config(v4_spec["weight"], v4_spec["input_tensors"])
+    assert v4_config.get_quant_method_target("linear", LinearBase)[2] is (
+        QuarkLinearMethod
+    )
+    linear = TestLinear()
+    assert isinstance(v4_config.get_quant_method(linear, "linear"), QuarkLinearMethod)
+    assert isinstance(linear.scheme, QuarkW8A8Fp8PerBlock)
+
+
+def test_quant_method_dispatch_mxfp8_canonical(default_vllm_config):
+    """Shape A (canonical MX): per_group/group_size=32/e8m0 FP8 routes
+    to QuarkOCP_MX with scale_block_rows == 1.
+
+    This covers the canonical MXFP8 spelling (no 2-D block). Validated by
+    synthetic dispatch test only — no canonical-MXFP8 checkpoint on hand.
+    """
+    from vllm.model_executor.layers.quantization.quark.schemes import QuarkOCP_MX
+    from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Static
+
+    default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
+    mxfp8_canonical_spec = {
+        "weight": {
+            "dtype": "fp8_e4m3",
+            "qscheme": "per_group",
+            "group_size": 32,
+            "scale_format": "e8m0",
+            "is_dynamic": False,
+        },
+        "input_tensors": {
+            "dtype": "fp8_e4m3",
+            "qscheme": "per_group",
+            "group_size": 32,
+            "scale_format": "e8m0",
+            "is_dynamic": True,
+        },
+    }
+    config = QuarkConfig(
+        {
+            "global_quant_config": mxfp8_canonical_spec,
+            "layer_type_quant_config": {},
+            "layer_quant_config": {},
+            "exclude": [],
+        }
+    )
+
+    wk, ak, mcls = config.get_quant_method_target("linear", LinearBase)
+    assert wk == kMxfp8Static
+    assert mcls is QuarkLinearMethod
+
+    class TestLinear(LinearBase):
+        def __init__(self):
+            torch.nn.Module.__init__(self)
+
+    linear = TestLinear()
+    method = config.get_quant_method(linear, "linear")
+    assert isinstance(method, QuarkLinearMethod)
+    assert isinstance(linear.scheme, QuarkOCP_MX)
+    assert linear.scheme.weight_quant_key == kMxfp8Static
+    assert linear.scheme.scale_block_rows == 1
+
+
+def test_quant_method_dispatch_mxfp8_moe_raises(default_vllm_config):
+    """MXFP8 in a MoE config raises ValueError — experts are unsupported."""
+    from vllm.model_executor.layers.quantization.quark.quark_moe import (
+        QuarkOCP_MX_MoEMethod,
+    )
+
+    default_vllm_config.model_config = SimpleNamespace(dtype=torch.bfloat16)
+    mxfp8_spec = {
+        "weight": {
+            "dtype": "fp8_e4m3",
+            "qscheme": "per_group",
+            "group_size": 32,
+            "scale_format": "e8m0",
+            "is_dynamic": False,
+        },
+        "input_tensors": {
+            "dtype": "fp8_e4m3",
+            "qscheme": "per_group",
+            "group_size": 32,
+            "scale_format": "e8m0",
+            "is_dynamic": True,
+        },
+    }
+    config = QuarkConfig(
+        {
+            "global_quant_config": mxfp8_spec,
+            "layer_type_quant_config": {},
+            "layer_quant_config": {},
+            "exclude": [],
+        }
+    )
+    wk, ak, mcls = config.get_quant_method_target("experts", RoutedExperts)
+    assert mcls is QuarkOCP_MX_MoEMethod
+    # The OCP MX MoE constructor should fail loudly for MXFP8.
+    fake_moe_config = MagicMock()
+    with pytest.raises(ValueError, match="MXFP8 experts are not supported"):
+        QuarkOCP_MX_MoEMethod(fake_moe_config, wk, ak)
+
+
 @pytest.mark.parametrize(
     ("weight", "input_tensors"),
     [