Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions python/sglang/srt/layers/quantization/expert_pack.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__()
Expand Down
13 changes: 10 additions & 3 deletions python/sglang/srt/models/kimi_k3.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Expand Down Expand Up @@ -1898,9 +1905,9 @@ def __init__(
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
Expand Down
11 changes: 11 additions & 0 deletions test/registered/expert_pack/test_kimi_k3_gguf.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,18 +10,29 @@
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",),
Expand Down
Loading