From 22012ba1bc2166f2280be2ad648ba732a0ff382b Mon Sep 17 00:00:00 2001 From: ori <39351881+froststeam@users.noreply.github.com> Date: Thu, 14 May 2026 02:06:37 +0800 Subject: [PATCH] Add BF16 support to EP-MoE for DeepGEMM (#17392) Co-authored-by: zhiguo.qin --- .../layers/deep_gemm_wrapper/compile_utils.py | 56 ++++++ .../layers/deep_gemm_wrapper/entrypoint.py | 34 ++++ .../sglang/srt/layers/moe/ep_moe/kernels.py | 163 ++++++++++++++--- python/sglang/srt/layers/moe/ep_moe/layer.py | 5 + .../srt/layers/moe/fused_moe_triton/layer.py | 5 +- .../srt/layers/moe/moe_runner/deep_gemm.py | 172 ++++++++++++++++-- .../srt/layers/moe/token_dispatcher/deepep.py | 34 ++-- .../compressed_tensors/compressed_tensors.py | 3 +- .../sglang/srt/layers/quantization/unquant.py | 25 ++- python/sglang/srt/server_args.py | 5 + 10 files changed, 450 insertions(+), 52 deletions(-) diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py index 02c31cf15..f0cc1a44b 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/compile_utils.py @@ -97,6 +97,8 @@ def update_deep_gemm_config(gpu_id: int, server_args: ServerArgs): class DeepGemmKernelType(IntEnum): GROUPED_GEMM_NT_F8F8BF16_MASKED = auto() GROUPED_GEMM_NT_F8F8BF16_CONTIG = auto() + GROUPED_GEMM_NT_BF16_MASKED = auto() + GROUPED_GEMM_NT_BF16_CONTIG = auto() GEMM_NT_F8F8BF16 = auto() GEMM_NT_BF16BF16F32 = auto() @@ -164,6 +166,9 @@ def _compile_deep_gemm_one_type_all( if kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG: m_alignment = deep_gemm.get_mk_alignment_for_contiguous_layout() m_list = sorted(list(set(m for m in m_list if m % m_alignment == 0))) + elif kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_BF16_CONTIG: + m_alignment = deep_gemm.get_mk_alignment_for_contiguous_layout() + m_list = sorted(list(set(m for m in m_list if m % m_alignment == 0))) # Here the precompilation is only run on the first rank, so gpu_id should be 0 memory_budget = get_available_gpu_memory(device="cuda", gpu_id=0) @@ -226,6 +231,8 @@ class _BaseWarmupExecutor: DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG: _GroupedContWarmupExecutor, DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED: _GroupedMaskedWarmupExecutor, DeepGemmKernelType.GEMM_NT_BF16BF16F32: _BF16F32WarmupExecutor, + DeepGemmKernelType.GROUPED_GEMM_NT_BF16_CONTIG: _BF16GroupedContWarmupExecutor, + DeepGemmKernelType.GROUPED_GEMM_NT_BF16_MASKED: _BF16GroupedMaskedWarmupExecutor, }[kernel_type](**kwargs) @staticmethod @@ -238,6 +245,10 @@ class _BaseWarmupExecutor: return (max_m * k + n * k + max_m * n * 2) / _GB elif kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_CONTIG: return (max_m * k + num_groups * n * k + max_m * 4 + max_m * n * 2) / _GB + elif kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_BF16_CONTIG: + return ( + max_m * k * 2 + num_groups * n * k * 2 + max_m * 4 + max_m * n * 2 + ) / _GB elif kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_F8F8BF16_MASKED: return ( num_groups * max_m * k @@ -248,6 +259,13 @@ class _BaseWarmupExecutor: elif kernel_type == DeepGemmKernelType.GEMM_NT_BF16BF16F32: # bf16 lhs + bf16 rhs + fp32 out return (max_m * k * 2 + n * k * 2 + max_m * n * 4) / _GB + elif kernel_type == DeepGemmKernelType.GROUPED_GEMM_NT_BF16_MASKED: + return ( + num_groups * max_m * k * 2 + + num_groups * n * k * 2 + + num_groups * 4 + + num_groups * max_m * n * 2 + ) / _GB else: raise ValueError(f"Invalid kernel type: {kernel_type}") @@ -310,6 +328,22 @@ class _GroupedContWarmupExecutor(_BaseWarmupExecutor): ) +class _BF16GroupedContWarmupExecutor(_BaseWarmupExecutor): + def __init__(self, max_m: int, n: int, k: int, num_groups: int): + self.a = torch.empty((max_m, k), device="cuda", dtype=torch.bfloat16) + self.b = torch.empty((num_groups, n, k), device="cuda", dtype=torch.bfloat16) + self.m_indices = torch.zeros((max_m,), device="cuda", dtype=torch.int32) + self.out = torch.empty((max_m, n), device="cuda", dtype=torch.bfloat16) + + def execute(self, m): + deep_gemm.m_grouped_bf16_gemm_nt_contiguous( + self.a[:m], + self.b, + self.out[:m], + m_indices=self.m_indices[:m], + ) + + class _GroupedMaskedWarmupExecutor(_BaseWarmupExecutor): def __init__(self, max_m: int, n: int, k: int, num_groups: int): self.lhs_q, self.lhs_s = _empty_token_fp8((num_groups, max_m, k)) @@ -340,6 +374,28 @@ class _BF16F32WarmupExecutor(_BaseWarmupExecutor): deep_gemm.bf16_gemm_nt(self.lhs[:m], self.rhs, self.out[:m]) +class _BF16GroupedMaskedWarmupExecutor(_BaseWarmupExecutor): + def __init__(self, max_m: int, n: int, k: int, num_groups: int): + self.a = torch.empty( + (num_groups, max_m, k), device="cuda", dtype=torch.bfloat16 + ) + self.b = torch.empty((num_groups, n, k), device="cuda", dtype=torch.bfloat16) + self.masked_m = torch.zeros((num_groups,), device="cuda", dtype=torch.int32) + self.out = torch.empty( + (num_groups, max_m, n), device="cuda", dtype=torch.bfloat16 + ) + + def execute(self, m): + deep_gemm.m_grouped_bf16_gemm_nt_masked( + self.a, + self.b, + self.out, + masked_m=self.masked_m, + # DeepGEMM uses `expect_m` instead of input shape for `get_best_config` + expected_m=m, + ) + + def deep_gemm_execution_hook( m: int, n: int, k: int, num_groups: int, kernel_type: DeepGemmKernelType ): diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index 37499c524..764b5345b 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py @@ -86,6 +86,29 @@ def _ensure_cuda( ) +def grouped_gemm_nt_bf16_masked( + a: torch.Tensor, + b: torch.Tensor, + d: torch.Tensor, + masked_m: torch.Tensor, + expected_m: int, +): + num_groups, _, k = a.shape + _, n, _ = b.shape + kernel_type = compile_utils.DeepGemmKernelType.GROUPED_GEMM_NT_BF16_MASKED + + with compile_utils.deep_gemm_execution_hook( + expected_m, n, k, num_groups, kernel_type + ): + return deep_gemm.m_grouped_bf16_gemm_nt_masked( + a, + b, + d, + masked_m, + expected_m, + ) + + def grouped_gemm_nt_f8f8bf16_contig( lhs: Tuple[torch.Tensor, torch.Tensor], rhs: Tuple[torch.Tensor, torch.Tensor], @@ -116,6 +139,17 @@ def grouped_gemm_nt_f8f8bf16_contig( ) +def grouped_gemm_nt_bf16_contig( + a: torch.Tensor, b: torch.Tensor, d: torch.Tensor, m_indices: torch.Tensor +): + m, k = a.shape + num_groups, n, _ = b.shape + kernel_type = compile_utils.DeepGemmKernelType.GROUPED_GEMM_NT_BF16_CONTIG + + with compile_utils.deep_gemm_execution_hook(m, n, k, num_groups, kernel_type): + deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a, b, d, m_indices) + + def gemm_nt_f8f8bf16( lhs: Tuple[torch.Tensor, torch.Tensor], rhs: Tuple[torch.Tensor, torch.Tensor], diff --git a/python/sglang/srt/layers/moe/ep_moe/kernels.py b/python/sglang/srt/layers/moe/ep_moe/kernels.py index 40de48e72..7def543f8 100644 --- a/python/sglang/srt/layers/moe/ep_moe/kernels.py +++ b/python/sglang/srt/layers/moe/ep_moe/kernels.py @@ -431,6 +431,114 @@ def silu_and_mul_masked_post_quant_fwd( return +@triton.jit +def _silu_and_mul_kernel( + input_ptr, + stride_input_0, + stride_input_1, + stride_input_2, + output_ptr, + stride_output_0, + stride_output_1, + stride_output_2, + masked_m_ptr, + size_n, + BLOCK_N: tl.constexpr, + NUM_STAGE: tl.constexpr, +): + expert_id = tl.program_id(2) + token_id = tl.program_id(1) + hidden_dim_block_index = tl.program_id(0) + + block_num_per_expert = tl.num_programs(1) + + token_num_cur_expert = tl.load(masked_m_ptr + expert_id) + + stride_input_0 = tl.cast(stride_input_0, dtype=tl.int64) + stride_output_0 = tl.cast(stride_output_0, dtype=tl.int64) + stride_input_1 = tl.cast(stride_input_1, dtype=tl.int64) + stride_output_1 = tl.cast(stride_output_1, dtype=tl.int64) + + offs_in_d = hidden_dim_block_index * BLOCK_N + tl.arange(0, BLOCK_N) + input_ptr_offs = input_ptr + expert_id * stride_input_0 + offs_in_d + output_ptr_offs = output_ptr + expert_id * stride_output_0 + offs_in_d + + for token_index in tl.range( + token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE + ): + gate = tl.load( + input_ptr_offs + token_index * stride_input_1, + mask=offs_in_d < size_n, + other=0.0, + ).to(tl.float32) + up = tl.load( + input_ptr_offs + token_index * stride_input_1 + size_n, + mask=offs_in_d < size_n, + other=0.0, + ) + gate = gate / (1 + tl.exp(-gate)) + gate = gate.to(input_ptr.dtype.element_ty) + gate_up = up * gate + tl.store( + output_ptr_offs + token_index * stride_output_1, + gate_up, + mask=offs_in_d < size_n, + ) + + +def silu_and_mul_masked_fwd( + input: torch.Tensor, + output: torch.Tensor, + masked_m: torch.Tensor, +): + """ + input shape [expert_num, token_num_padded, hidden_dim], dtype bf16 + output shape [expert_num, token_num_padded, hidden_dim // 2], dtype bf16 + masked_m shape [expert_num] + """ + + assert input.is_contiguous() + assert output.dtype == torch.bfloat16 + assert input.dtype == torch.bfloat16 + assert output.is_contiguous() + assert len(input.shape) == 3 + assert input.shape[0] == masked_m.shape[0] + assert input.shape[-1] % 2 == 0 + + size_n = input.shape[-1] // 2 + expert_num = len(masked_m) + + if expert_num < 4: + BLOCK_NUM_PER_EXPERT = 64 + else: + BLOCK_NUM_PER_EXPERT = 32 + + BLOCK_N = 128 + num_warps = 4 + NUM_STAGES = 4 + + hidden_dim_split_block_num = triton.cdiv(size_n, BLOCK_N) + + grid = ( + hidden_dim_split_block_num, + BLOCK_NUM_PER_EXPERT, + expert_num, + ) + + _silu_and_mul_kernel[grid]( + input, + *input.stride(), + output, + *output.stride(), + masked_m, + size_n, + BLOCK_N=BLOCK_N, + NUM_STAGE=NUM_STAGES, + num_warps=num_warps, + ) + return output + + @triton.jit def silu_mul_static_tensorwise_quant_triton_kernel_for_cutlass_moe( input_ptr, @@ -669,6 +777,7 @@ def _fwd_kernel_ep_scatter_2( SCALE_HIDDEN_SIZE_PAD: tl.constexpr, # Platform-specific semaphore for atomic_add performance tuning ATOMIC_ADD_SEM: tl.constexpr, + IS_FP8: tl.constexpr, ): start_token_id = tl.program_id(0) grid_num = tl.num_programs(0) @@ -682,12 +791,13 @@ def _fwd_kernel_ep_scatter_2( for token_id_int32 in range(start_token_id, total_token_num, grid_num): token_id = token_id_int32.to(tl.int64) to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask) - to_copy_s = tl.load( - recv_x_scale - + token_id * recv_x_scale_stride0 - + index_in_s * recv_x_scale_stride1, - mask=mask_s, - ) + if IS_FP8: + to_copy_s = tl.load( + recv_x_scale + + token_id * recv_x_scale_stride0 + + index_in_s * recv_x_scale_stride1, + mask=mask_s, + ) for topk_idx_int32 in tl.range(0, topk_num, 1, num_stages=4): topk_index = topk_idx_int32.to(tl.int64) @@ -705,15 +815,18 @@ def _fwd_kernel_ep_scatter_2( output_tensor_ptr = ( output_tensor + dest_token_index * output_tensor_stride0 ) - output_tensor_scale_ptr = ( - output_tensor_scale + dest_token_index * output_tensor_scale_stride0 - ) tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask) - tl.store( - output_tensor_scale_ptr + index_in_s * output_tensor_scale_stride1, - to_copy_s, - mask=mask_s, - ) + if IS_FP8: + output_tensor_scale_ptr = ( + output_tensor_scale + + dest_token_index * output_tensor_scale_stride0 + ) + tl.store( + output_tensor_scale_ptr + + index_in_s * output_tensor_scale_stride1, + to_copy_s, + mask=mask_s, + ) # copy from https://github.com/ModelTC/lightllm/blob/main/lightllm/common/fused_moe/deepep_scatter_gather.py @@ -745,10 +858,15 @@ def ep_scatter( scale_hidden_size = ceil_div(scale_hidden_size, 4) assert m_indices.shape[0] % BLOCK_E == 0 - assert ( - recv_x_scale.dtype == output_tensor_scale.dtype - ), f"recv_x_scale.dtype: {recv_x_scale.dtype}, output_tensor_scale.dtype: {output_tensor_scale.dtype}" - assert recv_x_scale.shape[1] == output_tensor_scale.shape[1] == scale_hidden_size + + is_fp8 = recv_x_scale is not None and recv_x.dtype != torch.bfloat16 + if is_fp8: + assert ( + recv_x_scale.dtype == output_tensor_scale.dtype + ), f"recv_x_scale.dtype: {recv_x_scale.dtype}, output_tensor_scale.dtype: {output_tensor_scale.dtype}" + assert ( + recv_x_scale.shape[1] == output_tensor_scale.shape[1] == scale_hidden_size + ) _fwd_kernel_ep_scatter_1[(grid,)]( num_recv_tokens_per_expert, @@ -769,8 +887,8 @@ def ep_scatter( recv_x.stride(0), recv_x.stride(1), recv_x_scale, - recv_x_scale.stride(0), - recv_x_scale.stride(1), + recv_x_scale.stride(0) if is_fp8 else 0, + recv_x_scale.stride(1) if is_fp8 else 0, recv_topk, recv_topk.stride(0), recv_topk.stride(1), @@ -778,8 +896,8 @@ def ep_scatter( output_tensor.stride(0), output_tensor.stride(1), output_tensor_scale, - output_tensor_scale.stride(0), - output_tensor_scale.stride(1), + output_tensor_scale.stride(0) if is_fp8 else 0, + output_tensor_scale.stride(1) if is_fp8 else 0, output_index, output_index.stride(0), output_index.stride(1), @@ -791,6 +909,7 @@ def ep_scatter( SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size), # XXX (MUSA): Atomic add with "relaxed" semaphore on musa backend for better performance ATOMIC_ADD_SEM=None if not _is_musa else "relaxed", + IS_FP8=is_fp8, ) return diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index f201d453a..15f5d0847 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -110,6 +110,11 @@ class DeepEPMoE(FusedMoE): quant_config, Fp8Config ): self.deprecate_flag = True + elif ( + deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + and envs.SGLANG_DEEPEP_BF16_DISPATCH.get() + ): + self.deprecate_flag = True else: self.deprecate_flag = False 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 82543626a..f61b08be5 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -225,6 +225,7 @@ class FusedMoE(torch.nn.Module): get_moe_runner_backend().is_flashinfer_trtllm() or get_moe_runner_backend().is_flashinfer_trtllm_routed() ) + self.use_deep_gemm = get_moe_runner_backend().is_deep_gemm() # flashinfer_trtllm kernel requires intermediate_size to be a multiple of 128 # Pad the intermediate_size_per_partition if necessary @@ -282,7 +283,9 @@ class FusedMoE(torch.nn.Module): self.quant_method = quant_config.get_quant_method(self, prefix) if self.quant_method is None: self.quant_method = UnquantizedFusedMoEMethod( - self.use_triton_kernels, self.use_flashinfer_trtllm_moe + self.use_triton_kernels, + self.use_flashinfer_trtllm_moe, + self.use_deep_gemm, ) self.quant_method.create_weights( diff --git a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py index cfdee4757..da6f13fcd 100644 --- a/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py +++ b/python/sglang/srt/layers/moe/moe_runner/deep_gemm.py @@ -51,6 +51,8 @@ _is_musa = is_musa() # Imported only for the SGLANG_OPT_FIX_MEGA_MOE_MEMORY=False fallback path. if not (_is_npu or _is_hip) and _is_cuda: from sglang.jit_kernel.activation import silu_and_mul as _legacy_silu_and_mul +elif _is_musa: + _silu_and_mul_musa = torch.nn.SwishGLU() else: _legacy_silu_and_mul = None @@ -139,14 +141,25 @@ class DeepGemmRunnerCore(MoeRunnerCore): running_state: dict, hooks: Optional[Any] = None, ) -> DeepGemmRunnerOutput: + weight_dtype = quant_info.w13_weight.dtype if not runner_input.use_masked_gemm: - hidden_states = self._run_contiguous_gemm( - runner_input, quant_info, running_state - ) + if weight_dtype == torch.bfloat16: + hidden_states = self._run_bf16_contiguous_gemm( + runner_input, quant_info, running_state + ) + else: + hidden_states = self._run_contiguous_gemm( + runner_input, quant_info, running_state + ) else: - hidden_states = self._run_masked_gemm( - runner_input, quant_info, running_state - ) + if weight_dtype == torch.bfloat16: + hidden_states = self._run_masked_bf16_gemm( + runner_input, quant_info, running_state + ) + else: + hidden_states = self._run_masked_gemm( + runner_input, quant_info, running_state + ) return DeepGemmRunnerOutput(hidden_states=hidden_states) def _run_contiguous_gemm( @@ -243,12 +256,15 @@ class DeepGemmRunnerCore(MoeRunnerCore): gateup_output, swiglu_limit=self.swiglu_limit ) - down_input = torch.empty( - (all_tokens, N // 2), - device=gateup_output.device, - dtype=torch.bfloat16, - ) - _legacy_silu_and_mul(gateup_output.view(-1, N), down_input) + if not _is_musa: + down_input = torch.empty( + (all_tokens, N // 2), + device=gateup_output.device, + dtype=torch.bfloat16, + ) + _legacy_silu_and_mul(gateup_output.view(-1, N), down_input) + else: + down_input = _silu_and_mul_musa(gateup_output.view(-1, N)) del gateup_output down_input_fp8, down_input_scale = sglang_per_token_group_quant_fp8( @@ -279,6 +295,71 @@ class DeepGemmRunnerCore(MoeRunnerCore): return down_output + def _run_bf16_contiguous_gemm( + self, + runner_input: DeepGemmRunnerInput, + quant_info: DeepGemmMoeQuantInfo, + running_state: dict, + ) -> torch.Tensor: + + hidden_states = runner_input.hidden_states + all_tokens = running_state["all_tokens"] + hidden_states_device = running_state["hidden_states_device"] + hidden_states_shape = running_state["hidden_states_shape"] + m_indices = runner_input.m_indices + + N = quant_info.w13_weight.size(1) + K = hidden_states_shape[1] + + w13_weight = quant_info.w13_weight + w2_weight = quant_info.w2_weight + + # GroupGemm-1: (M, K) (E, N, K) -> (M, N) + gateup_output = torch.empty( + (all_tokens, N), + device=hidden_states_device, + dtype=torch.bfloat16, + ) + + deep_gemm_wrapper.grouped_gemm_nt_bf16_contig( + hidden_states, + w13_weight, + gateup_output, + m_indices, + ) + + dispose_tensor(hidden_states) + + # Act: (M, N) -> (M, N/2) + if not _is_musa: + down_input = torch.empty( + ( + all_tokens, + N // 2, + ), + device=gateup_output.device, + dtype=torch.bfloat16, + ) + _legacy_silu_and_mul(gateup_output.view(-1, N), down_input) + else: + down_input = _silu_and_mul_musa(gateup_output.view(-1, N)) + del gateup_output + + # GroupGemm-2: (M, N/2) (E, K, N/2) -> (M, K) + down_output = torch.empty( + (all_tokens, K), + device=hidden_states_device, + dtype=torch.bfloat16, + ) + deep_gemm_wrapper.grouped_gemm_nt_bf16_contig( + down_input, + w2_weight, + down_output, + m_indices, + ) + + return down_output + def _run_masked_gemm( self, runner_input: DeepGemmRunnerInput, @@ -413,6 +494,70 @@ class DeepGemmRunnerCore(MoeRunnerCore): return down_output + def _run_masked_bf16_gemm( + self, + runner_input: DeepGemmRunnerInput, + quant_info: DeepGemmMoeQuantInfo, + running_state: dict, + ) -> torch.Tensor: + from sglang.srt.layers import deep_gemm_wrapper + from sglang.srt.layers.moe.ep_moe.kernels import silu_and_mul_masked_fwd + + hidden_states = runner_input.hidden_states + masked_m = runner_input.masked_m + expected_m = runner_input.expected_m + + w13_weight = quant_info.w13_weight + w2_weight = quant_info.w2_weight + + hidden_states_device = running_state["hidden_states_device"] + + # GroupGemm-0 + num_groups, m, k = hidden_states.shape + n = w13_weight.size(1) + gateup_output = torch.empty( + (num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16 + ) + deep_gemm_wrapper.grouped_gemm_nt_bf16_masked( + hidden_states, + w13_weight, + gateup_output, + masked_m, + expected_m, + ) + dispose_tensor(hidden_states) + + down_input = torch.empty( + ( + gateup_output.shape[0], + gateup_output.shape[1], + gateup_output.shape[2] // 2, + ), + device=hidden_states_device, + dtype=torch.bfloat16, + ) + + # Act + silu_and_mul_masked_fwd(gateup_output, down_input, masked_m) + del gateup_output + + # GroupGemm-1 + n = w2_weight.shape[1] + + down_output = torch.empty( + (num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16 + ) + deep_gemm_wrapper.grouped_gemm_nt_bf16_masked( + down_input, + w2_weight, + down_output, + masked_m, + expected_m, + ) + # Note: BF16 masked gemm doesn't support overlap_args, so no return value unpack + + return down_output + @property def runner_backend(self) -> MoeRunnerBackend: return MoeRunnerBackend.DEEP_GEMM @@ -632,7 +777,8 @@ def pre_permute_deepep_normal_to_deep_gemm( scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0, ) dispose_tensor(hidden_states) - dispose_tensor(hidden_states_scale) + if hidden_states_scale is not None: + dispose_tensor(hidden_states_scale) running_state["output_index"] = output_index diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index 790bc5b01..a6d0d754c 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -397,11 +397,14 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase): ): topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids topk_ids = topk_ids.to(torch.int64) - if ( - deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM - and not get_moe_runner_backend().is_cutlass() - and not envs.SGLANG_DEEPEP_BF16_DISPATCH.get() - ): + backend = get_moe_runner_backend() + # BF16 dispatch is needed when: + # - cutlass backend (uses different kernel) + # - deep_gemm backend with SGLANG_DEEPEP_BF16_DISPATCH enabled + need_bf16_dispatch = backend.is_cutlass() or ( + backend.is_deep_gemm() and envs.SGLANG_DEEPEP_BF16_DISPATCH.get() + ) + if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and not need_bf16_dispatch: # TODO hard code 128 block quant,use fp8 communication hidden_states = sglang_per_token_group_quant_fp8( hidden_states, @@ -624,14 +627,19 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): input_global_scale = self.quant_config.get("input_global_scale", None) if input_global_scale is not None: use_nvfp4 = True - elif not get_moe_runner_backend().is_flashinfer_cutedsl() and ( - not _is_npu or not envs.SGLANG_DEEPEP_BF16_DISPATCH.get() - ): - # flashinfer_cutedsl expects BF16 dispatch when NVFP4 dispatch is - # off; its kernel quantizes to NVFP4 internally. - # SGLANG_DEEPEP_BF16_DISPATCH forces BF16 dispatch for NPU - # where INT8 input + BF16 weight GMM is not supported. - use_fp8 = True + else: + backend = get_moe_runner_backend() + # BF16 dispatch is needed when: + # - flashinfer_cutedsl: kernel quantizes to NVFP4 internally + # - NPU with SGLANG_DEEPEP_BF16_DISPATCH: INT8 input + BF16 weight GMM not supported + # - deep_gemm with SGLANG_DEEPEP_BF16_DISPATCH: user requests BF16 dispatch + need_bf16_dispatch = ( + backend.is_flashinfer_cutedsl() + or (_is_npu and envs.SGLANG_DEEPEP_BF16_DISPATCH.get()) + or (backend.is_deep_gemm() and envs.SGLANG_DEEPEP_BF16_DISPATCH.get()) + ) + if not need_bf16_dispatch: + use_fp8 = True # round_scale / use_ue8m0 are FP8-DeepGEMM specific; they cause DeepEP # to return int32-packed UE8M0 scales that don't feed the flashinfer 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 f276fca11..6eff5999f 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -182,8 +182,9 @@ class CompressedTensorsConfig(QuantizationConfig): use_flashinfer_trtllm_moe = ( get_moe_runner_backend().is_flashinfer_trtllm() ) + use_deep_gemm = get_moe_runner_backend().is_deep_gemm() return UnquantizedFusedMoEMethod( - use_triton_kernels, use_flashinfer_trtllm_moe + use_triton_kernels, use_flashinfer_trtllm_moe, use_deep_gemm ) return CompressedTensorsFusedMoEMethod(self) return None diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index a3a381171..e6247fce9 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -9,6 +9,7 @@ import torch import torch.nn.functional as F from torch.nn.parameter import Parameter +from sglang.srt.environ import envs from sglang.srt.layers.amx_utils import ( CPUQuantMethod, _amx_process_weight_after_loading, @@ -163,13 +164,17 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): """MoE method without quantization.""" def __init__( - self, use_triton_kernels: bool = False, use_flashinfer_trtllm_moe: bool = False + self, + use_triton_kernels: bool = False, + use_flashinfer_trtllm_moe: bool = False, + use_deep_gemm: bool = False, ): super().__init__() self.use_flashinfer_cutlass = get_moe_runner_backend().is_flashinfer_cutlass() self.use_triton_kernels = use_triton_kernels self.with_bias = False self.use_flashinfer_trtllm_moe = use_flashinfer_trtllm_moe + self.use_deep_gemm = use_deep_gemm self._cache_permute_indices = dict({}) def create_weights( @@ -375,6 +380,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): if get_moe_runner_backend().is_flashinfer_trtllm_routed() else MoeRunnerBackend.FLASHINFER_TRTLLM ) + elif self.use_deep_gemm: + backend = MoeRunnerBackend.DEEP_GEMM elif self.use_triton_kernels: backend = MoeRunnerBackend.TRITON_KERNELS else: @@ -416,7 +423,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output moe_runner_config = self.moe_runner_config @@ -433,7 +439,22 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): w2_bias=getattr(layer, "w2_weight_bias", None), ) return self.runner.run(dispatch_output, quant_info) + elif self.runner.runner_backend.is_deep_gemm(): + w13_weight = layer.w13_weight + w2_weight = layer.w2_weight + from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmMoeQuantInfo + + # Only use_fp8=False when SGLANG_DEEPEP_BF16_DISPATCH is true, + # otherwise use_fp8=True for FP8 dispatch path + use_fp8 = not envs.SGLANG_DEEPEP_BF16_DISPATCH.get() + quant_info = DeepGemmMoeQuantInfo( + w13_weight=w13_weight, + w2_weight=w2_weight, + use_fp8=use_fp8, + ) + return self.runner.run(dispatch_output, quant_info) elif self.use_flashinfer_cutlass: + topk_output = dispatch_output.topk_output output = flashinfer_cutlass_fused_moe( input=x, token_selected_experts=topk_output.topk_ids, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 0391dd2f0..afec6de73 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2060,6 +2060,11 @@ class ServerArgs: logger.warning( "Detected ROCm with SGLANG_USE_AITER for GPT-OSS bf16 model, using triton MOE kernel." ) + elif is_musa() and envs.SGLANG_DEEPEP_BF16_DISPATCH.get(): + self.moe_runner_backend = "deep_gemm" + logger.warning( + "Detected MUSA with SGLANG_DEEPEP_BF16_DISPATCH for bf16 model, using deep_gemm kernel." + ) elif ( self.ep_size == 1 and is_triton_kernels_available()