diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 65318161d..faa1f6225 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -29,6 +29,9 @@ from sglang.srt.layers.moe.token_dispatcher.moriep import ( ) from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsFusedMoEMethod, +) from sglang.srt.layers.quantization.compressed_tensors.schemes import ( NPUCompressedTensorsW4A16Int4DynamicMoE, ) @@ -383,7 +386,11 @@ class DeepEPMoE(FusedMoE): else: input_quant = get_bool_env_var("DEEP_NORMAL_MODE_USE_INT8_QUANT") if not input_quant and not isinstance( - self.quant_method, NPUCompressedTensorsW4A16Int4DynamicMoE + self.quant_method, + ( + NPUCompressedTensorsW4A16Int4DynamicMoE, + CompressedTensorsFusedMoEMethod, + ), ): hidden_states, hidden_states_scale = torch_npu.npu_dynamic_quant( hidden_states diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py index 4cbfed6f9..a13c53af4 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -201,6 +201,7 @@ class CompressedTensorsConfig(QuantizationConfig): ): return self.target_scheme_map["FusedMoE"] = self.target_scheme_map["Linear"] + self.target_scheme_map["DeepEPMoE"] = self.target_scheme_map["Linear"] @property def weight_block_size(self) -> Optional[List[int]]: diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index 5e35a0d6f..f1b38b1fe 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -9,6 +9,7 @@ from torch import nn from transformers import activations from sglang.srt.configs.kimi_k25 import KimiK25Config, KimiK25VisionConfig +from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.managers.mm_utils import ( MultiModalityDataPaddingPatternMultimodalTokens, @@ -37,13 +38,15 @@ from sglang.srt.models.kimi_vl_moonvit import MLP2 from sglang.srt.models.utils import WeightsMapper from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import add_prefix +from sglang.srt.utils import add_prefix, is_npu KIMIV_VT_INFER_MAX_PATCH_NUM = 16328 logger = logging.getLogger(__name__) from sglang.srt.layers.dp_attention import is_dp_attention_enabled +_is_npu = is_npu() + def apply_rope( xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor, x_shape=None @@ -197,7 +200,7 @@ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): @get_rope_shape_decorate -@torch.compile(dynamic=True) +@torch.compile(dynamic=True, disable=_is_npu) def get_rope_shape(org, interpolation_mode, shape): return ( F.interpolate( @@ -774,5 +777,14 @@ class KimiK25ForConditionalGeneration(nn.Module): if language_weights: self.language_model.load_weights(language_weights) + @classmethod + def get_model_config_for_expert_location(cls, config: KimiK25Config): + text_config = config.text_config + return ModelConfigForExpertLocation( + num_layers=text_config.num_hidden_layers, + num_logical_experts=text_config.n_routed_experts, + num_groups=text_config.n_group, + ) + EntryClass = [KimiK25ForConditionalGeneration]