support Kimi-K2.5-w4a8 on ascend

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]

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.

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)
+        if isinstance(self.quant_config, ModelSlimConfig):
-        x = self.activation(x)
+            x = x.flatten(0, 1)
-        return self.fc1(x)
+            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):