From 4e2394c581ada7b4dd853af3762839e7d530f82e Mon Sep 17 00:00:00 2001 From: zhuyuhua-v Date: Wed, 28 Oct 2026 07:09:11 +0000 Subject: [PATCH] [MiniMax-M3] DSpark draft follows --online_quant_config The MiniMax-M3.1 DSpark draft built its QuantizationConfig from its own MXFP8 checkpoint block only, so --online_quant_config re-quantized the target and never the draft: under ptpc_fp8 the draft linears still ran MXFP8's triton block-scale GEMM plus a separate activation quant per linear. Pass the server's online quant config to _draft_quant_config, so the draft is re-quantized at load the way the target is. Without --online_quant_config nothing changes. --- atom/models/minimax_m3/dspark.py | 12 +++++++++--- tests/test_dspark.py | 26 ++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 3 deletions(-) diff --git a/atom/models/minimax_m3/dspark.py b/atom/models/minimax_m3/dspark.py index 19c0c6fd05..daab317a02 100644 --- a/atom/models/minimax_m3/dspark.py +++ b/atom/models/minimax_m3/dspark.py @@ -93,7 +93,9 @@ def _draft_rope_theta(config) -> float: return float(theta) -def _draft_quant_config(config) -> QuantizationConfig: +def _draft_quant_config( + config, online_quant_config: dict | None = None +) -> QuantizationConfig: """The draft's own quantization, independent of the target's. The draft checkpoint is MXFP8 (``quant_method: mxfp8``, ``[1, 32]`` blocks) @@ -102,8 +104,12 @@ def _draft_quant_config(config) -> QuantizationConfig: copies onto the draft, cannot be trusted to describe these weights. The checkpoint's own block goes through the registered parser here; a draft config without one is BF16. + + ``--online_quant_config`` re-quantizes the draft at load as it does the + target: under ``ptpc_fp8`` the draft linears run the per-token FP8 GEMM + instead of MXFP8's triton block-scale GEMM plus a quant per linear. """ - draft_qc = QuantizationConfig(config) + draft_qc = QuantizationConfig(config, online_quant_config) spec = draft_qc.global_spec if draft_qc.quant_method and (spec.quant_type, spec.quant_dtype) != ( QuantType.per_1x32, @@ -448,7 +454,7 @@ def __init__(self, atom_config: "Config", layer_offset: int = 0) -> None: ) config = atom_config.hf_config self.hf_config = config - quant_config = _draft_quant_config(config) + quant_config = _draft_quant_config(config, atom_config.online_quant_config) self.hidden_size = config.hidden_size self.noise_token_id = int(config.dspark_noise_token_id) diff --git a/tests/test_dspark.py b/tests/test_dspark.py index 79b3870c0f..2bf4a24609 100644 --- a/tests/test_dspark.py +++ b/tests/test_dspark.py @@ -2545,3 +2545,29 @@ def test_minimax_m31_dspark_refuses_tp_below_its_kv_heads(monkeypatch, tp): config = types.SimpleNamespace(num_attention_heads=64, num_key_value_heads=4) with pytest.raises(ValueError, match="needs TP >= 4"): m3_dspark.M3DSparkAttention(None, config, None, layer_num=0) + + +@pytest.mark.parametrize("online", [None, {"global_quant_config": "ptpc_fp8"}]) +def test_minimax_m31_draft_follows_online_quant_config(online): + """The MiniMax-M3.1 draft keeps its MXFP8 checkpoint quant and takes the + server's --online_quant_config on top, as the target does: under ptpc_fp8 + its linears are re-quantized per channel at load; without one they stay + MXFP8.""" + from aiter import QuantType + + from atom.models.minimax_m3.dspark import _draft_quant_config + + config = types.SimpleNamespace( + torch_dtype=torch.bfloat16, + quantization_config={ + "quant_method": "mxfp8", + "activation_scheme": "dynamic", + "weight_block_size": [1, 32], + }, + ) + qc = _draft_quant_config(config, online) + assert qc.quant_method == "mxfp8" + assert qc.global_spec.quant_type == QuantType.per_1x32 + assert qc.online_quant is (online is not None) + if online is not None: + assert qc.online_global_spec.quant_type == QuantType.per_Token