diff --git a/lmdeploy/turbomind/model_loader.py b/lmdeploy/turbomind/model_loader.py index ad792cd90d..3917c4cbe7 100644 --- a/lmdeploy/turbomind/model_loader.py +++ b/lmdeploy/turbomind/model_loader.py @@ -50,6 +50,15 @@ def _bind_runtime(self): model_tp = ParallelGroup(ec.attn_tp_size * ec.attn_cp_size, [mc.model_tp_rank(g) for g in range(self.gpu_count)]) + # Dense (non-expert) FFN TP: node-local — one node's ranks within + # the comm domain (shard index = inner_rank % domain_size, + # inner_rank = ep_rank * mlp_tp_size + mlp_tp_rank). + dense_size = min(mlp_tp.size * ep.size, self.gpu_count) + dense_tp = ParallelGroup( + dense_size, + [(e * mlp_tp.size + m) % dense_size + for e, m in zip(ep.ranks, mlp_tp.ranks)]) + self.model.bind_runtime( ctx=ctx, root_handles=[mc.root(g) for g in range(self.gpu_count)], @@ -57,6 +66,7 @@ def _bind_runtime(self): mlp_tp=mlp_tp, ep=ep, model_tp=model_tp, + dense_tp=dense_tp, ) def export(self): diff --git a/lmdeploy/turbomind/models/glm4_moe_lite.py b/lmdeploy/turbomind/models/glm4_moe_lite.py index f97ad9fbc4..167aedb9ef 100644 --- a/lmdeploy/turbomind/models/glm4_moe_lite.py +++ b/lmdeploy/turbomind/models/glm4_moe_lite.py @@ -100,14 +100,14 @@ def attn(self, pfx): # FFN / MoE factories # ------------------------------------------------------------------ - def ffn(self, pfx, inter_size, is_expert=False): + def ffn(self, pfx, inter_size, is_expert=False, *, tp): w1, w3, w2 = [self._linear(pfx + f'{x}_proj') for x in ('gate', 'up', 'down')] cfg = self._ffn_cfg.clone() cfg.inter_size = inter_size cfg.is_expert = is_expert - m = FfnBuilder(cfg, self._ctx, tp=self._mlp_tp) + m = FfnBuilder(cfg, self._ctx, tp=tp) m.add_ffn(w1, w2, w3) return m.build() @@ -124,13 +124,15 @@ def moe(self, pfx): 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) + self.cfg.moe_intermediate_size, + is_expert=True, tp=self._mlp_tp) m.experts = experts.build() - shared = self.ffn(pfx + 'shared_experts', - self.cfg.intermediate_size * self.cfg.n_shared_experts) + m.shared = self.ffn(pfx + 'shared_experts', + self.cfg.intermediate_size * self.cfg.n_shared_experts, + tp=self._mlp_tp) - return m.build(), shared + return m.build() def layers(self, pfx): layers = ModuleListBuilder(ModuleListConfig(), self._ctx) @@ -140,8 +142,9 @@ def layers(self, pfx): d.attention = self.attn(p + 'self_attn') d.ffn_norm = self.norm(p + 'post_attention_layernorm') if self.cfg.mlp_layer_types[i] == 'sparse': - d.moe_ffn, d.feed_forward = self.moe(p + 'mlp') + d.moe_ffn = self.moe(p + 'mlp') else: - d.feed_forward = self.ffn(p + 'mlp', self.cfg.intermediate_size) + d.feed_forward = self.ffn(p + 'mlp', self.cfg.intermediate_size, + tp=self._dense_tp) layers[i] = d.build() return layers.build() diff --git a/lmdeploy/turbomind/models/interns2_mobius.py b/lmdeploy/turbomind/models/interns2_mobius.py index 71c3c13c2e..1f40be69b0 100644 --- a/lmdeploy/turbomind/models/interns2_mobius.py +++ b/lmdeploy/turbomind/models/interns2_mobius.py @@ -118,6 +118,12 @@ def build_meta_moe_layer(text_model, pfx, layer_id: int): m = MoeBuilder(cfg, text_model._ctx, ep=text_model._ep) m.add_gate('shared_gate', text_model._linear(pfx + 'shared_expert_gate')) + # Attach the shared expert before build(): post-build child assignment + # raises, and the C++ MoeWeight owns it as the `shared` child. + m.shared = text_model.ffn( + pfx + 'shared_expert', + text_model.cfg.shared_expert_intermediate_size, + tp=text_model._mlp_tp) moe = m.build() pack = text_model._meta_pack_modules[layer_id % text_model._n_meta_groups] @@ -126,10 +132,7 @@ def build_meta_moe_layer(text_model, pfx, layer_id: int): with text_model._ctx.devices[i]: moe_h.set_meta_pack(pack_h) - shared = text_model.ffn( - pfx + 'shared_expert', - text_model.cfg.shared_expert_intermediate_size) - return moe, shared + return moe class InternS2MobiusTextModel(Qwen3_5TextModel): @@ -178,7 +181,7 @@ def layers(self, pfx): d.linear_attn = self.linear_attn(p + 'linear_attn') else: d.attention = self.attn(p + 'self_attn') - d.moe_ffn, d.feed_forward = build_meta_moe_layer(self, p + 'mlp', i) + d.moe_ffn = build_meta_moe_layer(self, p + 'mlp', i) d.attention_norm = self.norm(p + 'input_layernorm', zero_centered=True) d.ffn_norm = self.norm(p + 'post_attention_layernorm', zero_centered=True) layers[i] = d.build() diff --git a/lmdeploy/turbomind/models/internvl.py b/lmdeploy/turbomind/models/internvl.py index d15a3ccc2b..6ba893ce0c 100644 --- a/lmdeploy/turbomind/models/internvl.py +++ b/lmdeploy/turbomind/models/internvl.py @@ -433,7 +433,7 @@ def __init__(self, cfg: PretrainedConfig, *, resolver, raise ValueError(f'InternVL TurboMind vision architecture {arch!r} is not supported.') def bind_runtime(self, *, ctx, root_handles, - attn_tp, mlp_tp, ep, model_tp): + attn_tp, mlp_tp, ep, model_tp, dense_tp): self.text_model.bind_runtime( ctx=ctx, root_handles=root_handles, @@ -441,6 +441,7 @@ def bind_runtime(self, *, ctx, root_handles, mlp_tp=mlp_tp, ep=ep, model_tp=model_tp, + dense_tp=dense_tp, ) if self.vision_model is not None: vision_ctx = Context( diff --git a/lmdeploy/turbomind/models/qwen2.py b/lmdeploy/turbomind/models/qwen2.py index 73d50a249d..df19604ff8 100644 --- a/lmdeploy/turbomind/models/qwen2.py +++ b/lmdeploy/turbomind/models/qwen2.py @@ -96,14 +96,14 @@ def reorder(x): return m.build() - def ffn(self, pfx, inter_size, is_expert=False): + def ffn(self, pfx, inter_size, is_expert=False, *, tp): w1, w3, w2 = [self._linear(pfx + f'{x}_proj') for x in ('gate', 'up', 'down')] cfg = self._ffn_cfg.clone() cfg.inter_size = inter_size cfg.is_expert = is_expert - m = FfnBuilder(cfg, self._ctx, tp=self._mlp_tp) + m = FfnBuilder(cfg, self._ctx, tp=tp) m.add_ffn(w1, w2, w3) return m.build() @@ -118,14 +118,15 @@ def moe(self, pfx): for e in m.range(self.cfg.num_experts): experts[e] = self.ffn(pfx + 'experts' + e, self.cfg.moe_intermediate_size, - is_expert=True) + is_expert=True, tp=self._mlp_tp) 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) + m.shared = self.ffn(pfx + 'shared_expert', + self.cfg.shared_expert_intermediate_size, + tp=self._mlp_tp) - return m.build(), shared + return m.build() # ------------------------------------------------------------------ # layers() — layer dispatch loop @@ -137,9 +138,10 @@ def layers(self, pfx): d = DecoderLayerBuilder(DecoderLayerConfig(), self._ctx) d.attention = self.attn(p + 'self_attn') if self._n_experts > 0: - d.moe_ffn, d.feed_forward = self.moe(p + 'mlp') + d.moe_ffn = self.moe(p + 'mlp') else: - d.feed_forward = self.ffn(p + 'mlp', self.cfg.intermediate_size) + d.feed_forward = self.ffn(p + 'mlp', self.cfg.intermediate_size, + tp=self._mlp_tp) d.attention_norm = self.norm(p + 'input_layernorm') d.ffn_norm = self.norm(p + 'post_attention_layernorm') layers[i] = d.build() diff --git a/lmdeploy/turbomind/models/qwen2_vl.py b/lmdeploy/turbomind/models/qwen2_vl.py index 3430772840..33c7f05f47 100644 --- a/lmdeploy/turbomind/models/qwen2_vl.py +++ b/lmdeploy/turbomind/models/qwen2_vl.py @@ -413,7 +413,7 @@ def __init__(self, cfg, *, resolver, vision_resolver=None, self._vision_data_type = vision_data_type def bind_runtime(self, *, ctx, root_handles, - attn_tp, mlp_tp, ep, model_tp): + attn_tp, mlp_tp, ep, model_tp, dense_tp): self.text_model.bind_runtime( ctx=ctx, root_handles=root_handles, @@ -421,6 +421,7 @@ def bind_runtime(self, *, ctx, root_handles, mlp_tp=mlp_tp, ep=ep, model_tp=model_tp, + dense_tp=dense_tp, ) if self.vision_model is not None: vision_ctx = Context( diff --git a/lmdeploy/turbomind/models/qwen3.py b/lmdeploy/turbomind/models/qwen3.py index 9d240eacb5..f41264334b 100644 --- a/lmdeploy/turbomind/models/qwen3.py +++ b/lmdeploy/turbomind/models/qwen3.py @@ -97,13 +97,13 @@ def reorder(x): return m.build() - def ffn(self, pfx, is_expert=False): + def ffn(self, pfx, is_expert=False, *, tp): w1, w3, w2 = [self._linear(pfx + f'{x}_proj') for x in ('gate', 'up', 'down')] cfg = self._ffn_cfg.clone() cfg.is_expert = is_expert - m = FfnBuilder(cfg, self._ctx, tp=self._mlp_tp) + m = FfnBuilder(cfg, self._ctx, tp=tp) m.add_ffn(w1, w2, w3) return m.build() @@ -115,21 +115,30 @@ def moe(self, pfx): experts = ModuleListBuilder(ModuleListConfig(), self._ctx) for e in m.range(self.cfg.num_experts): - experts[e] = self.ffn(pfx + 'experts' + e, is_expert=True) + experts[e] = self.ffn(pfx + 'experts' + e, is_expert=True, + tp=self._mlp_tp) m.experts = experts.build() return m.build() def layers(self, pfx): + mlp_only = set(getattr(self.cfg, 'mlp_only_layers', None) or []) + # Dense layers in a MoE model shard node-locally (_dense_tp); a + # pure-dense model keeps the historical mlp_tp sharding. Keyed on + # whether any layer actually builds MoE — the same fact the C++ + # decoder keys the dense reduce group on. + any_moe = self._n_experts > 0 and any( + i not in mlp_only for i in range(self.cfg.num_hidden_layers)) + dense_tp = self._dense_tp if any_moe else self._mlp_tp layers = ModuleListBuilder(ModuleListConfig(), self._ctx) for i, p in pfx.slices(0, self.cfg.num_hidden_layers): d = DecoderLayerBuilder(DecoderLayerConfig(), self._ctx) d.attention_norm = self.norm(p + 'input_layernorm') d.attention = self.attn(p + 'self_attn') d.ffn_norm = self.norm(p + 'post_attention_layernorm') - if self._n_experts: + if self._n_experts and i not in mlp_only: d.moe_ffn = self.moe(p + 'mlp') else: - d.feed_forward = self.ffn(p + 'mlp') + d.feed_forward = self.ffn(p + 'mlp', tp=dense_tp) layers[i] = d.build() return layers.build() diff --git a/lmdeploy/turbomind/models/qwen3_5.py b/lmdeploy/turbomind/models/qwen3_5.py index 5f60236877..41d1fdaa07 100644 --- a/lmdeploy/turbomind/models/qwen3_5.py +++ b/lmdeploy/turbomind/models/qwen3_5.py @@ -204,7 +204,7 @@ def linear_attn(self, pfx): # FFN / MoE factories # ------------------------------------------------------------------ - def ffn(self, pfx, inter_size, is_expert=False): + def ffn(self, pfx, inter_size, is_expert=False, *, tp): try: w1, w3, w2 = [self._linear(pfx + f'{x}_proj') for x in ('gate', 'up', 'down')] @@ -215,7 +215,7 @@ def ffn(self, pfx, inter_size, is_expert=False): cfg.inter_size = inter_size cfg.is_expert = is_expert - m = FfnBuilder(cfg, self._ctx, tp=self._mlp_tp) + m = FfnBuilder(cfg, self._ctx, tp=tp) m.add_ffn(w1, w2, w3) return m.build() @@ -234,9 +234,13 @@ def moe(self, pfx): 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) + shared = self.ffn(pfx + 'shared_expert', + self.cfg.shared_expert_intermediate_size, + tp=self._mlp_tp) + if shared is not None: + m.shared = shared - return m.build(), shared + return m.build() def _packed_moe_ffn(self, experts_pfx, expert_idx, inter_size): w1, w2, w3 = read_packed_moe_expert( @@ -254,7 +258,8 @@ def _packed_moe_ffn(self, experts_pfx, expert_idx, inter_size): def _moe_expert_ffn(self, experts_pfx, expert_idx, inter_size): expert_pfx = experts_pfx + expert_idx - return (self.ffn(expert_pfx, inter_size, is_expert=True) + return (self.ffn(expert_pfx, inter_size, is_expert=True, + tp=self._mlp_tp) or self._packed_moe_ffn(experts_pfx, expert_idx, inter_size)) # ------------------------------------------------------------------ @@ -270,9 +275,10 @@ def layers(self, pfx): else: d.attention = self.attn(p + 'self_attn') if self._n_experts > 0: - d.moe_ffn, d.feed_forward = self.moe(p + 'mlp') + d.moe_ffn = self.moe(p + 'mlp') else: - d.feed_forward = self.ffn(p + 'mlp', self.cfg.intermediate_size) + d.feed_forward = self.ffn(p + 'mlp', self.cfg.intermediate_size, + tp=self._mlp_tp) d.attention_norm = self.norm( p + 'input_layernorm', zero_centered=True, @@ -602,7 +608,7 @@ def __init__(self, cfg: Qwen3_5Config | Qwen3_5MoeConfig, *, resolver, self._vision_data_type = vision_data_type def bind_runtime(self, *, ctx, root_handles, - attn_tp, mlp_tp, ep, model_tp): + attn_tp, mlp_tp, ep, model_tp, dense_tp): self.text_model.bind_runtime( ctx=ctx, root_handles=root_handles, @@ -610,6 +616,7 @@ def bind_runtime(self, *, ctx, root_handles, mlp_tp=mlp_tp, ep=ep, model_tp=model_tp, + dense_tp=dense_tp, ) if self.vision_model is not None: diff --git a/lmdeploy/turbomind/text_model.py b/lmdeploy/turbomind/text_model.py index e1f0aaf1fe..641696d1aa 100644 --- a/lmdeploy/turbomind/text_model.py +++ b/lmdeploy/turbomind/text_model.py @@ -41,13 +41,17 @@ def _vocab_size(self) -> int: return self.cfg.vocab_size def bind_runtime(self, *, ctx, root_handles, - attn_tp, mlp_tp, ep, model_tp): + attn_tp, mlp_tp, ep, model_tp, dense_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 + # TP group for dense layers inside MoE models — sharded node-locally + # by ModelLoader (C++ reduce group: d_node_group). Pure-dense models + # use the mlp_tp group. + self._dense_tp = dense_tp def _linear(self, pfx: Prefix, *, optional: bool = False) -> Linear | None: diff --git a/src/turbomind/comm/nccl/deepep/moe_a2a_utils.cu b/src/turbomind/comm/nccl/deepep/moe_a2a_utils.cu index 6605505349..3cda05b454 100644 --- a/src/turbomind/comm/nccl/deepep/moe_a2a_utils.cu +++ b/src/turbomind/comm/nccl/deepep/moe_a2a_utils.cu @@ -184,16 +184,23 @@ void invokeMoeA2AMapping(int* f2n, template __global__ void MoeA2ASharedCombineKernel(T* output, // const T* routed, + const T* shared, // this rank's shared FFN rows, or nullptr const float* shared_scales, - int hidden_dim, - float shared_scale) + int hidden_dim) { const int token_idx = blockIdx.x; output += (int64_t)token_idx * hidden_dim; routed += (int64_t)token_idx * hidden_dim; - if (shared_scales) { + // Nullability lives in the kernel, not the caller: this kernel is also + // the routed write-back (the Store below is unconditional), so the + // caller must stay unconditional and pass nullptr when there is no + // shared expert. + const T* shared_row = shared ? shared + (int64_t)token_idx * hidden_dim : nullptr; + + float shared_scale = 1.f; + if (shared_scales) { // null for gate-less shared experts shared_scale *= fdividef(1.f, 1.f + expf(-__ldg(shared_scales + token_idx))); } @@ -203,21 +210,20 @@ __global__ void MoeA2ASharedCombineKernel(T* output, // Vec routed_vec; Load(routed_vec, routed + i); auto result = cast(routed_vec); - - if (shared_scale != 0.f) { + if (shared_row) { Vec shared_vec; - Load(shared_vec, output + i); + Load(shared_vec, shared_row + i); // was: in-place read from output using namespace ops; result = result + cast(shared_vec) * shared_scale; } - Store(output + i, cast(result)); + Store(output + i, cast(result)); // ALWAYS executed — the routed write-back } } void invokeMoeA2ASharedCombine(core::Tensor& output, const core::Tensor& routed, + const core::Tensor& shared, const float* shared_scales, - float shared_scale, cudaStream_t stream) { const int tokens = output.shape(0); @@ -231,7 +237,7 @@ void invokeMoeA2ASharedCombine(core::Tensor& output, constexpr int vec_size = 16 / sizeof(T); constexpr int block_dim = 256; MoeA2ASharedCombineKernel<<>>( - output.data(), routed.data(), shared_scales, hidden_dim, shared_scale); + output.data(), routed.data(), shared.data_or((T*)nullptr), shared_scales, hidden_dim); TM_CUDA_CHECK(cudaGetLastError()); }; diff --git a/src/turbomind/comm/nccl/deepep/moe_a2a_utils.h b/src/turbomind/comm/nccl/deepep/moe_a2a_utils.h index 7c2019d2e7..e2d8072ef5 100644 --- a/src/turbomind/comm/nccl/deepep/moe_a2a_utils.h +++ b/src/turbomind/comm/nccl/deepep/moe_a2a_utils.h @@ -46,14 +46,19 @@ void invokeMoeA2AMapping(int* f2n, int num_local_experts, cudaStream_t stream); -// Merge the routed result with the in-place shared FFN result: -// output = routed + shared_scale * sigmoid(shared_scales) * output +// Merge the routed result with the shared FFN rows and write back to output. +// This call is ALSO the routed write-back: when `shared` is empty (nullptr) +// it degrades to output = routed, so callers must keep it unconditional. // -// shared_scales is optional. When it is null, shared_scale is applied directly. +// output = routed + sigmoid(shared_scales) * shared (shared present) +// output = routed (shared == nullptr) +// +// shared_scales is optional (gate-less shared experts); when null, the shared +// term carries weight 1. void invokeMoeA2ASharedCombine(Tensor& output, // const Tensor& routed, + const Tensor& shared, const float* shared_scales, - float shared_scale, cudaStream_t stream); } // namespace turbomind diff --git a/src/turbomind/comm/nccl/deepep/symm_ctx.h b/src/turbomind/comm/nccl/deepep/symm_ctx.h index f7c8c17a6e..2fc9b060f8 100644 --- a/src/turbomind/comm/nccl/deepep/symm_ctx.h +++ b/src/turbomind/comm/nccl/deepep/symm_ctx.h @@ -56,16 +56,22 @@ struct NCCLSymmetricMemoryContext { ncclCommProperties props = NCCL_COMM_PROPERTIES_INITIALIZER; NCCLCHECK(ncclCommQueryProperties(comm, &props)); ncclDevCommRequirements_t reqs = NCCL_DEV_COMM_REQUIREMENTS_INITIALIZER; + // Rail GIN is only needed for RDMA scale-out signaling; single-node + // runs (no scale-out) work without it. if (num_ranks > 1) { - TM_CHECK(props.railedGinType != NCCL_GIN_TYPE_NONE); - - reqs.ginContextCount = num_allocated_qps_; - reqs.ginExclusiveContexts = true; - reqs.ginQueueDepth = kGinQPDepth; - reqs.ginTrafficClass = GetEnv(); - // Customized RDMA barrier needs extra signals - reqs.ginSignalCount = num_ranks + 2 * 2; - reqs.ginConnectionType = NCCL_GIN_CONNECTION_RAIL; + if (props.railedGinType != NCCL_GIN_TYPE_NONE) { + reqs.ginContextCount = num_allocated_qps_; + reqs.ginExclusiveContexts = true; + reqs.ginQueueDepth = kGinQPDepth; + reqs.ginTrafficClass = GetEnv(); + // Customized RDMA barrier needs extra signals + reqs.ginSignalCount = num_ranks + 2 * 2; + reqs.ginConnectionType = NCCL_GIN_CONNECTION_RAIL; + } + else { + TM_LOG_WARN("NCCL rail GIN is unavailable; proceeding without GIN contexts. " + "Single-node (no RDMA) runs do not need GIN."); + } } #if NCCL_VERSION_CODE >= NCCL_VERSION(2, 31, 0) reqs.useRuntimeVersion = true; diff --git a/src/turbomind/kernels/gemm/moe_utils_v2.cu b/src/turbomind/kernels/gemm/moe_utils_v2.cu index 11a60aee7b..44df2649d8 100644 --- a/src/turbomind/kernels/gemm/moe_utils_v2.cu +++ b/src/turbomind/kernels/gemm/moe_utils_v2.cu @@ -1396,28 +1396,41 @@ void invokeMoeDispatchScales(Ref out_, } template -__global__ void MoeReduceKernel(T* dst, // [ n, d] +__global__ void MoeReduceKernel(T* dst, // [ n, d] write-only const T* src, // [e*n, d] const T* bias, // [ E, d] const float* scales, // [ e, n] const int* en2f, // [ e, n] :: (e,n) -> e*n const int* f2E, // [ e* n] - const float* dst_scales, // [n] + const float* dst_scales, // [n] INERT without shared_src + const T* shared_src, // [ n, d] or nullptr + int shared_begin, + int shared_end, int dim, int tokens, - T bscale, - float dst_scale) + T bscale) { if constexpr (TURBOMIND_ARCH_DTYPE_GUARD(data_type_v)) { const int64_t ti = blockIdx.x; dst += (int64_t)dim * ti; + // 1.f when a shared expert exists (every in-tree caller guards the + // whole shared block on moe.shared); sigmoid(gate) folds in below. + float dst_scale = 1.f; if (dst_scales) { const float scale = dst_scales[ti]; dst_scale *= fdividef(1.f, 1.f + expf(-scale)); } + // Block-uniform ownership test: the null pointer carries "not owned". + // Unowned rows contribute no shared term — the same net effect as + // today's zero-filled complement rows, without any Clear. + const T* shared_row = nullptr; + if (shared_src && ti >= shared_begin && ti < shared_end) { + shared_row = shared_src + (int64_t)dim * ti; + } + // Should be warp uniforms const T* src_[exp_k]{}; const T* bias_[exp_k]{}; @@ -1440,9 +1453,9 @@ __global__ void MoeReduceKernel(T* dst, // [ n, d] for (int i = threadIdx.x * vec_size; i < dim; i += block_dim * vec_size) { Array accum{}; - if (dst_scale) { + if (shared_row) { Vec v; - Load(v, &dst[i]); + Load(v, &shared_row[i]); using namespace ops; accum = cast(v) * dst_scale; } @@ -1478,11 +1491,13 @@ void invokeMoeReduce(T* dst, const int* en2f, const int* f2E, const float* dst_scales, + const T* shared_src, + int shared_begin, + int shared_end, int tokens, int experts_per_token, int dim, T bscale, - float dst_scale, cudaStream_t st) { const auto invoke = [&](auto e) { @@ -1497,10 +1512,12 @@ void invokeMoeReduce(T* dst, en2f, f2E, dst_scales, + shared_src, + shared_begin, + shared_end, dim, tokens, - bscale, - dst_scale); + bscale); TM_CUDA_CHECK(cudaGetLastError()); }; @@ -1538,13 +1555,22 @@ void invokeMoeCombine(Ref out_, const float* dst_scales, int experts_per_token, float bscale, - float dst_scale, + const Tensor& shared_src, + int shared_begin, + int shared_end, cudaStream_t st) { auto& out = out_.get(); const int tokens = out.shape(0); TM_CHECK_EQ(src.shape(0), tokens * experts_per_token); + // dst_scales only scales the shared-base term; passing gate logits with + // no shared source is almost certainly a caller bug. + TM_CHECK(!dst_scales || shared_src); + if (shared_src) { + TM_CHECK_EQ(shared_src.shape(0), tokens); + TM_CHECK_EQ(shared_src.shape(1), src.shape(1)); + } auto invoke = [&](auto has_bias, auto t) { using T = decltype(t); @@ -1555,11 +1581,13 @@ void invokeMoeCombine(Ref out_, en2f, f2E, dst_scales, + shared_src.data_or((T*)nullptr), + shared_begin, + shared_end, tokens, experts_per_token, src.shape(1), (T)bscale, - dst_scale, st); }; diff --git a/src/turbomind/kernels/gemm/moe_utils_v2.h b/src/turbomind/kernels/gemm/moe_utils_v2.h index 28cc7f6213..32548169e6 100644 --- a/src/turbomind/kernels/gemm/moe_utils_v2.h +++ b/src/turbomind/kernels/gemm/moe_utils_v2.h @@ -73,6 +73,10 @@ void invokeMoeDispatchScales(Ref out_, // const int* num_valid_tokens, cudaStream_t st); +// dst is write-only. dst_scales ([tokens] gate logits) scales the shared-base +// term only and is inert without shared_src (TM_CHECK'd). shared_src / +// shared_begin / shared_end describe the rank-owned rows of the dense shared +// expert output; nullptr/0/0 means "no shared expert". void invokeMoeCombine(Ref out_, const Tensor& src, const Tensor& bias, @@ -82,7 +86,9 @@ void invokeMoeCombine(Ref out_, const float* dst_scales, int experts_per_token, float bscale, - float dst_scale, + const Tensor& shared_src, + int shared_begin, + int shared_end, cudaStream_t st); void invokeMoeSoftmaxMaskTopKGroups( diff --git a/src/turbomind/models/llama/context.h b/src/turbomind/models/llama/context.h index 3904b4beec..be0f860c57 100644 --- a/src/turbomind/models/llama/context.h +++ b/src/turbomind/models/llama/context.h @@ -28,6 +28,7 @@ struct Communicators { int d_cp_group; int d_dp_group; int d_mlp_group; + int d_node_group; }; // Execution context for the model diff --git a/src/turbomind/models/llama/moe_ffn_layer.cc b/src/turbomind/models/llama/moe_ffn_layer.cc index f32127d3c4..7f93b26b2a 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.cc +++ b/src/turbomind/models/llama/moe_ffn_layer.cc @@ -16,6 +16,7 @@ #include "src/turbomind/comm/nccl/deepep/token_dispatcher.h" #endif +#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" @@ -30,20 +31,17 @@ namespace turbomind { class MoeFfnLayerImpl { public: - explicit MoeFfnLayerImpl(const Context& ctx): linear_(*ctx.linear) {} + explicit MoeFfnLayerImpl(const Context& ctx): linear_(*ctx.linear), ffn_(ctx) {} virtual ~MoeFfnLayerImpl() = default; virtual void Forward(MoeFfnLayer::ForwardParam& p) = 0; - virtual void Combine(MoeFfnLayer::ForwardParam& p) = 0; - - virtual Tensor GetShardFfnInput(Tensor& global_hidden_states, const std::vector& local_token_nums) = 0; - protected: Tensor_ Gate(const Tensor& input, const LinearWeight& gate); - LlamaLinear& linear_; + LlamaLinear& linear_; + LlamaFfnLayer ffn_; // dense executor for the shared expert }; Tensor_ MoeFfnLayerImpl::Gate(const Tensor& input, const LinearWeight& gate) @@ -63,10 +61,6 @@ class MoeFfnDefaultImpl final: public MoeFfnLayerImpl { void Forward(MoeFfnLayer::ForwardParam& p) override; - void Combine(MoeFfnLayer::ForwardParam& p) override; - - Tensor GetShardFfnInput(Tensor& global_hidden_states, const std::vector& local_token_nums) override; - private: void Init(const MoeWeight& weights); @@ -85,16 +79,6 @@ class MoeFfnDefaultImpl final: public MoeFfnLayerImpl { Buffer_ scales_; Buffer_ accum_; Buffer_ offsets_; - - 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, const MoeWeight& weights): @@ -105,11 +89,6 @@ 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(); - } Init(weights); } @@ -245,13 +224,13 @@ void MoeFfnDefaultImpl::Forward(MoeFfnLayer::ForwardParam& p) } } - temp_ = Tensor{{tokens * experts_per_token, hidden_dim}, p.input.dtype(), p.input.device()}; + Tensor temp{{tokens * experts_per_token, hidden_dim}, p.input.dtype(), p.input.device()}; // Masked-out tokens compact the routing tables even for ep_size == 1, so the // routed count is device-side offsets_[local_expert_num] in all cases. const int* num_valid_tokens = offsets_.data() + local_expert_num; - auto indices = f2n_.slice(0, temp_.shape(0)); + auto indices = f2n_.slice(0, temp.shape(0)); auto offsets = offsets_.slice(0, local_expert_num + 1); if (block.w1w3) { @@ -262,7 +241,7 @@ void MoeFfnDefaultImpl::Forward(MoeFfnLayer::ForwardParam& p) TM_SCOPE_CALL( linear_.Forward(p.input, unused_in_scales, *block.w1w3, indices, offsets, inter, inter_scales)); TM_SCOPE_CALL(linear_.Forward( - inter.slice({0, 0}, {-1, inter_size}), inter_scales, *block.w2, {}, offsets, temp_, unused_in_scales)); + inter.slice({0, 0}, {-1, inter_size}), inter_scales, *block.w2, {}, offsets, temp, unused_in_scales)); } else { TM_SCOPE_CALL(linear_.Forward(p.input, *block.w1w3, indices, offsets, inter)); @@ -271,7 +250,7 @@ void MoeFfnDefaultImpl::Forward(MoeFfnLayer::ForwardParam& p) TM_SCOPE_CALL(Activation(inter, block.w1w3->bias, f2E_, block.act_type, num_valid_tokens, st)); } - TM_SCOPE_CALL(linear_.Forward(inter.slice({0, 0}, {-1, inter_size}), *block.w2, {}, offsets, temp_)); + TM_SCOPE_CALL(linear_.Forward(inter.slice({0, 0}, {-1, inter_size}), *block.w2, {}, offsets, temp)); } } else { @@ -284,73 +263,45 @@ void MoeFfnDefaultImpl::Forward(MoeFfnLayer::ForwardParam& p) TM_SCOPE_CALL(Activation(gating, up, block.act_type, num_valid_tokens, st)); - TM_SCOPE_CALL(linear_.Forward(gating, *block.w2, {}, offsets, temp_)); + 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; - - if (clear_pending_) { - core::Context::stream().Wait(clear_done_event_); - clear_pending_ = false; + // Shared expert: dense FFN on a private buffer, folded into the combine. + Tensor shared_out; + Tensor_ shared_logits; + int shared_begin = 0, shared_end = 0; + if (moe.shared) { + if (moe.shared_gate) { + shared_logits = Gate(p.input, *moe.shared_gate); + } + shared_out = {{tokens, hidden_dim}, p.input.dtype(), p.input.device()}; + // Sequence parallel under EP (geometry identical to the retired + // GetShardFfnInput); degenerates to the full batch when ep_size_ == 1. + const int per = tokens / ep_size_, rem = tokens % ep_size_; + const int num = per + (ep_rank_ < rem); + shared_begin = ep_rank_ * per + std::min(ep_rank_, rem); + shared_end = shared_begin + num; + if (num > 0) { + ffn_.forward( + {p.input.slice(shared_begin, num), shared_out.slice(shared_begin, num), moe.shared.get(), p.layer_id}); + } + // No Clear: the fold reads only [shared_begin, shared_end), which the + // FFN just wrote. num == 0 ranks pass begin == end and own nothing. } TM_SCOPE_CALL(invokeMoeCombine(p.output, - temp_, - TM_CHECK_NOTNULL(moe.block())->w2->bias, + temp, + block.w2->bias, scales_.data(), en2f_.data(), f2E_.data(), - shared_scales_.data_or((float*)nullptr), - moe.experts_per_token, + shared_logits.data_or((float*)nullptr), + experts_per_token, 1.f / tp_size_, - p.scale, + shared_out, + shared_begin, + shared_end, core::Context::stream().handle())); - - temp_ = {}; - shared_scales_ = {}; -} - -Tensor MoeFfnDefaultImpl::GetShardFfnInput(Tensor& global_hidden_states, const std::vector&) -{ - TM_FUNCTION_SCOPE(); - if (ep_size_ == 1) { - 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_; - 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 || 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); } #ifdef USE_NCCL @@ -361,10 +312,6 @@ class MoeFfnA2AImpl final: public MoeFfnLayerImpl { void Forward(MoeFfnLayer::ForwardParam& p) override; - void Combine(MoeFfnLayer::ForwardParam& p) override; - - Tensor GetShardFfnInput(Tensor& global_hidden_states, const std::vector& local_token_nums) override; - private: void Init(const MoeWeight& weights); @@ -388,13 +335,6 @@ class MoeFfnA2AImpl final: public MoeFfnLayerImpl { Buffer_ topk_weights_; Buffer_ topk_indices_; - Tensor input_; - Tensor output_; - Tensor temp_; - Tensor_ shared_scales_; - MoeA2AInputPartition partition_{}; - int max_token_num_per_rank_{}; - Stream scales_copy_stream_; Event scales_copy_ready_event_; Event scales_copy_done_event_; @@ -465,10 +405,10 @@ void MoeFfnA2AImpl::CompileKernels(const MoeWeight& weights) Tensor input = {{token_num, weights.block()->hidden_dim}, kBfloat16, kDEVICE}; Tensor topk_indices = {topk_indices_, {token_num, weights.experts_per_token}}; Tensor topk_weights = {topk_weights_, {token_num, weights.experts_per_token}}; - Tensor scales_t, out_moe; + Tensor scales_t, output, out_moe; TM_SCOPE_CALL(dispatcher_->Dispatch( - input, topk_indices, topk_weights, token_num, output_, scales_t, f2n_, f2E_, en2f_, offsets_)); - TM_SCOPE_CALL(dispatcher_->Combine(output_, out_moe)); + input, topk_indices, topk_weights, token_num, output, scales_t, f2n_, f2E_, en2f_, offsets_)); + TM_SCOPE_CALL(dispatcher_->Combine(output, out_moe)); } } @@ -476,18 +416,18 @@ void MoeFfnA2AImpl::Forward(MoeFfnLayer::ForwardParam& p) { TM_FUNCTION_SCOPE(); - partition_ = GetInputPartition(p.input, p.local_token_num); - TM_CHECK_LE(partition_.max_tokens_per_rank, max_token_per_rank_num_); - input_ = p.input.slice(partition_.begin, partition_.size); + const MoeA2AInputPartition partition = GetInputPartition(p.input, p.local_token_num); + TM_CHECK_LE(partition.max_tokens_per_rank, max_token_per_rank_num_); + Tensor input = p.input.slice(partition.begin, partition.size); // slice token mask for current rank - const bool* token_mask = TM_CHECK_NOTNULL(p.token_mask) + partition_.begin; + const bool* token_mask = TM_CHECK_NOTNULL(p.token_mask) + partition.begin; // deepep jit treat the max_tokens_per_rank parameter as template parameter, round it to the nearest power of two to // reduce the number of functions that need to be compiled. - max_token_num_per_rank_ = std::min(CeilPowerOfTwo(partition_.max_tokens_per_rank), max_token_per_rank_num_); + const int max_token_num_per_rank = std::min(CeilPowerOfTwo(partition.max_tokens_per_rank), max_token_per_rank_num_); const auto& moe = *p.weights; const auto& block = *TM_CHECK_NOTNULL(moe.block()); - const int token_num = partition_.size; + const int token_num = partition.size; const int hidden_dim = block.hidden_dim; const int inter_size = block.inter_size; const int expert_num = moe.num_experts(); @@ -495,7 +435,7 @@ void MoeFfnA2AImpl::Forward(MoeFfnLayer::ForwardParam& p) const int experts_per_token = moe.experts_per_token; const auto st = core::Context::stream().handle(); - Tensor_ logits = token_num ? Gate(input_, *moe.gate.get()) : Tensor_{{0, expert_num}, kDEVICE}; + Tensor_ logits = token_num ? Gate(input, *moe.gate.get()) : Tensor_{{0, expert_num}, kDEVICE}; if (p.weights->topk_method == "noaux_tc") { TM_CHECK_EQ(moe.n_group, 1); @@ -548,22 +488,23 @@ void MoeFfnA2AImpl::Forward(MoeFfnLayer::ForwardParam& p) Tensor topk_indices = {topk_indices_, {token_num, experts_per_token}}; Tensor topk_weights = {topk_weights_, {token_num, experts_per_token}}; Tensor scales_t; + Tensor output; // dispatcher-owned view; not an allocation of ours TM_SCOPE_CALL(dispatcher_->Dispatch( - input_, topk_indices, topk_weights, max_token_num_per_rank_, output_, scales_t, f2n_, f2E_, en2f_, offsets_)); + input, topk_indices, topk_weights, max_token_num_per_rank, output, scales_t, f2n_, f2E_, en2f_, offsets_)); // transpose scales to [experts_per_token, token_num] for later combine Tensor scales_src = scales_t.t(); - Tensor scales_dst = {scales_, {experts_per_token, output_.shape(0)}}; + Tensor scales_dst = {scales_, {experts_per_token, output.shape(0)}}; auto& stream = core::Context::stream(); scales_copy_ready_event_.Record(stream); scales_copy_stream_.Wait(scales_copy_ready_event_); core::GenericCopy(scales_src, scales_dst, scales_copy_stream_.handle()); scales_copy_done_event_.Record(scales_copy_stream_); - temp_ = {{output_.shape(0) * p.weights->experts_per_token, hidden_dim}, output_.dtype(), output_.device()}; + Tensor temp{{output.shape(0) * p.weights->experts_per_token, hidden_dim}, output.dtype(), output.device()}; const int* num_valid_tokens = offsets_.data() + local_expert_num; - auto indices = f2n_.slice(0, temp_.shape(0)); + auto indices = f2n_.slice(0, temp.shape(0)); auto offsets = offsets_.slice(0, local_expert_num + 1); if (block.w1w3) { @@ -572,79 +513,67 @@ void MoeFfnA2AImpl::Forward(MoeFfnLayer::ForwardParam& p) Tensor unused_in_scales; if (block.is_fused_silu && block.w1w3->output_dtype() == kFloat8_e4m3) { TM_SCOPE_CALL( - linear_.Forward(output_, unused_in_scales, *block.w1w3, indices, offsets, inter, inter_scales)); + linear_.Forward(output, unused_in_scales, *block.w1w3, indices, offsets, inter, inter_scales)); TM_SCOPE_CALL(linear_.Forward( - inter.slice({0, 0}, {-1, inter_size}), inter_scales, *block.w2, {}, offsets, temp_, unused_in_scales)); + inter.slice({0, 0}, {-1, inter_size}), inter_scales, *block.w2, {}, offsets, temp, unused_in_scales)); } else { - TM_SCOPE_CALL(linear_.Forward(output_, *block.w1w3, indices, offsets, inter)); + TM_SCOPE_CALL(linear_.Forward(output, *block.w1w3, indices, offsets, inter)); if (!block.is_fused_silu) { TM_SCOPE_CALL(Activation(inter, block.w1w3->bias, f2E_, block.act_type, num_valid_tokens, st)); } - TM_SCOPE_CALL(linear_.Forward(inter.slice({0, 0}, {-1, inter_size}), *block.w2, {}, offsets, temp_)); + 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(output_, *block.w1, indices, offsets, gating)); + TM_SCOPE_CALL(linear_.Forward(output, *block.w1, indices, offsets, gating)); Tensor up; - TM_SCOPE_CALL(linear_.Forward(output_, *block.w3, indices, offsets, up)); + TM_SCOPE_CALL(linear_.Forward(output, *block.w3, indices, offsets, up)); TM_SCOPE_CALL(Activation(gating, up, block.act_type, num_valid_tokens, st)); - TM_SCOPE_CALL(linear_.Forward(gating, *block.w2, {}, offsets, temp_)); + TM_SCOPE_CALL(linear_.Forward(gating, *block.w2, {}, offsets, temp)); } - if (moe.shared_gate && token_num > 0) { - shared_scales_ = Gate(input_, *moe.shared_gate); + // Shared expert: dense FFN on a private buffer. Declarations live + // outside the guard — the tail call below is unconditional. + Tensor_ shared_logits; + Tensor shared_out; + if (moe.shared && token_num > 0) { + if (moe.shared_gate) { + shared_logits = Gate(input, *moe.shared_gate); + } + shared_out = {{token_num, hidden_dim}, input.dtype(), input.device()}; + ffn_.forward({input, shared_out, moe.shared.get(), p.layer_id}); } -} - -void MoeFfnA2AImpl::Combine(MoeFfnLayer::ForwardParam& p) -{ - TM_FUNCTION_SCOPE(); - const auto& moe = *p.weights; - const auto& block = *TM_CHECK_NOTNULL(moe.block()); - auto& stream = core::Context::stream(); stream.Wait(scales_copy_done_event_); - TM_SCOPE_CALL(invokeMoeCombine(output_, - temp_, + // The A2A impl never folds via the reduce — its shared fold happens in + // invokeMoeA2ASharedCombine below — so the shared args are null/zero. + TM_SCOPE_CALL(invokeMoeCombine(output, + temp, block.w2->bias, scales_.data(), en2f_.data(), f2E_.data(), nullptr, - moe.experts_per_token, + experts_per_token, 1.f / mlp_tp_size_, - 0.f, + Tensor{}, + 0, + 0, stream.handle())); Tensor out_moe; - TM_SCOPE_CALL(dispatcher_->Combine(output_, out_moe)); + TM_SCOPE_CALL(dispatcher_->Combine(output, out_moe)); + // Unconditional: with no shared expert this degrades to input = out_moe, + // which is the only place the MoE output lands in the residual stream. TM_SCOPE_CALL( - invokeMoeA2ASharedCombine(input_, out_moe, shared_scales_.data_or((float*)nullptr), p.scale, stream.handle())); - - input_ = {}; - output_ = {}; - temp_ = {}; - shared_scales_ = {}; - partition_ = {}; -} - -Tensor MoeFfnA2AImpl::GetShardFfnInput(Tensor& global_hidden_states, const std::vector& local_token_nums) -{ - TM_FUNCTION_SCOPE(); - - if (partition_.max_tokens_per_rank != 0) { - return global_hidden_states.slice(partition_.begin, partition_.size); - } - - const auto partition = GetInputPartition(global_hidden_states, local_token_nums); - return global_hidden_states.slice(partition.begin, partition.size); + invokeMoeA2ASharedCombine(input, out_moe, shared_out, shared_logits.data_or((float*)nullptr), stream.handle())); } #endif @@ -669,19 +598,9 @@ MoeFfnLayer::MoeFfnLayer(const EngineParam& engine, const Context& ctx, const Mo MoeFfnLayer::~MoeFfnLayer() = default; -void MoeFfnLayer::Forward(ForwardParam& p) +void MoeFfnLayer::Forward(ForwardParam p) { impl_->Forward(p); } -void MoeFfnLayer::Combine(ForwardParam& p) -{ - impl_->Combine(p); -} - -Tensor MoeFfnLayer::GetShardFfnInput(Tensor& global_hidden_states, const std::vector& local_token_nums) -{ - return impl_->GetShardFfnInput(global_hidden_states, local_token_nums); -} - } // namespace turbomind diff --git a/src/turbomind/models/llama/moe_ffn_layer.h b/src/turbomind/models/llama/moe_ffn_layer.h index ffb016a5ab..efe4ca56b5 100644 --- a/src/turbomind/models/llama/moe_ffn_layer.h +++ b/src/turbomind/models/llama/moe_ffn_layer.h @@ -24,16 +24,11 @@ class MoeFfnLayer { Tensor output; std::vector local_token_num; const MoeWeight* weights; - float scale; int layer_id; const bool* token_mask; // [tokens]; invalid tokens route nowhere }; - void Forward(ForwardParam& p); - - void Combine(ForwardParam& p); - - Tensor GetShardFfnInput(Tensor& global_hidden_states, const std::vector& local_token_nums); + void Forward(ForwardParam p); private: std::unique_ptr impl_; diff --git a/src/turbomind/models/llama/unified_decoder.cc b/src/turbomind/models/llama/unified_decoder.cc index 166dcd9b21..3eee72c496 100644 --- a/src/turbomind/models/llama/unified_decoder.cc +++ b/src/turbomind/models/llama/unified_decoder.cc @@ -1,7 +1,6 @@ #include -#include #include @@ -49,6 +48,7 @@ UnifiedDecoder::UnifiedDecoder(CacheRegistry& registry, mlp_tp_size_(engine.mlp_tp_size), attn_tp_group_(ctx.comm.d_tp_group), mlp_group_(ctx.comm.d_mlp_group), + node_group_(ctx.comm.d_node_group), d_comm_(ctx.comm.d_comm), tune_layer_num_(engine.tune_layer_num), is_warm_up_{*ctx.is_warm_up} @@ -78,6 +78,23 @@ UnifiedDecoder::UnifiedDecoder(CacheRegistry& registry, moe_ffn_layer_ = std::make_unique(engine, ctx, moe_weights.front()); } + // Per-layer FFN group, used on both sides of the FFN: the pre-FFN gather + // assembles exactly the chunks the group's members own, and the post-FFN + // reduce sums exactly those chunks. Everything runs over mlp_group_ (the + // MoE combine excludes the ep dimension; pure-dense models get the + // historical whole-domain allreduce) — except dense layers in a MoE + // model, whose weights shard node-locally (Python: _dense_tp) and whose + // node collectively holds every row it reduces, so they never pay + // cross-node traffic. + ffn_group_.assign(model_weight.num_layer, mlp_group_); + if (moe_ffn_layer_) { + for (int i = 0; i < model_weight.num_layer; ++i) { + if (!model_weight.layer(i)->moe_ffn) { + ffn_group_[i] = node_group_; + } + } + } + if (!ffn_weights.empty()) { ffn_layer_ = std::make_unique(ctx); } @@ -278,6 +295,10 @@ void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vectorattention->wo->bias; } + // Per-layer FFN group, precomputed in the ctor — the same group is + // used for the pre-FFN gather (here) and the post-FFN reduce. + const int ffn_group = ffn_group_[layer]; + AllreduceResidualRMSnorm(global_hidden_states, local_residual, out_bias, @@ -286,7 +307,7 @@ void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vectorffn_norm->zero_centered_, local_token_num, attn_tp_group_, - mlp_group_, + ffn_group, local_token_nums.data(), local_token_nums.size()); @@ -296,31 +317,24 @@ 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, - local_token_nums, - weights.at(layer)->moe_ffn.get(), - weights.at(layer)->feed_forward ? 1.f : 0.f, - layer, - (const bool*)args.at("token_mask").buffer().raw_data()}; - moe_ffn_layer_->Forward(*moe_fwd_param); + moe_ffn_layer_->Forward({global_hidden_states, + global_hidden_states, + local_token_nums, + weights.at(layer)->moe_ffn.get(), + (int)layer, + (const bool*)args.at("token_mask").buffer().raw_data()}); } if (ffn_layer_ && weights.at(layer)->feed_forward) { - auto ffn_input_shared = moe_ffn_layer_ ? - moe_ffn_layer_->GetShardFfnInput(global_hidden_states, local_token_nums) : - 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) { - moe_ffn_layer_->Combine(*moe_fwd_param); + // Staging invariant: a layer with both slots means old Python + // wiring still parks the shared expert in feed_forward while the + // C++ side already reduced into global_hidden_states — fail + // loudly instead of running the shared FFN on the combine result. + TM_CHECK(!weights.at(layer)->moe_ffn) << "layer " << layer << " has both moe_ffn and feed_forward — " + << "the shared expert belongs in moe.shared"; + ffn_layer_->forward( + {global_hidden_states, global_hidden_states, weights.at(layer)->feed_forward.get(), (int)layer}); } TM_DEBUG_TENSOR(global_hidden_states, Concat("ffn_block", layer), 2); @@ -338,7 +352,7 @@ void UnifiedDecoder::Forward(int phase, TensorMap& args, const std::vectorffn_norm->norm_eps_, scale_zero_centered, local_token_num, - mlp_group_, + ffn_group, attn_tp_group_, local_token_nums.data(), local_token_nums.size()); diff --git a/src/turbomind/models/llama/unified_decoder.h b/src/turbomind/models/llama/unified_decoder.h index f016d6d038..651e5a9f5d 100644 --- a/src/turbomind/models/llama/unified_decoder.h +++ b/src/turbomind/models/llama/unified_decoder.h @@ -40,6 +40,10 @@ class UnifiedDecoder { const int attn_tp_group_; const int mlp_group_; + const int node_group_; + + // Per-layer post-FFN reduce group, precomputed in the constructor. + std::vector ffn_group_; comm::DeviceCommImpl* const d_comm_; diff --git a/src/turbomind/models/moe_weight.h b/src/turbomind/models/moe_weight.h index 4be714d53e..ea7a29328b 100644 --- a/src/turbomind/models/moe_weight.h +++ b/src/turbomind/models/moe_weight.h @@ -74,6 +74,7 @@ class MoeWeight: public core::Module { #define MOE_WEIGHT_CHILDREN(X) \ X(LinearWeight, gate) \ X(LinearWeight, shared_gate) \ + X(FfnWeight, shared) \ X(core::ModuleList, experts) #define MOE_WEIGHT_PARAMS(X) X(score_correction_bias) diff --git a/src/turbomind/python/linear_bind.cpp b/src/turbomind/python/linear_bind.cpp index 118546e87a..dc3ef56a89 100644 --- a/src/turbomind/python/linear_bind.cpp +++ b/src/turbomind/python/linear_bind.cpp @@ -435,12 +435,19 @@ void bind_linear(py::module_& m) std::shared_ptr dst_scales, int experts_per_token, float bscale, - float dst_scale) { + std::shared_ptr shared_src, + int shared_begin, + int shared_end) { core::Tensor o = TensorOrEmpty(out); const float* scales_ptr = (scales && *scales) ? scales->data() : nullptr; const int* en2f_ptr = (en2f && *en2f) ? (const int*)en2f->raw_data() : nullptr; const int* f2E_ptr = (f2E && *f2E) ? (const int*)f2E->raw_data() : nullptr; const float* dst_scales_ptr = (dst_scales && *dst_scales) ? dst_scales->data() : nullptr; + // dst is write-only; dst_scales (gate logits) scales the + // shared-base term only, and is inert without shared_src. + // shared_src folds over [shared_begin, shared_end) — passing + // shared_src without an explicit range (defaults 0, 0) folds + // nothing. dtype mismatches abort in Buffer::data(). invokeMoeCombine(o, TensorFromShared(src, "src"), TensorOrEmpty(bias), @@ -450,7 +457,9 @@ void bind_linear(py::module_& m) dst_scales_ptr, experts_per_token, bscale, - dst_scale, + shared_src ? TensorOrEmpty(shared_src) : core::Tensor{}, + shared_begin, + shared_end, DefaultStream()); return std::make_shared(o); }, @@ -462,8 +471,10 @@ void bind_linear(py::module_& m) py::arg("f2E") = py::none(), py::arg("dst_scales") = py::none(), py::arg("experts_per_token"), - py::arg("bscale") = 1.f, - py::arg("dst_scale") = 0.f, + py::arg("bscale") = 1.f, + py::arg("shared_src") = py::none(), + py::arg("shared_begin") = 0, + py::arg("shared_end") = 0, py::call_guard()); } diff --git a/src/turbomind/turbomind.cc b/src/turbomind/turbomind.cc index e384fff205..c73d4b3ce4 100644 --- a/src/turbomind/turbomind.cc +++ b/src/turbomind/turbomind.cc @@ -245,9 +245,10 @@ void TurboMind::Impl::CreateContext(int index) if (comm_size_ > 1) { c.d_comm = CreateDeviceCommunicator(communicator_type_, comm_size_, inner_rank, c.h_comm); - c.d_tp_group = 0; - c.d_cp_group = 0; - c.d_dp_group = 0; + c.d_tp_group = 0; + c.d_cp_group = 0; + c.d_dp_group = 0; + c.d_node_group = 0; if (p.attn_dp_size > 1) { // has attn_dp c.d_tp_group = c.d_comm->Split(tp_color, 0, 0); @@ -264,6 +265,19 @@ void TurboMind::Impl::CreateContext(int index) 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.mlp_tp_size; + + // Node domain: ranks of one node within the comm domain. Dense FFN + // shards node-locally (Python: dense_tp), so its partials reduce + // within the node and never pay cross-node traffic. When the comm + // domain fits in one node, group 0 — the whole domain — serves. + const int node_size = (int)p.devices.size(); + if (comm_size_ > node_size) { + // Node-local dense sharding requires the comm domain to tile the + // node: shard index is inner_rank % node_size, so a partial node + // tail would reduce an incomplete shard set. + TM_CHECK(comm_size_ % node_size == 0); + c.d_node_group = c.d_comm->Split(inner_rank / node_size, 0, 0); + } // Layout: (outer, ep, mlp_tp) c.d_mlp_group = 0; if (p.ep_size > 1) {