From e3a23457b798c4dc4034b389166a0f59e7626b33 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Tue, 3 Nov 2026 12:42:29 +0000 Subject: [PATCH] fix_deepep --- .../fused_moe/fused_moe_weight.py | 56 +- .../fused_moe/impl/deepgemm_impl.py | 122 +++-- .../deepep_expanded_layout_kernels.py | 138 +++++ .../grouped_fused_moe_ep/__init__.py | 44 ++ .../fused_moe/grouped_fused_moe_ep/common.py | 61 +++ .../deepgemm.py} | 479 +++++------------- .../grouped_fused_moe_ep/fp4_mega.py | 88 ++++ .../prefill_decode_sync.py | 261 ++++++++++ .../moe_silu_and_mul_mix_quant_ep.py | 174 +++++++ lightllm/distributed/communication_op.py | 45 +- .../layer_infer/transformer_layer_infer.py | 69 +-- .../layer_infer/transformer_layer_infer.py | 38 +- .../layer_infer/transformer_layer_infer.py | 64 ++- lightllm/server/api_start.py | 22 + lightllm/utils/envs_utils.py | 44 +- .../fused_moe/test_deepgemm_activation.py | 2 +- .../fused_moe/test_ep_reduce_decode_output.py | 36 ++ .../test_grouped_fused_moe_ep_sms.py | 11 + .../test_moe_silu_and_mul_mix_quant_ep.py | 41 ++ .../test_ue8m0_fp8_quantization.py | 12 +- unit_tests/server/test_api_start_deepep.py | 45 ++ unit_tests/utils/test_envs_utils.py | 66 +++ 22 files changed, 1408 insertions(+), 510 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/__init__.py create mode 100644 lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/common.py rename lightllm/common/basemodel/triton_kernel/fused_moe/{grouped_fused_moe_ep.py => grouped_fused_moe_ep/deepgemm.py} (51%) create mode 100644 lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/fp4_mega.py create mode 100644 lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/prefill_decode_sync.py create mode 100644 unit_tests/common/fused_moe/test_ep_reduce_decode_output.py create mode 100644 unit_tests/common/fused_moe/test_grouped_fused_moe_ep_sms.py create mode 100644 unit_tests/server/test_api_start_deepep.py diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 72311d22e4..8279b6d65a 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -190,16 +190,17 @@ def experts_with_topk( clamp_up_add_one=clamp_up_add_one, ) - def low_latency_dispatch( + def decode_dispatch( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, ): - assert self.enable_ep_moe, "low_latency_dispatch is only supported when enable_ep_moe is True" - return self.fuse_moe_impl.low_latency_dispatch( + assert self.enable_ep_moe, "decode_dispatch requires enable_ep_moe" + return self.fuse_moe_impl.decode_dispatch( hidden_states=hidden_states, router_logits=router_logits, e_score_correction_bias=self.e_score_correction_bias, + w13=self.w13, use_grouped_topk=self.use_grouped_topk, num_experts_per_tok=self.num_experts_per_tok, norm_topk_prob=self.norm_topk_prob, @@ -208,15 +209,16 @@ def low_latency_dispatch( scoring_func=self.scoring_func, ) - def low_latency_dispatch_with_topk( + def decode_dispatch_with_topk( self, hidden_states: torch.Tensor, topk_idx: torch.Tensor, topk_weights: torch.Tensor, ): - assert self.enable_ep_moe, "low_latency_dispatch_with_topk is only supported when enable_ep_moe is True" - return self.fuse_moe_impl.low_latency_dispatch_with_topk( + assert self.enable_ep_moe, "decode_dispatch_with_topk requires enable_ep_moe" + return self.fuse_moe_impl.decode_dispatch_with_topk( hidden_states=hidden_states, + w13=self.w13, topk_idx=topk_idx, topk_weights=topk_weights, ) @@ -244,37 +246,37 @@ def select_experts_and_quant_input( scoring_func=self.scoring_func, ) - def dispatch( + def prefill_dispatch( self, qinput_tensor: Tuple[torch.Tensor], topk_idx: torch.Tensor, topk_weights: torch.Tensor, overlap_event: Optional[Any] = None, ): - assert self.enable_ep_moe, "dispatch is only supported when enable_ep_moe is True" - return self.fuse_moe_impl.dispatch( + assert self.enable_ep_moe, "prefill_dispatch is only supported when enable_ep_moe is True" + return self.fuse_moe_impl.prefill_dispatch( qinput_tensor=qinput_tensor, topk_idx=topk_idx, topk_weights=topk_weights, overlap_event=overlap_event, ) - def masked_group_gemm( + def decode_masked_group_gemm( self, recv_x: Tuple[torch.Tensor], - masked_m: torch.Tensor, + ep_handle: Any, dtype: torch.dtype, expected_m: int, alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, ): - assert self.enable_ep_moe, "masked_group_gemm is only supported when enable_ep_moe is True" - return self.fuse_moe_impl.masked_group_gemm( + assert self.enable_ep_moe, "decode_masked_group_gemm is only supported when enable_ep_moe is True" + return self.fuse_moe_impl.decode_masked_group_gemm( recv_x=recv_x, w13=self.w13, w2=self.w2, - masked_m=masked_m, + ep_handle=ep_handle, dtype=dtype, expected_m=expected_m, alpha=alpha, @@ -313,29 +315,27 @@ def prefilled_group_gemm( clamp_up_add_one=clamp_up_add_one, ) - def low_latency_combine( + def decode_combine( self, - gemm_out_b: torch.Tensor, - topk_idx: torch.Tensor, - topk_weights: torch.Tensor, - handle: Any, + expert_output: torch.Tensor, + ep_handle: Any, + recv_topk_weights: torch.Tensor, ): - assert self.enable_ep_moe, "low_latency_combine is only supported when enable_ep_moe is True" - return self.fuse_moe_impl.low_latency_combine( - gemm_out_b=gemm_out_b, - topk_idx=topk_idx, - topk_weights=topk_weights, - handle=handle, + assert self.enable_ep_moe, "decode_combine is only supported when enable_ep_moe is True" + return self.fuse_moe_impl.decode_combine( + expert_output=expert_output, + ep_handle=ep_handle, + recv_topk_weights=recv_topk_weights, ) - def combine( + def prefill_combine( self, gemm_out_b: torch.Tensor, handle: Any, overlap_event: Optional[Any] = None, ): - assert self.enable_ep_moe, "combine is only supported when enable_ep_moe is True" - return self.fuse_moe_impl.combine( + assert self.enable_ep_moe, "prefill_combine is only supported when enable_ep_moe is True" + return self.fuse_moe_impl.prefill_combine( gemm_out_b=gemm_out_b, handle=handle, overlap_event=overlap_event, diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 8e4220e4b2..62dbe29092 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -6,14 +6,17 @@ from lightllm.utils.envs_utils import ( get_deepep_num_max_dispatch_tokens_per_rank_prefill, get_deepep_num_max_dispatch_tokens_per_rank_decode, - get_env_use_ue8m0_scales, ) from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( fused_experts, get_ep_num_sms, - masked_group_gemm, + decode_masked_group_gemm, chunked_expanded_moe_forward, quantize_fused_experts_input, + get_mk_alignment_for_contiguous_layout, +) +from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_expanded_layout_kernels import ( + ep_reduce_decode_output, ) from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType @@ -97,11 +100,12 @@ def _fused_experts( ) return output - def low_latency_dispatch( + def decode_dispatch( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, e_score_correction_bias: torch.Tensor, + w13: WeightPack, use_grouped_topk: bool, num_experts_per_tok: int, norm_topk_prob: bool, @@ -120,33 +124,45 @@ def low_latency_dispatch( num_expert_group=n_group, scoring_func=scoring_func, ) - - return self.low_latency_dispatch_with_topk( + return self.decode_dispatch_with_topk( hidden_states=hidden_states, + w13=w13, topk_idx=topk_idx, topk_weights=topk_weights, ) - def low_latency_dispatch_with_topk( + def decode_dispatch_with_topk( self, hidden_states: torch.Tensor, + w13: WeightPack, topk_idx: torch.Tensor, topk_weights: torch.Tensor, ): topk_idx = topk_idx.to(torch.long) - num_max_dispatch_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank_decode() - use_fp8_w8a8 = self.quant_method.method_name != "none" - recv_x, masked_m, handle, event, hook = dist_group_manager.ep_low_latency_buffer.low_latency_dispatch( + qinput_tensor = quantize_fused_experts_input(hidden_states, w13, self.quant_method) + recv_x, _, recv_topk_weights, ep_handle, event = dist_group_manager.ep_buffer.dispatch( + qinput_tensor, topk_idx=topk_idx, - x=hidden_states, - num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank, + topk_weights=topk_weights, num_experts=self.total_expert_num_contain_redundancy, - use_fp8=use_fp8_w8a8, - round_scale=get_env_use_ue8m0_scales(), - async_finish=False, - return_recv_hook=True, + num_max_tokens_per_rank=get_deepep_num_max_dispatch_tokens_per_rank_decode(), + expert_alignment=get_mk_alignment_for_contiguous_layout(), + num_sms=get_ep_num_sms(overlap_with_compute=True), + async_with_compute_stream=True, + # Decode 不传 previous_event，输出直接归属计算流；DeepEP 会为 + # 通信流的异步访问登记 allocator 生命周期。 + allocate_on_comm_stream=False, + do_cpu_sync=False, + do_handle_copy=False, + do_expand=True, + do_zero_padding=True, + use_tma_aligned_col_major_sf=True, ) - return recv_x, masked_m, topk_idx, topk_weights, handle, hook + + def hook(): + event.current_stream_wait() + + return recv_x, ep_handle, recv_topk_weights, hook def quantize_dispatch_input(self, hidden_states: torch.Tensor, w13: WeightPack): return quantize_fused_experts_input(hidden_states, w13, self.quant_method) @@ -178,7 +194,7 @@ def select_experts_and_quant_input( qinput_tensor = quantize_fused_experts_input(hidden_states, w13, self.quant_method) return topk_weights, topk_idx.to(torch.long), qinput_tensor - def dispatch( + def prefill_dispatch( self, qinput_tensor: Tuple[torch.Tensor], topk_idx: torch.Tensor, @@ -193,8 +209,8 @@ def dispatch( topk_weights=topk_weights, num_experts=self.total_expert_num_contain_redundancy, num_max_tokens_per_rank=num_max_tokens_per_rank, - expert_alignment=128, - num_sms=get_ep_num_sms(), + expert_alignment=get_mk_alignment_for_contiguous_layout(), + num_sms=get_ep_num_sms(overlap_with_compute=True), previous_event=overlap_event, async_with_compute_stream=True, allocate_on_comm_stream=True, @@ -209,12 +225,12 @@ def hook(): return recv_x, recv_topk_idx, recv_topk_weights, handle.num_recv_tokens_per_expert_list, handle, hook - def masked_group_gemm( + def decode_masked_group_gemm( self, recv_x: Tuple[torch.Tensor], w13: WeightPack, w2: WeightPack, - masked_m: torch.Tensor, + ep_handle: Any, dtype: torch.dtype, expected_m: int, alpha: Optional[float] = None, @@ -223,14 +239,15 @@ def masked_group_gemm( ): w13_weight, w13_scale = w13.weight, w13.weight_scale w2_weight, w2_scale = w2.weight, w2.weight_scale - return masked_group_gemm( - recv_x, - masked_m, - dtype, - w13_weight, - w13_scale, - w2_weight, - w2_scale, + return decode_masked_group_gemm( + recv_x=recv_x, + expert_token_psum=ep_handle.psum_num_recv_tokens_per_expert, + expert_alignment=ep_handle.expert_alignment, + dtype=dtype, + w1=w13_weight, + w1_scale=w13_scale, + w2=w2_weight, + w2_scale=w2_scale, expected_m=expected_m, alpha=alpha, limit=limit, @@ -297,33 +314,60 @@ def prefilled_group_gemm( clamp_up_add_one=clamp_up_add_one, ) _gemm_out_a, _silu_out = None, None + # 只为占显存主体的 recv_x 登记计算流；较小的路由 metadata 继续由 + # handle 和当前调用栈维持生命周期。 + compute_stream = torch.cuda.current_stream() + recv_x[0].record_stream(compute_stream) + recv_x[1].record_stream(compute_stream) del recv_x return gather_out - def low_latency_combine( + def decode_combine( self, - gemm_out_b: torch.Tensor, - topk_idx: torch.Tensor, - topk_weights: torch.Tensor, - handle: Any, + expert_output: torch.Tensor, + ep_handle: Any, + recv_topk_weights: torch.Tensor, ): - combined_x, event_overlap, hook = dist_group_manager.ep_low_latency_buffer.low_latency_combine( - gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=True + # 阶段 1：把按 expert 展开的输出归约回去重接收 token 布局。 + # psum_num_recv_tokens_per_scaleup_rank 的最后一项是去重接收 token 数， + # 不是按 expert 展开后的有效行数。 + dense_output, compact_metadata = ep_reduce_decode_output( + expert_output=expert_output, + route_weights=recv_topk_weights, + recv_src_metadata=ep_handle.recv_src_metadata, + num_valid_recv_tokens=ep_handle.psum_num_recv_tokens_per_scaleup_rank[-1:], ) + ep_handle.recv_src_metadata = compact_metadata + + # 阶段 2：compact metadata 的第一个 top-k 槽位指向同序 dense row， + # 其余槽位置为 -1。异步 combine 返回 event，由调用方在消费输出前等待。 + combined_x, _, event = dist_group_manager.ep_buffer.combine( + dense_output, + ep_handle, + topk_weights=None, + num_sms=get_ep_num_sms(overlap_with_compute=True), + async_with_compute_stream=True, + # Decode combine 没有 previous_event，输出归属计算流；通信流的异步 + # 生命周期由 DeepEP 内部统一登记。 + allocate_on_comm_stream=False, + ) + + def hook(): + event.current_stream_wait() + return combined_x, hook - def combine( + def prefill_combine( self, gemm_out_b: torch.Tensor, handle: Any, overlap_event: Optional[Any] = None, ): - # normal combine combined_x, _, event = dist_group_manager.ep_buffer.combine( gemm_out_b, handle, topk_weights=None, - num_sms=get_ep_num_sms(), + num_sms=get_ep_num_sms(overlap_with_compute=True), previous_event=overlap_event, async_with_compute_stream=True, allocate_on_comm_stream=True, diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_expanded_layout_kernels.py b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_expanded_layout_kernels.py index ab54ea5b79..5f0267b644 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_expanded_layout_kernels.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_expanded_layout_kernels.py @@ -247,6 +247,144 @@ def ep_gather_chunk( ) +@triton.jit +def _ep_reduce_decode_output_kernel( + expert_output, + expert_output_stride_m, + expert_output_stride_k, + route_weights, + recv_src_metadata, + metadata_stride_m, + metadata_stride_k, + num_valid_recv_tokens_ptr, + dense_output, + dense_output_stride_m, + dense_output_stride_k, + compact_metadata, + compact_metadata_stride_m, + compact_metadata_stride_k, + hidden_size, + TOPK: tl.constexpr, + NUM_STAGE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_METADATA: tl.constexpr, + NEED_HIDDEN_MASK: tl.constexpr, +): + hidden_block_id = tl.program_id(0) + first_recv_token_id = tl.program_id(1) + recv_token_step = tl.num_programs(1) + num_valid_recv_tokens = tl.load(num_valid_recv_tokens_ptr) + + # 行号可能乘以较大的 hidden-size stride。先将行 stride 和后续使用的 + # 行索引提升为 int64，避免地址偏移在 int32 乘法阶段溢出；列偏移只在 + # 单行的 hidden/top-k 范围内变化，不需要额外提升。 + expert_output_stride_m = expert_output_stride_m.to(tl.int64) + metadata_stride_m = metadata_stride_m.to(tl.int64) + dense_output_stride_m = dense_output_stride_m.to(tl.int64) + compact_metadata_stride_m = compact_metadata_stride_m.to(tl.int64) + + hidden_offsets = hidden_block_id * BLOCK_D + tl.arange(0, BLOCK_D) + if NEED_HIDDEN_MASK: + hidden_mask = hidden_offsets < hidden_size + other = 0.0 + else: + hidden_mask = None + other = None + + for recv_token_id in range(first_recv_token_id, num_valid_recv_tokens, recv_token_step): + recv_token_row = recv_token_id.to(tl.int64) + accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) + for topk_id in tl.range(0, TOPK, num_stages=NUM_STAGE): + expert_row = tl.load( + recv_src_metadata + recv_token_row * metadata_stride_m + (topk_id + 2) * metadata_stride_k + ) + # -1 表示该 token 的这个 top-k expert 不在当前接收 rank。 + # 只为有效的本地 expert 行读取输出和权重。 + if expert_row >= 0: + expert_output_ptrs = ( + expert_output + + expert_row.to(tl.int64) * expert_output_stride_m + + hidden_offsets * expert_output_stride_k + ) + value = tl.load(expert_output_ptrs, mask=hidden_mask, other=other) + weight = tl.load(route_weights + expert_row) + accumulator += value.to(tl.float32) * weight + + dense_output_ptrs = ( + dense_output + recv_token_row * dense_output_stride_m + hidden_offsets * dense_output_stride_k + ) + tl.store(dense_output_ptrs, accumulator, mask=hidden_mask) + + if hidden_block_id == 0: + metadata_offsets = tl.arange(0, BLOCK_METADATA) + source_header = tl.load( + recv_src_metadata + recv_token_row * metadata_stride_m + metadata_offsets * metadata_stride_k, + mask=metadata_offsets < 2, + other=0, + ) + compact_values = tl.where( + metadata_offsets < 2, + source_header, + tl.where(metadata_offsets == 2, recv_token_id, -1), + ) + tl.store( + compact_metadata + + recv_token_row * compact_metadata_stride_m + + metadata_offsets * compact_metadata_stride_k, + compact_values, + mask=metadata_offsets < TOPK + 2, + ) + + +@torch.no_grad() +def ep_reduce_decode_output( + expert_output: torch.Tensor, # [num_expanded_rows, hidden_size] + route_weights: torch.Tensor, # [num_expanded_rows] + recv_src_metadata: torch.Tensor, # [num_recv_tokens, topk + 2] + num_valid_recv_tokens: torch.Tensor, # GPU scalar +) -> tuple[torch.Tensor, torch.Tensor]: + """将 expanded expert 行加权归约，并构造 dense combine metadata。""" + num_recv_tokens = recv_src_metadata.shape[0] + hidden_size = expert_output.shape[1] + dense_output = torch.empty( + (num_recv_tokens, hidden_size), + dtype=expert_output.dtype, + device=expert_output.device, + ) + compact_metadata = torch.empty_like(recv_src_metadata) + + block_d = 1024 + num_stage = 1 + # num_recv_tokens 是接收缓冲区容量，可能大于实际有效 token 数。1024 在 + # 大并发下能提供更充分的行并行度，小并发下额外空 program 的开销很小。 + token_programs = min(num_recv_tokens, 1024) + grid = (triton.cdiv(hidden_size, block_d), token_programs) + _ep_reduce_decode_output_kernel[grid]( + expert_output=expert_output, + expert_output_stride_m=expert_output.stride(0), + expert_output_stride_k=expert_output.stride(1), + route_weights=route_weights, + recv_src_metadata=recv_src_metadata, + metadata_stride_m=recv_src_metadata.stride(0), + metadata_stride_k=recv_src_metadata.stride(1), + num_valid_recv_tokens_ptr=num_valid_recv_tokens, + dense_output=dense_output, + dense_output_stride_m=dense_output.stride(0), + dense_output_stride_k=dense_output.stride(1), + compact_metadata=compact_metadata, + compact_metadata_stride_m=compact_metadata.stride(0), + compact_metadata_stride_k=compact_metadata.stride(1), + hidden_size=hidden_size, + TOPK=recv_src_metadata.shape[1] - 2, + NUM_STAGE=num_stage, + BLOCK_D=block_d, + BLOCK_METADATA=triton.next_power_of_2(recv_src_metadata.shape[1]), + NEED_HIDDEN_MASK=hidden_size % block_d != 0, + num_warps=2, + ) + return dense_output, compact_metadata + + @triton.jit def _ep_compact_metadata_kernel( recv_src_metadata, diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/__init__.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/__init__.py new file mode 100644 index 0000000000..9d959d0dc2 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/__init__.py @@ -0,0 +1,44 @@ +"""Fused MoE kernel.""" + +from lightllm.distributed import dist_group_manager + +from .common import ( + SUPPORTED_EP_EXPERT_DTYPES, + check_ep_expert_dtype, + get_ep_num_sms, + quantize_fused_experts_input, +) +from .deepgemm import ( + HAS_DEEPGEMM, + _get_max_chunk_rows, + chunked_expanded_moe_forward, + decode_masked_group_gemm, + deepgemm_grouped_fp8_fp4_nt_contiguous, + deepgemm_grouped_fp8_nt_contiguous, + get_mk_alignment_for_contiguous_layout, + masked_group_gemm, + set_mk_alignment_for_contiguous_layout, +) +from .fp4_mega import mega_moe_impl, use_sm100_mega_moe +from .prefill_decode_sync import fused_experts, fused_experts_impl + +__all__ = [ + "HAS_DEEPGEMM", + "SUPPORTED_EP_EXPERT_DTYPES", + "_get_max_chunk_rows", + "check_ep_expert_dtype", + "chunked_expanded_moe_forward", + "decode_masked_group_gemm", + "deepgemm_grouped_fp8_fp4_nt_contiguous", + "deepgemm_grouped_fp8_nt_contiguous", + "dist_group_manager", + "fused_experts", + "fused_experts_impl", + "get_ep_num_sms", + "get_mk_alignment_for_contiguous_layout", + "masked_group_gemm", + "mega_moe_impl", + "quantize_fused_experts_input", + "set_mk_alignment_for_contiguous_layout", + "use_sm100_mega_moe", +] diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/common.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/common.py new file mode 100644 index 0000000000..19c2a275f9 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/common.py @@ -0,0 +1,61 @@ +import torch +from typing import Any + +from lightllm.distributed import dist_group_manager +from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import ( + per_token_group_quant_fp8, +) +from lightllm.utils.device_utils import is_sm100_gpu +from lightllm.utils.envs_utils import get_env_use_ue8m0_scales + +from .fp4_mega import use_sm100_mega_moe + +SUPPORTED_EP_EXPERT_DTYPES = ("fp8w8a8-b128-deepgemm", "fp4fp8-b32-deepgemm") + + +def get_ep_num_sms(overlap_with_compute: bool) -> int: + """根据通信是否与计算重叠，返回对应的 DeepEP SM 配额。""" + attribute = "ep_num_sms" if overlap_with_compute else "ep_non_overlap_num_sms" + return getattr(dist_group_manager, attribute, None) or 0 + + +def check_ep_expert_dtype(quant_method: Any): + expert_dtype = getattr(quant_method, "method_name", None) + if expert_dtype not in SUPPORTED_EP_EXPERT_DTYPES: + raise ValueError( + "EP MoE requires --expert_dtype to be one of ['fp8', 'fp4'], " + f"but the resolved fused_moe quant method is `{expert_dtype}`. " + "Please start with --expert_dtype fp8 or --expert_dtype fp4. " + "Note that --expert_dtype fp4 is only supported on SM100 GPUs." + ) + if expert_dtype == "fp4fp8-b32-deepgemm" and not is_sm100_gpu(): + raise RuntimeError( + "--expert_dtype fp4 requires an SM100 GPU for EP MoE; " "please use --expert_dtype fp8 on non-SM100 GPUs." + ) + + +def quantize_fused_experts_input( + hidden_states: torch.Tensor, + w13: Any, + quant_method: Any, +): + check_ep_expert_dtype(quant_method) + if use_sm100_mega_moe(quant_method): + from deep_gemm.utils import per_token_cast_to_fp8 + + return per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=quant_method.block_size, + use_packed_ue8m0=True, + ) + + block_size_k = quant_method.block_size + assert block_size_k == 128, "block_size_k must be 128" + # DeepEP keeps FP32 scales; UE8M0 only rounds their values. + return per_token_group_quant_fp8( + hidden_states, + block_size_k, + dtype=w13.weight.dtype, + use_ue8m0_scales=get_env_use_ue8m0_scales(), + ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/deepgemm.py similarity index 51% rename from lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py rename to lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/deepgemm.py index 61cd81556a..eea00e1bc2 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/deepgemm.py @@ -1,14 +1,12 @@ -"""Fused MoE kernel.""" - import torch import triton -import triton.language as tl -from typing import Any, Callable, Dict, List, Optional, Tuple -from lightllm.distributed import dist_group_manager +from typing import List, Optional, Tuple + from lightllm.utils.log_utils import init_logger from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul_mix_quant_ep import ( silu_and_mul_masked_post_quant_fwd, + silu_and_mul_psum_post_quant_fwd, ) from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import ( per_token_group_quant_fp8, @@ -19,53 +17,152 @@ ep_gather_chunk, ep_zero_padding, ) -from lightllm.utils.envs_utils import ( - get_deepep_num_max_dispatch_tokens_per_rank_prefill, - get_deepep_num_max_dispatch_tokens_per_rank_decode, - get_env_use_ue8m0_scales, -) +from lightllm.utils.envs_utils import get_env_use_ue8m0_scales from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType -from lightllm.utils.device_utils import is_sm100_gpu from lightllm.utils.sgl_utils import HAS_SGL_KERNEL from lightllm.utils.tensor_buffer_manager import TensorBufferManager logger = init_logger(__name__) -_MEGA_MOE_STATES: Dict[Tuple[int, int, int, int], Dict[str, Any]] = {} -SUPPORTED_EP_EXPERT_DTYPES = ("fp8w8a8-b128-deepgemm", "fp4fp8-b32-deepgemm") - +deep_gemm = None try: - from deep_ep import Buffer, EventOverlap import deep_gemm HAS_DEEPGEMM = True -except: +except Exception: logger.warning("no deepep or deep_gemm") HAS_DEEPGEMM = False -def get_ep_num_sms() -> int: - return getattr(dist_group_manager, "ep_num_sms", None) or 0 +def get_mk_alignment_for_contiguous_layout() -> int: + """返回 DeepGEMM contiguous grouped GEMM 使用的 M 对齐值。""" + assert HAS_DEEPGEMM + return deep_gemm.get_mk_alignment_for_contiguous_layout() -def use_sm100_mega_moe(quant_method: Any) -> bool: - return is_sm100_gpu() and quant_method.method_name == "fp4fp8-b32-deepgemm" +def set_mk_alignment_for_contiguous_layout(alignment: int) -> None: + """设置 DeepGEMM contiguous grouped GEMM 使用的 M 对齐值。""" + assert HAS_DEEPGEMM + deep_gemm.set_mk_alignment_for_contiguous_layout(alignment) -def check_ep_expert_dtype(quant_method: Any): - expert_dtype = getattr(quant_method, "method_name", None) - if expert_dtype not in SUPPORTED_EP_EXPERT_DTYPES: - raise ValueError( - "EP MoE requires --expert_dtype to be one of ['fp8', 'fp4'], " - f"but the resolved fused_moe quant method is `{expert_dtype}`. " - "Please start with --expert_dtype fp8 or --expert_dtype fp4. " - "Note that --expert_dtype fp4 is only supported on SM100 GPUs." +def deepgemm_grouped_fp8_nt_contiguous( + input_tuple: Tuple[torch.Tensor, torch.Tensor], + w_tuple: Tuple[torch.Tensor, torch.Tensor], + out: torch.Tensor, + m_indices: torch.Tensor, + use_psum_layout: bool = False, + expected_m_for_psum_layout: Optional[int] = None, +): + """调用不同版本的 DeepGEMM contiguous grouped GEMM。 + + ``m_indices`` 在两种布局下的含义不同： + + - 普通布局：shape 为 ``[M]``，逐行记录输入使用的 expert ID。 + - psum 布局：shape 为 ``[expert_num]``，第 i 项记录 expert i 最后一个 + 有效行的后一位置。expert 0 的起点为 0；其余 expert 的起点是前一个 + expert 的结束位置向上对齐到 DeepGEMM M 对齐边界后的结果。 + """ + if not HAS_DEEPGEMM: + raise RuntimeError("deep_gemm does not provide grouped_gemm_fp8 NT contiguous GEMM kernel in this version") + + if hasattr(deep_gemm, "m_grouped_gemm_fp8_fp8_bf16_nt_contiguous"): + if not use_psum_layout: + return deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous( + lhs=input_tuple, + rhs=w_tuple, + out=out, + m_indices=m_indices, + ) + return deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous( + lhs=input_tuple, + rhs=w_tuple, + out=out, + m_indices=m_indices, + use_psum_layout=True, + expected_m_for_psum_layout=expected_m_for_psum_layout, ) - if expert_dtype == "fp4fp8-b32-deepgemm" and not is_sm100_gpu(): - raise RuntimeError( - "--expert_dtype fp4 requires an SM100 GPU for EP MoE; " "please use --expert_dtype fp8 on non-SM100 GPUs." + + if hasattr(deep_gemm, "m_grouped_fp8_gemm_nt_contiguous"): + return deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + a=input_tuple, + b=w_tuple, + d=out, + grouped_layout=m_indices, + use_psum_layout=use_psum_layout, + expected_m_for_psum_layout=expected_m_for_psum_layout, ) + raise RuntimeError("deep_gemm does not provide grouped_gemm_fp8 NT contiguous GEMM kernel in this version") + + +def decode_masked_group_gemm( + recv_x: Tuple[torch.Tensor, torch.Tensor], + expert_token_psum: torch.Tensor, + expert_alignment: int, + dtype: torch.dtype, + w1: torch.Tensor, + w1_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + expected_m: int, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, +): + """在 ElasticBuffer 的 expert 连续布局上执行两次 decode grouped GEMM。 + + ``recv_x`` 的容量按最坏情况固定，便于 CUDA Graph 重放；GPU 上的 + ``expert_token_psum`` 给出每个 expert 的有效结束位置，DeepGEMM 和中间 + SwiGLU 只处理这些位置覆盖的行，不需要把计数同步回 CPU。 + """ + assert (limit is None and alpha is None) or (limit is not None and alpha is not None) + max_rows, hidden_size = recv_x[0].shape + intermediate_twice = w1.shape[1] + intermediate_size = intermediate_twice // 2 + block_size = 128 + + gemm_out_a = torch.empty((max_rows, intermediate_twice), device=recv_x[0].device, dtype=dtype) + deepgemm_grouped_fp8_nt_contiguous( + recv_x, + (w1, w1_scale), + gemm_out_a, + expert_token_psum, + use_psum_layout=True, + expected_m_for_psum_layout=expected_m, + ) + + qsilu_out = torch.empty((max_rows, intermediate_size), dtype=w1.dtype, device=recv_x[0].device) + scale_storage = torch.empty( + (intermediate_size // block_size, triton.cdiv(max_rows, 4) * 4), + dtype=torch.float32, + device=recv_x[0].device, + ) + qsilu_out_scale = scale_storage[:, :max_rows].T + silu_and_mul_psum_post_quant_fwd( + input=gemm_out_a, + output=qsilu_out, + output_scale=qsilu_out_scale, + expert_token_psum=expert_token_psum, + expert_alignment=expert_alignment, + quant_group_size=block_size, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + ) + del gemm_out_a + + gemm_out_b = torch.empty((max_rows, hidden_size), device=recv_x[0].device, dtype=dtype) + deepgemm_grouped_fp8_nt_contiguous( + (qsilu_out, qsilu_out_scale), + (w2, w2_scale), + gemm_out_b, + expert_token_psum, + use_psum_layout=True, + expected_m_for_psum_layout=expected_m, + ) + return gemm_out_b + def masked_group_gemm( recv_x: Tuple[torch.Tensor, torch.Tensor], @@ -106,322 +203,6 @@ def masked_group_gemm( return gemm_out_b -def _get_mega_moe_cache_state(w13: Any, w2: Any): - state_key = ( - w13.weight.data_ptr(), - w13.weight_scale.data_ptr(), - w2.weight.data_ptr(), - w2.weight_scale.data_ptr(), - ) - return _MEGA_MOE_STATES.setdefault(state_key, {}) - - -def _get_mega_moe_weights(w13: Any, w2: Any, state: Dict[str, Any]): - if "weight_cache" not in state: - state["weight_cache"] = deep_gemm.transform_weights_for_mega_moe( - (w13.weight, w13.weight_scale), - (w2.weight, w2.weight_scale), - ) - return state["weight_cache"] - - -def _get_mega_moe_cumulative_stats(num_local_experts: int, device: torch.device, state: Dict[str, Any]): - stats = state.get("stats") - if stats is None or stats.numel() != num_local_experts or stats.device != device: - stats = torch.zeros((num_local_experts,), device=device, dtype=torch.int32) - state["stats"] = stats - return stats - - -def mega_moe_impl( - hidden_states: torch.Tensor, - w13: Any, - w2: Any, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - quant_method: Any, -): - if not (HAS_DEEPGEMM and hasattr(deep_gemm, "fp8_fp4_mega_moe")): - raise RuntimeError("deep_gemm does not provide fp8-fp4 Mega MoE kernel") - - from deep_gemm.utils import per_token_cast_to_fp8 - - buffer = getattr(dist_group_manager, "ep_mega_moe_buffer", None) - if buffer is None: - raise RuntimeError("SM100 Mega MoE requires dist_group_manager.ep_mega_moe_buffer to be initialized") - - num_tokens = hidden_states.shape[0] - if num_tokens > buffer.num_max_tokens_per_rank: - raise RuntimeError( - f"Mega MoE got {num_tokens} tokens, exceeding num_max_tokens_per_rank={buffer.num_max_tokens_per_rank}" - ) - - qinput_tensor = per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=quant_method.block_size, - use_packed_ue8m0=True, - ) - state = _get_mega_moe_cache_state(w13, w2) - l1_weights, l2_weights = _get_mega_moe_weights(w13, w2, state) - stats = _get_mega_moe_cumulative_stats(w13.weight.shape[0], hidden_states.device, state) - buffer.x[:num_tokens].copy_(qinput_tensor[0]) - buffer.x_sf[:num_tokens].copy_(qinput_tensor[1]) - buffer.topk_idx[:num_tokens].copy_(topk_ids) - buffer.topk_weights[:num_tokens].copy_(topk_weights) - - output = torch.empty_like(hidden_states) - deep_gemm.fp8_fp4_mega_moe( - output, - l1_weights, - l2_weights, - buffer, - cumulative_local_expert_recv_stats=stats, - ) - return output - - -def quantize_fused_experts_input( - hidden_states: torch.Tensor, - w13: Any, - quant_method: Any, -): - check_ep_expert_dtype(quant_method) - if use_sm100_mega_moe(quant_method): - from deep_gemm.utils import per_token_cast_to_fp8 - - return per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=quant_method.block_size, - use_packed_ue8m0=True, - ) - - block_size_k = quant_method.block_size - assert block_size_k == 128, "block_size_k must be 128" - # DeepEP keeps FP32 scales; UE8M0 only rounds their values. - return per_token_group_quant_fp8( - hidden_states, - block_size_k, - dtype=w13.weight.dtype, - use_ue8m0_scales=get_env_use_ue8m0_scales(), - ) - - -def fused_experts( - hidden_states: torch.Tensor, - w13: Any, - w2: Any, - topk_weights: torch.Tensor, - topk_idx: torch.Tensor, - num_experts: int, - quant_method: Any, - is_prefill: Optional[bool], - previous_event: Optional[Any] = None, - alpha: Optional[float] = None, - limit: Optional[float] = None, - clamp_up_add_one: bool = True, -): - assert (limit is None and alpha is None) or (limit is not None and alpha is not None) - check_ep_expert_dtype(quant_method) - if use_sm100_mega_moe(quant_method): - if limit is not None: - raise NotImplementedError("FP4 Mega MoE does not support clamped SwiGLU") - return mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method) - - buffer = dist_group_manager.ep_buffer if is_prefill else dist_group_manager.ep_low_latency_buffer - return fused_experts_impl( - hidden_states=hidden_states, - w1=w13.weight, - w2=w2.weight, - topk_weights=topk_weights, - topk_idx=topk_idx, - num_experts=num_experts, - block_size_k=quant_method.block_size, - buffer=buffer, - is_prefill=is_prefill, - use_fp8_w8a8=True, - use_fp8_all2all=True, - use_int8_w8a16=False, - w1_scale=w13.weight_scale, - w2_scale=w2.weight_scale, - previous_event=previous_event, - alpha=alpha, - limit=limit, - clamp_up_add_one=clamp_up_add_one, - ) - - -def fused_experts_impl( - hidden_states: torch.Tensor, # [M, K] - w1: torch.Tensor, # [group, N, K] - w2: torch.Tensor, # [group, K, N/2] - topk_weights: torch.Tensor, # [M, topk] - topk_idx: torch.Tensor, # [M, topk] - num_experts: int, - buffer: Any, - is_prefill: bool, - use_fp8_w8a8: bool = False, - use_fp8_all2all: bool = False, - use_int8_w8a16: bool = False, - w1_scale: Optional[torch.Tensor] = None, - w2_scale: Optional[torch.Tensor] = None, - previous_event: Optional[Any] = None, - alpha: Optional[float] = None, - limit: Optional[float] = None, - clamp_up_add_one: bool = True, - block_size_k: int = 128, -): - # Check constraints. - assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" - assert topk_weights.shape == topk_idx.shape, "topk shape mismatch" - assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" - assert w1.is_contiguous(), "Expert weights1 must be contiguous" - assert w2.is_contiguous(), "Expert weights2 must be contiguous" - assert hidden_states.dtype in [torch.float32, torch.float16, torch.bfloat16] - - # qaunt hidden_states - assert use_fp8_w8a8 and use_fp8_all2all, "use_fp8_w8a8 and use_fp8_all2all must be True" - - assert block_size_k == 128, "block_size_k must be 128" - - combined_x = None - if is_prefill: - qinput_tensor, input_scale = per_token_group_quant_fp8( - hidden_states, block_size_k, dtype=w1.dtype, use_ue8m0_scales=get_env_use_ue8m0_scales() - ) - allocate_on_comm_stream = previous_event is not None - # Expanded dispatch directly produces expert-contiguous, alignment-padded inputs: - # recv_x[0]: [num_expanded_tokens, hidden] - # recv_x[1]: [num_expanded_tokens, hidden // block_size_k], with a - # TMA-aligned column-major physical layout - # recv_topk_weights: [num_expanded_tokens] - # Here, num_expanded_tokens is the sum of each local expert's token count padded to expert_alignment. - # handle.num_recv_tokens_per_expert_list: a Python list of length num_local_experts; - # each value is the expert's token count padded to expert_alignment, and - # their sum is num_expanded_tokens - # handle.num_unaligned_recv_tokens_per_expert: [num_local_experts], the actual - # token counts before alignment padding - # handle.recv_src_metadata: [num_recv_tokens, topk + 2]; the last topk columns - # map each deduplicated receive token to rows in the expanded tensors - recv_x, _, recv_topk_weights, handle, _ = buffer.dispatch( - (qinput_tensor, input_scale), - topk_idx=topk_idx, - topk_weights=topk_weights, - num_experts=num_experts, - num_max_tokens_per_rank=get_deepep_num_max_dispatch_tokens_per_rank_prefill(), - expert_alignment=128, - previous_event=previous_event, - allocate_on_comm_stream=allocate_on_comm_stream, - do_cpu_sync=True, - do_handle_copy=False, - do_expand=True, - use_tma_aligned_col_major_sf=True, - ) - # Dispatch is synchronous in this path. Its FP8 source is no longer - # needed once the received tensors have been produced. - del qinput_tensor, input_scale - - all_tokens = sum(handle.num_recv_tokens_per_expert_list) - if all_tokens > 0: - gather_out = chunked_expanded_moe_forward( - num_recv_tokens_per_expert_list=handle.num_recv_tokens_per_expert_list, - num_unaligned_recv_tokens_per_expert=handle.num_unaligned_recv_tokens_per_expert, - recv_x=recv_x, - recv_topk_weights=recv_topk_weights, - recv_src_metadata=handle.recv_src_metadata, - w1=w1, - w1_scale=w1_scale, - w2=w2, - w2_scale=w2_scale, - block_size_k=block_size_k, - workspace=dist_group_manager.get_deep_ep_prefill_moe_workspace(), - hidden_dtype=hidden_states.dtype, - alpha=alpha, - limit=limit, - clamp_up_add_one=clamp_up_add_one, - ) - else: - gather_out = torch.empty( - (handle.recv_src_metadata.shape[0], w2.shape[1]), - device=recv_x[0].device, - dtype=hidden_states.dtype, - ) - ######################################## warning ################################################## - # A rank may receive no tokens during autotune warmup. Run one dummy token through - # silu_and_mul_fwd so the empty rank matches the first kernel call made by non-empty ranks. - # This branch does not synchronize additional calls caused by different positive chunk counts. - if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL): - N = w1.shape[1] - _gemm_out_a = torch.zeros((1, N), device=hidden_states.device, dtype=hidden_states.dtype) - _silu_out = torch.zeros((1, N // 2), device=hidden_states.device, dtype=hidden_states.dtype) - silu_and_mul_fwd( - _gemm_out_a.view(-1, N), - _silu_out, - alpha=alpha, - limit=limit, - clamp_up_add_one=clamp_up_add_one, - ) - _gemm_out_a, _silu_out = None, None - del recv_x - - # normal combine - combined_x, _, event = buffer.combine( - gather_out, - handle, - topk_weights=None, - previous_event=previous_event, - allocate_on_comm_stream=allocate_on_comm_stream, - ) - else: - # low latency dispatch - num_max_dispatch_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank_decode() - expected_m = triton.cdiv(hidden_states.shape[0] * buffer.group_size * topk_idx.shape[1], num_experts) - recv_x, masked_m, handle, event, hook = buffer.low_latency_dispatch( - hidden_states, - topk_idx, - num_max_dispatch_tokens_per_rank, - num_experts, - use_fp8=use_fp8_w8a8, - round_scale=get_env_use_ue8m0_scales(), - async_finish=False, - return_recv_hook=False, - ) - # deepgemm - gemm_out_b = masked_group_gemm( - recv_x, - masked_m, - hidden_states.dtype, - w1, - w1_scale, - w2, - w2_scale, - expected_m, - alpha=alpha, - limit=limit, - clamp_up_add_one=clamp_up_add_one, - ) - # low latency combine - combined_x, event_overlap, hook = buffer.low_latency_combine( - gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=False - ) - return combined_x - - -def deepgemm_grouped_fp8_nt_contiguous( - input_tuple: Tuple[torch.Tensor, torch.Tensor], - w_tuple: Tuple[torch.Tensor, torch.Tensor], - out: torch.Tensor, - m_indices: torch.Tensor, -): - if HAS_DEEPGEMM: - if hasattr(deep_gemm, "m_grouped_gemm_fp8_fp8_bf16_nt_contiguous"): - return deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(input_tuple, w_tuple, out, m_indices) - if hasattr(deep_gemm, "m_grouped_fp8_gemm_nt_contiguous"): - return deep_gemm.m_grouped_fp8_gemm_nt_contiguous(input_tuple, w_tuple, out, m_indices) - raise RuntimeError("deep_gemm does not provide grouped_gemm_fp8 NT contiguous GEMM kernel in this version") - - def _get_max_chunk_rows( workspace: torch.Tensor, gather_rows: int, @@ -504,7 +285,7 @@ def can_allocate(chunk_rows: int) -> bool: def chunked_expanded_moe_forward( - num_recv_tokens_per_expert_list: List[int], # [num_local_experts], 128-aligned token counts + num_recv_tokens_per_expert_list: List[int], # [num_local_experts], alignment-padded token counts num_unaligned_recv_tokens_per_expert: torch.Tensor, # [num_local_experts], actual token counts recv_x: Tuple[ torch.Tensor, torch.Tensor # [fp8, scale] @@ -523,7 +304,7 @@ def chunked_expanded_moe_forward( clamp_up_add_one: bool = True, ): """Run bounded expanded MoE and rewrite metadata for dense DeepEP combine.""" - alignment = 128 + alignment = get_mk_alignment_for_contiguous_layout() all_tokens, intermediate_twice = recv_x[0].shape[0], w1.shape[1] intermediate_size, hidden_size = intermediate_twice // 2, w2.shape[1] assert all_tokens == sum(num_recv_tokens_per_expert_list) and all_tokens % alignment == 0 diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/fp4_mega.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/fp4_mega.py new file mode 100644 index 0000000000..d4b74f20a4 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/fp4_mega.py @@ -0,0 +1,88 @@ +import torch +from typing import Any, Dict + +from lightllm.distributed import dist_group_manager +from lightllm.utils.device_utils import is_sm100_gpu + +from .deepgemm import HAS_DEEPGEMM, deep_gemm + +_MEGA_MOE_STATES: Dict[tuple, Dict[str, Any]] = {} + + +def use_sm100_mega_moe(quant_method: Any) -> bool: + return is_sm100_gpu() and quant_method.method_name == "fp4fp8-b32-deepgemm" + + +def mega_moe_impl( + hidden_states: torch.Tensor, + w13: Any, + w2: Any, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + quant_method: Any, +): + if not (HAS_DEEPGEMM and hasattr(deep_gemm, "fp8_fp4_mega_moe")): + raise RuntimeError("deep_gemm does not provide fp8-fp4 Mega MoE kernel") + + from deep_gemm.utils import per_token_cast_to_fp8 + + buffer = getattr(dist_group_manager, "ep_mega_moe_buffer", None) + if buffer is None: + raise RuntimeError("SM100 Mega MoE requires dist_group_manager.ep_mega_moe_buffer to be initialized") + + num_tokens = hidden_states.shape[0] + if num_tokens > buffer.num_max_tokens_per_rank: + raise RuntimeError( + f"Mega MoE got {num_tokens} tokens, exceeding num_max_tokens_per_rank={buffer.num_max_tokens_per_rank}" + ) + + qinput_tensor = per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=quant_method.block_size, + use_packed_ue8m0=True, + ) + state = _get_mega_moe_cache_state(w13, w2) + l1_weights, l2_weights = _get_mega_moe_weights(w13, w2, state) + stats = _get_mega_moe_cumulative_stats(w13.weight.shape[0], hidden_states.device, state) + buffer.x[:num_tokens].copy_(qinput_tensor[0]) + buffer.x_sf[:num_tokens].copy_(qinput_tensor[1]) + buffer.topk_idx[:num_tokens].copy_(topk_ids) + buffer.topk_weights[:num_tokens].copy_(topk_weights) + + output = torch.empty_like(hidden_states) + deep_gemm.fp8_fp4_mega_moe( + output, + l1_weights, + l2_weights, + buffer, + cumulative_local_expert_recv_stats=stats, + ) + return output + + +def _get_mega_moe_cache_state(w13: Any, w2: Any): + state_key = ( + w13.weight.data_ptr(), + w13.weight_scale.data_ptr(), + w2.weight.data_ptr(), + w2.weight_scale.data_ptr(), + ) + return _MEGA_MOE_STATES.setdefault(state_key, {}) + + +def _get_mega_moe_weights(w13: Any, w2: Any, state: Dict[str, Any]): + if "weight_cache" not in state: + state["weight_cache"] = deep_gemm.transform_weights_for_mega_moe( + (w13.weight, w13.weight_scale), + (w2.weight, w2.weight_scale), + ) + return state["weight_cache"] + + +def _get_mega_moe_cumulative_stats(num_local_experts: int, device: torch.device, state: Dict[str, Any]): + stats = state.get("stats") + if stats is None or stats.numel() != num_local_experts or stats.device != device: + stats = torch.zeros((num_local_experts,), device=device, dtype=torch.int32) + state["stats"] = stats + return stats diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/prefill_decode_sync.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/prefill_decode_sync.py new file mode 100644 index 0000000000..3379bb17f8 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/prefill_decode_sync.py @@ -0,0 +1,261 @@ +import torch +import triton +from typing import Any, Optional + +from lightllm.distributed import dist_group_manager +from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd +from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import ( + per_token_group_quant_fp8, +) +from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_expanded_layout_kernels import ( + ep_reduce_decode_output, +) +from lightllm.utils.envs_utils import ( + get_deepep_num_max_dispatch_tokens_per_rank_prefill, + get_deepep_num_max_dispatch_tokens_per_rank_decode, + get_env_use_ue8m0_scales, +) +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + +from .common import check_ep_expert_dtype, get_ep_num_sms +from .deepgemm import ( + chunked_expanded_moe_forward, + decode_masked_group_gemm, + get_mk_alignment_for_contiguous_layout, +) +from .fp4_mega import mega_moe_impl, use_sm100_mega_moe + + +def fused_experts( + hidden_states: torch.Tensor, + w13: Any, + w2: Any, + topk_weights: torch.Tensor, + topk_idx: torch.Tensor, + num_experts: int, + quant_method: Any, + is_prefill: Optional[bool], + previous_event: Optional[Any] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, +): + assert (limit is None and alpha is None) or (limit is not None and alpha is not None) + check_ep_expert_dtype(quant_method) + if use_sm100_mega_moe(quant_method): + if limit is not None: + raise NotImplementedError("FP4 Mega MoE does not support clamped SwiGLU") + return mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method) + + return fused_experts_impl( + hidden_states=hidden_states, + w1=w13.weight, + w2=w2.weight, + topk_weights=topk_weights, + topk_idx=topk_idx, + num_experts=num_experts, + block_size_k=quant_method.block_size, + buffer=dist_group_manager.ep_buffer, + is_prefill=is_prefill, + use_fp8_w8a8=True, + use_fp8_all2all=True, + use_int8_w8a16=False, + w1_scale=w13.weight_scale, + w2_scale=w2.weight_scale, + previous_event=previous_event, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + ) + + +def fused_experts_impl( + hidden_states: torch.Tensor, # [M, K] + w1: torch.Tensor, # [group, N, K] + w2: torch.Tensor, # [group, K, N/2] + topk_weights: torch.Tensor, # [M, topk] + topk_idx: torch.Tensor, # [M, topk] + num_experts: int, + buffer: Any, + is_prefill: bool, + use_fp8_w8a8: bool = False, + use_fp8_all2all: bool = False, + use_int8_w8a16: bool = False, + w1_scale: Optional[torch.Tensor] = None, + w2_scale: Optional[torch.Tensor] = None, + previous_event: Optional[Any] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, + block_size_k: int = 128, +): + # Check constraints. + assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" + assert topk_weights.shape == topk_idx.shape, "topk shape mismatch" + assert hidden_states.is_contiguous(), "Hidden_states must be contiguous" + assert w1.is_contiguous(), "Expert weights1 must be contiguous" + assert w2.is_contiguous(), "Expert weights2 must be contiguous" + assert hidden_states.dtype in [torch.float32, torch.float16, torch.bfloat16] + + # qaunt hidden_states + assert use_fp8_w8a8 and use_fp8_all2all, "use_fp8_w8a8 and use_fp8_all2all must be True" + assert block_size_k == 128, "block_size_k must be 128" + + if is_prefill: + qinput_tensor, input_scale = per_token_group_quant_fp8( + hidden_states, block_size_k, dtype=w1.dtype, use_ue8m0_scales=get_env_use_ue8m0_scales() + ) + allocate_on_comm_stream = previous_event is not None + # Expanded dispatch directly produces expert-contiguous, alignment-padded inputs: + # recv_x[0]: [num_expanded_tokens, hidden] + # recv_x[1]: [num_expanded_tokens, hidden // block_size_k], with a + # TMA-aligned column-major physical layout + # recv_topk_weights: [num_expanded_tokens] + # Here, num_expanded_tokens is the sum of each local expert's token count padded to expert_alignment. + # handle.num_recv_tokens_per_expert_list: a Python list of length num_local_experts; + # each value is the expert's token count padded to expert_alignment, and + # their sum is num_expanded_tokens + # handle.num_unaligned_recv_tokens_per_expert: [num_local_experts], the actual + # token counts before alignment padding + # handle.recv_src_metadata: [num_recv_tokens, topk + 2]; the last topk columns + # map each deduplicated receive token to rows in the expanded tensors + recv_x, _, recv_topk_weights, handle, _ = buffer.dispatch( + (qinput_tensor, input_scale), + topk_idx=topk_idx, + topk_weights=topk_weights, + num_experts=num_experts, + num_max_tokens_per_rank=get_deepep_num_max_dispatch_tokens_per_rank_prefill(), + expert_alignment=get_mk_alignment_for_contiguous_layout(), + # 当前 prefill 路径同步等待通信完成，不与 grouped GEMM 重叠， + # 因此使用非 overlap 配额，让通信可以占用更多 SM。 + num_sms=get_ep_num_sms(overlap_with_compute=False), + previous_event=previous_event, + async_with_compute_stream=False, + allocate_on_comm_stream=allocate_on_comm_stream, + do_cpu_sync=True, + do_handle_copy=False, + do_expand=True, + use_tma_aligned_col_major_sf=True, + ) + # Dispatch is synchronous in this path. Its FP8 source is no longer + # needed once the received tensors have been produced. + del qinput_tensor, input_scale + + all_tokens = sum(handle.num_recv_tokens_per_expert_list) + if all_tokens > 0: + gather_out = chunked_expanded_moe_forward( + num_recv_tokens_per_expert_list=handle.num_recv_tokens_per_expert_list, + num_unaligned_recv_tokens_per_expert=handle.num_unaligned_recv_tokens_per_expert, + recv_x=recv_x, + recv_topk_weights=recv_topk_weights, + recv_src_metadata=handle.recv_src_metadata, + w1=w1, + w1_scale=w1_scale, + w2=w2, + w2_scale=w2_scale, + block_size_k=block_size_k, + workspace=dist_group_manager.get_deep_ep_prefill_moe_workspace(), + hidden_dtype=hidden_states.dtype, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + ) + else: + gather_out = torch.empty( + (handle.recv_src_metadata.shape[0], w2.shape[1]), + device=recv_x[0].device, + dtype=hidden_states.dtype, + ) + ######################################## warning ################################################## + # A rank may receive no tokens during autotune warmup. Run one dummy token through + # silu_and_mul_fwd so the empty rank matches the first kernel call made by non-empty ranks. + # This branch does not synchronize additional calls caused by different positive chunk counts. + if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL): + N = w1.shape[1] + _gemm_out_a = torch.zeros((1, N), device=hidden_states.device, dtype=hidden_states.dtype) + _silu_out = torch.zeros((1, N // 2), device=hidden_states.device, dtype=hidden_states.dtype) + silu_and_mul_fwd( + _gemm_out_a.view(-1, N), + _silu_out, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + ) + _gemm_out_a, _silu_out = None, None + del recv_x + + # normal combine + combined_x, _, _ = buffer.combine( + gather_out, + handle, + topk_weights=None, + # 与 dispatch 保持相同的串行通信配额，避免退回 handle 中的隐式值。 + num_sms=get_ep_num_sms(overlap_with_compute=False), + previous_event=previous_event, + async_with_compute_stream=False, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + return combined_x + + qinput_tensor, input_scale = per_token_group_quant_fp8( + hidden_states, block_size_k, dtype=w1.dtype, use_ue8m0_scales=get_env_use_ue8m0_scales() + ) + # 同一 token 命中目标 rank 上的多个 expert 时，hidden vector 只跨卡 + # 发送一次；接收端再将其展开成 expert 连续布局。这里的对齐值必须与 + # DeepGEMM contiguous grouped GEMM 的 M 对齐要求完全一致。 + recv_x, _, recv_topk_weights, ep_handle, _ = buffer.dispatch( + (qinput_tensor, input_scale), + topk_idx=topk_idx, + topk_weights=topk_weights, + num_experts=num_experts, + num_max_tokens_per_rank=get_deepep_num_max_dispatch_tokens_per_rank_decode(), + expert_alignment=get_mk_alignment_for_contiguous_layout(), + # 串行 decode 会在 dispatch 内部等待通信完成，因此输出 tensor + # 直接归属计算流，并使用非 overlap SM 配额缩短通信时间。 + num_sms=get_ep_num_sms(overlap_with_compute=False), + async_with_compute_stream=False, + allocate_on_comm_stream=False, + do_cpu_sync=False, + do_handle_copy=False, + do_expand=True, + do_zero_padding=True, + use_tma_aligned_col_major_sf=True, + ) + del qinput_tensor, input_scale + + expected_m = triton.cdiv(hidden_states.shape[0] * buffer.num_ranks * topk_idx.shape[1], num_experts) + expert_output = decode_masked_group_gemm( + recv_x=recv_x, + expert_token_psum=ep_handle.psum_num_recv_tokens_per_expert, + expert_alignment=ep_handle.expert_alignment, + dtype=hidden_states.dtype, + w1=w1, + w1_scale=w1_scale, + w2=w2, + w2_scale=w2_scale, + expected_m=expected_m, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + ) + # 阶段 1：把按 expert 展开的输出归约回去重接收 token 布局。 + dense_output, compact_metadata = ep_reduce_decode_output( + expert_output=expert_output, + route_weights=recv_topk_weights, + recv_src_metadata=ep_handle.recv_src_metadata, + num_valid_recv_tokens=ep_handle.psum_num_recv_tokens_per_scaleup_rank[-1:], + ) + del expert_output + ep_handle.recv_src_metadata = compact_metadata + + # 阶段 2：compact metadata 的第一个 top-k 槽位指向同序 dense row， + # 其余槽位置为 -1。串行 combine 在内部等待通信完成，不返回有效 event。 + combined_x, _, _ = buffer.combine( + dense_output, + ep_handle, + topk_weights=None, + num_sms=get_ep_num_sms(overlap_with_compute=False), + async_with_compute_stream=False, + allocate_on_comm_stream=False, + ) + return combined_x diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py index 8c94f088c2..a63d40a842 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py @@ -165,3 +165,177 @@ def silu_and_mul_masked_post_quant_fwd( num_warps=num_warps, ) return + + +@triton.jit +def _silu_and_mul_psum_post_quant_kernel( + input_ptr, + output_ptr, + output_scale_ptr, + expert_token_psum_ptr, + input_stride_m, + output_stride_m, + output_scale_stride_m, + output_scale_stride_k, + size_n, + fp8_max, + fp8_min, + EXPERT_ALIGNMENT: tl.constexpr, + BLOCK_N: tl.constexpr, + NUM_STAGE: tl.constexpr, + NEED_MASK: tl.constexpr, + USE_TANH_APPROXIMATE_GELU: tl.constexpr = False, + USE_LIMIT_AND_ALPHA: tl.constexpr = False, + alpha: tl.constexpr = None, + limit: tl.constexpr = None, + CLAMP_UP_ADD_ONE: tl.constexpr = True, + USE_UE8M0_SCALE: tl.constexpr = False, +): + hidden_block_id = tl.program_id(0).to(tl.int64) + row_offset_in_expert = tl.program_id(1).to(tl.int64) + row_step = tl.num_programs(1).to(tl.int64) + expert_id = tl.program_id(2).to(tl.int64) + + # psum 保存每个 expert 最后一个有效行的后一位置；下一个 expert 的 + # 起点是前一个结束位置向上对齐到 EXPERT_ALIGNMENT 后的位置。 + expert_end = tl.load(expert_token_psum_ptr + expert_id).to(tl.int64) + previous_expert_end = tl.load( + expert_token_psum_ptr + expert_id - 1, + mask=expert_id > 0, + other=0, + ).to(tl.int64) + expert_start = tl.cdiv(previous_expert_end, EXPERT_ALIGNMENT) * EXPERT_ALIGNMENT + + # 行号与 stride 相乘后才用于地址计算；必须在乘法前提升为 int64， + # 否则大 tensor 的元素偏移可能先在 int32 中溢出。 + input_stride_m = tl.cast(input_stride_m, tl.int64) + output_stride_m = tl.cast(output_stride_m, tl.int64) + output_scale_stride_m = tl.cast(output_scale_stride_m, tl.int64) + output_scale_stride_k = tl.cast(output_scale_stride_k, tl.int64) + + hidden_offsets = hidden_block_id * BLOCK_N + tl.arange(0, BLOCK_N) + input_offsets = hidden_offsets + output_offsets = hidden_offsets + scale_offsets = hidden_block_id * output_scale_stride_k + if NEED_MASK: + hidden_mask = hidden_offsets < size_n + other = 0.0 + else: + hidden_mask = None + other = None + + first_row = expert_start + row_offset_in_expert + for row_index in tl.range(first_row, expert_end, row_step, num_stages=NUM_STAGE): + input_row_offsets = row_index * input_stride_m + gate = tl.load( + input_ptr + input_row_offsets + input_offsets, + mask=hidden_mask, + other=other, + ).to(tl.float32) + up = tl.load( + input_ptr + input_row_offsets + input_offsets + size_n, + mask=hidden_mask, + other=other, + ) + if USE_LIMIT_AND_ALPHA: + gate = tl.minimum(gate, limit) + up = tl.minimum(tl.maximum(up, -limit), limit) + gate = gate / (1 + tl.exp(-gate * alpha)) + if CLAMP_UP_ADD_ONE: + up += 1 + elif USE_TANH_APPROXIMATE_GELU: + gate_cubed = gate * gate * gate + tanh_arg = 0.7978845608028654 * (gate + 0.044715 * gate_cubed) + tanh_val = 2.0 / (1.0 + tl.exp(-2.0 * tanh_arg)) - 1.0 + gate = 0.5 * gate * (1.0 + tanh_val) + else: + gate = gate / (1 + tl.exp(-gate)) + gate = gate.to(input_ptr.dtype.element_ty) + gate_up = up * gate + if USE_LIMIT_AND_ALPHA: + gate_up = gate_up.to(input_ptr.dtype.element_ty).to(tl.float32) + _absmax = tl.maximum(tl.max(tl.abs(gate_up)), 1e-10) + output_s = _absmax / fp8_max + if USE_UE8M0_SCALE: + output_s = tl.exp2(tl.ceil(tl.log2(output_s))) + output_q = tl.clamp(gate_up / output_s, fp8_min, fp8_max).to(output_ptr.dtype.element_ty) + tl.store( + output_ptr + row_index * output_stride_m + output_offsets, + output_q, + mask=hidden_mask, + ) + tl.store(output_scale_ptr + row_index * output_scale_stride_m + scale_offsets, output_s) + + +def silu_and_mul_psum_post_quant_fwd( + input: torch.Tensor, + output: torch.Tensor, + output_scale: torch.Tensor, + expert_token_psum: torch.Tensor, + expert_alignment: int, + quant_group_size: int, + alpha=None, + limit=None, + clamp_up_add_one=True, +): + """按 expert 的 psum 区间处理有效行，跳过 expert 之间的对齐填充。""" + assert (limit is None and alpha is None) or (limit is not None and alpha is not None) + assert input.ndim == 2 and input.is_contiguous() + assert output.ndim == 2 and output.is_contiguous() + assert output_scale.ndim == 2 + assert output.dtype == torch.float8_e4m3fn + assert expert_token_psum.ndim == 1 and expert_token_psum.numel() > 0 + assert expert_alignment > 0 + + assert input.shape[0] == output.shape[0] == output_scale.shape[0] + assert input.shape[1] == output.shape[1] * 2 + + size_n = output.shape[1] + assert size_n % quant_group_size == 0 + assert output_scale.shape[1] == size_n // quant_group_size + + BLOCK_N = quant_group_size + num_warps = 1 + NUM_STAGES = 6 + hidden_dim_split_block_num = triton.cdiv(size_n, BLOCK_N) + assert BLOCK_N == quant_group_size + NEED_MASK = (size_n % BLOCK_N) != 0 + + num_experts = expert_token_psum.numel() + # 行方向总并行度以 256 为目标，并平均分给所有 expert。常见配置为： + # 4/8/16/32/64/128/256 个 expert 分别得到 64/32/16/8/4/2/1。 + # 单 expert 最多使用 64 个 program；expert 超过 256 时仍至少使用 1 个。 + target_row_parallelism = 256 + max_row_parallelism_per_expert = 64 + row_parallelism = target_row_parallelism // num_experts + row_parallelism = max(row_parallelism, 1) + row_parallelism = min(row_parallelism, max_row_parallelism_per_expert) + grid = (hidden_dim_split_block_num, row_parallelism, num_experts) + + finfo = torch.finfo(torch.float8_e4m3fn) + fp8_max = finfo.max + fp8_min = -fp8_max + _silu_and_mul_psum_post_quant_kernel[grid]( + input_ptr=input, + output_ptr=output, + output_scale_ptr=output_scale, + expert_token_psum_ptr=expert_token_psum, + input_stride_m=input.stride(0), + output_stride_m=output.stride(0), + output_scale_stride_m=output_scale.stride(0), + output_scale_stride_k=output_scale.stride(1), + size_n=size_n, + fp8_max=fp8_max, + fp8_min=fp8_min, + EXPERT_ALIGNMENT=expert_alignment, + BLOCK_N=BLOCK_N, + NUM_STAGE=NUM_STAGES, + NEED_MASK=NEED_MASK, + USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(), + USE_LIMIT_AND_ALPHA=limit is not None, + alpha=alpha, + limit=limit, + CLAMP_UP_ADD_ONE=clamp_up_add_one, + USE_UE8M0_SCALE=get_env_use_ue8m0_scales(), + num_warps=num_warps, + ) diff --git a/lightllm/distributed/communication_op.py b/lightllm/distributed/communication_op.py index 93c603212d..5bb4030d89 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -28,6 +28,7 @@ from lightllm.utils.device_utils import has_nvlink from lightllm.utils.envs_utils import ( get_env_start_args, + get_deepep_num_max_dispatch_tokens_per_rank, get_deepep_num_max_dispatch_tokens_per_rank_prefill, get_deepep_num_max_dispatch_tokens_per_rank_decode, get_redundancy_expert_num, @@ -112,6 +113,7 @@ def __init__(self): self.ep_low_latency_buffer = None self.ep_mega_moe_buffer = None self.ep_num_sms = None + self.ep_non_overlap_num_sms = None def __len__(self): return len(self.groups) @@ -167,7 +169,7 @@ def new_deepep_group( ``expert_quant_method_names`` 是各 MoE 层最终绑定的 quant method 名称集合。 同一个模型可能逐层混用 FP4 和 FP8：SM100 FP4 层走 Mega MoE，其他层走 - DeepEP legacy low-latency 路径。这里只为实际存在的执行路径分配 buffer， + DeepEP Elastic 路径。这里只为实际存在的执行路径分配 buffer， 避免为未使用的路径长期占用显存。 """ enable_ep_moe = get_env_start_args().enable_ep_moe @@ -178,9 +180,17 @@ def new_deepep_group( self.ep_low_latency_buffer = None self.ep_mega_moe_buffer = None self.ep_num_sms = None + self.ep_non_overlap_num_sms = None return assert HAS_DEEPEP, "deep_ep is required for expert parallelism" + from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep + + # LightLLM 的 expanded layout 按 128 行构造 expert 分段。初始化阶段同步 + # 设置 DeepGEMM 的全局 M 对齐值，后续 prefill 和 decode 都通过 getter + # 读取它，避免通信布局与 GEMM 分块不一致。 + grouped_fused_moe_ep.set_mk_alignment_for_contiguous_layout(128) + global_world_size = get_global_world_size() deepep_group = dist.new_group(list(range(global_world_size))) # DeepEP reuses this group's NCCL communicator via _comm_ptr(). Because the @@ -190,7 +200,7 @@ def new_deepep_group( group=deepep_group, device_ids=[torch.cuda.current_device()], ) - self.ll_num_tokens = prefill_num_max_dispatch_tokens_per_rank + self.ll_num_tokens = get_deepep_num_max_dispatch_tokens_per_rank() self.ll_decode_num_tokens = decode_num_max_dispatch_tokens_per_rank self.ll_hidden = hidden_size self.ll_num_experts = n_routed_experts + get_redundancy_expert_num() * global_world_size @@ -212,13 +222,14 @@ def new_deepep_group( is_sm100 = is_sm100_gpu() # Buffer 选择规则： - # 1. 非 SM100 不支持 Mega MoE，只初始化 legacy low-latency buffer； + # 1. 非 SM100 不支持 Mega MoE，只保留 prefill workspace 用的 RDMA buffer； # 2. SM100 全部 MoE 层为 FP4，只初始化 Mega MoE buffer； - # 3. SM100 全部 MoE 层为 FP8，只初始化 legacy low-latency buffer； + # 3. SM100 全部 MoE 层为 FP8，只保留 prefill workspace buffer； # 4. SM100 逐层混合 FP4/FP8，两套 buffer 都要初始化。 + # Decode 计算走 ElasticBuffer 的 expanded 布局，不再调用 low-latency dispatch。 if is_sm100: # 只要存在一个 FP4 MoE 层，就需要 Mega MoE buffer；只要存在一个非 FP4 - # MoE 层，就需要 legacy low-latency buffer。FP4/FP8 逐层混用时两者都会初始化。 + # MoE 层，prefill 就需要复用 RDMA workspace。FP4/FP8 逐层混用时两者都会初始化。 has_mega_moe_layer = mega_moe_quant_method in expert_quant_method_names has_legacy_moe_layer = any( method_name != mega_moe_quant_method for method_name in expert_quant_method_names @@ -230,8 +241,9 @@ def new_deepep_group( enable_low_latency_buffer = True if enable_low_latency_buffer: - # FP8 MoE 的 decode 使用 legacy low-latency buffer；prefill 阶段还会将其 - # 空闲的本地 RDMA storage 复用为分块 grouped GEMM 的临时 workspace。 + # Prefill expanded grouped GEMM 复用这块 RDMA storage 作为 workspace。 + # Decode 不再走 low-latency kernel，容量仍取 decode hint 与 prefill + # workspace 的较大值，避免后续若再次使用该 buffer 时容量不足。 decode_size_hint = deep_ep.Buffer.get_low_latency_rdma_size_hint( self.ll_decode_num_tokens, self.ll_hidden, global_world_size, self.ll_num_experts ) @@ -282,6 +294,25 @@ def new_deepep_group( theoretical_sms = self.ep_buffer.get_theoretical_num_sms(self.ll_num_experts, num_experts_per_tok) self._set_num_sms_for_deep_gemm(theoretical_sms) + # 串行路径会在通信后立即等待，不需要为并行计算预留大部分 SM。 + # DeepEP 要求 num_sms 为偶数，因此将设备总 SM 数的 60% 向下取偶数。 + device_sms = get_device_sm_count() + non_overlap_num_sms = int(device_sms * 0.6) // 2 * 2 + self.set_ep_num_sms_for_non_overlap(non_overlap_num_sms) + logger.info( + "DeepEP SM allocation: overlap=%s, non_overlap=%s, device=%s", + self.ep_num_sms, + self.ep_non_overlap_num_sms, + device_sms, + ) + + def set_ep_num_sms_for_non_overlap(self, num_sms: int) -> None: + """设置非 overlap 通信使用的 SM 数量，不改变 DeepGEMM 的 SM 配额。""" + device_sms = get_device_sm_count() + assert 4 <= num_sms <= device_sms + assert num_sms % 2 == 0 + self.ep_non_overlap_num_sms = num_sms + def _set_num_sms_for_deep_gemm(self, deepep_sms: int): try: try: diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index 986b8a703f..ca1e86a9bc 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -229,7 +229,6 @@ def _get_o( def _moe_ffn_tp( self, input, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight ) -> torch.Tensor: - hidden_states = input.view(-1, self.embed_dim_) num_tokens, hidden_dim = hidden_states.shape @@ -258,7 +257,6 @@ def _moe_ffn_tp( def _moe_ffn_edp( self, input, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight ) -> torch.Tensor: - hidden_states = input token_num, hidden_dim = hidden_states.shape if self.n_shared_experts is not None: @@ -340,14 +338,10 @@ def overlap_tpsp_token_forward( _0_shared_output = LlamaTransformerLayerInfer._ffn_tp(self, _0_input1, infer_state, layer_weight) # 0 dispatch - ( - _0_recv_x, - _0_masked_m, - _0_topk_idx, - _0_topk_weight, - _0_handle, - _0_hook, - ) = layer_weight.experts.low_latency_dispatch(_0_input1, _0_router_logits) + _0_recv_x, _0_handle, _0_recv_topk_weights, _0_hook = layer_weight.experts.decode_dispatch( + _0_input1, _0_router_logits + ) + del _0_input1, _0_router_logits infer_state.hook = _0_hook # 1 attention @@ -375,21 +369,20 @@ def overlap_tpsp_token_forward( _1_shared_output = LlamaTransformerLayerInfer._ffn_tp(self, _1_input1, infer_state1, layer_weight) # 1 dispatch - ( - _1_recv_x, - _1_masked_m, - _1_topk_idx, - _1_topk_weight, - _1_handle, - _1_hook, - ) = layer_weight.experts.low_latency_dispatch(_1_input1, _1_router_logits) + _1_recv_x, _1_handle, _1_recv_topk_weights, _1_hook = layer_weight.experts.decode_dispatch( + _1_input1, _1_router_logits + ) + del _1_input1, _1_router_logits infer_state1.hook = _1_hook # moe calu expected_m = triton.cdiv( input_embdings.shape[0] * get_global_world_size() * self.num_experts_per_tok, self.n_routed_experts ) - _0_moe_out = layer_weight.experts.masked_group_gemm(_0_recv_x, _0_masked_m, input_embdings.dtype, expected_m) + _0_moe_out = layer_weight.experts.decode_masked_group_gemm( + _0_recv_x, _0_handle, input_embdings.dtype, expected_m + ) + del _0_recv_x # 1 hook if getattr(infer_state1, "hook", None) is not None: @@ -397,27 +390,30 @@ def overlap_tpsp_token_forward( infer_state1.hook = None # 0 combine - _0_ffn_out, _0_hook = layer_weight.experts.low_latency_combine( - _0_moe_out, _0_topk_idx, _0_topk_weight, _0_handle - ) + _0_ffn_out, _0_hook = layer_weight.experts.decode_combine(_0_moe_out, _0_handle, _0_recv_topk_weights) + del _0_moe_out infer_state.hook = _0_hook # to do moe caclue - _1_moe_out = layer_weight.experts.masked_group_gemm(_1_recv_x, _1_masked_m, input_embdings1.dtype, expected_m) + _1_moe_out = layer_weight.experts.decode_masked_group_gemm( + _1_recv_x, _1_handle, input_embdings1.dtype, expected_m + ) + del _1_recv_x # 0 hook if getattr(infer_state, "hook", None) is not None: infer_state.hook() if self.n_shared_experts is not None: _0_ffn_out.add_(_0_shared_output) + del _0_shared_output input_embdings.add_(_0_ffn_out.view(-1, self.embed_dim_)) + del _0_ffn_out infer_state.hook = None # 1 combine - _1_ffn_out, _1_hook = layer_weight.experts.low_latency_combine( - _1_moe_out, _1_topk_idx, _1_topk_weight, _1_handle - ) + _1_ffn_out, _1_hook = layer_weight.experts.decode_combine(_1_moe_out, _1_handle, _1_recv_topk_weights) + del _1_moe_out def _1_hook_post(): _1_hook() @@ -465,6 +461,7 @@ def overlap_tpsp_context_forward( _0_topk_weight, _0_topk_idx, _0_qinput_tensor = layer_weight.experts.select_experts_and_quant_input( _0_input1, _0_router_logits ) + del _0_router_logits from deep_ep import ElasticBuffer _0_overlap_event = ElasticBuffer.capture() @@ -493,7 +490,10 @@ def overlap_tpsp_context_forward( _0_num_recv_tokens_per_expert_list, _0_handle, _0_hook, - ) = layer_weight.experts.dispatch(_0_qinput_tensor, _0_topk_idx, _0_topk_weight, overlap_event=_0_overlap_event) + ) = layer_weight.experts.prefill_dispatch( + _0_qinput_tensor, _0_topk_idx, _0_topk_weight, overlap_event=_0_overlap_event + ) + del _0_qinput_tensor, _0_topk_idx, _0_topk_weight infer_state.hook = _0_hook # wait 0 dispatch @@ -504,15 +504,18 @@ def overlap_tpsp_context_forward( _1_topk_weight, _1_topk_idx, _1_qinput_tensor = layer_weight.experts.select_experts_and_quant_input( _1_input1, _1_router_logits ) + del _1_router_logits _1_overlap_event = ElasticBuffer.capture() # 0 shared expert if self.n_shared_experts is not None: _0_shared_output = LlamaTransformerLayerInfer._ffn_tp(self, _0_input1, infer_state, layer_weight) + del _0_input1 # 1 shared expert if self.n_shared_experts is not None: _1_shared_output = LlamaTransformerLayerInfer._ffn_tp(self, _1_input1, infer_state1, layer_weight) + del _1_input1 # 0 moe calu _0_moe_out = layer_weight.experts.prefilled_group_gemm( @@ -524,6 +527,7 @@ def overlap_tpsp_context_forward( _0_recv_topk_weight, microbatch_index=0, ) + del _0_recv_x # 1 dispatch execute ( @@ -533,7 +537,10 @@ def overlap_tpsp_context_forward( _1_num_recv_tokens_per_expert_list, _1_handle, _1_hook, - ) = layer_weight.experts.dispatch(_1_qinput_tensor, _1_topk_idx, _1_topk_weight, overlap_event=_1_overlap_event) + ) = layer_weight.experts.prefill_dispatch( + _1_qinput_tensor, _1_topk_idx, _1_topk_weight, overlap_event=_1_overlap_event + ) + del _1_qinput_tensor, _1_topk_idx, _1_topk_weight infer_state1.hook = _1_hook # wait 1 dispatch @@ -543,7 +550,8 @@ def overlap_tpsp_context_forward( _0_combine_event = ElasticBuffer.capture() # 0 combine execute - _0_ffn_out, _0_hook = layer_weight.experts.combine(_0_moe_out, _0_handle, _0_combine_event) + _0_ffn_out, _0_hook = layer_weight.experts.prefill_combine(_0_moe_out, _0_handle, _0_combine_event) + del _0_moe_out infer_state.hook = _0_hook # 1 moe calc @@ -556,6 +564,7 @@ def overlap_tpsp_context_forward( _1_recv_topk_weight, microbatch_index=1, ) + del _1_recv_x # wait 0 combine if getattr(infer_state, "hook", None) is not None: @@ -569,7 +578,7 @@ def overlap_tpsp_context_forward( input_embdings.add_(_0_ffn_out.view(-1, self.embed_dim_)) # 1 combine execute - _1_ffn_out, _1_hook = layer_weight.experts.combine(_1_moe_out, _1_handle, _1_combine_event) + _1_ffn_out, _1_hook = layer_weight.experts.prefill_combine(_1_moe_out, _1_handle, _1_combine_event) def _1_hook_post(): _1_hook() diff --git a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py index 0df3c618fe..65dd507775 100644 --- a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py @@ -173,24 +173,26 @@ def overlap_tpsp_context_forward( x1 = self._tpsp_allgather(x1.view(-1, self.embed_dim_), infer_state1) weights1, indices1 = self._select_experts(x1, infer_state1, layer_weight) - recv_x0, recv_indices0, recv_weights0, recv_count0, handle0, dispatch_hook0 = experts.dispatch( + recv_x0, recv_indices0, recv_weights0, recv_count0, handle0, dispatch_hook0 = experts.prefill_dispatch( qinput0, indices0, weights0, overlap_event=dispatch_event0, ) + del qinput0, indices0, weights0 dispatch_hook0() indices1 = indices1.to(torch.long) qinput1 = experts.quantize_dispatch_input(x1) dispatch_event1 = ElasticBuffer.capture() - recv_x1, recv_indices1, recv_weights1, recv_count1, handle1, dispatch_hook1 = experts.dispatch( + recv_x1, recv_indices1, recv_weights1, recv_count1, handle1, dispatch_hook1 = experts.prefill_dispatch( qinput1, indices1, weights1, overlap_event=dispatch_event1, ) + del qinput1, indices1, weights1 moe_out0 = experts.prefilled_group_gemm( recv_count0, @@ -205,9 +207,11 @@ def overlap_tpsp_context_forward( limit=self.swiglu_limit, clamp_up_add_one=False, ) + del recv_x0 combine_event0 = ElasticBuffer.capture() - routed0, combine_hook0 = experts.combine(moe_out0, handle0, overlap_event=combine_event0) + routed0, combine_hook0 = experts.prefill_combine(moe_out0, handle0, overlap_event=combine_event0) + del moe_out0 dispatch_hook1() moe_out1 = experts.prefilled_group_gemm( @@ -223,9 +227,11 @@ def overlap_tpsp_context_forward( limit=self.swiglu_limit, clamp_up_add_one=False, ) + del recv_x1 combine_event1 = ElasticBuffer.capture() - routed1, combine_hook1 = experts.combine(moe_out1, handle1, overlap_event=combine_event1) + routed1, combine_hook1 = experts.prefill_combine(moe_out1, handle1, overlap_event=combine_event1) + del moe_out1 shared0 = self._ffn_tp(x0, infer_state, layer_weight) combine_hook0() @@ -269,9 +275,7 @@ def overlap_tpsp_token_forward( infer_state1.call_overlap_hook() shared0 = self._ffn_tp(x0, infer_state, layer_weight) - recv_x0, masked_m0, indices0, weights0, handle0, dispatch_hook0 = experts.low_latency_dispatch_with_topk( - x0, indices0, weights0 - ) + recv_x0, handle0, recv_weights0, dispatch_hook0 = experts.decode_dispatch_with_topk(x0, indices0, weights0) x1, residual1, residual_mix1, post_mix1 = self._hc_attn_in(input_embdings1, layer_weight) x1 = self.token_attention_forward(x1, infer_state1, layer_weight) @@ -281,42 +285,44 @@ def overlap_tpsp_token_forward( dispatch_hook0() shared1 = self._ffn_tp(x1, infer_state1, layer_weight) - recv_x1, masked_m1, indices1, weights1, handle1, dispatch_hook1 = experts.low_latency_dispatch_with_topk( - x1, indices1, weights1 - ) + recv_x1, handle1, recv_weights1, dispatch_hook1 = experts.decode_dispatch_with_topk(x1, indices1, weights1) expected_m = triton.cdiv( x0.shape[0] * get_global_world_size() * self.num_experts_per_tok, layer_weight.n_routed_experts, ) - moe_out0 = experts.masked_group_gemm( + moe_out0 = experts.decode_masked_group_gemm( recv_x0, - masked_m0, + handle0, x0.dtype, expected_m, alpha=1.0, limit=self.swiglu_limit, clamp_up_add_one=False, ) + del recv_x0 dispatch_hook1() - routed0, combine_hook0 = experts.low_latency_combine(moe_out0, indices0, weights0, handle0) + routed0, combine_hook0 = experts.decode_combine(moe_out0, handle0, recv_weights0) + del moe_out0 - moe_out1 = experts.masked_group_gemm( + moe_out1 = experts.decode_masked_group_gemm( recv_x1, - masked_m1, + handle1, x1.dtype, expected_m, alpha=1.0, limit=self.swiglu_limit, clamp_up_add_one=False, ) + del recv_x1 combine_hook0() routed0.add_(shared0) output0 = self._hc_ffn_out(routed0, residual0, residual_mix0, post_mix0) - routed1, combine_hook1 = experts.low_latency_combine(moe_out1, indices1, weights1, handle1) + routed1, combine_hook1 = experts.decode_combine(moe_out1, handle1, recv_weights1) + del moe_out1 if self.is_last_layer: combine_hook1() routed1.add_(shared1) diff --git a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py index 7311c4d141..3a709af80f 100644 --- a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py @@ -160,14 +160,10 @@ def overlap_tpsp_token_forward( infer_state1.hook = None # 0 dispatch - ( - _0_recv_x, - _0_masked_m, - _0_topk_idx, - _0_topk_weight, - _0_handle, - _0_hook, - ) = layer_weight.experts.low_latency_dispatch(_0_input1, _0_router_logits) + _0_recv_x, _0_handle, _0_recv_topk_weights, _0_hook = layer_weight.experts.decode_dispatch( + _0_input1, _0_router_logits + ) + del _0_input1, _0_router_logits infer_state.hook = _0_hook # 1 attention @@ -190,21 +186,20 @@ def overlap_tpsp_token_forward( infer_state.hook = None # 1 dispatch - ( - _1_recv_x, - _1_masked_m, - _1_topk_idx, - _1_topk_weight, - _1_handle, - _1_hook, - ) = layer_weight.experts.low_latency_dispatch(_1_input1, _1_router_logits) + _1_recv_x, _1_handle, _1_recv_topk_weights, _1_hook = layer_weight.experts.decode_dispatch( + _1_input1, _1_router_logits + ) + del _1_input1, _1_router_logits infer_state1.hook = _1_hook # moe calu expected_m = triton.cdiv( input_embdings.shape[0] * get_global_world_size() * self.num_experts_per_tok, self.n_routed_experts ) - _0_moe_out = layer_weight.experts.masked_group_gemm(_0_recv_x, _0_masked_m, input_embdings.dtype, expected_m) + _0_moe_out = layer_weight.experts.decode_masked_group_gemm( + _0_recv_x, _0_handle, input_embdings.dtype, expected_m + ) + del _0_recv_x # 1 hook if getattr(infer_state1, "hook", None) is not None: @@ -212,25 +207,27 @@ def overlap_tpsp_token_forward( infer_state1.hook = None # 0 combine - _0_ffn_out, _0_hook = layer_weight.experts.low_latency_combine( - _0_moe_out, _0_topk_idx, _0_topk_weight, _0_handle - ) + _0_ffn_out, _0_hook = layer_weight.experts.decode_combine(_0_moe_out, _0_handle, _0_recv_topk_weights) + del _0_moe_out infer_state.hook = _0_hook # to do moe caclue - _1_moe_out = layer_weight.experts.masked_group_gemm(_1_recv_x, _1_masked_m, input_embdings1.dtype, expected_m) + _1_moe_out = layer_weight.experts.decode_masked_group_gemm( + _1_recv_x, _1_handle, input_embdings1.dtype, expected_m + ) + del _1_recv_x # 0 hook if getattr(infer_state, "hook", None) is not None: infer_state.hook() input_embdings.add_(_0_ffn_out.view(-1, self.embed_dim_)) + del _0_ffn_out infer_state.hook = None # 1 combine - _1_ffn_out, _1_hook = layer_weight.experts.low_latency_combine( - _1_moe_out, _1_topk_idx, _1_topk_weight, _1_handle - ) + _1_ffn_out, _1_hook = layer_weight.experts.decode_combine(_1_moe_out, _1_handle, _1_recv_topk_weights) + del _1_moe_out def _1_hook_post(): _1_hook() @@ -275,6 +272,7 @@ def overlap_tpsp_context_forward( _0_topk_weight, _0_topk_idx, _0_qinput_tensor = layer_weight.experts.select_experts_and_quant_input( _0_input1, _0_router_logits ) + del _0_router_logits, _0_input1 from deep_ep import ElasticBuffer _0_overlap_event = ElasticBuffer.capture() @@ -302,7 +300,10 @@ def overlap_tpsp_context_forward( _0_num_recv_tokens_per_expert_list, _0_handle, _0_hook, - ) = layer_weight.experts.dispatch(_0_qinput_tensor, _0_topk_idx, _0_topk_weight, overlap_event=_0_overlap_event) + ) = layer_weight.experts.prefill_dispatch( + _0_qinput_tensor, _0_topk_idx, _0_topk_weight, overlap_event=_0_overlap_event + ) + del _0_qinput_tensor, _0_topk_idx, _0_topk_weight infer_state.hook = _0_hook # wait 0 dispatch @@ -313,6 +314,7 @@ def overlap_tpsp_context_forward( _1_topk_weight, _1_topk_idx, _1_qinput_tensor = layer_weight.experts.select_experts_and_quant_input( _1_input1, _1_router_logits ) + del _1_router_logits, _1_input1 _1_overlap_event = ElasticBuffer.capture() # 0 moe calu @@ -325,6 +327,7 @@ def overlap_tpsp_context_forward( _0_recv_topk_weight, microbatch_index=0, ) + del _0_recv_x # 1 dispatch execute ( @@ -334,7 +337,10 @@ def overlap_tpsp_context_forward( _1_num_recv_tokens_per_expert_list, _1_handle, _1_hook, - ) = layer_weight.experts.dispatch(_1_qinput_tensor, _1_topk_idx, _1_topk_weight, overlap_event=_1_overlap_event) + ) = layer_weight.experts.prefill_dispatch( + _1_qinput_tensor, _1_topk_idx, _1_topk_weight, overlap_event=_1_overlap_event + ) + del _1_qinput_tensor, _1_topk_idx, _1_topk_weight infer_state1.hook = _1_hook # wait 1 dispatch @@ -344,7 +350,8 @@ def overlap_tpsp_context_forward( _0_combine_event = ElasticBuffer.capture() # 0 combine execute - _0_ffn_out, _0_hook = layer_weight.experts.combine(_0_moe_out, _0_handle, _0_combine_event) + _0_ffn_out, _0_hook = layer_weight.experts.prefill_combine(_0_moe_out, _0_handle, _0_combine_event) + del _0_moe_out infer_state.hook = _0_hook # 1 moe calc @@ -357,6 +364,7 @@ def overlap_tpsp_context_forward( _1_recv_topk_weight, microbatch_index=1, ) + del _1_recv_x # wait 0 combine if getattr(infer_state, "hook", None) is not None: @@ -368,7 +376,7 @@ def overlap_tpsp_context_forward( input_embdings.add_(_0_ffn_out.view(-1, self.embed_dim_)) # 1 combine execute - _1_ffn_out, _1_hook = layer_weight.experts.combine(_1_moe_out, _1_handle, _1_combine_event) + _1_ffn_out, _1_hook = layer_weight.experts.prefill_combine(_1_moe_out, _1_handle, _1_combine_event) def _1_hook_post(): _1_hook() diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 3e2fb44b5b..a42acef884 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -169,6 +169,12 @@ def _launch_subprocesses(args: StartArgs): if args.enable_prefill_microbatch_overlap or args.enable_decode_microbatch_overlap: args.enable_tpsp_mix_mode = True + if args.enable_ep_moe and "PYTORCH_ALLOC_CONF" not in os.environ: + # 在启动模型子进程前设置，使 CUDA allocator 第一次初始化时即可 + # 使用 CUDA Graph 拓扑分析 record_stream 张量的真实依赖关系。在确认 + # 前一张量已经完成所有跨流访问后，Graph 私有内存池可以更早复用它的 + # 显存，避免 overlap 路径的临时张量在多层捕获过程中重复保留。 + os.environ["PYTORCH_ALLOC_CONF"] = "graph_capture_record_stream_reuse:True" if args.enable_prefill_decode_mixed: assert args.run_mode == "normal", "--enable_prefill_decode_mixed only supports run_mode normal" @@ -297,6 +303,22 @@ def _launch_subprocesses(args: StartArgs): f"but got {args.batch_max_tokens}, {args.chunked_prefill_size}" ) + if args.run_mode == "decode": + # Decode 节点的启动 warmup 仍会构造 batch_max_tokens 个 token 做一次 + # prefill-style model forward。Decode 节点没有必要沿用较大的 prefill + # batch 上限，但需要覆盖每个请求的一个 target token 和 mtp_step 个 + # draft token。 + decode_batch_max_tokens = int(args.running_max_req_size) * (1 + int(args.mtp_step)) + logger.info( + "Decode node overrides batch_max_tokens: configured=%s, running_max_req_size=%s, " + "mtp_step=%s, decode_batch_max_tokens=%s", + args.batch_max_tokens, + args.running_max_req_size, + args.mtp_step, + decode_batch_max_tokens, + ) + args.batch_max_tokens = decode_batch_max_tokens + args.disable_chunked_mm_input = args.disable_chunked_mm_input or get_default_disable_chunked_mm_input( args.model_dir ) diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index d3ef067c9b..9d8490b69e 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -85,21 +85,53 @@ def is_stop_matching_in_reasoning_enabled() -> bool: @lru_cache(maxsize=None) def get_deepep_num_max_dispatch_tokens_per_rank_prefill(): - # 该参数需要大于单卡最大batch size，且是8的倍数。该参数与显存占用直接相关，值越大，显存占用越大。 - # 如果未显式配置，则默认至少覆盖当前进程的 `batch_max_tokens`，避免 DeepEP V2 在 autotune - # warmup 或大 prefill batch 时因为 buffer 上界过小而报错。 + # 在单卡最大 prefill batch 之外额外保留 128 个 token，并向上对齐到 8， + # 避免 autotune warmup 或调度边界波动使实际输入超过 DeepEP buffer 上限。 + # 显式设置 NUM_MAX_DISPATCH_TOKENS_PER_RANK_PREFILL 时仍使用该覆盖值。 configured = os.getenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_PREFILL", None) if configured is not None: return int(configured) batch_max_tokens = get_env_start_args().batch_max_tokens or 256 - return ((int(batch_max_tokens) + 7) // 8) * 8 + capacity = ((int(batch_max_tokens) + 128 + 7) // 8) * 8 + logger.info( + "DeepEP prefill buffer capacity: batch_max_tokens=%s, safety_margin=128, capacity=%s", + batch_max_tokens, + capacity, + ) + return capacity @lru_cache(maxsize=None) def get_deepep_num_max_dispatch_tokens_per_rank_decode(): - # 该参数需要大于单卡最大batch size，且是8的倍数。该参数与显存占用直接相关，值越大，显存占用越大，如果出现显存不足，可以尝试调小该值 - return int(os.getenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", 256)) + # 每个请求最多产生 mtp_step + 1 个 verify token；额外保留 12 个 token + # 处理调度和 CUDA Graph 的边界波动，最后向上对齐到 DeepEP 要求的 8。 + # 显式设置 NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE 时仍使用该覆盖值。 + configured = os.getenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", None) + if configured is not None: + return int(configured) + + args = get_env_start_args() + required_tokens = int(args.running_max_req_size) * (int(args.mtp_step) + 1) + capacity = ((required_tokens + 12 + 7) // 8) * 8 + logger.info( + "DeepEP decode buffer capacity: running_max_req_size=%s, mtp_step=%s, " + "required_tokens=%s, safety_margin=12, capacity=%s", + args.running_max_req_size, + args.mtp_step, + required_tokens, + capacity, + ) + return capacity + + +@lru_cache(maxsize=None) +def get_deepep_num_max_dispatch_tokens_per_rank() -> int: + """返回同时覆盖 prefill 和 decode 的单 rank DeepEP buffer 容量。""" + return max( + get_deepep_num_max_dispatch_tokens_per_rank_prefill(), + get_deepep_num_max_dispatch_tokens_per_rank_decode(), + ) @lru_cache(maxsize=None) diff --git a/unit_tests/common/fused_moe/test_deepgemm_activation.py b/unit_tests/common/fused_moe/test_deepgemm_activation.py index e9ba2ab26e..3862e2ba5f 100644 --- a/unit_tests/common/fused_moe/test_deepgemm_activation.py +++ b/unit_tests/common/fused_moe/test_deepgemm_activation.py @@ -39,7 +39,7 @@ def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ metadata = torch.zeros(5, 3, device="cuda", dtype=torch.int32) metadata[:, 2] = torch.tensor([0, 1, 2, padded, padded + 1], device="cuda") # Force two chunks so both expert groups execute the requested activation. - monkeypatch.setattr(ep, "_get_max_chunk_rows", lambda **kwargs: padded) + monkeypatch.setattr(ep.deepgemm, "_get_max_chunk_rows", lambda **kwargs: padded) actual = ep.chunked_expanded_moe_forward( num_recv_tokens_per_expert_list=[padded, padded], num_unaligned_recv_tokens_per_expert=counts, diff --git a/unit_tests/common/fused_moe/test_ep_reduce_decode_output.py b/unit_tests/common/fused_moe/test_ep_reduce_decode_output.py new file mode 100644 index 0000000000..8b2805e3a3 --- /dev/null +++ b/unit_tests/common/fused_moe/test_ep_reduce_decode_output.py @@ -0,0 +1,36 @@ +import pytest +import torch + +from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_expanded_layout_kernels import ep_reduce_decode_output + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def test_ep_reduce_decode_output_weights_local_rows_and_compacts_metadata(): + hidden = 8 + expert_output = torch.tensor( + [[1.0] * hidden, [2.0] * hidden, [4.0] * hidden], + device="cuda", + dtype=torch.bfloat16, + ) + route_weights = torch.tensor([0.5, 0.25, 0.0], device="cuda", dtype=torch.float32) + recv_src_metadata = torch.tensor( + [[7, 9, 0, 1], [3, 4, -1, -1]], + device="cuda", + dtype=torch.int32, + ) + num_valid_recv_tokens = torch.tensor([1], device="cuda", dtype=torch.int32) + + dense_output, compact_metadata = ep_reduce_decode_output( + expert_output=expert_output, + route_weights=route_weights, + recv_src_metadata=recv_src_metadata, + num_valid_recv_tokens=num_valid_recv_tokens, + ) + + expected = expert_output[0].float() * 0.5 + expert_output[1].float() * 0.25 + torch.testing.assert_close(dense_output[0].float(), expected) + assert compact_metadata[0, :2].tolist() == [7, 9] + assert compact_metadata[0, 2].item() == 0 + assert compact_metadata[0, 3].item() == -1 diff --git a/unit_tests/common/fused_moe/test_grouped_fused_moe_ep_sms.py b/unit_tests/common/fused_moe/test_grouped_fused_moe_ep_sms.py new file mode 100644 index 0000000000..35b6ce7104 --- /dev/null +++ b/unit_tests/common/fused_moe/test_grouped_fused_moe_ep_sms.py @@ -0,0 +1,11 @@ +from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep.common import get_ep_num_sms + + +def test_get_ep_num_sms_selects_overlap_quota(monkeypatch): + import lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep.common as common + + monkeypatch.setattr(common.dist_group_manager, "ep_num_sms", 20, raising=False) + monkeypatch.setattr(common.dist_group_manager, "ep_non_overlap_num_sms", 64, raising=False) + + assert get_ep_num_sms(overlap_with_compute=True) == 20 + assert get_ep_num_sms(overlap_with_compute=False) == 64 diff --git a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py b/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py index 591a62f75a..b569ec8429 100644 --- a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py +++ b/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py @@ -15,6 +15,7 @@ def is_fp8_native_supported(): from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul_mix_quant_ep import ( silu_and_mul_masked_post_quant_fwd, + silu_and_mul_psum_post_quant_fwd, ) from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 @@ -151,5 +152,45 @@ def test_clamped_swiglu_quantization_matches_reference(monkeypatch, dtype, use_u assert torch.all(scales[expert, count:] == 7) +def test_silu_and_mul_psum_skips_alignment_padding(): + hidden = 128 + alignment = 8 + # expert0 使用 3 行，expert1 从对齐后的第 8 行开始再使用 2 行。 + expert_token_psum = torch.tensor([3, 10], device="cuda", dtype=torch.int32) + rows = 16 + gate = torch.randn((rows, hidden), device="cuda", dtype=torch.bfloat16) + up = torch.randn((rows, hidden), device="cuda", dtype=torch.bfloat16) + packed = torch.cat([gate, up], dim=1).contiguous() + output = torch.zeros((rows, hidden), device="cuda", dtype=torch.float8_e4m3fn) + scales = torch.zeros((rows, hidden // 128), device="cuda", dtype=torch.float32) + + silu_and_mul_psum_post_quant_fwd( + input=packed, + output=output, + output_scale=scales, + expert_token_psum=expert_token_psum, + expert_alignment=alignment, + quant_group_size=128, + ) + + def reference(row): + # Match the kernel: SiLU in fp32, then round the gate back to bf16 before multiplying. + activated_gate = torch.nn.functional.silu(gate[row].float()).to(torch.bfloat16) + activated = activated_gate.float() * up[row].float() + scale = activated.abs().amax().clamp(min=1e-10) / torch.finfo(torch.float8_e4m3fn).max + quantized = (activated / scale).clamp(-448, 448).to(torch.float8_e4m3fn) + return quantized, scale + + for row in (0, 1, 2, 8, 9): + expected_q, expected_scale = reference(row) + torch.testing.assert_close(scales[row], expected_scale.view(1), atol=1e-4, rtol=1e-3) + actual = output[row].float() * scales[row] + expected = expected_q.float() * expected_scale + # FP8 E4M3 bins are coarse, so compare the dequantized activation within one bin. + torch.testing.assert_close(actual, expected, atol=0.25, rtol=0.15) + # expert0 与 expert1 之间的对齐 padding 不应被写成有效值。 + assert torch.count_nonzero(output[3:8].float()) == 0 + + if __name__ == "__main__": pytest.main() diff --git a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py index c08a8a209b..527ad4e3ae 100644 --- a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py +++ b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py @@ -61,14 +61,14 @@ def record_call(*args, **kwargs): assert scales.dtype == (torch.int32 if use_packed_ue8m0 else torch.float32) if use_ue8m0_scales: actual = _unpack_scales(scales, 4) if use_packed_ue8m0 else scales - torch.testing.assert_close(actual, torch.full_like(actual, 2.0 ** -8), rtol=0, atol=0) + torch.testing.assert_close(actual, torch.full_like(actual, 2.0**-8), rtol=0, atol=0) @pytest.mark.parametrize("use_packed_ue8m0", [False, True]) def test_ue8m0_rounding_at_power_of_two_boundaries(monkeypatch, use_packed_ue8m0): monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr(activation, "HAS_SGL_KERNEL", False) - powers = torch.tensor([2.0 ** exponent for exponent in (-16, -8, 0, 8, 16, 118)], device="cuda") + powers = torch.tensor([2.0**exponent for exponent in (-16, -8, 0, 8, 16, 118)], device="cuda") above = torch.nextafter(powers, torch.full_like(powers, float("inf"))) below = torch.nextafter(powers, torch.zeros_like(powers)) magnitudes = torch.cat((powers, above, below, powers.new_zeros(1))) * 448.0 @@ -152,7 +152,7 @@ def test_packed_ue8m0_scales(monkeypatch, rows, groups, group_size, use_sgl): x[0].zero_() if rows > 1: x[1].fill_(1e-12) - x[-1, -group_size:].fill_(2.0 ** 16) + x[-1, -group_size:].fill_(2.0**16) q, packed_scales = activation.per_token_group_quant_fp8(x, group_size, use_ue8m0_scales=True, use_packed_ue8m0=True) assert packed_scales.dtype == torch.int32 @@ -212,9 +212,9 @@ def check_gemm_inputs(a, b, out): seen.append(a[1].dtype) assert b[1].dtype == (torch.int32 if sm100 else torch.float32) scales = _unpack_scales(a[1], 4) if sm100 else a[1] - torch.testing.assert_close(scales, torch.full_like(scales, 2.0 ** -8), rtol=0, atol=0) + torch.testing.assert_close(scales, torch.full_like(scales, 2.0**-8), rtol=0, atol=0) weight_scales = _unpack_scales(b[1], 4) if sm100 else b[1] - torch.testing.assert_close(weight_scales, torch.full_like(weight_scales, 2.0 ** -8), rtol=0, atol=0) + torch.testing.assert_close(weight_scales, torch.full_like(weight_scales, 2.0**-8), rtol=0, atol=0) out.zero_() # Exercise format selection and real quantization on Hopper; SM100 GEMM needs a separate machine. @@ -500,7 +500,7 @@ def check_dispatch(**kwargs): ) return x - monkeypatch.setattr(ep, "fused_experts_impl", check_dispatch) + monkeypatch.setattr(ep.prefill_decode_sync, "fused_experts_impl", check_dispatch) ep.fused_experts( x, w13, diff --git a/unit_tests/server/test_api_start_deepep.py b/unit_tests/server/test_api_start_deepep.py new file mode 100644 index 0000000000..949021cd34 --- /dev/null +++ b/unit_tests/server/test_api_start_deepep.py @@ -0,0 +1,45 @@ +import json + +import pytest + +from lightllm.server.api_start import _launch_subprocesses +from lightllm.server.core.objs.start_args_type import StartArgs + + +def test_decode_node_batch_max_tokens_covers_mtp_rows(monkeypatch, tmp_path): + for name in ( + "_set_envs_and_config", + "auto_set_max_req_total_len", + "auto_set_fused_shared_experts", + "set_unique_server_name", + "auto_set_page_size", + ): + monkeypatch.setattr(f"lightllm.server.api_start.{name}", lambda args: None) + + def stop(args): + raise RuntimeError("startup validation passed") + + monkeypatch.setattr("lightllm.server.api_start.auto_set_response_parsers", stop) + (tmp_path / "config.json").write_text(json.dumps({"model_type": "llama"})) + args = StartArgs( + run_mode="decode", + model_dir=str(tmp_path), + enable_ep_moe=True, + running_max_req_size=256, + mtp_mode="vanilla_no_att", + mtp_step=2, + max_req_total_len=4096, + batch_max_tokens=4096, + chunked_prefill_size=2048, + eos_id=[0], + data_type="float16", + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + router_token_ratio=0.0, + ) + + with pytest.raises(RuntimeError, match="startup validation passed"): + _launch_subprocesses(args) + + assert args.batch_max_tokens == 768 diff --git a/unit_tests/utils/test_envs_utils.py b/unit_tests/utils/test_envs_utils.py index a06e779392..f3b766f41d 100644 --- a/unit_tests/utils/test_envs_utils.py +++ b/unit_tests/utils/test_envs_utils.py @@ -1,4 +1,10 @@ +from types import SimpleNamespace + +from lightllm.utils import envs_utils from lightllm.utils.envs_utils import ( + get_deepep_num_max_dispatch_tokens_per_rank, + get_deepep_num_max_dispatch_tokens_per_rank_decode, + get_deepep_num_max_dispatch_tokens_per_rank_prefill, is_stop_matching_in_reasoning_enabled, get_pd_cache_high_priority_max_age_seconds, get_pd_cache_high_priority_min_prompt_tokens, @@ -111,3 +117,63 @@ def test_pd_node_busy_retry_timeout_reads_environment_variable(monkeypatch): assert get_pd_node_busy_retry_timeout_seconds() == 60 get_pd_node_busy_retry_timeout_seconds.cache_clear() + + +def test_deepep_decode_buffer_capacity_covers_mtp_rows_and_safety_margin(monkeypatch): + monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", raising=False) + monkeypatch.setattr( + envs_utils, + "get_env_start_args", + lambda: SimpleNamespace(running_max_req_size=256, mtp_step=2), + ) + get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() + + assert get_deepep_num_max_dispatch_tokens_per_rank_decode() == 784 + + get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() + + +def test_deepep_prefill_buffer_capacity_adds_safety_margin(monkeypatch): + monkeypatch.delenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_PREFILL", raising=False) + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: SimpleNamespace(batch_max_tokens=6000)) + get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() + + assert get_deepep_num_max_dispatch_tokens_per_rank_prefill() == 6128 + + get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() + + +def test_deepep_buffer_capacity_uses_larger_phase_capacity(monkeypatch): + cases = [ + (4096, 256, 4096), + (256, 512, 512), + ] + + for prefill_capacity, decode_capacity, expected_capacity in cases: + monkeypatch.setattr( + envs_utils, + "get_deepep_num_max_dispatch_tokens_per_rank_prefill", + lambda prefill_capacity=prefill_capacity: prefill_capacity, + ) + monkeypatch.setattr( + envs_utils, + "get_deepep_num_max_dispatch_tokens_per_rank_decode", + lambda decode_capacity=decode_capacity: decode_capacity, + ) + get_deepep_num_max_dispatch_tokens_per_rank.cache_clear() + assert get_deepep_num_max_dispatch_tokens_per_rank() == expected_capacity + + get_deepep_num_max_dispatch_tokens_per_rank.cache_clear() + + +def test_deepep_buffer_capacity_keeps_explicit_environment_override(monkeypatch): + monkeypatch.setenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_PREFILL", "1024") + monkeypatch.setenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", "128") + get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() + get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear() + + assert get_deepep_num_max_dispatch_tokens_per_rank_prefill() == 1024 + assert get_deepep_num_max_dispatch_tokens_per_rank_decode() == 128 + + get_deepep_num_max_dispatch_tokens_per_rank_prefill.cache_clear() + get_deepep_num_max_dispatch_tokens_per_rank_decode.cache_clear()