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 7a03a67418..8d3c17b263 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 @@ -8,7 +8,7 @@ get_col_slice_mixin, SliceMixinTpl, ) -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import select_fuse_moe_impl +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import create_fuse_moe_impl from lightllm.common.basemodel.moe_route_info_manager import get_moe_capture_callback from lightllm.common.quantization.quantize_method import QuantizationMethod from lightllm.utils.envs_utils import get_redundancy_expert_ids, get_redundancy_expert_num, get_env_start_args @@ -57,17 +57,19 @@ def __init__( self.num_fused_shared_experts = num_fused_shared_experts self._init_config(network_config) self._init_redundancy_expert_params() - self._init_parallel_params() - self.fuse_moe_impl = select_fuse_moe_impl(self.quant_method, self.enable_ep_moe)( + self.fuse_moe_impl = create_fuse_moe_impl( n_routed_experts=self.n_routed_experts, num_fused_shared_experts=self.num_fused_shared_experts, routed_scaling_factor=self.routed_scaling_factor, quant_method=self.quant_method, - redundancy_expert_num=self.redundancy_expert_num, - redundancy_expert_ids_tensor=self.redundancy_expert_ids_tensor, - routed_expert_counter_tensor=self.routed_expert_counter_tensor, - auto_update_redundancy_expert=self.auto_update_redundancy_expert, + enable_ep_moe=self.enable_ep_moe, + layer_index=self.layer_num_, ) + # 冗余专家的 CUDA 状态由 DeepGEMM 实现持有。自动更新仍通过权重对象读写同一份 tensor。 + if self.enable_ep_moe: + self.redundancy_expert_ids_tensor = self.fuse_moe_impl.redundancy_expert_ids_tensor + self.routed_expert_counter_tensor = self.fuse_moe_impl.routed_expert_counter_tensor + self._init_weight_partition() self.lock = threading.Lock() self._create_weight() @@ -84,12 +86,10 @@ def _init_redundancy_expert_params(self): self.redundancy_expert_num = get_redundancy_expert_num() self.redundancy_expert_ids = get_redundancy_expert_ids(self.layer_num_) self.auto_update_redundancy_expert: bool = get_env_start_args().auto_update_redundancy_expert - self.redundancy_expert_ids_tensor = torch.tensor(self.redundancy_expert_ids, dtype=torch.int64, device="cuda") - self.routed_expert_counter_tensor = torch.zeros((self.n_routed_experts,), dtype=torch.int64, device="cuda") # TODO: find out the reason of failure of deepep when redundancy_expert_num is 1. assert self.redundancy_expert_num != 1, "redundancy_expert_num can not be 1 for some unknown hang of deepep." - def _init_parallel_params(self): + def _init_weight_partition(self): if self.enable_ep_moe: self.tp_rank_ = 0 self.tp_world_size_ = 1 @@ -103,27 +103,14 @@ def _init_parallel_params(self): self.split_inter_size = self.moe_intermediate_size // self.tp_world_size_ if self.enable_ep_moe: assert self.num_fused_shared_experts == 0, "num_fused_shared_experts must be 0 when enable_ep_moe" + self.local_logic_expert_ids_list = self.fuse_moe_impl.local_logics_expert_ids_list logger.debug( f"global_rank {self.global_rank_} layerindex {self.layer_num_} " - f"redundancy_expertids: {self.redundancy_expert_ids}" - ) - self.local_n_routed_experts = self.n_routed_experts // self.global_world_size + self.redundancy_expert_num - n_experts_per_rank = self.n_routed_experts // self.global_world_size - start_expert_id = self.global_rank_ * n_experts_per_rank - self.local_expert_ids = ( - list(range(start_expert_id, start_expert_id + n_experts_per_rank)) + self.redundancy_expert_ids + f"local_logic_expert_ids_list: {self.local_logic_expert_ids_list}" ) - self.expert_idx_to_local_idx = { - expert_idx: expert_idx - start_expert_id for expert_idx in self.local_expert_ids[:n_experts_per_rank] - } - self.redundancy_expert_idx_to_local_idx = { - redundancy_expert_idx: n_experts_per_rank + i - for (i, redundancy_expert_idx) in enumerate(self.redundancy_expert_ids) - } + self.local_n_routed_experts = len(self.local_logic_expert_ids_list) else: - self.local_expert_ids = list(range(self.n_routed_experts + self.num_fused_shared_experts)) - self.expert_idx_to_local_idx = {expert_idx: i for (i, expert_idx) in enumerate(self.local_expert_ids)} - self.rexpert_idx_to_local_idx = {} + self.local_logic_expert_ids_list = list(range(self.n_routed_experts + self.num_fused_shared_experts)) def experts( self, @@ -134,13 +121,14 @@ def experts( use_grouped_topk: bool, topk_group: int, num_expert_group: int, - is_prefill: Optional[bool] = None, + is_prefill: bool, infer_state=None, shared_expert_gate: Optional[torch.Tensor] = None, alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, ) -> torch.Tensor: + assert is_prefill is not None, "is_prefill must be explicitly specified for fused MoE execution" # Captures MoE topk expert ids for routed-experts metadata when enabled. moe_capture_callback = get_moe_capture_callback(infer_state, self.layer_num_) return self.fuse_moe_impl( @@ -169,27 +157,22 @@ def experts_with_topk( input_tensor: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, - is_prefill: Optional[bool] = None, + is_prefill: bool, infer_state=None, alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, shared_expert_func=None, ) -> torch.Tensor: + assert is_prefill is not None, "is_prefill must be explicitly specified for fused MoE execution" moe_capture_callback = get_moe_capture_callback(infer_state, self.layer_num_) if moe_capture_callback is not None: moe_capture_callback(topk_ids) - if self.num_fused_shared_experts > 0: - from lightllm.common.basemodel.triton_kernel.fused_moe.append_shared_expert_topk import ( - append_fused_shared_experts, - ) - - topk_weights, topk_ids = append_fused_shared_experts( - topk_weights=topk_weights, - topk_ids=topk_ids, - shared_expert_start_id=self.n_routed_experts, - num_fused_shared_experts=self.num_fused_shared_experts, - ) + topk_weights, topk_ids = self.fuse_moe_impl._prepare_expert_execution( + topk_weights=topk_weights, + topk_ids=topk_ids, + is_prefill=is_prefill, + ) # The shared-expert callback is only implemented on the DeepGEMM path. run = self.fuse_moe_impl.fused_experts_with_topk extra = {} @@ -209,16 +192,17 @@ def experts_with_topk( **extra, ) - 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, @@ -227,15 +211,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, ) @@ -263,37 +248,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, @@ -310,7 +295,6 @@ def prefilled_group_gemm( recv_topk_idx: torch.Tensor, recv_topk_weights: torch.Tensor, hidden_dtype=torch.bfloat16, - microbatch_index: int = 0, alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, @@ -326,35 +310,32 @@ def prefilled_group_gemm( w13=self.w13, w2=self.w2, hidden_dtype=hidden_dtype, - microbatch_index=microbatch_index, alpha=alpha, limit=limit, 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, @@ -364,9 +345,7 @@ def load_hf_weights(self, weights): # Load bias self._load_e_score_correction_bias(weights) self._load_per_expert_scale(weights) - self._load_weight(self.expert_idx_to_local_idx, weights) - if self.redundancy_expert_num > 0: - self._load_weight(self.redundancy_expert_idx_to_local_idx, weights) + self._load_weight(self.local_logic_expert_ids_list, weights) def verify_load(self): weight_load_ok = all(all(_weight_pack.load_ok) for _weight_pack in self.w1_list + self.w2_list + self.w3_list) @@ -435,8 +414,8 @@ def _get_expert_weight_list(self, weight_pack: WeightPack): weight_list.append(expert_weight) return weight_list - def _load_weight(self, expert_idx_to_local_idx: Dict[int, int], weights: Dict[str, torch.Tensor]): - for expert_idx, local_expert_idx in expert_idx_to_local_idx.items(): + def _load_weight(self, local_logic_expert_ids_list: List[int], weights: Dict[str, torch.Tensor]): + for local_expert_idx, expert_idx in enumerate(local_logic_expert_ids_list): with self.lock: self._load_expert(expert_idx, local_expert_idx, weights) self._load_expert_scale( diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py index 67bb90e4ef..05cfc82a17 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py @@ -1,14 +1,35 @@ +from typing import Optional + from lightllm.common.quantization.quantize_method import QuantizationMethod from .triton_impl import FuseMoeTriton from .marlin_impl import FuseMoeMarlin from .deepgemm_impl import FuseMoeDeepGEMM -def select_fuse_moe_impl(quant_method: QuantizationMethod, enable_ep_moe: bool): - if enable_ep_moe: - return FuseMoeDeepGEMM +def create_fuse_moe_impl( + n_routed_experts: int, + num_fused_shared_experts: int, + routed_scaling_factor: float, + quant_method: QuantizationMethod, + enable_ep_moe: bool = False, + layer_index: Optional[int] = None, +): + """创建持有自身路由运行态的 MoE 执行实现。 - if quant_method.method_name == "awq_marlin": - return FuseMoeMarlin + 直接返回完成初始化的对象，而不是仅返回实现类，使冗余专家计数器等后端专属 + 状态与使用它们的 kernel 保持在同一个实现对象中。 + """ + if enable_ep_moe: + impl_cls = FuseMoeDeepGEMM + elif quant_method.method_name == "awq_marlin": + impl_cls = FuseMoeMarlin else: - return FuseMoeTriton + impl_cls = FuseMoeTriton + + return impl_cls( + n_routed_experts=n_routed_experts, + num_fused_shared_experts=num_fused_shared_experts, + routed_scaling_factor=routed_scaling_factor, + quant_method=quant_method, + layer_index=layer_index, + ) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py index e242870d69..8acaf6efb6 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py @@ -1,53 +1,53 @@ import torch -from abc import abstractmethod -from typing import Callable, Optional +from abc import ABC, abstractmethod +from typing import Callable, Optional, Tuple from lightllm.common.quantization.quantize_method import ( WeightPack, QuantizationMethod, ) -from lightllm.utils.dist_utils import ( - get_global_rank, - get_global_world_size, -) -class FuseMoeBaseImpl: +class FuseMoeBaseImpl(ABC): + """将逻辑专家路由与实际执行布局分离。 + + 融合 MoE 的调用流程如下:: + + _select_experts + -> topk_weights + logical_topk_ids + -> moe_capture_callback(logical_topk_ids) + -> _prepare_expert_execution + -> 追加 shared expert，或者 + -> 将逻辑专家 ID 修复为包含冗余专家的执行 ID + -> _fused_experts(topk_weights, execution_topk_ids) + + ``_select_experts`` 对所有实现都遵循同一套稳定接口：只在模型的逻辑专家 + 空间中选择专家，并返回原始 logical ID。该阶段不能应用任何与实际执行相关 + 的布局转换，例如追加 shared expert 行，或者改写冗余专家 ID。 + + capture callback 在任何布局转换之前执行，因此采集到的路由元数据始终描述 + 模型的 logical expert。所有实现都必须将 logical ID 视为只读数据。 + ``_prepare_expert_execution`` 可以为实际执行布局分配新的 ID tensor，但不能 + 原地修改 logical ID tensor。 + + 因此，``_fused_experts`` 接收的是 execution ID：普通路径中仍是未经修改的 + logical ID；融合 shared expert 路径中是追加了 shared expert 行的 ID；冗余 + 专家路径中则是完成修复后的执行 ID。 + """ + def __init__( self, n_routed_experts: int, num_fused_shared_experts: int, routed_scaling_factor: float, quant_method: QuantizationMethod, - redundancy_expert_num: int, - redundancy_expert_ids_tensor: torch.Tensor, - routed_expert_counter_tensor: torch.Tensor, - auto_update_redundancy_expert: bool, + layer_index: Optional[int] = None, ): self.n_routed_experts = n_routed_experts self.num_fused_shared_experts = num_fused_shared_experts self.routed_scaling_factor = routed_scaling_factor self.quant_method = quant_method - self.global_rank_ = get_global_rank() - self.global_world_size_ = get_global_world_size() - self.ep_n_routed_experts = self.n_routed_experts // self.global_world_size_ - self.total_expert_num_contain_redundancy = ( - self.n_routed_experts + redundancy_expert_num * self.global_world_size_ - ) + self.layer_index = layer_index - # redundancy expert related - self.redundancy_expert_num = redundancy_expert_num - self.redundancy_expert_ids_tensor = redundancy_expert_ids_tensor - self.routed_expert_counter_tensor = routed_expert_counter_tensor - self.auto_update_redundancy_expert = auto_update_redundancy_expert - - # workspace for kernel optimization - self.workspace = self.create_workspace() - - @abstractmethod - def create_workspace(self): - pass - - @abstractmethod def __call__( self, input_tensor: torch.Tensor, @@ -61,14 +61,147 @@ def __call__( use_grouped_topk: bool, topk_group: int, num_expert_group: int, - is_prefill: Optional[bool] = None, + is_prefill: bool, # Callback to capture MoE topk expert ids (routed experts metadata). moe_capture_callback: Optional[Callable[[torch.Tensor], None]] = None, per_expert_scale: Optional[torch.Tensor] = None, - # Qwen3.5 uses this gate to control fused shared expert aggregation weights. + # Qwen3Next/Qwen3.5-MoE 在 TP 模式下将 shared expert 融合进 routed MoE。 + # 该参数是 shared_expert_gate(hidden_states) 产生的逐 token 门控 logit； + # 追加 shared expert 时使用 sigmoid(logit) 作为其聚合权重。 shared_expert_gate: Optional[torch.Tensor] = None, alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, ) -> torch.Tensor: + assert is_prefill is not None, "is_prefill must be explicitly specified for fused MoE execution" + topk_weights, topk_ids = self._select_experts( + input_tensor=input_tensor, + router_logits=router_logits, + correction_bias=correction_bias, + top_k=top_k, + renormalize=renormalize, + use_grouped_topk=use_grouped_topk, + topk_group=topk_group, + num_expert_group=num_expert_group, + scoring_func=scoring_func, + per_expert_scale=per_expert_scale, + ) + if moe_capture_callback is not None: + moe_capture_callback(topk_ids) + topk_weights, topk_ids = self._prepare_expert_execution( + topk_weights=topk_weights, + topk_ids=topk_ids, + shared_expert_gate=shared_expert_gate, + is_prefill=is_prefill, + ) + return self._fused_experts( + input_tensor=input_tensor, + w13=w13, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + is_prefill=is_prefill, + router_logits=router_logits, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + ) + + def fused_experts_with_topk( + self, + input_tensor: torch.Tensor, + w13: WeightPack, + w2: WeightPack, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + is_prefill: bool, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, + ): + return self._fused_experts( + input_tensor=input_tensor, + w13=w13, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + is_prefill=is_prefill, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + ) + + @abstractmethod + def _select_experts( + self, + input_tensor: torch.Tensor, + router_logits: torch.Tensor, + correction_bias: Optional[torch.Tensor], + top_k: int, + renormalize: bool, + use_grouped_topk: bool, + topk_group: int, + num_expert_group: int, + scoring_func: str, + per_expert_scale: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """在模型的逻辑专家空间中完成 top-k 路由选择。 + + 返回形状一致的 ``topk_weights`` 和 ``logical_topk_ids``。这里的 + ``topk_ids`` 必须始终表示模型配置中的原始 logical expert，不能追加 + shared expert，也不能改写冗余专家 ID。返回的 ID tensor 会先交给 + capture callback，后续实现必须将其视为只读数据。 + """ + pass + + @abstractmethod + def _prepare_expert_execution( + self, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + is_prefill: bool, + shared_expert_gate: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """将逻辑路由结果转换为 MoE kernel 实际需要的执行布局。 + + 该方法在 capture callback 之后执行。每个实现都必须明确处理自己的执行 + 布局：普通路径保持权重和 logical ID 不变，shared-expert 路径追加对应 + expert，冗余专家路径将 logical ID 修复为执行 ID。发生布局转换时应返回 + 新的 ID tensor，不能原地修改传入的 logical ``topk_ids``；如果专家数量 + 发生变化，``topk_weights`` 必须同步调整。 + + ``shared_expert_gate`` 当前由 Qwen3Next/Qwen3.5-MoE 使用，其内容是 + ``shared_expert_gate(hidden_states)`` 计算出的逐 token 门控 logit。在 TP + fused shared-expert 路径中，``sigmoid(logit)`` 会作为追加 shared expert + 的权重，使其输出按 token 动态参与 routed expert 输出的聚合；传入 ``None`` + 时，普通 fused shared expert 的追加权重为 1。EP 路径不使用该参数，而是 + 单独计算 shared expert，并在应用相同门控后与 routed MoE 输出相加。 + + ``is_prefill`` 必须由上层入口显式传入 ``True`` 或 ``False``。即使当前实现 + 尚未使用该信息，也不允许用 ``None`` 隐式表示执行阶段。 + """ + pass + + @abstractmethod + def _fused_experts( + self, + input_tensor: torch.Tensor, + w13: WeightPack, + w2: WeightPack, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + is_prefill: bool, + router_logits: Optional[torch.Tensor] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, + shared_expert_func=None, + ) -> torch.Tensor: + """根据准备完成的路由结果执行融合 MoE 计算。 + + 这里的 ``topk_ids`` 已处于实际执行所需的 ID 空间：普通路径为 logical + ID，shared expert 路径包含追加的专家 ID，冗余专家路径则为修复后的执行 + ID。实现只能读取路由 ID，不能原地修改其内容。``is_prefill`` 只用于选择 + 底层执行策略，不应再影响专家选择或冗余专家映射语义。 + """ pass 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 9c5f58bb7f..26ac1a9df5 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 @@ -1,26 +1,52 @@ import torch from typing import Optional, Tuple, Any -from .triton_impl import FuseMoeTriton +from .base_impl import FuseMoeBaseImpl from lightllm.distributed import dist_group_manager from lightllm.common.quantization.quantize_method import WeightPack 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, + get_env_start_args, + get_redundancy_expert_ids, + get_redundancy_expert_num, ) +from lightllm.utils.dist_utils import get_global_rank, get_global_world_size 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 from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair -class FuseMoeDeepGEMM(FuseMoeTriton): +class FuseMoeDeepGEMM(FuseMoeBaseImpl): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + world_size = get_global_world_size() + assert self.n_routed_experts % world_size == 0 + global_rank = get_global_rank() + self.global_rank_ = global_rank + self.global_world_size_ = world_size + self.ep_n_routed_experts = self.n_routed_experts // world_size + self.redundancy_expert_num = get_redundancy_expert_num() + self.auto_update_redundancy_expert = get_env_start_args().auto_update_redundancy_expert + redundancy_ids = list(get_redundancy_expert_ids(self.layer_index)) + self.redundancy_expert_ids_tensor = torch.tensor(redundancy_ids, dtype=torch.int64, device="cuda") + self.routed_expert_counter_tensor = torch.zeros((self.n_routed_experts,), dtype=torch.int64, device="cuda") + self.total_expert_num_contain_redundancy = self.n_routed_experts + self.redundancy_expert_num * world_size + first_local_expert_id = global_rank * self.ep_n_routed_experts + self.local_logics_expert_ids_list = ( + list(range(first_local_expert_id, first_local_expert_id + self.ep_n_routed_experts)) + redundancy_ids + ) + def _select_experts( self, input_tensor: torch.Tensor, @@ -33,10 +59,8 @@ def _select_experts( num_expert_group: int, scoring_func: str, per_expert_scale: Optional[torch.Tensor] = None, - shared_expert_gate: Optional[torch.Tensor] = None, ): - """Select experts and return topk weights and ids.""" - assert shared_expert_gate is None, "fused shared expert as MoE is not supported by DeepGEMM fused MoE" + """只选择逻辑专家，冗余专家 ID 留到 ``_prepare_expert_execution``。""" from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import select_experts topk_weights, topk_ids = select_experts( @@ -54,10 +78,21 @@ def _select_experts( topk_weights.mul_(self.routed_scaling_factor) if per_expert_scale is not None: topk_weights = topk_weights * per_expert_scale[topk_ids.to(torch.long)].to(topk_weights.dtype) - origin_topk_ids = topk_ids + return topk_weights, topk_ids + + def _prepare_expert_execution( + self, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + is_prefill: bool, + shared_expert_gate: Optional[torch.Tensor] = None, + ): + """在逻辑路由采集之后，生成 DeepEP 实际使用的专家 ID。""" + assert is_prefill is not None, "is_prefill must be explicitly specified for fused MoE execution" + assert shared_expert_gate is None, "fused shared expert as MoE is not supported by DeepGEMM fused MoE" if self.redundancy_expert_num > 0: - # 因为 redundancy_topk_ids_repair 会修改 topk_ids，所以需要先复制一份 - origin_topk_ids = topk_ids.clone() + # repair 会原地改写 ID。复制后再改，保留已经交给 capture 的逻辑 ID。 + topk_ids = topk_ids.clone() redundancy_topk_ids_repair( topk_ids=topk_ids, redundancy_expert_ids=self.redundancy_expert_ids_tensor, @@ -66,7 +101,7 @@ def _select_experts( expert_counter=self.routed_expert_counter_tensor, enable_counter=self.auto_update_redundancy_expert, ) - return topk_weights, topk_ids, origin_topk_ids + return topk_weights, topk_ids def _fused_experts( self, @@ -75,8 +110,8 @@ def _fused_experts( w2: WeightPack, topk_weights: torch.Tensor, topk_ids: torch.Tensor, + is_prefill: bool, router_logits: Optional[torch.Tensor] = None, - is_prefill: Optional[bool] = None, alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, @@ -98,11 +133,12 @@ def _fused_experts( shared_expert_func=shared_expert_func, ) - 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, @@ -110,7 +146,7 @@ def low_latency_dispatch( n_group: int, scoring_func: str, ): - topk_weights, topk_idx, _ = self._select_experts( + topk_weights, topk_idx = self._select_experts( input_tensor=hidden_states, router_logits=router_logits, correction_bias=e_score_correction_bias, @@ -121,33 +157,46 @@ def low_latency_dispatch( num_expert_group=n_group, scoring_func=scoring_func, ) - - return self.low_latency_dispatch_with_topk( + topk_weights, topk_idx = self._prepare_expert_execution(topk_weights, topk_idx, is_prefill=False) + 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) @@ -165,7 +214,7 @@ def select_experts_and_quant_input( n_group: int, scoring_func: str, ): - topk_weights, topk_idx, _ = self._select_experts( + topk_weights, topk_idx = self._select_experts( input_tensor=hidden_states, router_logits=router_logits, correction_bias=e_score_correction_bias, @@ -176,10 +225,11 @@ def select_experts_and_quant_input( num_expert_group=n_group, scoring_func=scoring_func, ) + topk_weights, topk_idx = self._prepare_expert_execution(topk_weights, topk_idx, is_prefill=True) 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, @@ -194,8 +244,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, @@ -210,12 +260,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, @@ -224,14 +274,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, @@ -249,7 +300,6 @@ def prefilled_group_gemm( w13: WeightPack, w2: WeightPack, hidden_dtype=torch.bfloat16, - microbatch_index: int = 0, alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, @@ -270,7 +320,6 @@ def prefilled_group_gemm( w2=w2_weight, w2_scale=w2_scale, block_size_k=self.quant_method.block_size, - workspace=dist_group_manager.get_deep_ep_prefill_moe_workspace(microbatch_index), hidden_dtype=hidden_dtype, alpha=alpha, limit=limit, @@ -298,33 +347,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/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py index 5ccdbb4e9a..6587b9f207 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py @@ -11,6 +11,10 @@ class FuseMoeMarlin(FuseMoeTriton): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.workspace = self.create_workspace() + def create_workspace(self): from lightllm.utils.vllm_utils import HAS_VLLM @@ -28,11 +32,12 @@ def _fused_experts( w2: WeightPack, topk_weights: torch.Tensor, topk_ids: torch.Tensor, + is_prefill: bool, router_logits: Optional[torch.Tensor] = None, - is_prefill: Optional[bool] = None, alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + shared_expert_func=None, ): if alpha is not None or limit is not None: raise NotImplementedError("FuseMoeMarlin does not support clamped SwiGLU") diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py index 506b711fcf..61a1befe65 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py @@ -1,13 +1,10 @@ import torch -from typing import Callable, Optional +from typing import Optional from lightllm.common.quantization.no_quant import WeightPack from .base_impl import FuseMoeBaseImpl class FuseMoeTriton(FuseMoeBaseImpl): - def create_workspace(self): - return None - def _select_experts( self, input_tensor: torch.Tensor, @@ -20,9 +17,8 @@ def _select_experts( num_expert_group: int, scoring_func: str, per_expert_scale: Optional[torch.Tensor] = None, - shared_expert_gate: Optional[torch.Tensor] = None, ): - """Select experts and return topk weights and ids.""" + """Select logical experts and return their weights and ids.""" from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import select_experts topk_weights, topk_ids = select_experts( @@ -40,7 +36,21 @@ def _select_experts( topk_weights.mul_(self.routed_scaling_factor) if per_expert_scale is not None: topk_weights = topk_weights * per_expert_scale[topk_ids.to(torch.long)].to(topk_weights.dtype) - origin_topk_ids = topk_ids + return topk_weights, topk_ids + + def _prepare_expert_execution( + self, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + is_prefill: bool, + shared_expert_gate: Optional[torch.Tensor] = None, + ): + """把逻辑路由结果转换成 kernel 使用的执行布局。 + + 普通路径原样返回。融合 shared expert 时追加新的权重和 ID 行，不修改 + 传入的逻辑 ``topk_ids``。 + """ + assert is_prefill is not None, "is_prefill must be explicitly specified for fused MoE execution" if self.num_fused_shared_experts > 0: from lightllm.common.basemodel.triton_kernel.fused_moe.append_shared_expert_topk import ( append_fused_shared_experts, @@ -53,7 +63,7 @@ def _select_experts( num_fused_shared_experts=self.num_fused_shared_experts, shared_expert_gate=shared_expert_gate, ) - return topk_weights, topk_ids, origin_topk_ids + return topk_weights, topk_ids def _fused_experts( self, @@ -62,11 +72,12 @@ def _fused_experts( w2: WeightPack, topk_weights: torch.Tensor, topk_ids: torch.Tensor, + is_prefill: bool, router_logits: Optional[torch.Tensor] = None, - is_prefill: bool = False, alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + shared_expert_func=None, ): w13_weight, w13_scale = w13.weight, w13.weight_scale w2_weight, w2_scale = w2.weight, w2.weight_scale @@ -89,80 +100,3 @@ def _fused_experts( clamp_up_add_one=clamp_up_add_one, ) return input_tensor - - def fused_experts_with_topk( - self, - input_tensor: torch.Tensor, - w13: WeightPack, - w2: WeightPack, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - is_prefill: Optional[bool] = None, - alpha: Optional[float] = None, - limit: Optional[float] = None, - clamp_up_add_one: bool = True, - ): - return self._fused_experts( - input_tensor=input_tensor, - w13=w13, - w2=w2, - topk_weights=topk_weights, - topk_ids=topk_ids, - is_prefill=is_prefill, - alpha=alpha, - limit=limit, - clamp_up_add_one=clamp_up_add_one, - ) - - def __call__( - self, - input_tensor: torch.Tensor, - router_logits: torch.Tensor, - w13: WeightPack, - w2: WeightPack, - correction_bias: Optional[torch.Tensor], - scoring_func: str, - top_k: int, - renormalize: bool, - use_grouped_topk: bool, - topk_group: int, - num_expert_group: int, - is_prefill: Optional[bool] = None, - # Callback to capture MoE topk expert ids (routed experts metadata). - moe_capture_callback: Optional[Callable[[torch.Tensor], None]] = None, - per_expert_scale: Optional[torch.Tensor] = None, - shared_expert_gate: Optional[torch.Tensor] = None, - alpha: Optional[float] = None, - limit: Optional[float] = None, - clamp_up_add_one: bool = True, - ): - topk_weights, topk_ids, origin_topk_ids = self._select_experts( - input_tensor=input_tensor, - router_logits=router_logits, - correction_bias=correction_bias, - top_k=top_k, - renormalize=renormalize, - use_grouped_topk=use_grouped_topk, - topk_group=topk_group, - num_expert_group=num_expert_group, - scoring_func=scoring_func, - per_expert_scale=per_expert_scale, - shared_expert_gate=shared_expert_gate, - ) - - if moe_capture_callback is not None: - moe_capture_callback(origin_topk_ids) - - output = self._fused_experts( - input_tensor=input_tensor, - w13=w13, - w2=w2, - topk_weights=topk_weights, - topk_ids=topk_ids, - router_logits=router_logits, - is_prefill=is_prefill, - alpha=alpha, - limit=limit, - clamp_up_add_one=clamp_up_add_one, - ) - return output 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..631402c397 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,145 @@ 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 + # TODO: 联合搜索 NUM_STAGE、BLOCK_D 和 token program 数量。 + 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.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py deleted file mode 100644 index a4d7e9f5c6..0000000000 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ /dev/null @@ -1,693 +0,0 @@ -"""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 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, -) -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_build_m_indices, - ep_compact_metadata, - 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.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") - - -try: - from deep_ep import Buffer, EventOverlap - import deep_gemm - - HAS_DEEPGEMM = True -except: - 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 use_sm100_mega_moe(quant_method: Any) -> bool: - return is_sm100_gpu() and quant_method.method_name == "fp4fp8-b32-deepgemm" - - -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 masked_group_gemm( - recv_x: Tuple[torch.Tensor, torch.Tensor], - masked_m: torch.Tensor, - 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, -): - padded_m = recv_x[0].shape[1] - E, N, _ = w1.shape - block_size = 128 - # groupgemm (masked layout) - gemm_out_a = torch.empty((E, padded_m, N), device=recv_x[0].device, dtype=dtype) - expected_m = min(expected_m, padded_m) - qsilu_out_scale = torch.empty((E, padded_m, N // 2 // block_size), device=recv_x[0].device, dtype=torch.float32) - qsilu_out = torch.empty((E, padded_m, N // 2), dtype=w1.dtype, device=recv_x[0].device) - _deepgemm_grouped_fp8_nt_masked(recv_x, (w1, w1_scale), gemm_out_a, masked_m, expected_m) - - silu_and_mul_masked_post_quant_fwd( - gemm_out_a, - qsilu_out, - qsilu_out_scale, - block_size, - masked_m, - alpha=alpha, - limit=limit, - clamp_up_add_one=clamp_up_add_one, - ) - del gemm_out_a - gemm_out_b = torch.empty_like(recv_x[0], device=recv_x[0].device, dtype=dtype) - _deepgemm_grouped_fp8_nt_masked((qsilu_out, qsilu_out_scale), (w2, w2_scale), gemm_out_b, masked_m, expected_m) - 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, - shared_expert_func: Optional[Callable[[], torch.Tensor]] = None, -): - 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") - # Mega MoE has no split combine, so the shared expert cannot overlap it. - shared = shared_expert_func() if shared_expert_func is not None else None - routed = mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method) - return routed if shared is None else routed + shared - - 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, - shared_expert_func=shared_expert_func, - ) - - -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, - shared_expert_func: Optional[Callable[[], torch.Tensor]] = None, -): - # 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" - - # TODO: split the shared expert into up and down. Overlap up with dispatch and down - # with combine. Folding the whole expert into combine leaves the dispatch transfer idle. - 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 - - overlap = shared_expert_func is not None - if overlap: - # Combine waits for the routed GEMM, then the shared expert overlaps the transfer. - from deep_ep import ElasticBuffer - - previous_event = ElasticBuffer.capture() - allocate_on_comm_stream = True - combined_x, _, event = buffer.combine( - gather_out, - handle, - topk_weights=None, - previous_event=previous_event, - async_with_compute_stream=overlap, - allocate_on_comm_stream=allocate_on_comm_stream, - ) - if overlap: - shared = shared_expert_func() - event.current_stream_wait() - combined_x.add_(shared) - 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, - ) - overlap = shared_expert_func is not None - combined_x, event_overlap, hook = buffer.low_latency_combine( - gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=overlap - ) - # The send is queued on this stream. The shared expert reuses its buffer, then the recv runs. - if overlap: - del gemm_out_b - shared = shared_expert_func() - hook() - combined_x = combined_x + shared - 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, - hidden_size: int, - intermediate_size: int, - intermediate_twice: int, - scale_cols: int, - hidden_dtype: torch.dtype, - quant_dtype: torch.dtype, - expert_alignment: int, -) -> int: - """计算并缓存当前 workspace 配置能够容纳的最大 chunk 行数。""" - if not hasattr(_get_max_chunk_rows, "cache"): - _get_max_chunk_rows.cache = {} - max_chunk_rows_cache = _get_max_chunk_rows.cache - - # 同一 1024 行区间共用一个缓存项，并按区间上界探测，保证复用结果不会高估可用空间。 - cached_gather_rows = (gather_rows + 1023) // 1024 * 1024 - cache_key = ( - workspace.numel(), - workspace.device, - cached_gather_rows, - hidden_size, - intermediate_size, - intermediate_twice, - scale_cols, - hidden_dtype, - quant_dtype, - expert_alignment, - HAS_SGL_KERNEL, - ) - if cache_key in max_chunk_rows_cache: - return max_chunk_rows_cache[cache_key] - - def can_allocate(chunk_rows: int) -> bool: - """按实际计算阶段的生命周期申请 buffer，探测该 chunk 是否能够执行。""" - try: - probe_manager = TensorBufferManager(workspace) - probe_manager.alloc((cached_gather_rows, hidden_size), hidden_dtype) - - # W1 阶段同时保存 GEMM 输出和 SwiGLU 输出。 - silu_out = probe_manager.alloc((chunk_rows, intermediate_size), hidden_dtype) - gemm_out_a = probe_manager.alloc((chunk_rows, intermediate_twice), hidden_dtype) - probe_manager.free(gemm_out_a) - - # 量化阶段复用 W1 输出空间，并继续保留 SwiGLU 输出。 - probe_manager.alloc((chunk_rows, intermediate_size), quant_dtype) - aligned_chunk_rows = (chunk_rows + 3) // 4 * 4 - if HAS_SGL_KERNEL: - probe_manager.alloc((scale_cols, aligned_chunk_rows), torch.float32) - else: - # LightLLM fallback 通过 alloc_func 申请 row-major scale；后续 TMA 转置不占用 workspace。 - probe_manager.alloc((chunk_rows, scale_cols), torch.float32) - probe_manager.free(silu_out) - - # W2 阶段释放 SwiGLU 输出后申请最终 GEMM 输出。 - probe_manager.alloc((chunk_rows, hidden_size), hidden_dtype) - except MemoryError: - return False - return True - - # W1 阶段必须同时保存 gemm_out_a 和 silu_out，可据此得到 chunk 数量的绝对上界。 - w1_row_bytes = (intermediate_twice + intermediate_size) * hidden_dtype.itemsize - left = 1 - right = workspace.numel() // w1_row_bytes // expert_alignment - max_chunk_rows = 0 - - while left <= right: - chunk_count = (left + right) // 2 - chunk_rows = chunk_count * expert_alignment - if can_allocate(chunk_rows): - max_chunk_rows = chunk_rows - left = chunk_count + 1 - else: - right = chunk_count - 1 - - max_chunk_rows_cache[cache_key] = max_chunk_rows - logger.info("cache DeepEP max_chunk_rows: key=%s, max_chunk_rows=%s", cache_key, max_chunk_rows) - return max_chunk_rows - - -def chunked_expanded_moe_forward( - num_recv_tokens_per_expert_list: List[int], # [num_local_experts], 128-aligned 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] - ], # ([num_expanded_tokens, hidden_size], [num_expanded_tokens, hidden_size // block_size_k]) - recv_topk_weights: torch.Tensor, # [num_expanded_tokens] - recv_src_metadata: torch.Tensor, # [num_recv_tokens, topk + 2] - w1: torch.Tensor, # [num_local_experts, 2 * intermediate_size, hidden_size] - w1_scale: torch.Tensor, # [num_local_experts, 2 * intermediate_size // block_size_k, hidden_size // block_size_k] - w2: torch.Tensor, # [num_local_experts, hidden_size, intermediate_size] - w2_scale: torch.Tensor, # [num_local_experts, hidden_size // block_size_k, intermediate_size // block_size_k] - block_size_k: int, - workspace: torch.Tensor, # [workspace_bytes], uint8 - hidden_dtype: torch.dtype, # scalar dtype descriptor - alpha: Optional[float] = None, - limit: Optional[float] = None, - clamp_up_add_one: bool = True, -): - """Run bounded expanded MoE and rewrite metadata for dense DeepEP combine.""" - alignment = 128 - 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 - assert all_tokens > 0, "chunked_expanded_moe_forward requires non-empty input" - assert workspace.dtype == torch.uint8 and workspace.ndim == 1 and workspace.is_contiguous() - - m_indices = torch.empty(all_tokens, device=recv_x[0].device, dtype=torch.int32) - # 与 m_indices 一一对应：0 表示真实 token，1 表示 expert 对齐产生的 padding 行。 - # padding 行必须在 grouped GEMM 前清零，避免无效数据参与计算。 - padding_mask = torch.empty_like(m_indices) - ep_build_m_indices(num_unaligned_recv_tokens_per_expert, m_indices, padding_mask, alignment) - ep_zero_padding( - recv_x[0], - recv_x[1], - recv_topk_weights, - padding_mask, - ) - del padding_mask - - gather_rows = recv_src_metadata.shape[0] - scale_cols = intermediate_size // block_size_k - max_chunk_rows = _get_max_chunk_rows( - workspace=workspace, - gather_rows=gather_rows, - hidden_size=hidden_size, - intermediate_size=intermediate_size, - intermediate_twice=intermediate_twice, - scale_cols=scale_cols, - hidden_dtype=hidden_dtype, - quant_dtype=w2.dtype, - expert_alignment=alignment, - ) - - if max_chunk_rows == 0: - raise RuntimeError( - f"DeepEP workspace with {workspace.numel()} bytes cannot hold the dense output and " - f"one {alignment}-row temporary chunk" - ) - max_chunk_rows = min(all_tokens, max_chunk_rows) - - workspace_manager = TensorBufferManager(workspace) - gather_out = workspace_manager.alloc((gather_rows, hidden_size), hidden_dtype) - gather_out.zero_() - - # 不同 rank 接收到的 token 数不同，因此实际 chunk 数也可能不同。Autotuner warmup - # 中的分布式通信要求各 rank 进入 autotuning 的次数一致，否则容易发生通信错位。 - # 所以只允许第一个 chunk 保持 autotuning；从第二个 chunk 开始临时关闭，循环结束 - # 后再恢复进入函数时的 warmup 状态。零 token rank 的首次调用由外层特殊分支补齐。 - is_autotune_warmup = Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) - try: - for chunk_index, chunk_start in enumerate(range(0, all_tokens, max_chunk_rows)): - if is_autotune_warmup and chunk_index == 1: - Autotuner.end_autotune_warmup() - - chunk_end = min(chunk_start + max_chunk_rows, all_tokens) - chunk_rows = chunk_end - chunk_start - silu_out = workspace_manager.alloc((chunk_rows, intermediate_size), hidden_dtype) - gemm_out_a = workspace_manager.alloc((chunk_rows, intermediate_twice), hidden_dtype) - deepgemm_grouped_fp8_nt_contiguous( - (recv_x[0][chunk_start:chunk_end], recv_x[1][chunk_start:chunk_end]), - (w1, w1_scale), - gemm_out_a, - m_indices[chunk_start:chunk_end], - ) - silu_and_mul_fwd(gemm_out_a, silu_out, alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one) - workspace_manager.free(gemm_out_a) - del gemm_out_a - - quant_buffers = [] - - def workspace_quant_alloc(shape, dtype, device): - if device != workspace.device: - raise RuntimeError(f"quant buffer must be allocated on {workspace.device}, got {device}") - quant_buffer = workspace_manager.alloc(shape, dtype) - quant_buffers.append(quant_buffer) - return quant_buffer - - qsilu_out, qsilu_out_scale = per_token_group_quant_fp8( - silu_out, - block_size_k, - dtype=w2.dtype, - column_major_scales=True, - scale_tma_aligned=True, - alloc_func=workspace_quant_alloc, - use_ue8m0_scales=get_env_use_ue8m0_scales(), - ) - workspace_manager.free(silu_out) - del silu_out - - gemm_out_b = workspace_manager.alloc((chunk_rows, hidden_size), hidden_dtype) - deepgemm_grouped_fp8_nt_contiguous( - (qsilu_out, qsilu_out_scale), - (w2, w2_scale), - gemm_out_b, - m_indices[chunk_start:chunk_end], - ) - del qsilu_out, qsilu_out_scale - for quant_buffer in quant_buffers: - workspace_manager.free(quant_buffer) - - ep_gather_chunk(gemm_out_b, chunk_start, recv_topk_weights, recv_src_metadata, gather_out) - workspace_manager.free(gemm_out_b) - finally: - if is_autotune_warmup: - Autotuner.start_autotune_warmup(AutotuneKernelType.GENERAL) - - ep_compact_metadata(recv_src_metadata) - return gather_out - - -def _deepgemm_grouped_fp8_nt_masked( - input_tuple: Tuple[torch.Tensor, torch.Tensor], - w_tuple: Tuple[torch.Tensor, torch.Tensor], - out: torch.Tensor, - masked_m: torch.Tensor, - expected_m: int, -): - if HAS_DEEPGEMM: - if hasattr(deep_gemm, "m_grouped_fp8_gemm_nt_masked"): - return deep_gemm.m_grouped_fp8_gemm_nt_masked(input_tuple, w_tuple, out, masked_m, expected_m) - if hasattr(deep_gemm, "m_grouped_gemm_fp8_fp8_bf16_nt_masked"): - return deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_masked(input_tuple, w_tuple, out, masked_m, expected_m) - raise RuntimeError("deep_gemm does not provide grouped_gemm_fp8 NT contiguous GEMM kernel in this version") - - -def deepgemm_grouped_fp8_fp4_nt_contiguous( - input_tuple: Tuple[torch.Tensor, torch.Tensor], - w_tuple: Tuple[torch.Tensor, torch.Tensor], - out: torch.Tensor, - grouped_layout: torch.Tensor, - use_psum_layout: bool = False, -): - if HAS_DEEPGEMM and hasattr(deep_gemm, "m_grouped_fp8_fp4_gemm_nt_contiguous"): - return deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( - input_tuple, - w_tuple, - out, - grouped_layout, - use_psum_layout=use_psum_layout, - recipe=(1, 1, 32), - ) - raise RuntimeError("deep_gemm does not provide grouped fp8-fp4 NT contiguous GEMM kernel") 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..6598ea1f8f --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/__init__.py @@ -0,0 +1,40 @@ +"""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, + 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, + 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", + "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", + "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/deepgemm.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/deepgemm.py new file mode 100644 index 0000000000..c4b8998f20 --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/deepgemm.py @@ -0,0 +1,290 @@ +import torch +import triton +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_psum_post_quant_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_build_m_indices, + ep_compact_metadata, + ep_gather_chunk, + ep_zero_padding, +) +from lightllm.utils.envs_utils import get_env_use_ue8m0_scales +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType + +logger = init_logger(__name__) + +deep_gemm = None +try: + import deep_gemm + + HAS_DEEPGEMM = True +except Exception: + logger.warning("no deepep or deep_gemm") + HAS_DEEPGEMM = False + + +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 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 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 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 chunked_expanded_moe_forward( + 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] + ], # ([num_expanded_tokens, hidden_size], [num_expanded_tokens, hidden_size // block_size_k]) + recv_topk_weights: torch.Tensor, # [num_expanded_tokens] + recv_src_metadata: torch.Tensor, # [num_recv_tokens, topk + 2] + w1: torch.Tensor, # [num_local_experts, 2 * intermediate_size, hidden_size] + w1_scale: torch.Tensor, # [num_local_experts, 2 * intermediate_size // block_size_k, hidden_size // block_size_k] + w2: torch.Tensor, # [num_local_experts, hidden_size, intermediate_size] + w2_scale: torch.Tensor, # [num_local_experts, hidden_size // block_size_k, intermediate_size // block_size_k] + block_size_k: int, + hidden_dtype: torch.dtype, # scalar dtype descriptor + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, +): + """以最多 32K 行为一个 chunk 执行 expanded MoE，并重写 combine metadata。""" + 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 + assert all_tokens > 0, "chunked_expanded_moe_forward requires non-empty input" + assert 32768 % alignment == 0 + + m_indices = torch.empty(all_tokens, device=recv_x[0].device, dtype=torch.int32) + # 与 m_indices 一一对应：0 表示真实 token，1 表示 expert 对齐产生的 padding 行。 + # padding 行必须在 grouped GEMM 前清零，避免无效数据参与计算。 + padding_mask = torch.empty_like(m_indices) + ep_build_m_indices(num_unaligned_recv_tokens_per_expert, m_indices, padding_mask, alignment) + ep_zero_padding( + recv_x[0], + recv_x[1], + recv_topk_weights, + padding_mask, + ) + del padding_mask + + max_chunk_rows = min(all_tokens, 32768) + gather_out = torch.zeros( + (recv_src_metadata.shape[0], hidden_size), + dtype=hidden_dtype, + device=recv_x[0].device, + ) + + # 不同 rank 接收到的 token 数不同，因此实际 chunk 数也可能不同。Autotuner warmup + # 中的分布式通信要求各 rank 进入 autotuning 的次数一致，否则容易发生通信错位。 + # 所以只允许第一个 chunk 保持 autotuning；从第二个 chunk 开始临时关闭，循环结束 + # 后再恢复进入函数时的 warmup 状态。零 token rank 的首次调用由外层特殊分支补齐。 + is_autotune_warmup = Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL) + try: + for chunk_index, chunk_start in enumerate(range(0, all_tokens, max_chunk_rows)): + if is_autotune_warmup and chunk_index == 1: + Autotuner.end_autotune_warmup() + + chunk_end = min(chunk_start + max_chunk_rows, all_tokens) + chunk_rows = chunk_end - chunk_start + silu_out = torch.empty( + (chunk_rows, intermediate_size), + dtype=hidden_dtype, + device=recv_x[0].device, + ) + gemm_out_a = torch.empty( + (chunk_rows, intermediate_twice), + dtype=hidden_dtype, + device=recv_x[0].device, + ) + deepgemm_grouped_fp8_nt_contiguous( + (recv_x[0][chunk_start:chunk_end], recv_x[1][chunk_start:chunk_end]), + (w1, w1_scale), + gemm_out_a, + m_indices[chunk_start:chunk_end], + ) + silu_and_mul_fwd(gemm_out_a, silu_out, alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one) + del gemm_out_a + + qsilu_out, qsilu_out_scale = per_token_group_quant_fp8( + silu_out, + block_size_k, + dtype=w2.dtype, + column_major_scales=True, + scale_tma_aligned=True, + use_ue8m0_scales=get_env_use_ue8m0_scales(), + ) + del silu_out + + gemm_out_b = torch.empty( + (chunk_rows, hidden_size), + dtype=hidden_dtype, + device=recv_x[0].device, + ) + deepgemm_grouped_fp8_nt_contiguous( + (qsilu_out, qsilu_out_scale), + (w2, w2_scale), + gemm_out_b, + m_indices[chunk_start:chunk_end], + ) + del qsilu_out, qsilu_out_scale + + ep_gather_chunk(gemm_out_b, chunk_start, recv_topk_weights, recv_src_metadata, gather_out) + del gemm_out_b + finally: + if is_autotune_warmup: + Autotuner.start_autotune_warmup(AutotuneKernelType.GENERAL) + + ep_compact_metadata(recv_src_metadata) + return gather_out + + +def deepgemm_grouped_fp8_fp4_nt_contiguous( + input_tuple: Tuple[torch.Tensor, torch.Tensor], + w_tuple: Tuple[torch.Tensor, torch.Tensor], + out: torch.Tensor, + grouped_layout: torch.Tensor, + use_psum_layout: bool = False, +): + if HAS_DEEPGEMM and hasattr(deep_gemm, "m_grouped_fp8_fp4_gemm_nt_contiguous"): + return deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( + input_tuple, + w_tuple, + out, + grouped_layout, + use_psum_layout=use_psum_layout, + recipe=(1, 1, 32), + ) + raise RuntimeError("deep_gemm does not provide grouped fp8-fp4 NT contiguous GEMM kernel") 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..3a4909261d --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/prefill_decode_sync.py @@ -0,0 +1,299 @@ +import torch +import triton +from typing import Any, Callable, 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, + shared_expert_func: Optional[Callable[[], torch.Tensor]] = None, +): + 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") + # Mega MoE has no split combine, so the shared expert cannot overlap it. + shared = shared_expert_func() if shared_expert_func is not None else None + routed = mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method) + return routed if shared is None else routed + shared + + 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, + shared_expert_func=shared_expert_func, + ) + + +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, + shared_expert_func: Optional[Callable[[], torch.Tensor]] = None, +): + # 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" + # TODO: 共享专家目前整段只和 combine 重叠。可以拆开：dispatch 期间算 up， + # combine 期间算 down，把两段通信的空闲都盖住。 + + 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, + 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 + + combined_x = _combine_with_optional_shared_expert( + buffer, + gather_out, + handle, + previous_event, + allocate_on_comm_stream, + shared_expert_func, + ) + 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 + + # compact metadata 的第一个 top-k 槽位指向同序 dense row，其余槽位置为 -1。 + return _combine_with_optional_shared_expert( + buffer, + dense_output, + ep_handle, + previous_event=None, + allocate_on_comm_stream=False, + shared_expert_func=shared_expert_func, + ) + + +def _combine_with_optional_shared_expert( + buffer, + expert_output: torch.Tensor, + handle, + previous_event, + allocate_on_comm_stream: bool, + shared_expert_func: Optional[Callable[[], torch.Tensor]], +): + """Launch combine, optionally running the shared expert while the transfer is in flight.""" + overlap = shared_expert_func is not None + if overlap: + # Combine waits for the routed GEMM, then the shared expert overlaps the transfer. + from deep_ep import ElasticBuffer + + previous_event = ElasticBuffer.capture() + allocate_on_comm_stream = True + combined_x, _, event = buffer.combine( + expert_output, + handle, + topk_weights=None, + num_sms=get_ep_num_sms(overlap_with_compute=overlap), + previous_event=previous_event, + async_with_compute_stream=overlap, + allocate_on_comm_stream=allocate_on_comm_stream, + ) + if allocate_on_comm_stream and combined_x.is_cuda: + # 输出在通信流上分配，后续在计算流上使用。再登记一次计算流， + # 引用归零时分配器才会等这条流上已经排队的 kernel 结束。 + combined_x.record_stream(torch.cuda.current_stream()) + if overlap: + shared = shared_expert_func() + event.current_stream_wait() + combined_x.add_(shared) + 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..7e51d8d736 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -20,7 +20,6 @@ import os import torch -import triton import torch.distributed as dist from torch.distributed import ReduceOp, ProcessGroup from typing import List, Dict, Optional, Set, Union @@ -28,8 +27,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_prefill, - get_deepep_num_max_dispatch_tokens_per_rank_decode, + get_deepep_num_max_dispatch_tokens_per_rank, get_redundancy_expert_num, ) from lightllm.utils.dist_utils import ( @@ -109,9 +107,9 @@ class DistributeGroupManager: def __init__(self): self.groups = [] self.ep_buffer = None - 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,20 +165,28 @@ 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 - prefill_num_max_dispatch_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank_prefill() - decode_num_max_dispatch_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank_decode() if not enable_ep_moe: self.ep_buffer = None - 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" + # ElasticBuffer 按当前服务真正会执行的阶段选择容量。 + num_max_dispatch_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank() + + 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,8 +196,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_decode_num_tokens = decode_num_max_dispatch_tokens_per_rank + self.ll_num_tokens = 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 self.ep_buffer = deep_ep.ElasticBuffer( @@ -203,63 +208,15 @@ def new_deepep_group( allow_multiple_reduction=True, ) self.ep_mega_moe_buffer = None - self.ep_low_latency_buffer = None if not expert_quant_method_names: raise ValueError("No valid MoE quant method was found while initializing DeepEP buffers") mega_moe_quant_method = "fp4fp8-b32-deepgemm" - is_sm100 = is_sm100_gpu() - - # Buffer 选择规则： - # 1. 非 SM100 不支持 Mega MoE，只初始化 legacy low-latency buffer； - # 2. SM100 全部 MoE 层为 FP4，只初始化 Mega MoE buffer； - # 3. SM100 全部 MoE 层为 FP8，只初始化 legacy low-latency buffer； - # 4. SM100 逐层混合 FP4/FP8，两套 buffer 都要初始化。 - if is_sm100: - # 只要存在一个 FP4 MoE 层，就需要 Mega MoE buffer；只要存在一个非 FP4 - # MoE 层，就需要 legacy low-latency buffer。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 - ) - enable_mega_moe_buffer = has_mega_moe_layer - enable_low_latency_buffer = has_legacy_moe_layer - else: - enable_mega_moe_buffer = False - enable_low_latency_buffer = True - - if enable_low_latency_buffer: - # FP8 MoE 的 decode 使用 legacy low-latency buffer；prefill 阶段还会将其 - # 空闲的本地 RDMA storage 复用为分块 grouped GEMM 的临时 workspace。 - 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 - ) - microbatch_count = len(self.groups) - min_prefill_reuse_buffer_bytes = _calculate_min_chunked_expanded_moe_reuse_buffer_bytes( - global_world_size=global_world_size, - prefill_tokens_per_rank=prefill_num_max_dispatch_tokens_per_rank, - hidden_size=hidden_size, - moe_intermediate_size=moe_intermediate_size, - hidden_dtype=get_torch_dtype(get_env_start_args().data_type), - microbatch_count=microbatch_count, - ) - # Decode 和 prefill 不会同时使用 local RDMA storage，容量取两条路径的较大值。 - num_rdma_bytes = max(decode_size_hint, min_prefill_reuse_buffer_bytes) - # DeepEP 返回的 decode hint 不保证能被 microbatch 均分；最终再对齐一次， - # 确保每个 workspace slice 的容量及起始位置仍保持 256-byte 对齐。 - rdma_alignment = microbatch_count * 256 - num_rdma_bytes = triton.cdiv(num_rdma_bytes, rdma_alignment) * rdma_alignment - self.ep_low_latency_buffer = deep_ep.Buffer( - deepep_group, - num_rdma_bytes=num_rdma_bytes, - low_latency_mode=True, - num_qps_per_rank=(self.ll_num_experts // global_world_size), - ) + # SM100 FP4 层额外通过 DeepGEMM Mega MoE 完成通信和计算。 + enable_mega_moe_buffer = is_sm100_gpu() and mega_moe_quant_method in expert_quant_method_names if enable_mega_moe_buffer: - # SM100 FP4 层通过 DeepGEMM Mega MoE 完成通信和计算，不使用 legacy - # low-latency buffer，因此纯 FP4 模型无需承担后者的大块 RDMA 显存。 if moe_intermediate_size is None: raise ValueError("SM100 Mega MoE requires moe_intermediate_size or intermediate_size in model config") @@ -274,14 +231,35 @@ def new_deepep_group( moe_intermediate_size, ) logger.info( - "Initialize DeepEP MoE buffers: low_latency=%s, mega_moe=%s, expert_quant_method_names=%s", - enable_low_latency_buffer, + "Initialize DeepEP MoE buffers: elastic=True, mega_moe=%s, " + "num_max_dispatch_tokens_per_rank=%s, ll_num_experts=%s, expert_quant_method_names=%s", enable_mega_moe_buffer, + self.ll_num_tokens, + self.ll_num_experts, sorted(expert_quant_method_names), ) 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: @@ -292,51 +270,15 @@ def _set_num_sms_for_deep_gemm(self, deepep_sms: int): device_sms = get_device_sm_count() deepep_sms = max(0, min(deepep_sms, max(device_sms - 2, 0))) self.ep_num_sms = deepep_sms - if self.ep_low_latency_buffer is not None: - deep_ep.Buffer.set_num_sms(deepep_sms - deepep_sms % 2) set_num_sms(max(device_sms - deepep_sms, 2)) except BaseException as e: logger.warning(f"set num sms for deep_gemm failed: {e}") - def get_deep_ep_prefill_moe_workspace(self, microbatch_index: int = 0) -> torch.Tensor: - """Return a slice of the workspace reused by DeepEP prefill MoE kernels. - - DeepEP's low-latency RDMA buffer is idle during prefill, so its local - storage is reused as temporary workspace for the expanded MoE compute - path to reduce peak GPU memory. With one communication group, the - default ``microbatch_index=0`` receives the whole workspace. With - multiple groups, the workspace is split into ``len(self.groups)`` - equal slices and each in-flight microbatch uses the slice matching its - group index. - - Args: - microbatch_index: Zero-based microbatch and communication-group - index assigned to this in-flight prefill computation. - - Returns: - A one-dimensional uint8 tensor view over the selected workspace - slice; no additional GPU memory is allocated. - - This workspace is only valid after the DeepEP group has been - initialized. The same returned slice must not be used concurrently by - overlapping calls. - """ - assert self.ep_low_latency_buffer is not None, "DeepEP low-latency buffer is not initialized" - workspace = self.ep_low_latency_buffer.get_local_buffer_tensor(torch.uint8, use_rdma_buffer=True) - microbatch_count = len(self.groups) - assert 0 <= microbatch_index < microbatch_count - workspace_size = workspace.numel() // microbatch_count - return workspace.narrow(0, microbatch_index * workspace_size, workspace_size) - - def clear_deepep_buffer(self): - """ - Prefill MoE compute reuses the low-latency RDMA buffer as workspace. - Clean it before the buffer is used by low-latency decode kernels. - """ - if self.ep_low_latency_buffer is not None: - self.ep_low_latency_buffer.clean_low_latency_buffer( - self.ll_decode_num_tokens, self.ll_hidden, self.ll_num_experts - ) + def clear_deepep_buffer(self) -> None: + """清理通信组件在当前请求中产生的临时状态。""" + # 当前通信组件不需要显式清理。保留统一入口，后续接入需要清理的 + # 通信实现时，无需再修改模型执行流程。 + pass def all_reduce( @@ -424,63 +366,4 @@ def _is_single_group(group: Optional[Union[ProcessGroup, CustomProcessGroup]]) - return dist.get_world_size(group=group) == 1 -def _calculate_min_chunked_expanded_moe_reuse_buffer_bytes( - global_world_size: int, - prefill_tokens_per_rank: int, - hidden_size: int, - moe_intermediate_size: int, - hidden_dtype: torch.dtype, - microbatch_count: int, -) -> int: - """计算 ``chunked_expanded_moe_forward`` 极端情况下所需的最小复用 buffer。 - - Prefill 阶段会将空闲的 legacy DeepEP local RDMA storage 复用为 expanded - grouped GEMM 的 workspace。该函数只计算这条 prefill 路径的容量下限，不包含 - low-latency decode 自身所需的 RDMA 容量。 - - 每个 prefill workspace 由全程常驻的 dense gather 输出和分块计算的临时峰值组成。 - gather 按 ``global_world_size * prefill_tokens_per_rank`` 估算最大可能接收行数， - 并向 1024 行对齐；这与运行期 workspace 探测使用的保守分桶保持一致。临时计算 - 按最小 128 行 chunk 估算：W1 阶段同时需要 ``gemm_out_a`` 和 ``silu_out``；后续 - 量化/W2 阶段涉及 ``silu_out``、FP8 量化结果、scale 及 ``gemm_out_b``。这里有意 - 沿用保守公式，避免因生命周期估计不足而低估空间。 - - 多个 microbatch/communication group 并行时，每组独占一个 workspace slice，故 - prefill 容量乘以 ``microbatch_count``。最后按 ``microbatch_count * 256`` 对齐， - 保证均分后每个 slice 仍满足 256-byte 对齐。 - """ - - def align(value: int, alignment: int) -> int: - return triton.cdiv(value, alignment) * alignment - - def tensor_bytes(rows: int, columns: int, itemsize: int) -> int: - # 沿用TensorBufferManager的256对齐规则 - return align(rows * columns * itemsize, 256) - - chunk_rows = 128 - hidden_itemsize = hidden_dtype.itemsize - scale_cols = moe_intermediate_size // 128 - - silu_bytes = tensor_bytes(chunk_rows, moe_intermediate_size, hidden_itemsize) - gemm_out_a_bytes = tensor_bytes(chunk_rows, 2 * moe_intermediate_size, hidden_itemsize) - quant_bytes = tensor_bytes(chunk_rows, moe_intermediate_size, 1) - scale_bytes = tensor_bytes(chunk_rows, scale_cols, torch.float32.itemsize) - gemm_out_b_bytes = tensor_bytes(chunk_rows, hidden_size, hidden_itemsize) - - w1_peak_bytes = silu_bytes + gemm_out_a_bytes - - # silu 释放后，B 较小时可复用其 first-fit 空洞；B 更大则需另行预留完整空间。 - quant_w2_peak_bytes = ( - silu_bytes + quant_bytes + scale_bytes + (gemm_out_b_bytes if gemm_out_b_bytes > silu_bytes else 0) - ) - temporary_peak_bytes = max(w1_peak_bytes, quant_w2_peak_bytes) - - # 按 _get_max_chunk_rows 的 1024 行对齐规则 - gather_bytes = tensor_bytes(align(global_world_size * prefill_tokens_per_rank, 1024), hidden_size, hidden_itemsize) - per_workspace_bytes = gather_bytes + temporary_peak_bytes - - reuse_buffer_bytes = per_workspace_bytes * microbatch_count - return align(reuse_buffer_bytes, microbatch_count * 256) - - dist_group_manager = DistributeGroupManager() diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index 986b8a703f..a9d23a6de5 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 @@ -247,6 +246,7 @@ def _moe_ffn_tp( use_grouped_topk=self.n_group, topk_group=self.topk_group, num_expert_group=self.n_group, + is_prefill=infer_state.is_prefill, infer_state=infer_state, ) @@ -258,7 +258,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 +339,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 +370,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 +391,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 +462,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 +491,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 +505,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( @@ -522,8 +526,8 @@ def overlap_tpsp_context_forward( _0_recv_x, _0_recv_topk_idx, _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 @@ -554,8 +562,8 @@ def overlap_tpsp_context_forward( _1_recv_x, _1_recv_topk_idx, _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 +577,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 e0212859cd..21825e68e7 100644 --- a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py @@ -181,24 +181,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, @@ -208,14 +210,15 @@ def overlap_tpsp_context_forward( recv_indices0, recv_weights0, hidden_dtype=x0.dtype, - microbatch_index=infer_state.microbatch_index, alpha=1.0, 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( @@ -226,14 +229,15 @@ def overlap_tpsp_context_forward( recv_indices1, recv_weights1, hidden_dtype=x1.dtype, - microbatch_index=infer_state1.microbatch_index, alpha=1.0, 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_shared_expert(x0, infer_state, layer_weight) combine_hook0() @@ -277,9 +281,7 @@ def overlap_tpsp_token_forward( infer_state1.call_overlap_hook() shared0 = self._ffn_shared_expert(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) @@ -289,42 +291,44 @@ def overlap_tpsp_token_forward( dispatch_hook0() shared1 = self._ffn_shared_expert(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/gpt_oss/layer_infer/transformer_layer_infer.py b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py index b26025c753..b6ebc52984 100644 --- a/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/gpt_oss/layer_infer/transformer_layer_infer.py @@ -53,6 +53,7 @@ def _ffn(self, input, infer_state, layer_weight: GptOssTransformerLayerWeight) - use_grouped_topk=False, topk_group=None, num_expert_group=None, + is_prefill=infer_state.is_prefill, infer_state=infer_state, ) hidden_states = hidden_states.view(num_tokens, hidden_dim) diff --git a/lightllm/models/mixtral/layer_infer/transformer_layer_infer.py b/lightllm/models/mixtral/layer_infer/transformer_layer_infer.py index 8134dc266d..4e00659e80 100644 --- a/lightllm/models/mixtral/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/mixtral/layer_infer/transformer_layer_infer.py @@ -26,6 +26,7 @@ def _ffn(self, input, infer_state: InferStateInfo, layer_weight: MixtralTransfor use_grouped_topk=False, topk_group=None, num_expert_group=None, + is_prefill=infer_state.is_prefill, infer_state=infer_state, ) return hidden_states.view(num_tokens, hidden_dim) 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..0c8bd398c5 100644 --- a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py @@ -88,6 +88,7 @@ def _moe_ffn_tp( use_grouped_topk=False, topk_group=None, num_expert_group=None, + is_prefill=infer_state.is_prefill, infer_state=infer_state, ) return hidden_states.view(num_tokens, hidden_dim) @@ -160,14 +161,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 +187,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 +208,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 +273,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 +301,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 +315,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 @@ -323,8 +326,8 @@ def overlap_tpsp_context_forward( _0_recv_x, _0_recv_topk_idx, _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 @@ -355,8 +362,8 @@ def overlap_tpsp_context_forward( _1_recv_x, _1_recv_topk_idx, _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 +375,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/qwen3next/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py index 39f9c42919..1bc4ef914b 100644 --- a/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py @@ -94,6 +94,7 @@ def _moe_ffn_tp( use_grouped_topk=False, topk_group=None, num_expert_group=None, + is_prefill=infer_state.is_prefill, infer_state=infer_state, shared_expert_gate=shared_expert_gate, ) @@ -183,7 +184,6 @@ def _linear_prefill_cuda_graph_wrapper( infer_state: Qwen3NextInferStateInfo, layer_weight: Qwen3NextTransformerLayerWeight, ) -> Tuple[torch.Tensor, torch.Tensor]: - backend = infer_state.prefill_att_state1.backend mixed_qkvzba = mixed_qkvzba.contiguous() 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/test/benchmark/kernel/benchmark_shared_expert_overlap.py b/test/benchmark/kernel/benchmark_shared_expert_overlap.py index 5439a23524..9c92d72129 100644 --- a/test/benchmark/kernel/benchmark_shared_expert_overlap.py +++ b/test/benchmark/kernel/benchmark_shared_expert_overlap.py @@ -83,16 +83,17 @@ def worker(rank, args, port): group = dist.new_group(backend="gloo") dist.barrier(device_ids=[rank]) experts = args.experts_per_rank * args.world_size - size = deep_ep.Buffer.get_low_latency_rdma_size_hint(args.max_rows, args.hidden, args.world_size, experts) - print(rank, "creating DeepEP buffer", size, flush=True) - buffer = deep_ep.Buffer( + print(rank, "creating DeepEP ElasticBuffer", args.max_rows, args.hidden, flush=True) + buffer = deep_ep.ElasticBuffer( dist.group.WORLD, - num_rdma_bytes=size, - low_latency_mode=True, - num_qps_per_rank=args.experts_per_rank, + num_max_tokens_per_rank=args.max_rows, + hidden=args.hidden, + num_topk=args.topk, + use_fp8_dispatch=True, + allow_multiple_reduction=True, explicitly_destroy=True, ) - dist_group_manager.ep_low_latency_buffer = buffer + dist_group_manager.ep_buffer = buffer print(rank, "DeepEP ready", flush=True) try: torch.manual_seed(1013 + rank) @@ -109,7 +110,7 @@ def pack(shape): owner.layer_num_ = 0 owner.w13 = pack((args.experts_per_rank, 2 * args.intermediate, args.hidden)) owner.w2 = pack((args.experts_per_rank, args.hidden, args.intermediate)) - owner.fuse_moe_impl = FuseMoeDeepGEMM(experts, 0, 1.0, quant, 0, None, None, False) + owner.fuse_moe_impl = FuseMoeDeepGEMM(experts, 0, 1.0, quant, layer_index=0) def shared_weight(shape): w = MMWeightTpl.__new__(MMWeightTpl) diff --git a/unit_tests/common/fused_moe/test_activation_config.py b/unit_tests/common/fused_moe/test_activation_config.py index 8972dd88f0..c59ab2712d 100644 --- a/unit_tests/common/fused_moe/test_activation_config.py +++ b/unit_tests/common/fused_moe/test_activation_config.py @@ -85,35 +85,32 @@ def test_call_parameters_reach_expert_activation(monkeypatch, activation): expert_out = (gate.bfloat16() * up.bfloat16()).bfloat16() expected[row] += (expert_out.float() * probs[row, choice]).bfloat16().float() if kwargs: - default_output = weight.experts(x.clone(), router, 2, True, False, 0, 0) - actual = weight.experts(x.clone(), router, 2, True, False, 0, 0, **kwargs) + default_output = weight.experts(x.clone(), router, 2, True, False, 0, 0, False) + actual = weight.experts(x.clone(), router, 2, True, False, 0, 0, False, **kwargs) actual_with_topk = weight.experts_with_topk( x.clone(), topk_weights=probs, topk_ids=top.indices, + is_prefill=False, **kwargs, ) torch.testing.assert_close(actual, expected.bfloat16(), atol=0.125, rtol=0.01) torch.testing.assert_close(actual_with_topk, expected.bfloat16(), atol=0.125, rtol=0.01) if kwargs: # A clamped call must not change subsequent calls on the same weight. - actual_default = weight.experts(x.clone(), router, 2, True, False, 0, 0) + actual_default = weight.experts(x.clone(), router, 2, True, False, 0, 0, False) torch.testing.assert_close(actual_default, default_output, atol=0, rtol=0) def test_marlin_rejects_clamp_at_call(monkeypatch): backend = FuseMoeMarlin monkeypatch.setattr(backend, "create_workspace", lambda self: None) - monkeypatch.setattr(backend, "_select_experts", lambda *args, **kwargs: (None, None, None)) + monkeypatch.setattr(backend, "_select_experts", lambda *args, **kwargs: (None, None)) impl = backend( n_routed_experts=4, num_fused_shared_experts=0, routed_scaling_factor=1.0, quant_method=None, - redundancy_expert_num=0, - redundancy_expert_ids_tensor=None, - routed_expert_counter_tensor=None, - auto_update_redundancy_expert=False, ) with pytest.raises(NotImplementedError, match="does not support clamped SwiGLU"): impl( @@ -128,6 +125,7 @@ def test_marlin_rejects_clamp_at_call(monkeypatch): use_grouped_topk=False, topk_group=0, num_expert_group=0, + is_prefill=False, alpha=1.0, limit=10.0, clamp_up_add_one=False, diff --git a/unit_tests/common/fused_moe/test_append_shared_expert_topk.py b/unit_tests/common/fused_moe/test_append_shared_expert_topk.py index dbc8977fac..f071a187f3 100644 --- a/unit_tests/common/fused_moe/test_append_shared_expert_topk.py +++ b/unit_tests/common/fused_moe/test_append_shared_expert_topk.py @@ -4,6 +4,7 @@ import torch from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.triton_impl import FuseMoeTriton from lightllm.common.basemodel.triton_kernel.fused_moe.append_shared_expert_topk import ( append_fused_shared_experts, ) @@ -23,13 +24,17 @@ def run(**kwargs): seen.update(kwargs) return kwargs["input_tensor"] - weight = SimpleNamespace( + impl = SimpleNamespace( num_fused_shared_experts=1, n_routed_experts=4, + fused_experts_with_topk=run, + ) + impl._prepare_expert_execution = lambda **kwargs: FuseMoeTriton._prepare_expert_execution(impl, **kwargs) + weight = SimpleNamespace( layer_num_=0, w13=torch.ones(1, device="cuda"), w2=torch.ones(1, device="cuda"), - fuse_moe_impl=SimpleNamespace(fused_experts_with_topk=run), + fuse_moe_impl=impl, ) topk_ids = torch.tensor([[0, 2], [1, 3]], dtype=torch.int64, device="cuda") topk_weights = torch.tensor([[0.2, 0.8], [0.4, 0.6]], device="cuda") @@ -38,11 +43,85 @@ def run(**kwargs): torch.ones(2, 4, device="cuda"), topk_weights, topk_ids, + is_prefill=False, ) assert torch.equal(seen["topk_ids"], torch.tensor([[0, 2, 4], [1, 3, 4]], device="cuda")) assert torch.allclose(seen["topk_weights"], torch.tensor([[0.2, 0.8, 1.0], [0.4, 0.6, 1.0]], device="cuda")) +def test_capture_observes_logical_ids_before_execution_mapping(): + logical_ids = torch.tensor([[0, 2]], dtype=torch.int64, device="cuda") + logical_weights = torch.tensor([[0.2, 0.8]], device="cuda") + captured = {} + executed = {} + impl = FuseMoeTriton.__new__(FuseMoeTriton) + impl.n_routed_experts = 4 + impl.num_fused_shared_experts = 1 + impl._select_experts = lambda **kwargs: (logical_weights, logical_ids) + + def fused(**kwargs): + executed.update(kwargs) + return kwargs["input_tensor"] + + impl._fused_experts = fused + hidden = torch.ones(1, 4, device="cuda") + impl( + input_tensor=hidden, + router_logits=hidden, + w13=None, + w2=None, + correction_bias=None, + scoring_func="softmax", + top_k=2, + renormalize=False, + use_grouped_topk=False, + topk_group=0, + num_expert_group=0, + is_prefill=False, + moe_capture_callback=lambda ids: captured.setdefault("ids", ids), + ) + + assert captured["ids"] is logical_ids + assert torch.equal(logical_ids, torch.tensor([[0, 2]], device="cuda")) + assert torch.equal(executed["topk_ids"], torch.tensor([[0, 2, 4]], device="cuda")) + + +def test_redundancy_repair_does_not_overwrite_logical_ids(monkeypatch): + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl import ( + FuseMoeDeepGEMM, + ) + + logical_ids = torch.tensor([[1, 3]], dtype=torch.int64, device="cuda") + seen = {} + + def repair(topk_ids, **kwargs): + seen["ptr"] = topk_ids.data_ptr() + topk_ids.fill_(-1) + + monkeypatch.setattr( + "lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl.redundancy_topk_ids_repair", + repair, + ) + impl = FuseMoeDeepGEMM.__new__(FuseMoeDeepGEMM) + impl.redundancy_expert_num = 1 + impl.redundancy_expert_ids_tensor = None + impl.ep_n_routed_experts = 2 + impl.global_rank_ = 0 + impl.routed_expert_counter_tensor = None + impl.auto_update_redundancy_expert = False + impl.num_fused_shared_experts = 0 + impl.n_routed_experts = 4 + _, execution_ids = FuseMoeDeepGEMM._prepare_expert_execution( + impl, + torch.ones(1, 2, device="cuda"), + logical_ids, + is_prefill=False, + ) + assert seen["ptr"] != logical_ids.data_ptr() + assert torch.equal(logical_ids, torch.tensor([[1, 3]], device="cuda")) + assert torch.equal(execution_ids, torch.full_like(logical_ids, -1)) + + def test_append_fused_shared_experts_without_gate(): topk_ids = torch.tensor([[0, 2], [1, 3], [2, 0]], dtype=torch.int32, device="cuda") topk_weights = torch.tensor([[0.2, 0.8], [0.4, 0.6], [0.7, 0.3]], dtype=torch.float32, device="cuda") diff --git a/unit_tests/common/fused_moe/test_deepgemm_activation.py b/unit_tests/common/fused_moe/test_deepgemm_activation.py index e9ba2ab26e..528d9a1ade 100644 --- a/unit_tests/common/fused_moe/test_deepgemm_activation.py +++ b/unit_tests/common/fused_moe/test_deepgemm_activation.py @@ -31,15 +31,13 @@ def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ w2_scale = torch.ones(2, 1, 1, device="cuda") counts = torch.tensor([3, 2], device="cuda", dtype=torch.int32) activation_args = dict(alpha=alpha, limit=limit, clamp_up_add_one=add_one) - recv = torch.zeros(2, padded, dim, device="cuda", dtype=torch.float8_e4m3fn) - recv[0, :3] = x[:3] - recv[1, :2] = x[3:] if is_prefill: + recv = torch.zeros(2, padded, dim, device="cuda", dtype=torch.float8_e4m3fn) + recv[0, :3] = x[:3] + recv[1, :2] = x[3:] 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) actual = ep.chunked_expanded_moe_forward( num_recv_tokens_per_expert_list=[padded, padded], num_unaligned_recv_tokens_per_expert=counts, @@ -51,14 +49,29 @@ def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ w2=w2, w2_scale=w2_scale, block_size_k=128, - workspace=torch.empty(4 * 1024 * 1024, device="cuda", dtype=torch.uint8), hidden_dtype=torch.bfloat16, **activation_args, ) else: - output = ep.masked_group_gemm( - (recv, torch.ones(2, padded, 1, device="cuda")), - counts, + alignment = ep.get_mk_alignment_for_contiguous_layout() + starts = [] + psum_values = [] + cursor = 0 + for count in counts.tolist(): + if starts: + cursor = (cursor + alignment - 1) // alignment * alignment + starts.append(cursor) + cursor += count + psum_values.append(cursor) + psum = torch.tensor(psum_values, device="cuda", dtype=torch.int32) + max_rows = (int(psum[-1]) + alignment - 1) // alignment * alignment + recv = torch.zeros(max_rows, dim, device="cuda", dtype=torch.float8_e4m3fn) + recv[starts[0] : starts[0] + 3] = x[:3] + recv[starts[1] : starts[1] + 2] = x[3:] + output = ep.decode_masked_group_gemm( + (recv, torch.ones(max_rows, 1, device="cuda")), + psum, + alignment, torch.bfloat16, w1, w1_scale, @@ -67,7 +80,7 @@ def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ expected_m=3, **activation_args, ) - actual = torch.cat([output[0, :3], output[1, :2]]) + actual = torch.cat([output[starts[0] : starts[0] + 3], output[starts[1] : starts[1] + 2]]) gate = (x.float() * 2).bfloat16().float() up = (x.float() * 3).bfloat16().float() 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..f1c9b228a2 100644 --- a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py +++ b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py @@ -489,7 +489,6 @@ def test_ep_moe_rounds_scales_and_preserves_logical_block_size(monkeypatch, is_p torch.testing.assert_close(sf, expected_scales) seen = [] monkeypatch.setattr(ep.dist_group_manager, "ep_buffer", object(), raising=False) - monkeypatch.setattr(ep.dist_group_manager, "ep_low_latency_buffer", object(), raising=False) def check_dispatch(**kwargs): seen.append(kwargs["block_size_k"]) @@ -500,7 +499,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/models/deepseek_v4/test_shared_expert_overlap.py b/unit_tests/models/deepseek_v4/test_shared_expert_overlap.py index 13b5c46170..c088d811eb 100644 --- a/unit_tests/models/deepseek_v4/test_shared_expert_overlap.py +++ b/unit_tests/models/deepseek_v4/test_shared_expert_overlap.py @@ -3,8 +3,10 @@ from types import SimpleNamespace import pytest import torch -import lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep as ep_module -from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import fused_experts_impl +from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( + fused_experts_impl, + prefill_decode_sync as sync, +) from lightllm.models.deepseek_v4.layer_infer import transformer_layer_infer as layer_module from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import DeepseekV4TransformerLayerWeight @@ -63,26 +65,45 @@ def routed(**kwargs): def test_decode_runs_shared_expert_between_combine_send_and_recv(monkeypatch): events = [] hidden, topk_weights, topk_idx, w1, w2 = _moe_inputs() + import deep_ep + + monkeypatch.setattr(deep_ep.ElasticBuffer, "capture", staticmethod(lambda: events.append("event") or "event")) + monkeypatch.setattr(sync, "per_token_group_quant_fp8", lambda *args, **kwargs: (hidden, hidden)) + monkeypatch.setattr(sync, "get_mk_alignment_for_contiguous_layout", lambda: 128) + monkeypatch.setattr(sync, "get_deepep_num_max_dispatch_tokens_per_rank_decode", lambda: 16) + monkeypatch.setattr(sync, "get_ep_num_sms", lambda overlap_with_compute: 8) def dispatch(*args, **kwargs): events.append("dispatch") - return (hidden, hidden), torch.ones(2, dtype=torch.int32), "handle", None, None + handle = SimpleNamespace( + psum_num_recv_tokens_per_expert=torch.tensor([3, 3]), + expert_alignment=128, + recv_src_metadata=torch.zeros(3, 4, dtype=torch.int32), + psum_num_recv_tokens_per_scaleup_rank=torch.tensor([3]), + ) + return (hidden, hidden), None, topk_weights, handle, None - def gemm(*args, **kwargs): - assert args[7] == 3 # ceil(3 tokens * group 8 * topk 2 / 16 experts) + def gemm(**kwargs): + assert kwargs["expected_m"] == 3 # ceil(3 tokens * 8 ranks * topk 2 / 16 experts) events.append("gemm") return hidden - def combine(*args, **kwargs): - assert kwargs["return_recv_hook"] is True - events.append("combine_send") + def reduce(**kwargs): + events.append("reduce") + return hidden, kwargs["recv_src_metadata"] - def receive(): + class _Event: + def current_stream_wait(self): events.append("combine_recv") - return torch.full((3, 4), 3.0), None, receive + def combine(*args, **kwargs): + assert kwargs["previous_event"] == "event" + assert kwargs["async_with_compute_stream"] is True + events.append("combine_send") + return torch.full((3, 4), 3.0), None, _Event() - monkeypatch.setattr(ep_module, "masked_group_gemm", gemm) + monkeypatch.setattr(sync, "decode_masked_group_gemm", gemm) + monkeypatch.setattr(sync, "ep_reduce_decode_output", reduce) result = fused_experts_impl( hidden, w1, @@ -90,13 +111,13 @@ def receive(): topk_weights, topk_idx, num_experts=16, - buffer=SimpleNamespace(group_size=8, low_latency_dispatch=dispatch, low_latency_combine=combine), + buffer=SimpleNamespace(num_ranks=8, dispatch=dispatch, combine=combine), is_prefill=False, use_fp8_w8a8=True, use_fp8_all2all=True, shared_expert_func=lambda: events.append("shared") or torch.full((3, 4), 2.0), ) - assert events == ["dispatch", "gemm", "combine_send", "shared", "combine_recv"] + assert events == ["dispatch", "gemm", "reduce", "event", "combine_send", "shared", "combine_recv"] assert torch.all(result == 5) @@ -106,8 +127,10 @@ def test_prefill_overlaps_shared_expert_with_combine(monkeypatch): import deep_ep monkeypatch.setattr(deep_ep.ElasticBuffer, "capture", staticmethod(lambda: events.append("event") or "event")) - monkeypatch.setattr(ep_module, "get_deepep_num_max_dispatch_tokens_per_rank_prefill", lambda: 16) - monkeypatch.setattr(ep_module, "per_token_group_quant_fp8", lambda *args, **kwargs: (hidden, hidden)) + monkeypatch.setattr(sync, "get_deepep_num_max_dispatch_tokens_per_rank_prefill", lambda: 16) + monkeypatch.setattr(sync, "per_token_group_quant_fp8", lambda *args, **kwargs: (hidden, hidden)) + monkeypatch.setattr(sync, "get_mk_alignment_for_contiguous_layout", lambda: 128) + monkeypatch.setattr(sync, "get_ep_num_sms", lambda overlap_with_compute: 8) def dispatch(*args, **kwargs): events.append("dispatch") 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()