diff --git a/python/sglang/srt/layers/quantization/expert_pack.py b/python/sglang/srt/layers/quantization/expert_pack.py index 7ab5e52da..50951e914 100644 --- a/python/sglang/srt/layers/quantization/expert_pack.py +++ b/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__() diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 8c8e87900..6064795f3 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/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( - quant_config, "supports_kimi_k3_quantized_latent_projections", False - ) + # ModelSlim can quantize K3 latent projections while still storing + # 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 diff --git a/test/registered/expert_pack/test_kimi_k3_gguf.py b/test/registered/expert_pack/test_kimi_k3_gguf.py index ce530b4c9..ccdcf094f 100644 --- a/test/registered/expert_pack/test_kimi_k3_gguf.py +++ b/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",),