From b85a09b72c8a35b2e9f6b9b23cd0e246ebaac3aa Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Wed, 27 May 2026 19:27:23 +0800 Subject: [PATCH 01/17] [Feature] Support MegaMoE --- custom_ops/gpu_ops/mega_moe_pre_dispatch.cu | 320 ++++++++++++++++++ custom_ops/setup_ops.py | 5 +- .../layers/moe/fused_moe_backend_base.py | 39 ++- .../layers/moe/fused_moe_deepgemm_backend.py | 259 +++++++++++++- .../layers/quantization/fp8_utils.py | 26 ++ tests/operators/test_mega_moe_pre_dispatch.py | 177 ++++++++++ 6 files changed, 805 insertions(+), 21 deletions(-) create mode 100644 custom_ops/gpu_ops/mega_moe_pre_dispatch.cu create mode 100644 tests/operators/test_mega_moe_pre_dispatch.py diff --git a/custom_ops/gpu_ops/mega_moe_pre_dispatch.cu b/custom_ops/gpu_ops/mega_moe_pre_dispatch.cu new file mode 100644 index 00000000000..b5629e54c14 --- /dev/null +++ b/custom_ops/gpu_ops/mega_moe_pre_dispatch.cu @@ -0,0 +1,320 @@ +// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "paddle/extension.h" +#include "helper.h" + +#include +#include +#include + +#include +#include +#include + +#ifndef PD_BUILD_STATIC_OP +#define PD_BUILD_STATIC_OP(name) PD_BUILD_OP(static_op_##name) +#endif + +namespace { + +constexpr float kFP8E4M3Max = 448.0f; +constexpr uint32_t kVecElems = 8; + +template +__device__ __forceinline__ float WarpReduceMax(float value) { + static_assert(kNumThreads >= 1 && kNumThreads <= WARP_SIZE, + "kNumThreads must be in [1, 32]"); + static_assert((kNumThreads & (kNumThreads - 1)) == 0, + "kNumThreads must be a power of 2"); +#pragma unroll + for (int mask = kNumThreads / 2; mask > 0; mask >>= 1) { + value = fmaxf(value, __shfl_xor_sync(0xffffffffu, value, mask, WARP_SIZE)); + } + return value; +} + +__device__ __forceinline__ uint32_t CastToUE8M0(float value) { + value = fabsf(value); + uint32_t bits = __float_as_uint(value); + uint32_t exp = (bits >> 23) & 0xffu; + const uint32_t mantissa = bits & 0x7fffffu; + exp += mantissa != 0; + exp = min(max(exp, 1u), 254u); + return exp; +} + +struct MegaMoEPreDispatchParams { + const __nv_bfloat16* __restrict__ x; + const int64_t* __restrict__ topk_idx; + const float* __restrict__ topk_weights; + + phi::dtype::float8_e4m3fn* __restrict__ buf_x; + int32_t* __restrict__ buf_x_sf; + int64_t* __restrict__ buf_topk_idx; + float* __restrict__ buf_topk_weights; + + uint32_t num_tokens; + uint32_t padded_max; + uint32_t hidden; + uint32_t num_groups; + uint32_t top_k; +}; + +template +__global__ __launch_bounds__(1024, 2) void MegaMoEPreDispatchKernel( + const MegaMoEPreDispatchParams params) { + static_assert(kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, + "unsupported group_size"); + static_assert(kGroupSize % kVecElems == 0, + "group_size must be a multiple of 8"); + constexpr uint32_t kThreadsPerGroup = kGroupSize / kVecElems; + + const uint32_t bid = blockIdx.x; + const uint32_t tid = threadIdx.x; + + if (bid < params.num_tokens) { + const uint32_t token_id = bid; + const __nv_bfloat16* token_in = + params.x + static_cast(token_id) * params.hidden; + phi::dtype::float8_e4m3fn* token_out = + params.buf_x + static_cast(token_id) * params.hidden; + + const uint32_t base = tid * kVecElems; + float vals[kVecElems]; + float local_max = 0.0f; + +#pragma unroll + for (uint32_t i = 0; i < kVecElems; ++i) { + const float v = __bfloat162float(token_in[base + i]); + vals[i] = v; + local_max = fmaxf(local_max, fabsf(v)); + } + + local_max = WarpReduceMax(local_max); + + const float absmax = fmaxf(local_max, 1e-10f); + const float raw_scale = absmax / kFP8E4M3Max; + const uint32_t ue8m0_exp = CastToUE8M0(raw_scale); + const float inv_scale = __uint_as_float((127u + 127u - ue8m0_exp) << 23); + +#pragma unroll + for (uint32_t i = 0; i < kVecElems; ++i) { + token_out[base + i] = phi::dtype::float8_e4m3fn(vals[i] * inv_scale); + } + + const uint32_t group_id = tid / kThreadsPerGroup; + const uint32_t within_group_id = tid % kThreadsPerGroup; + if (within_group_id == 0 && group_id < params.num_groups) { + const uint32_t byte_off = token_id * params.num_groups + group_id; + reinterpret_cast(params.buf_x_sf)[byte_off] = + static_cast(ue8m0_exp); + } + + if (tid < params.top_k) { + const uint32_t off = token_id * params.top_k + tid; + params.buf_topk_idx[off] = static_cast(params.topk_idx[off]); + params.buf_topk_weights[off] = params.topk_weights[off]; + } + } +} + +void CheckShape2D(const paddle::Tensor& tensor, const char* name) { + PD_CHECK(tensor.shape().size() == 2, name, " must be a 2D tensor"); +} + +void CheckSameShape(const paddle::Tensor& lhs, + const paddle::Tensor& rhs, + const char* lhs_name, + const char* rhs_name) { + PD_CHECK(lhs.shape() == rhs.shape(), lhs_name, " shape must equal ", rhs_name, + " shape"); +} + +template +void LaunchMegaMoEPreDispatch(const MegaMoEPreDispatchParams& params, + uint32_t num_total_blocks, + uint32_t num_threads, + cudaStream_t stream) { + MegaMoEPreDispatchKernel + <<>>(params); +} + +} // namespace + +void MegaMoePreDispatch( + const paddle::Tensor& x, + const paddle::Tensor& topk_idx, + const paddle::Tensor& topk_weights, + const paddle::Tensor& buf_x, + const paddle::Tensor& buf_x_sf, + const paddle::Tensor& buf_topk_idx, + const paddle::Tensor& buf_topk_weights, + int64_t num_max_tokens_per_rank, + int64_t group_size) { + CheckShape2D(x, "x"); + CheckShape2D(topk_idx, "topk_idx"); + CheckShape2D(topk_weights, "topk_weights"); + CheckShape2D(buf_x, "buf_x"); + CheckShape2D(buf_x_sf, "buf_x_sf"); + CheckShape2D(buf_topk_idx, "buf_topk_idx"); + CheckShape2D(buf_topk_weights, "buf_topk_weights"); + CheckSameShape(topk_idx, topk_weights, "topk_idx", "topk_weights"); + CheckSameShape(buf_topk_idx, buf_topk_weights, "buf_topk_idx", + "buf_topk_weights"); + + PD_CHECK(x.dtype() == paddle::DataType::BFLOAT16, + "x must be bfloat16, but got ", x.dtype()); + PD_CHECK(topk_idx.dtype() == paddle::DataType::INT64, + "topk_idx must be int64, but got ", topk_idx.dtype()); + PD_CHECK(topk_weights.dtype() == paddle::DataType::FLOAT32, + "topk_weights must be float32, but got ", topk_weights.dtype()); + PD_CHECK(buf_x.dtype() == paddle::DataType::FLOAT8_E4M3FN, + "buf_x must be float8_e4m3fn, but got ", buf_x.dtype()); + PD_CHECK(buf_x_sf.dtype() == paddle::DataType::INT32, + "buf_x_sf must be int32, but got ", buf_x_sf.dtype()); + PD_CHECK(buf_topk_idx.dtype() == paddle::DataType::INT64, + "buf_topk_idx must be int64, but got ", buf_topk_idx.dtype()); + PD_CHECK(buf_topk_weights.dtype() == paddle::DataType::FLOAT32, + "buf_topk_weights must be float32, but got ", + buf_topk_weights.dtype()); + + const int64_t num_tokens_i64 = x.shape()[0]; + const int64_t hidden_i64 = x.shape()[1]; + const int64_t top_k_i64 = topk_idx.shape()[1]; + const int64_t padded_max_i64 = buf_x.shape()[0]; + + PD_CHECK(num_max_tokens_per_rank <= padded_max_i64, + "num_max_tokens_per_rank must not exceed buf_x.shape[0], but got ", + num_max_tokens_per_rank, " vs ", padded_max_i64); + PD_CHECK(num_tokens_i64 == topk_idx.shape()[0], + "x.shape[0] must equal topk_idx.shape[0]"); + PD_CHECK(buf_x.shape()[1] == hidden_i64, + "buf_x.shape[1] must equal hidden, but got ", buf_x.shape()[1], + " vs ", hidden_i64); + PD_CHECK(buf_topk_idx.shape()[0] == padded_max_i64, + "buf_topk_idx.shape[0] must equal padded_max"); + PD_CHECK(buf_topk_idx.shape()[1] == top_k_i64, + "buf_topk_idx.shape[1] must equal top_k"); + + PD_CHECK(group_size == 32 || group_size == 64 || group_size == 128, + "unsupported group_size: ", group_size); + PD_CHECK(num_tokens_i64 <= num_max_tokens_per_rank, + "num_tokens must not exceed padded_max"); + PD_CHECK(hidden_i64 % group_size == 0, + "hidden must be a multiple of group_size"); + const int64_t num_groups_i64 = hidden_i64 / group_size; + PD_CHECK(num_groups_i64 % 4 == 0, "num_groups must be a multiple of 4"); + PD_CHECK(buf_x_sf.shape()[0] == padded_max_i64, + "buf_x_sf.shape[0] must equal padded_max"); + PD_CHECK(buf_x_sf.shape()[1] == num_groups_i64 / 4, + "buf_x_sf.shape[1] must equal hidden/group_size/4, but got ", + buf_x_sf.shape()[1], " vs ", num_groups_i64 / 4); + PD_CHECK(hidden_i64 % static_cast(kVecElems) == 0, + "hidden must be a multiple of 8 (16B bf16 loads)"); + const int64_t num_threads_i64 = hidden_i64 / static_cast(kVecElems); + PD_CHECK(num_threads_i64 <= 1024, + "hidden too large for single-block-per-row quant"); + PD_CHECK(num_threads_i64 >= top_k_i64, "top_k must fit into one quant CTA"); + + const uint32_t num_tokens = static_cast(num_tokens_i64); + const uint32_t padded_max = static_cast(padded_max_i64); + const uint32_t hidden = static_cast(hidden_i64); + const uint32_t num_groups = static_cast(num_groups_i64); + const uint32_t top_k = static_cast(top_k_i64); + const uint32_t num_threads = static_cast(num_threads_i64); + const uint32_t num_total_blocks = num_tokens; + + const MegaMoEPreDispatchParams params{ + reinterpret_cast(x.data()), + topk_idx.data(), + topk_weights.data(), + const_cast( + buf_x.data()), + const_cast(buf_x_sf.data()), + const_cast(buf_topk_idx.data()), + const_cast(buf_topk_weights.data()), + num_tokens, + padded_max, + hidden, + num_groups, + top_k, + }; + + if (num_total_blocks > 0) { + auto stream = x.stream(); + switch (group_size) { + case 32: + LaunchMegaMoEPreDispatch<32>(params, num_total_blocks, num_threads, + stream); + break; + case 64: + LaunchMegaMoEPreDispatch<64>(params, num_total_blocks, num_threads, + stream); + break; + case 128: + LaunchMegaMoEPreDispatch<128>(params, num_total_blocks, num_threads, + stream); + break; + default: + PD_THROW("unsupported group_size: ", group_size); + } + } + + // return {buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights}; +} + +std::vector MegaMoePreDispatchInferDtype( + const paddle::DataType& x_dtype, + const paddle::DataType& topk_idx_dtype, + const paddle::DataType& topk_weights_dtype, + const paddle::DataType& buf_x_dtype, + const paddle::DataType& buf_x_sf_dtype, + const paddle::DataType& buf_topk_idx_dtype, + const paddle::DataType& buf_topk_weights_dtype) { + return {buf_x_dtype, buf_x_sf_dtype, buf_topk_idx_dtype, buf_topk_weights_dtype}; +} + +std::vector> MegaMoePreDispatchInferShape( + const std::vector& x_shape, + const std::vector& topk_idx_shape, + const std::vector& topk_weights_shape, + const std::vector& buf_x_shape, + const std::vector& buf_x_sf_shape, + const std::vector& buf_topk_idx_shape, + const std::vector& buf_topk_weights_shape) { + return {buf_x_shape, buf_x_sf_shape, buf_topk_idx_shape, buf_topk_weights_shape}; +} + + +PD_BUILD_STATIC_OP(mega_moe_pre_dispatch) + .Inputs({"x", + "topk_idx", + "topk_weights", + "buf_x", + "buf_x_sf", + "buf_topk_idx", + "buf_topk_weights"}) + .Outputs({"buf_x_out", + "buf_x_sf_out", + "buf_topk_idx_out", + "buf_topk_weights_out"}) + .Attrs({"num_max_tokens_per_rank: int64_t", "group_size: int64_t"}) + .SetInplaceMap({{"buf_x", "buf_x_out"}, + {"buf_x_sf", "buf_x_sf_out"}, + {"buf_topk_idx", "buf_topk_idx_out"}, + {"buf_topk_weights", "buf_topk_weights_out"}}) + .SetKernelFn(PD_KERNEL(MegaMoePreDispatch)) + .SetInferShapeFn(PD_INFER_SHAPE(MegaMoePreDispatchInferShape)) + .SetInferDtypeFn(PD_INFER_DTYPE(MegaMoePreDispatchInferDtype)); diff --git a/custom_ops/setup_ops.py b/custom_ops/setup_ops.py index a7b42f3ae74..9c2a1b03beb 100644 --- a/custom_ops/setup_ops.py +++ b/custom_ops/setup_ops.py @@ -522,7 +522,10 @@ def find_end_files(directory, end_str): # Add SM100 specific sources if any, e.g., for new hardware intrinsics # sources += ["gpu_ops/cutlass_kernels/w8a8/c4x_sm100.cu"] # Example - pass # No SM100 specific sources identified yet beyond what CUTLASS handles + sources += [ + "gpu_ops/mega_moe_pre_dispatch.cu" + ] + if has_generic_fp8: # For SM89 (Ada) or other architectures without dedicated paths diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py b/fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py index 7633ca79b1e..0d2edff3fdb 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py @@ -30,7 +30,7 @@ from fastdeploy.platforms import current_platform from ..quantization.quant_base import QuantMethodBase - +from fastdeploy import envs class MoEMethodBase(QuantMethodBase): """ """ @@ -232,24 +232,29 @@ def apply( """ if layer.ep_size > 1: is_moe_start_layer = layer.layer_idx == layer.fd_config.model_config.moe_layer_start_index - if layer.fd_config.model_config.moe_phase.phase == "prefill": - if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer: - self.ep_prefill_runner.clean_low_latency_buffer() - return self.apply_ep_prefill( - layer, - x, - gate, - topk_ids_hookfunc, - shared_experts, - fc1_latent_proj, - fc2_latent_proj, - ) - else: - if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer: - self.ep_decoder_runner.clean_low_latency_buffer() - return self.apply_ep_decode( + if envs.FD_ENABLE_MAGE_MOE: + return self.apply_mage_moe( layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj ) + else: + if layer.fd_config.model_config.moe_phase.phase == "prefill": + if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer: + self.ep_prefill_runner.clean_low_latency_buffer() + return self.apply_ep_prefill( + layer, + x, + gate, + topk_ids_hookfunc, + shared_experts, + fc1_latent_proj, + fc2_latent_proj, + ) + else: + if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer: + self.ep_decoder_runner.clean_low_latency_buffer() + return self.apply_ep_decode( + layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj + ) else: return self.apply_tp(layer, x, gate, topk_ids_hookfunc, fc1_latent_proj, fc2_latent_proj) diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py index 2ab8ccd5f10..a13de524b95 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py @@ -28,17 +28,24 @@ from fastdeploy.model_executor.layers.quantization.fp8_utils import ( deep_gemm, paddlefleet_ops, + _interleave_weights, + _transpose_sf_for_utccp, ) from fastdeploy.model_executor.layers.utils import get_tensor from fastdeploy.model_executor.ops.gpu import ( count_tokens_per_expert_func, depermute_prefill_combine, prefill_permute_to_masked_gemm, + mega_moe_pre_dispatch, ) from fastdeploy.platforms import current_platform -from fastdeploy.utils import register_custom_python_op +from fastdeploy.utils import register_custom_python_op, singleton from fastdeploy.worker.tbo import let_another_thread_run - +from fastdeploy.model_executor.utils import ( + free_tensor, + process_weight_transpose, + weight_fully_copied, +) from .fused_moe_backend_base import MoEMethodBase from .fused_moe_triton_backend import BlockWiseFP8MoEMethod @@ -208,11 +215,66 @@ def m_grouped_fp8_gemm_nt_contiguous_custom_python_op( return ffn_out +@singleton +class MegaMoEBuffer: + """ + A wrapper class for DeepEP engine. + Manages buffer lifecycle based on role and phase. + """ + + def __init__( + self, + ep_group, + num_experts: int, + num_max_tokens_per_rank: int, + top_k: int, + hidden_size: int, + moe_intermediate_size: int, + ): + self.buffer = deep_gemm.get_symm_buffer_for_mega_moe( + ep_group, + num_experts, + num_max_tokens_per_rank, + top_k, + hidden_size, + moe_intermediate_size, + ) + class DeepGemmFusedMoeMethod(MoEMethodBase): """ DeepGemmFusedMoeMethod is a class that implements the MoEMethodBase interface for DeepGemm backend. """ + def init_ep(self, layer: nn.Layer) -> None: + """ + Initialize EP (Expert Parallel) related modules. + MegaMoE 下取消初始化 EP buffer 以及初始化 mega buffer + """ + if fastdeploy.envs.FD_ENABLE_MAGE_MOE: + if layer.ep_size <= 1: + return + + config = layer.fd_config + splitwise_role = config.scheduler_config.splitwise_role + + if splitwise_role == "mixed" or splitwise_role == "prefill": + self.num_max_tokens_per_rank = config.scheduler_config.max_num_batched_tokens + elif splitwise_role == "decode": + self.num_max_tokens_per_rank = config.model_config.num_max_dispatch_tokens_per_rank + else: + raise ValueError(f"Unsupported splitwise role: {splitwise_role}") + + self.mega_moe_buffer = MegaMoEBuffer( + layer.fd_config.parallel_config.ep_group, + layer.num_experts, + self.num_max_tokens_per_rank, + layer.top_k, + layer.hidden_size, + layer.moe_intermediate_size, + ).buffer + else: + super().init_ep(layer) + def create_weights(self, layer: nn.Layer, **extra_weight_attrs): """ deepgemm create weight process. @@ -221,7 +283,88 @@ def create_weights(self, layer: nn.Layer, **extra_weight_attrs): def process_weights_after_loading(self, layer): """ """ - BlockWiseFP8MoEMethod.process_weights_after_loading(self, layer) + if not fastdeploy.envs.FD_ENABLE_MAGE_MOE: + BlockWiseFP8MoEMethod.process_weights_after_loading(self, layer) + else: + def _process_quantize_mega_moe(weight_idx): + def cast_grouped_weights_to_fp4(bf16_weights: paddle.Tensor): + num_groups, n, k = bf16_weights.shape + w = paddle.empty((num_groups, n, k // 2), dtype=paddle.int8) + w_sf = paddle.empty((num_groups, n, k // 32), dtype=paddle.float32) + for i in range(num_groups): + w[i], w_sf[i] = deep_gemm.per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=32) + w = w.contiguous() + w_sf = w_sf.contiguous() + + w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups) + return w, w_sf + + # weight + weight_name = self.added_weight_attrs[weight_idx] + unquantized_weight_name = weight_name.replace("quant_weight", "weight") + + weight_shape = self.up_gate_proj_weight_shape if weight_type == "gate_up" else self.down_proj_weight_shape + weight_dtype = paddle.bfloat16 + # scale + scale_name = self.added_scale_attrs[weight_idx] + + # 2.create tmp tensor and 3.quantize weight + weight = getattr(layer, unquantized_weight_name).transpose([0, 2, 1]) # [num_experts, 2 * moe_intermediate_size, hidden_size] + if list(weight.shape) != list(weight_shape): + raise ValueError( + f"MegaMoE weight shape mismatch for {unquantized_weight_name}: " + f"got {list(weight.shape)}, expected {list(weight_shape)}" + ) + if weight.dtype != weight_dtype: + weight = weight.astype(weight_dtype) + weight = weight.contiguous() + + weight_quantized = cast_grouped_weights_to_fp4(weight) + + if weight_type == "gate_up": + l1_interleaved = _interleave_weights(weight_quantized) + weight_quantized, scale = (l1_interleaved[0], _transpose_sf_for_utccp(l1_interleaved[1])) + else: + weight_quantized, scale = (weight_quantized[0], _transpose_sf_for_utccp(weight_quantized[1])) + + free_tensor(getattr(layer, weight_name)) + free_tensor(getattr(layer, unquantized_weight_name)) + setattr( + layer, + weight_name, + layer.create_parameter( + shape=weight_quantized.shape, + dtype=paddle.int8, + default_initializer=paddle.nn.initializer.Constant(0), + ), + ) + setattr( + layer, + scale_name, + layer.create_parameter( + shape=scale.shape, + dtype=paddle.int32, + default_initializer=paddle.nn.initializer.Constant(0), + ).as_strided(scale.shape, scale.stride()), + ) + getattr(layer, weight_name).copy_(weight_quantized, False) + getattr(layer, scale_name).copy_(scale, False) + + if self.quant_config.is_checkpoint_bf16: + # dynamic quantize + weight_id_map = {"gate_up": 0, "down": 1} + if weight_fully_copied(layer.up_gate_proj_weight): + weight_type = "gate_up" + else: + weight_type = "down" + if self.model_format == "torch": + # pt model + unquantized_weight_name = self.added_weight_attrs[weight_id_map[weight_type]].replace( + "quant_weight", "weight" + ) + process_weight_transpose(layer, unquantized_weight_name) + + _process_quantize_mega_moe(weight_id_map[weight_type]) def process_loaded_weights(self, layer: nn.Layer, state_dict): """ @@ -904,3 +1047,113 @@ def apply_tp( 1.0, ) return tmp_ffn_out + + + def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor): + if layer.redundant_table_manger is not None: + ( + ep_rank_to_expert_id_list, + expert_id_to_ep_rank_array, + expert_in_rank_num_list, + tokens_per_expert_stats_list, + ) = layer.redundant_table_manger.get_ep_rank_to_expert_id_list_by_layer(layer.layer_idx) + + if layer.topk_method == "noaux_tc": + from .moe import get_moe_scores + + score, topk_weights, topk_idx = get_moe_scores( + gate_out, + layer.n_group, + layer.topk_group, + layer.top_k, + layer.routed_scaling_factor, + layer.gate_correction_bias, + getattr(layer, "renormalize", True), + expert_id_to_ep_rank_array=expert_id_to_ep_rank_array, + expert_in_rank_num_list=expert_in_rank_num_list, + tokens_per_expert_stats_list=tokens_per_expert_stats_list, + redundant_ep_rank_num_plus_one=layer.fd_config.eplb_config.redundant_experts_num + 1, + topk_reduce_func=getattr(layer, "topk_reduce_func", None), + ) + else: + topk_idx, topk_weights = fastdeploy.model_executor.ops.gpu.moe_redundant_topk_select( + gating_logits=gate_out, + expert_id_to_ep_rank_array=expert_id_to_ep_rank_array, + expert_in_rank_num_list=expert_in_rank_num_list, + tokens_per_expert_stats_list=tokens_per_expert_stats_list, + bias=layer.gate_correction_bias, + moe_topk=layer.top_k, + apply_norm_weight=True, + enable_softmax_top_k_fused=False, + redundant_ep_rank_num_plus_one=layer.fd_config.eplb_config.redundant_experts_num + 1, + ) + else: + if layer.topk_method == "noaux_tc": + from fastdeploy.model_executor.layers.moe.moe import get_moe_scores + + score, topk_weights, topk_idx = get_moe_scores( + gate_out, + layer.n_group, + layer.topk_group, + layer.top_k, + layer.routed_scaling_factor, + layer.gate_correction_bias, + getattr(layer, "renormalize", True), + topk_reduce_func=getattr(layer, "topk_reduce_func", None), + ) + else: + topk_idx, topk_weights = fastdeploy.model_executor.ops.gpu.moe_topk_select( + gate_out, + layer.gate_correction_bias, + layer.top_k, + True, + False, + ) + return topk_idx, topk_weights + + def apply_mage_moe(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj): + gate_out = gate(x).cast("float32") + + hidden_size = layer.hidden_size + num_tokens = x.shape[0] + + # 1. Select topk experts and weights. + topk_idx, topk_weights = self.moe_select(layer, gate_out) + + mega_moe_pre_dispatch( + x, + topk_idx, + topk_weights, + self.mega_moe_buffer.x, + self.mega_moe_buffer.x_sf, + self.mega_moe_buffer.topk_idx, + self.mega_moe_buffer.topk_weights, + self.num_max_tokens_per_rank, + 32, # group_size + ) + + buffer_capacity = self.mega_moe_buffer.x.shape[0] + if num_tokens > buffer_capacity: + raise ValueError( + f"MegaMoE buffer capacity exceeded: num_tokens={num_tokens}, capacity={buffer_capacity}" + ) + + l1_weight = getattr(layer, self.added_weight_attrs[0]) + l1_scale = getattr(layer, self.added_scale_attrs[0]) + l2_weight = getattr(layer, self.added_weight_attrs[1]) + l2_scale = getattr(layer, self.added_scale_attrs[1]) + y = paddle.empty((max(num_tokens, 1), hidden_size), dtype=paddle.bfloat16) + + swiglu_limit = getattr(layer.fd_config.model_config, "swiglu_limit", None) + deep_gemm.fp8_fp4_mega_moe( + y, + (l1_weight, l1_scale), + (l2_weight, l2_scale), + self.mega_moe_buffer, + recipe=(1, 1, 32), + activation="swiglu", + activation_clamp=swiglu_limit, + fast_math=True, + ) + + return y diff --git a/fastdeploy/model_executor/layers/quantization/fp8_utils.py b/fastdeploy/model_executor/layers/quantization/fp8_utils.py index 89b9467ecc6..311c12b324c 100644 --- a/fastdeploy/model_executor/layers/quantization/fp8_utils.py +++ b/fastdeploy/model_executor/layers/quantization/fp8_utils.py @@ -236,3 +236,29 @@ def fused_stack_transpose_quant(expert_weight_list, use_ue8m0=False): raise RuntimeError("'fuse_stack_transpose_fp8_quant' is not available in the current paddlefleet_ops.") return w, scale + + +def _interleave_weights(l1_weights): + # [gate: 0..7, up: 0..7, gate: 8..15, up: 8..15, ...] instead of [gate | up] + def interleave(t, gran: int = 8) -> paddle.Tensor: + g, n, *rest = t.shape + half = n // 2 + gate = t[:, :half].reshape(g, half // gran, gran, *rest) + up = t[:, half:].reshape(g, half // gran, gran, *rest) + return paddle.stack([gate, up], dim=2).reshape(g, n, *rest).contiguous() + + return interleave(l1_weights[0]), interleave(l1_weights[1]) + + +def _transpose_sf_for_utccp(sf: paddle.Tensor) -> paddle.Tensor: + num_groups, mn, packed_sf_k = sf.shape + assert sf.dtype == paddle.int and mn % 128 == 0 + # sf is MN-major: strides [mn*packed_sf_k, 1, mn] + # We need to do the 4x32 transpose in data while preserving MN-major strides + sf_c = sf.contiguous() # make C-contiguous for reshape/transpose + result_c = (sf_c.reshape(num_groups, -1, 4, 32, packed_sf_k) + .transpose(2, 3) + .reshape(num_groups, mn, packed_sf_k) + .contiguous()) + # Convert back to MN-major layout: transpose last two dims, make contiguous, transpose back + return result_c.transpose(1, 2).contiguous().transpose(1, 2) diff --git a/tests/operators/test_mega_moe_pre_dispatch.py b/tests/operators/test_mega_moe_pre_dispatch.py new file mode 100644 index 00000000000..f5984309665 --- /dev/null +++ b/tests/operators/test_mega_moe_pre_dispatch.py @@ -0,0 +1,177 @@ +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import numpy as np +import paddle +import paddle.distributed as dist +from paddle.distributed import fleet + +from ernie5_serving.mm_custom_ops import mega_moe_pre_dispatch +from fastdeploy.model_executor.layers.moe.fused_moe_deepgemm_backend import MegaMoEBuffer + + +def ceil_div(x: int, y: int) -> int: + return (x + y - 1) // y + + +def align(x: int, y: int) -> int: + return ceil_div(x, y) * y + + +def ceil_to_ue8m0(x: paddle.Tensor): + bits = x.abs().astype("float32").view(paddle.int32) + mask_ff = paddle.to_tensor(0xFF, dtype=paddle.int32) + mask_mantissa = paddle.to_tensor(0x7FFFFF, dtype=paddle.int32) + exp = ((bits >> 23) & mask_ff) + ((bits & mask_mantissa) != 0).astype("int32") + return (exp.clip(1, 254) << 23).view(paddle.float32) + + +def pack_ue8m0_to_int(x: paddle.Tensor): + assert x.dtype == paddle.float32 and x.shape[-1] % 4 == 0 + x_bits = x.view(paddle.int32) + mantissa_mask = paddle.to_tensor((1 << 23) - 1, dtype=paddle.int32) + assert bool(((x_bits & mantissa_mask) == 0).all()) + return (x_bits >> 23).astype(paddle.uint8).view(paddle.int32) + + +def per_token_cast_to_fp8( + x: paddle.Tensor, + use_ue8m0: bool, + gran_k: int = 128, + use_packed_ue8m0: bool = False, +): + assert len(x.shape) == 2 + m, n = x.shape + padded_n = align(n, gran_k) + x_padded = paddle.zeros((m, padded_n), dtype=x.dtype) + x_padded[:, :n] = x + x_view = x_padded.reshape([m, padded_n // gran_k, gran_k]) + x_amax = x_view.abs().astype("float32").amax(axis=2).reshape([m, padded_n // gran_k]).clip(min=1e-4) + sf = x_amax / 448.0 + sf = ceil_to_ue8m0(sf) if use_ue8m0 else sf + x_fp8 = (x_view * (1.0 / sf.unsqueeze(2))).astype(paddle.float8_e4m3fn).reshape([m, padded_n])[:, :n] + return x_fp8.contiguous(), pack_ue8m0_to_int(sf) if use_packed_ue8m0 else sf + + +class TestMegaMoEPreDispatch(unittest.TestCase): + @classmethod + def setUpClass(cls): + paddle.seed(2025) + strategy = fleet.DistributedStrategy() + cls.expert_parallel_size = 8 + strategy.hybrid_configs = { + "dp_degree": 1, + "mp_degree": cls.expert_parallel_size, + "pp_degree": 1, + "sharding_degree": 1, + } + fleet.init(is_collective=True, strategy=strategy) + cls.ep_group = dist.new_group(range(cls.expert_parallel_size)) + + def setUp(self): + self.num_experts = 160 + self.num_max_tokens_per_rank = 8192 + self.top_k = 6 + self.hidden_size = 7168 + self.moe_intermediate_size = 3584 + self.group_size = 32 + self.num_tokens = 128 + + self.x = paddle.randn([self.num_tokens, self.hidden_size], dtype=paddle.bfloat16) + scores = paddle.randn((self.num_tokens, self.num_experts), dtype=paddle.float32) + self.topk_weights, self.topk_idx = paddle.topk(scores, self.top_k, axis=-1, largest=True, sorted=False) + self.topk_idx = self.topk_idx.astype("int32") + self.topk_weights = self.topk_weights.astype("float32") + + def _new_buffer(self): + return MegaMoEBuffer( + self.ep_group, + self.num_experts, + self.num_max_tokens_per_rank, + self.top_k, + self.hidden_size, + self.moe_intermediate_size, + ).buffer + + def mega_moe_pre_dispatch_ref(self, x: paddle.Tensor, topk_idx: paddle.Tensor, topk_weights: paddle.Tensor): + num_tokens = x.shape[0] + x_fp8, x_scale_tensor = per_token_cast_to_fp8( + x, use_ue8m0=True, gran_k=self.group_size, use_packed_ue8m0=True + ) + return ( + x_fp8, + x_scale_tensor, + topk_idx.astype("int64"), + topk_weights.astype("float32"), + ) + + def test_mega_moe_pre_dispatch(self): + buffer = self._new_buffer() + buffer.topk_idx[self.num_tokens :] = -2 + buffer.topk_weights[self.num_tokens :] = 3.0 + + mega_moe_pre_dispatch( + self.x, + self.topk_idx, + self.topk_weights, + buffer.x, + buffer.x_sf, + buffer.topk_idx, + buffer.topk_weights, + self.num_max_tokens_per_rank, + self.group_size, + ) + paddle.device.synchronize() + + x_ref, x_sf_ref, topk_idx_ref, topk_weights_ref = self.mega_moe_pre_dispatch_ref( + self.x, self.topk_idx, self.topk_weights + ) + + np.testing.assert_allclose( + buffer.x[: self.num_tokens].astype("float32").numpy(), + x_ref.astype("float32").numpy(), + rtol=0, + atol=0, + ) + np.testing.assert_array_equal( + buffer.x_sf[: self.num_tokens].numpy(), + x_sf_ref.numpy(), + ) + np.testing.assert_array_equal( + buffer.topk_idx[: self.num_tokens].numpy(), + topk_idx_ref.numpy(), + ) + np.testing.assert_allclose( + buffer.topk_weights[: self.num_tokens].numpy(), + topk_weights_ref.numpy(), + rtol=0, + atol=0, + ) + padded_max = buffer.x.shape[0] + np.testing.assert_array_equal( + buffer.topk_idx[self.num_tokens :].numpy(), + np.full((padded_max - self.num_tokens, self.top_k), -2, dtype=np.int64), + ) + np.testing.assert_allclose( + buffer.topk_weights[self.num_tokens :].numpy(), + np.full((padded_max - self.num_tokens, self.top_k), 3.0, dtype=np.float32), + rtol=0, + atol=0, + ) + + +if __name__ == "__main__": + unittest.main() From 0baf6037637c2fee47c30c75750b129b78669366 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Mon, 8 Jun 2026 16:29:21 +0800 Subject: [PATCH 02/17] update code --- fastdeploy/config.py | 3 +- fastdeploy/engine/args_utils.py | 11 + fastdeploy/engine/engine.py | 1 + .../layers/backends/xpu/moe/fused_moe.py | 5 +- fastdeploy/model_executor/layers/moe/ep.py | 24 +- .../layers/moe/fused_moe_backend_base.py | 37 +- .../layers/moe/fused_moe_blackwell_backend.py | 6 +- .../layers/moe/fused_moe_cutlass_backend.py | 5 +- .../layers/moe/fused_moe_deepgemm_backend.py | 655 ++++++++++++------ .../layers/quantization/__init__.py | 35 + .../layers/quantization/fp8_utils.py | 2 - .../layers/quantization/nvfp4.py | 5 +- .../layers/quantization/wfp4afp8.py | 66 ++ fastdeploy/worker/worker_process.py | 7 + tests/model_executor/test_ep.py | 7 +- 15 files changed, 618 insertions(+), 251 deletions(-) create mode 100644 fastdeploy/model_executor/layers/quantization/wfp4afp8.py diff --git a/fastdeploy/config.py b/fastdeploy/config.py index dfb1e3c530d..6d9e3d58038 100644 --- a/fastdeploy/config.py +++ b/fastdeploy/config.py @@ -651,7 +651,8 @@ def __init__( self.enable_expert_parallel = False self.enable_chunked_moe = False self.chunked_moe_size = 256 - + self.enable_mega_moe = False + self.local_data_parallel_id = 0 # Engine worker queue port self.engine_worker_queue_port: Union[int, str, list] = None diff --git a/fastdeploy/engine/args_utils.py b/fastdeploy/engine/args_utils.py index 9c3f3c0f10e..8d126998d46 100644 --- a/fastdeploy/engine/args_utils.py +++ b/fastdeploy/engine/args_utils.py @@ -369,6 +369,11 @@ class EngineArgs: Whether use chunked moe. """ + enable_mega_moe: bool = False + """ + Whether use MegaMoE wfp4afp8 for MoE and block_wise_fp8 for dense Linear. + """ + chunked_moe_size: int = 256 """ Chunk size of moe input. @@ -1176,6 +1181,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: default=EngineArgs.enable_chunked_moe, help="Use chunked moe.", ) + parallel_group.add_argument( + "--enable-mega-moe", + action="store_true", + default=EngineArgs.enable_mega_moe, + help="Use MegaMoE wfp4afp8 for MoE and block_wise_fp8 for dense Linear.", + ) parallel_group.add_argument( "--chunked-moe-size", type=int, diff --git a/fastdeploy/engine/engine.py b/fastdeploy/engine/engine.py index b0e7d82018b..8614d29999b 100644 --- a/fastdeploy/engine/engine.py +++ b/fastdeploy/engine/engine.py @@ -685,6 +685,7 @@ def _start_worker_service(self): worker_store_true_flag = { "enable_expert_parallel": self.cfg.parallel_config.enable_expert_parallel, "enable_chunked_moe": self.cfg.parallel_config.enable_chunked_moe, + "enable_mega_moe": self.cfg.parallel_config.enable_mega_moe, "enable_prefix_caching": self.cfg.cache_config.enable_prefix_caching, "enable_chunked_prefill": self.cfg.cache_config.enable_chunked_prefill, "do_profile": self.do_profile, diff --git a/fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py b/fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py index 6f313170a74..d58a85a64fa 100644 --- a/fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py +++ b/fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py @@ -39,6 +39,7 @@ free_tensor, set_weight_attrs, ) +from fastdeploy.model_executor.layers.moe.ep import EPRunner from .utils import get_moe_scores @@ -423,7 +424,7 @@ def apply_ep_prefill( """ gate_out = gate(x.cast("float32")) # 1. Select topk experts and weights - topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) # 2. Dynamic compute blockwise quantization scales if "a_tokenwise_int8" in self.xpu_moe_quant_type: @@ -518,7 +519,7 @@ def apply_ep_decode( gate_out = gate(x.cast("float32")) # 1. Select topk experts and weights - topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) # 2. EP Dispatch if "a_tokenwise_int8" in self.xpu_moe_quant_type: diff --git a/fastdeploy/model_executor/layers/moe/ep.py b/fastdeploy/model_executor/layers/moe/ep.py index 967c2a2fd02..ca215098761 100644 --- a/fastdeploy/model_executor/layers/moe/ep.py +++ b/fastdeploy/model_executor/layers/moe/ep.py @@ -490,7 +490,8 @@ def __init__( top_k=self.top_k, ) - def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor): + @staticmethod + def moe_select(layer: nn.Layer, gate_out: paddle.Tensor): if layer.redundant_table_manger is not None: ( ep_rank_to_expert_id_list, @@ -523,7 +524,7 @@ def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor): expert_in_rank_num_list=expert_in_rank_num_list, tokens_per_expert_stats_list=tokens_per_expert_stats_list, bias=layer.gate_correction_bias, - moe_topk=self.top_k, + moe_topk=layer.top_k, apply_norm_weight=True, enable_softmax_top_k_fused=False, redundant_ep_rank_num_plus_one=layer.fd_config.eplb_config.redundant_experts_num + 1, @@ -550,7 +551,7 @@ def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor): topk_idx, topk_weights = fastdeploy.model_executor.ops.gpu.moe_topk_select( gate_out, layer.gate_correction_bias, - self.top_k, + layer.top_k, True, False, ) @@ -788,3 +789,20 @@ def combine(self, ffn_out, topk_idx, topk_weights, handle, **kwargs): combine_hook() return combined_hidden_states + + +class FakeEPRunner: + """ """ + def __init__(self, *args, **kwargs): + pass + + def dispatch(self, *args, **kwargs): + """ """ + pass + + def combine(self, *args, **kwargs): + """ """ + pass + + def clean_low_latency_buffer(self): + pass \ No newline at end of file diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py b/fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py index 0d2edff3fdb..ac105043397 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py @@ -232,29 +232,24 @@ def apply( """ if layer.ep_size > 1: is_moe_start_layer = layer.layer_idx == layer.fd_config.model_config.moe_layer_start_index - if envs.FD_ENABLE_MAGE_MOE: - return self.apply_mage_moe( - layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj + if layer.fd_config.model_config.moe_phase.phase == "prefill": + if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer: + self.ep_prefill_runner.clean_low_latency_buffer() + return self.apply_ep_prefill( + layer, + x, + gate, + topk_ids_hookfunc, + shared_experts, + fc1_latent_proj, + fc2_latent_proj, ) else: - if layer.fd_config.model_config.moe_phase.phase == "prefill": - if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer: - self.ep_prefill_runner.clean_low_latency_buffer() - return self.apply_ep_prefill( - layer, - x, - gate, - topk_ids_hookfunc, - shared_experts, - fc1_latent_proj, - fc2_latent_proj, - ) - else: - if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer: - self.ep_decoder_runner.clean_low_latency_buffer() - return self.apply_ep_decode( - layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj - ) + if layer.fd_config.scheduler_config.splitwise_role == "mixed" and is_moe_start_layer: + self.ep_decoder_runner.clean_low_latency_buffer() + return self.apply_ep_decode( + layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj + ) else: return self.apply_tp(layer, x, gate, topk_ids_hookfunc, fc1_latent_proj, fc2_latent_proj) diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py index 274deda8b69..681ccb38c3b 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py @@ -23,7 +23,7 @@ from paddleformers.utils.log import logger import fastdeploy -from fastdeploy.model_executor.layers.moe.ep import deep_ep +from fastdeploy.model_executor.layers.moe.ep import deep_ep, EPRunner from fastdeploy.model_executor.layers.quantization.fp8_utils import ( deep_gemm, paddlefleet_ops, @@ -642,7 +642,7 @@ def apply_ep_prefill( getattr(layer, "renormalize", True), ) else: - topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) if topk_ids_hookfunc is not None: topk_ids_hookfunc(topk_ids=topk_idx) @@ -964,7 +964,7 @@ def apply_ep_decode( gate_out = gate(x) gate_out = gate_out.cast("float32") # 1. Select topk experts and weights - topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) if topk_ids_hookfunc is not None: topk_ids_hookfunc(topk_ids=topk_idx) diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py index 5334114b70e..8ec0305c4cf 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py @@ -52,6 +52,7 @@ set_weight_attrs, weight_fully_copied, ) +from fastdeploy.model_executor.layers.moe.ep import EPRunner def m_grouped_bf16_gemm_nn_contiguous(x, y, expert_idx_per_token): @@ -141,7 +142,7 @@ def apply_ep_prefill( if fc1_latent_proj is not None: x = fc1_latent_proj(x) # 1. Select topk experts and weights - topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) if layer.routed_scaling_factor_learnable: safe_topk_indices = paddle.clip(topk_idx, min=0) @@ -304,7 +305,7 @@ def apply_ep_decode( estimate_total_token_nums = gate_out.shape[0] * layer.top_k # 1. Select topk experts and weights - topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) if layer.routed_scaling_factor_learnable: safe_topk_indices = paddle.clip(topk_idx, min=0) diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py index a13de524b95..ee9acf76af5 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py @@ -22,9 +22,8 @@ import paddle.nn.functional as F from paddle import nn from paddleformers.utils.log import logger - import fastdeploy -from fastdeploy.model_executor.layers.moe.ep import deep_ep +from fastdeploy.model_executor.layers.moe.ep import deep_ep, EPRunner, FakeEPRunner from fastdeploy.model_executor.layers.quantization.fp8_utils import ( deep_gemm, paddlefleet_ops, @@ -39,12 +38,14 @@ mega_moe_pre_dispatch, ) from fastdeploy.platforms import current_platform -from fastdeploy.utils import register_custom_python_op, singleton +from fastdeploy.utils import ceil_div, register_custom_python_op, singleton from fastdeploy.worker.tbo import let_another_thread_run from fastdeploy.model_executor.utils import ( + TensorTracker, + set_weight_attrs, free_tensor, - process_weight_transpose, weight_fully_copied, + get_sm_version, ) from .fused_moe_backend_base import MoEMethodBase from .fused_moe_triton_backend import BlockWiseFP8MoEMethod @@ -215,66 +216,11 @@ def m_grouped_fp8_gemm_nt_contiguous_custom_python_op( return ffn_out -@singleton -class MegaMoEBuffer: - """ - A wrapper class for DeepEP engine. - Manages buffer lifecycle based on role and phase. - """ - - def __init__( - self, - ep_group, - num_experts: int, - num_max_tokens_per_rank: int, - top_k: int, - hidden_size: int, - moe_intermediate_size: int, - ): - self.buffer = deep_gemm.get_symm_buffer_for_mega_moe( - ep_group, - num_experts, - num_max_tokens_per_rank, - top_k, - hidden_size, - moe_intermediate_size, - ) - class DeepGemmFusedMoeMethod(MoEMethodBase): """ DeepGemmFusedMoeMethod is a class that implements the MoEMethodBase interface for DeepGemm backend. """ - def init_ep(self, layer: nn.Layer) -> None: - """ - Initialize EP (Expert Parallel) related modules. - MegaMoE 下取消初始化 EP buffer 以及初始化 mega buffer - """ - if fastdeploy.envs.FD_ENABLE_MAGE_MOE: - if layer.ep_size <= 1: - return - - config = layer.fd_config - splitwise_role = config.scheduler_config.splitwise_role - - if splitwise_role == "mixed" or splitwise_role == "prefill": - self.num_max_tokens_per_rank = config.scheduler_config.max_num_batched_tokens - elif splitwise_role == "decode": - self.num_max_tokens_per_rank = config.model_config.num_max_dispatch_tokens_per_rank - else: - raise ValueError(f"Unsupported splitwise role: {splitwise_role}") - - self.mega_moe_buffer = MegaMoEBuffer( - layer.fd_config.parallel_config.ep_group, - layer.num_experts, - self.num_max_tokens_per_rank, - layer.top_k, - layer.hidden_size, - layer.moe_intermediate_size, - ).buffer - else: - super().init_ep(layer) - def create_weights(self, layer: nn.Layer, **extra_weight_attrs): """ deepgemm create weight process. @@ -283,88 +229,7 @@ def create_weights(self, layer: nn.Layer, **extra_weight_attrs): def process_weights_after_loading(self, layer): """ """ - if not fastdeploy.envs.FD_ENABLE_MAGE_MOE: - BlockWiseFP8MoEMethod.process_weights_after_loading(self, layer) - else: - def _process_quantize_mega_moe(weight_idx): - def cast_grouped_weights_to_fp4(bf16_weights: paddle.Tensor): - num_groups, n, k = bf16_weights.shape - w = paddle.empty((num_groups, n, k // 2), dtype=paddle.int8) - w_sf = paddle.empty((num_groups, n, k // 32), dtype=paddle.float32) - for i in range(num_groups): - w[i], w_sf[i] = deep_gemm.per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=32) - w = w.contiguous() - w_sf = w_sf.contiguous() - - w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups) - return w, w_sf - - # weight - weight_name = self.added_weight_attrs[weight_idx] - unquantized_weight_name = weight_name.replace("quant_weight", "weight") - - weight_shape = self.up_gate_proj_weight_shape if weight_type == "gate_up" else self.down_proj_weight_shape - weight_dtype = paddle.bfloat16 - # scale - scale_name = self.added_scale_attrs[weight_idx] - - # 2.create tmp tensor and 3.quantize weight - weight = getattr(layer, unquantized_weight_name).transpose([0, 2, 1]) # [num_experts, 2 * moe_intermediate_size, hidden_size] - if list(weight.shape) != list(weight_shape): - raise ValueError( - f"MegaMoE weight shape mismatch for {unquantized_weight_name}: " - f"got {list(weight.shape)}, expected {list(weight_shape)}" - ) - if weight.dtype != weight_dtype: - weight = weight.astype(weight_dtype) - weight = weight.contiguous() - - weight_quantized = cast_grouped_weights_to_fp4(weight) - - if weight_type == "gate_up": - l1_interleaved = _interleave_weights(weight_quantized) - weight_quantized, scale = (l1_interleaved[0], _transpose_sf_for_utccp(l1_interleaved[1])) - else: - weight_quantized, scale = (weight_quantized[0], _transpose_sf_for_utccp(weight_quantized[1])) - - free_tensor(getattr(layer, weight_name)) - free_tensor(getattr(layer, unquantized_weight_name)) - setattr( - layer, - weight_name, - layer.create_parameter( - shape=weight_quantized.shape, - dtype=paddle.int8, - default_initializer=paddle.nn.initializer.Constant(0), - ), - ) - setattr( - layer, - scale_name, - layer.create_parameter( - shape=scale.shape, - dtype=paddle.int32, - default_initializer=paddle.nn.initializer.Constant(0), - ).as_strided(scale.shape, scale.stride()), - ) - getattr(layer, weight_name).copy_(weight_quantized, False) - getattr(layer, scale_name).copy_(scale, False) - - if self.quant_config.is_checkpoint_bf16: - # dynamic quantize - weight_id_map = {"gate_up": 0, "down": 1} - if weight_fully_copied(layer.up_gate_proj_weight): - weight_type = "gate_up" - else: - weight_type = "down" - if self.model_format == "torch": - # pt model - unquantized_weight_name = self.added_weight_attrs[weight_id_map[weight_type]].replace( - "quant_weight", "weight" - ) - process_weight_transpose(layer, unquantized_weight_name) - - _process_quantize_mega_moe(weight_id_map[weight_type]) + BlockWiseFP8MoEMethod.process_weights_after_loading(self, layer) def process_loaded_weights(self, layer: nn.Layer, state_dict): """ @@ -488,7 +353,7 @@ def apply_ep_prefill( hidden_size = layer.hidden_size # 1. Select topk experts and weights - topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) if layer.routed_scaling_factor_learnable: safe_topk_indices = paddle.clip(topk_idx, min=0) @@ -818,7 +683,7 @@ def apply_ep_decode( gate_out = gate(x) gate_out = gate_out.cast("float32") # 1. Select topk experts and weights - topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) if layer.routed_scaling_factor_learnable: safe_topk_indices = paddle.clip(topk_idx, min=0) @@ -1049,77 +914,448 @@ def apply_tp( return tmp_ffn_out - def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor): - if layer.redundant_table_manger is not None: - ( - ep_rank_to_expert_id_list, - expert_id_to_ep_rank_array, - expert_in_rank_num_list, - tokens_per_expert_stats_list, - ) = layer.redundant_table_manger.get_ep_rank_to_expert_id_list_by_layer(layer.layer_idx) - - if layer.topk_method == "noaux_tc": - from .moe import get_moe_scores - - score, topk_weights, topk_idx = get_moe_scores( - gate_out, - layer.n_group, - layer.topk_group, - layer.top_k, - layer.routed_scaling_factor, - layer.gate_correction_bias, - getattr(layer, "renormalize", True), - expert_id_to_ep_rank_array=expert_id_to_ep_rank_array, - expert_in_rank_num_list=expert_in_rank_num_list, - tokens_per_expert_stats_list=tokens_per_expert_stats_list, - redundant_ep_rank_num_plus_one=layer.fd_config.eplb_config.redundant_experts_num + 1, - topk_reduce_func=getattr(layer, "topk_reduce_func", None), - ) +@singleton +class MegaMoEBuffer: + """ + A wrapper class for DeepEP engine. + Manages buffer lifecycle based on role and phase. + """ + + def __init__( + self, + ep_group, + num_experts: int, + num_max_tokens_per_rank: int, + top_k: int, + hidden_size: int, + moe_intermediate_size: int, + ): + self.buffer = deep_gemm.get_symm_buffer_for_mega_moe( + ep_group, + num_experts, + num_max_tokens_per_rank, + top_k, + hidden_size, + moe_intermediate_size, + ) + + +class DeepGemmMegaMoEMethod(DeepGemmFusedMoeMethod): + def __init__(self, quant_config): + if not get_sm_version() >= 100: + raise ValueError("MegaMoE now only support sm100+ devices.") + super().__init__(quant_config) + self.added_scale_attrs = ["up_gate_proj_weight_scale", "down_proj_weight_scale"] + self.quant_config.deepgemm_scale_ue8m0 = True + self.gran_k = 32 + + def create_weights(self, layer: nn.Layer, **extra_weight_attrs): + """ + Triton MoE create weight process. + """ + logger.info("mega create_weights") + self.model_format = extra_weight_attrs.get("model_format") + self.up_gate_proj_quant_weight_shape = [ + layer.num_local_experts, + layer.moe_intermediate_size * 2, + layer.hidden_size, + ] + self.down_proj_quant_weight_shape = [ + layer.num_local_experts, + layer.hidden_size, + layer.moe_intermediate_size, + ] + self.up_gate_proj_pretranspose_weight_shape = [ + layer.num_local_experts, + layer.hidden_size, + layer.moe_intermediate_size * 2, + ] + self.down_proj_pretranspose_weight_shape = [ + layer.num_local_experts, + layer.moe_intermediate_size, + layer.hidden_size, + ] + if self.model_format != "torch": + self.up_gate_proj_bf16_weight_shape = self.up_gate_proj_pretranspose_weight_shape + self.down_proj_bf16_weight_shape = self.down_proj_pretranspose_weight_shape + else: + self.up_gate_proj_bf16_weight_shape = self.up_gate_proj_quant_weight_shape + self.down_proj_bf16_weight_shape = self.down_proj_quant_weight_shape + self.up_gate_proj_packed_weight_shape = [ + layer.num_local_experts, + layer.moe_intermediate_size * 2, + layer.hidden_size // 2, # 4-bit packing + ] + self.down_proj_packed_weight_shape = [ + layer.num_local_experts, + layer.hidden_size, + layer.moe_intermediate_size // 2, # 4-bit packing + ] + up_num_scales = ceil_div(layer.hidden_size, self.gran_k) + down_num_scales = ceil_div(layer.moe_intermediate_size, self.gran_k) + self.up_gate_proj_scale_shape = [ + layer.num_local_experts, + layer.moe_intermediate_size * 2, + (up_num_scales + 3) // 4, + ] + self.down_proj_scale_shape = [ + layer.num_local_experts, + layer.hidden_size, + (down_num_scales + 3) // 4, + ] + self.up_gate_proj_weight_shape = self.up_gate_proj_quant_weight_shape + self.down_proj_weight_shape = self.down_proj_quant_weight_shape + + logger.info( + "MegaMoE create_weights: " + f"is_checkpoint_bf16={self.quant_config.is_checkpoint_bf16}, " + f"load_choices={layer.fd_config.load_config.load_choices}, " + f"model_format={self.model_format}, " + f"up_gate_bf16_shape={self.up_gate_proj_bf16_weight_shape}, " + f"down_bf16_shape={self.down_proj_bf16_weight_shape}, " + f"up_gate_quant_shape={self.up_gate_proj_quant_weight_shape}, " + f"down_quant_shape={self.down_proj_quant_weight_shape}, " + f"up_gate_packed_shape={self.up_gate_proj_packed_weight_shape}, " + f"down_packed_shape={self.down_proj_packed_weight_shape}, " + f"up_gate_scale_shape={self.up_gate_proj_scale_shape}, " + f"down_scale_shape={self.down_proj_scale_shape}" + ) + + if self.quant_config.is_checkpoint_bf16 and layer.fd_config.load_config.load_choices == "default_v1": + if self.model_format != "torch": + up_gate_proj_attrs = { + **extra_weight_attrs, + "tensor_track": TensorTracker(shape=self.up_gate_proj_bf16_weight_shape, output_dim=True), + "SHARD_ID_TO_SHARDED_DIM": {"gate": 1, "down": 0, "up": 1}, + } + down_proj_attrs = { + **extra_weight_attrs, + "tensor_track": TensorTracker(shape=self.down_proj_bf16_weight_shape, output_dim=False), + "SHARD_ID_TO_SHARDED_DIM": {"gate": 1, "down": 0, "up": 1}, + } else: - topk_idx, topk_weights = fastdeploy.model_executor.ops.gpu.moe_redundant_topk_select( - gating_logits=gate_out, - expert_id_to_ep_rank_array=expert_id_to_ep_rank_array, - expert_in_rank_num_list=expert_in_rank_num_list, - tokens_per_expert_stats_list=tokens_per_expert_stats_list, - bias=layer.gate_correction_bias, - moe_topk=layer.top_k, - apply_norm_weight=True, - enable_softmax_top_k_fused=False, - redundant_ep_rank_num_plus_one=layer.fd_config.eplb_config.redundant_experts_num + 1, - ) + up_gate_proj_attrs = { + **extra_weight_attrs, + "tensor_track": TensorTracker(shape=self.up_gate_proj_bf16_weight_shape, output_dim=False), + "SHARD_ID_TO_SHARDED_DIM": {"gate": 0, "down": 1, "up": 0}, + } + down_proj_attrs = { + **extra_weight_attrs, + "tensor_track": TensorTracker(shape=self.down_proj_bf16_weight_shape, output_dim=True), + "SHARD_ID_TO_SHARDED_DIM": {"gate": 0, "down": 1, "up": 0}, + } + layer.up_gate_proj_weight = layer.create_parameter( + shape=self.up_gate_proj_bf16_weight_shape, + dtype=layer.weight_dtype, + default_initializer=paddle.nn.initializer.Constant(0), + ) + + layer.down_proj_weight = layer.create_parameter( + shape=self.down_proj_bf16_weight_shape, + dtype=layer.weight_dtype, + default_initializer=paddle.nn.initializer.Constant(0), + ) + + set_weight_attrs( + layer.up_gate_proj_weight, + up_gate_proj_attrs, + ) + set_weight_attrs( + layer.down_proj_weight, + down_proj_attrs, + ) + else: + # offline quant + self.up_gate_proj_weight_shape = self.up_gate_proj_packed_weight_shape + self.down_proj_weight_shape = self.down_proj_packed_weight_shape + up_gate_proj_attrs = {} + down_proj_attrs = {} + + self.weight_dtype = paddle.int8 + up_gate_proj_weight_name = self.added_weight_attrs[0] + down_proj_weight_name = self.added_weight_attrs[1] + up_gate_proj_scale_name = self.added_scale_attrs[0] + down_proj_scale_name = self.added_scale_attrs[1] + + setattr( + layer, + up_gate_proj_weight_name, + layer.create_parameter( + shape=self.up_gate_proj_packed_weight_shape, + dtype=self.weight_dtype, + default_initializer=paddle.nn.initializer.Constant(0), + ), + ) + setattr( + layer, + down_proj_weight_name, + layer.create_parameter( + shape=self.down_proj_packed_weight_shape, + dtype=self.weight_dtype, + default_initializer=paddle.nn.initializer.Constant(0), + ), + ) + # weight_scale + setattr( + layer, + up_gate_proj_scale_name, + layer.create_parameter( + shape=self.up_gate_proj_scale_shape, + dtype="int32", + default_initializer=paddle.nn.initializer.Constant(0), + ), + ) + setattr( + layer, + down_proj_scale_name, + layer.create_parameter( + shape=self.down_proj_scale_shape, + dtype="int32", + default_initializer=paddle.nn.initializer.Constant(0), + ), + ) + + set_weight_attrs( + getattr(layer, up_gate_proj_weight_name), + up_gate_proj_attrs, + ) + set_weight_attrs( + getattr(layer, up_gate_proj_scale_name), + up_gate_proj_attrs, + ) + + set_weight_attrs( + getattr(layer, down_proj_weight_name), + down_proj_attrs, + ) + set_weight_attrs( + getattr(layer, down_proj_scale_name), + down_proj_attrs, + ) + + def init_ep(self, layer: nn.Layer) -> None: + logger.info("USE MegaMoE backend") + if layer.ep_size <= 1: + return + + config = layer.fd_config + splitwise_role = config.scheduler_config.splitwise_role + + if splitwise_role == "mixed" or splitwise_role == "prefill": + self.num_max_tokens_per_rank = config.scheduler_config.max_num_batched_tokens + elif splitwise_role == "decode": + num_spec_tokens = config.speculative_config.num_speculative_tokens + self.num_max_tokens_per_rank = config.scheduler_config.max_num_seqs * (num_spec_tokens + 1) else: - if layer.topk_method == "noaux_tc": - from fastdeploy.model_executor.layers.moe.moe import get_moe_scores - - score, topk_weights, topk_idx = get_moe_scores( - gate_out, - layer.n_group, - layer.topk_group, - layer.top_k, - layer.routed_scaling_factor, - layer.gate_correction_bias, - getattr(layer, "renormalize", True), - topk_reduce_func=getattr(layer, "topk_reduce_func", None), + raise ValueError(f"Unsupported splitwise role: {splitwise_role}") + + self.mega_moe_buffer = MegaMoEBuffer( + layer.fd_config.parallel_config.ep_group, + layer.num_experts, + self.num_max_tokens_per_rank, + layer.top_k, + layer.hidden_size, + layer.moe_intermediate_size, + ).buffer + self.num_max_tokens_per_rank = self.mega_moe_buffer.num_max_tokens_per_rank + self.cumulative_local_expert_recv_stats = paddle.zeros( + (layer.num_local_experts,), dtype=paddle.int32 + ) + + self.ep_prefill_runner = FakeEPRunner() + self.ep_decoder_runner = FakeEPRunner() + + def process_weights_after_loading(self, layer): + def cast_grouped_weights_to_fp4(bf16_weights: paddle.Tensor): + num_groups, n, k = bf16_weights.shape + w = paddle.empty((num_groups, n, k // 2), dtype=paddle.int8) + w_sf = paddle.empty((num_groups, n, k // self.gran_k), dtype=paddle.float32) + for i in range(num_groups): + w[i], w_sf[i] = deep_gemm.per_token_cast_to_fp4( + bf16_weights[i], use_ue8m0=True, gran_k=self.gran_k ) - else: - topk_idx, topk_weights = fastdeploy.model_executor.ops.gpu.moe_topk_select( - gate_out, - layer.gate_correction_bias, - layer.top_k, - True, - False, + w = w.contiguous() + w_sf = w_sf.contiguous() + w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, self.gran_k), num_groups) + return w, w_sf + + def _process_quantize_mega_moe(weight_type): + weight_idx = 0 if weight_type == "gate_up" else 1 + weight_name = self.added_weight_attrs[weight_idx] + scale_name = self.added_scale_attrs[weight_idx] + weight = getattr(layer, weight_name) + if not hasattr(weight, "tensor_track") or weight.tensor_track is None: + return + + if self.model_format != "torch": + weight = weight.transpose([0, 2, 1]).contiguous() + + expected_weight_shape = ( + self.up_gate_proj_quant_weight_shape if weight_type == "gate_up" else self.down_proj_quant_weight_shape + ) + expected_packed_shape = ( + self.up_gate_proj_packed_weight_shape if weight_type == "gate_up" else self.down_proj_packed_weight_shape + ) + expected_scale_shape = self.up_gate_proj_scale_shape if weight_type == "gate_up" else self.down_proj_scale_shape + + if list(weight.shape) != list(expected_weight_shape): + raise ValueError( + f"MegaMoE {weight_type} BF16 weight shape mismatch for {weight_name}: " + f"got {list(weight.shape)}, expected {list(expected_weight_shape)}" ) - return topk_idx, topk_weights + if weight.dtype != paddle.bfloat16: + weight = weight.astype(paddle.bfloat16) + weight = weight.contiguous() - def apply_mage_moe(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj): - gate_out = gate(x).cast("float32") + logger.info( + f"MegaMoE dynamic quantize {weight_type}: input shape={list(weight.shape)}, dtype={weight.dtype}" + ) + weight_quantized, scale = cast_grouped_weights_to_fp4(weight) + if weight_type == "gate_up": + weight_quantized, scale = _interleave_weights((weight_quantized, scale)) + scale = _transpose_sf_for_utccp(scale) + + if list(weight_quantized.shape) != list(expected_packed_shape): + raise ValueError( + f"MegaMoE {weight_type} packed weight shape mismatch: " + f"got {list(weight_quantized.shape)}, expected {list(expected_packed_shape)}" + ) + if list(scale.shape) != list(expected_scale_shape): + raise ValueError( + f"MegaMoE {weight_type} scale shape mismatch: " + f"got {list(scale.shape)}, expected {list(expected_scale_shape)}" + ) + if weight_quantized.dtype != paddle.int8: + raise ValueError(f"MegaMoE {weight_type} packed weight dtype mismatch: got {weight_quantized.dtype}") + if scale.dtype != paddle.int32: + raise ValueError(f"MegaMoE {weight_type} scale dtype mismatch: got {scale.dtype}") + + free_tensor(getattr(layer, weight_name)) + setattr( + layer, + weight_name, + layer.create_parameter( + shape=weight_quantized.shape, + dtype=paddle.int8, + default_initializer=paddle.nn.initializer.Constant(0), + ), + ) + setattr( + layer, + scale_name, + layer.create_parameter( + shape=scale.shape, + dtype=paddle.int32, + default_initializer=paddle.nn.initializer.Constant(0), + ).as_strided(scale.shape, scale.stride()), + ) + getattr(layer, weight_name).copy_(weight_quantized, False) + getattr(layer, scale_name).copy_(scale, False) + logger.info( + f"MegaMoE dynamic quantize {weight_type}: packed weight shape={list(weight_quantized.shape)}, " + f"scale shape={list(scale.shape)}" + ) + + if not self.quant_config.is_checkpoint_bf16: + return + if hasattr(layer, "up_gate_proj_weight") and weight_fully_copied(layer.up_gate_proj_weight): + _process_quantize_mega_moe("gate_up") + if hasattr(layer, "down_proj_weight") and weight_fully_copied(layer.down_proj_weight): + _process_quantize_mega_moe("down") + + def process_prequanted_weights(self, layer: nn.Layer, state_dict, is_rearrange: bool = False): + """ + Paddle cutlass process prequanted weights. + """ + logger.info(f"start process_prequanted_weights in megamoe") + up_gate_proj_expert_weight_key = layer.weight_key_map.get("up_gate_proj_expert_weight_key", None) + down_proj_expert_weight_key = layer.weight_key_map.get("down_proj_expert_weight_key", None) + up_gate_proj_expert_weight_scale_key = layer.weight_key_map.get("up_gate_proj_expert_weight_scale_key", None) + down_proj_expert_weight_scale_key = layer.weight_key_map.get("down_proj_expert_weight_scale_key", None) + + up_gate_proj_weights, down_proj_weights, logical_expert_ids, _ = layer.load_experts_weight( + state_dict, up_gate_proj_expert_weight_key, down_proj_expert_weight_key, is_rearrange + ) + # self.check(layer, up_gate_proj_weights, down_proj_weights) + up_gate_proj_weight_scale = [] + down_proj_weight_scale = [] + + if isinstance(state_dict, list): + state_dict = dict(state_dict) + + for expert_idx in logical_expert_ids: + up_gate_proj_expert_weight_scale_key_name = up_gate_proj_expert_weight_scale_key.format(expert_idx) + down_proj_expert_weight_scale_key_name = down_proj_expert_weight_scale_key.format(expert_idx) + + up_gate_weight_scale = get_tensor( + ( + state_dict.pop(up_gate_proj_expert_weight_scale_key_name) + if up_gate_proj_expert_weight_scale_key_name in state_dict + else up_gate_proj_expert_weight_scale_key_name + ), + layer.fd_config.model_config.model, + ) + down_weight_scale = get_tensor( + ( + state_dict.pop(down_proj_expert_weight_scale_key_name) + if down_proj_expert_weight_scale_key_name in state_dict + else down_proj_expert_weight_scale_key_name + ), + layer.fd_config.model_config.model, + ) + + up_gate_proj_weight_scale.append( + up_gate_weight_scale + ) + down_proj_weight_scale.append( + down_weight_scale + ) + + + up_gate_proj_weight = ( + paddle.stack(up_gate_proj_weights, axis=0) + ) + down_proj_weight = ( + paddle.stack(down_proj_weights, axis=0) + ) + up_gate_proj_weight_scale = paddle.stack(up_gate_proj_weight_scale, axis=0).transpose([0, 2, 1]) + down_proj_weight_scale = paddle.stack(down_proj_weight_scale, axis=0).transpose([0, 2, 1]) + name_tensor_map = { + self.added_weight_attrs[0]: up_gate_proj_weight, + self.added_weight_attrs[1]: down_proj_weight, + self.added_scale_attrs[0]: up_gate_proj_weight_scale, + self.added_scale_attrs[1]: down_proj_weight_scale, + } + for name, tensor in name_tensor_map.items(): + getattr(layer, name).data = tensor + + def apply_ep_prefill(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj): + return self.apply_mage_moe( + layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj + ) + + def apply_ep_decode(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj): + return self.apply_mage_moe( + layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj + ) + + def apply_mage_moe(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj): hidden_size = layer.hidden_size num_tokens = x.shape[0] + gate_out = gate(x).cast("float32") + # 1. Select topk experts and weights. - topk_idx, topk_weights = self.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + buffer_capacity = self.mega_moe_buffer.x.shape[0] + if num_tokens > buffer_capacity: + raise ValueError( + f"MegaMoE buffer capacity exceeded: num_tokens={num_tokens}, capacity={buffer_capacity}" + ) + + # copy x, topk_idx, topk_weights to mega_moe_buffer and quantization. mega_moe_pre_dispatch( x, topk_idx, @@ -1132,25 +1368,20 @@ def apply_mage_moe(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_ 32, # group_size ) - buffer_capacity = self.mega_moe_buffer.x.shape[0] - if num_tokens > buffer_capacity: - raise ValueError( - f"MegaMoE buffer capacity exceeded: num_tokens={num_tokens}, capacity={buffer_capacity}" - ) - l1_weight = getattr(layer, self.added_weight_attrs[0]) l1_scale = getattr(layer, self.added_scale_attrs[0]) l2_weight = getattr(layer, self.added_weight_attrs[1]) l2_scale = getattr(layer, self.added_scale_attrs[1]) - y = paddle.empty((max(num_tokens, 1), hidden_size), dtype=paddle.bfloat16) + y = paddle.empty((num_tokens, hidden_size), dtype=paddle.bfloat16) - swiglu_limit = getattr(layer.fd_config.model_config, "swiglu_limit", None) + swiglu_limit = getattr(layer.fd_config.model_config, "swiglu_limit", 10) deep_gemm.fp8_fp4_mega_moe( y, (l1_weight, l1_scale), (l2_weight, l2_scale), self.mega_moe_buffer, - recipe=(1, 1, 32), + cumulative_local_expert_recv_stats=self.cumulative_local_expert_recv_stats, + recipe=(1, 1, self.gran_k), activation="swiglu", activation_clamp=swiglu_limit, fast_math=True, diff --git a/fastdeploy/model_executor/layers/quantization/__init__.py b/fastdeploy/model_executor/layers/quantization/__init__.py index 5780edc1d1d..915141a90b8 100644 --- a/fastdeploy/model_executor/layers/quantization/__init__.py +++ b/fastdeploy/model_executor/layers/quantization/__init__.py @@ -30,6 +30,7 @@ "weight_only", "block_wise_fp8", "w4afp8", + "wfp4afp8", "w8a8", "w4a8", "wfp8afp8", @@ -67,11 +68,43 @@ def _is_full_quantization_config(quantization_dict): return True return False +def _is_mega_moe_quantization_config(quantization_config): + return ( + isinstance(quantization_config, dict) + and quantization_config.get("moe_quant_type") == "wfp4afp8" + ) + +def _get_mega_moe_quantization_config(): + return { + "quantization": "mix_quant", + "kv_cache_quant_type": "block_wise_fp8", + "dense_quant_type": "block_wise_fp8", + "moe_quant_type": "wfp4afp8", + "is_quantized": False, + } + def parse_quant_config(args, model_config, is_ernie, is_v1_loader): if args.quantization is not None and isinstance(args.quantization, str): args.quantization = parse_quantization(args.quantization) + enable_mega_moe = getattr(args, "enable_mega_moe", False) + if enable_mega_moe: + mega_moe_quantization_config = _get_mega_moe_quantization_config() + + if args.quantization is None and model_config.quantization_config is None: + args.quantization = mega_moe_quantization_config + if args.quantization is not None and not _is_mega_moe_quantization_config(args.quantization): + raise ValueError( + "--enable-mega-moe requires moe_quant_type=wfp4afp8." + ) + if model_config.quantization_config is not None and not _is_mega_moe_quantization_config( + model_config.quantization_config + ): + raise ValueError( + "--enable-mega-moe conflicts with model quantization_config. It requires moe_quant_type=wfp4afp8." + ) + # Determine whether CLI --quantization is a simple method name or a full JSON quantization_config cli_quantization = args.quantization cli_is_full_config = ( @@ -249,6 +282,7 @@ def get_quantization_config(quantization: str) -> Type[QuantConfigBase]: from .tensor_wise_fp8 import TensorWiseFP8Config from .w4a8 import W4A8Config from .w4afp8 import W4AFP8Config + from .wfp4afp8 import WFP4AFP8Config from .w8a8 import W8A8Config from .weight_only import WeightOnlyConfig, WINT4Config, WINT8Config from .wfp8afp8 import WFP8AFP8Config @@ -267,6 +301,7 @@ def get_quantization_config(quantization: str) -> Type[QuantConfigBase]: "w8a8": W8A8Config, "w4a8": W4A8Config, "wfp8afp8": WFP8AFP8Config, + "wfp4afp8": WFP4AFP8Config, "tensor_wise_fp8": TensorWiseFP8Config, "kvcache": KvCacheQuantConfig, "mix_quant": MixQuantConfig, diff --git a/fastdeploy/model_executor/layers/quantization/fp8_utils.py b/fastdeploy/model_executor/layers/quantization/fp8_utils.py index 311c12b324c..9c4cdefbf98 100644 --- a/fastdeploy/model_executor/layers/quantization/fp8_utils.py +++ b/fastdeploy/model_executor/layers/quantization/fp8_utils.py @@ -237,7 +237,6 @@ def fused_stack_transpose_quant(expert_weight_list, use_ue8m0=False): return w, scale - def _interleave_weights(l1_weights): # [gate: 0..7, up: 0..7, gate: 8..15, up: 8..15, ...] instead of [gate | up] def interleave(t, gran: int = 8) -> paddle.Tensor: @@ -249,7 +248,6 @@ def interleave(t, gran: int = 8) -> paddle.Tensor: return interleave(l1_weights[0]), interleave(l1_weights[1]) - def _transpose_sf_for_utccp(sf: paddle.Tensor) -> paddle.Tensor: num_groups, mn, packed_sf_k = sf.shape assert sf.dtype == paddle.int and mn % 128 == 0 diff --git a/fastdeploy/model_executor/layers/quantization/nvfp4.py b/fastdeploy/model_executor/layers/quantization/nvfp4.py index 627901ddf20..86f32d39683 100644 --- a/fastdeploy/model_executor/layers/quantization/nvfp4.py +++ b/fastdeploy/model_executor/layers/quantization/nvfp4.py @@ -33,6 +33,7 @@ set_weight_attrs, ) from fastdeploy.worker.tbo import let_another_thread_run +from fastdeploy.model_executor.layers.moe.ep import EPRunner from .quant_base import QuantConfigBase, QuantMethodBase, is_nvfp4_supported @@ -672,7 +673,7 @@ def apply_ep_prefill( # 1. top experts and weights gate_out = gate(x.cast("float32")) - topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) hidden_size = x.shape[1] if topk_ids_hookfunc is not None: @@ -861,7 +862,7 @@ def apply_ep_decode( ) -> paddle.Tensor: gate_out = gate(x.cast("float32")) - topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) if topk_ids_hookfunc is not None: topk_ids_hookfunc(topk_ids=topk_idx) diff --git a/fastdeploy/model_executor/layers/quantization/wfp4afp8.py b/fastdeploy/model_executor/layers/quantization/wfp4afp8.py new file mode 100644 index 00000000000..8f66f56467a --- /dev/null +++ b/fastdeploy/model_executor/layers/quantization/wfp4afp8.py @@ -0,0 +1,66 @@ +""" +# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" + +from typing import Optional + +import paddle + +import fastdeploy + +from ..moe import FusedMoE +from .quant_base import QuantConfigBase, QuantMethodBase +from paddleformers.utils.log import logger + +QUANT_SCALING_FACTOR = 6 + + +class WFP4AFP8Config(QuantConfigBase): + """ + quantization config for weight 4bits and activation fp8 + """ + + def __init__(self, weight_scale_dict, act_scale_dict, is_permuted, is_quantized) -> None: + super().__init__() + self.weight_scale_dict = weight_scale_dict + self.act_scale_dict = act_scale_dict + self.quant_max_bound = 6 + self.quant_min_bound = -6 + self.quant_round_type = 1 + self.is_permuted = is_permuted + self.is_quantized = is_quantized + self.is_checkpoint_bf16 = not is_quantized + + def name(self) -> str: + return "wfp4afp8" + + @classmethod + def from_config(cls, config: dict) -> "WFP4AFP8Config": + weight_scale_dict = config.get("weight_scale_dict", None) + act_scale_dict = config.get("act_scale_dict", None) + is_permuted = config.get("is_permuted", True) + is_quantized = config.get("is_quantized", False) + return cls(weight_scale_dict, act_scale_dict, is_permuted, is_quantized) + + def get_quant_method(self, layer) -> Optional[QuantMethodBase]: + logger.info("Currently only support DeepGEMMMegaMoE for wfp4afp8") + if isinstance(layer, FusedMoE): + from fastdeploy.model_executor.layers.moe.fused_moe_deepgemm_backend import ( + DeepGemmMegaMoEMethod + ) + return DeepGemmMegaMoEMethod(self) + else: + raise NotImplementedError(f"wfp4afp8 quant method not supported for {type(layer)}") + diff --git a/fastdeploy/worker/worker_process.py b/fastdeploy/worker/worker_process.py index e97949da7ce..c7a7e7ea375 100644 --- a/fastdeploy/worker/worker_process.py +++ b/fastdeploy/worker/worker_process.py @@ -866,6 +866,13 @@ def parse_args(): action="store_true", help="enable chunked moe", ) + parser.add_argument( + "--enable_mega_moe", + "--enable-mega-moe", + action="store_true", + dest="enable_mega_moe", + help="enable MegaMoE wfp4afp8 for MoE and block_wise_fp8 for dense Linear", + ) parser.add_argument( "--chunked_moe_size", type=int, diff --git a/tests/model_executor/test_ep.py b/tests/model_executor/test_ep.py index b099c7ad57e..6b34104c10a 100644 --- a/tests/model_executor/test_ep.py +++ b/tests/model_executor/test_ep.py @@ -21,6 +21,7 @@ from fastdeploy.config import MoEPhase from fastdeploy.model_executor.layers.moe import ep +from fastdeploy.model_executor.layers.moe.ep import EPRunner class FakeConfig: @@ -423,7 +424,7 @@ def fake_get_moe_scores(*_args, **_kwargs): ) gate_out = paddle.randn([1, 4], dtype="float32") - topk_idx, topk_weights = runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) assert list(topk_idx.shape) == [1, 1] assert list(topk_weights.shape) == [1, 1] assert paddle.allclose(topk_idx, paddle.to_tensor([[1]], dtype="int64")) @@ -465,7 +466,7 @@ def get_ep_rank_to_expert_id_list_by_layer(self, _layer_idx): ) gate_out = paddle.randn([1, 4], dtype="float32") - topk_idx, topk_weights = runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) assert list(topk_idx.shape) == [1, 1] assert list(topk_weights.shape) == [1, 1] assert paddle.allclose(topk_idx, paddle.to_tensor([[2]], dtype="int64")) @@ -497,7 +498,7 @@ def fake_topk_select(*_args, **_kwargs): ) gate_out = paddle.randn([1, 4], dtype="float32") - topk_idx, topk_weights = runner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) assert list(topk_idx.shape) == [1, 1] assert list(topk_weights.shape) == [1, 1] assert paddle.allclose(topk_idx, paddle.to_tensor([[3]], dtype="int64")) From f3ad8a0bf11d5559207d2a29d6b7506171b3f475 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Mon, 8 Jun 2026 17:04:28 +0800 Subject: [PATCH 03/17] fix code style --- custom_ops/gpu_ops/mega_moe_pre_dispatch.cu | 82 ++++++++------ custom_ops/setup_ops.py | 5 +- fastdeploy/config.py | 2 +- .../layers/backends/xpu/moe/fused_moe.py | 2 +- fastdeploy/model_executor/layers/moe/ep.py | 5 +- .../layers/moe/fused_moe_backend_base.py | 2 +- .../layers/moe/fused_moe_blackwell_backend.py | 2 +- .../layers/moe/fused_moe_cutlass_backend.py | 2 +- .../layers/moe/fused_moe_deepgemm_backend.py | 105 ++++++------------ .../layers/quantization/__init__.py | 13 +-- .../layers/quantization/fp8_utils.py | 12 +- .../layers/quantization/nvfp4.py | 2 +- .../layers/quantization/wfp4afp8.py | 8 +- tests/operators/test_mega_moe_pre_dispatch.py | 10 +- 14 files changed, 117 insertions(+), 135 deletions(-) diff --git a/custom_ops/gpu_ops/mega_moe_pre_dispatch.cu b/custom_ops/gpu_ops/mega_moe_pre_dispatch.cu index b5629e54c14..30194210827 100644 --- a/custom_ops/gpu_ops/mega_moe_pre_dispatch.cu +++ b/custom_ops/gpu_ops/mega_moe_pre_dispatch.cu @@ -138,7 +138,10 @@ void CheckSameShape(const paddle::Tensor& lhs, const paddle::Tensor& rhs, const char* lhs_name, const char* rhs_name) { - PD_CHECK(lhs.shape() == rhs.shape(), lhs_name, " shape must equal ", rhs_name, + PD_CHECK(lhs.shape() == rhs.shape(), + lhs_name, + " shape must equal ", + rhs_name, " shape"); } @@ -153,16 +156,15 @@ void LaunchMegaMoEPreDispatch(const MegaMoEPreDispatchParams& params, } // namespace -void MegaMoePreDispatch( - const paddle::Tensor& x, - const paddle::Tensor& topk_idx, - const paddle::Tensor& topk_weights, - const paddle::Tensor& buf_x, - const paddle::Tensor& buf_x_sf, - const paddle::Tensor& buf_topk_idx, - const paddle::Tensor& buf_topk_weights, - int64_t num_max_tokens_per_rank, - int64_t group_size) { +void MegaMoePreDispatch(const paddle::Tensor& x, + const paddle::Tensor& topk_idx, + const paddle::Tensor& topk_weights, + const paddle::Tensor& buf_x, + const paddle::Tensor& buf_x_sf, + const paddle::Tensor& buf_topk_idx, + const paddle::Tensor& buf_topk_weights, + int64_t num_max_tokens_per_rank, + int64_t group_size) { CheckShape2D(x, "x"); CheckShape2D(topk_idx, "topk_idx"); CheckShape2D(topk_weights, "topk_weights"); @@ -171,21 +173,27 @@ void MegaMoePreDispatch( CheckShape2D(buf_topk_idx, "buf_topk_idx"); CheckShape2D(buf_topk_weights, "buf_topk_weights"); CheckSameShape(topk_idx, topk_weights, "topk_idx", "topk_weights"); - CheckSameShape(buf_topk_idx, buf_topk_weights, "buf_topk_idx", - "buf_topk_weights"); + CheckSameShape( + buf_topk_idx, buf_topk_weights, "buf_topk_idx", "buf_topk_weights"); PD_CHECK(x.dtype() == paddle::DataType::BFLOAT16, - "x must be bfloat16, but got ", x.dtype()); + "x must be bfloat16, but got ", + x.dtype()); PD_CHECK(topk_idx.dtype() == paddle::DataType::INT64, - "topk_idx must be int64, but got ", topk_idx.dtype()); + "topk_idx must be int64, but got ", + topk_idx.dtype()); PD_CHECK(topk_weights.dtype() == paddle::DataType::FLOAT32, - "topk_weights must be float32, but got ", topk_weights.dtype()); + "topk_weights must be float32, but got ", + topk_weights.dtype()); PD_CHECK(buf_x.dtype() == paddle::DataType::FLOAT8_E4M3FN, - "buf_x must be float8_e4m3fn, but got ", buf_x.dtype()); + "buf_x must be float8_e4m3fn, but got ", + buf_x.dtype()); PD_CHECK(buf_x_sf.dtype() == paddle::DataType::INT32, - "buf_x_sf must be int32, but got ", buf_x_sf.dtype()); + "buf_x_sf must be int32, but got ", + buf_x_sf.dtype()); PD_CHECK(buf_topk_idx.dtype() == paddle::DataType::INT64, - "buf_topk_idx must be int64, but got ", buf_topk_idx.dtype()); + "buf_topk_idx must be int64, but got ", + buf_topk_idx.dtype()); PD_CHECK(buf_topk_weights.dtype() == paddle::DataType::FLOAT32, "buf_topk_weights must be float32, but got ", buf_topk_weights.dtype()); @@ -197,19 +205,24 @@ void MegaMoePreDispatch( PD_CHECK(num_max_tokens_per_rank <= padded_max_i64, "num_max_tokens_per_rank must not exceed buf_x.shape[0], but got ", - num_max_tokens_per_rank, " vs ", padded_max_i64); + num_max_tokens_per_rank, + " vs ", + padded_max_i64); PD_CHECK(num_tokens_i64 == topk_idx.shape()[0], "x.shape[0] must equal topk_idx.shape[0]"); PD_CHECK(buf_x.shape()[1] == hidden_i64, - "buf_x.shape[1] must equal hidden, but got ", buf_x.shape()[1], - " vs ", hidden_i64); + "buf_x.shape[1] must equal hidden, but got ", + buf_x.shape()[1], + " vs ", + hidden_i64); PD_CHECK(buf_topk_idx.shape()[0] == padded_max_i64, "buf_topk_idx.shape[0] must equal padded_max"); PD_CHECK(buf_topk_idx.shape()[1] == top_k_i64, "buf_topk_idx.shape[1] must equal top_k"); PD_CHECK(group_size == 32 || group_size == 64 || group_size == 128, - "unsupported group_size: ", group_size); + "unsupported group_size: ", + group_size); PD_CHECK(num_tokens_i64 <= num_max_tokens_per_rank, "num_tokens must not exceed padded_max"); PD_CHECK(hidden_i64 % group_size == 0, @@ -220,7 +233,9 @@ void MegaMoePreDispatch( "buf_x_sf.shape[0] must equal padded_max"); PD_CHECK(buf_x_sf.shape()[1] == num_groups_i64 / 4, "buf_x_sf.shape[1] must equal hidden/group_size/4, but got ", - buf_x_sf.shape()[1], " vs ", num_groups_i64 / 4); + buf_x_sf.shape()[1], + " vs ", + num_groups_i64 / 4); PD_CHECK(hidden_i64 % static_cast(kVecElems) == 0, "hidden must be a multiple of 8 (16B bf16 loads)"); const int64_t num_threads_i64 = hidden_i64 / static_cast(kVecElems); @@ -256,16 +271,16 @@ void MegaMoePreDispatch( auto stream = x.stream(); switch (group_size) { case 32: - LaunchMegaMoEPreDispatch<32>(params, num_total_blocks, num_threads, - stream); + LaunchMegaMoEPreDispatch<32>( + params, num_total_blocks, num_threads, stream); break; case 64: - LaunchMegaMoEPreDispatch<64>(params, num_total_blocks, num_threads, - stream); + LaunchMegaMoEPreDispatch<64>( + params, num_total_blocks, num_threads, stream); break; case 128: - LaunchMegaMoEPreDispatch<128>(params, num_total_blocks, num_threads, - stream); + LaunchMegaMoEPreDispatch<128>( + params, num_total_blocks, num_threads, stream); break; default: PD_THROW("unsupported group_size: ", group_size); @@ -283,7 +298,8 @@ std::vector MegaMoePreDispatchInferDtype( const paddle::DataType& buf_x_sf_dtype, const paddle::DataType& buf_topk_idx_dtype, const paddle::DataType& buf_topk_weights_dtype) { - return {buf_x_dtype, buf_x_sf_dtype, buf_topk_idx_dtype, buf_topk_weights_dtype}; + return { + buf_x_dtype, buf_x_sf_dtype, buf_topk_idx_dtype, buf_topk_weights_dtype}; } std::vector> MegaMoePreDispatchInferShape( @@ -294,10 +310,10 @@ std::vector> MegaMoePreDispatchInferShape( const std::vector& buf_x_sf_shape, const std::vector& buf_topk_idx_shape, const std::vector& buf_topk_weights_shape) { - return {buf_x_shape, buf_x_sf_shape, buf_topk_idx_shape, buf_topk_weights_shape}; + return { + buf_x_shape, buf_x_sf_shape, buf_topk_idx_shape, buf_topk_weights_shape}; } - PD_BUILD_STATIC_OP(mega_moe_pre_dispatch) .Inputs({"x", "topk_idx", diff --git a/custom_ops/setup_ops.py b/custom_ops/setup_ops.py index 9c2a1b03beb..7a789dc2dad 100644 --- a/custom_ops/setup_ops.py +++ b/custom_ops/setup_ops.py @@ -522,10 +522,7 @@ def find_end_files(directory, end_str): # Add SM100 specific sources if any, e.g., for new hardware intrinsics # sources += ["gpu_ops/cutlass_kernels/w8a8/c4x_sm100.cu"] # Example - sources += [ - "gpu_ops/mega_moe_pre_dispatch.cu" - ] - + sources += ["gpu_ops/mega_moe_pre_dispatch.cu"] if has_generic_fp8: # For SM89 (Ada) or other architectures without dedicated paths diff --git a/fastdeploy/config.py b/fastdeploy/config.py index 6d9e3d58038..1b080680e6f 100644 --- a/fastdeploy/config.py +++ b/fastdeploy/config.py @@ -652,7 +652,7 @@ def __init__( self.enable_chunked_moe = False self.chunked_moe_size = 256 self.enable_mega_moe = False - + self.local_data_parallel_id = 0 # Engine worker queue port self.engine_worker_queue_port: Union[int, str, list] = None diff --git a/fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py b/fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py index d58a85a64fa..4c1c7c233ef 100644 --- a/fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py +++ b/fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py @@ -22,6 +22,7 @@ from paddleformers.utils.log import logger from fastdeploy import envs +from fastdeploy.model_executor.layers.moe.ep import EPRunner from fastdeploy.model_executor.layers.moe.fused_moe_backend_base import MoEMethodBase from fastdeploy.model_executor.layers.quantization.weight_only import WeightOnlyConfig from fastdeploy.model_executor.layers.utils import get_tensor @@ -39,7 +40,6 @@ free_tensor, set_weight_attrs, ) -from fastdeploy.model_executor.layers.moe.ep import EPRunner from .utils import get_moe_scores diff --git a/fastdeploy/model_executor/layers/moe/ep.py b/fastdeploy/model_executor/layers/moe/ep.py index ca215098761..0253dc21c33 100644 --- a/fastdeploy/model_executor/layers/moe/ep.py +++ b/fastdeploy/model_executor/layers/moe/ep.py @@ -793,6 +793,7 @@ def combine(self, ffn_out, topk_idx, topk_weights, handle, **kwargs): class FakeEPRunner: """ """ + def __init__(self, *args, **kwargs): pass @@ -803,6 +804,6 @@ def dispatch(self, *args, **kwargs): def combine(self, *args, **kwargs): """ """ pass - + def clean_low_latency_buffer(self): - pass \ No newline at end of file + pass diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py b/fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py index ac105043397..7633ca79b1e 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_backend_base.py @@ -30,7 +30,7 @@ from fastdeploy.platforms import current_platform from ..quantization.quant_base import QuantMethodBase -from fastdeploy import envs + class MoEMethodBase(QuantMethodBase): """ """ diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py index 681ccb38c3b..dc20492d0f4 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py @@ -23,7 +23,7 @@ from paddleformers.utils.log import logger import fastdeploy -from fastdeploy.model_executor.layers.moe.ep import deep_ep, EPRunner +from fastdeploy.model_executor.layers.moe.ep import EPRunner, deep_ep from fastdeploy.model_executor.layers.quantization.fp8_utils import ( deep_gemm, paddlefleet_ops, diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py index 8ec0305c4cf..e97b16ee4f3 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py @@ -43,6 +43,7 @@ except: logger.warning("import w4afp8_gemm_scale_permute Failed!") +from fastdeploy.model_executor.layers.moe.ep import EPRunner from fastdeploy.model_executor.layers.moe.moe import get_moe_scores from fastdeploy.model_executor.layers.quantization.fp8_utils import paddlefleet_ops from fastdeploy.model_executor.utils import ( @@ -52,7 +53,6 @@ set_weight_attrs, weight_fully_copied, ) -from fastdeploy.model_executor.layers.moe.ep import EPRunner def m_grouped_bf16_gemm_nn_contiguous(x, y, expert_idx_per_token): diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py index ee9acf76af5..9a41bb27a55 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py @@ -22,31 +22,33 @@ import paddle.nn.functional as F from paddle import nn from paddleformers.utils.log import logger + import fastdeploy -from fastdeploy.model_executor.layers.moe.ep import deep_ep, EPRunner, FakeEPRunner +from fastdeploy.model_executor.layers.moe.ep import EPRunner, FakeEPRunner, deep_ep from fastdeploy.model_executor.layers.quantization.fp8_utils import ( - deep_gemm, - paddlefleet_ops, _interleave_weights, _transpose_sf_for_utccp, + deep_gemm, + paddlefleet_ops, ) from fastdeploy.model_executor.layers.utils import get_tensor from fastdeploy.model_executor.ops.gpu import ( count_tokens_per_expert_func, depermute_prefill_combine, - prefill_permute_to_masked_gemm, mega_moe_pre_dispatch, + prefill_permute_to_masked_gemm, ) -from fastdeploy.platforms import current_platform -from fastdeploy.utils import ceil_div, register_custom_python_op, singleton -from fastdeploy.worker.tbo import let_another_thread_run from fastdeploy.model_executor.utils import ( TensorTracker, - set_weight_attrs, free_tensor, - weight_fully_copied, get_sm_version, + set_weight_attrs, + weight_fully_copied, ) +from fastdeploy.platforms import current_platform +from fastdeploy.utils import ceil_div, register_custom_python_op, singleton +from fastdeploy.worker.tbo import let_another_thread_run + from .fused_moe_backend_base import MoEMethodBase from .fused_moe_triton_backend import BlockWiseFP8MoEMethod @@ -953,7 +955,7 @@ def create_weights(self, layer: nn.Layer, **extra_weight_attrs): """ Triton MoE create weight process. """ - logger.info("mega create_weights") + logger.info("MegaMoE create_weights...") self.model_format = extra_weight_attrs.get("model_format") self.up_gate_proj_quant_weight_shape = [ layer.num_local_experts, @@ -984,12 +986,12 @@ def create_weights(self, layer: nn.Layer, **extra_weight_attrs): self.up_gate_proj_packed_weight_shape = [ layer.num_local_experts, layer.moe_intermediate_size * 2, - layer.hidden_size // 2, # 4-bit packing + layer.hidden_size // 2, # 4-bit packing ] self.down_proj_packed_weight_shape = [ layer.num_local_experts, layer.hidden_size, - layer.moe_intermediate_size // 2, # 4-bit packing + layer.moe_intermediate_size // 2, # 4-bit packing ] up_num_scales = ceil_div(layer.hidden_size, self.gran_k) down_num_scales = ceil_div(layer.moe_intermediate_size, self.gran_k) @@ -1006,21 +1008,6 @@ def create_weights(self, layer: nn.Layer, **extra_weight_attrs): self.up_gate_proj_weight_shape = self.up_gate_proj_quant_weight_shape self.down_proj_weight_shape = self.down_proj_quant_weight_shape - logger.info( - "MegaMoE create_weights: " - f"is_checkpoint_bf16={self.quant_config.is_checkpoint_bf16}, " - f"load_choices={layer.fd_config.load_config.load_choices}, " - f"model_format={self.model_format}, " - f"up_gate_bf16_shape={self.up_gate_proj_bf16_weight_shape}, " - f"down_bf16_shape={self.down_proj_bf16_weight_shape}, " - f"up_gate_quant_shape={self.up_gate_proj_quant_weight_shape}, " - f"down_quant_shape={self.down_proj_quant_weight_shape}, " - f"up_gate_packed_shape={self.up_gate_proj_packed_weight_shape}, " - f"down_packed_shape={self.down_proj_packed_weight_shape}, " - f"up_gate_scale_shape={self.up_gate_proj_scale_shape}, " - f"down_scale_shape={self.down_proj_scale_shape}" - ) - if self.quant_config.is_checkpoint_bf16 and layer.fd_config.load_config.load_choices == "default_v1": if self.model_format != "torch": up_gate_proj_attrs = { @@ -1134,7 +1121,7 @@ def create_weights(self, layer: nn.Layer, **extra_weight_attrs): ) def init_ep(self, layer: nn.Layer) -> None: - logger.info("USE MegaMoE backend") + logger.info("Use MegaMoE backend") if layer.ep_size <= 1: return @@ -1158,24 +1145,25 @@ def init_ep(self, layer: nn.Layer) -> None: layer.moe_intermediate_size, ).buffer self.num_max_tokens_per_rank = self.mega_moe_buffer.num_max_tokens_per_rank - self.cumulative_local_expert_recv_stats = paddle.zeros( - (layer.num_local_experts,), dtype=paddle.int32 - ) + self.cumulative_local_expert_recv_stats = paddle.zeros((layer.num_local_experts,), dtype=paddle.int32) self.ep_prefill_runner = FakeEPRunner() self.ep_decoder_runner = FakeEPRunner() - + def process_weights_after_loading(self, layer): def cast_grouped_weights_to_fp4(bf16_weights: paddle.Tensor): num_groups, n, k = bf16_weights.shape w = paddle.empty((num_groups, n, k // 2), dtype=paddle.int8) w_sf = paddle.empty((num_groups, n, k // self.gran_k), dtype=paddle.float32) for i in range(num_groups): - w[i], w_sf[i] = deep_gemm.per_token_cast_to_fp4( - bf16_weights[i], use_ue8m0=True, gran_k=self.gran_k - ) + w[i], w_sf[i] = deep_gemm.per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=self.gran_k) w = w.contiguous() w_sf = w_sf.contiguous() + + # pack four scales into one and transform to specific stride. + # for example: + # shape: [48, 6144, 224] -> [48, 6144, 56] + # stride: [1376256, 224, 1] -> [344064, 1, 6144] w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, self.gran_k), num_groups) return w, w_sf @@ -1194,9 +1182,13 @@ def _process_quantize_mega_moe(weight_type): self.up_gate_proj_quant_weight_shape if weight_type == "gate_up" else self.down_proj_quant_weight_shape ) expected_packed_shape = ( - self.up_gate_proj_packed_weight_shape if weight_type == "gate_up" else self.down_proj_packed_weight_shape + self.up_gate_proj_packed_weight_shape + if weight_type == "gate_up" + else self.down_proj_packed_weight_shape + ) + expected_scale_shape = ( + self.up_gate_proj_scale_shape if weight_type == "gate_up" else self.down_proj_scale_shape ) - expected_scale_shape = self.up_gate_proj_scale_shape if weight_type == "gate_up" else self.down_proj_scale_shape if list(weight.shape) != list(expected_weight_shape): raise ValueError( @@ -1207,9 +1199,6 @@ def _process_quantize_mega_moe(weight_type): weight = weight.astype(paddle.bfloat16) weight = weight.contiguous() - logger.info( - f"MegaMoE dynamic quantize {weight_type}: input shape={list(weight.shape)}, dtype={weight.dtype}" - ) weight_quantized, scale = cast_grouped_weights_to_fp4(weight) if weight_type == "gate_up": weight_quantized, scale = _interleave_weights((weight_quantized, scale)) @@ -1251,10 +1240,6 @@ def _process_quantize_mega_moe(weight_type): ) getattr(layer, weight_name).copy_(weight_quantized, False) getattr(layer, scale_name).copy_(scale, False) - logger.info( - f"MegaMoE dynamic quantize {weight_type}: packed weight shape={list(weight_quantized.shape)}, " - f"scale shape={list(scale.shape)}" - ) if not self.quant_config.is_checkpoint_bf16: return @@ -1267,7 +1252,6 @@ def process_prequanted_weights(self, layer: nn.Layer, state_dict, is_rearrange: """ Paddle cutlass process prequanted weights. """ - logger.info(f"start process_prequanted_weights in megamoe") up_gate_proj_expert_weight_key = layer.weight_key_map.get("up_gate_proj_expert_weight_key", None) down_proj_expert_weight_key = layer.weight_key_map.get("down_proj_expert_weight_key", None) up_gate_proj_expert_weight_scale_key = layer.weight_key_map.get("up_gate_proj_expert_weight_scale_key", None) @@ -1304,20 +1288,11 @@ def process_prequanted_weights(self, layer: nn.Layer, state_dict, is_rearrange: layer.fd_config.model_config.model, ) - up_gate_proj_weight_scale.append( - up_gate_weight_scale - ) - down_proj_weight_scale.append( - down_weight_scale - ) - + up_gate_proj_weight_scale.append(up_gate_weight_scale) + down_proj_weight_scale.append(down_weight_scale) - up_gate_proj_weight = ( - paddle.stack(up_gate_proj_weights, axis=0) - ) - down_proj_weight = ( - paddle.stack(down_proj_weights, axis=0) - ) + up_gate_proj_weight = paddle.stack(up_gate_proj_weights, axis=0) + down_proj_weight = paddle.stack(down_proj_weights, axis=0) up_gate_proj_weight_scale = paddle.stack(up_gate_proj_weight_scale, axis=0).transpose([0, 2, 1]) down_proj_weight_scale = paddle.stack(down_proj_weight_scale, axis=0).transpose([0, 2, 1]) @@ -1331,14 +1306,10 @@ def process_prequanted_weights(self, layer: nn.Layer, state_dict, is_rearrange: getattr(layer, name).data = tensor def apply_ep_prefill(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj): - return self.apply_mage_moe( - layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj - ) + return self.apply_mage_moe(layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj) def apply_ep_decode(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj): - return self.apply_mage_moe( - layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj - ) + return self.apply_mage_moe(layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj) def apply_mage_moe(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj): hidden_size = layer.hidden_size @@ -1351,9 +1322,7 @@ def apply_mage_moe(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_ buffer_capacity = self.mega_moe_buffer.x.shape[0] if num_tokens > buffer_capacity: - raise ValueError( - f"MegaMoE buffer capacity exceeded: num_tokens={num_tokens}, capacity={buffer_capacity}" - ) + raise ValueError(f"MegaMoE buffer capacity exceeded: num_tokens={num_tokens}, capacity={buffer_capacity}") # copy x, topk_idx, topk_weights to mega_moe_buffer and quantization. mega_moe_pre_dispatch( @@ -1365,7 +1334,7 @@ def apply_mage_moe(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_ self.mega_moe_buffer.topk_idx, self.mega_moe_buffer.topk_weights, self.num_max_tokens_per_rank, - 32, # group_size + 32, # group_size ) l1_weight = getattr(layer, self.added_weight_attrs[0]) diff --git a/fastdeploy/model_executor/layers/quantization/__init__.py b/fastdeploy/model_executor/layers/quantization/__init__.py index 915141a90b8..217ae18aa39 100644 --- a/fastdeploy/model_executor/layers/quantization/__init__.py +++ b/fastdeploy/model_executor/layers/quantization/__init__.py @@ -68,11 +68,10 @@ def _is_full_quantization_config(quantization_dict): return True return False + def _is_mega_moe_quantization_config(quantization_config): - return ( - isinstance(quantization_config, dict) - and quantization_config.get("moe_quant_type") == "wfp4afp8" - ) + return isinstance(quantization_config, dict) and quantization_config.get("moe_quant_type") == "wfp4afp8" + def _get_mega_moe_quantization_config(): return { @@ -95,9 +94,7 @@ def parse_quant_config(args, model_config, is_ernie, is_v1_loader): if args.quantization is None and model_config.quantization_config is None: args.quantization = mega_moe_quantization_config if args.quantization is not None and not _is_mega_moe_quantization_config(args.quantization): - raise ValueError( - "--enable-mega-moe requires moe_quant_type=wfp4afp8." - ) + raise ValueError("--enable-mega-moe requires moe_quant_type=wfp4afp8.") if model_config.quantization_config is not None and not _is_mega_moe_quantization_config( model_config.quantization_config ): @@ -282,9 +279,9 @@ def get_quantization_config(quantization: str) -> Type[QuantConfigBase]: from .tensor_wise_fp8 import TensorWiseFP8Config from .w4a8 import W4A8Config from .w4afp8 import W4AFP8Config - from .wfp4afp8 import WFP4AFP8Config from .w8a8 import W8A8Config from .weight_only import WeightOnlyConfig, WINT4Config, WINT8Config + from .wfp4afp8 import WFP4AFP8Config from .wfp8afp8 import WFP8AFP8Config from .wint2 import WINT2Config diff --git a/fastdeploy/model_executor/layers/quantization/fp8_utils.py b/fastdeploy/model_executor/layers/quantization/fp8_utils.py index 9c4cdefbf98..094308a6ed7 100644 --- a/fastdeploy/model_executor/layers/quantization/fp8_utils.py +++ b/fastdeploy/model_executor/layers/quantization/fp8_utils.py @@ -237,6 +237,7 @@ def fused_stack_transpose_quant(expert_weight_list, use_ue8m0=False): return w, scale + def _interleave_weights(l1_weights): # [gate: 0..7, up: 0..7, gate: 8..15, up: 8..15, ...] instead of [gate | up] def interleave(t, gran: int = 8) -> paddle.Tensor: @@ -248,15 +249,18 @@ def interleave(t, gran: int = 8) -> paddle.Tensor: return interleave(l1_weights[0]), interleave(l1_weights[1]) + def _transpose_sf_for_utccp(sf: paddle.Tensor) -> paddle.Tensor: num_groups, mn, packed_sf_k = sf.shape assert sf.dtype == paddle.int and mn % 128 == 0 # sf is MN-major: strides [mn*packed_sf_k, 1, mn] # We need to do the 4x32 transpose in data while preserving MN-major strides sf_c = sf.contiguous() # make C-contiguous for reshape/transpose - result_c = (sf_c.reshape(num_groups, -1, 4, 32, packed_sf_k) - .transpose(2, 3) - .reshape(num_groups, mn, packed_sf_k) - .contiguous()) + result_c = ( + sf_c.reshape(num_groups, -1, 4, 32, packed_sf_k) + .transpose(2, 3) + .reshape(num_groups, mn, packed_sf_k) + .contiguous() + ) # Convert back to MN-major layout: transpose last two dims, make contiguous, transpose back return result_c.transpose(1, 2).contiguous().transpose(1, 2) diff --git a/fastdeploy/model_executor/layers/quantization/nvfp4.py b/fastdeploy/model_executor/layers/quantization/nvfp4.py index 86f32d39683..0e5722a343a 100644 --- a/fastdeploy/model_executor/layers/quantization/nvfp4.py +++ b/fastdeploy/model_executor/layers/quantization/nvfp4.py @@ -25,6 +25,7 @@ import fastdeploy from fastdeploy import envs from fastdeploy.model_executor.layers.moe import FusedMoE +from fastdeploy.model_executor.layers.moe.ep import EPRunner from fastdeploy.model_executor.layers.moe.fused_moe_backend_base import MoEMethodBase from fastdeploy.model_executor.utils import ( create_parameter_and_copy, @@ -33,7 +34,6 @@ set_weight_attrs, ) from fastdeploy.worker.tbo import let_another_thread_run -from fastdeploy.model_executor.layers.moe.ep import EPRunner from .quant_base import QuantConfigBase, QuantMethodBase, is_nvfp4_supported diff --git a/fastdeploy/model_executor/layers/quantization/wfp4afp8.py b/fastdeploy/model_executor/layers/quantization/wfp4afp8.py index 8f66f56467a..07a0680e9dd 100644 --- a/fastdeploy/model_executor/layers/quantization/wfp4afp8.py +++ b/fastdeploy/model_executor/layers/quantization/wfp4afp8.py @@ -16,13 +16,11 @@ from typing import Optional -import paddle +from paddleformers.utils.log import logger -import fastdeploy from ..moe import FusedMoE from .quant_base import QuantConfigBase, QuantMethodBase -from paddleformers.utils.log import logger QUANT_SCALING_FACTOR = 6 @@ -58,9 +56,9 @@ def get_quant_method(self, layer) -> Optional[QuantMethodBase]: logger.info("Currently only support DeepGEMMMegaMoE for wfp4afp8") if isinstance(layer, FusedMoE): from fastdeploy.model_executor.layers.moe.fused_moe_deepgemm_backend import ( - DeepGemmMegaMoEMethod + DeepGemmMegaMoEMethod, ) + return DeepGemmMegaMoEMethod(self) else: raise NotImplementedError(f"wfp4afp8 quant method not supported for {type(layer)}") - diff --git a/tests/operators/test_mega_moe_pre_dispatch.py b/tests/operators/test_mega_moe_pre_dispatch.py index f5984309665..e6196b8d899 100644 --- a/tests/operators/test_mega_moe_pre_dispatch.py +++ b/tests/operators/test_mega_moe_pre_dispatch.py @@ -17,10 +17,12 @@ import numpy as np import paddle import paddle.distributed as dist +from ernie5_serving.mm_custom_ops import mega_moe_pre_dispatch from paddle.distributed import fleet -from ernie5_serving.mm_custom_ops import mega_moe_pre_dispatch -from fastdeploy.model_executor.layers.moe.fused_moe_deepgemm_backend import MegaMoEBuffer +from fastdeploy.model_executor.layers.moe.fused_moe_deepgemm_backend import ( + MegaMoEBuffer, +) def ceil_div(x: int, y: int) -> int: @@ -108,9 +110,7 @@ def _new_buffer(self): def mega_moe_pre_dispatch_ref(self, x: paddle.Tensor, topk_idx: paddle.Tensor, topk_weights: paddle.Tensor): num_tokens = x.shape[0] - x_fp8, x_scale_tensor = per_token_cast_to_fp8( - x, use_ue8m0=True, gran_k=self.group_size, use_packed_ue8m0=True - ) + x_fp8, x_scale_tensor = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=self.group_size, use_packed_ue8m0=True) return ( x_fp8, x_scale_tensor, From 4d3cb82893e0cc06a93004d7d6d3ec960dd0e56e Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Mon, 8 Jun 2026 17:06:14 +0800 Subject: [PATCH 04/17] fix code style --- .../layers/quantization/wfp4afp8.py | 1 - tests/model_executor/test_ep.py | 24 ------------------- tests/operators/test_mega_moe_pre_dispatch.py | 1 - 3 files changed, 26 deletions(-) diff --git a/fastdeploy/model_executor/layers/quantization/wfp4afp8.py b/fastdeploy/model_executor/layers/quantization/wfp4afp8.py index 07a0680e9dd..6dabdba7a7e 100644 --- a/fastdeploy/model_executor/layers/quantization/wfp4afp8.py +++ b/fastdeploy/model_executor/layers/quantization/wfp4afp8.py @@ -18,7 +18,6 @@ from paddleformers.utils.log import logger - from ..moe import FusedMoE from .quant_base import QuantConfigBase, QuantMethodBase diff --git a/tests/model_executor/test_ep.py b/tests/model_executor/test_ep.py index 6b34104c10a..f069d9dc624 100644 --- a/tests/model_executor/test_ep.py +++ b/tests/model_executor/test_ep.py @@ -403,14 +403,6 @@ def fake_get_moe_scores(*_args, **_kwargs): monkeypatch.setattr(moe_module, "get_moe_scores", fake_get_moe_scores, raising=True) - runner = ep.EPPrefillRunner( - top_k=2, - hidden_size=4, - num_experts=2, - splitwise_role="prefill", - num_max_dispatch_tokens_per_rank=1, - ) - layer = SimpleNamespace( redundant_table_manger=None, topk_method="noaux_tc", @@ -441,14 +433,6 @@ def fake_redundant_topk_select(**_kwargs): monkeypatch.setattr(gpu_ops, "moe_redundant_topk_select", fake_redundant_topk_select, raising=True) - runner = ep.EPPrefillRunner( - top_k=2, - hidden_size=4, - num_experts=2, - splitwise_role="prefill", - num_max_dispatch_tokens_per_rank=1, - ) - class FakeRedundantTableManager: def get_ep_rank_to_expert_id_list_by_layer(self, _layer_idx): return [0], paddle.to_tensor([0], dtype="int64"), [1], [1] @@ -483,14 +467,6 @@ def fake_topk_select(*_args, **_kwargs): monkeypatch.setattr(gpu_ops, "moe_topk_select", fake_topk_select, raising=True) - runner = ep.EPPrefillRunner( - top_k=2, - hidden_size=4, - num_experts=2, - splitwise_role="prefill", - num_max_dispatch_tokens_per_rank=1, - ) - layer = SimpleNamespace( redundant_table_manger=None, topk_method="aux", diff --git a/tests/operators/test_mega_moe_pre_dispatch.py b/tests/operators/test_mega_moe_pre_dispatch.py index e6196b8d899..2f39ae83510 100644 --- a/tests/operators/test_mega_moe_pre_dispatch.py +++ b/tests/operators/test_mega_moe_pre_dispatch.py @@ -109,7 +109,6 @@ def _new_buffer(self): ).buffer def mega_moe_pre_dispatch_ref(self, x: paddle.Tensor, topk_idx: paddle.Tensor, topk_weights: paddle.Tensor): - num_tokens = x.shape[0] x_fp8, x_scale_tensor = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=self.group_size, use_packed_ue8m0=True) return ( x_fp8, From 06c347a0676d96dfce8fd57eca542f4955b8ca67 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Mon, 8 Jun 2026 17:07:39 +0800 Subject: [PATCH 05/17] fix test --- tests/operators/test_mega_moe_pre_dispatch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/operators/test_mega_moe_pre_dispatch.py b/tests/operators/test_mega_moe_pre_dispatch.py index 2f39ae83510..14ba3e91724 100644 --- a/tests/operators/test_mega_moe_pre_dispatch.py +++ b/tests/operators/test_mega_moe_pre_dispatch.py @@ -17,7 +17,7 @@ import numpy as np import paddle import paddle.distributed as dist -from ernie5_serving.mm_custom_ops import mega_moe_pre_dispatch +from fastdeploy.model_executor.ops.gpu import mega_moe_pre_dispatch from paddle.distributed import fleet from fastdeploy.model_executor.layers.moe.fused_moe_deepgemm_backend import ( From f36da3fafbf5965f00a6ea19b16d54d056016341 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Mon, 8 Jun 2026 17:18:19 +0800 Subject: [PATCH 06/17] fix test --- tests/operators/test_mega_moe_pre_dispatch.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/operators/test_mega_moe_pre_dispatch.py b/tests/operators/test_mega_moe_pre_dispatch.py index 14ba3e91724..a4f45e46ac1 100644 --- a/tests/operators/test_mega_moe_pre_dispatch.py +++ b/tests/operators/test_mega_moe_pre_dispatch.py @@ -73,7 +73,7 @@ class TestMegaMoEPreDispatch(unittest.TestCase): def setUpClass(cls): paddle.seed(2025) strategy = fleet.DistributedStrategy() - cls.expert_parallel_size = 8 + cls.expert_parallel_size = 2 strategy.hybrid_configs = { "dp_degree": 1, "mp_degree": cls.expert_parallel_size, @@ -95,7 +95,7 @@ def setUp(self): self.x = paddle.randn([self.num_tokens, self.hidden_size], dtype=paddle.bfloat16) scores = paddle.randn((self.num_tokens, self.num_experts), dtype=paddle.float32) self.topk_weights, self.topk_idx = paddle.topk(scores, self.top_k, axis=-1, largest=True, sorted=False) - self.topk_idx = self.topk_idx.astype("int32") + self.topk_idx = self.topk_idx.astype("int64") self.topk_weights = self.topk_weights.astype("float32") def _new_buffer(self): From f4b400d686024f74fd2e9ad64d02e3c01028f766 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Mon, 8 Jun 2026 17:31:25 +0800 Subject: [PATCH 07/17] fix test --- tests/model_executor/test_ep.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/model_executor/test_ep.py b/tests/model_executor/test_ep.py index f069d9dc624..83972476be4 100644 --- a/tests/model_executor/test_ep.py +++ b/tests/model_executor/test_ep.py @@ -471,6 +471,7 @@ def fake_topk_select(*_args, **_kwargs): redundant_table_manger=None, topk_method="aux", gate_correction_bias=None, + top_k=2, ) gate_out = paddle.randn([1, 4], dtype="float32") From 910dee5e24271e6b7d72228890c95cd52238bf33 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Mon, 8 Jun 2026 18:07:14 +0800 Subject: [PATCH 08/17] fix test --- custom_ops/setup_ops.py | 3 +- .../layers/moe/fused_moe_deepgemm_backend.py | 6 +-- tests/operators/test_mega_moe_pre_dispatch.py | 50 ++++++++++--------- 3 files changed, 31 insertions(+), 28 deletions(-) diff --git a/custom_ops/setup_ops.py b/custom_ops/setup_ops.py index 7a789dc2dad..6568dbe37b4 100644 --- a/custom_ops/setup_ops.py +++ b/custom_ops/setup_ops.py @@ -348,6 +348,7 @@ def find_end_files(directory, end_str): "gpu_ops/gelu_tanh.cu", "gpu_ops/reasoning_phase_token_constraint.cu", "gpu_ops/get_attn_mask_q.cu", + "gpu_ops/mega_moe_pre_dispatch.cu" ] sm_versions = get_sm_version(archs) # Some kernels in this file require SM75+ instructions. Exclude them when building SM70 (V100). @@ -522,7 +523,7 @@ def find_end_files(directory, end_str): # Add SM100 specific sources if any, e.g., for new hardware intrinsics # sources += ["gpu_ops/cutlass_kernels/w8a8/c4x_sm100.cu"] # Example - sources += ["gpu_ops/mega_moe_pre_dispatch.cu"] + pass # No SM100 specific sources identified yet beyond what CUTLASS handles if has_generic_fp8: # For SM89 (Ada) or other architectures without dedicated paths diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py index 9a41bb27a55..af663c67233 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py @@ -1306,12 +1306,12 @@ def process_prequanted_weights(self, layer: nn.Layer, state_dict, is_rearrange: getattr(layer, name).data = tensor def apply_ep_prefill(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj): - return self.apply_mage_moe(layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj) + return self.apply_mega_moe(layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj) def apply_ep_decode(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj): - return self.apply_mage_moe(layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj) + return self.apply_mega_moe(layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj) - def apply_mage_moe(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj): + def apply_mega_moe(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_latent_proj, fc2_latent_proj): hidden_size = layer.hidden_size num_tokens = x.shape[0] diff --git a/tests/operators/test_mega_moe_pre_dispatch.py b/tests/operators/test_mega_moe_pre_dispatch.py index a4f45e46ac1..f6ea393c0f0 100644 --- a/tests/operators/test_mega_moe_pre_dispatch.py +++ b/tests/operators/test_mega_moe_pre_dispatch.py @@ -16,13 +16,17 @@ import numpy as np import paddle -import paddle.distributed as dist + from fastdeploy.model_executor.ops.gpu import mega_moe_pre_dispatch -from paddle.distributed import fleet +from dataclasses import dataclass + +@dataclass +class FakeBuffer: + x: paddle.Tensor + x_sf: paddle.Tensor + topk_idx: paddle.Tensor + topk_weights: paddle.Tensor -from fastdeploy.model_executor.layers.moe.fused_moe_deepgemm_backend import ( - MegaMoEBuffer, -) def ceil_div(x: int, y: int) -> int: @@ -72,16 +76,6 @@ class TestMegaMoEPreDispatch(unittest.TestCase): @classmethod def setUpClass(cls): paddle.seed(2025) - strategy = fleet.DistributedStrategy() - cls.expert_parallel_size = 2 - strategy.hybrid_configs = { - "dp_degree": 1, - "mp_degree": cls.expert_parallel_size, - "pp_degree": 1, - "sharding_degree": 1, - } - fleet.init(is_collective=True, strategy=strategy) - cls.ep_group = dist.new_group(range(cls.expert_parallel_size)) def setUp(self): self.num_experts = 160 @@ -99,17 +93,25 @@ def setUp(self): self.topk_weights = self.topk_weights.astype("float32") def _new_buffer(self): - return MegaMoEBuffer( - self.ep_group, - self.num_experts, - self.num_max_tokens_per_rank, - self.top_k, - self.hidden_size, - self.moe_intermediate_size, - ).buffer + x = paddle.zeros([self.num_max_tokens_per_rank, self.hidden_size], dtype=paddle.bfloat16).astype("float8_e4m3fn") + x_sf = paddle.zeros([self.num_max_tokens_per_rank, self.hidden_size // self.group_size // 4], paddle.int32) + topk_idx = paddle.zeros([self.num_max_tokens_per_rank, self.top_k], dtype=paddle.int64) + topk_weights = paddle.zeros([self.num_max_tokens_per_rank, self.top_k], dtype=paddle.float32) + fake_buffer = FakeBuffer( + x=x, + x_sf=x_sf, + topk_idx=topk_idx, + topk_weights=topk_weights + ) + + return fake_buffer + def mega_moe_pre_dispatch_ref(self, x: paddle.Tensor, topk_idx: paddle.Tensor, topk_weights: paddle.Tensor): - x_fp8, x_scale_tensor = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=self.group_size, use_packed_ue8m0=True) + num_tokens = x.shape[0] + x_fp8, x_scale_tensor = per_token_cast_to_fp8( + x, use_ue8m0=True, gran_k=self.group_size, use_packed_ue8m0=True + ) return ( x_fp8, x_scale_tensor, From 5cfd2814c14771fa102e83724e70fb0688e98eda Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Mon, 8 Jun 2026 18:08:55 +0800 Subject: [PATCH 09/17] fix typo --- .../model_executor/layers/moe/fused_moe_deepgemm_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py index af663c67233..21b384acd04 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py @@ -1334,7 +1334,7 @@ def apply_mega_moe(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_ self.mega_moe_buffer.topk_idx, self.mega_moe_buffer.topk_weights, self.num_max_tokens_per_rank, - 32, # group_size + self.gran_k, # group_size ) l1_weight = getattr(layer, self.added_weight_attrs[0]) From be143c5e5c784e38b5d31f7640312f34d9a62f54 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Mon, 8 Jun 2026 18:11:19 +0800 Subject: [PATCH 10/17] fix code style --- custom_ops/setup_ops.py | 4 ++-- tests/operators/test_mega_moe_pre_dispatch.py | 21 +++++++------------ 2 files changed, 9 insertions(+), 16 deletions(-) diff --git a/custom_ops/setup_ops.py b/custom_ops/setup_ops.py index 6568dbe37b4..cc6834b56df 100644 --- a/custom_ops/setup_ops.py +++ b/custom_ops/setup_ops.py @@ -348,7 +348,7 @@ def find_end_files(directory, end_str): "gpu_ops/gelu_tanh.cu", "gpu_ops/reasoning_phase_token_constraint.cu", "gpu_ops/get_attn_mask_q.cu", - "gpu_ops/mega_moe_pre_dispatch.cu" + "gpu_ops/mega_moe_pre_dispatch.cu", ] sm_versions = get_sm_version(archs) # Some kernels in this file require SM75+ instructions. Exclude them when building SM70 (V100). @@ -523,7 +523,7 @@ def find_end_files(directory, end_str): # Add SM100 specific sources if any, e.g., for new hardware intrinsics # sources += ["gpu_ops/cutlass_kernels/w8a8/c4x_sm100.cu"] # Example - pass # No SM100 specific sources identified yet beyond what CUTLASS handles + pass # No SM100 specific sources identified yet beyond what CUTLASS handles if has_generic_fp8: # For SM89 (Ada) or other architectures without dedicated paths diff --git a/tests/operators/test_mega_moe_pre_dispatch.py b/tests/operators/test_mega_moe_pre_dispatch.py index f6ea393c0f0..9d80c7ddee2 100644 --- a/tests/operators/test_mega_moe_pre_dispatch.py +++ b/tests/operators/test_mega_moe_pre_dispatch.py @@ -13,12 +13,13 @@ # limitations under the License. import unittest +from dataclasses import dataclass import numpy as np import paddle from fastdeploy.model_executor.ops.gpu import mega_moe_pre_dispatch -from dataclasses import dataclass + @dataclass class FakeBuffer: @@ -28,7 +29,6 @@ class FakeBuffer: topk_weights: paddle.Tensor - def ceil_div(x: int, y: int) -> int: return (x + y - 1) // y @@ -93,25 +93,18 @@ def setUp(self): self.topk_weights = self.topk_weights.astype("float32") def _new_buffer(self): - x = paddle.zeros([self.num_max_tokens_per_rank, self.hidden_size], dtype=paddle.bfloat16).astype("float8_e4m3fn") + x = paddle.zeros([self.num_max_tokens_per_rank, self.hidden_size], dtype=paddle.bfloat16).astype( + "float8_e4m3fn" + ) x_sf = paddle.zeros([self.num_max_tokens_per_rank, self.hidden_size // self.group_size // 4], paddle.int32) topk_idx = paddle.zeros([self.num_max_tokens_per_rank, self.top_k], dtype=paddle.int64) topk_weights = paddle.zeros([self.num_max_tokens_per_rank, self.top_k], dtype=paddle.float32) - fake_buffer = FakeBuffer( - x=x, - x_sf=x_sf, - topk_idx=topk_idx, - topk_weights=topk_weights - ) + fake_buffer = FakeBuffer(x=x, x_sf=x_sf, topk_idx=topk_idx, topk_weights=topk_weights) return fake_buffer - def mega_moe_pre_dispatch_ref(self, x: paddle.Tensor, topk_idx: paddle.Tensor, topk_weights: paddle.Tensor): - num_tokens = x.shape[0] - x_fp8, x_scale_tensor = per_token_cast_to_fp8( - x, use_ue8m0=True, gran_k=self.group_size, use_packed_ue8m0=True - ) + x_fp8, x_scale_tensor = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=self.group_size, use_packed_ue8m0=True) return ( x_fp8, x_scale_tensor, From 4f2b4809ac724fee7e47b099c1b764378bd21ebc Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Tue, 9 Jun 2026 17:02:06 +0800 Subject: [PATCH 11/17] fix test --- fastdeploy/model_executor/layers/moe/ep.py | 3 +- .../layers/moe/fused_moe_blackwell_backend.py | 6 ++-- .../layers/moe/fused_moe_cutlass_backend.py | 5 ++- tests/engine/test_engine.py | 1 + tests/model_executor/test_ep.py | 31 ++++++++++++++++--- 5 files changed, 34 insertions(+), 12 deletions(-) diff --git a/fastdeploy/model_executor/layers/moe/ep.py b/fastdeploy/model_executor/layers/moe/ep.py index 0253dc21c33..11c98463fcf 100644 --- a/fastdeploy/model_executor/layers/moe/ep.py +++ b/fastdeploy/model_executor/layers/moe/ep.py @@ -490,8 +490,7 @@ def __init__( top_k=self.top_k, ) - @staticmethod - def moe_select(layer: nn.Layer, gate_out: paddle.Tensor): + def moe_select(self, layer: nn.Layer, gate_out: paddle.Tensor): if layer.redundant_table_manger is not None: ( ep_rank_to_expert_id_list, diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py index dc20492d0f4..274deda8b69 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_blackwell_backend.py @@ -23,7 +23,7 @@ from paddleformers.utils.log import logger import fastdeploy -from fastdeploy.model_executor.layers.moe.ep import EPRunner, deep_ep +from fastdeploy.model_executor.layers.moe.ep import deep_ep from fastdeploy.model_executor.layers.quantization.fp8_utils import ( deep_gemm, paddlefleet_ops, @@ -642,7 +642,7 @@ def apply_ep_prefill( getattr(layer, "renormalize", True), ) else: - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out) if topk_ids_hookfunc is not None: topk_ids_hookfunc(topk_ids=topk_idx) @@ -964,7 +964,7 @@ def apply_ep_decode( gate_out = gate(x) gate_out = gate_out.cast("float32") # 1. Select topk experts and weights - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out) if topk_ids_hookfunc is not None: topk_ids_hookfunc(topk_ids=topk_idx) diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py index e97b16ee4f3..5334114b70e 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_cutlass_backend.py @@ -43,7 +43,6 @@ except: logger.warning("import w4afp8_gemm_scale_permute Failed!") -from fastdeploy.model_executor.layers.moe.ep import EPRunner from fastdeploy.model_executor.layers.moe.moe import get_moe_scores from fastdeploy.model_executor.layers.quantization.fp8_utils import paddlefleet_ops from fastdeploy.model_executor.utils import ( @@ -142,7 +141,7 @@ def apply_ep_prefill( if fc1_latent_proj is not None: x = fc1_latent_proj(x) # 1. Select topk experts and weights - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out) if layer.routed_scaling_factor_learnable: safe_topk_indices = paddle.clip(topk_idx, min=0) @@ -305,7 +304,7 @@ def apply_ep_decode( estimate_total_token_nums = gate_out.shape[0] * layer.top_k # 1. Select topk experts and weights - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out) if layer.routed_scaling_factor_learnable: safe_topk_indices = paddle.clip(topk_idx, min=0) diff --git a/tests/engine/test_engine.py b/tests/engine/test_engine.py index b465169df57..a2a513e4cfe 100644 --- a/tests/engine/test_engine.py +++ b/tests/engine/test_engine.py @@ -36,6 +36,7 @@ def _make_cfg(**ov): pc = ns(tensor_parallel_size=1, tensor_parallel_rank=0, device_ids="0", data_parallel_size=1) pc.expert_parallel_size, pc.chunked_moe_size, pc.engine_worker_queue_port = 1, 0, [6778] pc.enable_expert_parallel = pc.enable_chunked_moe = pc.disable_custom_all_reduce = False + pc.enable_mega_moe = False pc.use_internode_ll_two_stage = pc.disable_sequence_parallel_moe = False pc.shutdown_comm_group_if_worker_idle = False pc.ep_prefill_use_worst_num_tokens = False diff --git a/tests/model_executor/test_ep.py b/tests/model_executor/test_ep.py index 83972476be4..a795249ef55 100644 --- a/tests/model_executor/test_ep.py +++ b/tests/model_executor/test_ep.py @@ -21,7 +21,6 @@ from fastdeploy.config import MoEPhase from fastdeploy.model_executor.layers.moe import ep -from fastdeploy.model_executor.layers.moe.ep import EPRunner class FakeConfig: @@ -403,6 +402,14 @@ def fake_get_moe_scores(*_args, **_kwargs): monkeypatch.setattr(moe_module, "get_moe_scores", fake_get_moe_scores, raising=True) + runner = ep.EPPrefillRunner( + top_k=2, + hidden_size=4, + num_experts=2, + splitwise_role="prefill", + num_max_dispatch_tokens_per_rank=1, + ) + layer = SimpleNamespace( redundant_table_manger=None, topk_method="noaux_tc", @@ -416,7 +423,7 @@ def fake_get_moe_scores(*_args, **_kwargs): ) gate_out = paddle.randn([1, 4], dtype="float32") - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = runner.moe_select(layer, gate_out) assert list(topk_idx.shape) == [1, 1] assert list(topk_weights.shape) == [1, 1] assert paddle.allclose(topk_idx, paddle.to_tensor([[1]], dtype="int64")) @@ -433,6 +440,14 @@ def fake_redundant_topk_select(**_kwargs): monkeypatch.setattr(gpu_ops, "moe_redundant_topk_select", fake_redundant_topk_select, raising=True) + runner = ep.EPPrefillRunner( + top_k=2, + hidden_size=4, + num_experts=2, + splitwise_role="prefill", + num_max_dispatch_tokens_per_rank=1, + ) + class FakeRedundantTableManager: def get_ep_rank_to_expert_id_list_by_layer(self, _layer_idx): return [0], paddle.to_tensor([0], dtype="int64"), [1], [1] @@ -450,7 +465,7 @@ def get_ep_rank_to_expert_id_list_by_layer(self, _layer_idx): ) gate_out = paddle.randn([1, 4], dtype="float32") - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = runner.moe_select(layer, gate_out) assert list(topk_idx.shape) == [1, 1] assert list(topk_weights.shape) == [1, 1] assert paddle.allclose(topk_idx, paddle.to_tensor([[2]], dtype="int64")) @@ -467,6 +482,14 @@ def fake_topk_select(*_args, **_kwargs): monkeypatch.setattr(gpu_ops, "moe_topk_select", fake_topk_select, raising=True) + runner = ep.EPPrefillRunner( + top_k=2, + hidden_size=4, + num_experts=2, + splitwise_role="prefill", + num_max_dispatch_tokens_per_rank=1, + ) + layer = SimpleNamespace( redundant_table_manger=None, topk_method="aux", @@ -475,7 +498,7 @@ def fake_topk_select(*_args, **_kwargs): ) gate_out = paddle.randn([1, 4], dtype="float32") - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = runner.moe_select(layer, gate_out) assert list(topk_idx.shape) == [1, 1] assert list(topk_weights.shape) == [1, 1] assert paddle.allclose(topk_idx, paddle.to_tensor([[3]], dtype="int64")) From 93105300d77f89dc1f336cbb4c710f1dfc086ae7 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Tue, 9 Jun 2026 19:24:13 +0800 Subject: [PATCH 12/17] fix xpu test --- .../model_executor/layers/backends/xpu/moe/fused_moe.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py b/fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py index 4c1c7c233ef..6f313170a74 100644 --- a/fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py +++ b/fastdeploy/model_executor/layers/backends/xpu/moe/fused_moe.py @@ -22,7 +22,6 @@ from paddleformers.utils.log import logger from fastdeploy import envs -from fastdeploy.model_executor.layers.moe.ep import EPRunner from fastdeploy.model_executor.layers.moe.fused_moe_backend_base import MoEMethodBase from fastdeploy.model_executor.layers.quantization.weight_only import WeightOnlyConfig from fastdeploy.model_executor.layers.utils import get_tensor @@ -424,7 +423,7 @@ def apply_ep_prefill( """ gate_out = gate(x.cast("float32")) # 1. Select topk experts and weights - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out) # 2. Dynamic compute blockwise quantization scales if "a_tokenwise_int8" in self.xpu_moe_quant_type: @@ -519,7 +518,7 @@ def apply_ep_decode( gate_out = gate(x.cast("float32")) # 1. Select topk experts and weights - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out) # 2. EP Dispatch if "a_tokenwise_int8" in self.xpu_moe_quant_type: From 29802da72f3e0dc27118ba77df8d2a24dc08c384 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Tue, 9 Jun 2026 21:13:40 +0800 Subject: [PATCH 13/17] fix test --- .../model_executor/layers/moe/fused_moe_deepgemm_backend.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py index 21b384acd04..ff71d1fb9d7 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py @@ -355,7 +355,7 @@ def apply_ep_prefill( hidden_size = layer.hidden_size # 1. Select topk experts and weights - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out) if layer.routed_scaling_factor_learnable: safe_topk_indices = paddle.clip(topk_idx, min=0) @@ -685,7 +685,7 @@ def apply_ep_decode( gate_out = gate(x) gate_out = gate_out.cast("float32") # 1. Select topk experts and weights - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out) if layer.routed_scaling_factor_learnable: safe_topk_indices = paddle.clip(topk_idx, min=0) @@ -1318,7 +1318,7 @@ def apply_mega_moe(self, layer, x, gate, topk_ids_hookfunc, shared_experts, fc1_ gate_out = gate(x).cast("float32") # 1. Select topk experts and weights. - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = EPRunner.moe_select(None, layer, gate_out) buffer_capacity = self.mega_moe_buffer.x.shape[0] if num_tokens > buffer_capacity: From 561280ccd6e2ec9f865d5b43cfbc1f0100ebdf92 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Wed, 10 Jun 2026 14:34:21 +0800 Subject: [PATCH 14/17] fix test --- .../model_executor/layers/moe/fused_moe_deepgemm_backend.py | 2 +- fastdeploy/model_executor/layers/quantization/wfp4afp8.py | 2 +- fastdeploy/worker/worker_process.py | 1 - 3 files changed, 2 insertions(+), 3 deletions(-) diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py index ff71d1fb9d7..5552ef9a93b 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py @@ -1123,7 +1123,7 @@ def create_weights(self, layer: nn.Layer, **extra_weight_attrs): def init_ep(self, layer: nn.Layer) -> None: logger.info("Use MegaMoE backend") if layer.ep_size <= 1: - return + raise ValueError("Ep size must be greater than 1 when use MegaMoE backend. Please set --enable-expert-parallel") config = layer.fd_config splitwise_role = config.scheduler_config.splitwise_role diff --git a/fastdeploy/model_executor/layers/quantization/wfp4afp8.py b/fastdeploy/model_executor/layers/quantization/wfp4afp8.py index 6dabdba7a7e..4c8a462d476 100644 --- a/fastdeploy/model_executor/layers/quantization/wfp4afp8.py +++ b/fastdeploy/model_executor/layers/quantization/wfp4afp8.py @@ -52,7 +52,7 @@ def from_config(cls, config: dict) -> "WFP4AFP8Config": return cls(weight_scale_dict, act_scale_dict, is_permuted, is_quantized) def get_quant_method(self, layer) -> Optional[QuantMethodBase]: - logger.info("Currently only support DeepGEMMMegaMoE for wfp4afp8") + logger.debug("Currently only support DeepGEMMMegaMoE for wfp4afp8") if isinstance(layer, FusedMoE): from fastdeploy.model_executor.layers.moe.fused_moe_deepgemm_backend import ( DeepGemmMegaMoEMethod, diff --git a/fastdeploy/worker/worker_process.py b/fastdeploy/worker/worker_process.py index c7a7e7ea375..9b5cb67f731 100644 --- a/fastdeploy/worker/worker_process.py +++ b/fastdeploy/worker/worker_process.py @@ -868,7 +868,6 @@ def parse_args(): ) parser.add_argument( "--enable_mega_moe", - "--enable-mega-moe", action="store_true", dest="enable_mega_moe", help="enable MegaMoE wfp4afp8 for MoE and block_wise_fp8 for dense Linear", From 15b7a125e9794f77f86a2f503368dfa34f4441bb Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Wed, 10 Jun 2026 14:38:00 +0800 Subject: [PATCH 15/17] fix code style --- .../model_executor/layers/moe/fused_moe_deepgemm_backend.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py index 5552ef9a93b..a02acc88c1d 100644 --- a/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py +++ b/fastdeploy/model_executor/layers/moe/fused_moe_deepgemm_backend.py @@ -1123,7 +1123,9 @@ def create_weights(self, layer: nn.Layer, **extra_weight_attrs): def init_ep(self, layer: nn.Layer) -> None: logger.info("Use MegaMoE backend") if layer.ep_size <= 1: - raise ValueError("Ep size must be greater than 1 when use MegaMoE backend. Please set --enable-expert-parallel") + raise ValueError( + "Ep size must be greater than 1 when use MegaMoE backend. Please set --enable-expert-parallel" + ) config = layer.fd_config splitwise_role = config.scheduler_config.splitwise_role From f7bc99d4f3ad7798b40bbf267ea133a62ac80522 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Wed, 10 Jun 2026 14:44:52 +0800 Subject: [PATCH 16/17] fix typo --- fastdeploy/model_executor/layers/quantization/nvfp4.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/fastdeploy/model_executor/layers/quantization/nvfp4.py b/fastdeploy/model_executor/layers/quantization/nvfp4.py index 0e5722a343a..c5c5cfd95f9 100644 --- a/fastdeploy/model_executor/layers/quantization/nvfp4.py +++ b/fastdeploy/model_executor/layers/quantization/nvfp4.py @@ -673,7 +673,7 @@ def apply_ep_prefill( # 1. top experts and weights gate_out = gate(x.cast("float32")) - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = self.ep_prefill_runner.moe_select(layer, gate_out) hidden_size = x.shape[1] if topk_ids_hookfunc is not None: @@ -862,7 +862,7 @@ def apply_ep_decode( ) -> paddle.Tensor: gate_out = gate(x.cast("float32")) - topk_idx, topk_weights = EPRunner.moe_select(layer, gate_out) + topk_idx, topk_weights = self.ep_decoder_runner.moe_select(layer, gate_out) if topk_ids_hookfunc is not None: topk_ids_hookfunc(topk_ids=topk_idx) From 083a651ecaf5c18f40c5d5182a0ced1aebbac7c8 Mon Sep 17 00:00:00 2001 From: Wanglongzhi2001 <583087864@qq.com> Date: Wed, 10 Jun 2026 15:17:14 +0800 Subject: [PATCH 17/17] fix typo --- fastdeploy/model_executor/layers/quantization/nvfp4.py | 1 - 1 file changed, 1 deletion(-) diff --git a/fastdeploy/model_executor/layers/quantization/nvfp4.py b/fastdeploy/model_executor/layers/quantization/nvfp4.py index c5c5cfd95f9..627901ddf20 100644 --- a/fastdeploy/model_executor/layers/quantization/nvfp4.py +++ b/fastdeploy/model_executor/layers/quantization/nvfp4.py @@ -25,7 +25,6 @@ import fastdeploy from fastdeploy import envs from fastdeploy.model_executor.layers.moe import FusedMoE -from fastdeploy.model_executor.layers.moe.ep import EPRunner from fastdeploy.model_executor.layers.moe.fused_moe_backend_base import MoEMethodBase from fastdeploy.model_executor.utils import ( create_parameter_and_copy,