fix(kimi-k3): preserve dense ModelSlim MLA weights (#36603)

python/sglang/srt/layers/quantization/expert_pack.py

@@ -49,6 +49,7 @@ class ExpertPackConfig(GGUFConfig):
 
     is_fp4_experts = True
     supports_kimi_k3_quantized_latent_projections = True
+    supports_kimi_k3_split_gguf_kv_b = True
 
     def __init__(self, store: ExpertPackStore) -> None:
         super().__init__()

python/sglang/srt/models/kimi_k3.py

@@ -149,6 +149,13 @@ def _uses_modelopt_fp8_pb_wo(
     return resolver is not None and resolver(prefix) == "FP8_PB_WO"
 
 
+def _uses_split_gguf_kv_b(
+    quant_config: Optional[QuantizationConfig],
+) -> bool:
+    """Whether a K3 checkpoint stores MLA K/V as separate GGUF tensors."""
+    return bool(getattr(quant_config, "supports_kimi_k3_split_gguf_kv_b", False))
+
+
 def _maybe_map_fp8_pb_scale_name(name: str, params_dict: dict) -> str:
     """Map ModelOpt FP8_PB_WO scale keys to SGLang block-FP8 params."""
     if name.endswith(".weight_scale"):
@@ -1898,9 +1905,9 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA):
         alt_stream: Optional[torch.cuda.Stream] = None,
         gate_alt_stream: Optional[torch.cuda.Stream] = None,
     ) -> None:
-        split_gguf_kv_b = getattr(
+        # ModelSlim can quantize K3 latent projections while still storing
-            quant_config, "supports_kimi_k3_quantized_latent_projections", False
+        # MLA kv_b_proj as one dense tensor; only GGUF expert packs split K/V.
-        )
+        split_gguf_kv_b = _uses_split_gguf_kv_b(quant_config)
         self.all_reduce_fusion = all_reduce_fusion
         self.use_output_gate = getattr(config, "mla_use_output_gate", False)
         # The fused Ascend split+RMSNorm path is not numerically equivalent for

test/registered/expert_pack/test_kimi_k3_gguf.py

@@ -10,18 +10,29 @@ from sglang.srt.layers.quantization.gguf import (
     GGUFLinearMethod,
     _ordered_gguf_shard_ids,
 )
+from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
 from sglang.srt.model_loader.kimi_k3_gguf import (
     _kda_a_log_target_value,
     _residual_target_value,
     kimi_k3_checkpoint_targets,
     routed_expert_tensor,
 )
+from sglang.srt.models.kimi_k3 import _uses_split_gguf_kv_b
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
 
 
 class TestKimiK3GGUFMapping(unittest.TestCase):
+    def test_split_kv_capability_is_expert_pack_specific(self) -> None:
+        self.assertTrue(ModelSlimConfig.supports_kimi_k3_quantized_latent_projections)
+        self.assertFalse(_uses_split_gguf_kv_b(ModelSlimConfig))
+        self.assertTrue(
+            _uses_split_gguf_kv_b(
+                SimpleNamespace(supports_kimi_k3_split_gguf_kv_b=True)
+            )
+        )
+
     def test_maps_dense_kda_mla_moe_and_residual_tensors(self) -> None:
         cases = {
             "token_embd.weight": ("model.embed_tokens.weight",),