[diffusion] fix: fix quantized qkv scales and missing-param policy for minimax-h3 (#35740)

python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py

@@ -176,6 +176,39 @@ def _reorder_grouped_qkv_to_qkv(
     )
 
 
+def _install_qkv_row_reorder(
+    param: torch.Tensor,
+    reorder: Callable[[torch.Tensor], torch.Tensor],
+    qkv_rows: int,
+) -> None:
+    """Reorder a per-output-row qkv parameter the same way its rows are reordered.
+
+    Applied to quantization metadata rather than the weight itself. Anything whose
+    leading dim is not the checkpoint's qkv row count is passed through: per-tensor
+    scales are scalars, and a swizzled block-scale layout is not row-indexed.
+    """
+
+    def _maybe_reorder(loaded_weight: torch.Tensor) -> torch.Tensor:
+        if loaded_weight.dim() >= 2 and loaded_weight.shape[0] == qkv_rows:
+            return reorder(loaded_weight)
+        return loaded_weight
+
+    base_loader = (
+        param._weight_loader
+        if hasattr(param, "_weight_loader")
+        else param.weight_loader
+    )
+
+    def _weight_loader(p: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        base_loader(p, _maybe_reorder(loaded_weight))
+
+    if hasattr(param, "_weight_loader"):
+        param._weight_loader = _weight_loader
+    else:
+        param.weight_loader = _weight_loader
+    param.rank_local_weight_transform = _maybe_reorder
+
+
 def _copy_grouped_qkv_tp_shard(
     param: torch.Tensor,
     loaded_weight: torch.Tensor,
@@ -662,6 +695,17 @@ class MiniMaxH3Attention(nn.Module):
         # rank-local FSDP must reorder grouped QKV before selecting each shard
         weight.rank_local_weight_transform = _reorder_checkpoint_weight
 
+        # A quantized checkpoint stores metadata indexed by output row next to the
+        # rows themselves (NVFP4 block scales, fp8 per-channel scales). Those rows
+        # are permuted above, so the per-row metadata has to be permuted the same
+        # way. Row count is the gate: a swizzled scale layout is not row-indexed,
+        # and per-tensor scales are scalars, so both are passed through untouched.
+        qkv_rows = 3 * arch.num_attention_heads * arch.attention_head_dim
+        for name, param in self.qkv_proj.named_parameters(recurse=False):
+            if name == "weight":
+                continue
+            _install_qkv_row_reorder(param, _reorder_checkpoint_weight, qkv_rows)
+
     def _forward_mps_streamed_attention(
         self,
         x: torch.Tensor,
@@ -1816,7 +1860,11 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
 
     def _mark_missing_params_required(self) -> None:
         for _, param in self.named_parameters():
-            param.missing_param_init = "error"
+            # A quant method's create_weights() declares its own policy for scales
+            # it can synthesize (weight-only NVFP4 has no input_scale and marks it
+            # "ones"); claiming only undeclared params keeps that intact.
+            if getattr(param, "missing_param_init", None) is None:
+                param.missing_param_init = "error"
 
     def post_load_weights(self) -> None:
         fp32_param_names = list(_MINIMAX_H3_FP32_PARAM_NAMES_IN_MODEL_ORDER)

python/sglang/multimodal_gen/test/unit/test_minimax_h3_qkv_scale_reorder.py

@@ -0,0 +1,60 @@
+import unittest
+
+import torch
+
+from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import (
+    _install_qkv_row_reorder,
+)
+
+QKV_ROWS = 12
+
+
+def _install(reorder=lambda w: w.flip(0)):
+    param = torch.zeros(QKV_ROWS)
+    seen = []
+    param.weight_loader = lambda p, loaded_weight: seen.append(loaded_weight)
+    _install_qkv_row_reorder(param, reorder, QKV_ROWS)
+    return param, seen
+
+
+class TestMiniMaxH3QkvScaleReorder(unittest.TestCase):
+    """MiniMax-H3 permutes the fused qkv weight's output rows on load.
+
+    Quantization metadata indexed by output row (NVFP4 block scales, fp8
+    per-channel scales) has to move with those rows, or every scale lands on the
+    wrong row: the model loads and runs, and renders noise. Metadata that is not
+    row-indexed must be left alone, so the row count is the gate.
+    """
+
+    def test_row_indexed_metadata_is_reordered(self):
+        param, seen = _install()
+        scales = torch.arange(QKV_ROWS * 2, dtype=torch.float32).reshape(QKV_ROWS, 2)
+        param.weight_loader(param, scales)
+        self.assertTrue(torch.equal(seen[0], scales.flip(0)))
+
+    def test_metadata_with_another_row_count_is_passed_through(self):
+        param, seen = _install()
+        swizzled = torch.arange(QKV_ROWS, dtype=torch.float32).reshape(QKV_ROWS // 3, 3)
+        param.weight_loader(param, swizzled)
+        self.assertTrue(torch.equal(seen[0], swizzled))
+
+    def test_per_tensor_scale_is_passed_through(self):
+        param, seen = _install()
+        scale = torch.tensor(0.5)
+        param.weight_loader(param, scale)
+        self.assertTrue(torch.equal(seen[0], scale))
+
+    def test_rank_local_transform_follows_the_same_gate(self):
+        # rank-local FSDP reorders before selecting its shard, so it needs the
+        # transform rather than the wrapped loader.
+        param, _ = _install()
+        scales = torch.arange(QKV_ROWS * 2, dtype=torch.float32).reshape(QKV_ROWS, 2)
+        self.assertTrue(
+            torch.equal(param.rank_local_weight_transform(scales), scales.flip(0))
+        )
+        scale = torch.tensor(0.5)
+        self.assertTrue(torch.equal(param.rank_local_weight_transform(scale), scale))
+
+
+if __name__ == "__main__":
+    unittest.main()