diff --git a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py index 7e63d5c88..6ce5a3ddc 100644 --- a/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py +++ b/python/sglang/srt/layers/moe/cutlass_w4a8_moe.py @@ -25,7 +25,6 @@ from sglang.srt.layers.moe.ep_moe.kernels import ( deepep_permute_triton_kernel, deepep_post_reorder_triton_kernel, deepep_run_moe_deep_preprocess, - fp8_per_token_to_per_tensor_quant_triton, post_reorder_for_cutlass_moe, pre_reorder_for_cutlass_moe, silu_and_mul_masked_post_per_tensor_quant_fwd, @@ -412,8 +411,7 @@ def cutlass_w4a8_moe_deepep_normal( def cutlass_w4a8_moe_deepep_ll( - a_states: torch.Tensor, - a_scales: torch.Tensor, + a: torch.Tensor, w1_q: torch.Tensor, w2_q: torch.Tensor, w1_scale: torch.Tensor, @@ -475,7 +473,7 @@ def cutlass_w4a8_moe_deepep_ll( """ assert w1_q.dtype == torch.int8 assert w2_q.dtype == torch.int8 - assert a_states.shape[2] // 2 == w1_q.shape[2], "Hidden size mismatch w1" + assert a.shape[2] // 2 == w1_q.shape[2], "Hidden size mismatch w1" assert w1_q.shape[2] * 2 == w2_q.shape[1], "Hidden size mismatch w2" assert w1_q.shape[0] == w2_q.shape[0], "Expert number mismatch" assert w1_q.shape[0] == w1_scale.shape[0], "w1 scales expert number mismatch" @@ -486,12 +484,12 @@ def cutlass_w4a8_moe_deepep_ll( assert a_strides2.shape[0] == w2_q.shape[0], "A Strides 2 expert number mismatch" assert b_strides2.shape[0] == w2_q.shape[0], "B Strides 2 expert number mismatch" num_experts = w1_q.size(0) - m = a_states.size(1) + m = a.size(1) k = w1_q.size(2) * 2 # w1_q is transposed and packed n = w2_q.size(2) * 2 # w2_q is transposed and packed topk = topk_ids_.size(1) - device = a_states.device + device = a.device problem_sizes1, problem_sizes2 = deepep_ll_get_cutlass_w4a8_moe_mm_data( masked_m, @@ -502,14 +500,8 @@ def cutlass_w4a8_moe_deepep_ll( k, ) - gateup_input = torch.empty(a_states.shape, dtype=torch.float8_e4m3fn, device=device) - fp8_per_token_to_per_tensor_quant_triton( - x=a_states, - x_scale=a_scales, - masked_m=masked_m, - output_scale=a1_scale, - output=gateup_input, - ) + gateup_input = torch.empty(a.shape, dtype=torch.float8_e4m3fn, device=device) + per_tensor_quant_fp8(a, gateup_input, a1_scale.float(), True) c1 = torch.empty((num_experts, m, n * 2), device=device, dtype=torch.bfloat16) c2 = torch.empty((num_experts, m, k), device=device, dtype=torch.bfloat16) @@ -530,7 +522,7 @@ def cutlass_w4a8_moe_deepep_ll( ) intermediate_q = torch.empty( - (num_experts, m, n), device=a_states.device, dtype=torch.float8_e4m3fn + (num_experts, m, n), device=a.device, dtype=torch.float8_e4m3fn ) silu_and_mul_masked_post_per_tensor_quant_fwd( c1, intermediate_q, masked_m, a2_scale diff --git a/python/sglang/srt/layers/moe/ep_moe/kernels.py b/python/sglang/srt/layers/moe/ep_moe/kernels.py index 4cd4b4f81..044c590f2 100644 --- a/python/sglang/srt/layers/moe/ep_moe/kernels.py +++ b/python/sglang/srt/layers/moe/ep_moe/kernels.py @@ -1381,76 +1381,3 @@ def silu_and_mul_masked_post_per_tensor_quant_fwd( NUM_STAGE=NUM_STAGES, ) return output - - -@triton.jit -def _fp8_per_token_quant_to_per_tensor_quant_kernel( - x_ptr, - x_scale_ptr, - x_scale_stride0, - x_scale_stride1, - x_scale_stride2, - masked_m_ptr, - output_scale_ptr, - output_ptr, - m, - k, - K_SCALE_BLOCK_SIZE: tl.constexpr, - K_BLOCK_SIZE: tl.constexpr, -): - pid_k, pid_m, pid_e = ( - tl.program_id(axis=0), - tl.program_id(axis=1), - tl.program_id(axis=2), - ) - pid_m_dim = tl.num_programs(1) - - token_id = pid_m - last_effective_id = tl.load(masked_m_ptr + pid_e) - - if token_id >= last_effective_id: - return - output_scale_val_inv = 1.0 / tl.load(output_scale_ptr).to(tl.float32) - k_offsets = pid_k * K_BLOCK_SIZE + tl.arange(0, K_BLOCK_SIZE) - scale_offsets = (k_offsets // K_SCALE_BLOCK_SIZE) * x_scale_stride2 - - x_ptrs = x_ptr + pid_e * m * k + k_offsets - output_ptrs = output_ptr + pid_e * m * k + k_offsets - x_scale_ptrs = x_scale_ptr + pid_e * x_scale_stride0 + scale_offsets - - for tok_idx in tl.range(token_id, last_effective_id, pid_m_dim): - hidden = tl.load(x_ptrs + tok_idx * k).to(tl.float32) - scale_fp32 = tl.load(x_scale_ptrs + tok_idx * x_scale_stride1).to(tl.float32) - hidden = hidden * scale_fp32 * output_scale_val_inv - tl.store(output_ptrs + tok_idx * k, hidden.to(output_ptr.dtype.element_ty)) - - -def fp8_per_token_to_per_tensor_quant_triton( - x: torch.Tensor, - x_scale: torch.Tensor, - masked_m: torch.Tensor, - output_scale: torch.Tensor, - output: torch.Tensor, -): - K_SCALE_BLOCK_SIZE = 128 - assert len(x.shape) == 3 and x.size(2) % K_SCALE_BLOCK_SIZE == 0 - assert x.is_contiguous() - assert x_scale.size(2) == x.size(2) // K_SCALE_BLOCK_SIZE - assert output_scale.numel() == 1 - - K_BLOCK_SIZE = 1024 - assert x.size(2) % K_BLOCK_SIZE == 0 - grid = (x.size(2) // K_BLOCK_SIZE, 32, x.size(0)) - _fp8_per_token_quant_to_per_tensor_quant_kernel[grid]( - x, - x_scale, - *x_scale.stride(), - masked_m, - output_scale, - output, - x.size(1), - x.size(2), - K_SCALE_BLOCK_SIZE=K_SCALE_BLOCK_SIZE, - K_BLOCK_SIZE=K_BLOCK_SIZE, - num_warps=8, - ) diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 10f825d4f..ef2225ecc 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -331,9 +331,6 @@ class DeepEPMoE(FusedMoE): ): assert self.moe_runner_config.activation == "silu" assert isinstance(self.quant_method, W4AFp8MoEMethod) - assert ( - envs.SGLANG_DEEPEP_BF16_DISPATCH.get() - ), "W4AFP8 does not support FP8 normal dispatch; please set SGLANG_DEEPEP_BF16_DISPATCH=1." return self.quant_method.apply_deepep_normal( layer=self, dispatch_output=dispatch_output, @@ -345,6 +342,9 @@ class DeepEPMoE(FusedMoE): ): assert self.moe_runner_config.activation == "silu" assert isinstance(self.quant_method, W4AFp8MoEMethod) + assert ( + envs.SGLANG_DEEPEP_BF16_DISPATCH.get() + ), "W4AFP8 does not support FP8 dispatch; please set SGLANG_DEEPEP_BF16_DISPATCH=1." return self.quant_method.apply_deepep_ll( layer=self, dispatch_output=dispatch_output, diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py index a4339808b..8539639d5 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -609,7 +609,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): input_global_scale = self.quant_config.get("input_global_scale", None) if input_global_scale is not None: use_nvfp4 = True - else: + elif not envs.SGLANG_DEEPEP_BF16_DISPATCH.get(): use_fp8 = True buffer = self._get_buffer() diff --git a/python/sglang/srt/layers/quantization/w4afp8.py b/python/sglang/srt/layers/quantization/w4afp8.py index 42b1dc24c..83869289e 100644 --- a/python/sglang/srt/layers/quantization/w4afp8.py +++ b/python/sglang/srt/layers/quantization/w4afp8.py @@ -334,11 +334,10 @@ class W4AFp8MoEMethod(FusedMoEMethodBase): from sglang.srt.layers.moe.cutlass_w4a8_moe import cutlass_w4a8_moe_deepep_ll - hidden_states, hidden_scales, topk_ids, _, masked_m, _ = dispatch_output + hidden_states, _, topk_ids, _, masked_m, _ = dispatch_output output = cutlass_w4a8_moe_deepep_ll( hidden_states, - hidden_scales, layer.w13_weight, layer.w2_weight, layer.w13_weight_scale_inv,