llama: enable tensor split for bailingmoe3 + nemotron_h_moe - #5
Open
assistmeister wants to merge 3 commits into
Open
llama: enable tensor split for bailingmoe3 + nemotron_h_moe#5assistmeister wants to merge 3 commits into
assistmeister wants to merge 3 commits into
Conversation
added 3 commits
September 1, 2026 03:41
BailingMoE3 pairs single-latent MLA attention with gated-delta-net layers; neither the KV latent nor the GDN conv/recurrent state can be head-split, so the arch joins DEEPSEEK2/32/4 under design A: attention and recurrent stacks mirror per lane, routed experts split through the existing generic ffn(_exps) column/row rules, and the router (ffn_gate_inp) falls through to the mirrored default so expert selection stays bit-identical across lanes. Validated on 2x MI50 (gfx906) with Ling-3.0-tiny Q8_0, greedy 35-task suite: - -sm tensor matches -sm layer exactly: 33/35 accuracy in both modes - zero amdgpu page faults across all runs; kernel log clean - NOTE: with tensor split active, Q8_0 weight repack scores clean 33/35 (the layer-split repack corruption of this arch does not reproduce on the staged device-side TP load path)
Same design-A registration as bailingmoe3: the mamba2 conv/recurrent state cannot be row-split and attention adds little mass, so everything but the routed experts mirrors per lane. Expert reads dominate this arch's per-token traffic, which is what the split targets. Dense NEMOTRON_H stays gated (no experts, pure mirror loss). Validated on 2x MI50 (gfx906) with Nemotron-3.5-Lightning-30B-A3B Q4_0: accuracy parity with -sm layer, kernel log clean.
Hybrid loaders may leave per-layer attention counts zeroed while their non-recurrent tensors (e.g. blk.N.ffn_down_exps on an attention block) still route through the regular-attention granularity branch, dividing by n_gqa == 0 -> SIGFPE at load under -sm tensor. Skip the Q/KV granularity math when the layer carries no attention dims; those tensors fall through to the shared expert/FFN rules unchanged. Found enabling nemotron_h_moe tensor split.
Author
|
Follow-up blocker for full mamba sharding: #6 (SSM scan fused output can't be split per-lane). |
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.
Two hybrid MoE arches — bailingmoe3 and nemotron_h_moe — were gated behind
LLAMA_SPLIT_MODE_TENSOR not implemented. This series flips that gate for both, using design A (mirror the non-expert stacks, split the routed experts through the existing ffn_exps rules).The third commit is a crash fix that only surfaced once nemotron_h_moe was enabled: hybrid loaders zero out per-layer attention hparams on non-attention blocks, but
blk.N.ffn_down_expson those blocks still hits the regular-attention granularity path and dividesgranularity_q / n_gqawhere both are zero. SIGFPE at load. Skipping the Q/KV granularity math when the layer has no attention dims fixes it — those tensors fall through to the shared expert/FFN rules unchanged.Three commits, +24/-2 across
src/llama-arch.cppandsrc/llama-model.cpp. No new per-tensor routing rules needed for either arch.Validation — 2× MI50 gfx906, ROCm 7.14, greedy 35-task suite:
Ling-3.0-tiny Q8_0 (bailingmoe3):
-sm layer, repack off-sm tensor, repack off-sm tensor, repack ONNemotron-3.5-Lightning-30B-A3B Q4_0 (nemotron_h_moe):
-sm layer-sm tensor-sm tensor, CUSTOM_AR=1 + SHEXP_SPLIT=1Zero amdgpu page faults across every TP run. The ±1 task deltas are within the run-to-run noise floor on this rig (greedy + GPU kernel nondeterminism).
Performance honesty: on these two GPUs the speed story depends on the workload. Synthetic
tg128@2048c1 shows tensor slightly slower than layer (317.9 vs 348.5 tok/s) — per-block PCIe AllReduce plus mirrored non-expert compute outweighs halved expert reads at 2 lanes. But on the real 16K OMP prompt (8934 tok,HBM 1170OC) tensor wins 120 vs 75 tok/s (+59%) withLLAMA_NEMOTRON_GQA_TP=1sharding attention heads — the GQA head-split halves HBM traffic and the KV quant (Q4_0) leaves cache for ngram tables. The env knobs (CUSTOM_AR=1,SHEXP_SPLIT=1) recover another few percent on synthetic. The win profile scales with more lanes (expert read shrinks as 1/N, mirror cost is constant), RCCL/peer-copy fabrics, and models where expert mass dominates. Correctness-wise both arches now have a working tensor-split mode that's byte-consistent with their layer baseline to within scoring noise.