[Fix] Support Kimi-K3 ModelOpt mixed NVFP4/FP8 checkpoint (#35077)

test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_scales.py

@@ -291,6 +291,26 @@ class TestG1ScaleC(CustomTestCase):
         self.assertTrue(g1_scale_c.is_contiguous())
         torch.testing.assert_close(g1_scale_c, torch.full((num_experts,), 20.0))
 
+    def test_situ_keeps_both_dequant_scales_inside_activation(self):
+        # SiTU keeps both GEMM1 scales before tanh; scale_c contains only the
+        # GEMM2 input requantization.
+        num_experts = GATED_CONFIGS[0][1]
+        w2_input_scale_quant = torch.tensor(20.0)
+        gate = _global_scales(num_experts, 1, seed=8)
+        up = _global_scales(num_experts, 1, seed=9)
+
+        g1_scale_c = _compute_g1_scale_c(
+            w2_input_scale_quant,
+            gate,
+            up,
+            is_gated=True,
+            activation="situ",
+        )
+
+        self.assertEqual(g1_scale_c.shape, (num_experts,))
+        self.assertTrue(g1_scale_c.is_contiguous())
+        torch.testing.assert_close(g1_scale_c, torch.full((num_experts,), 20.0))
+
     def test_scale_c_is_float32(self):
         # Lower-precision inputs are upcast to fp32 for the kernel.
         num_experts = GATED_CONFIGS[0][1]

test/registered/unit/model_loader/test_modelopt_loader.py

@@ -696,6 +696,28 @@ class TestModelOptFp4LoaderSelection(CustomTestCase):
 
 
 class TestModelOptMixedPrecisionConfig(CustomTestCase):
+    def test_fp8_pb_wo_dispatches_to_native_block_fp8(self):
+        quant_config = ModelOptMixedPrecisionConfig.from_config(
+            {
+                "quant_algo": "MIXED_PRECISION",
+                "quantized_layers": {
+                    "model.layers.0.self_attn.q_proj": {"quant_algo": "FP8_PB_WO"},
+                },
+                "packed_modules_mapping": {},
+            }
+        )
+
+        # Type dispatch only needs a LinearBase instance; skip GPU weight setup.
+        linear = ReplicatedLinear.__new__(ReplicatedLinear)
+        method = quant_config.get_quant_method(
+            linear, "model.layers.0.self_attn.q_proj"
+        )
+
+        self.assertIsInstance(method, Fp8LinearMethod)
+        self.assertEqual(method.quant_config.weight_block_size, [128, 128])
+        self.assertTrue(method.quant_config.is_checkpoint_fp8_serialized)
+        self.assertEqual(method.quant_config.activation_scheme, "dynamic")
+
     def test_incomplete_inline_config_falls_back_to_hf_quant_config_file(self):
         packed_modules_mapping = {
             "qkv_proj": ["q_proj", "k_proj", "v_proj"],

test/registered/unit/models/test_kimi_k3_bfa_overlap.py

@@ -8,7 +8,10 @@ from unittest.mock import patch
 
 import torch
 
-from sglang.srt.models.kimi_k3 import KimiK3DeltaAttention
+from sglang.srt.models.kimi_k3 import (
+    KimiK3DeltaAttention,
+    _get_k3_dense_weight,
+)
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -39,9 +42,9 @@ def _make_owner(with_stream: bool):
     owner = SimpleNamespace(
         use_full_rank_gate=True,
         _bfa_w=_randn(_BFA_W_ROWS, _H).contiguous(),
+        _bfa_f_b_w=_randn(1536, _N_FA).contiguous(),
         _bfa_fa_size=_N_FA,
         _bfa_b_size=_N_B,
-        f_b_proj=SimpleNamespace(weight=_randn(1536, _N_FA).contiguous()),
         fused_qkvg_proj=fused_qkvg_proj,
         split_sizes=[3 * 1536, 1536],
         _bfa_alt_stream=torch.cuda.Stream() if with_stream else None,
@@ -97,6 +100,39 @@ class TestKimiK3BfaOverlap(CustomTestCase):
         for got, ref in zip(overlap, serial):
             self.assertTrue(torch.equal(got, ref))
 
+    def test_block_fp8_weight_is_dequantized_for_tiny_gemm(self):
+        module = SimpleNamespace(
+            weight=torch.nn.Parameter(
+                torch.ones((130, 129), device="cuda", dtype=torch.float8_e4m3fn),
+                requires_grad=False,
+            ),
+            weight_scale_inv=torch.nn.Parameter(
+                torch.tensor([[1.0, 2.0], [3.0, 4.0]], device="cuda"),
+                requires_grad=False,
+            ),
+            quant_method=SimpleNamespace(weight_block_size=[128, 128]),
+            params_dtype=torch.bfloat16,
+        )
+
+        weight = _get_k3_dense_weight(module)
+
+        self.assertEqual(weight.dtype, torch.bfloat16)
+        torch.testing.assert_close(
+            weight[[0, 0, 128, 128], [0, 128, 0, 128]].float(),
+            torch.tensor([1.0, 2.0, 3.0, 4.0], device="cuda"),
+        )
+
+    def test_per_tensor_fp8_weight_is_not_block_dequantized(self):
+        weight = torch.nn.Parameter(
+            torch.ones((2, 2), device="cuda", dtype=torch.float8_e4m3fn),
+            requires_grad=False,
+        )
+        module = SimpleNamespace(
+            weight=weight, weight_scale=torch.ones(1, device="cuda")
+        )
+
+        self.assertEqual(_get_k3_dense_weight(module).data_ptr(), weight.data_ptr())
+
 
 if __name__ == "__main__":
     unittest.main()