Intern-S2-Mobius meta-MoE support, MoE gate v2, CP attention fixes - #4835
Open
lzhangzz wants to merge 9 commits into
Open
Intern-S2-Mobius meta-MoE support, MoE gate v2, CP attention fixes#4835lzhangzz wants to merge 9 commits into
lzhangzz wants to merge 9 commits into
Conversation
… GEMM MoE gate v2: - v8 kernel specializations for e256_k8 and e2560_k8, selection mask widened beyond 32 items/thread - perf: hierarchical multi-CTA scan for large-E tile counts, fused mask clear for E<=512, 2-row shared histogram - pybind bridge with optional preallocated outs; pytest coverage and steady-state latency benchmark Meta-MoE (InternS2Mobius): - InternS2Mobius Python model owning the meta-MoE loading path - encapsulated meta_moe prepare/alias unit in C++; LinearWeight::prepare idempotent for shared buffers; strided n_groups layer mapping FP8 GEMM determinism at partial M tiles: - fp8_wa indexed-gather producer: skip cp.async ZFILL for out-of-range slots at TILE_M=8 - fp8_wa/v3 fused epilogues: clamp per-row output-scale stores to global M for ungrouped GEMM (was INT_MAX, OOB past scales alloc)
MoeBuilder was constructed without the EP group, so every rank baked ep_size=1 into the MoE weight config and the donor loaded all experts while the runtime indexed rank-local slices — EP=8 produced garbage. Pass ep=text_model._ep and build the donor's local slice via m.range(). alias_routed_moe iterated the global expert_num and named dst children with 0-based local indices, but MoeWeight::expert() looks up local_expert_offset()+i. Iterate num_local_experts() and name children with the global expert index; assert donor/dst EP rank match. Verified: Mobius coherent at TP8/EP8 and CP8/EP8 (was gibberish).
For CP > WarpCnt (4), each warp iterates over multiple CP ranks (cp_i = warp_id + i * WarpCnt), but frag_M was selected by warp id alone, so ranks 0..WarpCnt-1 had their local max overwritten by a later rank's max. The wrong maximum was then used to rescale the rank's partial attention output before the all-reduce, corrupting outputs for any CP > 4 (deterministic garbage at CP=8). Select frag_M only from the iteration whose cp_i equals this rank.
In mixed prefill/decode passes the two CpPost AllGathers of a CP group could run concurrently on different streams (prefill on the aux stream, decode on the main stream). Concurrent collectives on one communicator are unordered and can pair inconsistently across ranks. CpPost now brackets each gather with events against the producing/consuming stream and runs the AllGather itself on a dedicated cp_stream owned by CpPostContext, so all gathers of the group serialize in host call order on every rank while prefill/decode compute keeps overlapping. The stream and events are only created when cp_size > 1, and CpPost asserts the stream exists.
Build a per-token validity mask and use it end to end:
- attention reduce: mask finished rows in the CP softmax reduction so
stale partial outputs never enter the AllReduce
- attention layer owns a DP-local per-token validity mask derived from
q_offsets and finished, cleared/set once per forward pass
- LanguageModel builds the mask globally: symmetric buffers for
offset_q/finished, AllGather across DP ranks, and a
ClearTokenMaskKernel (invokeBuildTokenMask) producing the flat global
mask matching global_hidden_states
- MoE gate (both invokeMoeGate_V2 and invokeMoeGate_NoAuxTC) takes a
required token_mask; invalid tokens produce no routing entries, so
scan, expert GEMM, dispatch and combine never touch them. This also
removes the garbage shared_exp_id out-of-bounds write seen with
Mobius CP8/EP8
- moe_ffn_layer: clear en2f to -1 and pass num_valid_tokens
unconditionally (EP=1 included)
- unified_attention_layer: drop tmp_attn_ entirely; attn is now a
per-pass {q_count, dim} tensor
- tests: token_mask is a required binding input, harness defaults to
all-true; add deterministic masked-NaN routing test; benchmark
preallocates the mask outside the timed loop
The checkpoint stores meta-MoE packs once at model level; the weight
tree now mirrors that instead of cloning them into every layer:
- ModelWeight gains a meta_experts ModuleList of full routed MoeWeight
packs, loaded once by the Python loader and prepared explicitly
before the layers in ModelWeight::prepare()
- Per-layer MoE weights carry only their own shared_gate and are wired
to their pack via MoeWeight::set_meta_pack (non-owning pointer, set
by the loader through a new pybind method)
- MoeWeight::prepare() aliases the pack's routed gate/expert tensors
with shallow tensor-sharing copies (AliasLinear/AliasFfn); link_block
remains the only linking path
- Delete meta_moe.{cc,h} (PrepareMetaMoe, ModelHasMetaMoe,
alias_routed_moe), the prepare fork in ModelWeight::prepare, the
is_meta_donor/meta_group config fields, and prepare_routed_linears
Runtime (moe_ffn_layer, unified_decoder) is untouched: after prepare
each layer's moe_ffn presents the same gate/experts/block/shared_gate
as before.
Verified: Intern-S2-Mobius-FP8 CP8/EP8 22/22 coherent responses,
Qwen3.5-35B-A3B TP=1 22/22 (non-meta regression).
… config The shape() accessor existed only for infer_meta_geometry to read the expert count from checkpoint tensors, but cfg.num_experts already provides it via Qwen3_5TextModel.__init__. infer_meta_groups now only counts the meta packs with has().
Absolute tests.turbomind.moe_gate imports only resolve when the repo root is on sys.path, which is not true under CI's pytest collection (ModuleNotFoundError: No module named 'tests.turbomind'). The package has __init__.py at both levels, so relative imports work under pytest's prepend import mode regardless.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds TurboMind support for Intern-S2-Mobius (
interns2_mobius), a Qwen3.5-topology checkpoint whose routed MoE weights are stored as meta-MoE packs shared across layer groups, along with the kernel and parallel-path fixes needed to run it correctly at CP8/EP8. Verified: Mobius CP8/EP8 22/22 coherent responses, Qwen3.5-35B-A3B (TP=1, CP8/EP8, dp2/tp4/ep4) 22/22 — no regressions on non-meta models.Meta-MoE weight wiring
meta_mlp.{g}.experts.*/meta_mlp.{g}.gate.weightis mapped onto model-levelmeta_experts.{g}packs, each loaded exactly once as aModelWeight::meta_expertschild (lmdeploy/turbomind/models/interns2_mobius.py).MoeWeights, prepared explicitly before the layers inModelWeight::prepare(). Per-layer MoE weights carry only their ownshared_gate/shared_expertand are wired to their pack viaMoeWeight::set_meta_pack(non-owning pointer set by the loader);MoeWeight::prepare()aliases the pack's gate/expert tensors with shallow copies, andlink_block()remains the only linking path.moe_ffn_layer,unified_decoder) is untouched: after prepare, each layer'smoe_ffnpresents the same gate/experts/block/shared_gate as a regular MoE model.MoE gate v2 + invalid-token routing
invokeMoeGate_V2/invokeMoeGate_NoAuxTCkernels with a requiredtoken_mask: invalid tokens produce no routing entries, so the scan, expert GEMM, dispatch, and combine never touch them. This also eliminates the garbageshared_exp_idout-of-bounds write that crashed Mobius at CP8/EP8.LanguageModellevel: symmetric buffers foroffset_q/finished, AllGather across DP ranks, andClearTokenMaskKernelbuilding a flat mask matchingglobal_hidden_states(DP-local fast path whenattn_dp == 1).en2fcleared to -1 and routed-count pointer passed unconditionally (EP=1 included), so activation/combine never read compacted-out rows.tmp_attn_removed; attention output is a per-pass{q_count, dim}tensor.tests/turbomind/moe_gate/(reference implementation, 20 correctness tests incl. deterministic masked-NaN routing, allocation-free benchmark harness);token_maskis a required binding input.CP attention fixes
reduce.cu).CpPostAllGathers (prefill/decode) are serialized on a dedicated CP stream — no event/stream sync on the main stream, and no stream/event creation at CP=1.Deterministic / fused-scale FP8 GEMM fixes
gemm_universal_sm90_fp8_wa.h,gemm_universal_sm90_v3.h): the ungroupedrow_endwas effectively infinite while the scale allocation only covered the logical M extent.TM_GEMM_EXPORT/TM_GEMM_IMPORTdeterministic GEMM algo cache support.Test plan
scripts/test_turbomind_model.py, 22 mixed-length prompts (/tmp/prompts_mixed.json):python -m pytest tests/turbomind/moe_gate/test_moe_gate_v2.py -q: 20 passed.