[MiniMax-M3.1] DSpark draft follows --online_quant_config

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)

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