Skip to content

feat(qwen4): add Qwen3.8-Flash-Next MoE block golden - #1154

Draft
sysuTobo wants to merge 8 commits into
pegainfer-project:mainfrom
sysuTobo:feat/qwen38fn-moe-golden
Draft

sysuTobo wants to merge 8 commits into
pegainfer-project:mainfrom
sysuTobo:feat/qwen38fn-moe-golden

Conversation

@sysuTobo

@sysuTobo sysuTobo commented Oct 12, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Reference golden for the MoE block of Qwen3.8-Flash-Next (#1103).

  • tools/accuracy/dump_qwen38fn_moe_golden.py — dumps Qwen4ExpTextSparseMoeBlock (transformers 5.16.1 qwen4_exp) from the frozen checkpoint config: router logits, renormalized top-k weights, selected experts, and the routed-expert / shared / block outputs.
  • test_data/qwen38fn-moe-golden.safetensors — 3 cases (T=8 and T=12 prefill, T=1 decode), 26 tensors, 12.3 MiB. Provenance lives in one metadata key (format_version 2).
  • Geometry: real frozen-config geometry — hidden 2560, 512 experts top-10, moe_intermediate_size 640, shared 640, and the upstream norm_topk_prob=True default that config.json does not carry. The fused 512-expert stack is NOT stored (~4.7 GiB / 5 GB bf16, 2.52e9 params): the reference's routed-expert output is recorded per case as an activation, and only the weights a gate can check directly (router gate_weight 2.6 MB, shared projections ~10 MB, scalar shared gate) are stored.
  • The data carries: the router renorm (absent from config.json), the fused-3D expert layout (recorded as an activation), and the shared expert's learned scalar sigmoid gate summed with the routed output.

Consumer

pegainfer-qwen4/tests/moe_block_replay.rs reads the fixture: a device-free replay of Qwen4ExpTextSparseMoeBlock on the host — no CUDA, no Triton, no checkpoint, so it runs in any CPU CI job and from a fresh clone (the fixture is a real 12.3 MiB blob, not LFS). half and safetensors were already in Cargo.lock; the dev-dep addition is additive (cargo metadata --locked clean).

$ cargo test --release --locked -p pegainfer-qwen4 --test moe_block_replay -- --nocapture
fixture provenance: format_version 2, config_sha256 889658f2508e8c61d409b02e70e0d78d8d4452ec65aaafbe129805d213d2e74b,
                    seed 3543236481410, transformers 5.16.1, torch 2.14.1+cpu, x86_64
fixture geometry: hidden 2560, 512 experts, top-10 routed + 1 shared, shared intermediate 640; case token counts [8, 12, 1]
...
test result: ok. 10 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.5s

Provenance is asserted on load and never skipped: the metadata key must parse, format_version must be 2, and config_sha256 must equal sha256(pegainfer-qwen4/tests/frozen_config.json) — whose text_config also supplies the geometry every tensor shape is checked against (including hidden_act == "silu"). A missing fixture or a moved config panics.

What it proves

step replay (host, f32 between bf16 boundaries) tolerance measured worst
router bf16(hidden · gateᵀ) 1 bf16 ULP 1 ULP on 8 of 10752 logits, rest bitwise
routing fp32 softmax over 512 → top-10 → renormalize bitwise bitwise on 210/210 weights
router→routing chain the same, on the replay's own bf16 logits bitwise bitwise on 210/210 weights
shared expert down(silu(gate)·up) + the scalar sigmoid gate 1 ULP + a derived floor 0.0625 abs, 0.41 of the floor
combine expert_out + sigmoid(gate_logit) · shared_out bitwise bitwise on 53760/53760

Tolerance justification — these are the reference's own rounding, not slack for a port:

  • The model is bf16, so the replay rounds where the reference rounds (both projections, the activation, the product, the down projection, the gate, the sum) and keeps f32 only between boundaries, like torch's opmath. The residual is then accumulation order alone. Measured, not assumed: an f32 sum of 2560 products lands 8 of 10752 logits 1 ULP from the golden and the rest bitwise, so 1 ULP is the router's bound; the routing weights read the golden's stored bf16 logits and reproduce all 210 bitwise, so theirs is bitwise.
  • Emulating the roundings is also the tighter choice: one unrounded f32 pass through the same four shared-branch operators lands 0.125 max abs away with only ~40% of elements bitwise, versus 0.0625 with 99.8%+ bitwise when rounding where the reference does.
  • The shared branch needs one absolute term because cancellation in the 640-term down projection makes the relative error large on near-zero outputs (worst element 89 ULP but only 0.0625 absolute). That term is derived per element from the replay's own intermediate — 2 · 2^-7 · √Σ(h·w)², one bf16 ULP of h propagated through down_proj in quadrature, doubled for sign correlation — not a magic constant. Measured worst element sits at 0.41 of it.
  • The renormalized weights must sum to 1 per token: derived bound 2^-9 ≈ 0.002 (ten bf16 roundings of values summing to 1), measured worst 0.0017, gate at 0.01; the un-renormalized sum lands 0.34–0.47 away.

The traps, each with a control that asserts the comparison FAILS

control the bug it stands for observed
control_skipping_topk_renorm_fails norm_topk_prob is absent from config.json and defaults to True upstream, so a port reading only the checkpoint config skips the top-k renormalization 210/210 weights move, up to 0.2383 abs; token-0 top-10 mass 0.53–0.66 instead of 1
control_bf16_softmax_fails dropping dtype=torch.float from the reference's softmax (a bf16 probability buffer) 56/210 weights move, up to 1 bf16 ULP
control_dropping_shared_sigmoid_gate_fails summing the raw shared output, i.e. forgetting the learned per-token scalar gate 45069/53760 block outputs move, up to 22.25 abs
control_gelu_instead_of_silu_fails gelu_pytorch_tanh (the vision config's activation) in the text MLP 41951/53760 shared outputs move, up to 1.91 abs
control_unrounded_router_logits_fails keeping the router's logits in f32 instead of storing bf16 171/210 weights move, up to 5 bf16 ULP

The last one is why the chain row exists: step 1's 1-ULP logit bound also passes an unrounded f32 logit (it sits within half an ULP), and step 2 reads the golden's stored logits by design, so without the chain a port could keep the router in f32 and clear every other gate while routing differently.

Top-k tie rule (measured, not assumed): torch.topk leaves the order of an exact tie undefined, and this fixture has ties — 5 of case0's 8 rows repeat a probability inside the top-10, and 3 of case1's 12 rows tie at the k-th boundary (row 0 takes expert 393 where index order would take 151, both at p = 0.027970612). So selected_experts is compared as a weight-matched set, not positionally: (1) the stored row is a valid top-k set — no selected expert is beaten by an unselected one; (2) its probability sequence equals the top-k's descending sequence, bitwise slot by slot; (3) each stored weight is the renormalized probability of the expert stored beside it, bitwise, which binds slot to expert; (4) the expert set equals the canonical top-k set (descending probability, ascending index inside a tie) wherever the boundary is untied. pegainfer-gemma4/src/moe/tests.rs:359 can assert its picks positionally because the other side there is its own kernel; here the other side is torch.

What remains ungated

  • expert_out is NOT recomputed — it is a recorded activation. The fused 512-expert stack (2.52e9 parameters, ~4.7 GiB bf16) is deliberately not stored, so the routed-expert GEMM itself (the fused 3D gate_up_proj / down_proj layout and its scatter/gather) has no parity check yet. That lands with the [Model] Qwen3.8-Flash-Next: 512-expert MoE #1107 implementation, which should replay this same fixture.
  • The GPU router kernel is likewise untouched: this gate is host-only and pins the contract a kernel must meet, not a kernel.
  • The gate therefore does not execute production code yet. What it owns is the part a port gets silently wrong — router precision, renorm, shared branch, gate, combine — and it holds config's norm_topk_prob default (already unit-tested at pegainfer-qwen4/src/config.rs:617) to the numbers that default implies.
  • Follow-up left out of this PR on purpose: wiring -p pegainfer-qwen4 --test moe_block_replay into the CPU unit tests job (it needs no device and takes 0.5s; pegainfer-sim --test frontend_e2e is the precedent for running one integration target there).

Also in this push: docs/models/qwen4/bring-up.md now points the norm_topk_prob paragraph at the gate that prices the trap.

Device path

The device half of this consumer now lives in pegainfer-kernels, behind a new additive qwen4 feature of that crate: the model crate still carries no CUDA dependency at all, the default build does not compile the new translation unit (cargo check --locked -p pegainfer-kernels stays green and the build script prints that the Qwen4 kernels are disabled), and the device-free suite above is unchanged and still green. This retires the two host-only caveats in "What remains ungated" above — the GPU router kernel is now gated, and the fixture now executes production code (on device, behind a feature no CI job enables yet).

Two new kernels in pegainfer-kernels/csrc/qwen4/qwen4_moe_route.cu (sm_89-common subset: no TMA, no wgmma, no cluster), FFI in src/ffi/qwen4.rs, wrappers in src/ops/qwen4.rs, gate in tests/qwen4_moe_block_parity.rs:

  • qwen4_moe_route_cuda — one block per token: fp32 softmax over all 512 logits (bf16 in, exactly the bits the router GEMM stored) → top-10 by K masked block-argmax passes in the canonical order (descending probability, ascending expert index inside a tie) → optional norm_topk_prob renorm, summing the selected probabilities in that same descending order → bf16 weights + i32 ids. Fixed shuffle trees, so the picks and the softmax denominator are deterministic run to run. A non-finite logit fails its whole row closed through a block-uniform __syncthreads_or before the softmax (NaN weights, slot-index placeholder ids) — gemma4's policy by the opposite mechanism: early return rather than flooding NaN through the reduction and catching it at the sentinel.
  • qwen4_moe_shared_combine_cuda — out = bf16(expert_out + bf16(bf16(sigmoid(gate_logit)) · shared_out)), rounding at the gate, the product and the sum where the reference rounds.

Everything else is composition, so the gate grades the production path rather than a parallel one: the router GEMM, the shared expert's three projections and its SwiGLU all go through gemm_into_checked / silu_mul_batch_into.

Reuse survey — nothing in-tree carries "fp32 softmax over 512 → top-10 → optional renorm → bf16 weights", so the routing kernel is new:

existing router its own contract why it cannot serve this one
csrc/gemma4/moe.cu (behind gemma4) 128 experts, __expf softmax, renorm unconditional, × a per-expert scale, f32 weights out expert count is a compile-time constant, there is no renorm switch, and the weights are f32
csrc/glm52/glm52_router.cu 256 experts / top-8, scores are sigmoid + bias, renorm by route_scale / Σ not a softmax; both counts hard-pinned
csrc/kimi_k2/kimi_router.cu 384 experts / top-8, e_score_correction_bias, f32 weights bias-corrected sigmoid scoring, own counts
FlashInfer FilteredTopK (via ops/glm52/topk.rs) f32 logits → top-k idx/weights no softmax, no renorm, no bf16 weights

Each would need re-templating behind another line's feature, which is neither additive nor minimal. The combine is the one near-duplicate in this landing: its two-rounding sum is add_scaled_bf16_kernel's arithmetic (csrc/shared/elementwise.cu:26-42, wrapped by ops::add_scaled_bf16_into), which takes its scale as a single host f32 uniform over the buffer and so cannot carry a per-token device sigmoid gate — reusing it would cost a D2H of the gate logits plus one launch per token (and would round the sigmoid on the host instead of where the reference rounds), or an edit to an op glm52's step body already calls. docs/models/qwen4/bring-up.md records both conclusions.

Commands (RTX 5060 sm_120, CUDA 12.9.1, pegainfer-dev:cu132):

./docker/dev.sh run cargo test --release --locked -p pegainfer-kernels \
  --features qwen4 --test qwen4_moe_block_parity -- --nocapture   # 11 passed, 0.45 s
./docker/dev.sh run cargo test --release --locked -p pegainfer-kernels --features qwen4 --lib
./docker/dev.sh run cargo clippy --release --locked -p pegainfer-kernels --features qwen4 \
  --all-targets -- -D warnings
./docker/dev.sh run cargo check --locked -p pegainfer-kernels        # feature off: the .cu is not compiled
cargo test --release --locked -p pegainfer-qwen4 --test moe_block_replay   # device-free, 10 passed

Measured, every bound derived first and then checked against the measurement (the gate prints per-case worst-abs and bound fraction):

step bound measured
router GEMM (gemm_into_checked) 1 bf16 ULP + a 0.0625 absolute floor (cuBLAS's f32 accumulation order at K=2560) worst 0.03125 abs = 0.333 of the bound; the T=1 decode row bitwise
routing on the golden's stored logits 1 ULP bitwise 210/210 weights, self-consistent with the host canonical route over the same bits
router→routing chain derived e^{2a} − 1 = 0.1332 relative, a = the router floor worst rel 0.0397 = 0.298 of the bound; bitwise against the host canonical route over the device's own logits
shared branch (production ops) 1 ULP + the derived quadrature floor × 6 headroom 0.241 / 0.290 / 0.000 of the bound for case0 / case1 / case2
combine on stored operands 2 ULP + `2^-7· shared_out
combine on the device's own branch the shared floor, gate-scaled ≤ 0.371 of the bound

The three controls stay non-vacuous on device — each drives a production entry point and asserts FAILURE: no-renorm moves 210/210 weights at 146× the tolerance; the bf16 softmax moves 56/210 (26.7%) with no row re-ordering its picks, so it is not a tie-order confound; dropping the shared gate moves 40919/53760 (76%) block outputs, up to 22.25 absolute, at 125× the bound. Two refusal contracts sit beside them: device-free geometry / null-pointer / flag refusal with an empty extent let through (the positive control against a guard that refuses everything), and on device one poisoned row failing closed while the untouched rows of the same launch still route bitwise canonical.

Fixture facts the constants are sized from — properties of the stored golden, not of the port: the smallest nonzero 10th-vs-11th probability gap over the 21 tokens is 4.157e-4 (case0 token 1), which is what puts the ordering epsilon 41× under it; exact k-th boundary ties do occur (3 rows, all in case1), which is why selected_experts is graded as a weight-matched set rather than positionally; and the worst renormalized |Σw − 1| is 0.00171, inside the derived 2^-9 ≈ 0.00195.

Still ungated: expert_out remains a recorded activation, so the routed-expert GEMM (the fused 3-D gate_up_proj / down_proj layout and its scatter/gather) still waits for #1107, which should replay this same fixture; and no CI job enables the kernels qwen4 feature yet, so the device gate is the manual command above — the device-free moe_block_replay suite remains the CI cross-check on the same fixture.

Validation

python tools/accuracy/dump_qwen38fn_moe_golden.py --out test_data/qwen38fn-moe-golden.safetensors

The dump asserts nothing — what it stores is the reference's own output at the real geometry, so there is nothing to re-check:

case0: T=8 experts_hit=78/512 max_weight=0.3887
case1: T=12 experts_hit=108/512 max_weight=0.6133
case2: T=1 experts_hit=10/512 max_weight=0.4160
wrote test_data/qwen38fn-moe-golden.safetensors (12.31 MiB, 26 tensors)

No runtime code changes: fixture, dump script, and the device-free test that consumes it (pegainfer-qwen4/tests/moe_block_replay.rs, additive dev-deps only).

Refs #1103

@github-actions github-actions Bot added the enhancement New feature or request label Oct 12, 2026
@sysuTobo

Copy link
Copy Markdown
Contributor Author

Code review

范围:2 文件(dump 脚本 162 行 + fixture 26 张量),无运行时数据流。逐文件认领:脚本对照 transformers 5.16.1 qwen4_exp 参考实现逐行核对;fixture header 解析后与脚本写入逐张量核对(26/26 名称、shape、dtype 一致)。
正向路径:语义正确已验证——router_logits 存 pre-softmax logits(#1103 要求 oracle 分开比较 pre-softmax logits / shared 分支 / routed 分支 / 合成结果,本 fixture 四项分存,dump_qwen38fn_moe_golden.py:147-152);router 重推导(:79-86)是独立实现的 fp32 softmax → top-k → renorm,与模块输出逐位比较(:132-133);sum-to-1(:139-141)是独立于镜像推导的性质断言,把 norm_topk_prob 关掉它会红,renorm 这条陷阱有真守卫;本机复跑与 committed fixture 逐字节一致。

ID 级别 问题 证据
M1 M skewed / empty experts 无 case 覆盖。#1103 MoE spike 明确要求覆盖 "skewed and empty experts";现有 2 case 随机输入下路由近均匀(实测 experts_hit=15/16、16/16),参考 eager 循环的"空专家跳过"分支与"单专家吃下全部 token"的 grouped-GEMM 形态一次都没被触到,未来 kernel gate 回放这两个 case 抓不住空专家/倾斜类错误。另无 S=1 decode 形态 case(#1103 要求 prefill and decode;MoE 块逐 token 无关位置,但单行形态是 decode kernel 的独立入口)。 tools/accuracy/dump_qwen38fn_moe_golden.py:50,153
M2 M block=experts+shared 的 atol=1e-2 无标定依据,且实测可收紧为逐位相等。两侧输入是同一批确定性 bf16 张量,模块 forward 做的正是这一步 bf16 逐元素加,实测两 case max_abs_diff=0.0,torch.equal 直接成立。1e-2 相对实测 max|final|≈1.23 约 0.8%,能放过逐位相等本可抓住的合成级偏差(gate 用错分支、小系数漂移);对照臂可以精确相等时不应使用带容差的度量。 tools/accuracy/dump_qwen38fn_moe_golden.py:135-137
L1 L logits 断言用 allclose,与 PR body "matches the module bit-for-bit" 的声明不一致。r_logits 与 logits 同式同权(实测 bit-equal=True),torch.equal 零成本且更严格;scores/selected 已用 equal,logits 单独放宽没有理由。 tools/accuracy/dump_qwen38fn_moe_golden.py:134
L2 L config_digest 漏 hidden_act。silu 是影响输出的常量:改激活函数后 digest 不变而 fixture 字节变 = 假 provenance。姊妹脚本(#1150 最终态)的 canon 收了同类影响输出的常量(rope_theta、partial_rotary_factor)。seed / case 表与 #1150 一致不收(权重与输入全部随 fixture 存储,seed 只影响再生成、不影响消费者回放),不算分歧。 tools/accuracy/dump_qwen38fn_moe_golden.py:62,67-72
L3 L 融合 3D 布局的切分语义没写进 fixture 契约。experts_gate_up_proj 是 [E, 2I, H],gate=前 I 行、up=后 I 行(参考实现 chunk(2, dim=-1) 的语义);experts_down_proj 是 [E, H, I](F.linear 的 out-first 布局,与 nn.Linear 权重转置);且 expert_out/shared_out 存扁平 [T,H] 而 hidden/out 存 [1,T,H]。消费者要靠读 transformers 源码才能确定切分方向,切错方向 replay 必红但定位成本高;#1150 的 docstring 有完整 fixture 契约段(NaN padding、-1 padding 都写明)。 tools/accuracy/dump_qwen38fn_moe_golden.py:18-21,110-111
L4 L 家族级:provenance 只记 transformers 版本,不记 torch 版本。权重字节依赖 torch 全局 RNG 流(模块构造先消耗 RNG、再 reseed),跨 torch 版本再生成会静默变字节而 ref_transformers 不变。#1150 同缺口,属 #1150-#1155 系列的统一决定,本 PR 不单独改,记录为已知边界。 tools/accuracy/dump_qwen38fn_moe_golden.py:108

最小加固:修 M1(追加 S=1 decode case 与对齐倾斜 case——实测 4.0 * gate.weight[0] 作输入可造出 4 专家命中 / 12 空专家、renorm 权重 [0.984, 0.006, 0.006, 0.005],并对倾斜 case 断言条件真实成立,防止后续 seed/scale 调整后静默退化为均匀)、修 M2 / L1(收紧为 torch.equal,实测成立)、修 L2(canon 补 act,需重生成 fixture,本机已验证逐字节可复现)、修 L3(docstring 契约段补切分语义与扁平形状约定);L4 记录为家族级已知边界,不改。

@sysuTobo

Copy link
Copy Markdown
Contributor Author

修改措施

Pushed c8a91fc5(追加 commit,未 force push,保持 draft)。

  • M1 已修:CASES 改为 (seq_len, mode) 表并追加两个 case——case2 S=1(decode 行数形态)、case3 S=16 倾斜(全部 token 与 expert 0 的 router 行对齐,SKEW_SCALE=4.0):实测 4 专家命中 / 12 空专家、renorm 权重 [0.9844, 0.0059, 0.0056, 0.0053]、selected [0,6,8,5],覆盖 [Model] Qwen3.8-Flash-Next roadmap #1103 MoE spike 要求的 skewed + empty experts。倾斜条件带守卫断言(n_hit == TOP_K、max_w > 0.5),防止后续 seed/scale 变动静默退化为均匀分布。新 case 追加在尾部,case0/case1 种子不变、张量字节稳定。tools/accuracy/dump_qwen38fn_moe_golden.py:74-80,150-152,178-184
  • M2 已修:block=experts+shared 从 allclose(atol=1e-2) 收紧为 torch.equal(bf16 逐元素加,两侧确定性同式,实测 4 case 全部逐位相等)。tools/accuracy/dump_qwen38fn_moe_golden.py:170-172
  • L1 已修:router logits 从 allclose 收紧为 torch.equal,与 body 的 bit-for-bit 声明一致(实测逐位相等)。tools/accuracy/dump_qwen38fn_moe_golden.py:169
  • L2 已修:hidden_act 提为常量 HIDDEN_ACT 并入 digest canon(...;norm_topk=True;act=silu),改激活函数不再静默沿用旧 provenance。seed / case 表维持与 feat(qwen4): add Qwen3.8-Flash-Next QSA indexer golden #1150 相同口径不入 canon。tools/accuracy/dump_qwen38fn_moe_golden.py:61,92,99-102
  • L3 已修:docstring 补 fixture 契约段——pre-softmax logits、renorm 权重 sum-to-1、experts_gate_up_proj [E,2I,H] gate=前 I 行 / up=后 I 行、experts_down_proj [E,H,I] F.linear 布局、shared_out 已含逐 token 标量 sigmoid 门、expert_out/shared_out 扁平 [T,H] 而 hidden/out 带 batch 维。同时对齐 feat(qwen4): add Qwen3.8-Flash-Next QSA indexer golden #1150 的 FORMAT_VERSION 注释与 config_digest docstring。tools/accuracy/dump_qwen38fn_moe_golden.py:18-36,64,98
  • L4 不修(家族级已知边界):provenance 不含 torch 版本,feat(qwen4): add Qwen3.8-Flash-Next QSA indexer golden #1150 同缺口;是否统一加 ref_torch 张量留给 feat(qwen4): add Qwen3.8-Flash-Next QSA indexer golden #1150-feat(qwen4): add Qwen3.8-Flash-Next PLE n-gram golden #1155 系列一并决定,本 PR 不单独偏离家族布局。

复验:fixture 重生成 42 张量(10 共享 + 4 case × 8),provenance = transformers 5.16.1 / config_sha256=41374b44…(canon 重算一致)/ format_version=1。4 case 重推导全过:rederivation_ok=True,experts_hit = 15 / 16 / 4 / 4(case3 的 4 = 12 空专家),max_w = 0.7695 / 0.8438 / 0.4414 / 0.9844;sum-to-1、isfinite、倾斜守卫断言全绿。case0/case1 与全部权重张量对旧 fixture 逐字节相同(torch.equal 全 True)。本机重跑与 committed fixture 逐字节一致(环境:torch 2.14.1+cpu / transformers 5.16.1 / CPU)。

@sysuTobo
sysuTobo marked this pull request as ready for review October 14, 2026 14:40
@sysuTobo
sysuTobo marked this pull request as draft October 16, 2026 13:45
@github-actions github-actions Bot added qwen4 Qwen4-architecture model crate (pegainfer-qwen4) needs-rebase Conflicts with main, needs a rebase labels Oct 16, 2026
@github-actions

Copy link
Copy Markdown

This pull request has merge conflicts with main. Please rebase it so it can be reviewed and merged.

@sysuTobo

Copy link
Copy Markdown
Contributor Author

Code review

自查对象:本次推送的 070d9404(replay gate + dev-deps)与 57d33e8b(doc 指针),基线 4e30b8a9。
覆盖清单(全部认领):pegainfer-qwen4/tests/moe_block_replay.rs(新增 1394 行)、pegainfer-qwen4/Cargo.toml(+2 dev-deps)、Cargo.lock(+2 行,仅本 crate 依赖表)、docs/models/qwen4/bring-up.md(+3/-1)。
证据文件缩写:replay = pegainfer-qwen4/tests/moe_block_replay.rs,host_ops = pegainfer-deepseek-v2-lite/src/host_ops.rs,gemma4 moe = pegainfer-gemma4/src/moe/tests.rs。

ID 级别 一句话 证据(文件:行)
R1 H 链路缺口:把 router 留在 f32(不落 bf16)的实现能过原先全部 8 个测试——step1 的 1-ULP 上界看不见它(未取整值落在半个 ULP 内),step2 又按设计读 golden 存的 logits 变异实测 171/210 权重偏移、最多 5 ULP;补 replay:1088(chain 正向)+ replay:1358(控制 e)
R2 M 度量盲区:bf16_ulp 在 bf16 最大有限值处 bits+1 变 inf,该元素 limit 饱和成 inf,检查静默失效 replay:206-213(已加 is_infinite 回退到相对间距)
R3 M 校验空转 + 越界:加载期只断言 top_k <= experts;相等时「无未选专家」使 min_selected >= max_unselected 恒真,且 sorted[top_k] 越界 replay:576-586(已改为严格 <,并写明两条后果)
R4 M D1 重复:host 侧 softmax / router GEMM / canonical tie rule 与既有生产实现是同一逻辑的第二份,且原先没有命名这层关系 replay:666-675 ↔ host_ops:388(topk_softmax_routes)、host_ops:414(gate_logits_host)、host_ops:434(softmax)
R5 M 与 gemma4 同类检查的两处取舍没引用先例:为什么 step2 读 golden 存的 logits、为什么不能像 gemma4 那样位置相等断言 picks replay:716-723 ↔ gemma4 moe:286-290;replay:971-977 ↔ gemma4 moe:359
R6 M 选择契约被两个测试需要(存量 logits 与链路),若各写一份就是同一判据两份实现 replay:891(抽成单一 owner assert_selection,两个测试共用)
R7 L 文档数字用理论值或旧值:头部写「f32 累加 ~1e-5」而非实测、FLOOR_SAFETY 注释写 0.20(实测 0.41)、权重和阈值 0.01 无推导、「逐边界取整更紧」无实测 replay:43-56、replay:97、replay:733-738、replay:1010-1017
R8 L D3 冗余执行:控制 (d) 为取「正向路径的 floor」把 silu 分支又算了一遍 replay:1329(保留,理由见下)
R9 L D4 同一事实两个读者:几何由测试用 serde_json 直读 frozen_config.json,而 src/config.rs 已有类型化解析并单测钉住 norm_topk_prob 默认值 replay:468 ↔ pegainfer-qwen4/src/config.rs:617
R10 L 本 gate 不执行任何生产代码;expert_out 是录制激活,routed GEMM 与 GPU router kernel 未被门禁 replay:58-64 + 每个 case 的测试输出 + PR body「What remains ungated」
R11 L 未被 CI 覆盖:CPU job 只跑 4 个 crate 的 --lib 加一个 --test frontend_e2e,本 gate 无设备依赖却不在其中 .github/workflows/ci.yml:142-156、:223

正向数据流逐条模拟(代入 frozen_config.json 真实几何:hidden 2560 / 512 experts / top-10 / shared intermediate 640;cases T=8,12,1):

# 数据流 变换 判定
DF1 caseN_hidden [T,2560] × gate_weight [512,2560] f32 累加 → 单次 bf16 取整(F.linear 的 bf16 语义,参考实现 modeling_qwen4_exp.py:909) ✓ 8/10752 元素差 1 ULP,其余 bitwise;worst_slack 恰为 1.0000,即上界是紧的不是松的
DF2 存量 router_logits → softmax(fp32, 512) → top-10 → /= sum → bf16 对应 modeling_qwen4_exp.py:910-914 ✓ 210/210 bitwise;每 token 权重和 ∈ [0.998291, 1.001709](推导上界 2^-9≈0.002)
DF3 DF1 的输出接 DF2(链路) 同上,输入换成重放出的 bf16 logits ✓ 210/210 bitwise;case1 有 2/12 行专家集合漂移,全部落在 3 个边界 tie 行内
DF4 hidden × shared_{gate,up}_proj [640,2560] → silu(a)·b → × shared_down_proj [2560,640] 每个算子后按参考实现取整(modeling_qwen4_exp.py:854) ✓ worst_abs 0.0625(case1,worst 89 ULP 但都是近零元素),最差元素为推导 floor 的 0.41
DF5 sigmoid(hidden · shared_expert_gateᵀ [1,2560]) → × shared_out → + expert_out 逐 token 标量门 + bf16 乘 + bf16 加(modeling_qwen4_exp.py:934-936) ✓ 53760/53760 bitwise(含门本身,说明 sigmoid 的 bf16 语义复刻到位)
DF6 元数据 qwen4_moe_golden → format_version / config / config_sha256 与 sha256(frozen_config.json) 比对 ✓ 全部 panic 而非 skip;实测 sha256 889658f2…e74b 一致

Top-k tie 规则(实测,非假设):torch.topk 对完全相等的值不定义顺序,本 fixture 确实有 tie——case0 8 行中 5 行在 top-10 内部有重复概率,case1 12 行中 3 行在第 k 位边界打平(row0 取 393,按 index 序会取 151,两者 p = 0.027970612)。因此按权重匹配的集合比,而不是按位置比:① 存量行是合法 top-k 集合(没有未选专家胜过任一已选专家);② 其概率序列逐 slot bitwise 等于 top-k 的降序序列;③ 每个存量权重 bitwise 等于它旁边那个专家的归一化概率(把 slot 与 expert 绑死,仅保留权重多重集的置换会失败);④ 边界未打平时专家集合必须等于 canonical 集合(降序概率、tie 内升序 index)。

修改措施

  • R1 已修:新增 chained_router_to_routing_replay(replay:1088,链路权重 210/210 bitwise)与第 5 个控制 control_unrounded_router_logits_fails(replay:1358,实测 171/210 偏移、最多 5 ULP,断言 violations*2 >= total);头部表格加「router→routing chain」一行,trap 列表由两条改为三条(replay:20、replay:37-41)。
  • R2 已修:bf16_ulp 在下一个 bit pattern 为 inf 时回退到 mag * BF16_RELATIVE_ULP,并写明不回退就会让该元素的检查变真空(replay:206-213)。
  • R3 已修:加载期断言改为 top_k < experts,注释点明两条后果(合法性判据恒真、边界 tie 探测越界)(replay:576-586)。
  • R4 就地更正(命名而非合并):不能复用——pegainfer-deepseek-v2-lite 是 CUDA-linked + feature-gated,helper 是 crate-private,契约也不同(它把 picks 按 index 重排以对齐 HF 的 unsorted gate 路径,renorm 语义亦不同),而本 crate 在 src/lib.rs:5-10 明确声明零 CUDA 依赖。已在 Routing 的文档注释里点名先例、写清不复用的理由,并写明 [Model] Qwen3.8-Flash-Next: 512-expert MoE #1107 是收敛点(应由共享代码重放本 fixture,而不是再抄第三份)(replay:666-675)。sha256_hex 的第三份拷贝同理保留(pegainfer-gemma4/src/testkit.rs:69 与 pegainfer-sim/src/profile.rs:599 均为 crate-private,跨 crate 不可达)。
  • R5 已修:两处取舍都补上先例与理由——step2 读存量 logits 引 gemma4 moe:286-290(同一做法同一理由:契约应在 kernel 真正消费的那些 bit 上判定);不做位置相等引 gemma4 moe:359(那边对端是自家 kernel,实现同一条 canonical 规则,这边对端是 torch)(replay:716-723、replay:971-977)。
  • R6 已修:把每 token 的选择契约抽成 assert_selection + print_selection(replay:891、replay:1030),两个测试共用一份;census 增加 set_drifts 计数,使「集合漂移只发生在 tie 行」成为读数而非口头结论。
  • R7 就地更正:头部「f32 累加 ~1e-5」换成实测「8/10752 落 1 ULP,其余 bitwise」(replay:48);FLOOR_SAFETY 注释 0.20 → 实测 0.41(replay:97);权重和阈值补完整推导(2^-9≈0.002 推导上界、实测 0.0017、gate 取 0.01 留 5x、无 renorm 偏 0.34-0.47 即 34x 之外)(replay:1010-1017);「逐边界取整更紧」补实测对照(不取整的 f32 单遍 0.125 abs、仅 ~40% bitwise)(replay:733-738)。
  • R8 明确保留:控制臂必须与正向臂用同一条线,所以 floor 只能取自 silu 路径;代价是每个 case 多算一遍 shared 分支(全套 10 个测试 0.5s),已在代码里写明理由(replay:1327-1329)。
  • R9 明确留给 [Model] Qwen3.8-Flash-Next: 512-expert MoE #1107:集成测试拿不到 pub(crate) 的 config::Config,本次不改生产 API 可见性;已在模块文档写明分工(config 的单测钉默认值,本 gate 钉该默认值蕴含的数值),并把「[Model] Qwen3.8-Flash-Next: 512-expert MoE #1107 落地后改为经由 crate 的 Config 取几何,使 gate 同时覆盖生产解析」写进 PR body 的 follow-up。
  • R10 已在三处写明:模块文档 replay:58-64、每个 case 的测试输出(expert_out itself is the fixture's recorded activation, not a replay)、PR body 的「What remains ungated」。
  • R11 明确留给维护者:改 CI 不在本次范围(会动到所有 PR 共用的 workflow)。已核实本 gate 具备进 CI 的全部条件(无设备依赖、0.5s、fixture 是仓库内真实 blob 而非 LFS、cargo metadata --locked 干净),并在 PR body 的 follow-up 里点名 pegainfer-sim --test frontend_e2e 作为「在 CPU job 里跑单个集成 target」的先例。

验证回执(全部实跑):cargo fmt --all -- --check 干净;cargo clippy --release --locked -p pegainfer-qwen4 --all-targets -- -D warnings 0 error;cargo test --release --locked -p pegainfer-qwen4 = 40 lib + 10 replay 全绿;cargo metadata --locked --no-deps 通过。红线扫描(4 个改动文件)零命中。dump 脚本与 fixture 未改动、未重新生成。

@sysuTobo
sysuTobo force-pushed the feat/qwen38fn-moe-golden branch from 57d33e8 to 0a991a4 Compare October 16, 2026 16:24
@github-actions github-actions Bot removed the needs-rebase Conflicts with main, needs a rebase label Oct 16, 2026
@sysuTobo

Copy link
Copy Markdown
Contributor Author

Code review

对这次设备侧落地(pegainfer-kernels 的新 qwen4 feature:路由核 + 共享专家合成核 + FFI/ops 封装 + parity gate)做了一轮审计。正向路径结论是可信的:gate 11/11 通过(RTX 5060 sm_120,CUDA 12.9.1,0.45 s)——路由对存量 logits 逐位 210/210,合成对存量操作数逐位 53760/53760,链式相对误差 0.0397 对推导上界 0.1332(占比 0.298),三个负控制均非空阈值(去重整化 210/210 且 146× 于门槛、bf16 softmax 56/210、去共享门 76% 且 125×)。共 11 项发现,其中 1 项阻塞(会让 clippy -D warnings 变红):

ID 级别 一句话 证据(文件:行)
C1 H(阻塞) 失败关闭测试用 f32 精确 assert_eq!,却没带本文件自己的 float_cmp allow 惯用法,clippy -D warnings 会红 pegainfer-kernels/tests/qwen38fn_moe_block_parity.rs:1546(惯用法见 :366-369、:865-868、:1101-1104)
C2 M 三个 ops 封装只校验 HiddenStates 的标签,没有 crate 标准的 checked_extent,逻辑尺寸超出后备分配时会带着越界 stride 进 kernel pegainfer-kernels/src/ops/qwen38fn_moe.rs:126-131,288-302(标准见 src/tensor.rs:580,姊妹惯用法 src/ops/gemma4.rs:41-42,53)
C3 M 文档断言"没有重复实现 csrc/shared/elementwise.cu 已有的东西",与新合成核矛盾——它就是 add_scaled_bf16_kernel 的算术 docs/models/qwen4/bring-up.md:105 ↔ csrc/shared/elementwise.cu:26-42
C4 M WEIGHT_SUM_TOL 注释两处数字错:全 21 token 的最紧余量是 14×(0.1425 对门槛),不是 ">30x"(那只来自 token 0);"10·2^-9 bound sits at 0.002" 也不成立(10·2^-9 = 0.0195) pegainfer-kernels/tests/qwen38fn_moe_block_parity.rs:133-136
C5 M 替代匹配的 `.find( (d,_)
C6 L 头注释称三个控制"镜像 device-free gate 的陷阱",但那个 gate 有 5 个(gelu-vs-silu 与 unrounded-router-logit 未移植) pegainfer-kernels/tests/qwen38fn_moe_block_parity.rs:31-36 ↔ pegainfer-qwen4/tests/moe_block_replay.rs:1183,1228,1276,1317,1358
C7 L ORDER_EPS = 1e-5 的理由写成"~1e-6 softmax 噪声",真正的决定量是 fixture 的最小非零边界间隔 4.157e-4(→ 41× 余量) pegainfer-kernels/tests/qwen38fn_moe_block_parity.rs:112-117
C8 L ops 模块的 2 个单元测试没有出现在任何文档命令里 pegainfer-kernels/src/ops/qwen38fn_moe.rs:333-392
C9 L saturating_mul 会把溢出的行数饱和掉、让后面的长度校验静默通过,应按姊妹写法用 checked_mul + 显式溢出错误 pegainfer-kernels/src/ops/qwen38fn_moe.rs:132 ↔ src/ops/gemma4.rs:49-52
C10 L win_count < kPerThread 是可证明不可达的静默降级保护,但没写明其不可达性 pegainfer-kernels/csrc/qwen38fn/qwen38fn_moe_route.cu:268
C11 L crate 内现在有两个"毒行"机制相反的路由器(gemma4 靠 sentinel 传播 NaN,qwen38fn 提前返回),不对称没有被点名 pegainfer-kernels/csrc/qwen38fn/qwen38fn_moe_route.cu:188-198 ↔ csrc/gemma4/moe.cu:34-38,90-97

修改措施

全部 11 项已修,随 d058928c(feat(qwen4): route and combine the MoE golden's block on device)一并落地;未 amend、未 force-push。C3/C4/C6/C7/C8 属文档与注释中的数字/断言错误,均为就地更正(改原句,不追加 changelog)。

  • C1 已修:给 a_non_finite_logit_fails_its_row_closed 补上本文件的多行惯用法 #[allow(clippy::float_cmp, reason = "both sides are bf16-representable, so equality is bitwise equality")](修后 tests/qwen38fn_moe_block_parity.rs:1533-1537)。cargo clippy --release --locked -p pegainfer-kernels --features qwen4 --all-targets -- -D warnings 现 exit 0。
  • C2 已修:三个封装都补上 checked_extent——route 的 logits(src/ops/qwen38fn_moe.rs:136)、shared branch 的 hidden / shared_out(:277-278)、combine 的 expert_out / shared_out / out(:322-324),并注明"HiddenStates 的维度是标签、kernel 拿它当 stride 读分配"。gemm_into_checked 只断言标签,所以这两处校验不是重复的。
  • C3 已修(就地更正):docs/models/qwen4/bring-up.md:109 的句子改为"没有第二份 GEMM、没有第二份 SwiGLU",并点名唯一的近似重复就是合成核:add_scaled_bf16_kernel(csrc/shared/elementwise.cu:26-42,封装 ops::add_scaled_bf16_into)已有同样的两次取整算术 out = bf16(b + bf16(scale·a));它不能直接复用,因为其 scale 是整块缓冲区共用的一个 host f32,而本参考的共享专家门是逐 token 的设备值 bf16(sigmoid(bf16 gate_logit))——复用要么 D2H 门 logits 并每 token 一次 launch(且 sigmoid 会在 host 侧取整,而非参考取整的位置),要么给共享 kernel 加逐行 scale 指针,即改动 glm52 step body 已在调用的 op。:112 的 combine 条目与 src/ops/qwen38fn_moe.rs 模块头同步更正(该头原先也写着"只拥有无生产对应物的部分")。
  • C4 已修(就地更正):WEIGHT_SUM_TOL 注释(tests/qwen38fn_moe_block_parity.rs:148-158)改回 device-free gate 的真实推导:每个存量权重是 bf16(w_j),一次 bf16 取整最多移动半间距(相对 2^-9),十个 w_j 之和为 1,故十次取整最多移动 2^-9 ≈ 0.00195;实测最差 |Σw − 1| = 0.00171,门槛 0.01 留 ~5× 余量。去重整化的余量改为全 21 token 的最紧值 0.1425(case1 token 8)→ 14×,并写明 0.34–0.47 是三个 case 的 token 0(也就是控制里打印的那一行,是最宽而不是最紧的)。
  • C5 已修:替代者先收集成一个可消费的池(:1072-1078),并断言每行漂移 ≤ 1 个专家(:1079-1085);匹配改为 substitutes.pop()(:1102),池空即 panic。两个 golden 缺失槽再也不能共用同一个替代者,两专家漂移会直接失败而不是被当作两次独立的单专家并列。
  • C6 已修(就地更正):头注释(:31-44)改为"三个控制是 device-free gate 五个陷阱中的三个——即有设备侧开关的那三个",并点名未移植的两个(control_gelu_instead_of_silu_fails、control_unrounded_router_logits_fails)为何无法移植:设备侧 SwiGLU 就是生产 silu_mul_batch_into(SiLU 硬编码,无激活开关),router logits 就是生产 GEMM 的 bf16 存储(没有可误用的 f32-logit 变体),这两个错误在本 crate 的入口不可达,仍由 device-free replay 把关。
  • C7 已修(就地更正):ORDER_EPS 注释(:118-132)改为以实测为依据:决定量不是求和顺序噪声(~2^-20 相对,此量级 ~1e-8),而是 fixture 的最小非零 10th-vs-11th 概率间隔 4.157e-4(case0 token 1:0.0135116 对 0.0130959),故 1e-5 位于其下 41×——足以吸收求和顺序,又绝不足以放过真实的边界跨越。同时写明精确并列不归这个常量管:case1 有 3 行并列,走 grade_golden_token 的权重匹配集合规则;严格并列规则只在精确相等处生效。
  • C8 已修(就地更正):bring-up 文档的测试命令块补上 ./docker/dev.sh run cargo test --release --locked -p pegainfer-kernels --features qwen4 --lib(:104-107),注明它带着本模块的两个 device-free 单元测试(only_the_compiled_geometry_passes、production_semantics_upcast_the_softmax)。已实跑:21 passed / 6 ignored,两个都在内。
  • C9 已修:rows.saturating_mul(...) 换成 checked_mul + 显式溢出错误(:138-143),与 src/ops/gemma4.rs:49-52 一致——饱和值会让下面的 weights.len() >= picks 校验失去意义。
  • C10 已修:在 win_count < kPerThread 上写明不可达性(csrc/qwen38fn/qwen38fn_moe_route.cu:279-286):本线程恰好拥有 kPerThread 个专家,而每一轮都用 -INFINITY 掩掉自己刚拿走的赢家,所以它最多赢 kPerThread 轮,win_count 只在填满最后一格的那次自增时到达数组长度。判断保留而非删除,是为了 kTopK 或专家分配一旦变化时边界不会退化成静默越界写。
  • C11 已修:只在 qwen38fn 侧点名不对称(csrc/qwen38fn/qwen38fn_moe_route.cu:196-205,文件头 :31-33 同步指向):策略与 gemma4 相同、机制相反——csrc/gemma4/moe.cu 用 __any_sync 投票后把 held[] 灌成 NaN,让它穿过 softmax 与 argmax 树,最后在 taken >= experts 的 sentinel 分支里把 pick 改写成槽位号、权重改写成 NaN;这里投票后直接 return,NaN 从不进入算术,后续分支也不必识别它。两者可观测字节相同(NaN 权重 + 槽位号占位 id),但只有"灌 NaN"那条依赖 sentinel 在与 NaN 的比较中存活;新路由器应抄这里的提前返回,而不是抄灌 NaN。gemma4 的文件未改动。

验证(全部实跑)

命令 结果
./docker/dev.sh run cargo test --release --locked -p pegainfer-kernels --features qwen4 --test qwen38fn_moe_block_parity -- --nocapture 11 passed / 0 failed,0.45 s
./docker/dev.sh run cargo test --release --locked -p pegainfer-kernels --features qwen4 --lib 21 passed / 6 ignored(含本模块 2 个单元测试)
./docker/dev.sh run cargo clippy --release --locked -p pegainfer-kernels --features qwen4 --all-targets -- -D warnings exit 0(C1 修后)
./docker/dev.sh run cargo check --locked -p pegainfer-kernels exit 0;构建脚本打印 Qwen4 ... CUDA kernels disabled,即默认 feature 下新 .cu 未参与编译
./docker/dev.sh run cargo fmt --all -- --check exit 0
cargo test --release --locked -p pegainfer-qwen4 --test moe_block_replay 10 passed,device-free 套件未改动、仍绿

修后重跑的正向数字与审计一致:router GEMM 最差 0.03125 绝对(占界 0.333,T=1 解码行逐位);链式最差相对 0.0397(占界 0.298);共享分支占界 0.241 / 0.290 / 0.000;合成对存量操作数逐位 53760/53760,对设备分支最差占界 0.371;三个控制 146.5× / 56-of-210(26.7%,且 0 行重排 pick)/ 125.5×。

@github-actions github-actions Bot added the needs-rebase Conflicts with main, needs a rebase label Oct 20, 2026
@github-actions

Copy link
Copy Markdown

This pull request has merge conflicts with main. Please rebase it so it can be reviewed and merged.

Signed-off-by: bohaotu <bohaotu@pony.ai>
Address review on the MoE block golden:
- add a skewed-routing case (all tokens aligned with expert 0's router
  row: 4 experts hit / 12 empty, dominant renorm weight 0.984) with
  asserts that the skew condition actually holds, and an S=1 decode
  shape case, per the pegainfer-project#1103 MoE spike's skewed/empty-expert requirement
- tighten the router-logits and block-composition checks from allclose
  to torch.equal (both arms are bitwise-identical by construction;
  measured max_abs_diff was exactly 0.0)
- add hidden_act to the config digest canon so an activation change
  cannot leave provenance stale
- document the fixture contract in the dump docstring: fused gate_up
  split semantics, down_proj F.linear layout, flat vs batched shapes,
  pre-softmax logits and gated shared output

Regenerated the fixture (42 tensors); case0/case1 and all weight
tensors are byte-stable (cases are append-only), config_sha256 changed
with the canon.

Signed-off-by: bohaotu <bohaotu@pony.ai>
…etry

The fixture was dumped from a reduced config (hidden 512, 16 experts top-4,
intermediate 64). Every dimension now comes from the frozen checkpoint config:
hidden 2560, 512 experts top-10, moe_intermediate 640, shared 640, plus the
upstream `norm_topk_prob=True` default that config.json does not carry (the
dump asserts the field is still absent and that the reference renormalizes).

At real geometry the fused expert stack is 2.52e9 parameters (~5 GB bf16), so it
is not stored: the reference's `experts` return is recorded per case as the
activation entering the block sum, while the router `gate_weight` (2.6 MB) and
the shared expert's projections (~10 MB) stay stored weights. Provenance moves
into one metadata JSON key (format_version 2); the loose provenance tensors are
gone, and the unstored stack stays reproducible from the recorded seed.

Cases are token batches of 8, 12 and 1 (the block is row-wise, so a 12-row case
is what exercises the 512-way renorm and the shared gate). The dump keeps
re-deriving the router (fp32 softmax -> top-k -> renorm) and asserts it matches
the module bit-for-bit, asserts the block equals the routed output plus the
shared output under its scalar sigmoid gate, and asserts the top-k weights sum
to 1. After writing, the fixture is reloaded and every case replayed bit-for-bit
- from the stored tensors alone, and against a block rebuilt from the recorded
seed for the recorded `expert_out`.

Fixture: 12.31 MiB, 26 tensors.
Signed-off-by: bohaotu <bohaotu@pony.ai>
`reproduce_from_fixture` re-implemented the router forward (fp32 softmax ->
top-k -> `norm_topk_prob` renorm) and the block sum in torch purely to prove
the fixture reproduces itself from the file. Every value it recomputes is
captured from the reference module through forward hooks in `case_tensors`,
and the same semantics are already asserted in-run against the live module, so
the oracle does not depend on it. The reload pass that rebuilds the reference
from the recorded seed and pins the stored weight set stays.

Fixture byte-identical (26 tensors, same metadata): the pass only read the
written file.

Signed-off-by: bohaotu <bohaotu@pony.ai>
Drop every self-check: the router logits / routing weights / top-k
re-derivations, the block-composition and renorm asserts, the finiteness
check, the norm_topk_prob config pins and the stored-gate_weight shape pin,
plus the reload-and-replay pass (`reproduce_with_reference`) and the helpers
and dead imports that only served them. The fixture stores the reference's
own output, so there is nothing to re-check. Compress the module docstring to
match.

The stored tensor set is unchanged: regenerating the fixture is byte-identical.

Signed-off-by: bohaotu <bohaotu@pony.ai>
The MoE block golden had no reader: the fixture and its dumper landed, but
nothing consumed them, which is exactly what pegainfer-project#1103 forbids. Add the consumer —
a device-free replay of `Qwen4ExpTextSparseMoeBlock` over the committed
fixture, in f32 on the host, with no CUDA, no Triton and no checkpoint:

    cargo test --release --locked -p pegainfer-qwen4 --test moe_block_replay

Ten tests, five of them negative controls that assert the comparison FAILS.
Measured against the golden (transformers 5.16.1, torch 2.14.1+cpu):

  router   bf16(hidden · gateᵀ)                1 ULP   8 of 10752 logits off by it
  routing  fp32 softmax → top-10 → renorm      bitwise 210/210 weights
  chain    the same on the replayed logits     bitwise 210/210 weights
  shared   bf16-rounded SwiGLU + down proj     0.0625 abs, 0.41 of the floor
  combine  expert_out + sigmoid(gate)·shared   bitwise 53760/53760

Tolerances are the reference's own rounding, not slack. Every dtype boundary
rounds where the reference rounds, so the only residual is accumulation order;
the shared branch adds one derived term — a 1-ULP difference in its
`silu(gate)·up` intermediate, propagated through the 640-term down projection
in quadrature from the replay's own intermediate — instead of a magic constant.
Emulating the roundings is the tighter choice and the measured one: one
unrounded f32 pass lands 0.125 abs away with ~40% of elements bitwise.

Three silent traps get a control each, plus two more for precision:
skipping the top-k renorm (`norm_topk_prob` is absent from config.json and
defaults to true upstream) moves 210/210 weights by up to 0.24; a bf16 softmax
moves 56/210; dropping the shared expert's learned scalar sigmoid gate moves
the block output by up to 22.25; gelu in place of silu moves the shared output
by up to 1.91; and keeping the router's logits in f32 instead of storing bf16
moves 171/210 weights by up to 5 ULP — which neither the 1-ULP logit bound nor
the stored-logit routing replay can see, so the chain row above exists for it.

Top-k ties: `torch.topk` leaves the order of an exact tie undefined, and this
fixture has them (5 of case0's 8 rows repeat a probability inside the top-10;
3 of case1's 12 rows tie at the k-th boundary, where row 0 takes expert 393
over 151 at p = 0.027970612). So the selection is compared as a weight-matched
set — valid top-k set, descending probability sequence bitwise, each stored
weight the renormalized probability of the expert stored beside it bitwise, and
set equality wherever the boundary is untied — not positionally. gemma4 can
assert its picks positionally because its other side is its own kernel; here the
other side is torch.

Not gated, deliberately: `expert_out` is a recorded activation, since the fused
512-expert stack is 2.52e9 parameters (~4.7 GiB bf16) and is not stored. The
routed-expert GEMM and the GPU router kernel stay ungated until pegainfer-project#1107 replays
this fixture on device.

Provenance is asserted on load, never skipped: the metadata key must parse,
`format_version` must be 2, and `config_sha256` must equal
sha256(pegainfer-qwen4/tests/frozen_config.json), whose geometry also supplies
the shapes the tensors are checked against.

The host-side softmax / router GEMM / canonical tie rule already exist as
production code for another line (pegainfer-deepseek-v2-lite/src/host_ops.rs),
but that crate is CUDA-linked and feature-gated, its helpers are crate-private
and its contract differs, while this crate carries no CUDA dependency at all.
The comment on `Routing` names the copy and points pegainfer-project#1107 at converging them.

Dev-deps are additive: `half` and `safetensors` were already in Cargo.lock.

Signed-off-by: bohaotu <bohaotu@pony.ai>
The bring-up doc already records that `norm_topk_prob` is absent from
config.json and defaults to true upstream, so a port reading only the
checkpoint silently drops the router renormalization. It stops at the
mechanism; now that the MoE block golden has a host replay, name where the
cost is measured and what it is: all 210 stored routing weights move, by up to
0.24 absolute. One sentence, pointing at the device-free test that carries the
negative control.

Signed-off-by: bohaotu <bohaotu@pony.ai>
The Qwen3.8-Flash-Next MoE golden had a device-free consumer only, so it
pinned a contract without executing any production code. Add its device
half inside pegainfer-kernels, behind the line's additive `qwen4` feature, appending
to the module pegainfer-project#1160 opened, so the model crate stays CUDA-free and the
feature keeps the new translation unit out of default builds:

  csrc/qwen4/qwen4_moe_route.cu         routing + shared-expert combine
  src/ffi/qwen4.rs                      the two C entry points, appended
  src/ops/qwen4.rs                      geometry/semantics guards, wrappers
  tests/qwen4_moe_block_parity.rs       the gate, through the ops layer

Routing is fp32 softmax over 512 -> top-10 in the canonical order
(descending probability, ascending index inside a tie) -> optional
`norm_topk_prob` renorm summing in that same order -> bf16 weights and
i32 ids, one block per token over fixed shuffle trees so the picks are
deterministic run to run; a non-finite logit fails its row closed by a
block-uniform vote before the softmax, so no NaN reaches the arithmetic.
The combine folds the learned per-token sigmoid gate into the block sum,
rounding where the reference rounds. Everything else is composition: the
router GEMM, the shared expert's three projections and its SwiGLU go
through `gemm_into_checked` / `silu_mul_batch_into`, so the gate grades
the production path.

Nothing in-tree carries this routing contract, which is why the kernel is
new: gemma4's router is 128 experts with `__expf`, an unconditional
renorm and f32 weights; glm52 (256/top-8) and kimi (384/top-8) score with
sigmoid + bias; FlashInfer's FilteredTopK does no softmax at all. The
combine is the one near-duplicate -- `add_scaled_bf16_kernel`'s
arithmetic, which takes its scale as one host f32 and so cannot carry a
per-token device gate -- and the doc names it as such.

Measured on an RTX 5060 (sm_120), CUDA 12.9.1, every bound derived before
it was checked: routing bitwise 210/210 on the stored logits, the chain
0.0397 relative against the derived 0.1332, the combine bitwise
53760/53760 on stored operands and <=0.371 of its bound on the device
branch, and three controls that are not vacuous thresholds (no-renorm
210/210 at 146x, bf16 softmax 56/210 with no re-ordered picks, no shared
gate 76% of outputs at 125x). The device-free replay is untouched and
still green.

Signed-off-by: bohaotu <bohaotu@pony.ai>
@sysuTobo
sysuTobo force-pushed the feat/qwen38fn-moe-golden branch from d058928 to 47d66d7 Compare October 21, 2026 03:10
@github-actions github-actions Bot removed the needs-rebase Conflicts with main, needs a rebase label Oct 21, 2026
@sysuTobo

Copy link
Copy Markdown
Contributor Author

Rebase

已 rebase 到 main @ 1f17ddb3(含 #1160 QSA indexer、#1176 FlashInfer v0.7.2),并统一到 #1160 合入后的 kernels 布局:

  • 核文件 csrc/qwen38fn/qwen38fn_moe_route.cu → csrc/qwen4/qwen4_moe_route.cu,两个 C 入口 → qwen4_moe_route_cuda / qwen4_moe_shared_combine_cuda
  • FFI/ops 追加进 main 已有的 src/ffi/qwen4.rs / src/ops/qwen4.rs;is_qwen38fn_source 与 feature 行删除,复用 main 的 is_qwen4_source
  • gate 改名 tests/qwen4_moe_block_parity.rs;fixture 与设备无关的 moe_block_replay.rs 未动
  • 公开类型名随目录收敛:Qwen38fn* → Qwen4*

bring-up.md 的冲突按 main 重排后的文本保留,只把本 PR 新增的 MoE 设备段接在其后。重跑:gate 11 passed,lib 21 passed,cargo clippy --features qwen4 --all-targets -- -D warnings 绿,默认 feature 的 cargo check --all-targets 绿,cargo fmt --check 干净。

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request qwen4 Qwen4-architecture model crate (pegainfer-qwen4)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant