From 5f717913a04d74e10ee5ab7a98a6a9d3e2e57953 Mon Sep 17 00:00:00 2001 From: khalilzhk Date: Tue, 10 Mar 2026 14:43:27 +0800 Subject: [PATCH] support Kimi-K2.5-w4a8 on ascend --- .../quantization/modelslim/modelslim.py | 3 + python/sglang/srt/models/kimi_k25.py | 47 ++++++++++++++-- python/sglang/srt/models/kimi_vl_moonvit.py | 56 +++++++++++++++---- 3 files changed, 91 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index 1212a40d4..65433c57c 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -152,6 +152,9 @@ class ModelSlimConfig(QuantizationConfig): key = "vision_model" elif "visual" in prefix: key = "visual" + if "vision_tower" in prefix or "mm_projector" in prefix: + prefix = prefix.replace(r"attn.qkv_proj", r"wqkv") + prefix = prefix.replace(r"attn.proj", r"wo") packed_modules_mapping_subset = self.packed_modules_mapping.get(key, {}) prefix_in_quant_config = prefix proj_name = prefix.split(".")[-1] diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index bf931d5cc..7825389fb 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -26,6 +26,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.managers.schedule_batch import ( Modality, MultimodalDataItem, @@ -120,7 +121,13 @@ class MoonViTEncoderLayer(nn.Module): self.norm0 = nn.LayerNorm(hidden_dim) self.norm1 = nn.LayerNorm(hidden_dim) - self.mlp = MLP2([hidden_dim, mlp_dim, hidden_dim], activation) + + self.mlp = MLP2( + [hidden_dim, mlp_dim, hidden_dim], + activation, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) self.attn = VisionAttention( embed_dim=hidden_dim, @@ -430,6 +437,8 @@ class MoonViT3dEncoder(nn.Module): num_layers: int, block_cfg: dict, video_attn_type: str = "spatial_temporal", + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", ) -> None: super().__init__() @@ -441,7 +450,14 @@ class MoonViT3dEncoder(nn.Module): block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512 ) self.blocks = nn.ModuleList( - [MoonViTEncoderLayer(**block_cfg) for _ in range(num_layers)] + [ + MoonViTEncoderLayer( + **block_cfg, + quant_config=quant_config, + prefix=add_prefix(f"blocks.{layer_idx}", prefix), + ) + for layer_idx in range(num_layers) + ] ) self.final_layernorm = nn.LayerNorm(hidden_dim) @@ -480,7 +496,15 @@ class MoonViT3dPretrainedModel(nn.Module): _supports_flash_attn_2 = True _supports_sdpa = True - def __init__(self, config, *inputs, use_data_parallel: bool = False, **kwargs): + def __init__( + self, + config, + *inputs, + use_data_parallel: bool = False, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + **kwargs, + ): super().__init__() config = deepcopy(config) self.config = config @@ -509,6 +533,8 @@ class MoonViT3dPretrainedModel(nn.Module): "use_data_parallel": use_data_parallel, }, video_attn_type=config.video_attn_type, + quant_config=quant_config, + prefix=add_prefix("encoder", prefix), ) @property @@ -671,12 +697,23 @@ class KimiK25ForConditionalGeneration(nn.Module): self.use_data_parallel = get_global_server_args().mm_enable_dp_encoder # Create vision tower self.vision_tower = MoonViT3dPretrainedModel( - config.vision_config, use_data_parallel=self.use_data_parallel + config.vision_config, + use_data_parallel=self.use_data_parallel, + quant_config=( + quant_config if isinstance(quant_config, ModelSlimConfig) else None + ), + prefix="vision_tower", ) # Create mm projector self.mm_projector = K2VLMultiModalProjector(config.vision_config) - self.language_model = DeepseekV3ForCausalLM(config.text_config, quant_config) + self.language_model = DeepseekV3ForCausalLM( + config.text_config, + quant_config, + prefix=( + "language_model" if isinstance(quant_config, ModelSlimConfig) else "" + ), + ) # Ensure that the dtype of the vision_tower and mm_projector matches that of the language_model. # This solves the dtype mismatch issue when using device_map="auto" and torch_dtype. diff --git a/python/sglang/srt/models/kimi_vl_moonvit.py b/python/sglang/srt/models/kimi_vl_moonvit.py index 286e85772..9a0e5c405 100644 --- a/python/sglang/srt/models/kimi_vl_moonvit.py +++ b/python/sglang/srt/models/kimi_vl_moonvit.py @@ -58,6 +58,10 @@ except ImportError: flash_attn_varlen_func = None from sglang.srt.configs import MoonViTConfig +from sglang.srt.layers.linear import ReplicatedLinear +from sglang.srt.layers.quantization import QuantizationConfig +from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig +from sglang.srt.utils import add_prefix def multihead_attention( @@ -393,21 +397,53 @@ class MLP2(nn.Module): bias: whether to use bias in linear layer. """ - def __init__(self, dims: list[int], activation, bias=True): + def __init__( + self, + dims: list[int], + activation, + bias: bool = True, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): super().__init__() assert len(dims) == 3 - self.fc0 = nn.Linear(dims[0], dims[1], bias=bias) - self.fc1 = nn.Linear(dims[1], dims[2], bias=bias) + + self.quant_config = quant_config + if isinstance(self.quant_config, ModelSlimConfig): + self.fc0 = ReplicatedLinear( + dims[0], + dims[1], + bias=bias, + quant_config=quant_config, + prefix=add_prefix("fc0", prefix), + ) + self.fc1 = ReplicatedLinear( + dims[1], + dims[2], + bias=bias, + quant_config=quant_config, + prefix=add_prefix("fc1", prefix), + ) + else: + self.fc0 = nn.Linear(dims[0], dims[1], bias=bias) + self.fc1 = nn.Linear(dims[1], dims[2], bias=bias) + for m in [self.fc0, self.fc1]: + nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features)) + if m.bias is not None: + nn.init.zeros_(m.bias) self.activation = activation - for m in [self.fc0, self.fc1]: - nn.init.trunc_normal_(m.weight, std=math.sqrt(2 / m.in_features)) - if m.bias is not None: - nn.init.zeros_(m.bias) def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.fc0(x) - x = self.activation(x) - return self.fc1(x) + if isinstance(self.quant_config, ModelSlimConfig): + x = x.flatten(0, 1) + x, _ = self.fc0(x) + x = self.activation(x) + x, _ = self.fc1(x) + else: + x = self.fc0(x) + x = self.activation(x) + x = self.fc1(x) + return x class MoonVitEncoderLayer(nn.Module):