[AMD] Fix Kimi-K2.6 Quark MXFP4 loading prefix and packed module mapping (#23408)
This commit is contained in:
@@ -199,7 +199,12 @@ def _get_quantization_config(
|
|||||||
remap_prefix = getattr(model_class, "remap_prefix", None)
|
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
|
# TODO: we should remove this code and switch to the packed_modules_mapping declared inside the modeling files
|
||||||
if model_config.quantization == "quark":
|
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:
|
if _is_npu:
|
||||||
packed_modules_mapping.update(
|
packed_modules_mapping.update(
|
||||||
|
|||||||
@@ -28,6 +28,7 @@ except ImportError:
|
|||||||
from sglang.srt.layers.attention.vision import VisionAttention
|
from sglang.srt.layers.attention.vision import VisionAttention
|
||||||
from sglang.srt.layers.linear import ReplicatedLinear
|
from sglang.srt.layers.linear import ReplicatedLinear
|
||||||
from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
|
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 (
|
from sglang.srt.managers.schedule_batch import (
|
||||||
Modality,
|
Modality,
|
||||||
MultimodalDataItem,
|
MultimodalDataItem,
|
||||||
@@ -661,7 +662,7 @@ class KimiK25ForConditionalGeneration(nn.Module):
|
|||||||
quant_config,
|
quant_config,
|
||||||
prefix=(
|
prefix=(
|
||||||
"language_model"
|
"language_model"
|
||||||
if isinstance(quant_config, ModelSlimConfig)
|
if isinstance(quant_config, (ModelSlimConfig, QuarkConfig))
|
||||||
else ""
|
else ""
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user