[Lora] Lora kimi support (#22381)

python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py

@@ -682,6 +682,16 @@ class CompressedTensorsConfig(QuantizationConfig):
                     logger.info_once("Using CompressedTensorsWNA16TritonMoE (ROCm)")
                     return CompressedTensorsWNA16TritonMoE(self)
                 else:
+                    from sglang.srt.server_args import get_global_server_args
+
+                    server_args = get_global_server_args()
+                    if server_args and server_args.enable_lora:
+                        logger.info_once(
+                            "Using CompressedTensorsWNA16TritonMoEMethod "
+                            "(LoRA requires triton-compatible MoE weights)"
+                        )
+                        return CompressedTensorsWNA16TritonMoE(self)
+
                     logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
                     return CompressedTensorsWNA16MoE(self)
             else:
@@ -997,6 +1007,9 @@ class CompressedTensorsFusedMoEMethod(FusedMoEMethodBase):
     ):
         return layer.scheme.create_moe_runner(layer, moe_runner_config)
 
+    def get_triton_quant_info(self, layer: torch.nn.Module):
+        return layer.scheme.get_triton_quant_info(layer)
+
     def apply(
         self,
         layer: torch.nn.Module,

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py

@@ -448,18 +448,10 @@ class CompressedTensorsWNA16TritonMoE(CompressedTensorsWNA16MoE):
         self.moe_runner_config = moe_runner_config
         self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
 
-    def apply_weights(
-        self,
-        layer: torch.nn.Module,
-        dispatch_output: "StandardDispatchOutput",
-    ) -> "CombineInput":
+    def get_triton_quant_info(self, layer):
         from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
 
-        assert (
-            self.moe_runner_config.activation == "silu"
-        ), "Only SiLU activation is supported."
-
-        quant_info = TritonMoeQuantInfo(
+        return TritonMoeQuantInfo(
             w13_weight=layer.w13_weight_packed,
             w2_weight=layer.w2_weight_packed,
             use_int4_w4a16=True,
@@ -467,6 +459,17 @@ class CompressedTensorsWNA16TritonMoE(CompressedTensorsWNA16MoE):
             w2_scale=layer.w2_weight_scale,
             block_shape=[0, self.group_size],
         )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: "StandardDispatchOutput",
+    ) -> "CombineInput":
+        assert (
+            self.moe_runner_config.activation == "silu"
+        ), "Only SiLU activation is supported."
+
+        quant_info = self.get_triton_quant_info(layer)
         return self.runner.run(dispatch_output, quant_info)
 
 

python/sglang/srt/lora/layers.py

@@ -809,10 +809,17 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
         )
 
         # initialize triton_lora moe runner for batches with lora enabled
+        from sglang.srt.layers.moe import MoeRunnerBackend
         from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
 
+        qm = base_layer.quant_method
+        if hasattr(qm, "runner") and qm.runner is not None:
+            runner_backend = qm.runner.runner_backend
+        else:
+            runner_backend = MoeRunnerBackend.TRITON
+
         self._lora_runner = MoeRunner(
-            base_layer.quant_method.runner.runner_backend,
+            runner_backend,
             base_layer.moe_runner_config,
             lora_enabled=True,
         )

python/sglang/srt/lora/lora_manager.py

@@ -66,7 +66,10 @@ class LoRAManager:
         lora_paths: Optional[List[LoRARef]] = None,
     ):
         self.base_model: torch.nn.Module = base_model
-        self.base_hf_config: AutoConfig = base_hf_config
+        if hasattr(base_hf_config, "get_text_config"):
+            self.base_hf_config: AutoConfig = base_hf_config.get_text_config()
+        else:
+            self.base_hf_config: AutoConfig = base_hf_config
         self.max_loras_per_batch: int = max_loras_per_batch
         self.load_config: LoadConfig = load_config
         self.dtype: torch.dtype = dtype