diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 5f77356da..4f22ba798 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -548,8 +548,9 @@ def get_moe_impl_class(quant_config: Optional[QuantizationConfig]): quant_config is None or quant_config.get_name() == "fp8" or quant_config.get_name() == "modelopt_fp8" + or quant_config.get_name() == "compressed_tensors" ): - # FlashInferFusedMoE support bf16 and fp8 + # FlashInferFusedMoE support bf16, fp8 and compressed_tensors return FlashInferFusedMoE if get_moe_runner_backend().is_flashinfer_cutlass(): diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index b68ec3146..73fe410b8 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -1093,7 +1093,6 @@ class FlashInferFusedMoE(FusedMoE): else: - # FP8 Matrix multiply. final_hidden_states = self.quant_method.apply_with_router_logits( layer=self, dispatch_output=StandardDispatchOutput( diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py index 0e03c2057..b06bb0cc1 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors_moe.py @@ -11,10 +11,15 @@ import torch from compressed_tensors import CompressionFormat from compressed_tensors.quantization import QuantizationStrategy -from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.distributed import get_tensor_model_parallel_world_size, get_tp_group +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.moe.cutlass_moe_params import CutlassMoEParams, CutlassMoEType from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo +from sglang.srt.layers.moe.utils import get_moe_runner_backend from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase from sglang.srt.layers.quantization.compressed_tensors.schemes import ( WNA16_SUPPORTED_BITS, @@ -29,10 +34,18 @@ from sglang.srt.layers.quantization.marlin_utils import marlin_moe_permute_scale from sglang.srt.layers.quantization.utils import ( all_close_1d, per_tensor_dequantize, + prepare_static_weights_for_trtllm_fp4_moe, + reorder_w1w3_to_w3w1, replace_parameter, swizzle_blockscale, ) -from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip, set_weight_attrs +from sglang.srt.utils import ( + get_bool_env_var, + is_cuda, + is_hip, + next_power_of_2, + set_weight_attrs, +) if TYPE_CHECKING: from sglang.srt.layers.moe.fused_moe_triton import FusedMoE @@ -115,6 +128,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): ) self.quant_config = quant_config self.group_size = 16 + self.use_flashinfer_trtllm = get_moe_runner_backend().is_flashinfer_trtllm() def create_weights( self, @@ -127,7 +141,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): ): from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported - layer.num_experts = num_experts layer.params_dtype = params_dtype w13_weight = torch.nn.Parameter( @@ -240,6 +253,13 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): ) delattr(layer, "w2_weight_packed") + if self.use_flashinfer_trtllm: + w, s = reorder_w1w3_to_w3w1( + layer.w13_weight.data, layer.w13_weight_scale.data, dim=-2 + ) + layer.w13_weight = torch.nn.Parameter(w, requires_grad=False) + layer.w13_weight_scale = torch.nn.Parameter(s, requires_grad=False) + if not torch.allclose( layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1] ): @@ -258,9 +278,16 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): ) # w13 - w13_input_global_scale = layer.w13_input_global_scale.min(dim=1).values.to( - torch.float32 - ) + if self.use_flashinfer_trtllm: + w13_input_global_scale = ( + layer.w13_input_global_scale.min() + .to(torch.float32) + .expand(layer.num_local_experts) + ) + else: + w13_input_global_scale = layer.w13_input_global_scale.min(dim=1).values.to( + torch.float32 + ) layer.g1_alphas = torch.nn.Parameter( ((1 / w13_input_global_scale) * layer.w13_weight_scale_2), requires_grad=False, @@ -271,7 +298,14 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): ) # w2 - w2_input_global_scale = layer.w2_input_global_scale + if self.use_flashinfer_trtllm: + w2_input_global_scale = ( + layer.w2_input_global_scale.min() + .to(torch.float32) + .expand(layer.num_local_experts) + ) + else: + w2_input_global_scale = layer.w2_input_global_scale layer.g2_alphas = torch.nn.Parameter( ((1 / w2_input_global_scale) * layer.w2_weight_scale_2).to(torch.float32), @@ -282,22 +316,66 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): (w2_input_global_scale), requires_grad=False ) - # swizzle weight scales - layer.w13_weight_scale = torch.nn.Parameter( - swizzle_blockscale(layer.w13_weight_scale), requires_grad=False - ) + # TensorRT-LLM specific processing + if self.use_flashinfer_trtllm: + # Prepare static weights for TRT-LLM kernel + ( + gemm1_weights_fp4_shuffled, + gemm1_scales_fp4_shuffled, + gemm2_weights_fp4_shuffled, + gemm2_scales_fp4_shuffled, + ) = prepare_static_weights_for_trtllm_fp4_moe( + layer.w13_weight, + layer.w2_weight, + layer.w13_weight_scale, + layer.w2_weight_scale, + layer.w2_weight.size(-2), # hidden_size + layer.w13_weight.size(-2) // 2, # intermediate_size + layer.w13_weight.size(0), # num_experts + ) + logger.debug("Finished shuffling weights for TRT-LLM MOE") - layer.w2_weight_scale = torch.nn.Parameter( - swizzle_blockscale(layer.w2_weight_scale), requires_grad=False - ) + layer.gemm1_weights_fp4_shuffled = torch.nn.Parameter( + gemm1_weights_fp4_shuffled, requires_grad=False + ) + layer.gemm2_weights_fp4_shuffled = torch.nn.Parameter( + gemm2_weights_fp4_shuffled, requires_grad=False + ) + layer.gemm1_scales_fp4_shuffled = torch.nn.Parameter( + gemm1_scales_fp4_shuffled, requires_grad=False + ) + layer.gemm2_scales_fp4_shuffled = torch.nn.Parameter( + gemm2_scales_fp4_shuffled, requires_grad=False + ) - layer.cutlass_moe_params = CutlassMoEParams( - CutlassMoEType.BlockscaledFP4, - layer.w13_weight.device, - num_experts=layer.num_experts, - intermediate_size_per_partition=layer.w2_weight.shape[2] * 2, - hidden_size=layer.w13_weight.shape[2] * 2, - ) + # Additional parameter needed for TRT-LLM + layer.g1_scale_c = torch.nn.Parameter( + (layer.w2_input_scale_quant * layer.g1_alphas).to(torch.float32), + requires_grad=False, + ) + + # Clean up weights that won't be used by TRT-LLM + del layer.w2_weight + del layer.w2_weight_scale + del layer.w13_weight + del layer.w13_weight_scale + else: + # swizzle weight scales + layer.w13_weight_scale = torch.nn.Parameter( + swizzle_blockscale(layer.w13_weight_scale), requires_grad=False + ) + + layer.w2_weight_scale = torch.nn.Parameter( + swizzle_blockscale(layer.w2_weight_scale), requires_grad=False + ) + + layer.cutlass_moe_params = CutlassMoEParams( + CutlassMoEType.BlockscaledFP4, + layer.w13_weight.device, + num_experts=layer.num_experts, + intermediate_size_per_partition=layer.w2_weight.shape[2] * 2, + hidden_size=layer.w13_weight.shape[2] * 2, + ) def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig @@ -336,6 +414,100 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): return StandardCombineInput(hidden_states=output) + def apply_with_router_logits( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> torch.Tensor: + assert self.use_flashinfer_trtllm + + x = dispatch_output.hidden_states + topk_output = dispatch_output.topk_output + + from flashinfer import fp4_quantize, trtllm_fp4_block_scale_moe + + from sglang.srt.layers.moe.utils import RoutingMethodType + + router_logits = topk_output.router_logits + topk_config = topk_output.topk_config + + # Quantize input hidden states using fp4_quantize + hs_fp4_bytes, hs_sf_bytes = fp4_quantize( + x, + layer.w13_input_scale_quant, + self.group_size, # sf_vec_size + False, # use_ue8m0 + False, # is_sf_swizzled_layout + ) + hs_fp4 = hs_fp4_bytes.reshape(x.shape[0], x.shape[1] // 2) + hs_scale = hs_sf_bytes.view(torch.float8_e4m3fn).reshape(-1) + + correction_bias = ( + None + if topk_config.correction_bias is None + else topk_config.correction_bias.to(x.dtype) + ) + + assert layer.routing_method_type is not None + + # DeepSeekV3 style routing requires float32 router logits + if layer.routing_method_type == RoutingMethodType.DeepSeekV3: + router_logits = router_logits.to(torch.float32) + + routed_scaling_factor = self.moe_runner_config.routed_scaling_factor + routed_scaling_factor = ( + routed_scaling_factor if routed_scaling_factor is not None else 1.0 + ) + + with use_symmetric_memory( + get_tp_group(), disabled=not is_allocation_symmetric() + ): + num_tokens = hs_fp4.shape[0] + hidden_size = ( + hs_fp4.shape[-1] * 2 + if hs_fp4.dtype == torch.uint8 + else hs_fp4.shape[-1] + ) + symm_output = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_fp4.device + ) + + return trtllm_fp4_block_scale_moe( + routing_logits=router_logits, + routing_bias=correction_bias, + hidden_states=hs_fp4, + hidden_states_scale=hs_scale, + gemm1_weights=layer.gemm1_weights_fp4_shuffled, + gemm1_weights_scale=layer.gemm1_scales_fp4_shuffled.view( + torch.float8_e4m3fn + ), + gemm1_bias=None, + gemm1_alpha=None, + gemm1_beta=None, + gemm1_clamp_limit=None, + gemm2_weights=layer.gemm2_weights_fp4_shuffled, + gemm2_weights_scale=layer.gemm2_scales_fp4_shuffled.view( + torch.float8_e4m3fn + ), + gemm2_bias=None, + output1_scale_scalar=layer.g1_scale_c, + output1_scale_gate_scalar=layer.g1_alphas, + output2_scale_scalar=layer.g2_alphas, + num_experts=layer.num_experts, + top_k=topk_config.top_k, + n_group=topk_config.num_expert_group, + topk_group=topk_config.topk_group, + intermediate_size=layer.intermediate_size_per_partition, + local_expert_offset=layer.moe_ep_rank * layer.num_local_experts, + local_num_experts=layer.num_local_experts, + routed_scaling_factor=routed_scaling_factor, + tile_tokens_dim=None, + routing_method_type=layer.routing_method_type, + do_finalize=True, + tune_max_num_tokens=next_power_of_2(hs_fp4.shape[0]), + output=symm_output, + )[0] + class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 45ecd113d..d4faa2ddf 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -42,6 +42,7 @@ from sglang.srt.layers.quantization.utils import ( convert_to_channelwise, is_layer_skipped, per_tensor_dequantize, + prepare_static_weights_for_trtllm_fp4_moe, requantize_with_max_scale, swizzle_blockscale, ) @@ -1398,130 +1399,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): w2_input_scale._sglang_require_global_experts = True layer.register_parameter("w2_input_scale", w2_input_scale) - def prepare_static_weights_for_kernel( - self, - # args_dequant, - # args, - gemm1_weights, - gemm2_weights, - gemm1_scales_linear_fp4_bytes, - gemm2_scales_linear_fp4_bytes, - hidden_size, - intermediate_size, - num_experts, - ): - from flashinfer import nvfp4_block_scale_interleave - from flashinfer.fused_moe.core import ( - _maybe_get_cached_w3_w1_permute_indices, - get_w2_permute_indices_with_cache, - ) - - """Prepare quantized weights for kernel (done offline with weights).""" - epilogue_tile_m = 128 # FIXME: this depends on the kernel internals - - # Convert quantized weights to proper formats - gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape( - num_experts, 2 * intermediate_size, hidden_size // 2 - ) # packed fp4 - gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view( - torch.float8_e4m3fn - ).reshape( - num_experts, 2 * intermediate_size, hidden_size // 16 - ) # fp8 scaling factors - - gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape( - num_experts, hidden_size, intermediate_size // 2 - ) # packed fp4 - gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view( - torch.float8_e4m3fn - ).reshape( - num_experts, hidden_size, intermediate_size // 16 - ) # fp8 scaling factors - - gemm1_weights_fp4_shuffled = [] - gemm1_scales_fp4_shuffled = [] - gemm2_weights_fp4_shuffled = [] - gemm2_scales_fp4_shuffled = [] - for i in range(num_experts): - # Calculate the permute indices for the following: - # 1. Reorder rows of W1 and scales for fused gated activation - # 2. Shuffle weights and scaling factors for transposed mma output - # for both w3_w1 and w2 weights and scale factors - permute_indices = _maybe_get_cached_w3_w1_permute_indices( - self._cache_permute_indices, - gemm1_weights_fp4[i].view(torch.uint8), - epilogue_tile_m, - ) - gemm1_weights_fp4_shuffled.append( - gemm1_weights_fp4[i] - .view(torch.uint8)[permute_indices.to(gemm1_weights_fp4.device)] - .contiguous() - ) - - permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices( - self._cache_permute_indices, - gemm1_scales_linear_fp4[i].view(torch.uint8), - epilogue_tile_m, - num_elts_per_sf=16, - ) - gemm1_scales_fp4_shuffled.append( - nvfp4_block_scale_interleave( - gemm1_scales_linear_fp4[i] - .view(torch.uint8)[ - permute_sf_indices.to(gemm1_scales_linear_fp4.device) - ] - .contiguous() - ) - ) - - permute_indices = get_w2_permute_indices_with_cache( - self._cache_permute_indices, - gemm2_weights_fp4[i].view(torch.uint8), - epilogue_tile_m, - ) - gemm2_weights_fp4_shuffled.append( - gemm2_weights_fp4[i] - .view(torch.uint8)[permute_indices.to(gemm2_weights_fp4.device)] - .contiguous() - ) - - permute_sf_indices = get_w2_permute_indices_with_cache( - self._cache_permute_indices, - gemm2_scales_linear_fp4[i].view(torch.uint8), - epilogue_tile_m, - num_elts_per_sf=16, - ) - gemm2_scales_fp4_shuffled.append( - nvfp4_block_scale_interleave( - gemm2_scales_linear_fp4[i] - .view(torch.uint8)[ - permute_sf_indices.to(gemm2_scales_linear_fp4.device) - ] - .contiguous() - ) - ) - - # Stack weights for all experts - gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled) - gemm1_scales_fp4_shuffled = ( - torch.stack(gemm1_scales_fp4_shuffled) - .view(torch.float8_e4m3fn) - .reshape(num_experts, 2 * intermediate_size, hidden_size // 16) - ) - - gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled) - gemm2_scales_fp4_shuffled = ( - torch.stack(gemm2_scales_fp4_shuffled) - .view(torch.float8_e4m3fn) - .reshape(num_experts, hidden_size, intermediate_size // 16) - ) - return ( - gemm1_weights_fp4_shuffled, - gemm1_scales_fp4_shuffled, - gemm2_weights_fp4_shuffled, - gemm2_scales_fp4_shuffled, - ) - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: """Process FP4 MoE weights after loading from serialized checkpoint. @@ -1633,7 +1510,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): gemm1_scales_fp4_shuffled, gemm2_weights_fp4_shuffled, gemm2_scales_fp4_shuffled, - ) = self.prepare_static_weights_for_kernel( + ) = prepare_static_weights_for_trtllm_fp4_moe( layer.w13_weight, layer.w2_weight, layer.w13_weight_scale, diff --git a/python/sglang/srt/layers/quantization/utils.py b/python/sglang/srt/layers/quantization/utils.py index a2da44d00..d98381d36 100644 --- a/python/sglang/srt/layers/quantization/utils.py +++ b/python/sglang/srt/layers/quantization/utils.py @@ -592,3 +592,143 @@ def swizzle_blockscale(scale: torch.Tensor): if scale_ndim == 2 else swizzled_scale.reshape(B, M_padded, K_padded) ) + + +def reorder_w1w3_to_w3w1( + weight: torch.Tensor, scale: torch.Tensor, dim: int = -2 +) -> tuple[torch.Tensor, torch.Tensor]: + """Re-order the concatenated `[w1, w3]` tensors to `[w3, w1]`""" + size = weight.size(dim) + assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}" + half = size // 2 + + w1, w3 = weight.split(half, dim=dim) + s1, s3 = scale.split(half, dim=dim) + + return ( + torch.cat([w3, w1], dim=dim).contiguous(), + torch.cat([s3, s1], dim=dim).contiguous(), + ) + + +def prepare_static_weights_for_trtllm_fp4_moe( + gemm1_weights, + gemm2_weights, + gemm1_scales_linear_fp4_bytes, + gemm2_scales_linear_fp4_bytes, + hidden_size, + intermediate_size, + num_experts, +): + from flashinfer import nvfp4_block_scale_interleave + from flashinfer.fused_moe.core import ( + _maybe_get_cached_w3_w1_permute_indices, + get_w2_permute_indices_with_cache, + ) + + """Prepare quantized weights for kernel (done offline with weights).""" + _cache_permute_indices: dict[torch.Size, torch.Tensor] = {} + epilogue_tile_m = 128 # FIXME: this depends on the kernel internals + + # Convert quantized weights to proper formats + gemm1_weights_fp4 = gemm1_weights.view(torch.float8_e4m3fn).reshape( + num_experts, 2 * intermediate_size, hidden_size // 2 + ) # packed fp4 + gemm1_scales_linear_fp4 = gemm1_scales_linear_fp4_bytes.view( + torch.float8_e4m3fn + ).reshape( + num_experts, 2 * intermediate_size, hidden_size // 16 + ) # fp8 scaling factors + + gemm2_weights_fp4 = gemm2_weights.view(torch.float8_e4m3fn).reshape( + num_experts, hidden_size, intermediate_size // 2 + ) # packed fp4 + gemm2_scales_linear_fp4 = gemm2_scales_linear_fp4_bytes.view( + torch.float8_e4m3fn + ).reshape( + num_experts, hidden_size, intermediate_size // 16 + ) # fp8 scaling factors + + gemm1_weights_fp4_shuffled = [] + gemm1_scales_fp4_shuffled = [] + gemm2_weights_fp4_shuffled = [] + gemm2_scales_fp4_shuffled = [] + for i in range(num_experts): + # Calculate the permute indices for the following: + # 1. Reorder rows of W1 and scales for fused gated activation + # 2. Shuffle weights and scaling factors for transposed mma output + # for both w3_w1 and w2 weights and scale factors + permute_indices = _maybe_get_cached_w3_w1_permute_indices( + _cache_permute_indices, + gemm1_weights_fp4[i].view(torch.uint8), + epilogue_tile_m, + ) + gemm1_weights_fp4_shuffled.append( + gemm1_weights_fp4[i] + .view(torch.uint8)[permute_indices.to(gemm1_weights_fp4.device)] + .contiguous() + ) + + permute_sf_indices = _maybe_get_cached_w3_w1_permute_indices( + _cache_permute_indices, + gemm1_scales_linear_fp4[i].view(torch.uint8), + epilogue_tile_m, + num_elts_per_sf=16, + ) + gemm1_scales_fp4_shuffled.append( + nvfp4_block_scale_interleave( + gemm1_scales_linear_fp4[i] + .view(torch.uint8)[ + permute_sf_indices.to(gemm1_scales_linear_fp4.device) + ] + .contiguous() + ) + ) + + permute_indices = get_w2_permute_indices_with_cache( + _cache_permute_indices, + gemm2_weights_fp4[i].view(torch.uint8), + epilogue_tile_m, + ) + gemm2_weights_fp4_shuffled.append( + gemm2_weights_fp4[i] + .view(torch.uint8)[permute_indices.to(gemm2_weights_fp4.device)] + .contiguous() + ) + + permute_sf_indices = get_w2_permute_indices_with_cache( + _cache_permute_indices, + gemm2_scales_linear_fp4[i].view(torch.uint8), + epilogue_tile_m, + num_elts_per_sf=16, + ) + gemm2_scales_fp4_shuffled.append( + nvfp4_block_scale_interleave( + gemm2_scales_linear_fp4[i] + .view(torch.uint8)[ + permute_sf_indices.to(gemm2_scales_linear_fp4.device) + ] + .contiguous() + ) + ) + + # Stack weights for all experts + gemm1_weights_fp4_shuffled = torch.stack(gemm1_weights_fp4_shuffled) + gemm1_scales_fp4_shuffled = ( + torch.stack(gemm1_scales_fp4_shuffled) + .view(torch.float8_e4m3fn) + .reshape(num_experts, 2 * intermediate_size, hidden_size // 16) + ) + + gemm2_weights_fp4_shuffled = torch.stack(gemm2_weights_fp4_shuffled) + gemm2_scales_fp4_shuffled = ( + torch.stack(gemm2_scales_fp4_shuffled) + .view(torch.float8_e4m3fn) + .reshape(num_experts, hidden_size, intermediate_size // 16) + ) + return ( + gemm1_weights_fp4_shuffled, + gemm1_scales_fp4_shuffled, + gemm2_weights_fp4_shuffled, + gemm2_scales_fp4_shuffled, + ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 764048180..3c9cd983d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1790,8 +1790,9 @@ class ServerArgs: "modelopt_fp4", "fp8", "modelopt_fp8", + "compressed-tensors", None, - ], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'fp8', 'modelopt_fp8', or bfloat16 (None)." + ], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'fp8', 'modelopt_fp8', 'compressed-tensors', or bfloat16 (None)." self.disable_shared_experts_fusion = True logger.warning( "FlashInfer TRTLLM MoE is enabled. --disable-shared-experts-fusion is automatically set." diff --git a/python/sglang/srt/utils/mistral_utils.py b/python/sglang/srt/utils/mistral_utils.py index 7be1ba5bc..52f8769b3 100644 --- a/python/sglang/srt/utils/mistral_utils.py +++ b/python/sglang/srt/utils/mistral_utils.py @@ -81,8 +81,6 @@ def adapt_config_dict( config_dict = _remap_mistral_vision_args(config_dict) if is_audio: config_dict = _remap_mistral_audio_args(config_dict) - if is_eagle: - config_dict["routing_method_type"] = 1 # RoutingMethodType.Renormalize config = PretrainedConfig.from_dict(config_dict) @@ -234,6 +232,7 @@ def _remap_moe_args(config: dict) -> dict: config["topk_method"] = None config["scoring_func"] = "softmax" + config["routing_method_type"] = 1 # RoutingMethodType.Renormalize return config