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 f3fc5a544..72483f4ea 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -291,13 +291,18 @@ class FusedMoE(torch.nn.Module): self.quant_method.create_moe_runner(self, self.moe_runner_config) self.dispatcher = create_moe_dispatcher(self.moe_runner_config) - self.should_fuse_routed_scaling_factor_in_topk = isinstance( - self.quant_method, ModelOptNvFp4FusedMoEMethod - ) or ( - isinstance(self.quant_method, Fp8MoEMethod) - and ( - get_moe_runner_backend().is_cutlass() - or get_moe_runner_backend().is_flashinfer_trtllm_routed() + self.should_fuse_routed_scaling_factor_in_topk = ( + isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod) + or ( + isinstance(self.quant_method, Fp8MoEMethod) + and ( + get_moe_runner_backend().is_cutlass() + or get_moe_runner_backend().is_flashinfer_trtllm_routed() + ) + ) + or ( + isinstance(self.quant_method, UnquantizedFusedMoEMethod) + and get_moe_runner_backend().is_flashinfer_trtllm_routed() ) ) diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 3ccdfd66f..4494f195f 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -47,6 +47,10 @@ elif is_cuda_alike(): else: fp4_quantize = None +_flashinfer_trtllm_shuffle_row_indices_cache_mxfp8: dict[ + tuple, dict[str, torch.Tensor] +] = {} + def align_fp8_moe_weights_for_flashinfer_trtllm( layer: Module, swap_w13_halves: bool = False @@ -126,10 +130,13 @@ def align_fp8_moe_weights_for_flashinfer_trtllm( def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None: """Prepare MXFP8 MoE weights/scales for FlashInfer TRT-LLM kernels.""" - from flashinfer import ( - reorder_rows_for_gated_act_gemm, - shuffle_matrix_a, - shuffle_matrix_sf_a, + from flashinfer import block_scale_interleave + from flashinfer.fused_moe.core import ( + get_reorder_rows_for_gated_act_gemm_row_indices, + ) + from flashinfer.utils import ( + get_shuffle_matrix_a_row_indices, + get_shuffle_matrix_sf_a_row_indices, ) w13_weight = cast(torch.Tensor, layer.w13_weight).contiguous() @@ -144,52 +151,93 @@ def align_mxfp8_moe_weights_for_flashinfer_trtllm(layer: Module) -> None: _, hidden_size, _ = w2_weight.shape epilogue_tile_m = 128 - w13_interleaved = [ - reorder_rows_for_gated_act_gemm(w13_weight[i]) for i in range(num_experts) - ] - w13_scale_interleaved = [ - reorder_rows_for_gated_act_gemm(w13_scale[i]) for i in range(num_experts) - ] + # Reuse precomputed row-index transforms whenever shape/device are unchanged. + w13_weight_u8 = w13_weight.view(torch.uint8) + w2_weight_u8 = w2_weight.view(torch.uint8) + cache_key = ( + two_n, + hidden_size, + w2_weight.shape[-1], + w13_scale.shape[-1], + w2_scale.shape[-1], + epilogue_tile_m, + (w13_weight.device.type, w13_weight.device.index), + (w2_weight.device.type, w2_weight.device.index), + (w13_scale.device.type, w13_scale.device.index), + (w2_scale.device.type, w2_scale.device.index), + ) + cache = _flashinfer_trtllm_shuffle_row_indices_cache_mxfp8.get(cache_key) + if cache is None: + reorder_row_indices = get_reorder_rows_for_gated_act_gemm_row_indices( + w13_weight_u8[0] + ).to(w13_weight.device) + w13_shuffle_row_indices = get_shuffle_matrix_a_row_indices( + w13_weight_u8[0], epilogue_tile_m + ).to(w13_weight.device) + w2_shuffle_row_indices = get_shuffle_matrix_a_row_indices( + w2_weight_u8[0], epilogue_tile_m + ).to(w2_weight.device) + w13_scale_shuffle_row_indices = get_shuffle_matrix_sf_a_row_indices( + w13_scale[0].reshape(two_n, -1), epilogue_tile_m + ).to(w13_scale.device) + w2_scale_shuffle_row_indices = get_shuffle_matrix_sf_a_row_indices( + w2_scale[0].reshape(hidden_size, -1), epilogue_tile_m + ).to(w2_scale.device) + cache = { + "reorder_row_indices": reorder_row_indices, + "w13_shuffle_row_indices": w13_shuffle_row_indices, + "w2_shuffle_row_indices": w2_shuffle_row_indices, + "w13_scale_shuffle_row_indices": w13_scale_shuffle_row_indices, + "w2_scale_shuffle_row_indices": w2_scale_shuffle_row_indices, + } + _flashinfer_trtllm_shuffle_row_indices_cache_mxfp8[cache_key] = cache - w13_shuffled = [ - shuffle_matrix_a(w13_interleaved[i].view(torch.uint8), epilogue_tile_m) - for i in range(num_experts) - ] - w2_shuffled = [ - shuffle_matrix_a(w2_weight[i].view(torch.uint8), epilogue_tile_m) - for i in range(num_experts) - ] - w13_scale_shuffled = [ - shuffle_matrix_sf_a( - w13_scale_interleaved[i].view(torch.uint8).reshape(two_n, -1), - epilogue_tile_m, + reorder_row_indices = cache["reorder_row_indices"] + w13_shuffle_row_indices = cache["w13_shuffle_row_indices"] + w2_shuffle_row_indices = cache["w2_shuffle_row_indices"] + w13_scale_shuffle_row_indices = cache["w13_scale_shuffle_row_indices"] + w2_scale_shuffle_row_indices = cache["w2_scale_shuffle_row_indices"] + + w13_shuffled_u8 = torch.empty_like(w13_weight_u8) + w2_shuffled_u8 = torch.empty_like(w2_weight_u8) + w13_scale_shuffled = torch.empty_like(w13_scale) + w2_scale_shuffled = torch.empty_like(w2_scale) + + for i in range(num_experts): + w13_interleaved_u8 = w13_weight_u8[i].index_select(0, reorder_row_indices) + w13_scale_interleaved = w13_scale[i].index_select(0, reorder_row_indices) + + w13_shuffled_u8[i].copy_( + w13_interleaved_u8.index_select(0, w13_shuffle_row_indices) ) - for i in range(num_experts) - ] - w2_scale_shuffled = [ - shuffle_matrix_sf_a( - w2_scale[i].view(torch.uint8).reshape(hidden_size, -1), - epilogue_tile_m, + w2_shuffled_u8[i].copy_(w2_weight_u8[i].index_select(0, w2_shuffle_row_indices)) + + w13_scale_linear = w13_scale_interleaved.reshape(two_n, -1) + w13_scale_shuffled[i].copy_( + block_scale_interleave( + w13_scale_linear.index_select(0, w13_scale_shuffle_row_indices) + ).reshape_as(w13_scale_shuffled[i]) + ) + + w2_scale_linear = w2_scale[i].reshape(hidden_size, -1) + w2_scale_shuffled[i].copy_( + block_scale_interleave( + w2_scale_linear.index_select(0, w2_scale_shuffle_row_indices) + ).reshape_as(w2_scale_shuffled[i]) ) - for i in range(num_experts) - ] # Keep parameter identities stable for CUDA graph capture reuse. - copy_or_rebind_param( - layer, "w13_weight", torch.stack(w13_shuffled).view(torch.float8_e4m3fn) - ) - copy_or_rebind_param( - layer, "w2_weight", torch.stack(w2_shuffled).view(torch.float8_e4m3fn) - ) + copy_or_rebind_param(layer, "w13_weight", w13_shuffled_u8.view(torch.float8_e4m3fn)) + copy_or_rebind_param(layer, "w2_weight", w2_shuffled_u8.view(torch.float8_e4m3fn)) copy_or_rebind_param( layer, "w13_weight_scale_inv", - torch.stack(w13_scale_shuffled).reshape_as(w13_scale).contiguous(), + w13_scale_shuffled.contiguous(), ) copy_or_rebind_param( layer, "w2_weight_scale_inv", - torch.stack(w2_scale_shuffled).reshape_as(w2_scale).contiguous(), + w2_scale_shuffled.contiguous(), ) layer.w13_weight_scale_inv.format_ue8m0 = True layer.w2_weight_scale_inv.format_ue8m0 = True diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 9d0d15716..45ac15649 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -93,6 +93,7 @@ if TYPE_CHECKING: from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput from sglang.srt.layers.moe.topk import TopKOutput from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config + from sglang.srt.models.utils import WeightsMapper _is_hip = is_hip() _is_cuda = is_cuda() @@ -241,6 +242,12 @@ class Fp8Config(QuantizationConfig): def get_scaled_act_names(self) -> List[str]: return [] + def apply_weight_name_mapper(self, hf_to_sglang_mapper: "WeightsMapper"): + if self.ignored_layers: + self.ignored_layers = list( + dict.fromkeys(hf_to_sglang_mapper.apply_list(self.ignored_layers)) + ) + class Fp8LinearMethod(LinearMethodBase): """Linear method for FP8.