From 148586b727582289e95b5d640d82dc90a0fa7bf7 Mon Sep 17 00:00:00 2001 From: irexyc Date: Thu, 2 Jul 2026 04:32:57 +0000 Subject: [PATCH 01/19] support convert model with ep --- lmdeploy/cli/serve.py | 4 ++- lmdeploy/messages.py | 3 ++ lmdeploy/turbomind/builders/_base.py | 41 +++++++++++++++++++-- lmdeploy/turbomind/builders/moe.py | 42 +++++++++++++++++++++- lmdeploy/turbomind/model_loader.py | 3 ++ lmdeploy/turbomind/models/glm4_moe_lite.py | 11 +++--- lmdeploy/turbomind/models/gpt_oss.py | 7 ++-- lmdeploy/turbomind/models/internvl.py | 3 +- lmdeploy/turbomind/models/mixtral.py | 8 ++--- lmdeploy/turbomind/models/qwen2.py | 12 +++---- lmdeploy/turbomind/models/qwen3.py | 8 ++--- lmdeploy/turbomind/models/qwen3_5.py | 13 ++++--- lmdeploy/turbomind/text_model.py | 3 +- lmdeploy/turbomind/turbomind.py | 21 ++++++----- src/turbomind/engine/engine_config.h | 2 ++ src/turbomind/models/llama/llama_params.h | 1 + src/turbomind/models/moe_weight.cc | 14 +++++--- src/turbomind/models/moe_weight.h | 12 +++++++ src/turbomind/python/bind.cpp | 1 + src/turbomind/turbomind.cc | 10 ++++-- src/turbomind/turbomind.h | 3 ++ 21 files changed, 166 insertions(+), 56 deletions(-) diff --git a/lmdeploy/cli/serve.py b/lmdeploy/cli/serve.py index a801f93f89..0905e488b9 100644 --- a/lmdeploy/cli/serve.py +++ b/lmdeploy/cli/serve.py @@ -130,7 +130,7 @@ def add_parser_api_server(): hf_overrides = ArgumentHelper.hf_overrides(pt_group) disable_metrics = ArgumentHelper.disable_metrics(pt_group) dp = ArgumentHelper.dp(pt_group) - ArgumentHelper.ep(pt_group) + ep_act = ArgumentHelper.ep(pt_group) ArgumentHelper.enable_microbatch(pt_group) ArgumentHelper.enable_eplb(pt_group) ArgumentHelper.role(pt_group) @@ -158,6 +158,7 @@ def add_parser_api_server(): tb_group._group_actions.append(hf_overrides) tb_group._group_actions.append(disable_metrics) tb_group._group_actions.append(dp) + tb_group._group_actions.append(ep_act) tb_group._group_actions.append(disable_vision_encoder) ArgumentHelper.cp(tb_group) ArgumentHelper.rope_scaling_factor(tb_group) @@ -270,6 +271,7 @@ def api_server(args): tp=args.tp, dp=args.dp, cp=args.cp, + ep=args.ep, nnodes=args.nnodes, node_rank=args.node_rank, dist_init_addr=args.dist_init_addr, diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index 9486cc62b0..2bffcd2d92 100644 --- a/lmdeploy/messages.py +++ b/lmdeploy/messages.py @@ -280,6 +280,8 @@ class TurbomindEngineConfig: tp: int = 1 dp: int = 1 cp: int = 1 + ep: int = 1 + all2all_backend: str = 'allgather_reducescatter' device_num: int = None attn_tp_size: int = None attn_cp_size: int = None @@ -317,6 +319,7 @@ def __post_init__(self): """Check input validation.""" assert self.dtype in ['auto', 'float16', 'bfloat16'] assert self.tp >= 1, 'tp must be a positive integer' + assert self.ep >= 1, 'ep must be a positive integer' assert self.cache_max_entry_count > 0, 'invalid cache_max_entry_count' try: self.quant_policy = QuantPolicy(self.quant_policy) diff --git a/lmdeploy/turbomind/builders/_base.py b/lmdeploy/turbomind/builders/_base.py index b332b2b7c7..023c02faa8 100644 --- a/lmdeploy/turbomind/builders/_base.py +++ b/lmdeploy/turbomind/builders/_base.py @@ -142,10 +142,33 @@ class Context: def __init__(self, devices, data_type): self.devices = devices self.data_type = data_type + self._active_mask_stack = [(True,) * len(devices)] + + @property + def active_mask(self): + return self._active_mask_stack[-1] + + def active_mask_scope(self, active_mask): + return _ActiveMaskScope(self, active_mask) + + +class _ActiveMaskScope: + """Temporarily restrict builders created under the context.""" + + def __init__(self, ctx, active_mask): + self._ctx = ctx + self._active_mask = tuple(bool(x) for x in active_mask) + + def __enter__(self): + self._ctx._active_mask_stack.append(self._active_mask) + + def __exit__(self, exc_type, exc, tb): + self._ctx._active_mask_stack.pop() + return False class ParallelGroup: - """Bundle a parallelism size with per-GPU TP ranks.""" + """Bundle a parallelism size with per-GPU ranks.""" def __init__(self, size, ranks): self.size = size self.ranks = ranks @@ -181,6 +204,7 @@ def __init__(self, config, ctx): # __setattr__. self._built = False self._ctx = ctx + self._active_mask = ctx.active_mask self.tp = ParallelGroup(1, None) # default: no TP self.config = config if hasattr(self.config, 'data_type'): @@ -220,6 +244,9 @@ def _rank_for(self, gpu_idx: int) -> int: return self.tp.ranks[gpu_idx] return 0 + def _is_active(self, gpu_idx: int) -> bool: + return self._active_mask[gpu_idx] + # ------------------------------------------------------------------ # Add methods — stage into pending dicts (pre-build only) # ------------------------------------------------------------------ @@ -280,6 +307,9 @@ def _add_linear(self, name: str, linear: Linear, # --- Per-GPU: standalone creation + tensor copy -------------------- handles = [] for i, ctx in enumerate(self._ctx.devices): + if not self._is_active(i): + handles.append(None) + continue with ctx: rank = self._rank_for(i) if tp > 1 else 0 @@ -363,6 +393,9 @@ def _create_handles(self): """Create one C++ module per context via ``_tm.create_module(cfg)``.""" handles = [] for i, ctx in enumerate(self._ctx.devices): + if not self._is_active(i): + handles.append(None) + continue with ctx: cfg = self._cfg_for_rank(i) handle = _tm.create_module(cfg) @@ -380,7 +413,9 @@ def _cfg_for_rank(self, gpu_idx: int): def _commit_child(self, name: str, handles: list): """Attach pre-created per-GPU child handles to parent handles.""" for i, (parent_h, child_h) in enumerate( - zip(self._handles, handles)): + zip(self._handles, handles, strict=True)): + if parent_h is None or child_h is None: + continue with self._ctx.devices[i]: parent_h.add_child_raw(name, child_h) @@ -405,6 +440,8 @@ def _commit_tensor(self, name: str, tensor: torch.Tensor, split_dim = _SPLIT_SIDE_TO_DIM.get(split_side) if split_side else None for i, handle in enumerate(self._handles): + if handle is None: + continue with self._ctx.devices[i]: rank = self._rank_for(i) if tp > 1 else 0 shard = _shard(tensor, split_dim, tp, rank) diff --git a/lmdeploy/turbomind/builders/moe.py b/lmdeploy/turbomind/builders/moe.py index a4a1bcb36d..62ef65fcdc 100644 --- a/lmdeploy/turbomind/builders/moe.py +++ b/lmdeploy/turbomind/builders/moe.py @@ -1,7 +1,8 @@ # Copyright (c) OpenMMLab. All rights reserved. from __future__ import annotations -from ._base import Builder, SplitSide +from ._base import Builder, ParallelGroup, SplitSide +from .module_list import ModuleListBuilder, ModuleListConfig # --------------------------------------------------------------------------- # MoeBuilder -- gate, non-expert params @@ -11,6 +12,24 @@ class MoeBuilder(Builder): """MoE weight loading builder.""" + def __init__(self, config, ctx, ep: ParallelGroup | None = None): + super().__init__(config, ctx) + self.ep = ep or ParallelGroup(1, None) + if self.ep.size > 1 and config.expert_num % self.ep.size != 0: + raise ValueError( + f'num_experts={config.expert_num} must be divisible by ' + f'ep={self.ep.size}') + self.config.ep_size = self.ep.size + + def _cfg_for_rank(self, gpu_idx: int): + """Set this GPU context's EP rank when EP is active.""" + if self.ep.size <= 1: + return super()._cfg_for_rank(gpu_idx) + else: + cfg = self.config.clone() + cfg.ep_rank = self.ep.ranks[gpu_idx] + return cfg + def add_gate(self, name, linear): """Commit a gate linear (broadcast, no split).""" self._add_linear(name, linear, split_side=None) @@ -20,3 +39,24 @@ def add_param(self, name, tensor, split_side=None): if split_side is not None and not isinstance(split_side, SplitSide): split_side = None # specs may pass None for broadcast self._add_tensor(name, tensor, split_side) + + def add_experts(self, build_expert, name='experts'): + """Build and attach expert modules with contiguous EP ownership.""" + experts = ModuleListBuilder(ModuleListConfig(), self._ctx) + for expert_idx in range(self.config.expert_num): + active_mask = self._expert_active_mask(expert_idx) + if not any(active_mask): + continue + with self._ctx.active_mask_scope(active_mask): + experts[expert_idx] = build_expert(expert_idx) + setattr(self, name, experts.build()) + + def _expert_active_mask(self, expert_idx: int): + ep_size = self.ep.size + if ep_size <= 1: + return [True] * len(self._ctx.devices) + ranks = self.ep.ranks + assert ranks is not None + local = self.config.expert_num // ep_size + return [rank * local <= expert_idx < (rank + 1) * local + for rank in ranks] diff --git a/lmdeploy/turbomind/model_loader.py b/lmdeploy/turbomind/model_loader.py index c7f8851f01..d79406b456 100644 --- a/lmdeploy/turbomind/model_loader.py +++ b/lmdeploy/turbomind/model_loader.py @@ -36,6 +36,8 @@ def _bind_runtime(self): [mc.attn_tp_rank(g) for g in range(self.gpu_count)]) mlp_tp = ParallelGroup(ec.mlp_tp_size, [mc.mlp_tp_rank(g) for g in range(self.gpu_count)]) + ep = ParallelGroup(ec.ep, + [mc.ep_rank(g) for g in range(self.gpu_count)]) model_tp = ParallelGroup(ec.attn_tp_size * ec.attn_cp_size, [mc.model_tp_rank(g) for g in range(self.gpu_count)]) @@ -44,6 +46,7 @@ def _bind_runtime(self): root_handles=[mc.root(g) for g in range(self.gpu_count)], attn_tp=attn_tp, mlp_tp=mlp_tp, + ep=ep, model_tp=model_tp, ) diff --git a/lmdeploy/turbomind/models/glm4_moe_lite.py b/lmdeploy/turbomind/models/glm4_moe_lite.py index e779633460..3737b2b167 100644 --- a/lmdeploy/turbomind/models/glm4_moe_lite.py +++ b/lmdeploy/turbomind/models/glm4_moe_lite.py @@ -115,18 +115,17 @@ def ffn(self, pfx, inter_size, is_expert=False): def moe(self, pfx): cfg = self._moe_cfg.clone() - m = MoeBuilder(cfg, self._ctx) + m = MoeBuilder(cfg, self._ctx, ep=self._ep) m.add_gate('gate', self._linear(pfx + 'gate')) correction = pfx.pop('gate.e_score_correction_bias') m.add_param('score_correction_bias', correction) - experts = ModuleListBuilder(ModuleListConfig(), self._ctx) - for e in range(cfg.expert_num): - experts[e] = self.ffn(pfx + 'experts' + e, - self.cfg.moe_intermediate_size, is_expert=True) - m.experts = experts.build() + m.add_experts( + lambda e: self.ffn(pfx + 'experts' + e, + self.cfg.moe_intermediate_size, + is_expert=True)) shared = self.ffn(pfx + 'shared_experts', self.cfg.intermediate_size * self.cfg.n_shared_experts) diff --git a/lmdeploy/turbomind/models/gpt_oss.py b/lmdeploy/turbomind/models/gpt_oss.py index 5480eaee7b..4453cf101f 100644 --- a/lmdeploy/turbomind/models/gpt_oss.py +++ b/lmdeploy/turbomind/models/gpt_oss.py @@ -107,13 +107,10 @@ def reorder(x): def moe(self, pfx): cfg = self._moe_cfg.clone() - m = MoeBuilder(cfg, self._ctx) + m = MoeBuilder(cfg, self._ctx, ep=self._ep) m.add_gate('gate', self._linear(pfx + 'router')) experts_pfx = pfx + 'experts' - experts = ModuleListBuilder(ModuleListConfig(), self._ctx) - for e in range(cfg.expert_num): - experts[e] = self._packed_moe_ffn(experts_pfx, e) - m.experts = experts.build() + m.add_experts(lambda e: self._packed_moe_ffn(experts_pfx, e)) return m.build() def layers(self, pfx): diff --git a/lmdeploy/turbomind/models/internvl.py b/lmdeploy/turbomind/models/internvl.py index ee3f344850..022f9d91f8 100644 --- a/lmdeploy/turbomind/models/internvl.py +++ b/lmdeploy/turbomind/models/internvl.py @@ -60,12 +60,13 @@ def __init__(self, cfg: PretrainedConfig, *, resolver): self.vision_model = None def bind_runtime(self, *, ctx, root_handles, - attn_tp, mlp_tp, model_tp): + attn_tp, mlp_tp, ep, model_tp): self.text_model.bind_runtime( ctx=ctx, root_handles=root_handles, attn_tp=attn_tp, mlp_tp=mlp_tp, + ep=ep, model_tp=model_tp, ) diff --git a/lmdeploy/turbomind/models/mixtral.py b/lmdeploy/turbomind/models/mixtral.py index f6c41b4e3a..ffd1d75963 100644 --- a/lmdeploy/turbomind/models/mixtral.py +++ b/lmdeploy/turbomind/models/mixtral.py @@ -83,13 +83,11 @@ def ffn(self, pfx, is_expert=False): def moe(self, pfx): cfg = self._moe_cfg.clone() - m = MoeBuilder(cfg, self._ctx) + m = MoeBuilder(cfg, self._ctx, ep=self._ep) m.add_gate('gate', self._linear(pfx + 'gate')) - experts = ModuleListBuilder(ModuleListConfig(), self._ctx) - for e in range(self._n_experts): - experts[e] = self.ffn(pfx + 'experts' + e, is_expert=True) - m.experts = experts.build() + m.add_experts( + lambda e: self.ffn(pfx + 'experts' + e, is_expert=True)) return m.build() diff --git a/lmdeploy/turbomind/models/qwen2.py b/lmdeploy/turbomind/models/qwen2.py index 64b5accb92..7b73130605 100644 --- a/lmdeploy/turbomind/models/qwen2.py +++ b/lmdeploy/turbomind/models/qwen2.py @@ -110,16 +110,14 @@ def ffn(self, pfx, inter_size, is_expert=False): def moe(self, pfx): cfg = self._moe_cfg.clone() - m = MoeBuilder(cfg, self._ctx) + m = MoeBuilder(cfg, self._ctx, ep=self._ep) m.add_gate('gate', self._linear(pfx + 'gate')) - experts = ModuleListBuilder(ModuleListConfig(), self._ctx) - for e in range(self.cfg.num_experts): - experts[e] = self.ffn(pfx + 'experts' + e, - self.cfg.moe_intermediate_size, - is_expert=True) - m.experts = experts.build() + m.add_experts( + lambda e: self.ffn(pfx + 'experts' + e, + self.cfg.moe_intermediate_size, + is_expert=True)) m.add_gate('shared_gate', self._linear(pfx + 'shared_expert_gate')) shared = self.ffn(pfx + 'shared_expert', diff --git a/lmdeploy/turbomind/models/qwen3.py b/lmdeploy/turbomind/models/qwen3.py index f30f1d95b8..398ff38ca7 100644 --- a/lmdeploy/turbomind/models/qwen3.py +++ b/lmdeploy/turbomind/models/qwen3.py @@ -110,13 +110,11 @@ def ffn(self, pfx, is_expert=False): def moe(self, pfx): cfg = self._moe_cfg.clone() - m = MoeBuilder(cfg, self._ctx) + m = MoeBuilder(cfg, self._ctx, ep=self._ep) m.add_gate('gate', self._linear(pfx + 'gate')) - experts = ModuleListBuilder(ModuleListConfig(), self._ctx) - for e in range(self.cfg.num_experts): - experts[e] = self.ffn(pfx + 'experts' + e, is_expert=True) - m.experts = experts.build() + m.add_experts( + lambda e: self.ffn(pfx + 'experts' + e, is_expert=True)) return m.build() diff --git a/lmdeploy/turbomind/models/qwen3_5.py b/lmdeploy/turbomind/models/qwen3_5.py index 0d13f372ae..2c55ea6c9b 100644 --- a/lmdeploy/turbomind/models/qwen3_5.py +++ b/lmdeploy/turbomind/models/qwen3_5.py @@ -204,16 +204,14 @@ def ffn(self, pfx, inter_size, is_expert=False): def moe(self, pfx): cfg = self._moe_cfg.clone() - m = MoeBuilder(cfg, self._ctx) + m = MoeBuilder(cfg, self._ctx, ep=self._ep) m.add_gate('gate', self._linear(pfx + 'gate')) experts_pfx = pfx + 'experts' - experts = ModuleListBuilder(ModuleListConfig(), self._ctx) - for e in range(self._n_experts): - experts[e] = self._moe_expert_ffn( - experts_pfx, e, self.cfg.moe_intermediate_size) - m.experts = experts.build() + m.add_experts( + lambda e: self._moe_expert_ffn( + experts_pfx, e, self.cfg.moe_intermediate_size)) m.add_gate('shared_gate', self._linear(pfx + 'shared_expert_gate')) shared = self.ffn(pfx + 'shared_expert', self.cfg.shared_expert_intermediate_size) @@ -592,7 +590,7 @@ def __init__(self, cfg: Qwen3_5Config | Qwen3_5MoeConfig, *, resolver, vision_cfg, resolver=vision_resolver or resolver) def bind_runtime(self, *, ctx, root_handles, - attn_tp, mlp_tp, model_tp): + attn_tp, mlp_tp, ep, model_tp): for m in (self.text_model, self.vision_model): if m is not None: m.bind_runtime( @@ -600,6 +598,7 @@ def bind_runtime(self, *, ctx, root_handles, root_handles=root_handles, attn_tp=attn_tp, mlp_tp=mlp_tp, + ep=ep, model_tp=model_tp, ) diff --git a/lmdeploy/turbomind/text_model.py b/lmdeploy/turbomind/text_model.py index 9d86fd5cff..75842f6e07 100644 --- a/lmdeploy/turbomind/text_model.py +++ b/lmdeploy/turbomind/text_model.py @@ -40,11 +40,12 @@ def _vocab_size(self) -> int: return self.cfg.vocab_size def bind_runtime(self, *, ctx, root_handles, - attn_tp, mlp_tp, model_tp): + attn_tp, mlp_tp, ep, model_tp): self._ctx = ctx self._root_handles = root_handles self._attn_tp = attn_tp self._mlp_tp = mlp_tp + self._ep = ep self._model_tp = model_tp def _linear(self, pfx: Prefix, *, diff --git a/lmdeploy/turbomind/turbomind.py b/lmdeploy/turbomind/turbomind.py index 2df03cf340..dd924e704d 100644 --- a/lmdeploy/turbomind/turbomind.py +++ b/lmdeploy/turbomind/turbomind.py @@ -82,26 +82,29 @@ def complete_parallel_config(cfg: TurbomindEngineConfig): def update_parallel_config(cfg: TurbomindEngineConfig): cfg.device_num = len(cfg.devices) * cfg.nnodes if cfg.devices else cfg.device_num + assert cfg.ep == 1 or cfg.tp == 1 if not complete_parallel_config(cfg): - total = cfg.dp * cfg.tp + total = cfg.dp * cfg.ep * cfg.tp if not cfg.device_num: count = torch.cuda.device_count() * cfg.nnodes if total < count: count = total cfg.device_num = count assert total % cfg.device_num == 0 + size = max(cfg.ep, cfg.tp) overlap = total // cfg.device_num - attn_dp_size = overlap - mlp_tp_size = overlap - inner_tp_size = cfg.tp // mlp_tp_size - cfg.outer_dp_size = cfg.dp // attn_dp_size - cfg.attn_dp_size = attn_dp_size + inner_tp_size = size // overlap + cfg.outer_dp_size = cfg.dp // overlap + cfg.attn_dp_size = overlap cfg.attn_tp_size = inner_tp_size // cfg.cp cfg.attn_cp_size = cfg.cp cfg.mlp_dp_size = 1 - cfg.mlp_tp_size = mlp_tp_size * inner_tp_size - assert cfg.attn_dp_size * cfg.attn_tp_size * cfg.attn_cp_size == cfg.mlp_dp_size * cfg.mlp_tp_size + cfg.mlp_tp_size = total // cfg.ep if cfg.ep > 1 else overlap * inner_tp_size + if cfg.ep > 1: + assert cfg.communicator == 'nccl', f'{cfg.communicator} communicator does not support ep > 1' + assert cfg.mlp_tp_size == 1, 'Only support mlp_tp_size == 1 when ep > 1' assert cfg.attn_dp_size * cfg.attn_tp_size * cfg.attn_cp_size * cfg.outer_dp_size == cfg.device_num + assert cfg.attn_dp_size * cfg.attn_tp_size * cfg.attn_cp_size == cfg.mlp_dp_size * cfg.mlp_tp_size * cfg.ep # update devices cfg.devices = cfg.devices or list(range(cfg.device_num // cfg.nnodes)) cfg.devices = cfg.devices[:cfg.device_num // cfg.nnodes] @@ -247,6 +250,8 @@ def _from_hf(self, model_path: str, engine_config: TurbomindEngineConfig, ec.attn_tp_size = engine_config.attn_tp_size ec.attn_cp_size = engine_config.attn_cp_size ec.mlp_tp_size = engine_config.mlp_tp_size + ec.ep_size = engine_config.ep + ec.all2all_backend = engine_config.all2all_backend ec.devices = engine_config.devices ec.nnodes = engine_config.nnodes ec.node_rank = engine_config.node_rank diff --git a/src/turbomind/engine/engine_config.h b/src/turbomind/engine/engine_config.h index 2c0381e9c4..04aed14892 100644 --- a/src/turbomind/engine/engine_config.h +++ b/src/turbomind/engine/engine_config.h @@ -32,6 +32,8 @@ struct EngineConfig { X(int, attn_tp_size) \ X(int, attn_cp_size) \ X(int, mlp_tp_size) \ + X(int, ep_size, 1) \ + X(std::string, all2all_backend, "allgather_reducescatter") \ X(std::vector, devices) \ X(int, nnodes) \ X(int, node_rank) \ diff --git a/src/turbomind/models/llama/llama_params.h b/src/turbomind/models/llama/llama_params.h index 4ce0a586fa..23cc89d82a 100644 --- a/src/turbomind/models/llama/llama_params.h +++ b/src/turbomind/models/llama/llama_params.h @@ -15,6 +15,7 @@ struct EngineParam: EngineConfig { int attn_tp_rank = 0; int attn_cp_rank = 0; int mlp_tp_rank = 0; + int ep_rank = 0; int model_tp_rank = 0; // rank(d_tp_group), in [0, attn_tp_size × attn_cp_size) // Derived field (set in Impl ctor) diff --git a/src/turbomind/models/moe_weight.cc b/src/turbomind/models/moe_weight.cc index 2020e82e1c..3c34b2a183 100644 --- a/src/turbomind/models/moe_weight.cc +++ b/src/turbomind/models/moe_weight.cc @@ -23,6 +23,8 @@ MoeWeight::MoeWeight(const core::MoeConfig& cfg) act_type_ = static_cast(cfg.act_type); fuse_silu_act_ = cfg.fuse_silu; expert_num = cfg.expert_num; + ep_size = cfg.ep_size; + ep_rank = cfg.ep_rank; } // Adapted from LinkExperts for LinearWeight @@ -79,7 +81,7 @@ FfnWeight* MoeWeight::expert(int i) const if (!experts) { return nullptr; } - return static_cast(experts->child(std::to_string(i))); + return static_cast(experts->child(std::to_string(local_expert_offset() + i))); } void MoeWeight::prepare() @@ -87,6 +89,8 @@ void MoeWeight::prepare() // First prepare all children (experts, gate, etc.) Module::prepare(); + const int local_expert_num = num_local_experts(); + // Create batched block view for fused MoE path auto e0 = TM_CHECK_NOTNULL(expert(0)); // exemplar expert @@ -123,23 +127,23 @@ void MoeWeight::prepare() if (get_expert_w1w3(0)) { // Fused w1w3 path: experts have a single fused gate+up projection block_->add_child("w1w3", std::make_unique()); - LinkLinearExperts(get_expert_w1w3, expert_num, *block_->w1w3); + LinkLinearExperts(get_expert_w1w3, local_expert_num, *block_->w1w3); } else { // Separate w1/w3 path: link individually block_->add_child("w1", std::make_unique()); block_->add_child("w3", std::make_unique()); if (get_expert_w1(0)) { - LinkLinearExperts(get_expert_w1, expert_num, *block_->w1); + LinkLinearExperts(get_expert_w1, local_expert_num, *block_->w1); } if (get_expert_w3(0)) { - LinkLinearExperts(get_expert_w3, expert_num, *block_->w3); + LinkLinearExperts(get_expert_w3, local_expert_num, *block_->w3); } } block_->add_child("w2", std::make_unique()); if (get_expert_w2(0)) { - LinkLinearExperts(get_expert_w2, expert_num, *block_->w2); + LinkLinearExperts(get_expert_w2, local_expert_num, *block_->w2); } // Propagate the actual fused-silu state from the first expert to diff --git a/src/turbomind/models/moe_weight.h b/src/turbomind/models/moe_weight.h index 4b767e6710..9bdeea5d93 100644 --- a/src/turbomind/models/moe_weight.h +++ b/src/turbomind/models/moe_weight.h @@ -26,6 +26,8 @@ struct MoeConfig: ModuleConfig { X(int, n_group) \ X(int, router_n_groups) \ X(double, routed_scale) \ + X(int, ep_size, 1) \ + X(int, ep_rank, 0) \ X(DataType, data_type) MOE_FIELDS(TM_MEMBER) @@ -54,6 +56,14 @@ class MoeWeight: public core::Module { { return expert_num; } + int num_local_experts() const + { + return expert_num / ep_size; + } + int local_expert_offset() const + { + return ep_rank * num_local_experts(); + } // --- X-macro child members --- #define MOE_WEIGHT_CHILDREN(X) \ @@ -82,6 +92,8 @@ class MoeWeight: public core::Module { int n_group{}; std::string scoring_func; int router_n_groups{}; + int ep_size{1}; + int ep_rank{0}; private: ActivationType act_type_{}; diff --git a/src/turbomind/python/bind.cpp b/src/turbomind/python/bind.cpp index 1a7a107a7f..9d070bcc7d 100644 --- a/src/turbomind/python/bind.cpp +++ b/src/turbomind/python/bind.cpp @@ -793,5 +793,6 @@ PYBIND11_MODULE(_turbomind, m) .def("is_dummy_node", [](TurboMind* model) { return model->is_dummy_node(); }) .def("attn_tp_rank", &TurboMind::GetAttnTpRank, "index"_a) .def("mlp_tp_rank", &TurboMind::GetMlpTpRank, "index"_a) + .def("ep_rank", &TurboMind::GetEpRank, "index"_a) .def("model_tp_rank", &TurboMind::GetModelTpRank, "index"_a); } diff --git a/src/turbomind/turbomind.cc b/src/turbomind/turbomind.cc index 9ed314bff2..a3733c0390 100644 --- a/src/turbomind/turbomind.cc +++ b/src/turbomind/turbomind.cc @@ -167,7 +167,7 @@ TurboMind::Impl::Impl(string model_dir, EngineConfig config, FFICtxFactory ffi_c } comm_size_ = engine_param_.attn_dp_size * engine_param_.attn_tp_size * engine_param_.attn_cp_size; - TM_CHECK(engine_param_.mlp_tp_size == comm_size_); + TM_CHECK(engine_param_.mlp_tp_size * engine_param_.ep_size == comm_size_); communicator_type_ = std::move(config.communicator); @@ -243,7 +243,8 @@ void TurboMind::Impl::CreateContext(int index) p.model_tp_rank = c.d_comm->rank(c.d_tp_group); p.attn_tp_rank = p.model_tp_rank / p.attn_cp_size; - p.mlp_tp_rank = c.d_comm->rank(0); + p.mlp_tp_rank = inner_rank % p.mlp_tp_size; + p.ep_rank = inner_rank % p.ep_size; } if (c.h_tp_group->rank() == 0) { @@ -516,6 +517,11 @@ int TurboMind::GetMlpTpRank(int index) return impl_->engine_params_.at(index).mlp_tp_rank; } +int TurboMind::GetEpRank(int index) +{ + return impl_->engine_params_.at(index).ep_rank; +} + int TurboMind::GetModelTpRank(int index) { return impl_->engine_params_.at(index).model_tp_rank; diff --git a/src/turbomind/turbomind.h b/src/turbomind/turbomind.h index 4d19d12641..74d602493c 100644 --- a/src/turbomind/turbomind.h +++ b/src/turbomind/turbomind.h @@ -52,6 +52,9 @@ class TurboMind { /// MLP TP rank for GPU *index*. int GetMlpTpRank(int index); + /// Expert-parallel rank for GPU *index*. + int GetEpRank(int index); + /// Model-level TP rank (rank within d_tp_group) for GPU *index*. int GetModelTpRank(int index); From 4033613707ba719cd201e6e358b3c4ff2f5491ff Mon Sep 17 00:00:00 2001 From: irexyc Date: Thu, 2 Jul 2026 08:40:09 +0000 Subject: [PATCH 02/19] add MoeFfnAgRsImpl stub --- src/turbomind/comm/device_comm.h | 21 ++++ src/turbomind/comm/nccl/nccl.cu | 60 +++++++++- src/turbomind/models/llama/moe_ffn_layer.cc | 125 +++++++++++++++++++- src/turbomind/models/llama/moe_ffn_layer.h | 43 ++----- 4 files changed, 208 insertions(+), 41 deletions(-) diff --git a/src/turbomind/comm/device_comm.h b/src/turbomind/comm/device_comm.h index a6948762df..9ce0be0d12 100644 --- a/src/turbomind/comm/device_comm.h +++ b/src/turbomind/comm/device_comm.h @@ -5,6 +5,7 @@ #include #include +#include #include @@ -54,6 +55,16 @@ class DeviceCommImpl { int group, cudaStream_t stream) = 0; + virtual void AllGatherV(const void* sendbuff, + void* recvbuff, + const std::vector& counts, + DataType type, + int group, + cudaStream_t stream) + { + throw std::runtime_error("not implemented"); + } + virtual void ReduceScatter(const void* sendbuff, // void* recvbuff, size_t recvcount, @@ -64,6 +75,16 @@ class DeviceCommImpl { throw std::runtime_error("not implemented"); } + virtual void ReduceScatterV(const void* sendbuff, + void* recvbuff, + const std::vector& counts, + DataType type, + int group, + cudaStream_t stream) + { + throw std::runtime_error("not implemented"); + } + virtual void AllreduceResidualBiasRMSnorm(void* hidden, void* residual, const void* bias, diff --git a/src/turbomind/comm/nccl/nccl.cu b/src/turbomind/comm/nccl/nccl.cu index 67e0fef181..f0a7363854 100644 --- a/src/turbomind/comm/nccl/nccl.cu +++ b/src/turbomind/comm/nccl/nccl.cu @@ -262,6 +262,35 @@ public: NCCLCHECK(ncclGroupEnd()); } + void AllGatherV(const void* sendbuff, + void* recvbuff, + const std::vector& counts, + DataType type, + int group, + cudaStream_t stream) override + { + const size_t elem_size = byte_size(type); + const ncclDataType_t nccl_type = to_nccl_dtype(type); + const int n_ranks = this->n_ranks(group); + const int rank = this->rank(group); + ncclComm_t comm = groups_.at(group); + + TM_CHECK_EQ((int)counts.size(), n_ranks); + + NCCLCHECK(ncclGroupStart()); + size_t offset = 0; + for (int i = 0; i < n_ranks; ++i) { + const size_t count = counts[i]; + if (count) { + auto* recv = static_cast(recvbuff) + elem_size * offset; + const void* send = i == rank ? sendbuff : recv; + NCCLCHECK(ncclBroadcast(send, recv, count, nccl_type, i, comm, stream)); + } + offset += count; + } + NCCLCHECK(ncclGroupEnd()); + } + void ReduceScatter( const void* sendbuff, void* recvbuff, size_t recvcount, DataType type, int group, cudaStream_t stream) override { @@ -271,6 +300,35 @@ public: NCCLCHECK(ncclGroupEnd()); } + void ReduceScatterV(const void* sendbuff, + void* recvbuff, + const std::vector& counts, + DataType type, + int group, + cudaStream_t stream) override + { + const size_t elem_size = byte_size(type); + const ncclDataType_t nccl_type = to_nccl_dtype(type); + const int n_ranks = this->n_ranks(group); + const int rank = this->rank(group); + ncclComm_t comm = groups_.at(group); + + TM_CHECK_EQ((int)counts.size(), n_ranks); + + NCCLCHECK(ncclGroupStart()); + size_t offset = 0; + for (int i = 0; i < n_ranks; ++i) { + const size_t count = counts[i]; + if (count) { + const auto* send = static_cast(sendbuff) + elem_size * offset; + auto* recv = i == rank ? recvbuff : const_cast(send); + NCCLCHECK(ncclReduce(send, recv, count, nccl_type, ncclSum, i, comm, stream)); + } + offset += count; + } + NCCLCHECK(ncclGroupEnd()); + } + void AllreduceResidualBiasRMSnorm(void* hidden, void* residual, const void* bias, @@ -392,7 +450,7 @@ public: int group, cudaStream_t stream) override { - NCCLCHECK(ncclBroadcast(recvbuff, recvbuff, count, to_nccl_dtype(type), root, groups_.at(group), stream)); + NCCLCHECK(ncclBroadcast(sendbuff, recvbuff, count, to_nccl_dtype(type), root, groups_.at(group), stream)); } private: diff --git a/src/turbomind/models/llama/moe_ffn_layer.cc b/src/turbomind/models/llama/moe_ffn_layer.cc index ffff786e66..3ea5489d3b 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.cc +++ b/src/turbomind/models/llama/moe_ffn_layer.cc @@ -1,12 +1,17 @@ // Copyright (c) OpenMMLab. All rights reserved. +#include +#include + #include #include "src/turbomind/core/context.h" #include "src/turbomind/core/scope.h" #include "src/turbomind/kernels/activation.h" +#include "src/turbomind/kernels/gemm/moe_utils_v2.h" #include "src/turbomind/kernels/norm/rms_norm.h" +#include "src/turbomind/models/llama/LlamaFfnLayer.h" #include "src/turbomind/models/llama/LlamaLinear.h" #include "src/turbomind/models/llama/llama_utils.h" #include "src/turbomind/models/llama/moe_ffn_layer.h" @@ -19,16 +24,65 @@ namespace turbomind { -MoeFfnLayer::MoeFfnLayer(const EngineParam& engine, const Context& ctx): +class MoeFfnLayerImpl { +public: + explicit MoeFfnLayerImpl(const Context& ctx): linear_(*ctx.linear) {} + + virtual ~MoeFfnLayerImpl() = default; + + virtual void Forward(MoeFfnLayer::ForwardParam& p) = 0; + + virtual void Combine(MoeFfnLayer::ForwardParam& p) = 0; + +protected: + Tensor_ Gate(const Tensor& input, const LinearWeight& gate); + + LlamaLinear& linear_; +}; + +class MoeFfnTpImpl final: public MoeFfnLayerImpl { +public: + MoeFfnTpImpl(const EngineParam& engine, const Context& ctx); + + void Forward(MoeFfnLayer::ForwardParam& p) override; + + void Combine(MoeFfnLayer::ForwardParam& p) override; + +private: + void Init(MoeFfnLayer::ForwardParam& p); + + const int tp_size_; + const int max_token_num_; + int& is_warm_up_; + + std::unique_ptr expert_ffn_; + + bool initialized_ = false; + + Buffer_ h_offsets_; + + Buffer_ masks_; + Buffer_ f2n_; + Buffer_ f2E_; + Buffer_ en2f_; + Buffer_ scales_; + Buffer_ accum_; + Buffer_ offsets_; + + Tensor temp_; + Tensor_ shared_scales_; +}; + +MoeFfnTpImpl::MoeFfnTpImpl(const EngineParam& engine, const Context& ctx): + MoeFfnLayerImpl(ctx), tp_size_(engine.mlp_tp_size), max_token_num_(engine.max_forward_token_num * engine.attn_dp_size), is_warm_up_(*ctx.is_warm_up), - linear_(*ctx.linear), expert_ffn_(std::make_unique(ctx)) { } -void MoeFfnLayer::Init(ForwardParam& p) +void MoeFfnTpImpl::Init(MoeFfnLayer::ForwardParam& p) { const int expert_num = p.weights->num_experts(); const int experts_per_token = p.weights->experts_per_token; @@ -48,7 +102,7 @@ void MoeFfnLayer::Init(ForwardParam& p) initialized_ = true; } -Tensor_ MoeFfnLayer::Gate(const Tensor& input, const LinearWeight& gate) +Tensor_ MoeFfnLayerImpl::Gate(const Tensor& input, const LinearWeight& gate) { TM_FUNCTION_SCOPE(); @@ -61,7 +115,7 @@ Tensor_ MoeFfnLayer::Gate(const Tensor& input, const LinearWeight& gate) return logits; } -void MoeFfnLayer::Forward(ForwardParam& p) +void MoeFfnTpImpl::Forward(MoeFfnLayer::ForwardParam& p) { TM_FUNCTION_SCOPE(); if (!initialized_) { @@ -196,7 +250,7 @@ void MoeFfnLayer::Forward(ForwardParam& p) } } -void MoeFfnLayer::Combine(ForwardParam& p) +void MoeFfnTpImpl::Combine(MoeFfnLayer::ForwardParam& p) { TM_FUNCTION_SCOPE(); auto& moe = *p.weights; @@ -218,4 +272,63 @@ void MoeFfnLayer::Combine(ForwardParam& p) shared_scales_ = {}; } +class MoeFfnAgRsImpl final: public MoeFfnLayerImpl { +public: + MoeFfnAgRsImpl(const EngineParam& engine, const Context& ctx): + MoeFfnLayerImpl(ctx), ep_size_(engine.ep_size), ep_rank_(engine.ep_rank), backend_(engine.all2all_backend) + { + TM_CHECK_GT(ep_size_, 1); + TM_CHECK_EQ(engine.mlp_tp_size, 1); + TM_CHECK_EQ(backend_, "allgather_reducescatter"); + } + + void Forward(MoeFfnLayer::ForwardParam&) override + { + ReportUnsupported(); + } + + void Combine(MoeFfnLayer::ForwardParam&) override + { + ReportUnsupported(); + } + +private: + void ReportUnsupported() const + { + TM_LOG_FATAL( + "MoeFfnAgRsImpl is a stub: resolve AllreduceResidualRMSnorm and ag-rs ownership before enabling EP MoE"); + } + + int ep_size_; + int ep_rank_; + std::string backend_; +}; + +MoeFfnLayer::MoeFfnLayer(const EngineParam& engine, const Context& ctx) +{ + if (engine.ep_size <= 1) { + impl_ = std::make_unique(engine, ctx); + return; + } + + if (engine.all2all_backend == "allgather_reducescatter") { + impl_ = std::make_unique(engine, ctx); + return; + } + + TM_LOG_FATAL("Unsupported MoE EP all2all backend: {}", engine.all2all_backend); +} + +MoeFfnLayer::~MoeFfnLayer() = default; + +void MoeFfnLayer::Forward(ForwardParam& p) +{ + impl_->Forward(p); +} + +void MoeFfnLayer::Combine(ForwardParam& p) +{ + impl_->Combine(p); +} + } // namespace turbomind diff --git a/src/turbomind/models/llama/moe_ffn_layer.h b/src/turbomind/models/llama/moe_ffn_layer.h index d50bc4869b..39f68fa952 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.h +++ b/src/turbomind/models/llama/moe_ffn_layer.h @@ -2,18 +2,23 @@ #pragma once -#include "src/turbomind/kernels/gemm/context.h" -#include "src/turbomind/kernels/gemm/moe_utils_v2.h" -#include "src/turbomind/models/llama/LlamaFfnLayer.h" +#include + +#include "src/turbomind/core/core.h" +#include "src/turbomind/models/llama/context.h" #include "src/turbomind/models/llama/llama_params.h" #include "src/turbomind/models/moe_weight.h" namespace turbomind { +class MoeFfnLayerImpl; + class MoeFfnLayer { public: MoeFfnLayer(const EngineParam& engine, const Context& ctx); + ~MoeFfnLayer(); + struct ForwardParam { Tensor input; Tensor output; @@ -27,37 +32,7 @@ class MoeFfnLayer { void Combine(ForwardParam& p); private: - void Init(ForwardParam& p); - - Tensor_ Gate(const Tensor& input, const LinearWeight& gate); - - void dump_logits(int token_num, int layer_id, int expert_num); - - const int tp_size_; - const int max_token_num_; - int& is_warm_up_; - - LlamaLinear& linear_; - - std::unique_ptr expert_ffn_; - - bool initialized_ = false; - - /////////////////////////////////////////////////////// - /// runtime states - Buffer_ h_offsets_; - - Buffer_ masks_; - Buffer_ f2n_; - Buffer_ f2E_; - Buffer_ en2f_; - Buffer_ scales_; - Buffer_ accum_; - Buffer_ offsets_; - - Tensor temp_; - Tensor_ shared_scales_; - /////////////////////////////////////////////////////// + std::unique_ptr impl_; }; } // namespace turbomind From d2bd263acbca4cd535b01495e0d93a8ff8904774 Mon Sep 17 00:00:00 2001 From: irexyc Date: Thu, 2 Jul 2026 14:09:23 +0000 Subject: [PATCH 03/19] add ffn token partition --- src/turbomind/models/llama/moe_ffn_layer.cc | 6 +- src/turbomind/models/llama/unified_decoder.cc | 180 +++++++++++++++--- src/turbomind/models/llama/unified_decoder.h | 15 ++ 3 files changed, 172 insertions(+), 29 deletions(-) diff --git a/src/turbomind/models/llama/moe_ffn_layer.cc b/src/turbomind/models/llama/moe_ffn_layer.cc index 3ea5489d3b..101175e861 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.cc +++ b/src/turbomind/models/llama/moe_ffn_layer.cc @@ -280,6 +280,8 @@ class MoeFfnAgRsImpl final: public MoeFfnLayerImpl { TM_CHECK_GT(ep_size_, 1); TM_CHECK_EQ(engine.mlp_tp_size, 1); TM_CHECK_EQ(backend_, "allgather_reducescatter"); + // Decoder owns attention TP/CP-group FFN boundary communication. This backend consumes and produces + // this rank's compact token slice and owns expert-parallel allgather/reducescatter communication. } void Forward(MoeFfnLayer::ForwardParam&) override @@ -295,8 +297,8 @@ class MoeFfnAgRsImpl final: public MoeFfnLayerImpl { private: void ReportUnsupported() const { - TM_LOG_FATAL( - "MoeFfnAgRsImpl is a stub: resolve AllreduceResidualRMSnorm and ag-rs ownership before enabling EP MoE"); + TM_LOG_FATAL("MoeFfnAgRsImpl is a stub: decoder supplies a compact attention-TP/CP token slice; implement EP " + "allgather-reducescatter dispatch/combine before enabling EP MoE"); } int ep_size_; diff --git a/src/turbomind/models/llama/unified_decoder.cc b/src/turbomind/models/llama/unified_decoder.cc index 2562156634..efa3a88848 100644 --- a/src/turbomind/models/llama/unified_decoder.cc +++ b/src/turbomind/models/llama/unified_decoder.cc @@ -1,7 +1,8 @@ - +#include #include #include +#include #include @@ -25,6 +26,30 @@ namespace turbomind { +struct FfnTokenPartition { + int first_token{}; + int token_count{}; + std::vector elem_counts; + + static FfnTokenPartition + Make(const std::vector& local_token_nums, int dp_rank, int group_rank, int group_size, size_t hidden_units) + { + const int local_token_num = local_token_nums[dp_rank]; + const int base = local_token_num / group_size; + const int remainder = local_token_num % group_size; + + const int first_token = group_rank * base + std::min(group_rank, remainder); + const int token_count = base + (group_rank < remainder ? 1 : 0); + std::vector elem_counts(group_size); + for (int rank = 0; rank < group_size; ++rank) { + const int count = base + (rank < remainder ? 1 : 0); + elem_counts[rank] = static_cast(count) * hidden_units; + } + + return FfnTokenPartition{first_token, token_count, std::move(elem_counts)}; + } +}; + void UnifiedDecoder::Run(BatchOp op, int phase, TensorMap& env) { attn_layer_->Run(op, phase, env); @@ -44,6 +69,7 @@ UnifiedDecoder::UnifiedDecoder(const EngineParam& engine, attn_dp_rank_(engine.attn_dp_rank), mlp_tp_size_(engine.mlp_tp_size), attn_tp_group_(ctx.comm.d_tp_group), + ep_size_(engine.ep_size), d_comm_(ctx.comm.d_comm), tune_layer_num_(engine.tune_layer_num), is_warm_up_{*ctx.is_warm_up} @@ -156,6 +182,77 @@ void UnifiedDecoder::AllreduceResidualRMSnorm(Tensor& hidden_states, } } +void UnifiedDecoder::ReduceScatterVResidualRMSnorm(Tensor& local_hidden_states, + Tensor& local_residual, + const Tensor& bias, + const Tensor& weight, + float eps, + const FfnTokenPartition& partition) +{ + TM_CHECK(ep_size_ > 1); + TM_CHECK(d_comm_); + TM_CHECK_EQ(local_residual.shape(0), local_hidden_states.shape(0)); + TM_CHECK_EQ(partition.elem_counts.size(), static_cast(d_comm_->n_ranks(attn_tp_group_))); + TM_CHECK_EQ(std::accumulate(partition.elem_counts.begin(), partition.elem_counts.end(), size_t{}), + static_cast(local_hidden_states.size())); + + const auto dtype = local_hidden_states.dtype(); + const auto stream = core::Context::stream().handle(); + + Tensor hidden_slice = local_hidden_states.slice({partition.first_token, 0}, {partition.token_count, -1}); + Tensor residual_slice = local_residual.slice({partition.first_token, 0}, {partition.token_count, -1}); + + d_comm_->ReduceScatterV( + local_hidden_states.raw_data(), hidden_slice.raw_data(), partition.elem_counts, dtype, attn_tp_group_, stream); + TM_CUDA_CHECK(cudaGetLastError()); + + invokeResidualBiasRMSNorm(hidden_slice.raw_data(), + residual_slice.raw_data(), + weight.raw_data(), + bias.data_or((void*)nullptr), + dtype, + hidden_units_, + partition.token_count, + eps, + stream); + TM_CUDA_CHECK(cudaGetLastError()); +} + +void UnifiedDecoder::ResidualRMSnormAllGatherV(Tensor& local_hidden_states, + Tensor& local_residual, + const Tensor& weight, + float eps, + const FfnTokenPartition& partition) +{ + TM_CHECK(ep_size_ > 1); + TM_CHECK(d_comm_); + TM_CHECK_EQ(local_residual.shape(0), local_hidden_states.shape(0)); + TM_CHECK_EQ(partition.elem_counts.size(), static_cast(d_comm_->n_ranks(attn_tp_group_))); + TM_CHECK_EQ(std::accumulate(partition.elem_counts.begin(), partition.elem_counts.end(), size_t{}), + static_cast(local_hidden_states.size())); + + const auto dtype = local_hidden_states.dtype(); + const auto stream = core::Context::stream().handle(); + + Tensor hidden_slice = local_hidden_states.slice({partition.first_token, 0}, {partition.token_count, -1}); + Tensor residual_slice = local_residual.slice({partition.first_token, 0}, {partition.token_count, -1}); + + invokeResidualBiasRMSNorm(hidden_slice.raw_data(), + residual_slice.raw_data(), + weight.raw_data(), + nullptr, + dtype, + hidden_units_, + partition.token_count, + eps, + stream); + TM_CUDA_CHECK(cudaGetLastError()); + + d_comm_->AllGatherV( + hidden_slice.raw_data(), local_hidden_states.raw_data(), partition.elem_counts, dtype, attn_tp_group_, stream); + TM_CUDA_CHECK(cudaGetLastError()); +} + void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vector& weights) { TM_FUNCTION_SCOPE(); @@ -212,6 +309,14 @@ void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vector ffn_partition; + if (ep_size_ > 1) { + const int group_size = d_comm_->n_ranks(attn_tp_group_); + const int group_rank = d_comm_->rank(attn_tp_group_); + ffn_partition.emplace( + FfnTokenPartition::Make(local_token_nums, attn_dp_rank_, group_rank, group_size, hidden_units_)); + } + TM_LOG_DEBUG("local_token_num=%d, global_token_num=%d", (int)local_token_num, (int)global_token_num); TM_DEBUG_TENSOR(local_residual, "res", 1); @@ -269,15 +374,30 @@ void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vectorattention->wo->bias; } - AllreduceResidualRMSnorm(global_hidden_states, - local_residual, - out_bias, - weights.at(layer)->ffn_norm->weight, - weights.at(layer)->ffn_norm->norm_eps_, - local_token_num, - attn_tp_group_, - 0, - local_token_nums.data()); + Tensor ffn_hidden_states; + + if (ep_size_ > 1) { + ReduceScatterVResidualRMSnorm(local_hidden_states, + local_residual, + out_bias, + weights.at(layer)->ffn_norm->weight, + weights.at(layer)->ffn_norm->norm_eps_, + *ffn_partition); + ffn_hidden_states = + local_hidden_states.slice({ffn_partition->first_token, 0}, {ffn_partition->token_count, -1}); + } + else { + AllreduceResidualRMSnorm(global_hidden_states, + local_residual, + out_bias, + weights.at(layer)->ffn_norm->weight, + weights.at(layer)->ffn_norm->norm_eps_, + local_token_num, + attn_tp_group_, + 0, + local_token_nums.data()); + ffn_hidden_states = global_hidden_states; + } TM_DEBUG_TENSOR(local_residual, Concat("residual0", layer), 2); TM_DEBUG_TENSOR(local_hidden_states, Concat("norm1", layer), 2); @@ -288,38 +408,44 @@ void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vector moe_fwd_param; if (weights.at(layer)->moe_ffn) { - moe_fwd_param = MoeFfnLayer::ForwardParam{global_hidden_states, - global_hidden_states, - weights.at(layer)->moe_ffn.get(), - ffn_layer_ ? 1.f : 0.f, - layer}; + moe_fwd_param = MoeFfnLayer::ForwardParam{ + ffn_hidden_states, ffn_hidden_states, weights.at(layer)->moe_ffn.get(), ffn_layer_ ? 1.f : 0.f, layer}; moe_ffn_layer_->Forward(*moe_fwd_param); } - if (ffn_layer_ && weights.at(layer)->feed_forward) { + if (ffn_layer_ && weights.at(layer)->feed_forward && ffn_hidden_states.shape(0) > 0) { ffn_layer_->forward( - {global_hidden_states, global_hidden_states, weights.at(layer)->feed_forward.get(), (int)layer}); + {ffn_hidden_states, ffn_hidden_states, weights.at(layer)->feed_forward.get(), (int)layer}); } if (moe_fwd_param) { moe_ffn_layer_->Combine(*moe_fwd_param); } - TM_DEBUG_TENSOR(global_hidden_states, Concat("ffn_block", layer), 2); + TM_DEBUG_TENSOR(ffn_hidden_states, Concat("ffn_block", layer), 2); const bool last = layer == layer_num_ - 1; auto& scale_weight = !last ? weights.at(layer + 1)->attention_norm->weight : args.at("output_norm_weight"); - AllreduceResidualRMSnorm(global_hidden_states, - local_residual, - {}, - scale_weight, - weights.at(layer)->ffn_norm->norm_eps_, - local_token_num, - 0, - attn_tp_group_, - local_token_nums.data()); + if (ep_size_ > 1) { + ResidualRMSnormAllGatherV(local_hidden_states, + local_residual, + scale_weight, + weights.at(layer)->ffn_norm->norm_eps_, + *ffn_partition); + } + else { + AllreduceResidualRMSnorm(global_hidden_states, + local_residual, + {}, + scale_weight, + weights.at(layer)->ffn_norm->norm_eps_, + local_token_num, + 0, + attn_tp_group_, + local_token_nums.data()); + } TM_CUDA_CHECK(cudaGetLastError()); TM_DEBUG_TENSOR(local_residual, Concat("residual1", layer), 2); diff --git a/src/turbomind/models/llama/unified_decoder.h b/src/turbomind/models/llama/unified_decoder.h index 7c1f36af65..541866c99a 100644 --- a/src/turbomind/models/llama/unified_decoder.h +++ b/src/turbomind/models/llama/unified_decoder.h @@ -12,6 +12,7 @@ namespace turbomind { class ModelWeight; class DecoderLayerWeight; +struct FfnTokenPartition; class UnifiedDecoder { public: @@ -31,6 +32,7 @@ class UnifiedDecoder { const int attn_dp_size_; const int attn_dp_rank_; const int mlp_tp_size_; + const int ep_size_; const int attn_tp_group_; @@ -54,6 +56,19 @@ class UnifiedDecoder { int t0, int t1, const int* local_token_nums); + + void ReduceScatterVResidualRMSnorm(Tensor& local_hidden_states, + Tensor& local_residual, + const Tensor& bias, + const Tensor& weight, + float eps, + const FfnTokenPartition& partition); + + void ResidualRMSnormAllGatherV(Tensor& local_hidden_states, + Tensor& local_residual, + const Tensor& weight, + float eps, + const FfnTokenPartition& partition); }; } // namespace turbomind From 3695542d41265412b052935f13c78dca6369f984 Mon Sep 17 00:00:00 2001 From: irexyc Date: Sun, 5 Jul 2026 18:22:33 +0000 Subject: [PATCH 04/19] implement MoeFfnAgRsImpl --- src/turbomind/comm/device_comm.h | 3 + src/turbomind/comm/nccl/nccl.cu | 12 + src/turbomind/kernels/activation.cu | 127 +++- src/turbomind/kernels/activation.h | 12 + src/turbomind/kernels/gemm/CMakeLists.txt | 1 + src/turbomind/kernels/gemm/cublas.cu | 4 +- src/turbomind/kernels/gemm/moe_ep_utils.cu | 552 ++++++++++++++++++ src/turbomind/kernels/gemm/moe_ep_utils.h | 47 ++ src/turbomind/kernels/gemm/moe_utils_v2.cu | 72 ++- src/turbomind/kernels/gemm/moe_utils_v2.h | 21 +- src/turbomind/kernels/gemm/test/testbed_v3.h | 2 +- src/turbomind/models/llama/LlamaLinear.cu | 28 +- src/turbomind/models/llama/LlamaLinear.h | 3 +- src/turbomind/models/llama/moe_ffn_layer.cc | 238 +++++++- src/turbomind/models/llama/moe_ffn_layer.h | 6 + src/turbomind/models/llama/unified_decoder.cc | 77 ++- 16 files changed, 1103 insertions(+), 102 deletions(-) create mode 100644 src/turbomind/kernels/gemm/moe_ep_utils.cu create mode 100644 src/turbomind/kernels/gemm/moe_ep_utils.h diff --git a/src/turbomind/comm/device_comm.h b/src/turbomind/comm/device_comm.h index 9ce0be0d12..ba4d14c0da 100644 --- a/src/turbomind/comm/device_comm.h +++ b/src/turbomind/comm/device_comm.h @@ -138,6 +138,9 @@ class DeviceCommImpl { { throw std::runtime_error("not implemented"); } + + virtual void GroupStart() {} + virtual void GroupEnd() {} }; class DeviceComm { diff --git a/src/turbomind/comm/nccl/nccl.cu b/src/turbomind/comm/nccl/nccl.cu index f0a7363854..504b4a943d 100644 --- a/src/turbomind/comm/nccl/nccl.cu +++ b/src/turbomind/comm/nccl/nccl.cu @@ -44,6 +44,8 @@ static inline ncclDataType_t to_nccl_dtype(DataType type) return ncclBfloat16; case kUint8: return ncclUint8; + case kInt64: + return ncclInt64; default: throw std::runtime_error("not supported"); } @@ -453,6 +455,16 @@ public: NCCLCHECK(ncclBroadcast(sendbuff, recvbuff, count, to_nccl_dtype(type), root, groups_.at(group), stream)); } + void GroupStart() override + { + NCCLCHECK(ncclGroupStart()); + } + + void GroupEnd() override + { + NCCLCHECK(ncclGroupEnd()); + } + private: HostComm h_comm_; diff --git a/src/turbomind/kernels/activation.cu b/src/turbomind/kernels/activation.cu index cab79c6ee7..1406c9c78b 100644 --- a/src/turbomind/kernels/activation.cu +++ b/src/turbomind/kernels/activation.cu @@ -1,4 +1,6 @@ +#include + #include "src/turbomind/core/data_type.h" #include "src/turbomind/kernels/activation.h" @@ -26,9 +28,14 @@ struct Silu { } }; -template -__global__ void ActivationKernel( - T* gate_buf, const T* __restrict__ up_buf, Activation activation, int64_t stride, int token_num, int dim) +template +__global__ void ActivationKernel(T* gate_buf, + const T* __restrict__ up_buf, + Activation activation, + int64_t stride, + int token_num, + int dim, + const int* num_valid_tokens) { if constexpr (TURBOMIND_ARCH_DTYPE_GUARD(data_type_v)) { const int di = threadIdx.x + blockIdx.y * blockDim.x; @@ -36,6 +43,12 @@ __global__ void ActivationKernel( dim /= vec_size; + if constexpr (has_valid_tokens) { + if (ti >= __ldg(num_valid_tokens)) { + return; + } + } + if (di >= dim) { return; } @@ -61,6 +74,12 @@ __global__ void ActivationKernel( } void Activation(Ref gate_, const Tensor& up, ActivationType type, cudaStream_t stream) +{ + Activation(gate_, up, type, nullptr, stream); +} + +void Activation( + Ref gate_, const Tensor& up, ActivationType type, const int* num_valid_tokens, cudaStream_t stream) { auto& gate = gate_.get(); @@ -69,52 +88,75 @@ void Activation(Ref gate_, const Tensor& up, ActivationType type, cudaSt int num, dim; std::tie(num, dim) = gate.shapes(0, 1); - auto invoke = [&](auto t, auto act) { - using T = decltype(t); + auto invoke = [&](auto t, auto act, auto has_valid_tokens) { + using T = decltype(t); + constexpr bool kHasValidTokens = decltype(has_valid_tokens)::value; constexpr int vec_size = 4; constexpr int threads = 512; const dim3 blocks(num, cdiv(dim, threads * vec_size)); - - ActivationKernel<<>>(gate.data(), // - up.data(), - act, - gate.stride(0), - num, - dim); + const int* valid_tokens = kHasValidTokens ? num_valid_tokens : nullptr; + + ActivationKernel<<>>(gate.data(), // + up.data(), + act, + gate.stride(0), + num, + dim, + valid_tokens); }; - auto dispatch = [&](auto t) { + auto dispatch_act = [&](auto t, auto has_valid_tokens) { using T = decltype(t); if (type == ActivationType::kSilu) { - return invoke(t, Silu{}); + return invoke(t, Silu{}, has_valid_tokens); } else if (type == ActivationType::kSiluGptOss) { - return invoke(t, SiluGptOss{}); + return invoke(t, SiluGptOss{}, has_valid_tokens); } else { TM_LOG_FATAL("unknown activation type: {}", (int)type); } }; + auto dispatch = [&](auto t) { + if (num_valid_tokens) { + return dispatch_act(t, std::true_type{}); + } + return dispatch_act(t, std::false_type{}); + }; + TM_DISPATCH_PRIMARY_DTYPES(gate.dtype(), dispatch); TM_CUDA_CHECK(cudaGetLastError()); } -template -__global__ void ActivationKernel( - T* gate_up, const T* bias, const int* group_ids, int64_t stride, Activation activation, int token_num, int dim) +template +__global__ void ActivationKernel(T* gate_up, + const T* bias, + const int* group_ids, + int64_t stride, + Activation activation, + int token_num, + int dim, + const int* num_valid_tokens) { if constexpr (TURBOMIND_ARCH_DTYPE_GUARD(data_type_v)) { const int di = (threadIdx.x + blockIdx.y * blockDim.x) * vec_size; const int ti = blockIdx.x; - const int gi = group_ids ? group_ids[ti] : 0; + + if constexpr (has_valid_tokens) { + if (ti >= __ldg(num_valid_tokens)) { + return; + } + } if (di >= dim) { return; } + const int gi = group_ids ? group_ids[ti] : 0; + using Vec = Array; Vec gate_bias{}, up_bias{}; @@ -145,6 +187,16 @@ void Activation(Tensor& gate_up, // const Buffer_& group_ids, ActivationType type, cudaStream_t stream) +{ + Activation(gate_up, bias, group_ids, type, nullptr, stream); +} + +void Activation(Tensor& gate_up, // + const Tensor& bias, + const Buffer_& group_ids, + ActivationType type, + const int* num_valid_tokens, + cudaStream_t stream) { const int num = gate_up.shape(0); const int dim = gate_up.shape(1) / 2; @@ -153,42 +205,53 @@ void Activation(Tensor& gate_up, // Activation(gate_up.slice({0, 0}, {-1, dim}), // gate_up.slice({0, dim}, {-1, -1}), type, + num_valid_tokens, stream); return; } TM_CHECK_EQ(gate_up.shape(-1), bias.shape(-1)); - auto invoke = [&](auto t, auto act) { - using T = decltype(t); + auto invoke = [&](auto t, auto act, auto has_valid_tokens) { + using T = decltype(t); + constexpr bool kHasValidTokens = decltype(has_valid_tokens)::value; constexpr int vec_size = 4; constexpr int threads = 512; const dim3 blocks(num, cdiv(dim, threads * vec_size)); - - ActivationKernel<<>>(gate_up.data(), // - bias.data_or((T*)nullptr), - group_ids.data_or(nullptr), - gate_up.stride(0), - act, - num, - dim); + const int* valid_tokens = kHasValidTokens ? num_valid_tokens : nullptr; + + ActivationKernel<<>>(gate_up.data(), // + bias.data_or((T*)nullptr), + group_ids.data_or(nullptr), + gate_up.stride(0), + act, + num, + dim, + valid_tokens); }; - auto dispatch = [&](auto t) { + auto dispatch_act = [&](auto t, auto has_valid_tokens) { using T = decltype(t); if (type == ActivationType::kSilu) { - return invoke(t, Silu{}); + return invoke(t, Silu{}, has_valid_tokens); } else if (type == ActivationType::kSiluGptOss) { - return invoke(t, SiluGptOss{}); + return invoke(t, SiluGptOss{}, has_valid_tokens); } else { TM_LOG_FATAL("unknown activation type: {}", (int)type); } }; + auto dispatch = [&](auto t) { + if (num_valid_tokens) { + return dispatch_act(t, std::true_type{}); + } + return dispatch_act(t, std::false_type{}); + }; + TM_DISPATCH_PRIMARY_DTYPES(gate_up.dtype(), dispatch); TM_CUDA_CHECK(cudaGetLastError()); } diff --git a/src/turbomind/kernels/activation.h b/src/turbomind/kernels/activation.h index 2eceeb9d89..72947b054c 100644 --- a/src/turbomind/kernels/activation.h +++ b/src/turbomind/kernels/activation.h @@ -14,10 +14,22 @@ enum class ActivationType void Activation(Ref gate, const Tensor& up, ActivationType type, cudaStream_t stream); +// num_valid_tokens points to a device-side valid token count for padded token buffers. +void Activation( + Ref gate, const Tensor& up, ActivationType type, const int* num_valid_tokens, cudaStream_t stream); + +void Activation(Tensor& gate_up, // + const Tensor& bias, + const Buffer_& group_ids, + ActivationType type, + cudaStream_t stream); + +// num_valid_tokens points to a device-side valid token count for padded token buffers. void Activation(Tensor& gate_up, // const Tensor& bias, const Buffer_& group_ids, ActivationType type, + const int* num_valid_tokens, cudaStream_t stream); } // namespace turbomind diff --git a/src/turbomind/kernels/gemm/CMakeLists.txt b/src/turbomind/kernels/gemm/CMakeLists.txt index 7a1b210762..31ce211b1c 100644 --- a/src/turbomind/kernels/gemm/CMakeLists.txt +++ b/src/turbomind/kernels/gemm/CMakeLists.txt @@ -60,6 +60,7 @@ add_library(gemm2 ${GEMM2_KERNELS_SM80} cublas.cu moe_utils_v2.cu + moe_ep_utils.cu test/test_utils.cu ) diff --git a/src/turbomind/kernels/gemm/cublas.cu b/src/turbomind/kernels/gemm/cublas.cu index f809500122..461c0d5036 100644 --- a/src/turbomind/kernels/gemm/cublas.cu +++ b/src/turbomind/kernels/gemm/cublas.cu @@ -310,9 +310,9 @@ public: const int N = Bdesc.cols; const int K = Adesc.cols; - if (ptr_offsets[group_count] != Adesc.rows) { + if (ptr_offsets[group_count] > Adesc.rows) { fprintf(stderr, - "[TM][GEMM] CublasGrouped: offsets[%d]=%d != Adesc.rows=%d (would OOB)\n", + "[TM][GEMM] CublasGrouped: offsets[%d]=%d exceeds Adesc.rows=%d (would OOB)\n", group_count, ptr_offsets[group_count], Adesc.rows); diff --git a/src/turbomind/kernels/gemm/moe_ep_utils.cu b/src/turbomind/kernels/gemm/moe_ep_utils.cu new file mode 100644 index 0000000000..330f55d7ee --- /dev/null +++ b/src/turbomind/kernels/gemm/moe_ep_utils.cu @@ -0,0 +1,552 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +#include "src/turbomind/kernels/gemm/moe_ep_utils.h" + +#include "src/turbomind/core/check.h" +#include "src/turbomind/core/data_type.h" +#include "src/turbomind/kernels/core/array_ops.h" +#include "src/turbomind/kernels/core/common.h" +#include "src/turbomind/kernels/core/math.h" +#include "src/turbomind/kernels/gemm/moe_utils_v2.h" +#include "src/turbomind/utils/cuda_utils.h" + +#include +#include +#include + +namespace turbomind { + +template +__global__ void MoeGateKernel(float* topk_weights, // [n, topk] + int64_t* topk_idx, // [n, topk] + const float* logits, // [n,E] + int token_num, + int expert_num, + int top_k, + bool softmax, + bool norm_topk, + float routed_scale) +{ + constexpr int threads_per_token = max_expert_num / items_per_thread; + + static_assert(items_per_thread <= 32); + static_assert(threads_per_token <= 32); + static_assert((threads_per_token & (threads_per_token - 1)) == 0); + + const int thread_idx = threadIdx.x + blockIdx.x * blockDim.x; + + const int ti = thread_idx / threads_per_token; + const int ei = thread_idx % threads_per_token; + + const int warp_ti = threadIdx.x % WARP_SIZE / threads_per_token; + + float data[items_per_thread]; + int idxs[items_per_thread]; + + PRAGMA_UNROLL + for (int i = 0; i < items_per_thread; ++i) { + data[i] = -std::numeric_limits::infinity(); + idxs[i] = ei * items_per_thread + i; + } + if (ti < token_num) { + PRAGMA_UNROLL + for (int i = 0; i < items_per_thread; i += access_size) { + const int e = ei * items_per_thread + i; + if (e + access_size <= expert_num) { + Ldg((Array&)data[i], &logits[ti * expert_num + e]); + } + else { + PRAGMA_UNROLL + for (int j = 0; j < access_size; ++j) { + if (e + j < expert_num) { + data[i + j] = logits[ti * expert_num + e + j]; + } + } + } + } + } + + unsigned mask = (unsigned)-1; + float max_logit; + + const int warp_ti_offset = warp_ti * threads_per_token; + + int sel_item[max_top_k]; + + auto run = [&](int k) { + unsigned bit = 1; + unsigned max_bit = 0; + float max_val = -std::numeric_limits::infinity(); + PRAGMA_UNROLL + for (int i = 0; i < items_per_thread; ++i) { + if ((mask & bit) && data[i] > max_val) { + max_bit = bit; + max_val = data[i]; + } + asm("shl.b32 %0, %1, 1;\n" : "=r"(bit) : "r"(bit)); + } + + int g_max_ei = ei; + float g_max_val = max_val; + if constexpr (threads_per_token > 1) { + PRAGMA_UNROLL + for (int m = threads_per_token / 2; m >= 1; m /= 2) { + g_max_val = fmaxf(g_max_val, __shfl_xor_sync((uint32_t)-1, g_max_val, m)); + } + const auto active = __ballot_sync((uint32_t)-1, max_val == g_max_val); + g_max_ei = __ffs(active >> (unsigned)warp_ti_offset) - 1; + } + if (k == 0) { + max_logit = g_max_val; + } + int local_item = -1; + if (ei == g_max_ei) { + local_item = __ffs(max_bit) - 1; + mask -= max_bit; + } + sel_item[k] = local_item; + }; + + run(0); + + for (int k = 1; k < top_k; ++k) { + run(k); + } + + mask = ~mask; + + int used[items_per_thread]; + { + unsigned bit = 1; + PRAGMA_UNROLL + for (int i = 0; i < items_per_thread; ++i) { + used[i] = (mask & bit) > 0; + asm("shl.b32 %0, %1, 1;\n" : "=r"(bit) : "r"(bit)); + } + } + + float sum_prob{}; + + if (softmax) { + PRAGMA_UNROLL + for (int i = 0; i < items_per_thread; ++i) { + if (!norm_topk || used[i]) { + data[i] = expf(data[i] - max_logit); + sum_prob += data[i]; + } + } + PRAGMA_UNROLL + for (int m = threads_per_token / 2; m >= 1; m /= 2) { + sum_prob += __shfl_xor_sync((uint32_t)-1, sum_prob, m); + } + sum_prob = fdividef(1.f, sum_prob); + } + else { + sum_prob = 1.f; + } + + if (ti < token_num) { + PRAGMA_UNROLL + for (int k = 0; k < max_top_k; ++k) { + if (k < top_k && sel_item[k] >= 0) { + const int i = sel_item[k]; + topk_weights[ti * top_k + k] = data[i] * sum_prob * routed_scale; + topk_idx[ti * top_k + k] = idxs[i]; + } + } + } +} + +template +inline constexpr std::integral_constant _Int{}; + +void invokeMoeGateEp(float* topk_weights, + int64_t* topk_idx, + const float* logits, + int tokens, + int experts, + int experts_per_token, + bool softmax, + bool norm_topk, + float routed_scale, + cudaStream_t stream) +{ + if (tokens == 0) { + return; + } + + auto invoke = [&](auto max_expert_num, auto top_k, auto items_per_thread, auto vec_size) { + constexpr int thrs_per_tok = max_expert_num.value / items_per_thread.value; + constexpr int threads = 256; + const int blocks = ceil_div(tokens, threads / thrs_per_tok); + + MoeGateKernel + <<>>( + topk_weights, topk_idx, logits, tokens, experts, experts_per_token, softmax, norm_topk, routed_scale); + TM_CUDA_CHECK(cudaGetLastError()); + + return true; + }; + + if (!softmax && norm_topk) { + TM_CHECK(0) << softmax << " " << norm_topk; + } + + auto dispatch = [&] { + if (experts <= 8) { + if (experts_per_token <= 2) { + return invoke(_Int<8>, _Int<2>, _Int<8>, _Int<4>); + } + else { + return invoke(_Int<8>, _Int<8>, _Int<8>, _Int<4>); + } + } + else if (experts <= 64) { + if (experts_per_token <= 4) { + return invoke(_Int<64>, _Int<4>, _Int<16>, _Int<4>); + } + else if (experts_per_token <= 8) { + return invoke(_Int<64>, _Int<8>, _Int<16>, _Int<4>); + } + } + else if (experts <= 128) { + if (experts_per_token <= 8) { + return invoke(_Int<128>, _Int<8>, _Int<16>, _Int<4>); + } + } + else if (experts <= 160) { + if (experts_per_token <= 8) { + return invoke(_Int<160>, _Int<8>, _Int<10>, _Int<2>); + } + } + else if (experts <= 512) { + if (experts_per_token <= 10) { + return invoke(_Int<512>, _Int<10>, _Int<16>, _Int<4>); + } + } + return false; + }; + + auto success = dispatch(); + + TM_CHECK(success) << "unsupported moe config: expert_num=" << experts << ", top_k=" << experts_per_token + << ", softmax=" << softmax << ", norm_topk=" << norm_topk; +} + +template +__global__ void MoeAgRsBuildMasksKernel(int8_t* masks, // [num_local_experts, tokens_padded], pre-set to -1 + int* accum, // [num_local_experts, tiles], pre-set to 0 + const int64_t* topk_idx, // [tokens, topk] + int tokens, + int tokens_padded, + int topk, + int local_expert_offset, + int num_local_experts, + int log_tile, + int tiles) +{ + __shared__ int shared_accum[kMoeGateMaxTiles][max_local_experts + 1]; + + // Clear shared_accum + for (int i = threadIdx.x; i < tiles * num_local_experts; i += block_dim) { + shared_accum[i / num_local_experts][i % num_local_experts] = 0; + } + __syncthreads(); + + // Each thread processes one token + const int ti = threadIdx.x + blockIdx.x * blockDim.x; + if (ti < tokens) { + for (int k = 0; k < topk; ++k) { + const int global_eid = static_cast(topk_idx[ti * topk + k]); + const int local_eid = global_eid - local_expert_offset; + if (0 <= local_eid && local_eid < num_local_experts) { + masks[local_eid * tokens_padded + ti] = static_cast(k); + atomicAdd(&shared_accum[ti >> log_tile][local_eid], 1); + } + } + } + __syncthreads(); + + // Flush shared_accum to global accum + for (int i = threadIdx.x; i < tiles * num_local_experts; i += block_dim) { + int t = i / num_local_experts; + int e = i % num_local_experts; + if (t < tiles && shared_accum[t][e]) { + atomicAdd(&accum[e * tiles + t], shared_accum[t][e]); + } + } +} + +void invokeMoeBuildAgRsRoutingMap(int* f2n, + int* f2E, + int* en2f, + int* offsets, + int8_t* masks, + int* accum, + const int64_t* topk_idx, + int tokens, + int topk, + int local_expert_offset, + int num_local_experts, + cudaStream_t stream) +{ + if (tokens == 0) { + TM_CUDA_CHECK(cudaMemsetAsync(offsets, 0, sizeof(int) * (num_local_experts + 1), stream)); + return; + } + + const int tokens_padded = round_up(tokens, kMoeGateVecSize); + + // Compute log_tile / tiles (must match invokeMoeScanKernel's internal computation) + constexpr int base_log_tile = 9; + int log_tile = base_log_tile; + while (((tokens_padded + (1 << log_tile) - 1) >> log_tile) > kMoeGateMaxTiles) { + ++log_tile; + } + const int tiles = ceil_div(tokens_padded, 1 << log_tile); + + // Initialise buffers + TM_CUDA_CHECK(cudaMemsetAsync(masks, -1, sizeof(int8_t) * num_local_experts * tokens_padded, stream)); + TM_CUDA_CHECK(cudaMemsetAsync(accum, 0, sizeof(int) * num_local_experts * kMoeGateMaxTiles, stream)); + TM_CUDA_CHECK(cudaMemsetAsync(en2f, -1, sizeof(int) * tokens * topk, stream)); + + constexpr int block = 256; + const int grid = cdiv(tokens_padded, block); + + auto invoke = [&](auto max_local_experts) { + MoeAgRsBuildMasksKernel<<>>(masks, + accum, + topk_idx, + tokens, + tokens_padded, + topk, + local_expert_offset, + num_local_experts, + log_tile, + tiles); + return true; + }; + + auto dispatch = [&] { + if (num_local_experts <= 4) { + return invoke(_Int<4>); + } + else if (num_local_experts <= 8) { + return invoke(_Int<8>); + } + else if (num_local_experts <= 16) { + return invoke(_Int<16>); + } + else if (num_local_experts <= 32) { + return invoke(_Int<32>); + } + else if (num_local_experts <= 64) { + return invoke(_Int<64>); + } + else if (num_local_experts <= 128) { + return invoke(_Int<128>); + } + return false; + }; + + auto success = dispatch(); + TM_CHECK(success) << "unsupported num_local_experts: " << num_local_experts; + TM_CUDA_CHECK(cudaGetLastError()); + + invokeMoeScanKernel(f2n, f2E, en2f, offsets, masks, accum, tokens, tokens_padded, num_local_experts, stream); +} + +template +__global__ void MoeCombineKernel(T* dst, // [num_tokens, dim] + const T* src, // [expert_token_num, dim] + const T* bias, // [num_local_experts, dim] + const float* topk_weights, // [num_tokens, topk] + const int* en2f, // [topk, num_tokens] + const int* f2E, // [expert_token_num] + int dim, + int tokens) +{ + if constexpr (TURBOMIND_ARCH_DTYPE_GUARD(data_type_v)) { + for (int ti = blockIdx.x; ti < tokens; ti += gridDim.x) { + T* dst_row = dst + (int64_t)dim * ti; + + const T* src_[exp_k]{}; + const T* bias_[exp_k]{}; + float weight[exp_k]{}; + + PRAGMA_UNROLL + for (int e = 0; e < exp_k; ++e) { + const int fid = __ldg(&en2f[e * tokens + ti]); + if (fid >= 0) { + src_[e] = src + (int64_t)dim * fid; + weight[e] = __ldg(&topk_weights[ti * exp_k + e]); + if constexpr (has_bias) { + bias_[e] = bias + (int64_t)dim * __ldg(&f2E[fid]); + } + } + } + + using Vec = Array; + + for (int i = threadIdx.x * vec_size; i < dim; i += block_dim * vec_size) { + Array accum{}; + PRAGMA_UNROLL + for (int e = 0; e < exp_k; ++e) { + if (src_[e] == nullptr) { + continue; + } + Vec v; + Load(v, src_[e] + i); + if constexpr (has_bias) { + Vec b; + Load(b, bias_[e] + i); + PRAGMA_UNROLL + for (int j = 0; j < vec_size; ++j) { + v[j] = (T)((float)v[j] + (float)b[j]); + } + } + using namespace ops; + const auto x = cast(v) * weight[e]; + accum = accum + x; + } + Store(&dst_row[i], cast(accum)); + } + } + } +} + +void invokeMoeLocalCombineEp(Ref out_, + const Tensor& src, + const Tensor& bias, + const float* topk_weights, + const int* en2f, + const int* f2E, + int experts_per_token, + cudaStream_t st) +{ + auto& out = out_.get(); + + const int tokens = out.shape(0); + + if (tokens == 0) { + return; + } + + const int dim = src.shape(1); + + auto invoke = [&](auto t, auto e, auto has_bias_) { + using T = decltype(t); + constexpr int threads = 256; + constexpr int vec_size = sizeof(uint4) / sizeof(T); + constexpr int exp_per_tok = decltype(e)::value; + constexpr bool has_bias = decltype(has_bias_)::value; + + static const int sm_count = getSMCount(); + const int grid = std::min(tokens, sm_count * 4); + + MoeCombineKernel<<>>( + out.data(), src.data(), bias.data_or((T*)nullptr), topk_weights, en2f, f2E, dim, tokens); + TM_CUDA_CHECK(cudaGetLastError()); + }; + + auto dispatch_topk = [&](auto has_bias, auto t) { + switch (experts_per_token) { + case 1: + return invoke(t, std::integral_constant{}, has_bias); + case 2: + return invoke(t, std::integral_constant{}, has_bias); + case 4: + return invoke(t, std::integral_constant{}, has_bias); + case 6: + return invoke(t, std::integral_constant{}, has_bias); + case 8: + return invoke(t, std::integral_constant{}, has_bias); + case 10: + return invoke(t, std::integral_constant{}, has_bias); + default: + TM_CHECK(0) << "unsupported experts_per_token " << experts_per_token; + } + }; + + auto dispatch_dtype = [&](auto t) { + if (bias) { + TM_CHECK_NOTNULL(f2E); + return dispatch_topk(std::true_type{}, t); + } + else { + return dispatch_topk(std::false_type{}, t); + } + }; + + TM_DISPATCH_PRIMARY_DTYPES(src.dtype(), dispatch_dtype); +} + +template +__global__ void MoeCombineOutputEpKernel(T* dst, // [tokens, dim] + const T* shared, // [tokens, dim] + const float* shared_scales, // [tokens] or nullptr + int dim, + float scale) +{ + if constexpr (TURBOMIND_ARCH_DTYPE_GUARD(data_type_v)) { + const int ti = blockIdx.x; + + float dst_scale = scale; + if (shared_scales) { + dst_scale = __ldg(&shared_scales[ti]); + dst_scale = fdividef(1.f, 1.f + expf(-dst_scale)); + } + + dst += (int64_t)dim * ti; + shared += (int64_t)dim * ti; + + using Vec = Array; + + for (int i = threadIdx.x * vec_size; i < dim; i += block_dim * vec_size) { + Vec routed; + Load(routed, &dst[i]); + using namespace ops; + Array accum = cast(routed); + { + Vec v; + Load(v, &shared[i]); + accum = accum + cast(v) * dst_scale; + } + Store(&dst[i], cast(accum)); + } + } +} + +void invokeMoeCombineOutputEp( + Ref output, const Tensor& shared, const float* shared_scales, float scale, cudaStream_t st) +{ + auto& out = output.get(); + + TM_CHECK(shared); + TM_CHECK_EQ(shared.shape(0), out.shape(0)); + TM_CHECK_EQ(shared.shape(1), out.shape(1)); + + const int tokens = out.shape(0); + const int dim = out.shape(1); + + if (tokens == 0) { + return; + } + + if (shared_scales == nullptr && scale == 0) { + return; + } + + auto dispatch = [&](auto t) { + using T = decltype(t); + constexpr int threads = 256; + constexpr int vec_size = sizeof(uint4) / sizeof(T); + MoeCombineOutputEpKernel + <<>>(out.data(), shared.data(), shared_scales, dim, scale); + TM_CUDA_CHECK(cudaGetLastError()); + }; + + TM_DISPATCH_PRIMARY_DTYPES(out.dtype(), dispatch); +} + +} // namespace turbomind diff --git a/src/turbomind/kernels/gemm/moe_ep_utils.h b/src/turbomind/kernels/gemm/moe_ep_utils.h new file mode 100644 index 0000000000..221583bec8 --- /dev/null +++ b/src/turbomind/kernels/gemm/moe_ep_utils.h @@ -0,0 +1,47 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +#pragma once + +#include + +#include "src/turbomind/core/core.h" + +namespace turbomind { + +void invokeMoeGateEp(float* topk_weights, + int64_t* topk_idx, + const float* logits, + int tokens, + int experts, + int experts_per_token, + bool softmax, + bool norm_topk, + float routed_scale, + cudaStream_t stream); + +void invokeMoeBuildAgRsRoutingMap(int* f2n, + int* f2E, + int* en2f, + int* offsets, + int8_t* masks, + int* accum, + const int64_t* topk_idx, + int tokens, + int topk, + int local_expert_offset, + int num_local_experts, + cudaStream_t stream); + +void invokeMoeLocalCombineEp(Ref out, + const Tensor& src, + const Tensor& bias, + const float* topk_weights, + const int* en2f, + const int* f2E, + int experts_per_token, + cudaStream_t st); + +void invokeMoeCombineOutputEp( + Ref output, const Tensor& shared, const float* shared_scales, float scale, cudaStream_t st); + +} // namespace turbomind diff --git a/src/turbomind/kernels/gemm/moe_utils_v2.cu b/src/turbomind/kernels/gemm/moe_utils_v2.cu index 2a6187cc5f..070a89a52e 100644 --- a/src/turbomind/kernels/gemm/moe_utils_v2.cu +++ b/src/turbomind/kernels/gemm/moe_utils_v2.cu @@ -571,6 +571,32 @@ __global__ void MoeGateKernel_v8(float* scales, // [e,n] template inline constexpr std::integral_constant _Int{}; +void invokeMoeScanKernel(int* f2n, + int* f2E, + int* en2f, + int* offsets, + int8_t* masks, + const int* accum, + int tokens, + int tokens_padded, + int experts, + cudaStream_t st) +{ + constexpr int base_log_tile = 9; + int log_tile = base_log_tile; + while (((tokens_padded + (1 << log_tile) - 1) >> log_tile) > kMoeGateMaxTiles) { + ++log_tile; + } + const int tiles = ceil_div(tokens_padded, 1 << log_tile); + + constexpr int threads = (1 << base_log_tile) / kMoeGateVecSize; + const dim3 blocks(tiles, experts + 1); + + MoeScanKernel_v2<<>>( + f2n, f2E, en2f, offsets, masks, accum, log_tile, tiles, tokens, tokens_padded, experts); + TM_CUDA_CHECK(cudaGetLastError()); +} + void invokeMoeGate_V2(int* f2n, // [e*n] -> n int* f2E, // [e*n] -> E int* en2f, // [e,n] -> n*e @@ -885,8 +911,13 @@ template __global__ void MoeGatherKernel(T* dst, // [e*n, d] const T* src, // [ n, d] const int* f2n, // [e*n] :: e*n -> n + const int* num_valid_tokens, int dims) { + if (num_valid_tokens && blockIdx.x >= __ldg(num_valid_tokens)) { + return; + } + using Vec = Array; const int64_t bi = blockIdx.x; @@ -899,23 +930,32 @@ __global__ void MoeGatherKernel(T* dst, // [e*n, d] } } -void invokeMoeDispatch(Ref out_, const Tensor& src, const int* f2n, int expert_per_token, cudaStream_t st) +void invokeMoeDispatch(Ref out_, + const Tensor& src, + const int* f2n, + int num_worst_tokens, + const int* num_valid_tokens, + cudaStream_t st) { auto& out = out_.get(); auto invoke = [&](auto t) { using T = decltype(t); - auto [num, dim] = src.shapes(0, 1); + const int dim = src.shape(1); constexpr int threads = 256; constexpr int vec_size = 16 / sizeof(T); - // std::cout << num * expert_per_token << " " << dim << "\n"; - MoeGatherKernel<<>>( // + // f2n/out have num_worst_tokens rows; num_valid_tokens limits rows that read f2n. + MoeGatherKernel<<>>( // (T*)out.raw_data(), (const T*)src.raw_data(), f2n, + num_valid_tokens, dim / vec_size); TM_CUDA_CHECK(cudaGetLastError()); }; TM_CHECK_EQ(src.dtype(), out.dtype()); + if (num_worst_tokens == 0) { + return; + } const auto elem_size = byte_size(src.dtype()); if (elem_size == sizeof(uint16_t)) { invoke(uint16_t{}); @@ -958,10 +998,14 @@ __global__ void MoeDispatchScales( } template -__global__ void -MoeDispatchScalesNonaligned(T* dst, const T* src, int dst_stride, int src_stride, const int* f2n, int dim) +__global__ void MoeDispatchScalesNonaligned( + T* dst, const T* src, int dst_stride, int src_stride, const int* f2n, const int* num_valid_tokens, int dim) { const int bi = blockIdx.x; + if (num_valid_tokens && bi >= __ldg(num_valid_tokens)) { + return; + } + const int ti = f2n[bi]; if (threadIdx.x < dim) { @@ -969,14 +1013,20 @@ MoeDispatchScalesNonaligned(T* dst, const T* src, int dst_stride, int src_stride } } -void invokeMoeDispatchScales(Ref out_, const Tensor& src, const int* f2n, int expert_per_token, cudaStream_t st) +void invokeMoeDispatchScales(Ref out_, + const Tensor& src, + const int* f2n, + int num_worst_tokens, + const int* num_valid_tokens, + cudaStream_t st) { using T = float; constexpr int alignment = 16 / sizeof(T); - auto [dim, num] = src.shapes(0, 1); + const int dim = src.shape(0); - const int size = num * expert_per_token; + // Keep the scale layout aligned to num_worst_tokens; num_valid_tokens limits rows that read f2n. + const int size = num_worst_tokens; const int aligned_size = round_up(size, alignment); auto& out = out_.get(); @@ -991,6 +1041,9 @@ void invokeMoeDispatchScales(Ref out_, const Tensor& src, const int* f2n } TM_CHECK_LE(dim, 1024); + if (size == 0) { + return; + } const int threads = round_up(dim, WARP_SIZE); const int blocks = size; @@ -1001,6 +1054,7 @@ void invokeMoeDispatchScales(Ref out_, const Tensor& src, const int* f2n out.stride(0), src.stride(0), f2n, + num_valid_tokens, dim); TM_CUDA_CHECK(cudaGetLastError()); diff --git a/src/turbomind/kernels/gemm/moe_utils_v2.h b/src/turbomind/kernels/gemm/moe_utils_v2.h index ae72148c95..15b1b2c889 100644 --- a/src/turbomind/kernels/gemm/moe_utils_v2.h +++ b/src/turbomind/kernels/gemm/moe_utils_v2.h @@ -31,16 +31,33 @@ void invokeMoeGate_V2(int* f2n, float routed_scale, cudaStream_t st); +void invokeMoeScanKernel(int* f2n, + int* f2E, + int* en2f, + int* offsets, + int8_t* masks, + const int* accum, + int tokens, + int tokens_padded, + int experts, + cudaStream_t st); + +// num_worst_tokens is the output capacity / launch upper bound for f2n/out. +// If num_valid_tokens is set, it points to a device-side count of valid rows and +// rows >= *num_valid_tokens return before reading f2n. void invokeMoeDispatch(Ref out_, // const Tensor& src, const int* f2n, - int expert_per_token, + int num_worst_tokens, + const int* num_valid_tokens, cudaStream_t st); +// Same num_worst_tokens / num_valid_tokens contract as invokeMoeDispatch void invokeMoeDispatchScales(Ref out_, // const Tensor& src, const int* f2n, - int expert_per_token, + int num_worst_tokens, + const int* num_valid_tokens, cudaStream_t st); void invokeMoeCombine(Ref out_, diff --git a/src/turbomind/kernels/gemm/test/testbed_v3.h b/src/turbomind/kernels/gemm/test/testbed_v3.h index 805b9aca02..d633baa18b 100644 --- a/src/turbomind/kernels/gemm/test/testbed_v3.h +++ b/src/turbomind/kernels/gemm/test/testbed_v3.h @@ -378,7 +378,7 @@ struct Testbed_v3: Parameter { Tensor xe{{x.shape(0) * experts_per_token, input_dim}, data_type, kDEVICE}; Tensor de{{x.shape(0) * experts_per_token, output_dim}, data_type, kDEVICE}; - TM_SCOPE_CALL(invokeMoeDispatch(xe, x, f2n_.data(), experts_per_token, stream_)); + TM_SCOPE_CALL(invokeMoeDispatch(xe, x, f2n_.data(), x.shape(0) * experts_per_token, nullptr, stream_)); for (int i = 0; i < expert_num; ++i) { const int base = h_offsets_[i], size = h_offsets_[i + 1] - base; diff --git a/src/turbomind/models/llama/LlamaLinear.cu b/src/turbomind/models/llama/LlamaLinear.cu index 515a1eb8dd..7547bf0c23 100644 --- a/src/turbomind/models/llama/LlamaLinear.cu +++ b/src/turbomind/models/llama/LlamaLinear.cu @@ -63,8 +63,12 @@ struct LlamaLinear::Impl { return {B, desc_B, V, desc_V}; } - std::tuple - GetOperandA(const LinearWeight& weight, const Tensor& input, Buffer_ indices, const Buffer_& offsets) + std::tuple // + GetOperandA(const LinearWeight& weight, + const Tensor& input, + Buffer_ indices, + const Buffer_& offsets, + bool indices_padded) { auto st = core::Context::stream().handle(); @@ -84,13 +88,13 @@ struct LlamaLinear::Impl { // SM100+ grouped bf16/fp16: use chunk() weights so Activation() runs separately. const bool is_cublas_grouped = offsets && getSMVersion() == 100 && weight.weight_format.dtype == kBfloat16; if (indices && (A.dtype() == kFloat8_e4m3 || is_cublas_grouped)) { - const auto [bsz, k] = A.shapes(0, 1); - const int e = indices.size() / bsz; - Tensor A_e = {{m, k}, A.dtype(), kDEVICE}; - TM_SCOPE_CALL(invokeMoeDispatch(A_e, A, indices.data(), e, st)); + const int k = A.shape(1); + const int* num_valid_tokens = indices_padded ? offsets.data() + offsets.size() - 1 : nullptr; + Tensor A_e = {{m, k}, A.dtype(), kDEVICE}; + TM_SCOPE_CALL(invokeMoeDispatch(A_e, A, indices.data(), m, num_valid_tokens, st)); if (U) { Tensor U_e; - TM_SCOPE_CALL(invokeMoeDispatchScales(U_e, U, indices.data(), e, st)); + TM_SCOPE_CALL(invokeMoeDispatchScales(U_e, U, indices.data(), m, num_valid_tokens, st)); U = U_e; } A = A_e; @@ -117,7 +121,8 @@ struct LlamaLinear::Impl { const Tensor& input, // const LinearWeight& weight, const Buffer_& indices, - const Buffer_& offsets) + const Buffer_& offsets, + bool indices_padded) { TM_FUNCTION_SCOPE(); using namespace gemm; @@ -129,7 +134,7 @@ struct LlamaLinear::Impl { op.quant_b = MakeQuantDesc(weight.weight_format); op.batch_dim = 0; - auto&& [A, desc_A, U, desc_U] = GetOperandA(weight, input, indices, offsets); + auto&& [A, desc_A, U, desc_U] = GetOperandA(weight, input, indices, offsets, indices_padded); auto&& [B, desc_B, V, desc_V] = GetOperandB(weight); Tensor& D = output; @@ -195,7 +200,8 @@ void LlamaLinear::Forward(const Tensor& input, // const LinearWeight& weight, const Buffer_& indices, const Buffer_& offsets, - Ref output) + Ref output, + bool indices_padded) { Tensor in = input.view({-1, input.shape(-1)}); @@ -203,7 +209,7 @@ void LlamaLinear::Forward(const Tensor& input, // output.get() = output.get().view({-1, output.get().shape(-1)}); } - impl_->Forward(output.get(), in, weight, indices, offsets); + impl_->Forward(output.get(), in, weight, indices, offsets, indices_padded); } void LlamaLinear::set_measure(bool measure) diff --git a/src/turbomind/models/llama/LlamaLinear.h b/src/turbomind/models/llama/LlamaLinear.h index ee3315a73d..cbd1a73efa 100644 --- a/src/turbomind/models/llama/LlamaLinear.h +++ b/src/turbomind/models/llama/LlamaLinear.h @@ -22,7 +22,8 @@ class LlamaLinear { const LinearWeight& weight, const Buffer_& indices, const Buffer_& offsets, - Ref output); + Ref output, + bool indices_padded = false); void set_measure(bool measure); diff --git a/src/turbomind/models/llama/moe_ffn_layer.cc b/src/turbomind/models/llama/moe_ffn_layer.cc index 101175e861..b51ab89469 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.cc +++ b/src/turbomind/models/llama/moe_ffn_layer.cc @@ -8,6 +8,7 @@ #include "src/turbomind/core/context.h" #include "src/turbomind/core/scope.h" #include "src/turbomind/kernels/activation.h" +#include "src/turbomind/kernels/gemm/moe_ep_utils.h" #include "src/turbomind/kernels/gemm/moe_utils_v2.h" #include "src/turbomind/kernels/norm/rms_norm.h" @@ -274,37 +275,228 @@ void MoeFfnTpImpl::Combine(MoeFfnLayer::ForwardParam& p) class MoeFfnAgRsImpl final: public MoeFfnLayerImpl { public: - MoeFfnAgRsImpl(const EngineParam& engine, const Context& ctx): - MoeFfnLayerImpl(ctx), ep_size_(engine.ep_size), ep_rank_(engine.ep_rank), backend_(engine.all2all_backend) - { - TM_CHECK_GT(ep_size_, 1); - TM_CHECK_EQ(engine.mlp_tp_size, 1); - TM_CHECK_EQ(backend_, "allgather_reducescatter"); - // Decoder owns attention TP/CP-group FFN boundary communication. This backend consumes and produces - // this rank's compact token slice and owns expert-parallel allgather/reducescatter communication. + MoeFfnAgRsImpl(const EngineParam& engine, const Context& ctx); + + void Forward(MoeFfnLayer::ForwardParam& p) override; + + void Combine(MoeFfnLayer::ForwardParam& p) override; + +private: + void Init(MoeFfnLayer::ForwardParam& p); + + void ForwardFused(MoeFfnLayer::ForwardParam& p); + + const int ep_size_; + const int ep_rank_; + const int max_token_num_; + const std::string backend_; + comm::DeviceCommImpl* const d_comm_; + Allocator symm_allocator_; + + bool initialized_ = false; + + Buffer_ topk_weights_; + Buffer_ topk_idx_; + Buffer_ f2n_; + Buffer_ f2E_; + Buffer_ en2f_; + Buffer_ offsets_; + Buffer_ masks_; + Buffer_ accum_; + + Tensor temp_; + Tensor_ shared_scales_; +}; + +MoeFfnAgRsImpl::MoeFfnAgRsImpl(const EngineParam& engine, const Context& ctx): + MoeFfnLayerImpl(ctx), + ep_size_(engine.ep_size), + ep_rank_(engine.ep_rank), + max_token_num_(engine.max_forward_token_num * engine.attn_dp_size), + backend_(engine.all2all_backend), + d_comm_(ctx.comm.d_comm), + symm_allocator_(GetSymmAllocator(ctx.comm.d_comm)) +{ +} + +void MoeFfnAgRsImpl::Init(MoeFfnLayer::ForwardParam& p) +{ + const auto& moe = *p.weights; + const int local_expert_num = moe.num_local_experts(); + const int experts_per_token = moe.experts_per_token; + const int pad_token_num = (max_token_num_ + kMoeGateVecSize - 1) / kMoeGateVecSize * kMoeGateVecSize; + + TM_CHECK(p.weights->topk_method != "noaux_tc") + << "MoeFfnAgRsImpl does not support topk_method=noaux_tc because it requires noaux-specific scoring and " + "correction-bias handling"; + + topk_weights_ = {max_token_num_ * experts_per_token, symm_allocator_}; + topk_idx_ = {max_token_num_ * experts_per_token, symm_allocator_}; + masks_ = {local_expert_num * pad_token_num, kDEVICE}; + f2n_ = {max_token_num_ * experts_per_token, kDEVICE}; + f2E_ = {max_token_num_ * experts_per_token, kDEVICE}; + en2f_ = {max_token_num_ * experts_per_token, kDEVICE}; + offsets_ = {local_expert_num + 1, kDEVICE}; + accum_ = {local_expert_num * kMoeGateMaxTiles, kDEVICE}; + + initialized_ = true; +} + +void MoeFfnAgRsImpl::Forward(MoeFfnLayer::ForwardParam& p) +{ + TM_FUNCTION_SCOPE(); + + if (!initialized_) { + Init(p); } - void Forward(MoeFfnLayer::ForwardParam&) override - { - ReportUnsupported(); + const auto& moe = *p.weights; + const int tokens = p.global_hidden_states.shape(0); + const int experts_per_token = moe.experts_per_token; + const int expert_num = moe.num_experts(); + const int local_expert_num = moe.num_local_experts(); + const int hidden_dim = TM_CHECK_NOTNULL(moe.block())->hidden_dim; + const auto st = core::Context::stream().handle(); + + TM_CHECK_LE(tokens, max_token_num_); + + std::vector topk_elem_counts(ep_size_); + size_t local_topk_offset{}; + for (int rank = 0; rank < ep_size_; ++rank) { + topk_elem_counts[rank] = TM_CHECK_NOTNULL(p.ep_elem_counts)->at(rank) / hidden_dim * experts_per_token; + local_topk_offset += rank < ep_rank_ ? topk_elem_counts[rank] : 0; } - void Combine(MoeFfnLayer::ForwardParam&) override - { - ReportUnsupported(); + auto local_topk_weights = topk_weights_.data() + local_topk_offset; + auto local_topk_idx = topk_idx_.data() + local_topk_offset; + + if (auto local_tokens = p.input.shape(0); local_tokens > 0) { + auto logits = Gate(p.input, *moe.gate.get()); + TM_DEBUG_TENSOR(logits, "logits", 2); + + bool softmax = true; + if (p.weights->topk_method == "group_limited_greedy") { + invokeMoeSoftmaxMaskTopKGroups( + logits.data(), local_tokens, expert_num, expert_num / p.weights->n_group, p.weights->topk_group, st); + TM_CUDA_CHECK(cudaGetLastError()); + softmax = false; + } + + invokeMoeGateEp(local_topk_weights, + local_topk_idx, + logits.data(), + local_tokens, + expert_num, + experts_per_token, + softmax, + p.weights->norm_topk_prob, + p.weights->routed_scale, + st); + TM_CUDA_CHECK(cudaGetLastError()); } -private: - void ReportUnsupported() const - { - TM_LOG_FATAL("MoeFfnAgRsImpl is a stub: decoder supplies a compact attention-TP/CP token slice; implement EP " - "allgather-reducescatter dispatch/combine before enabling EP MoE"); + d_comm_->GroupStart(); + d_comm_->AllGatherV( + p.input.raw_data(), p.global_hidden_states.raw_data(), *p.ep_elem_counts, p.input.dtype(), 0, st); + d_comm_->AllGatherV(local_topk_weights, topk_weights_.data(), topk_elem_counts, kFloat32, 0, st); + d_comm_->AllGatherV(local_topk_idx, topk_idx_.data(), topk_elem_counts, kInt64, 0, st); + d_comm_->GroupEnd(); + TM_CUDA_CHECK(cudaGetLastError()); + + invokeMoeBuildAgRsRoutingMap(f2n_.data(), + f2E_.data(), + en2f_.data(), + offsets_.data(), + masks_.data(), + accum_.data(), + topk_idx_.data(), + tokens, + experts_per_token, + moe.local_expert_offset(), + local_expert_num, + st); + TM_CUDA_CHECK(cudaGetLastError()); + + temp_ = Tensor{{tokens * experts_per_token, hidden_dim}, p.input.dtype(), p.input.device()}; + ForwardFused(p); + + shared_scales_ = {}; + if (moe.shared_gate && p.input.shape(0) > 0) { + shared_scales_ = Gate(p.input, *moe.shared_gate); } +} - int ep_size_; - int ep_rank_; - std::string backend_; -}; +void MoeFfnAgRsImpl::ForwardFused(MoeFfnLayer::ForwardParam& p) +{ + TM_FUNCTION_SCOPE(); + + const auto& moe = *p.weights; + const auto& block = *TM_CHECK_NOTNULL(moe.block()); + const int local_expert_num = moe.num_local_experts(); + const auto st = core::Context::stream().handle(); + + auto input = p.global_hidden_states; + const int* num_valid_tokens = offsets_.data() + local_expert_num; + auto indices = f2n_.slice(0, temp_.shape(0)); + auto offsets = offsets_.slice(0, local_expert_num + 1); + const bool indices_padded = true; + + if (block.w1w3) { + Tensor inter; + TM_SCOPE_CALL(linear_.Forward(input, *block.w1w3, indices, offsets, inter, indices_padded)); + + if (!block.is_fused_silu) { + Activation(inter, block.w1w3->bias, f2E_, block.act_type, num_valid_tokens, st); + TM_CUDA_CHECK(cudaGetLastError()); + } + + TM_SCOPE_CALL(linear_.Forward(inter.slice({0, 0}, {-1, block.inter_size}), *block.w2, {}, offsets, temp_)); + } + else { + Tensor gating; + TM_SCOPE_CALL(linear_.Forward(input, *block.w1, indices, offsets, gating, indices_padded)); + + Tensor up; + TM_SCOPE_CALL(linear_.Forward(input, *block.w3, indices, offsets, up, indices_padded)); + + Activation(gating, up, block.act_type, num_valid_tokens, st); + TM_CUDA_CHECK(cudaGetLastError()); + + TM_SCOPE_CALL(linear_.Forward(gating, *block.w2, {}, offsets, temp_)); + } +} + +void MoeFfnAgRsImpl::Combine(MoeFfnLayer::ForwardParam& p) +{ + TM_FUNCTION_SCOPE(); + + const auto& moe = *p.weights; + const auto& block = *TM_CHECK_NOTNULL(moe.block()); + const int experts_per_token = moe.experts_per_token; + const auto st = core::Context::stream().handle(); + + invokeMoeLocalCombineEp(p.global_hidden_states, + temp_, + block.w2->bias, + topk_weights_.data(), + en2f_.data(), + f2E_.data(), + experts_per_token, + st); + TM_CUDA_CHECK(cudaGetLastError()); + + d_comm_->ReduceScatterV( + p.global_hidden_states.raw_data(), p.output.raw_data(), *p.ep_elem_counts, p.output.dtype(), 0, st); + TM_CUDA_CHECK(cudaGetLastError()); + + if (p.shared_output) { + invokeMoeCombineOutputEp(p.output, p.shared_output, shared_scales_.data_or((float*)nullptr), p.scale, st); + TM_CUDA_CHECK(cudaGetLastError()); + } + + temp_ = {}; + shared_scales_ = {}; +} MoeFfnLayer::MoeFfnLayer(const EngineParam& engine, const Context& ctx) { diff --git a/src/turbomind/models/llama/moe_ffn_layer.h b/src/turbomind/models/llama/moe_ffn_layer.h index 39f68fa952..f1c0371db9 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.h +++ b/src/turbomind/models/llama/moe_ffn_layer.h @@ -3,6 +3,7 @@ #pragma once #include +#include #include "src/turbomind/core/core.h" #include "src/turbomind/models/llama/context.h" @@ -25,6 +26,11 @@ class MoeFfnLayer { const MoeWeight* weights; float scale; int layer_id; + + // EP allgather-reducescatter path only. + Tensor global_hidden_states; + const std::vector* ep_elem_counts{}; + Tensor shared_output; }; void Forward(ForwardParam& p); diff --git a/src/turbomind/models/llama/unified_decoder.cc b/src/turbomind/models/llama/unified_decoder.cc index efa3a88848..f70a4554f5 100644 --- a/src/turbomind/models/llama/unified_decoder.cc +++ b/src/turbomind/models/llama/unified_decoder.cc @@ -29,24 +29,43 @@ namespace turbomind { struct FfnTokenPartition { int first_token{}; int token_count{}; - std::vector elem_counts; + std::vector tp_elem_counts; + std::vector ep_elem_counts; static FfnTokenPartition Make(const std::vector& local_token_nums, int dp_rank, int group_rank, int group_size, size_t hidden_units) { + auto split_count = [group_size](int local_token_num, int rank) { + const int base = local_token_num / group_size; + const int remainder = local_token_num % group_size; + return base + (rank < remainder ? 1 : 0); + }; + + auto first_token = [group_size](int local_token_num, int rank) { + const int base = local_token_num / group_size; + const int remainder = local_token_num % group_size; + return rank * base + std::min(rank, remainder); + }; + const int local_token_num = local_token_nums[dp_rank]; - const int base = local_token_num / group_size; - const int remainder = local_token_num % group_size; - const int first_token = group_rank * base + std::min(group_rank, remainder); - const int token_count = base + (group_rank < remainder ? 1 : 0); - std::vector elem_counts(group_size); + std::vector tp_elem_counts(group_size); for (int rank = 0; rank < group_size; ++rank) { - const int count = base + (rank < remainder ? 1 : 0); - elem_counts[rank] = static_cast(count) * hidden_units; + tp_elem_counts[rank] = static_cast(split_count(local_token_num, rank)) * hidden_units; + } + + std::vector ep_elem_counts; + ep_elem_counts.reserve(local_token_nums.size() * group_size); + for (const int dp_token_num : local_token_nums) { + for (int rank = 0; rank < group_size; ++rank) { + ep_elem_counts.push_back(static_cast(split_count(dp_token_num, rank)) * hidden_units); + } } - return FfnTokenPartition{first_token, token_count, std::move(elem_counts)}; + return FfnTokenPartition{first_token(local_token_num, group_rank), + split_count(local_token_num, group_rank), + std::move(tp_elem_counts), + std::move(ep_elem_counts)}; } }; @@ -192,8 +211,8 @@ void UnifiedDecoder::ReduceScatterVResidualRMSnorm(Tensor& loca TM_CHECK(ep_size_ > 1); TM_CHECK(d_comm_); TM_CHECK_EQ(local_residual.shape(0), local_hidden_states.shape(0)); - TM_CHECK_EQ(partition.elem_counts.size(), static_cast(d_comm_->n_ranks(attn_tp_group_))); - TM_CHECK_EQ(std::accumulate(partition.elem_counts.begin(), partition.elem_counts.end(), size_t{}), + TM_CHECK_EQ(partition.tp_elem_counts.size(), static_cast(d_comm_->n_ranks(attn_tp_group_))); + TM_CHECK_EQ(std::accumulate(partition.tp_elem_counts.begin(), partition.tp_elem_counts.end(), size_t{}), static_cast(local_hidden_states.size())); const auto dtype = local_hidden_states.dtype(); @@ -202,8 +221,12 @@ void UnifiedDecoder::ReduceScatterVResidualRMSnorm(Tensor& loca Tensor hidden_slice = local_hidden_states.slice({partition.first_token, 0}, {partition.token_count, -1}); Tensor residual_slice = local_residual.slice({partition.first_token, 0}, {partition.token_count, -1}); - d_comm_->ReduceScatterV( - local_hidden_states.raw_data(), hidden_slice.raw_data(), partition.elem_counts, dtype, attn_tp_group_, stream); + d_comm_->ReduceScatterV(local_hidden_states.raw_data(), + hidden_slice.raw_data(), + partition.tp_elem_counts, + dtype, + attn_tp_group_, + stream); TM_CUDA_CHECK(cudaGetLastError()); invokeResidualBiasRMSNorm(hidden_slice.raw_data(), @@ -227,8 +250,8 @@ void UnifiedDecoder::ResidualRMSnormAllGatherV(Tensor& local_hi TM_CHECK(ep_size_ > 1); TM_CHECK(d_comm_); TM_CHECK_EQ(local_residual.shape(0), local_hidden_states.shape(0)); - TM_CHECK_EQ(partition.elem_counts.size(), static_cast(d_comm_->n_ranks(attn_tp_group_))); - TM_CHECK_EQ(std::accumulate(partition.elem_counts.begin(), partition.elem_counts.end(), size_t{}), + TM_CHECK_EQ(partition.tp_elem_counts.size(), static_cast(d_comm_->n_ranks(attn_tp_group_))); + TM_CHECK_EQ(std::accumulate(partition.tp_elem_counts.begin(), partition.tp_elem_counts.end(), size_t{}), static_cast(local_hidden_states.size())); const auto dtype = local_hidden_states.dtype(); @@ -248,8 +271,12 @@ void UnifiedDecoder::ResidualRMSnormAllGatherV(Tensor& local_hi stream); TM_CUDA_CHECK(cudaGetLastError()); - d_comm_->AllGatherV( - hidden_slice.raw_data(), local_hidden_states.raw_data(), partition.elem_counts, dtype, attn_tp_group_, stream); + d_comm_->AllGatherV(hidden_slice.raw_data(), + local_hidden_states.raw_data(), + partition.tp_elem_counts, + dtype, + attn_tp_group_, + stream); TM_CUDA_CHECK(cudaGetLastError()); } @@ -408,14 +435,22 @@ void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vector moe_fwd_param; if (weights.at(layer)->moe_ffn) { - moe_fwd_param = MoeFfnLayer::ForwardParam{ - ffn_hidden_states, ffn_hidden_states, weights.at(layer)->moe_ffn.get(), ffn_layer_ ? 1.f : 0.f, layer}; + moe_fwd_param = MoeFfnLayer::ForwardParam{ffn_hidden_states, + ffn_hidden_states, + weights.at(layer)->moe_ffn.get(), + ffn_layer_ ? 1.f : 0.f, + layer, + ep_size_ > 1 ? global_hidden_states : Tensor{}, + ep_size_ > 1 ? &ffn_partition->ep_elem_counts : nullptr}; moe_ffn_layer_->Forward(*moe_fwd_param); } if (ffn_layer_ && weights.at(layer)->feed_forward && ffn_hidden_states.shape(0) > 0) { - ffn_layer_->forward( - {ffn_hidden_states, ffn_hidden_states, weights.at(layer)->feed_forward.get(), (int)layer}); + auto ffn_output = ffn_hidden_states; + if (ep_size_ > 1) { + ffn_output = moe_fwd_param->shared_output = core::empty_like(ffn_hidden_states); + } + ffn_layer_->forward({ffn_hidden_states, ffn_output, weights.at(layer)->feed_forward.get(), (int)layer}); } if (moe_fwd_param) { From 45c18c3a017ab151fbf8fa15421fe3994bbab5a9 Mon Sep 17 00:00:00 2001 From: irexyc Date: Mon, 6 Jul 2026 07:10:48 +0000 Subject: [PATCH 05/19] fix moe reduce int overflow when using dp --- src/turbomind/kernels/gemm/moe_utils_v2.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/turbomind/kernels/gemm/moe_utils_v2.cu b/src/turbomind/kernels/gemm/moe_utils_v2.cu index 070a89a52e..c74e9cc364 100644 --- a/src/turbomind/kernels/gemm/moe_utils_v2.cu +++ b/src/turbomind/kernels/gemm/moe_utils_v2.cu @@ -1076,7 +1076,7 @@ __global__ void MoeReduceKernel(T* dst, // [ n, d] if constexpr (TURBOMIND_ARCH_DTYPE_GUARD(data_type_v)) { const int64_t ti = blockIdx.x; - dst += dim * ti; + dst += (int64_t)dim * ti; if (dst_scales) { dst_scale = dst_scales[ti]; @@ -1092,9 +1092,9 @@ __global__ void MoeReduceKernel(T* dst, // [ n, d] PRAGMA_UNROLL for (int e = 0; e < exp_k; ++e) { int fid = __ldg(&en2f[e * tokens + ti]); - src_[e] = src + dim * fid; + src_[e] = src + (int64_t)dim * fid; if constexpr (has_bias) { - bias_[e] = bias + __ldg(&f2E[fid]) * dim; + bias_[e] = bias + __ldg(&f2E[fid]) * (int64_t)dim; } scale[e] = scales ? __ldg(&scales[e * tokens + ti]) : 1.f; } From b851ec3b195f91f7858f49c84ebc06e2bac7db34 Mon Sep 17 00:00:00 2001 From: irexyc Date: Mon, 6 Jul 2026 08:46:21 +0000 Subject: [PATCH 06/19] fix dp when using ep --- lmdeploy/turbomind/turbomind.py | 4 ++-- src/turbomind/models/llama/unified_decoder.cc | 16 ++++++++-------- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/lmdeploy/turbomind/turbomind.py b/lmdeploy/turbomind/turbomind.py index dd924e704d..df35bb296a 100644 --- a/lmdeploy/turbomind/turbomind.py +++ b/lmdeploy/turbomind/turbomind.py @@ -98,13 +98,13 @@ def update_parallel_config(cfg: TurbomindEngineConfig): cfg.attn_dp_size = overlap cfg.attn_tp_size = inner_tp_size // cfg.cp cfg.attn_cp_size = cfg.cp + comm_size = cfg.attn_dp_size * cfg.attn_tp_size * cfg.attn_cp_size cfg.mlp_dp_size = 1 - cfg.mlp_tp_size = total // cfg.ep if cfg.ep > 1 else overlap * inner_tp_size + cfg.mlp_tp_size = comm_size // cfg.ep if cfg.ep > 1 else overlap * inner_tp_size if cfg.ep > 1: assert cfg.communicator == 'nccl', f'{cfg.communicator} communicator does not support ep > 1' assert cfg.mlp_tp_size == 1, 'Only support mlp_tp_size == 1 when ep > 1' assert cfg.attn_dp_size * cfg.attn_tp_size * cfg.attn_cp_size * cfg.outer_dp_size == cfg.device_num - assert cfg.attn_dp_size * cfg.attn_tp_size * cfg.attn_cp_size == cfg.mlp_dp_size * cfg.mlp_tp_size * cfg.ep # update devices cfg.devices = cfg.devices or list(range(cfg.device_num // cfg.nnodes)) cfg.devices = cfg.devices[:cfg.device_num // cfg.nnodes] diff --git a/src/turbomind/models/llama/unified_decoder.cc b/src/turbomind/models/llama/unified_decoder.cc index f70a4554f5..a2dab7e9e0 100644 --- a/src/turbomind/models/llama/unified_decoder.cc +++ b/src/turbomind/models/llama/unified_decoder.cc @@ -221,16 +221,16 @@ void UnifiedDecoder::ReduceScatterVResidualRMSnorm(Tensor& loca Tensor hidden_slice = local_hidden_states.slice({partition.first_token, 0}, {partition.token_count, -1}); Tensor residual_slice = local_residual.slice({partition.first_token, 0}, {partition.token_count, -1}); - d_comm_->ReduceScatterV(local_hidden_states.raw_data(), - hidden_slice.raw_data(), + d_comm_->ReduceScatterV(local_hidden_states.data_or((void*)nullptr), + residual_slice.data_or((void*)nullptr), partition.tp_elem_counts, dtype, attn_tp_group_, stream); TM_CUDA_CHECK(cudaGetLastError()); - invokeResidualBiasRMSNorm(hidden_slice.raw_data(), - residual_slice.raw_data(), + invokeResidualBiasRMSNorm(hidden_slice.data_or((void*)nullptr), + residual_slice.data_or((void*)nullptr), weight.raw_data(), bias.data_or((void*)nullptr), dtype, @@ -260,8 +260,8 @@ void UnifiedDecoder::ResidualRMSnormAllGatherV(Tensor& local_hi Tensor hidden_slice = local_hidden_states.slice({partition.first_token, 0}, {partition.token_count, -1}); Tensor residual_slice = local_residual.slice({partition.first_token, 0}, {partition.token_count, -1}); - invokeResidualBiasRMSNorm(hidden_slice.raw_data(), - residual_slice.raw_data(), + invokeResidualBiasRMSNorm(hidden_slice.data_or((void*)nullptr), + residual_slice.data_or((void*)nullptr), weight.raw_data(), nullptr, dtype, @@ -271,8 +271,8 @@ void UnifiedDecoder::ResidualRMSnormAllGatherV(Tensor& local_hi stream); TM_CUDA_CHECK(cudaGetLastError()); - d_comm_->AllGatherV(hidden_slice.raw_data(), - local_hidden_states.raw_data(), + d_comm_->AllGatherV(hidden_slice.data_or((void*)nullptr), + local_hidden_states.data_or((void*)nullptr), partition.tp_elem_counts, dtype, attn_tp_group_, From cc3509141a73bf880c88144170a8bb27713fadd9 Mon Sep 17 00:00:00 2001 From: irexyc Date: Tue, 7 Jul 2026 03:31:39 +0000 Subject: [PATCH 07/19] update AllGatherV / ReduceScatterV --- src/turbomind/comm/nccl/nccl.cu | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/src/turbomind/comm/nccl/nccl.cu b/src/turbomind/comm/nccl/nccl.cu index 504b4a943d..f481378677 100644 --- a/src/turbomind/comm/nccl/nccl.cu +++ b/src/turbomind/comm/nccl/nccl.cu @@ -279,6 +279,17 @@ public: TM_CHECK_EQ((int)counts.size(), n_ranks); + bool equal_counts = true; + for (auto i = 1; i < counts.size(); ++i) { + if (counts[i] != counts[0]) { + equal_counts = false; + break; + } + } + if (equal_counts) { + return AllGather(sendbuff, recvbuff, counts[0], type, group, stream); + } + NCCLCHECK(ncclGroupStart()); size_t offset = 0; for (int i = 0; i < n_ranks; ++i) { @@ -317,6 +328,17 @@ public: TM_CHECK_EQ((int)counts.size(), n_ranks); + bool equal_counts = true; + for (auto i = 1; i < counts.size(); ++i) { + if (counts[i] != counts[0]) { + equal_counts = false; + break; + } + } + if (equal_counts) { + return ReduceScatter(sendbuff, recvbuff, counts[0], type, group, stream); + } + NCCLCHECK(ncclGroupStart()); size_t offset = 0; for (int i = 0; i < n_ranks; ++i) { From c40fa849065f7851a38901c0138b2e266afe2c54 Mon Sep 17 00:00:00 2001 From: irexyc Date: Wed, 8 Jul 2026 09:21:11 +0000 Subject: [PATCH 08/19] replace reduce_scatter/allgatherv with allreduce --- lmdeploy/turbomind/turbomind.py | 2 +- src/turbomind/kernels/gemm/moe_utils_v2.cu | 41 ++-- src/turbomind/kernels/gemm/moe_utils_v2.h | 2 + src/turbomind/models/llama/moe_ffn_layer.cc | 170 ++++++-------- src/turbomind/models/llama/moe_ffn_layer.h | 6 - src/turbomind/models/llama/unified_decoder.cc | 212 +++--------------- src/turbomind/models/llama/unified_decoder.h | 14 -- 7 files changed, 116 insertions(+), 331 deletions(-) diff --git a/lmdeploy/turbomind/turbomind.py b/lmdeploy/turbomind/turbomind.py index 8fa2e43c74..9fcf75cb30 100644 --- a/lmdeploy/turbomind/turbomind.py +++ b/lmdeploy/turbomind/turbomind.py @@ -102,7 +102,7 @@ def update_parallel_config(cfg: TurbomindEngineConfig): cfg.mlp_dp_size = 1 cfg.mlp_tp_size = comm_size // cfg.ep if cfg.ep > 1 else overlap * inner_tp_size if cfg.ep > 1: - assert cfg.communicator == 'nccl', f'{cfg.communicator} communicator does not support ep > 1' + assert cfg.nnodes == 1, 'ep > 1 is only supported in single-node mode' assert cfg.mlp_tp_size == 1, 'Only support mlp_tp_size == 1 when ep > 1' assert cfg.attn_dp_size * cfg.attn_tp_size * cfg.attn_cp_size * cfg.outer_dp_size == cfg.device_num # update devices diff --git a/src/turbomind/kernels/gemm/moe_utils_v2.cu b/src/turbomind/kernels/gemm/moe_utils_v2.cu index c74e9cc364..384845eca4 100644 --- a/src/turbomind/kernels/gemm/moe_utils_v2.cu +++ b/src/turbomind/kernels/gemm/moe_utils_v2.cu @@ -598,17 +598,19 @@ void invokeMoeScanKernel(int* f2n, } void invokeMoeGate_V2(int* f2n, // [e*n] -> n - int* f2E, // [e*n] -> E + int* f2E, // [e*n] -> local E int* en2f, // [e,n] -> n*e - int* offsets, // [E+1] + int* offsets, // [local E+1] float* scales, // [e,n] void* masks, // [E,n] - int* accum, // [E] - const float* logits, // [e,n] + int* accum, // [E,tiles] + const float* logits, // [n,E] int tokens, // n int tokens_padded, // round_up(n, 4) int experts, // E int experts_per_token, + int local_expert_offset, + int local_expert_num, bool softmax, bool norm_topk, float routed_scale, @@ -697,19 +699,19 @@ void invokeMoeGate_V2(int* f2n, // [e*n] -> n { constexpr int threads = (1 << base_log_tile) / kMoeGateVecSize; - const dim3 blocks(tiles, experts + 1); + const dim3 blocks(tiles, local_expert_num + 1); MoeScanKernel_v2<<>>(f2n, // f2E, en2f, offsets, - (int8_t*)masks, - accum, + (int8_t*)masks + local_expert_offset * tokens_padded, + accum + local_expert_offset * tiles, log_tile, tiles, tokens, tokens_padded, - experts); + local_expert_num); } TM_CUDA_CHECK(cudaGetLastError()); } @@ -1079,24 +1081,26 @@ __global__ void MoeReduceKernel(T* dst, // [ n, d] dst += (int64_t)dim * ti; if (dst_scales) { - dst_scale = dst_scales[ti]; - dst_scale = fdividef(1.f, 1.f + expf(-dst_scale)); + const float scale = dst_scales[ti]; + dst_scale *= fdividef(1.f, 1.f + expf(-scale)); } // Should be warp uniforms - const T* src_[exp_k]; - const T* bias_[exp_k]; + const T* src_[exp_k]{}; + const T* bias_[exp_k]{}; - float scale[exp_k]; + float scale[exp_k]{}; PRAGMA_UNROLL for (int e = 0; e < exp_k; ++e) { int fid = __ldg(&en2f[e * tokens + ti]); - src_[e] = src + (int64_t)dim * fid; - if constexpr (has_bias) { - bias_[e] = bias + __ldg(&f2E[fid]) * (int64_t)dim; + if (fid >= 0) { + src_[e] = src + (int64_t)dim * fid; + if constexpr (has_bias) { + bias_[e] = bias + __ldg(&f2E[fid]) * (int64_t)dim; + } + scale[e] = scales ? __ldg(&scales[e * tokens + ti]) : 1.f; } - scale[e] = scales ? __ldg(&scales[e * tokens + ti]) : 1.f; } using Vec = Array; @@ -1111,6 +1115,9 @@ __global__ void MoeReduceKernel(T* dst, // [ n, d] } PRAGMA_UNROLL for (int e = 0; e < exp_k; ++e) { + if (src_[e] == nullptr) { + continue; + } Vec v; Load(v, src_[e] + i); using namespace ops; diff --git a/src/turbomind/kernels/gemm/moe_utils_v2.h b/src/turbomind/kernels/gemm/moe_utils_v2.h index 15b1b2c889..3facd6f877 100644 --- a/src/turbomind/kernels/gemm/moe_utils_v2.h +++ b/src/turbomind/kernels/gemm/moe_utils_v2.h @@ -26,6 +26,8 @@ void invokeMoeGate_V2(int* f2n, int tokens_padded, int experts, int exp_per_tok, + int local_expert_offset, + int local_expert_num, bool softmax, bool norm_topk, float routed_scale, diff --git a/src/turbomind/models/llama/moe_ffn_layer.cc b/src/turbomind/models/llama/moe_ffn_layer.cc index b51ab89469..9c07764bf2 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.cc +++ b/src/turbomind/models/llama/moe_ffn_layer.cc @@ -8,7 +8,6 @@ #include "src/turbomind/core/context.h" #include "src/turbomind/core/scope.h" #include "src/turbomind/kernels/activation.h" -#include "src/turbomind/kernels/gemm/moe_ep_utils.h" #include "src/turbomind/kernels/gemm/moe_utils_v2.h" #include "src/turbomind/kernels/norm/rms_norm.h" @@ -193,6 +192,8 @@ void MoeFfnTpImpl::Forward(MoeFfnLayer::ForwardParam& p) padded, expert_num, p.weights->experts_per_token, + moe.local_expert_offset(), + moe.num_local_experts(), softmax, p.weights->norm_topk_prob, p.weights->routed_scale, @@ -286,42 +287,32 @@ class MoeFfnAgRsImpl final: public MoeFfnLayerImpl { void ForwardFused(MoeFfnLayer::ForwardParam& p); - const int ep_size_; - const int ep_rank_; - const int max_token_num_; - const std::string backend_; - comm::DeviceCommImpl* const d_comm_; - Allocator symm_allocator_; + const int ep_size_; + const int max_token_num_; bool initialized_ = false; - Buffer_ topk_weights_; - Buffer_ topk_idx_; - Buffer_ f2n_; - Buffer_ f2E_; - Buffer_ en2f_; - Buffer_ offsets_; - Buffer_ masks_; - Buffer_ accum_; + Buffer_ f2n_; + Buffer_ f2E_; + Buffer_ en2f_; + Buffer_ offsets_; + Buffer_ masks_; + Buffer_ accum_; + Buffer_ scales_; Tensor temp_; Tensor_ shared_scales_; }; MoeFfnAgRsImpl::MoeFfnAgRsImpl(const EngineParam& engine, const Context& ctx): - MoeFfnLayerImpl(ctx), - ep_size_(engine.ep_size), - ep_rank_(engine.ep_rank), - max_token_num_(engine.max_forward_token_num * engine.attn_dp_size), - backend_(engine.all2all_backend), - d_comm_(ctx.comm.d_comm), - symm_allocator_(GetSymmAllocator(ctx.comm.d_comm)) + MoeFfnLayerImpl(ctx), ep_size_(engine.ep_size), max_token_num_(engine.max_forward_token_num * engine.attn_dp_size) { } void MoeFfnAgRsImpl::Init(MoeFfnLayer::ForwardParam& p) { const auto& moe = *p.weights; + const int expert_num = moe.num_experts(); const int local_expert_num = moe.num_local_experts(); const int experts_per_token = moe.experts_per_token; const int pad_token_num = (max_token_num_ + kMoeGateVecSize - 1) / kMoeGateVecSize * kMoeGateVecSize; @@ -330,14 +321,13 @@ void MoeFfnAgRsImpl::Init(MoeFfnLayer::ForwardParam& p) << "MoeFfnAgRsImpl does not support topk_method=noaux_tc because it requires noaux-specific scoring and " "correction-bias handling"; - topk_weights_ = {max_token_num_ * experts_per_token, symm_allocator_}; - topk_idx_ = {max_token_num_ * experts_per_token, symm_allocator_}; - masks_ = {local_expert_num * pad_token_num, kDEVICE}; - f2n_ = {max_token_num_ * experts_per_token, kDEVICE}; - f2E_ = {max_token_num_ * experts_per_token, kDEVICE}; - en2f_ = {max_token_num_ * experts_per_token, kDEVICE}; - offsets_ = {local_expert_num + 1, kDEVICE}; - accum_ = {local_expert_num * kMoeGateMaxTiles, kDEVICE}; + masks_ = {expert_num * pad_token_num, kDEVICE}; + f2n_ = {experts_per_token * max_token_num_, kDEVICE}; + f2E_ = {experts_per_token * max_token_num_, kDEVICE}; + en2f_ = {experts_per_token * max_token_num_, kDEVICE}; + scales_ = {experts_per_token * max_token_num_, kDEVICE}; + offsets_ = {local_expert_num + 1, kDEVICE}; + accum_ = {expert_num * kMoeGateMaxTiles, kDEVICE}; initialized_ = true; } @@ -351,77 +341,51 @@ void MoeFfnAgRsImpl::Forward(MoeFfnLayer::ForwardParam& p) } const auto& moe = *p.weights; - const int tokens = p.global_hidden_states.shape(0); + const int tokens = p.input.shape(0); const int experts_per_token = moe.experts_per_token; const int expert_num = moe.num_experts(); - const int local_expert_num = moe.num_local_experts(); const int hidden_dim = TM_CHECK_NOTNULL(moe.block())->hidden_dim; const auto st = core::Context::stream().handle(); - TM_CHECK_LE(tokens, max_token_num_); - - std::vector topk_elem_counts(ep_size_); - size_t local_topk_offset{}; - for (int rank = 0; rank < ep_size_; ++rank) { - topk_elem_counts[rank] = TM_CHECK_NOTNULL(p.ep_elem_counts)->at(rank) / hidden_dim * experts_per_token; - local_topk_offset += rank < ep_rank_ ? topk_elem_counts[rank] : 0; - } - - auto local_topk_weights = topk_weights_.data() + local_topk_offset; - auto local_topk_idx = topk_idx_.data() + local_topk_offset; - - if (auto local_tokens = p.input.shape(0); local_tokens > 0) { - auto logits = Gate(p.input, *moe.gate.get()); - TM_DEBUG_TENSOR(logits, "logits", 2); - - bool softmax = true; - if (p.weights->topk_method == "group_limited_greedy") { - invokeMoeSoftmaxMaskTopKGroups( - logits.data(), local_tokens, expert_num, expert_num / p.weights->n_group, p.weights->topk_group, st); - TM_CUDA_CHECK(cudaGetLastError()); - softmax = false; - } + auto logits = Gate(p.input, *moe.gate.get()); + TM_DEBUG_TENSOR(logits, "logits", 2); - invokeMoeGateEp(local_topk_weights, - local_topk_idx, - logits.data(), - local_tokens, - expert_num, - experts_per_token, - softmax, - p.weights->norm_topk_prob, - p.weights->routed_scale, - st); + bool softmax = true; + if (p.weights->topk_method == "group_limited_greedy") { + invokeMoeSoftmaxMaskTopKGroups( + logits.data(), tokens, expert_num, expert_num / p.weights->n_group, p.weights->topk_group, st); TM_CUDA_CHECK(cudaGetLastError()); + softmax = false; } - d_comm_->GroupStart(); - d_comm_->AllGatherV( - p.input.raw_data(), p.global_hidden_states.raw_data(), *p.ep_elem_counts, p.input.dtype(), 0, st); - d_comm_->AllGatherV(local_topk_weights, topk_weights_.data(), topk_elem_counts, kFloat32, 0, st); - d_comm_->AllGatherV(local_topk_idx, topk_idx_.data(), topk_elem_counts, kInt64, 0, st); - d_comm_->GroupEnd(); - TM_CUDA_CHECK(cudaGetLastError()); + TM_CUDA_CHECK(cudaMemsetAsync(accum_.data(), 0, sizeof(int) * expert_num * kMoeGateMaxTiles, st)); + TM_CUDA_CHECK(cudaMemsetAsync(en2f_.data(), -1, sizeof(int) * tokens * experts_per_token, st)); - invokeMoeBuildAgRsRoutingMap(f2n_.data(), - f2E_.data(), - en2f_.data(), - offsets_.data(), - masks_.data(), - accum_.data(), - topk_idx_.data(), - tokens, - experts_per_token, - moe.local_expert_offset(), - local_expert_num, - st); + invokeMoeGate_V2(f2n_.data(), + f2E_.data(), + en2f_.data(), + offsets_.data(), + scales_.data(), + masks_.data(), + accum_.data(), + logits.data(), + tokens, + (tokens + kMoeGateVecSize - 1) / kMoeGateVecSize * kMoeGateVecSize, + expert_num, + experts_per_token, + moe.local_expert_offset(), + moe.num_local_experts(), + softmax, + p.weights->norm_topk_prob, + p.weights->routed_scale, + st); TM_CUDA_CHECK(cudaGetLastError()); temp_ = Tensor{{tokens * experts_per_token, hidden_dim}, p.input.dtype(), p.input.device()}; ForwardFused(p); shared_scales_ = {}; - if (moe.shared_gate && p.input.shape(0) > 0) { + if (moe.shared_gate) { shared_scales_ = Gate(p.input, *moe.shared_gate); } } @@ -435,7 +399,7 @@ void MoeFfnAgRsImpl::ForwardFused(MoeFfnLayer::ForwardParam& p) const int local_expert_num = moe.num_local_experts(); const auto st = core::Context::stream().handle(); - auto input = p.global_hidden_states; + auto input = p.input; const int* num_valid_tokens = offsets_.data() + local_expert_num; auto indices = f2n_.slice(0, temp_.shape(0)); auto offsets = offsets_.slice(0, local_expert_num + 1); @@ -470,30 +434,22 @@ void MoeFfnAgRsImpl::Combine(MoeFfnLayer::ForwardParam& p) { TM_FUNCTION_SCOPE(); - const auto& moe = *p.weights; - const auto& block = *TM_CHECK_NOTNULL(moe.block()); - const int experts_per_token = moe.experts_per_token; - const auto st = core::Context::stream().handle(); - - invokeMoeLocalCombineEp(p.global_hidden_states, - temp_, - block.w2->bias, - topk_weights_.data(), - en2f_.data(), - f2E_.data(), - experts_per_token, - st); - TM_CUDA_CHECK(cudaGetLastError()); + const auto& moe = *p.weights; + const auto& block = *TM_CHECK_NOTNULL(moe.block()); - d_comm_->ReduceScatterV( - p.global_hidden_states.raw_data(), p.output.raw_data(), *p.ep_elem_counts, p.output.dtype(), 0, st); + invokeMoeCombine(p.output, + temp_, + block.w2->bias, + scales_.data(), + en2f_.data(), + f2E_.data(), + shared_scales_.data_or((float*)nullptr), + moe.experts_per_token, + 1.f, + p.scale / ep_size_, + core::Context::stream().handle()); TM_CUDA_CHECK(cudaGetLastError()); - if (p.shared_output) { - invokeMoeCombineOutputEp(p.output, p.shared_output, shared_scales_.data_or((float*)nullptr), p.scale, st); - TM_CUDA_CHECK(cudaGetLastError()); - } - temp_ = {}; shared_scales_ = {}; } diff --git a/src/turbomind/models/llama/moe_ffn_layer.h b/src/turbomind/models/llama/moe_ffn_layer.h index f1c0371db9..39f68fa952 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.h +++ b/src/turbomind/models/llama/moe_ffn_layer.h @@ -3,7 +3,6 @@ #pragma once #include -#include #include "src/turbomind/core/core.h" #include "src/turbomind/models/llama/context.h" @@ -26,11 +25,6 @@ class MoeFfnLayer { const MoeWeight* weights; float scale; int layer_id; - - // EP allgather-reducescatter path only. - Tensor global_hidden_states; - const std::vector* ep_elem_counts{}; - Tensor shared_output; }; void Forward(ForwardParam& p); diff --git a/src/turbomind/models/llama/unified_decoder.cc b/src/turbomind/models/llama/unified_decoder.cc index a2dab7e9e0..86a7f81d38 100644 --- a/src/turbomind/models/llama/unified_decoder.cc +++ b/src/turbomind/models/llama/unified_decoder.cc @@ -1,8 +1,7 @@ -#include + #include #include -#include #include @@ -26,49 +25,6 @@ namespace turbomind { -struct FfnTokenPartition { - int first_token{}; - int token_count{}; - std::vector tp_elem_counts; - std::vector ep_elem_counts; - - static FfnTokenPartition - Make(const std::vector& local_token_nums, int dp_rank, int group_rank, int group_size, size_t hidden_units) - { - auto split_count = [group_size](int local_token_num, int rank) { - const int base = local_token_num / group_size; - const int remainder = local_token_num % group_size; - return base + (rank < remainder ? 1 : 0); - }; - - auto first_token = [group_size](int local_token_num, int rank) { - const int base = local_token_num / group_size; - const int remainder = local_token_num % group_size; - return rank * base + std::min(rank, remainder); - }; - - const int local_token_num = local_token_nums[dp_rank]; - - std::vector tp_elem_counts(group_size); - for (int rank = 0; rank < group_size; ++rank) { - tp_elem_counts[rank] = static_cast(split_count(local_token_num, rank)) * hidden_units; - } - - std::vector ep_elem_counts; - ep_elem_counts.reserve(local_token_nums.size() * group_size); - for (const int dp_token_num : local_token_nums) { - for (int rank = 0; rank < group_size; ++rank) { - ep_elem_counts.push_back(static_cast(split_count(dp_token_num, rank)) * hidden_units); - } - } - - return FfnTokenPartition{first_token(local_token_num, group_rank), - split_count(local_token_num, group_rank), - std::move(tp_elem_counts), - std::move(ep_elem_counts)}; - } -}; - void UnifiedDecoder::Run(BatchOp op, int phase, TensorMap& env) { attn_layer_->Run(op, phase, env); @@ -201,85 +157,6 @@ void UnifiedDecoder::AllreduceResidualRMSnorm(Tensor& hidden_states, } } -void UnifiedDecoder::ReduceScatterVResidualRMSnorm(Tensor& local_hidden_states, - Tensor& local_residual, - const Tensor& bias, - const Tensor& weight, - float eps, - const FfnTokenPartition& partition) -{ - TM_CHECK(ep_size_ > 1); - TM_CHECK(d_comm_); - TM_CHECK_EQ(local_residual.shape(0), local_hidden_states.shape(0)); - TM_CHECK_EQ(partition.tp_elem_counts.size(), static_cast(d_comm_->n_ranks(attn_tp_group_))); - TM_CHECK_EQ(std::accumulate(partition.tp_elem_counts.begin(), partition.tp_elem_counts.end(), size_t{}), - static_cast(local_hidden_states.size())); - - const auto dtype = local_hidden_states.dtype(); - const auto stream = core::Context::stream().handle(); - - Tensor hidden_slice = local_hidden_states.slice({partition.first_token, 0}, {partition.token_count, -1}); - Tensor residual_slice = local_residual.slice({partition.first_token, 0}, {partition.token_count, -1}); - - d_comm_->ReduceScatterV(local_hidden_states.data_or((void*)nullptr), - residual_slice.data_or((void*)nullptr), - partition.tp_elem_counts, - dtype, - attn_tp_group_, - stream); - TM_CUDA_CHECK(cudaGetLastError()); - - invokeResidualBiasRMSNorm(hidden_slice.data_or((void*)nullptr), - residual_slice.data_or((void*)nullptr), - weight.raw_data(), - bias.data_or((void*)nullptr), - dtype, - hidden_units_, - partition.token_count, - eps, - stream); - TM_CUDA_CHECK(cudaGetLastError()); -} - -void UnifiedDecoder::ResidualRMSnormAllGatherV(Tensor& local_hidden_states, - Tensor& local_residual, - const Tensor& weight, - float eps, - const FfnTokenPartition& partition) -{ - TM_CHECK(ep_size_ > 1); - TM_CHECK(d_comm_); - TM_CHECK_EQ(local_residual.shape(0), local_hidden_states.shape(0)); - TM_CHECK_EQ(partition.tp_elem_counts.size(), static_cast(d_comm_->n_ranks(attn_tp_group_))); - TM_CHECK_EQ(std::accumulate(partition.tp_elem_counts.begin(), partition.tp_elem_counts.end(), size_t{}), - static_cast(local_hidden_states.size())); - - const auto dtype = local_hidden_states.dtype(); - const auto stream = core::Context::stream().handle(); - - Tensor hidden_slice = local_hidden_states.slice({partition.first_token, 0}, {partition.token_count, -1}); - Tensor residual_slice = local_residual.slice({partition.first_token, 0}, {partition.token_count, -1}); - - invokeResidualBiasRMSNorm(hidden_slice.data_or((void*)nullptr), - residual_slice.data_or((void*)nullptr), - weight.raw_data(), - nullptr, - dtype, - hidden_units_, - partition.token_count, - eps, - stream); - TM_CUDA_CHECK(cudaGetLastError()); - - d_comm_->AllGatherV(hidden_slice.data_or((void*)nullptr), - local_hidden_states.data_or((void*)nullptr), - partition.tp_elem_counts, - dtype, - attn_tp_group_, - stream); - TM_CUDA_CHECK(cudaGetLastError()); -} - void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vector& weights) { TM_FUNCTION_SCOPE(); @@ -336,14 +213,6 @@ void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vector ffn_partition; - if (ep_size_ > 1) { - const int group_size = d_comm_->n_ranks(attn_tp_group_); - const int group_rank = d_comm_->rank(attn_tp_group_); - ffn_partition.emplace( - FfnTokenPartition::Make(local_token_nums, attn_dp_rank_, group_rank, group_size, hidden_units_)); - } - TM_LOG_DEBUG("local_token_num=%d, global_token_num=%d", (int)local_token_num, (int)global_token_num); TM_DEBUG_TENSOR(local_residual, "res", 1); @@ -401,30 +270,15 @@ void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vectorattention->wo->bias; } - Tensor ffn_hidden_states; - - if (ep_size_ > 1) { - ReduceScatterVResidualRMSnorm(local_hidden_states, - local_residual, - out_bias, - weights.at(layer)->ffn_norm->weight, - weights.at(layer)->ffn_norm->norm_eps_, - *ffn_partition); - ffn_hidden_states = - local_hidden_states.slice({ffn_partition->first_token, 0}, {ffn_partition->token_count, -1}); - } - else { - AllreduceResidualRMSnorm(global_hidden_states, - local_residual, - out_bias, - weights.at(layer)->ffn_norm->weight, - weights.at(layer)->ffn_norm->norm_eps_, - local_token_num, - attn_tp_group_, - 0, - local_token_nums.data()); - ffn_hidden_states = global_hidden_states; - } + AllreduceResidualRMSnorm(global_hidden_states, + local_residual, + out_bias, + weights.at(layer)->ffn_norm->weight, + weights.at(layer)->ffn_norm->norm_eps_, + local_token_num, + attn_tp_group_, + 0, + local_token_nums.data()); TM_DEBUG_TENSOR(local_residual, Concat("residual0", layer), 2); TM_DEBUG_TENSOR(local_hidden_states, Concat("norm1", layer), 2); @@ -435,52 +289,38 @@ void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vector moe_fwd_param; if (weights.at(layer)->moe_ffn) { - moe_fwd_param = MoeFfnLayer::ForwardParam{ffn_hidden_states, - ffn_hidden_states, + moe_fwd_param = MoeFfnLayer::ForwardParam{global_hidden_states, + global_hidden_states, weights.at(layer)->moe_ffn.get(), ffn_layer_ ? 1.f : 0.f, - layer, - ep_size_ > 1 ? global_hidden_states : Tensor{}, - ep_size_ > 1 ? &ffn_partition->ep_elem_counts : nullptr}; + layer}; moe_ffn_layer_->Forward(*moe_fwd_param); } - if (ffn_layer_ && weights.at(layer)->feed_forward && ffn_hidden_states.shape(0) > 0) { - auto ffn_output = ffn_hidden_states; - if (ep_size_ > 1) { - ffn_output = moe_fwd_param->shared_output = core::empty_like(ffn_hidden_states); - } - ffn_layer_->forward({ffn_hidden_states, ffn_output, weights.at(layer)->feed_forward.get(), (int)layer}); + if (ffn_layer_ && weights.at(layer)->feed_forward) { + ffn_layer_->forward( + {global_hidden_states, global_hidden_states, weights.at(layer)->feed_forward.get(), (int)layer}); } if (moe_fwd_param) { moe_ffn_layer_->Combine(*moe_fwd_param); } - TM_DEBUG_TENSOR(ffn_hidden_states, Concat("ffn_block", layer), 2); + TM_DEBUG_TENSOR(global_hidden_states, Concat("ffn_block", layer), 2); const bool last = layer == layer_num_ - 1; auto& scale_weight = !last ? weights.at(layer + 1)->attention_norm->weight : args.at("output_norm_weight"); - if (ep_size_ > 1) { - ResidualRMSnormAllGatherV(local_hidden_states, - local_residual, - scale_weight, - weights.at(layer)->ffn_norm->norm_eps_, - *ffn_partition); - } - else { - AllreduceResidualRMSnorm(global_hidden_states, - local_residual, - {}, - scale_weight, - weights.at(layer)->ffn_norm->norm_eps_, - local_token_num, - 0, - attn_tp_group_, - local_token_nums.data()); - } + AllreduceResidualRMSnorm(global_hidden_states, + local_residual, + {}, + scale_weight, + weights.at(layer)->ffn_norm->norm_eps_, + local_token_num, + 0, + attn_tp_group_, + local_token_nums.data()); TM_CUDA_CHECK(cudaGetLastError()); TM_DEBUG_TENSOR(local_residual, Concat("residual1", layer), 2); diff --git a/src/turbomind/models/llama/unified_decoder.h b/src/turbomind/models/llama/unified_decoder.h index 541866c99a..889a6dbee7 100644 --- a/src/turbomind/models/llama/unified_decoder.h +++ b/src/turbomind/models/llama/unified_decoder.h @@ -12,7 +12,6 @@ namespace turbomind { class ModelWeight; class DecoderLayerWeight; -struct FfnTokenPartition; class UnifiedDecoder { public: @@ -56,19 +55,6 @@ class UnifiedDecoder { int t0, int t1, const int* local_token_nums); - - void ReduceScatterVResidualRMSnorm(Tensor& local_hidden_states, - Tensor& local_residual, - const Tensor& bias, - const Tensor& weight, - float eps, - const FfnTokenPartition& partition); - - void ResidualRMSnormAllGatherV(Tensor& local_hidden_states, - Tensor& local_residual, - const Tensor& weight, - float eps, - const FfnTokenPartition& partition); }; } // namespace turbomind From 765caee5bb82e4ede48cd815f3730371643e9997 Mon Sep 17 00:00:00 2001 From: irexyc Date: Wed, 8 Jul 2026 14:23:11 +0000 Subject: [PATCH 09/19] merge implement of MoeFfnTpImpl and MoeFfnAgRsImpl --- src/turbomind/kernels/gemm/moe_utils_v2.cu | 18 +- src/turbomind/kernels/gemm/moe_utils_v2.h | 2 + src/turbomind/models/llama/moe_ffn_layer.cc | 257 +++++++++++++++++++- 3 files changed, 262 insertions(+), 15 deletions(-) diff --git a/src/turbomind/kernels/gemm/moe_utils_v2.cu b/src/turbomind/kernels/gemm/moe_utils_v2.cu index 384845eca4..84de014ae2 100644 --- a/src/turbomind/kernels/gemm/moe_utils_v2.cu +++ b/src/turbomind/kernels/gemm/moe_utils_v2.cu @@ -859,6 +859,8 @@ void invokeMoeGate_NoAuxTC(int* f2n, int tokens_padded, int experts, int exp_per_tok, + int local_expert_offset, + int local_expert_num, bool norm_topk_prob, float routed_scale, bool use_sigmoid, @@ -902,9 +904,19 @@ void invokeMoeGate_NoAuxTC(int* f2n, use_sigmoid); constexpr int scan_threads = (1 << base_log_tile) / kMoeGateVecSize; - const dim3 scan_blocks(tiles, experts + 1); - MoeScanKernel_v2<<>>( - f2n, f2E, en2f, offsets, (int8_t*)masks, accum, log_tile, tiles, tokens, tokens_padded, experts); + const dim3 scan_blocks(tiles, local_expert_num + 1); + MoeScanKernel_v2 + <<>>(f2n, + f2E, + en2f, + offsets, + (int8_t*)masks + local_expert_offset * tokens_padded, + accum + local_expert_offset * tiles, + log_tile, + tiles, + tokens, + tokens_padded, + local_expert_num); TM_CUDA_CHECK(cudaGetLastError()); } diff --git a/src/turbomind/kernels/gemm/moe_utils_v2.h b/src/turbomind/kernels/gemm/moe_utils_v2.h index 3facd6f877..5031febe3c 100644 --- a/src/turbomind/kernels/gemm/moe_utils_v2.h +++ b/src/turbomind/kernels/gemm/moe_utils_v2.h @@ -94,6 +94,8 @@ void invokeMoeGate_NoAuxTC(int* f2n, int tokens_padded, int experts, int exp_per_tok, + int local_expert_offset, + int local_expert_num, bool norm_topk_prob, float routed_scale, bool use_sigmoid, diff --git a/src/turbomind/models/llama/moe_ffn_layer.cc b/src/turbomind/models/llama/moe_ffn_layer.cc index 9c07764bf2..a10fe7e0db 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.cc +++ b/src/turbomind/models/llama/moe_ffn_layer.cc @@ -1,7 +1,6 @@ // Copyright (c) OpenMMLab. All rights reserved. #include -#include #include @@ -55,8 +54,6 @@ class MoeFfnTpImpl final: public MoeFfnLayerImpl { const int max_token_num_; int& is_warm_up_; - std::unique_ptr expert_ffn_; - bool initialized_ = false; Buffer_ h_offsets_; @@ -77,8 +74,7 @@ MoeFfnTpImpl::MoeFfnTpImpl(const EngineParam& engine, const Context& ctx): MoeFfnLayerImpl(ctx), tp_size_(engine.mlp_tp_size), max_token_num_(engine.max_forward_token_num * engine.attn_dp_size), - is_warm_up_(*ctx.is_warm_up), - expert_ffn_(std::make_unique(ctx)) + is_warm_up_(*ctx.is_warm_up) { } @@ -162,6 +158,8 @@ void MoeFfnTpImpl::Forward(MoeFfnLayer::ForwardParam& p) padded, expert_num, p.weights->experts_per_token, + moe.local_expert_offset(), + moe.num_local_experts(), p.weights->norm_topk_prob, p.weights->routed_scale, p.weights->scoring_func == "sigmoid", @@ -454,19 +452,254 @@ void MoeFfnAgRsImpl::Combine(MoeFfnLayer::ForwardParam& p) shared_scales_ = {}; } -MoeFfnLayer::MoeFfnLayer(const EngineParam& engine, const Context& ctx) +class MoeFfnDefaultImpl final: public MoeFfnLayerImpl { +public: + MoeFfnDefaultImpl(const EngineParam& engine, const Context& ctx); + + void Forward(MoeFfnLayer::ForwardParam& p) override; + + void Combine(MoeFfnLayer::ForwardParam& p) override; + +private: + void Init(MoeFfnLayer::ForwardParam& p); + + const int tp_size_; + const int ep_size_; + const int max_token_num_; + int& is_warm_up_; + + bool initialized_ = false; + + Buffer_ h_offsets_; + + Buffer_ masks_; + Buffer_ f2n_; + Buffer_ f2E_; + Buffer_ en2f_; + Buffer_ scales_; + Buffer_ accum_; + Buffer_ offsets_; + + Tensor temp_; + Tensor_ shared_scales_; +}; + +MoeFfnDefaultImpl::MoeFfnDefaultImpl(const EngineParam& engine, const Context& ctx): + MoeFfnLayerImpl(ctx), + tp_size_(engine.mlp_tp_size), + ep_size_(engine.ep_size), + max_token_num_(engine.max_forward_token_num * engine.attn_dp_size), + is_warm_up_(*ctx.is_warm_up) { - if (engine.ep_size <= 1) { - impl_ = std::make_unique(engine, ctx); - return; +} + +void MoeFfnDefaultImpl::Init(MoeFfnLayer::ForwardParam& p) +{ + const int expert_num = p.weights->num_experts(); + const int local_expert_num = p.weights->num_local_experts(); + const int experts_per_token = p.weights->experts_per_token; + + h_offsets_ = {local_expert_num + 1, kCPU}; + + const int pad_token_num = (max_token_num_ + kMoeGateVecSize - 1) / kMoeGateVecSize * kMoeGateVecSize; + + masks_ = {expert_num * pad_token_num, kDEVICE}; + f2n_ = {experts_per_token * max_token_num_, kDEVICE}; + f2E_ = {experts_per_token * max_token_num_, kDEVICE}; + en2f_ = {experts_per_token * max_token_num_, kDEVICE}; + scales_ = {experts_per_token * max_token_num_, kDEVICE}; + offsets_ = {local_expert_num + 1, kDEVICE}; + accum_ = {expert_num * kMoeGateMaxTiles, kDEVICE}; + + initialized_ = true; +} + +void MoeFfnDefaultImpl::Forward(MoeFfnLayer::ForwardParam& p) +{ + TM_FUNCTION_SCOPE(); + if (!initialized_) { + Init(p); + } + + const int tokens = p.input.shape(0); + const auto& moe = *p.weights; + + const auto& block = *TM_CHECK_NOTNULL(moe.block()); + + const int hidden_dim = block.hidden_dim; + const int inter_size = block.inter_size; + + const size_t padded = (tokens + kMoeGateVecSize - 1) / kMoeGateVecSize * kMoeGateVecSize; + + const int expert_num = moe.num_experts(); + const int local_expert_num = moe.num_local_experts(); + const int expert_offset = moe.local_expert_offset(); + const int experts_per_token = moe.experts_per_token; + + TM_CHECK(expert_num); + + auto logits = Gate(p.input, *moe.gate.get()); + + TM_DEBUG_TENSOR(logits, "logits", 2); + + const auto st = core::Context::stream().handle(); + + if (ep_size_ > 1) { + TM_CUDA_CHECK(cudaMemsetAsync(en2f_.data(), -1, sizeof(int) * tokens * experts_per_token, st)); } - if (engine.all2all_backend == "allgather_reducescatter") { - impl_ = std::make_unique(engine, ctx); + if (p.weights->topk_method == "noaux_tc") { + TM_CHECK_EQ(p.weights->n_group, 1); + TM_CHECK_EQ(p.weights->topk_group, 1); + const float* correction_bias = nullptr; + if (moe.score_correction_bias) { + correction_bias = moe.score_correction_bias.size() > 0 ? moe.score_correction_bias.data() : nullptr; + } + invokeMoeGate_NoAuxTC(f2n_.data(), + f2E_.data(), + en2f_.data(), + offsets_.data(), + scales_.data(), + masks_.data(), + accum_.data(), + logits.data(), + correction_bias, + tokens, + padded, + expert_num, + experts_per_token, + expert_offset, + local_expert_num, + p.weights->norm_topk_prob, + p.weights->routed_scale, + p.weights->scoring_func == "sigmoid", + st); + } + else { + // V2: accum must be cleared by caller; masks cleared internally + TM_CUDA_CHECK(cudaMemsetAsync(accum_.data(), 0, sizeof(int) * expert_num * kMoeGateMaxTiles, st)); + + bool softmax = true; + if (p.weights->topk_method == "group_limited_greedy") { + invokeMoeSoftmaxMaskTopKGroups( + logits.data(), tokens, expert_num, expert_num / p.weights->n_group, p.weights->topk_group, st); + TM_CUDA_CHECK(cudaGetLastError()); + softmax = false; + } + + /// TODO: fix illegal memory access even if NaN are present in logits + invokeMoeGate_V2(f2n_.data(), + f2E_.data(), + en2f_.data(), + offsets_.data(), + scales_.data(), + masks_.data(), + accum_.data(), + logits.data(), + tokens, + padded, + expert_num, + experts_per_token, + expert_offset, + local_expert_num, + softmax, + p.weights->norm_topk_prob, + p.weights->routed_scale, + st); + } + TM_CUDA_CHECK(cudaGetLastError()); + + if (is_warm_up_) { + std::mt19937 g; + const auto expert_ids = SampleUniform(tokens, local_expert_num, experts_per_token, g); + std::vector cnt(local_expert_num); + for (const auto& x : expert_ids) { + ++cnt[x]; + } + h_offsets_[0] = 0; + for (int i = 0; i < local_expert_num; ++i) { + h_offsets_[i + 1] = h_offsets_[i] + cnt[i]; + } + TM_CUDA_CHECK(cudaMemcpyAsync( + offsets_.data(), h_offsets_.data(), sizeof(int) * (local_expert_num + 1), cudaMemcpyDefault, st)); + + if (ep_size_ > 1) { + const auto entries = static_cast(tokens) * experts_per_token; + TM_CUDA_CHECK(cudaMemsetAsync(f2n_.data(), 0, sizeof(int) * entries, st)); + TM_CUDA_CHECK(cudaMemsetAsync(f2E_.data(), 0, sizeof(int) * entries, st)); + TM_CUDA_CHECK(cudaMemsetAsync(en2f_.data(), -1, sizeof(int) * entries, st)); + } + } + + temp_ = Tensor{{tokens * experts_per_token, hidden_dim}, p.input.dtype(), p.input.device()}; + + // For ep_size > 1, the valid tokens are less than tokens * experts_per_token + const bool indices_padded = ep_size_ > 1 ? true : false; + const int* num_valid_tokens = ep_size_ > 1 ? offsets_.data() + local_expert_num : nullptr; + + auto indices = f2n_.slice(0, temp_.shape(0)); + auto offsets = offsets_.slice(0, local_expert_num + 1); + + if (block.w1w3) { + Tensor inter; + TM_SCOPE_CALL(linear_.Forward(p.input, *block.w1w3, indices, offsets, inter, indices_padded)); + + if (!block.is_fused_silu) { + Activation(inter, block.w1w3->bias, f2E_, block.act_type, num_valid_tokens, st); + TM_CUDA_CHECK(cudaGetLastError()); + } + + TM_SCOPE_CALL(linear_.Forward(inter.slice({0, 0}, {-1, inter_size}), *block.w2, {}, offsets, temp_)); + } + else { + // Separate w1/w3 path + Tensor gating; + TM_SCOPE_CALL(linear_.Forward(p.input, *block.w1, indices, offsets, gating, indices_padded)); + + Tensor up; + TM_SCOPE_CALL(linear_.Forward(p.input, *block.w3, indices, offsets, up, indices_padded)); + + Activation(gating, up, block.act_type, num_valid_tokens, st); + TM_CUDA_CHECK(cudaGetLastError()); + + TM_SCOPE_CALL(linear_.Forward(gating, *block.w2, {}, offsets, temp_)); + } + + if (moe.shared_gate) { + shared_scales_ = Gate(p.input, *moe.shared_gate); + } +} + +void MoeFfnDefaultImpl::Combine(MoeFfnLayer::ForwardParam& p) +{ + TM_FUNCTION_SCOPE(); + auto& moe = *p.weights; + + invokeMoeCombine(p.output, + temp_, + TM_CHECK_NOTNULL(moe.block())->w2->bias, + scales_.data(), + en2f_.data(), + f2E_.data(), + shared_scales_.data_or((float*)nullptr), + moe.experts_per_token, + 1.f / tp_size_, + p.scale / ep_size_, + core::Context::stream().handle()); + TM_CUDA_CHECK(cudaGetLastError()); + + temp_ = {}; + shared_scales_ = {}; +} + +MoeFfnLayer::MoeFfnLayer(const EngineParam& engine, const Context& ctx) +{ + if (engine.ep_size <= 1 || engine.nnodes == 1) { + impl_ = std::make_unique(engine, ctx); return; } - TM_LOG_FATAL("Unsupported MoE EP all2all backend: {}", engine.all2all_backend); + TM_LOG_FATAL("Unsupported config for MoeFfnLayer"); } MoeFfnLayer::~MoeFfnLayer() = default; From 0ee86ecb95b3d944eef4252508d13fb918e61120 Mon Sep 17 00:00:00 2001 From: irexyc Date: Thu, 9 Jul 2026 03:39:42 +0000 Subject: [PATCH 10/19] remove unused --- lmdeploy/messages.py | 1 - lmdeploy/turbomind/turbomind.py | 1 - src/turbomind/comm/device_comm.h | 3 - src/turbomind/comm/nccl/nccl.cu | 10 - src/turbomind/engine/engine_config.h | 1 - src/turbomind/kernels/gemm/CMakeLists.txt | 1 - src/turbomind/kernels/gemm/moe_ep_utils.cu | 552 ------------------ src/turbomind/kernels/gemm/moe_ep_utils.h | 47 -- src/turbomind/models/llama/moe_ffn_layer.cc | 401 ------------- src/turbomind/models/llama/unified_decoder.cc | 1 - src/turbomind/models/llama/unified_decoder.h | 1 - 11 files changed, 1019 deletions(-) delete mode 100644 src/turbomind/kernels/gemm/moe_ep_utils.cu delete mode 100644 src/turbomind/kernels/gemm/moe_ep_utils.h diff --git a/lmdeploy/messages.py b/lmdeploy/messages.py index fd71cefdce..17873b218f 100644 --- a/lmdeploy/messages.py +++ b/lmdeploy/messages.py @@ -282,7 +282,6 @@ class TurbomindEngineConfig: dp: int = 1 cp: int = 1 ep: int = 1 - all2all_backend: str = 'allgather_reducescatter' device_num: int = None attn_tp_size: int = None attn_cp_size: int = None diff --git a/lmdeploy/turbomind/turbomind.py b/lmdeploy/turbomind/turbomind.py index 9fcf75cb30..d0dd35bfe3 100644 --- a/lmdeploy/turbomind/turbomind.py +++ b/lmdeploy/turbomind/turbomind.py @@ -251,7 +251,6 @@ def _from_hf(self, model_path: str, engine_config: TurbomindEngineConfig, ec.attn_cp_size = engine_config.attn_cp_size ec.mlp_tp_size = engine_config.mlp_tp_size ec.ep_size = engine_config.ep - ec.all2all_backend = engine_config.all2all_backend ec.devices = engine_config.devices ec.nnodes = engine_config.nnodes ec.node_rank = engine_config.node_rank diff --git a/src/turbomind/comm/device_comm.h b/src/turbomind/comm/device_comm.h index ba4d14c0da..9ce0be0d12 100644 --- a/src/turbomind/comm/device_comm.h +++ b/src/turbomind/comm/device_comm.h @@ -138,9 +138,6 @@ class DeviceCommImpl { { throw std::runtime_error("not implemented"); } - - virtual void GroupStart() {} - virtual void GroupEnd() {} }; class DeviceComm { diff --git a/src/turbomind/comm/nccl/nccl.cu b/src/turbomind/comm/nccl/nccl.cu index f481378677..a7f843b657 100644 --- a/src/turbomind/comm/nccl/nccl.cu +++ b/src/turbomind/comm/nccl/nccl.cu @@ -477,16 +477,6 @@ public: NCCLCHECK(ncclBroadcast(sendbuff, recvbuff, count, to_nccl_dtype(type), root, groups_.at(group), stream)); } - void GroupStart() override - { - NCCLCHECK(ncclGroupStart()); - } - - void GroupEnd() override - { - NCCLCHECK(ncclGroupEnd()); - } - private: HostComm h_comm_; diff --git a/src/turbomind/engine/engine_config.h b/src/turbomind/engine/engine_config.h index 04aed14892..efd126d52a 100644 --- a/src/turbomind/engine/engine_config.h +++ b/src/turbomind/engine/engine_config.h @@ -33,7 +33,6 @@ struct EngineConfig { X(int, attn_cp_size) \ X(int, mlp_tp_size) \ X(int, ep_size, 1) \ - X(std::string, all2all_backend, "allgather_reducescatter") \ X(std::vector, devices) \ X(int, nnodes) \ X(int, node_rank) \ diff --git a/src/turbomind/kernels/gemm/CMakeLists.txt b/src/turbomind/kernels/gemm/CMakeLists.txt index 31ce211b1c..7a1b210762 100644 --- a/src/turbomind/kernels/gemm/CMakeLists.txt +++ b/src/turbomind/kernels/gemm/CMakeLists.txt @@ -60,7 +60,6 @@ add_library(gemm2 ${GEMM2_KERNELS_SM80} cublas.cu moe_utils_v2.cu - moe_ep_utils.cu test/test_utils.cu ) diff --git a/src/turbomind/kernels/gemm/moe_ep_utils.cu b/src/turbomind/kernels/gemm/moe_ep_utils.cu deleted file mode 100644 index 330f55d7ee..0000000000 --- a/src/turbomind/kernels/gemm/moe_ep_utils.cu +++ /dev/null @@ -1,552 +0,0 @@ -// Copyright (c) OpenMMLab. All rights reserved. - -#include "src/turbomind/kernels/gemm/moe_ep_utils.h" - -#include "src/turbomind/core/check.h" -#include "src/turbomind/core/data_type.h" -#include "src/turbomind/kernels/core/array_ops.h" -#include "src/turbomind/kernels/core/common.h" -#include "src/turbomind/kernels/core/math.h" -#include "src/turbomind/kernels/gemm/moe_utils_v2.h" -#include "src/turbomind/utils/cuda_utils.h" - -#include -#include -#include - -namespace turbomind { - -template -__global__ void MoeGateKernel(float* topk_weights, // [n, topk] - int64_t* topk_idx, // [n, topk] - const float* logits, // [n,E] - int token_num, - int expert_num, - int top_k, - bool softmax, - bool norm_topk, - float routed_scale) -{ - constexpr int threads_per_token = max_expert_num / items_per_thread; - - static_assert(items_per_thread <= 32); - static_assert(threads_per_token <= 32); - static_assert((threads_per_token & (threads_per_token - 1)) == 0); - - const int thread_idx = threadIdx.x + blockIdx.x * blockDim.x; - - const int ti = thread_idx / threads_per_token; - const int ei = thread_idx % threads_per_token; - - const int warp_ti = threadIdx.x % WARP_SIZE / threads_per_token; - - float data[items_per_thread]; - int idxs[items_per_thread]; - - PRAGMA_UNROLL - for (int i = 0; i < items_per_thread; ++i) { - data[i] = -std::numeric_limits::infinity(); - idxs[i] = ei * items_per_thread + i; - } - if (ti < token_num) { - PRAGMA_UNROLL - for (int i = 0; i < items_per_thread; i += access_size) { - const int e = ei * items_per_thread + i; - if (e + access_size <= expert_num) { - Ldg((Array&)data[i], &logits[ti * expert_num + e]); - } - else { - PRAGMA_UNROLL - for (int j = 0; j < access_size; ++j) { - if (e + j < expert_num) { - data[i + j] = logits[ti * expert_num + e + j]; - } - } - } - } - } - - unsigned mask = (unsigned)-1; - float max_logit; - - const int warp_ti_offset = warp_ti * threads_per_token; - - int sel_item[max_top_k]; - - auto run = [&](int k) { - unsigned bit = 1; - unsigned max_bit = 0; - float max_val = -std::numeric_limits::infinity(); - PRAGMA_UNROLL - for (int i = 0; i < items_per_thread; ++i) { - if ((mask & bit) && data[i] > max_val) { - max_bit = bit; - max_val = data[i]; - } - asm("shl.b32 %0, %1, 1;\n" : "=r"(bit) : "r"(bit)); - } - - int g_max_ei = ei; - float g_max_val = max_val; - if constexpr (threads_per_token > 1) { - PRAGMA_UNROLL - for (int m = threads_per_token / 2; m >= 1; m /= 2) { - g_max_val = fmaxf(g_max_val, __shfl_xor_sync((uint32_t)-1, g_max_val, m)); - } - const auto active = __ballot_sync((uint32_t)-1, max_val == g_max_val); - g_max_ei = __ffs(active >> (unsigned)warp_ti_offset) - 1; - } - if (k == 0) { - max_logit = g_max_val; - } - int local_item = -1; - if (ei == g_max_ei) { - local_item = __ffs(max_bit) - 1; - mask -= max_bit; - } - sel_item[k] = local_item; - }; - - run(0); - - for (int k = 1; k < top_k; ++k) { - run(k); - } - - mask = ~mask; - - int used[items_per_thread]; - { - unsigned bit = 1; - PRAGMA_UNROLL - for (int i = 0; i < items_per_thread; ++i) { - used[i] = (mask & bit) > 0; - asm("shl.b32 %0, %1, 1;\n" : "=r"(bit) : "r"(bit)); - } - } - - float sum_prob{}; - - if (softmax) { - PRAGMA_UNROLL - for (int i = 0; i < items_per_thread; ++i) { - if (!norm_topk || used[i]) { - data[i] = expf(data[i] - max_logit); - sum_prob += data[i]; - } - } - PRAGMA_UNROLL - for (int m = threads_per_token / 2; m >= 1; m /= 2) { - sum_prob += __shfl_xor_sync((uint32_t)-1, sum_prob, m); - } - sum_prob = fdividef(1.f, sum_prob); - } - else { - sum_prob = 1.f; - } - - if (ti < token_num) { - PRAGMA_UNROLL - for (int k = 0; k < max_top_k; ++k) { - if (k < top_k && sel_item[k] >= 0) { - const int i = sel_item[k]; - topk_weights[ti * top_k + k] = data[i] * sum_prob * routed_scale; - topk_idx[ti * top_k + k] = idxs[i]; - } - } - } -} - -template -inline constexpr std::integral_constant _Int{}; - -void invokeMoeGateEp(float* topk_weights, - int64_t* topk_idx, - const float* logits, - int tokens, - int experts, - int experts_per_token, - bool softmax, - bool norm_topk, - float routed_scale, - cudaStream_t stream) -{ - if (tokens == 0) { - return; - } - - auto invoke = [&](auto max_expert_num, auto top_k, auto items_per_thread, auto vec_size) { - constexpr int thrs_per_tok = max_expert_num.value / items_per_thread.value; - constexpr int threads = 256; - const int blocks = ceil_div(tokens, threads / thrs_per_tok); - - MoeGateKernel - <<>>( - topk_weights, topk_idx, logits, tokens, experts, experts_per_token, softmax, norm_topk, routed_scale); - TM_CUDA_CHECK(cudaGetLastError()); - - return true; - }; - - if (!softmax && norm_topk) { - TM_CHECK(0) << softmax << " " << norm_topk; - } - - auto dispatch = [&] { - if (experts <= 8) { - if (experts_per_token <= 2) { - return invoke(_Int<8>, _Int<2>, _Int<8>, _Int<4>); - } - else { - return invoke(_Int<8>, _Int<8>, _Int<8>, _Int<4>); - } - } - else if (experts <= 64) { - if (experts_per_token <= 4) { - return invoke(_Int<64>, _Int<4>, _Int<16>, _Int<4>); - } - else if (experts_per_token <= 8) { - return invoke(_Int<64>, _Int<8>, _Int<16>, _Int<4>); - } - } - else if (experts <= 128) { - if (experts_per_token <= 8) { - return invoke(_Int<128>, _Int<8>, _Int<16>, _Int<4>); - } - } - else if (experts <= 160) { - if (experts_per_token <= 8) { - return invoke(_Int<160>, _Int<8>, _Int<10>, _Int<2>); - } - } - else if (experts <= 512) { - if (experts_per_token <= 10) { - return invoke(_Int<512>, _Int<10>, _Int<16>, _Int<4>); - } - } - return false; - }; - - auto success = dispatch(); - - TM_CHECK(success) << "unsupported moe config: expert_num=" << experts << ", top_k=" << experts_per_token - << ", softmax=" << softmax << ", norm_topk=" << norm_topk; -} - -template -__global__ void MoeAgRsBuildMasksKernel(int8_t* masks, // [num_local_experts, tokens_padded], pre-set to -1 - int* accum, // [num_local_experts, tiles], pre-set to 0 - const int64_t* topk_idx, // [tokens, topk] - int tokens, - int tokens_padded, - int topk, - int local_expert_offset, - int num_local_experts, - int log_tile, - int tiles) -{ - __shared__ int shared_accum[kMoeGateMaxTiles][max_local_experts + 1]; - - // Clear shared_accum - for (int i = threadIdx.x; i < tiles * num_local_experts; i += block_dim) { - shared_accum[i / num_local_experts][i % num_local_experts] = 0; - } - __syncthreads(); - - // Each thread processes one token - const int ti = threadIdx.x + blockIdx.x * blockDim.x; - if (ti < tokens) { - for (int k = 0; k < topk; ++k) { - const int global_eid = static_cast(topk_idx[ti * topk + k]); - const int local_eid = global_eid - local_expert_offset; - if (0 <= local_eid && local_eid < num_local_experts) { - masks[local_eid * tokens_padded + ti] = static_cast(k); - atomicAdd(&shared_accum[ti >> log_tile][local_eid], 1); - } - } - } - __syncthreads(); - - // Flush shared_accum to global accum - for (int i = threadIdx.x; i < tiles * num_local_experts; i += block_dim) { - int t = i / num_local_experts; - int e = i % num_local_experts; - if (t < tiles && shared_accum[t][e]) { - atomicAdd(&accum[e * tiles + t], shared_accum[t][e]); - } - } -} - -void invokeMoeBuildAgRsRoutingMap(int* f2n, - int* f2E, - int* en2f, - int* offsets, - int8_t* masks, - int* accum, - const int64_t* topk_idx, - int tokens, - int topk, - int local_expert_offset, - int num_local_experts, - cudaStream_t stream) -{ - if (tokens == 0) { - TM_CUDA_CHECK(cudaMemsetAsync(offsets, 0, sizeof(int) * (num_local_experts + 1), stream)); - return; - } - - const int tokens_padded = round_up(tokens, kMoeGateVecSize); - - // Compute log_tile / tiles (must match invokeMoeScanKernel's internal computation) - constexpr int base_log_tile = 9; - int log_tile = base_log_tile; - while (((tokens_padded + (1 << log_tile) - 1) >> log_tile) > kMoeGateMaxTiles) { - ++log_tile; - } - const int tiles = ceil_div(tokens_padded, 1 << log_tile); - - // Initialise buffers - TM_CUDA_CHECK(cudaMemsetAsync(masks, -1, sizeof(int8_t) * num_local_experts * tokens_padded, stream)); - TM_CUDA_CHECK(cudaMemsetAsync(accum, 0, sizeof(int) * num_local_experts * kMoeGateMaxTiles, stream)); - TM_CUDA_CHECK(cudaMemsetAsync(en2f, -1, sizeof(int) * tokens * topk, stream)); - - constexpr int block = 256; - const int grid = cdiv(tokens_padded, block); - - auto invoke = [&](auto max_local_experts) { - MoeAgRsBuildMasksKernel<<>>(masks, - accum, - topk_idx, - tokens, - tokens_padded, - topk, - local_expert_offset, - num_local_experts, - log_tile, - tiles); - return true; - }; - - auto dispatch = [&] { - if (num_local_experts <= 4) { - return invoke(_Int<4>); - } - else if (num_local_experts <= 8) { - return invoke(_Int<8>); - } - else if (num_local_experts <= 16) { - return invoke(_Int<16>); - } - else if (num_local_experts <= 32) { - return invoke(_Int<32>); - } - else if (num_local_experts <= 64) { - return invoke(_Int<64>); - } - else if (num_local_experts <= 128) { - return invoke(_Int<128>); - } - return false; - }; - - auto success = dispatch(); - TM_CHECK(success) << "unsupported num_local_experts: " << num_local_experts; - TM_CUDA_CHECK(cudaGetLastError()); - - invokeMoeScanKernel(f2n, f2E, en2f, offsets, masks, accum, tokens, tokens_padded, num_local_experts, stream); -} - -template -__global__ void MoeCombineKernel(T* dst, // [num_tokens, dim] - const T* src, // [expert_token_num, dim] - const T* bias, // [num_local_experts, dim] - const float* topk_weights, // [num_tokens, topk] - const int* en2f, // [topk, num_tokens] - const int* f2E, // [expert_token_num] - int dim, - int tokens) -{ - if constexpr (TURBOMIND_ARCH_DTYPE_GUARD(data_type_v)) { - for (int ti = blockIdx.x; ti < tokens; ti += gridDim.x) { - T* dst_row = dst + (int64_t)dim * ti; - - const T* src_[exp_k]{}; - const T* bias_[exp_k]{}; - float weight[exp_k]{}; - - PRAGMA_UNROLL - for (int e = 0; e < exp_k; ++e) { - const int fid = __ldg(&en2f[e * tokens + ti]); - if (fid >= 0) { - src_[e] = src + (int64_t)dim * fid; - weight[e] = __ldg(&topk_weights[ti * exp_k + e]); - if constexpr (has_bias) { - bias_[e] = bias + (int64_t)dim * __ldg(&f2E[fid]); - } - } - } - - using Vec = Array; - - for (int i = threadIdx.x * vec_size; i < dim; i += block_dim * vec_size) { - Array accum{}; - PRAGMA_UNROLL - for (int e = 0; e < exp_k; ++e) { - if (src_[e] == nullptr) { - continue; - } - Vec v; - Load(v, src_[e] + i); - if constexpr (has_bias) { - Vec b; - Load(b, bias_[e] + i); - PRAGMA_UNROLL - for (int j = 0; j < vec_size; ++j) { - v[j] = (T)((float)v[j] + (float)b[j]); - } - } - using namespace ops; - const auto x = cast(v) * weight[e]; - accum = accum + x; - } - Store(&dst_row[i], cast(accum)); - } - } - } -} - -void invokeMoeLocalCombineEp(Ref out_, - const Tensor& src, - const Tensor& bias, - const float* topk_weights, - const int* en2f, - const int* f2E, - int experts_per_token, - cudaStream_t st) -{ - auto& out = out_.get(); - - const int tokens = out.shape(0); - - if (tokens == 0) { - return; - } - - const int dim = src.shape(1); - - auto invoke = [&](auto t, auto e, auto has_bias_) { - using T = decltype(t); - constexpr int threads = 256; - constexpr int vec_size = sizeof(uint4) / sizeof(T); - constexpr int exp_per_tok = decltype(e)::value; - constexpr bool has_bias = decltype(has_bias_)::value; - - static const int sm_count = getSMCount(); - const int grid = std::min(tokens, sm_count * 4); - - MoeCombineKernel<<>>( - out.data(), src.data(), bias.data_or((T*)nullptr), topk_weights, en2f, f2E, dim, tokens); - TM_CUDA_CHECK(cudaGetLastError()); - }; - - auto dispatch_topk = [&](auto has_bias, auto t) { - switch (experts_per_token) { - case 1: - return invoke(t, std::integral_constant{}, has_bias); - case 2: - return invoke(t, std::integral_constant{}, has_bias); - case 4: - return invoke(t, std::integral_constant{}, has_bias); - case 6: - return invoke(t, std::integral_constant{}, has_bias); - case 8: - return invoke(t, std::integral_constant{}, has_bias); - case 10: - return invoke(t, std::integral_constant{}, has_bias); - default: - TM_CHECK(0) << "unsupported experts_per_token " << experts_per_token; - } - }; - - auto dispatch_dtype = [&](auto t) { - if (bias) { - TM_CHECK_NOTNULL(f2E); - return dispatch_topk(std::true_type{}, t); - } - else { - return dispatch_topk(std::false_type{}, t); - } - }; - - TM_DISPATCH_PRIMARY_DTYPES(src.dtype(), dispatch_dtype); -} - -template -__global__ void MoeCombineOutputEpKernel(T* dst, // [tokens, dim] - const T* shared, // [tokens, dim] - const float* shared_scales, // [tokens] or nullptr - int dim, - float scale) -{ - if constexpr (TURBOMIND_ARCH_DTYPE_GUARD(data_type_v)) { - const int ti = blockIdx.x; - - float dst_scale = scale; - if (shared_scales) { - dst_scale = __ldg(&shared_scales[ti]); - dst_scale = fdividef(1.f, 1.f + expf(-dst_scale)); - } - - dst += (int64_t)dim * ti; - shared += (int64_t)dim * ti; - - using Vec = Array; - - for (int i = threadIdx.x * vec_size; i < dim; i += block_dim * vec_size) { - Vec routed; - Load(routed, &dst[i]); - using namespace ops; - Array accum = cast(routed); - { - Vec v; - Load(v, &shared[i]); - accum = accum + cast(v) * dst_scale; - } - Store(&dst[i], cast(accum)); - } - } -} - -void invokeMoeCombineOutputEp( - Ref output, const Tensor& shared, const float* shared_scales, float scale, cudaStream_t st) -{ - auto& out = output.get(); - - TM_CHECK(shared); - TM_CHECK_EQ(shared.shape(0), out.shape(0)); - TM_CHECK_EQ(shared.shape(1), out.shape(1)); - - const int tokens = out.shape(0); - const int dim = out.shape(1); - - if (tokens == 0) { - return; - } - - if (shared_scales == nullptr && scale == 0) { - return; - } - - auto dispatch = [&](auto t) { - using T = decltype(t); - constexpr int threads = 256; - constexpr int vec_size = sizeof(uint4) / sizeof(T); - MoeCombineOutputEpKernel - <<>>(out.data(), shared.data(), shared_scales, dim, scale); - TM_CUDA_CHECK(cudaGetLastError()); - }; - - TM_DISPATCH_PRIMARY_DTYPES(out.dtype(), dispatch); -} - -} // namespace turbomind diff --git a/src/turbomind/kernels/gemm/moe_ep_utils.h b/src/turbomind/kernels/gemm/moe_ep_utils.h deleted file mode 100644 index 221583bec8..0000000000 --- a/src/turbomind/kernels/gemm/moe_ep_utils.h +++ /dev/null @@ -1,47 +0,0 @@ -// Copyright (c) OpenMMLab. All rights reserved. - -#pragma once - -#include - -#include "src/turbomind/core/core.h" - -namespace turbomind { - -void invokeMoeGateEp(float* topk_weights, - int64_t* topk_idx, - const float* logits, - int tokens, - int experts, - int experts_per_token, - bool softmax, - bool norm_topk, - float routed_scale, - cudaStream_t stream); - -void invokeMoeBuildAgRsRoutingMap(int* f2n, - int* f2E, - int* en2f, - int* offsets, - int8_t* masks, - int* accum, - const int64_t* topk_idx, - int tokens, - int topk, - int local_expert_offset, - int num_local_experts, - cudaStream_t stream); - -void invokeMoeLocalCombineEp(Ref out, - const Tensor& src, - const Tensor& bias, - const float* topk_weights, - const int* en2f, - const int* f2E, - int experts_per_token, - cudaStream_t st); - -void invokeMoeCombineOutputEp( - Ref output, const Tensor& shared, const float* shared_scales, float scale, cudaStream_t st); - -} // namespace turbomind diff --git a/src/turbomind/models/llama/moe_ffn_layer.cc b/src/turbomind/models/llama/moe_ffn_layer.cc index a10fe7e0db..033cdb237f 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.cc +++ b/src/turbomind/models/llama/moe_ffn_layer.cc @@ -10,7 +10,6 @@ #include "src/turbomind/kernels/gemm/moe_utils_v2.h" #include "src/turbomind/kernels/norm/rms_norm.h" -#include "src/turbomind/models/llama/LlamaFfnLayer.h" #include "src/turbomind/models/llama/LlamaLinear.h" #include "src/turbomind/models/llama/llama_utils.h" #include "src/turbomind/models/llama/moe_ffn_layer.h" @@ -39,65 +38,6 @@ class MoeFfnLayerImpl { LlamaLinear& linear_; }; -class MoeFfnTpImpl final: public MoeFfnLayerImpl { -public: - MoeFfnTpImpl(const EngineParam& engine, const Context& ctx); - - void Forward(MoeFfnLayer::ForwardParam& p) override; - - void Combine(MoeFfnLayer::ForwardParam& p) override; - -private: - void Init(MoeFfnLayer::ForwardParam& p); - - const int tp_size_; - const int max_token_num_; - int& is_warm_up_; - - bool initialized_ = false; - - Buffer_ h_offsets_; - - Buffer_ masks_; - Buffer_ f2n_; - Buffer_ f2E_; - Buffer_ en2f_; - Buffer_ scales_; - Buffer_ accum_; - Buffer_ offsets_; - - Tensor temp_; - Tensor_ shared_scales_; -}; - -MoeFfnTpImpl::MoeFfnTpImpl(const EngineParam& engine, const Context& ctx): - MoeFfnLayerImpl(ctx), - tp_size_(engine.mlp_tp_size), - max_token_num_(engine.max_forward_token_num * engine.attn_dp_size), - is_warm_up_(*ctx.is_warm_up) -{ -} - -void MoeFfnTpImpl::Init(MoeFfnLayer::ForwardParam& p) -{ - const int expert_num = p.weights->num_experts(); - const int experts_per_token = p.weights->experts_per_token; - - h_offsets_ = {expert_num + 1, kCPU}; - - const int pad_token_num = (max_token_num_ + kMoeGateVecSize - 1) / kMoeGateVecSize * kMoeGateVecSize; - - masks_ = {expert_num * pad_token_num, kDEVICE}; - f2n_ = {experts_per_token * max_token_num_, kDEVICE}; - f2E_ = {experts_per_token * max_token_num_, kDEVICE}; - en2f_ = {experts_per_token * max_token_num_, kDEVICE}; - scales_ = {experts_per_token * max_token_num_, kDEVICE}; - offsets_ = {expert_num + 1, kDEVICE}; - accum_ = {expert_num * kMoeGateMaxTiles, kDEVICE}; - - initialized_ = true; -} - Tensor_ MoeFfnLayerImpl::Gate(const Tensor& input, const LinearWeight& gate) { TM_FUNCTION_SCOPE(); @@ -111,347 +51,6 @@ Tensor_ MoeFfnLayerImpl::Gate(const Tensor& input, const LinearWeight& ga return logits; } -void MoeFfnTpImpl::Forward(MoeFfnLayer::ForwardParam& p) -{ - TM_FUNCTION_SCOPE(); - if (!initialized_) { - Init(p); - } - - const int tokens = p.input.shape(0); - const auto& moe = *p.weights; - - const auto& block = *TM_CHECK_NOTNULL(moe.block()); - - const int hidden_dim = block.hidden_dim; - const int inter_size = block.inter_size; - - const size_t padded = (tokens + kMoeGateVecSize - 1) / kMoeGateVecSize * kMoeGateVecSize; - const int expert_num = moe.num_experts(); - - TM_CHECK(expert_num); - - auto logits = Gate(p.input, *moe.gate.get()); - - TM_DEBUG_TENSOR(logits, "logits", 2); - - const auto st = core::Context::stream().handle(); - - if (p.weights->topk_method == "noaux_tc") { - // invokeMoeGate_NoAuxTC clears accum and masks internally - TM_CHECK_EQ(p.weights->n_group, 1); - TM_CHECK_EQ(p.weights->topk_group, 1); - const float* correction_bias = nullptr; - if (moe.score_correction_bias) { - correction_bias = moe.score_correction_bias.size() > 0 ? moe.score_correction_bias.data() : nullptr; - } - invokeMoeGate_NoAuxTC(f2n_.data(), - f2E_.data(), - en2f_.data(), - offsets_.data(), - scales_.data(), - masks_.data(), - accum_.data(), - logits.data(), - correction_bias, - tokens, - padded, - expert_num, - p.weights->experts_per_token, - moe.local_expert_offset(), - moe.num_local_experts(), - p.weights->norm_topk_prob, - p.weights->routed_scale, - p.weights->scoring_func == "sigmoid", - st); - } - else { - // V2: accum must be cleared by caller; masks cleared internally - TM_CUDA_CHECK(cudaMemsetAsync(accum_.data(), 0, sizeof(int) * expert_num * kMoeGateMaxTiles, st)); - - bool softmax = true; - if (p.weights->topk_method == "group_limited_greedy") { - invokeMoeSoftmaxMaskTopKGroups( - logits.data(), tokens, expert_num, expert_num / p.weights->n_group, p.weights->topk_group, st); - TM_CUDA_CHECK(cudaGetLastError()); - softmax = false; - } - - /// TODO: fix illegal memory access even if NaN are present in logits - invokeMoeGate_V2(f2n_.data(), - f2E_.data(), - en2f_.data(), - offsets_.data(), - scales_.data(), - masks_.data(), - accum_.data(), - logits.data(), - tokens, - padded, - expert_num, - p.weights->experts_per_token, - moe.local_expert_offset(), - moe.num_local_experts(), - softmax, - p.weights->norm_topk_prob, - p.weights->routed_scale, - st); - } - TM_CUDA_CHECK(cudaGetLastError()); - - if (is_warm_up_) { - std::mt19937 g; - const auto expert_ids = SampleUniform(tokens, expert_num, p.weights->experts_per_token, g); - std::vector cnt(expert_num); - for (const auto& x : expert_ids) { - ++cnt[x]; - } - h_offsets_[0] = 0; - for (int i = 0; i < expert_num; ++i) { - h_offsets_[i + 1] = h_offsets_[i] + cnt[i]; - } - TM_CUDA_CHECK( - cudaMemcpyAsync(offsets_.data(), h_offsets_.data(), sizeof(int) * (expert_num + 1), cudaMemcpyDefault, st)); - } - - temp_ = Tensor{{p.weights->experts_per_token * tokens, hidden_dim}, p.input.dtype(), p.input.device()}; - - auto indices = f2n_.slice(0, tokens * p.weights->experts_per_token); - auto offsets = offsets_.slice(0, expert_num + 1); - - if (block.w1w3) { - // Fused w1w3 path - Tensor inter; - TM_SCOPE_CALL(linear_.Forward(p.input, *block.w1w3, indices, offsets_, inter)); - - if (!block.is_fused_silu) { - Activation(inter, block.w1w3->bias, f2E_, block.act_type, st); - TM_CUDA_CHECK(cudaGetLastError()); - } - - TM_SCOPE_CALL(linear_.Forward(inter.slice({0, 0}, {-1, inter_size}), *block.w2, {}, offsets, temp_)); - } - else { - // Separate w1/w3 path - Tensor gating; - TM_SCOPE_CALL(linear_.Forward(p.input, *block.w1, indices, offsets_, gating)); - - Tensor up; - TM_SCOPE_CALL(linear_.Forward(p.input, *block.w3, indices, offsets_, up)); - - Activation(gating, up, block.act_type, st); - TM_CUDA_CHECK(cudaGetLastError()); - - TM_SCOPE_CALL(linear_.Forward(gating, *block.w2, {}, offsets, temp_)); - } - - if (moe.shared_gate) { - shared_scales_ = Gate(p.input, *moe.shared_gate); - } -} - -void MoeFfnTpImpl::Combine(MoeFfnLayer::ForwardParam& p) -{ - TM_FUNCTION_SCOPE(); - auto& moe = *p.weights; - - invokeMoeCombine(p.output, - temp_, - moe.block()->w2->bias, - scales_.data(), - en2f_.data(), - f2E_.data(), - shared_scales_.data_or((float*)nullptr), - p.weights->experts_per_token, - 1.f / tp_size_, - p.scale, - core::Context::stream().handle()); - TM_CUDA_CHECK(cudaGetLastError()); - - temp_ = {}; - shared_scales_ = {}; -} - -class MoeFfnAgRsImpl final: public MoeFfnLayerImpl { -public: - MoeFfnAgRsImpl(const EngineParam& engine, const Context& ctx); - - void Forward(MoeFfnLayer::ForwardParam& p) override; - - void Combine(MoeFfnLayer::ForwardParam& p) override; - -private: - void Init(MoeFfnLayer::ForwardParam& p); - - void ForwardFused(MoeFfnLayer::ForwardParam& p); - - const int ep_size_; - const int max_token_num_; - - bool initialized_ = false; - - Buffer_ f2n_; - Buffer_ f2E_; - Buffer_ en2f_; - Buffer_ offsets_; - Buffer_ masks_; - Buffer_ accum_; - Buffer_ scales_; - - Tensor temp_; - Tensor_ shared_scales_; -}; - -MoeFfnAgRsImpl::MoeFfnAgRsImpl(const EngineParam& engine, const Context& ctx): - MoeFfnLayerImpl(ctx), ep_size_(engine.ep_size), max_token_num_(engine.max_forward_token_num * engine.attn_dp_size) -{ -} - -void MoeFfnAgRsImpl::Init(MoeFfnLayer::ForwardParam& p) -{ - const auto& moe = *p.weights; - const int expert_num = moe.num_experts(); - const int local_expert_num = moe.num_local_experts(); - const int experts_per_token = moe.experts_per_token; - const int pad_token_num = (max_token_num_ + kMoeGateVecSize - 1) / kMoeGateVecSize * kMoeGateVecSize; - - TM_CHECK(p.weights->topk_method != "noaux_tc") - << "MoeFfnAgRsImpl does not support topk_method=noaux_tc because it requires noaux-specific scoring and " - "correction-bias handling"; - - masks_ = {expert_num * pad_token_num, kDEVICE}; - f2n_ = {experts_per_token * max_token_num_, kDEVICE}; - f2E_ = {experts_per_token * max_token_num_, kDEVICE}; - en2f_ = {experts_per_token * max_token_num_, kDEVICE}; - scales_ = {experts_per_token * max_token_num_, kDEVICE}; - offsets_ = {local_expert_num + 1, kDEVICE}; - accum_ = {expert_num * kMoeGateMaxTiles, kDEVICE}; - - initialized_ = true; -} - -void MoeFfnAgRsImpl::Forward(MoeFfnLayer::ForwardParam& p) -{ - TM_FUNCTION_SCOPE(); - - if (!initialized_) { - Init(p); - } - - const auto& moe = *p.weights; - const int tokens = p.input.shape(0); - const int experts_per_token = moe.experts_per_token; - const int expert_num = moe.num_experts(); - const int hidden_dim = TM_CHECK_NOTNULL(moe.block())->hidden_dim; - const auto st = core::Context::stream().handle(); - - auto logits = Gate(p.input, *moe.gate.get()); - TM_DEBUG_TENSOR(logits, "logits", 2); - - bool softmax = true; - if (p.weights->topk_method == "group_limited_greedy") { - invokeMoeSoftmaxMaskTopKGroups( - logits.data(), tokens, expert_num, expert_num / p.weights->n_group, p.weights->topk_group, st); - TM_CUDA_CHECK(cudaGetLastError()); - softmax = false; - } - - TM_CUDA_CHECK(cudaMemsetAsync(accum_.data(), 0, sizeof(int) * expert_num * kMoeGateMaxTiles, st)); - TM_CUDA_CHECK(cudaMemsetAsync(en2f_.data(), -1, sizeof(int) * tokens * experts_per_token, st)); - - invokeMoeGate_V2(f2n_.data(), - f2E_.data(), - en2f_.data(), - offsets_.data(), - scales_.data(), - masks_.data(), - accum_.data(), - logits.data(), - tokens, - (tokens + kMoeGateVecSize - 1) / kMoeGateVecSize * kMoeGateVecSize, - expert_num, - experts_per_token, - moe.local_expert_offset(), - moe.num_local_experts(), - softmax, - p.weights->norm_topk_prob, - p.weights->routed_scale, - st); - TM_CUDA_CHECK(cudaGetLastError()); - - temp_ = Tensor{{tokens * experts_per_token, hidden_dim}, p.input.dtype(), p.input.device()}; - ForwardFused(p); - - shared_scales_ = {}; - if (moe.shared_gate) { - shared_scales_ = Gate(p.input, *moe.shared_gate); - } -} - -void MoeFfnAgRsImpl::ForwardFused(MoeFfnLayer::ForwardParam& p) -{ - TM_FUNCTION_SCOPE(); - - const auto& moe = *p.weights; - const auto& block = *TM_CHECK_NOTNULL(moe.block()); - const int local_expert_num = moe.num_local_experts(); - const auto st = core::Context::stream().handle(); - - auto input = p.input; - const int* num_valid_tokens = offsets_.data() + local_expert_num; - auto indices = f2n_.slice(0, temp_.shape(0)); - auto offsets = offsets_.slice(0, local_expert_num + 1); - const bool indices_padded = true; - - if (block.w1w3) { - Tensor inter; - TM_SCOPE_CALL(linear_.Forward(input, *block.w1w3, indices, offsets, inter, indices_padded)); - - if (!block.is_fused_silu) { - Activation(inter, block.w1w3->bias, f2E_, block.act_type, num_valid_tokens, st); - TM_CUDA_CHECK(cudaGetLastError()); - } - - TM_SCOPE_CALL(linear_.Forward(inter.slice({0, 0}, {-1, block.inter_size}), *block.w2, {}, offsets, temp_)); - } - else { - Tensor gating; - TM_SCOPE_CALL(linear_.Forward(input, *block.w1, indices, offsets, gating, indices_padded)); - - Tensor up; - TM_SCOPE_CALL(linear_.Forward(input, *block.w3, indices, offsets, up, indices_padded)); - - Activation(gating, up, block.act_type, num_valid_tokens, st); - TM_CUDA_CHECK(cudaGetLastError()); - - TM_SCOPE_CALL(linear_.Forward(gating, *block.w2, {}, offsets, temp_)); - } -} - -void MoeFfnAgRsImpl::Combine(MoeFfnLayer::ForwardParam& p) -{ - TM_FUNCTION_SCOPE(); - - const auto& moe = *p.weights; - const auto& block = *TM_CHECK_NOTNULL(moe.block()); - - invokeMoeCombine(p.output, - temp_, - block.w2->bias, - scales_.data(), - en2f_.data(), - f2E_.data(), - shared_scales_.data_or((float*)nullptr), - moe.experts_per_token, - 1.f, - p.scale / ep_size_, - core::Context::stream().handle()); - TM_CUDA_CHECK(cudaGetLastError()); - - temp_ = {}; - shared_scales_ = {}; -} - class MoeFfnDefaultImpl final: public MoeFfnLayerImpl { public: MoeFfnDefaultImpl(const EngineParam& engine, const Context& ctx); diff --git a/src/turbomind/models/llama/unified_decoder.cc b/src/turbomind/models/llama/unified_decoder.cc index 86a7f81d38..2562156634 100644 --- a/src/turbomind/models/llama/unified_decoder.cc +++ b/src/turbomind/models/llama/unified_decoder.cc @@ -44,7 +44,6 @@ UnifiedDecoder::UnifiedDecoder(const EngineParam& engine, attn_dp_rank_(engine.attn_dp_rank), mlp_tp_size_(engine.mlp_tp_size), attn_tp_group_(ctx.comm.d_tp_group), - ep_size_(engine.ep_size), d_comm_(ctx.comm.d_comm), tune_layer_num_(engine.tune_layer_num), is_warm_up_{*ctx.is_warm_up} diff --git a/src/turbomind/models/llama/unified_decoder.h b/src/turbomind/models/llama/unified_decoder.h index 889a6dbee7..7c1f36af65 100644 --- a/src/turbomind/models/llama/unified_decoder.h +++ b/src/turbomind/models/llama/unified_decoder.h @@ -31,7 +31,6 @@ class UnifiedDecoder { const int attn_dp_size_; const int attn_dp_rank_; const int mlp_tp_size_; - const int ep_size_; const int attn_tp_group_; From 5af3a896b4930b7f340e8288f66f6c3415276a5f Mon Sep 17 00:00:00 2001 From: irexyc Date: Thu, 9 Jul 2026 08:47:53 +0000 Subject: [PATCH 11/19] remove unused --- src/turbomind/kernels/gemm/moe_utils_v2.cu | 26 ---------------------- src/turbomind/kernels/gemm/moe_utils_v2.h | 11 --------- 2 files changed, 37 deletions(-) diff --git a/src/turbomind/kernels/gemm/moe_utils_v2.cu b/src/turbomind/kernels/gemm/moe_utils_v2.cu index 84de014ae2..1045e1e2eb 100644 --- a/src/turbomind/kernels/gemm/moe_utils_v2.cu +++ b/src/turbomind/kernels/gemm/moe_utils_v2.cu @@ -571,32 +571,6 @@ __global__ void MoeGateKernel_v8(float* scales, // [e,n] template inline constexpr std::integral_constant _Int{}; -void invokeMoeScanKernel(int* f2n, - int* f2E, - int* en2f, - int* offsets, - int8_t* masks, - const int* accum, - int tokens, - int tokens_padded, - int experts, - cudaStream_t st) -{ - constexpr int base_log_tile = 9; - int log_tile = base_log_tile; - while (((tokens_padded + (1 << log_tile) - 1) >> log_tile) > kMoeGateMaxTiles) { - ++log_tile; - } - const int tiles = ceil_div(tokens_padded, 1 << log_tile); - - constexpr int threads = (1 << base_log_tile) / kMoeGateVecSize; - const dim3 blocks(tiles, experts + 1); - - MoeScanKernel_v2<<>>( - f2n, f2E, en2f, offsets, masks, accum, log_tile, tiles, tokens, tokens_padded, experts); - TM_CUDA_CHECK(cudaGetLastError()); -} - void invokeMoeGate_V2(int* f2n, // [e*n] -> n int* f2E, // [e*n] -> local E int* en2f, // [e,n] -> n*e diff --git a/src/turbomind/kernels/gemm/moe_utils_v2.h b/src/turbomind/kernels/gemm/moe_utils_v2.h index 5031febe3c..ca0053fddc 100644 --- a/src/turbomind/kernels/gemm/moe_utils_v2.h +++ b/src/turbomind/kernels/gemm/moe_utils_v2.h @@ -33,17 +33,6 @@ void invokeMoeGate_V2(int* f2n, float routed_scale, cudaStream_t st); -void invokeMoeScanKernel(int* f2n, - int* f2E, - int* en2f, - int* offsets, - int8_t* masks, - const int* accum, - int tokens, - int tokens_padded, - int experts, - cudaStream_t st); - // num_worst_tokens is the output capacity / launch upper bound for f2n/out. // If num_valid_tokens is set, it points to a device-side count of valid rows and // rows >= *num_valid_tokens return before reading f2n. From 2662b0c2b509ca21a305904147db85e33c18fc68 Mon Sep 17 00:00:00 2001 From: irexyc Date: Wed, 15 Jul 2026 10:13:15 +0000 Subject: [PATCH 12/19] update moe convert --- lmdeploy/turbomind/builders/moe.py | 12 ++++-------- lmdeploy/turbomind/models/glm4_moe_lite.py | 9 +++++---- lmdeploy/turbomind/models/gpt_oss.py | 5 ++++- lmdeploy/turbomind/models/mixtral.py | 6 ++++-- lmdeploy/turbomind/models/qwen2.py | 10 ++++++---- lmdeploy/turbomind/models/qwen3.py | 6 ++++-- lmdeploy/turbomind/models/qwen3_5.py | 8 +++++--- 7 files changed, 32 insertions(+), 24 deletions(-) diff --git a/lmdeploy/turbomind/builders/moe.py b/lmdeploy/turbomind/builders/moe.py index 62ef65fcdc..fe350d6006 100644 --- a/lmdeploy/turbomind/builders/moe.py +++ b/lmdeploy/turbomind/builders/moe.py @@ -2,7 +2,6 @@ from __future__ import annotations from ._base import Builder, ParallelGroup, SplitSide -from .module_list import ModuleListBuilder, ModuleListConfig # --------------------------------------------------------------------------- # MoeBuilder -- gate, non-expert params @@ -40,16 +39,13 @@ def add_param(self, name, tensor, split_side=None): split_side = None # specs may pass None for broadcast self._add_tensor(name, tensor, split_side) - def add_experts(self, build_expert, name='experts'): - """Build and attach expert modules with contiguous EP ownership.""" - experts = ModuleListBuilder(ModuleListConfig(), self._ctx) - for expert_idx in range(self.config.expert_num): - active_mask = self._expert_active_mask(expert_idx) + def range(self, num_experts): + for e in range(num_experts): + active_mask = self._expert_active_mask(e) if not any(active_mask): continue with self._ctx.active_mask_scope(active_mask): - experts[expert_idx] = build_expert(expert_idx) - setattr(self, name, experts.build()) + yield e def _expert_active_mask(self, expert_idx: int): ep_size = self.ep.size diff --git a/lmdeploy/turbomind/models/glm4_moe_lite.py b/lmdeploy/turbomind/models/glm4_moe_lite.py index 3737b2b167..76075ff051 100644 --- a/lmdeploy/turbomind/models/glm4_moe_lite.py +++ b/lmdeploy/turbomind/models/glm4_moe_lite.py @@ -122,10 +122,11 @@ def moe(self, pfx): correction = pfx.pop('gate.e_score_correction_bias') m.add_param('score_correction_bias', correction) - m.add_experts( - lambda e: self.ffn(pfx + 'experts' + e, - self.cfg.moe_intermediate_size, - is_expert=True)) + experts = ModuleListBuilder(ModuleListConfig(), self._ctx) + for e in m.range(cfg.expert_num): + experts[e] = self.ffn(pfx + 'experts' + e, + self.cfg.moe_intermediate_size, is_expert=True) + m.experts = experts.build() shared = self.ffn(pfx + 'shared_experts', self.cfg.intermediate_size * self.cfg.n_shared_experts) diff --git a/lmdeploy/turbomind/models/gpt_oss.py b/lmdeploy/turbomind/models/gpt_oss.py index 4453cf101f..018c4d4ba8 100644 --- a/lmdeploy/turbomind/models/gpt_oss.py +++ b/lmdeploy/turbomind/models/gpt_oss.py @@ -110,7 +110,10 @@ def moe(self, pfx): m = MoeBuilder(cfg, self._ctx, ep=self._ep) m.add_gate('gate', self._linear(pfx + 'router')) experts_pfx = pfx + 'experts' - m.add_experts(lambda e: self._packed_moe_ffn(experts_pfx, e)) + experts = ModuleListBuilder(ModuleListConfig(), self._ctx) + for e in m.range(cfg.expert_num): + experts[e] = self._packed_moe_ffn(experts_pfx, e) + m.experts = experts.build() return m.build() def layers(self, pfx): diff --git a/lmdeploy/turbomind/models/mixtral.py b/lmdeploy/turbomind/models/mixtral.py index ffd1d75963..07c90aa6a9 100644 --- a/lmdeploy/turbomind/models/mixtral.py +++ b/lmdeploy/turbomind/models/mixtral.py @@ -86,8 +86,10 @@ def moe(self, pfx): m = MoeBuilder(cfg, self._ctx, ep=self._ep) m.add_gate('gate', self._linear(pfx + 'gate')) - m.add_experts( - lambda e: self.ffn(pfx + 'experts' + e, is_expert=True)) + experts = ModuleListBuilder(ModuleListConfig(), self._ctx) + for e in m.range(self._n_experts): + experts[e] = self.ffn(pfx + 'experts' + e, is_expert=True) + m.experts = experts.build() return m.build() diff --git a/lmdeploy/turbomind/models/qwen2.py b/lmdeploy/turbomind/models/qwen2.py index 7b73130605..c23c767afa 100644 --- a/lmdeploy/turbomind/models/qwen2.py +++ b/lmdeploy/turbomind/models/qwen2.py @@ -114,10 +114,12 @@ def moe(self, pfx): m.add_gate('gate', self._linear(pfx + 'gate')) - m.add_experts( - lambda e: self.ffn(pfx + 'experts' + e, - self.cfg.moe_intermediate_size, - is_expert=True)) + experts = ModuleListBuilder(ModuleListConfig(), self._ctx) + for e in m.range(self.cfg.num_experts): + experts[e] = self.ffn(pfx + 'experts' + e, + self.cfg.moe_intermediate_size, + is_expert=True) + m.experts = experts.build() m.add_gate('shared_gate', self._linear(pfx + 'shared_expert_gate')) shared = self.ffn(pfx + 'shared_expert', diff --git a/lmdeploy/turbomind/models/qwen3.py b/lmdeploy/turbomind/models/qwen3.py index 398ff38ca7..d10820eb3f 100644 --- a/lmdeploy/turbomind/models/qwen3.py +++ b/lmdeploy/turbomind/models/qwen3.py @@ -113,8 +113,10 @@ def moe(self, pfx): m = MoeBuilder(cfg, self._ctx, ep=self._ep) m.add_gate('gate', self._linear(pfx + 'gate')) - m.add_experts( - lambda e: self.ffn(pfx + 'experts' + e, is_expert=True)) + experts = ModuleListBuilder(ModuleListConfig(), self._ctx) + for e in m.range(self._n_experts): + experts[e] = self.ffn(pfx + 'experts' + e, is_expert=True) + m.experts = experts.build() return m.build() diff --git a/lmdeploy/turbomind/models/qwen3_5.py b/lmdeploy/turbomind/models/qwen3_5.py index c6335170ba..f820d7455a 100644 --- a/lmdeploy/turbomind/models/qwen3_5.py +++ b/lmdeploy/turbomind/models/qwen3_5.py @@ -211,9 +211,11 @@ def moe(self, pfx): m.add_gate('gate', self._linear(pfx + 'gate')) experts_pfx = pfx + 'experts' - m.add_experts( - lambda e: self._moe_expert_ffn( - experts_pfx, e, self.cfg.moe_intermediate_size)) + experts = ModuleListBuilder(ModuleListConfig(), self._ctx) + for e in m.range(self._n_experts): + experts[e] = self._moe_expert_ffn( + experts_pfx, e, self.cfg.moe_intermediate_size) + m.experts = experts.build() m.add_gate('shared_gate', self._linear(pfx + 'shared_expert_gate')) shared = self.ffn(pfx + 'shared_expert', self.cfg.shared_expert_intermediate_size) From b320b5e490f4434cd4e2de4f37e0c2f9850ba5d8 Mon Sep 17 00:00:00 2001 From: irexyc Date: Wed, 15 Jul 2026 10:14:29 +0000 Subject: [PATCH 13/19] update moe convert --- lmdeploy/turbomind/models/qwen3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lmdeploy/turbomind/models/qwen3.py b/lmdeploy/turbomind/models/qwen3.py index d10820eb3f..98fcb8eee1 100644 --- a/lmdeploy/turbomind/models/qwen3.py +++ b/lmdeploy/turbomind/models/qwen3.py @@ -114,7 +114,7 @@ def moe(self, pfx): m.add_gate('gate', self._linear(pfx + 'gate')) experts = ModuleListBuilder(ModuleListConfig(), self._ctx) - for e in m.range(self._n_experts): + for e in m.range(self.cfg.num_experts): experts[e] = self.ffn(pfx + 'experts' + e, is_expert=True) m.experts = experts.build() From 7365e3f075823ce0131534e122dc0538b779a97b Mon Sep 17 00:00:00 2001 From: irexyc Date: Thu, 16 Jul 2026 08:04:56 +0000 Subject: [PATCH 14/19] fix test_moe_utils.cu --- src/turbomind/kernels/gemm/test/test_moe_utils.cu | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/turbomind/kernels/gemm/test/test_moe_utils.cu b/src/turbomind/kernels/gemm/test/test_moe_utils.cu index f5338cc4ad..37e7b79a85 100644 --- a/src/turbomind/kernels/gemm/test/test_moe_utils.cu +++ b/src/turbomind/kernels/gemm/test/test_moe_utils.cu @@ -233,6 +233,8 @@ bool test_moe_gate(int tokens, // tokens_padded, expert_num, experts_per_token, + 0, + expert_num, softmax, false, 1.f, From 676de7e91472477d3e6d4e0f957f599b6ef6295f Mon Sep 17 00:00:00 2001 From: irexyc Date: Thu, 16 Jul 2026 09:47:54 +0000 Subject: [PATCH 15/19] update --- src/turbomind/turbomind.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/turbomind/turbomind.cc b/src/turbomind/turbomind.cc index 37db8d53b2..1a9ff0e0dc 100644 --- a/src/turbomind/turbomind.cc +++ b/src/turbomind/turbomind.cc @@ -245,7 +245,7 @@ void TurboMind::Impl::CreateContext(int index) p.model_tp_rank = c.d_comm->rank(c.d_tp_group); p.attn_tp_rank = p.model_tp_rank / p.attn_cp_size; p.mlp_tp_rank = inner_rank % p.mlp_tp_size; - p.ep_rank = inner_rank % p.ep_size; + p.ep_rank = inner_rank / p.mlp_tp_size; } if (c.h_tp_group->rank() == 0) { From eed837ea2cef5a35b0e83111886874031077e474 Mon Sep 17 00:00:00 2001 From: irexyc Date: Fri, 17 Jul 2026 09:05:24 +0000 Subject: [PATCH 16/19] use shard ffn input --- src/turbomind/models/llama/moe_ffn_layer.cc | 37 ++++++++++++++++++- src/turbomind/models/llama/moe_ffn_layer.h | 2 + src/turbomind/models/llama/unified_decoder.cc | 12 ++++-- 3 files changed, 47 insertions(+), 4 deletions(-) diff --git a/src/turbomind/models/llama/moe_ffn_layer.cc b/src/turbomind/models/llama/moe_ffn_layer.cc index 033cdb237f..5f42cfbc52 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.cc +++ b/src/turbomind/models/llama/moe_ffn_layer.cc @@ -32,6 +32,8 @@ class MoeFfnLayerImpl { virtual void Combine(MoeFfnLayer::ForwardParam& p) = 0; + virtual Tensor GetShardFfnInput(Tensor& global_hidden_states) = 0; + protected: Tensor_ Gate(const Tensor& input, const LinearWeight& gate); @@ -59,11 +61,14 @@ class MoeFfnDefaultImpl final: public MoeFfnLayerImpl { void Combine(MoeFfnLayer::ForwardParam& p) override; + Tensor GetShardFfnInput(Tensor& global_hidden_states) override; + private: void Init(MoeFfnLayer::ForwardParam& p); const int tp_size_; const int ep_size_; + const int ep_rank_; const int max_token_num_; int& is_warm_up_; @@ -87,6 +92,7 @@ MoeFfnDefaultImpl::MoeFfnDefaultImpl(const EngineParam& engine, const Context& c MoeFfnLayerImpl(ctx), tp_size_(engine.mlp_tp_size), ep_size_(engine.ep_size), + ep_rank_(engine.ep_rank), max_token_num_(engine.max_forward_token_num * engine.attn_dp_size), is_warm_up_(*ctx.is_warm_up) { @@ -283,7 +289,7 @@ void MoeFfnDefaultImpl::Combine(MoeFfnLayer::ForwardParam& p) shared_scales_.data_or((float*)nullptr), moe.experts_per_token, 1.f / tp_size_, - p.scale / ep_size_, + p.scale, core::Context::stream().handle()); TM_CUDA_CHECK(cudaGetLastError()); @@ -291,6 +297,30 @@ void MoeFfnDefaultImpl::Combine(MoeFfnLayer::ForwardParam& p) shared_scales_ = {}; } +Tensor MoeFfnDefaultImpl::GetShardFfnInput(Tensor& global_hidden_states) +{ + TM_FUNCTION_SCOPE(); + if (ep_size_ == 1) { + return global_hidden_states; + } + + const int token_num = global_hidden_states.shape(0); + const int tokens_per_rank = token_num / ep_size_; + const int remainder = token_num % ep_size_; + const int local_token_num = tokens_per_rank + (ep_rank_ < remainder); + const int local_token_begin = ep_rank_ * tokens_per_rank + std::min(ep_rank_, remainder); + const int local_token_end = local_token_begin + local_token_num; + + if (local_token_begin > 0) { + Clear(global_hidden_states.slice(0, local_token_begin)); + } + if (local_token_end < token_num) { + Clear(global_hidden_states.slice(local_token_end, token_num - local_token_end)); + } + + return global_hidden_states.slice(local_token_begin, local_token_num); +} + MoeFfnLayer::MoeFfnLayer(const EngineParam& engine, const Context& ctx) { if (engine.ep_size <= 1 || engine.nnodes == 1) { @@ -313,4 +343,9 @@ void MoeFfnLayer::Combine(ForwardParam& p) impl_->Combine(p); } +Tensor MoeFfnLayer::GetShardFfnInput(Tensor& global_hidden_states) +{ + return impl_->GetShardFfnInput(global_hidden_states); +} + } // namespace turbomind diff --git a/src/turbomind/models/llama/moe_ffn_layer.h b/src/turbomind/models/llama/moe_ffn_layer.h index 39f68fa952..ea366440a6 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.h +++ b/src/turbomind/models/llama/moe_ffn_layer.h @@ -31,6 +31,8 @@ class MoeFfnLayer { void Combine(ForwardParam& p); + Tensor GetShardFfnInput(Tensor& global_hidden_states); + private: std::unique_ptr impl_; }; diff --git a/src/turbomind/models/llama/unified_decoder.cc b/src/turbomind/models/llama/unified_decoder.cc index ba7857c7c6..ba0fbfe743 100644 --- a/src/turbomind/models/llama/unified_decoder.cc +++ b/src/turbomind/models/llama/unified_decoder.cc @@ -96,6 +96,8 @@ UnifiedDecoder::UnifiedDecoder(CacheRegistry& registry, ctx, phases); } + + TM_CHECK(!(moe_weights.empty() && engine.ep_size > 1)) << "Dense model is not supported with ep_size > 1"; } void UnifiedDecoder::AllreduceResidualRMSnorm(Tensor& hidden_states, @@ -289,14 +291,18 @@ void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vectormoe_ffn.get(), - ffn_layer_ ? 1.f : 0.f, + weights.at(layer)->feed_forward ? 1.f : 0.f, layer}; moe_ffn_layer_->Forward(*moe_fwd_param); } if (ffn_layer_ && weights.at(layer)->feed_forward) { - ffn_layer_->forward( - {global_hidden_states, global_hidden_states, weights.at(layer)->feed_forward.get(), (int)layer}); + auto ffn_input_shared = + moe_ffn_layer_ ? moe_ffn_layer_->GetShardFfnInput(global_hidden_states) : global_hidden_states; + if (ffn_input_shared.shape(0) > 0) { + ffn_layer_->forward( + {ffn_input_shared, ffn_input_shared, weights.at(layer)->feed_forward.get(), (int)layer}); + } } if (moe_fwd_param) { From 555d17e36abdcf617b695064516ec0d6fa961bcb Mon Sep 17 00:00:00 2001 From: irexyc Date: Fri, 17 Jul 2026 10:28:54 +0000 Subject: [PATCH 17/19] remove indices_padded --- src/turbomind/models/llama/LlamaLinear.cu | 18 ++++++------------ src/turbomind/models/llama/LlamaLinear.h | 3 +-- src/turbomind/models/llama/moe_ffn_layer.cc | 7 +++---- 3 files changed, 10 insertions(+), 18 deletions(-) diff --git a/src/turbomind/models/llama/LlamaLinear.cu b/src/turbomind/models/llama/LlamaLinear.cu index 7547bf0c23..7102f97372 100644 --- a/src/turbomind/models/llama/LlamaLinear.cu +++ b/src/turbomind/models/llama/LlamaLinear.cu @@ -64,11 +64,7 @@ struct LlamaLinear::Impl { } std::tuple // - GetOperandA(const LinearWeight& weight, - const Tensor& input, - Buffer_ indices, - const Buffer_& offsets, - bool indices_padded) + GetOperandA(const LinearWeight& weight, const Tensor& input, Buffer_ indices, const Buffer_& offsets) { auto st = core::Context::stream().handle(); @@ -89,8 +85,8 @@ struct LlamaLinear::Impl { const bool is_cublas_grouped = offsets && getSMVersion() == 100 && weight.weight_format.dtype == kBfloat16; if (indices && (A.dtype() == kFloat8_e4m3 || is_cublas_grouped)) { const int k = A.shape(1); - const int* num_valid_tokens = indices_padded ? offsets.data() + offsets.size() - 1 : nullptr; Tensor A_e = {{m, k}, A.dtype(), kDEVICE}; + const int* num_valid_tokens = offsets ? offsets.data() + offsets.size() - 1 : nullptr; TM_SCOPE_CALL(invokeMoeDispatch(A_e, A, indices.data(), m, num_valid_tokens, st)); if (U) { Tensor U_e; @@ -121,8 +117,7 @@ struct LlamaLinear::Impl { const Tensor& input, // const LinearWeight& weight, const Buffer_& indices, - const Buffer_& offsets, - bool indices_padded) + const Buffer_& offsets) { TM_FUNCTION_SCOPE(); using namespace gemm; @@ -134,7 +129,7 @@ struct LlamaLinear::Impl { op.quant_b = MakeQuantDesc(weight.weight_format); op.batch_dim = 0; - auto&& [A, desc_A, U, desc_U] = GetOperandA(weight, input, indices, offsets, indices_padded); + auto&& [A, desc_A, U, desc_U] = GetOperandA(weight, input, indices, offsets); auto&& [B, desc_B, V, desc_V] = GetOperandB(weight); Tensor& D = output; @@ -200,8 +195,7 @@ void LlamaLinear::Forward(const Tensor& input, // const LinearWeight& weight, const Buffer_& indices, const Buffer_& offsets, - Ref output, - bool indices_padded) + Ref output) { Tensor in = input.view({-1, input.shape(-1)}); @@ -209,7 +203,7 @@ void LlamaLinear::Forward(const Tensor& input, // output.get() = output.get().view({-1, output.get().shape(-1)}); } - impl_->Forward(output.get(), in, weight, indices, offsets, indices_padded); + impl_->Forward(output.get(), in, weight, indices, offsets); } void LlamaLinear::set_measure(bool measure) diff --git a/src/turbomind/models/llama/LlamaLinear.h b/src/turbomind/models/llama/LlamaLinear.h index cbd1a73efa..ee3315a73d 100644 --- a/src/turbomind/models/llama/LlamaLinear.h +++ b/src/turbomind/models/llama/LlamaLinear.h @@ -22,8 +22,7 @@ class LlamaLinear { const LinearWeight& weight, const Buffer_& indices, const Buffer_& offsets, - Ref output, - bool indices_padded = false); + Ref output); void set_measure(bool measure); diff --git a/src/turbomind/models/llama/moe_ffn_layer.cc b/src/turbomind/models/llama/moe_ffn_layer.cc index 5f42cfbc52..68c7048399 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.cc +++ b/src/turbomind/models/llama/moe_ffn_layer.cc @@ -239,7 +239,6 @@ void MoeFfnDefaultImpl::Forward(MoeFfnLayer::ForwardParam& p) temp_ = Tensor{{tokens * experts_per_token, hidden_dim}, p.input.dtype(), p.input.device()}; // For ep_size > 1, the valid tokens are less than tokens * experts_per_token - const bool indices_padded = ep_size_ > 1 ? true : false; const int* num_valid_tokens = ep_size_ > 1 ? offsets_.data() + local_expert_num : nullptr; auto indices = f2n_.slice(0, temp_.shape(0)); @@ -247,7 +246,7 @@ void MoeFfnDefaultImpl::Forward(MoeFfnLayer::ForwardParam& p) if (block.w1w3) { Tensor inter; - TM_SCOPE_CALL(linear_.Forward(p.input, *block.w1w3, indices, offsets, inter, indices_padded)); + TM_SCOPE_CALL(linear_.Forward(p.input, *block.w1w3, indices, offsets, inter)); if (!block.is_fused_silu) { Activation(inter, block.w1w3->bias, f2E_, block.act_type, num_valid_tokens, st); @@ -259,10 +258,10 @@ void MoeFfnDefaultImpl::Forward(MoeFfnLayer::ForwardParam& p) else { // Separate w1/w3 path Tensor gating; - TM_SCOPE_CALL(linear_.Forward(p.input, *block.w1, indices, offsets, gating, indices_padded)); + TM_SCOPE_CALL(linear_.Forward(p.input, *block.w1, indices, offsets, gating)); Tensor up; - TM_SCOPE_CALL(linear_.Forward(p.input, *block.w3, indices, offsets, up, indices_padded)); + TM_SCOPE_CALL(linear_.Forward(p.input, *block.w3, indices, offsets, up)); Activation(gating, up, block.act_type, num_valid_tokens, st); TM_CUDA_CHECK(cudaGetLastError()); From 6dd7f0697f51ef117611fb75a4ddfe76af79c8b8 Mon Sep 17 00:00:00 2001 From: irexyc Date: Fri, 17 Jul 2026 10:31:35 +0000 Subject: [PATCH 18/19] remove // --- src/turbomind/models/llama/LlamaLinear.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/turbomind/models/llama/LlamaLinear.cu b/src/turbomind/models/llama/LlamaLinear.cu index 7102f97372..f0a3126c7e 100644 --- a/src/turbomind/models/llama/LlamaLinear.cu +++ b/src/turbomind/models/llama/LlamaLinear.cu @@ -63,7 +63,7 @@ struct LlamaLinear::Impl { return {B, desc_B, V, desc_V}; } - std::tuple // + std::tuple GetOperandA(const LinearWeight& weight, const Tensor& input, Buffer_ indices, const Buffer_& offsets) { auto st = core::Context::stream().handle(); From ae724cb5a110dc19087817b2020afdd70d4ba3bd Mon Sep 17 00:00:00 2001 From: irexyc Date: Sat, 18 Jul 2026 11:17:18 +0000 Subject: [PATCH 19/19] use seperate stream for ffn input clear --- src/turbomind/models/llama/moe_ffn_layer.cc | 38 ++++++++++++++++++--- 1 file changed, 33 insertions(+), 5 deletions(-) diff --git a/src/turbomind/models/llama/moe_ffn_layer.cc b/src/turbomind/models/llama/moe_ffn_layer.cc index 68c7048399..18a44874f2 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.cc +++ b/src/turbomind/models/llama/moe_ffn_layer.cc @@ -86,6 +86,13 @@ class MoeFfnDefaultImpl final: public MoeFfnLayerImpl { Tensor temp_; Tensor_ shared_scales_; + + // When a MOE model enables EP inference and the model includes dense layers or shared experts, an additional stream + // is used to perform sharding and cleaning of the FFN inputs. + Stream clear_stream_; + Event clear_ready_event_; + Event clear_done_event_; + bool clear_pending_ = false; }; MoeFfnDefaultImpl::MoeFfnDefaultImpl(const EngineParam& engine, const Context& ctx): @@ -96,6 +103,11 @@ MoeFfnDefaultImpl::MoeFfnDefaultImpl(const EngineParam& engine, const Context& c max_token_num_(engine.max_forward_token_num * engine.attn_dp_size), is_warm_up_(*ctx.is_warm_up) { + if (ep_size_ > 1) { + clear_stream_ = Stream::create(); + clear_ready_event_ = Event::create(); + clear_done_event_ = Event::create(); + } } void MoeFfnDefaultImpl::Init(MoeFfnLayer::ForwardParam& p) @@ -279,6 +291,11 @@ void MoeFfnDefaultImpl::Combine(MoeFfnLayer::ForwardParam& p) TM_FUNCTION_SCOPE(); auto& moe = *p.weights; + if (clear_pending_) { + core::Context::stream().Wait(clear_done_event_); + clear_pending_ = false; + } + invokeMoeCombine(p.output, temp_, TM_CHECK_NOTNULL(moe.block())->w2->bias, @@ -303,6 +320,7 @@ Tensor MoeFfnDefaultImpl::GetShardFfnInput(Tensor& global_hidden_states) return global_hidden_states; } + TM_CHECK(!clear_pending_); const int token_num = global_hidden_states.shape(0); const int tokens_per_rank = token_num / ep_size_; const int remainder = token_num % ep_size_; @@ -310,11 +328,21 @@ Tensor MoeFfnDefaultImpl::GetShardFfnInput(Tensor& global_hidden_states) const int local_token_begin = ep_rank_ * tokens_per_rank + std::min(ep_rank_, remainder); const int local_token_end = local_token_begin + local_token_num; - if (local_token_begin > 0) { - Clear(global_hidden_states.slice(0, local_token_begin)); - } - if (local_token_end < token_num) { - Clear(global_hidden_states.slice(local_token_end, token_num - local_token_end)); + if (local_token_begin > 0 || local_token_end < token_num) { + auto& stream = core::Context::stream(); + + clear_ready_event_.Record(stream); + clear_stream_.Wait(clear_ready_event_); + + if (local_token_begin > 0) { + Clear(global_hidden_states.slice(0, local_token_begin), clear_stream_); + } + if (local_token_end < token_num) { + Clear(global_hidden_states.slice(local_token_end, token_num - local_token_end), clear_stream_); + } + + clear_done_event_.Record(clear_stream_); + clear_pending_ = true; } return global_hidden_states.slice(local_token_begin, local_token_num);