From d49561b8ae9e2da22e9c58f1e947243d1f5b748a Mon Sep 17 00:00:00 2001 From: Colin Z <59755453+ColinZ22@users.noreply.github.com> Date: Sun, 26 Apr 2026 23:56:15 -0700 Subject: [PATCH] [AMD] Fix Kimi-K2.6 Quark MXFP4 loading prefix and packed module mapping (#23408) --- python/sglang/srt/model_loader/loader.py | 7 ++++++- python/sglang/srt/models/kimi_k25.py | 3 ++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 0d90c49cc..87dee392d 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -199,7 +199,12 @@ def _get_quantization_config( remap_prefix = getattr(model_class, "remap_prefix", None) # TODO: we should remove this code and switch to the packed_modules_mapping declared inside the modeling files if model_config.quantization == "quark": - packed_modules_mapping.update({"gate_up_proj": ["gate_proj", "up_proj"]}) + packed_modules_mapping.update( + { + "gate_up_proj": ["gate_proj", "up_proj"], + "fused_qkv_a_proj_with_mqa": ["q_a_proj", "kv_a_proj_with_mqa"], + } + ) if _is_npu: packed_modules_mapping.update( diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index b45490d25..571f9fbb4 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -28,6 +28,7 @@ except ImportError: from sglang.srt.layers.attention.vision import VisionAttention from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig +from sglang.srt.layers.quantization.quark.quark import QuarkConfig from sglang.srt.managers.schedule_batch import ( Modality, MultimodalDataItem, @@ -661,7 +662,7 @@ class KimiK25ForConditionalGeneration(nn.Module): quant_config, prefix=( "language_model" - if isinstance(quant_config, ModelSlimConfig) + if isinstance(quant_config, (ModelSlimConfig, QuarkConfig)) else "" ), )