Repository navigation
Conversation
Code review范围:2 文件(dump 脚本 162 行 + fixture 26 张量),无运行时数据流。逐文件认领:脚本对照 transformers 5.16.1
最小加固:修 M1(追加 S=1 decode case 与对齐倾斜 case——实测 |
修改措施Pushed
复验:fixture 重生成 42 张量(10 共享 + 4 case × 8),provenance = transformers 5.16.1 / |
|
This pull request has merge conflicts with main. Please rebase it so it can be reviewed and merged. |
Code review自查对象:本次推送的
正向数据流逐条模拟(代入
Top-k tie 规则(实测,非假设): 修改措施
验证回执(全部实跑): |
57d33e8 to
0a991a4
Compare
Code review对这次设备侧落地(
修改措施全部 11 项已修,随
验证(全部实跑)
修后重跑的正向数字与审计一致: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×。 |
|
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>
d058928 to
47d66d7
Compare
Rebase已 rebase 到
|
Summary
Reference golden for the MoE block of Qwen3.8-Flash-Next (#1103).
tools/accuracy/dump_qwen38fn_moe_golden.py— dumpsQwen4ExpTextSparseMoeBlock(transformers 5.16.1qwen4_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_version2).moe_intermediate_size640, shared 640, and the upstreamnorm_topk_prob=Truedefault thatconfig.jsondoes 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 (routergate_weight2.6 MB, shared projections ~10 MB, scalar shared gate) are stored.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.rsreads the fixture: a device-free replay ofQwen4ExpTextSparseMoeBlockon 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).halfandsafetensorswere already inCargo.lock; the dev-dep addition is additive (cargo metadata --lockedclean).Provenance is asserted on load and never skipped: the metadata key must parse,
format_versionmust be2, andconfig_sha256must equalsha256(pegainfer-qwen4/tests/frozen_config.json)— whosetext_configalso supplies the geometry every tensor shape is checked against (includinghidden_act == "silu"). A missing fixture or a moved config panics.What it proves
bf16(hidden · gateᵀ)down(silu(gate)·up)+ the scalar sigmoid gateexpert_out + sigmoid(gate_logit) · shared_outTolerance justification — these are the reference's own rounding, not slack for a port:
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.2 · 2^-7 · √Σ(h·w)², one bf16 ULP ofhpropagated throughdown_projin quadrature, doubled for sign correlation — not a magic constant. Measured worst element sits at 0.41 of it.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_skipping_topk_renorm_failsnorm_topk_probis absent fromconfig.jsonand defaults toTrueupstream, so a port reading only the checkpoint config skips the top-k renormalizationcontrol_bf16_softmax_failsdtype=torch.floatfrom the reference's softmax (a bf16 probability buffer)control_dropping_shared_sigmoid_gate_failscontrol_gelu_instead_of_silu_failsgelu_pytorch_tanh(the vision config's activation) in the text MLPcontrol_unrounded_router_logits_failsThe 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.topkleaves 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 atp = 0.027970612). Soselected_expertsis 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:359can assert its picks positionally because the other side there is its own kernel; here the other side is torch.What remains ungated
expert_outis 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 3Dgate_up_proj/down_projlayout 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.config'snorm_topk_probdefault (already unit-tested atpegainfer-qwen4/src/config.rs:617) to the numbers that default implies.-p pegainfer-qwen4 --test moe_block_replayinto theCPU unit testsjob (it needs no device and takes 0.5s;pegainfer-sim --test frontend_e2eis the precedent for running one integration target there).Also in this push:
docs/models/qwen4/bring-up.mdnow points thenorm_topk_probparagraph at the gate that prices the trap.Device path
The device half of this consumer now lives in
pegainfer-kernels, behind a new additiveqwen4feature 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-kernelsstays 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 insrc/ffi/qwen4.rs, wrappers insrc/ops/qwen4.rs, gate intests/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) → optionalnorm_topk_probrenorm, 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_orbefore 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:
csrc/gemma4/moe.cu(behindgemma4)__expfsoftmax, renorm unconditional, × a per-expert scale, f32 weights outcsrc/glm52/glm52_router.cusigmoid + bias, renorm byroute_scale / Σcsrc/kimi_k2/kimi_router.cue_score_correction_bias, f32 weightsFilteredTopK(viaops/glm52/topk.rs)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 byops::add_scaled_bf16_into), which takes its scale as a single hostf32uniform 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.mdrecords both conclusions.Commands (RTX 5060 sm_120, CUDA 12.9.1,
pegainfer-dev:cu132):Measured, every bound derived first and then checked against the measurement (the gate prints per-case worst-abs and bound fraction):
gemm_into_checked)e^{2a} − 1 = 0.1332relative,a= the router floorThe 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_expertsis graded as a weight-matched set rather than positionally; and the worst renormalized |Σw − 1| is 0.00171, inside the derived2^-9 ≈ 0.00195.Still ungated:
expert_outremains a recorded activation, so the routed-expert GEMM (the fused 3-Dgate_up_proj/down_projlayout and its scatter/gather) still waits for #1107, which should replay this same fixture; and no CI job enables the kernelsqwen4feature yet, so the device gate is the manual command above — the device-freemoe_block_replaysuite remains the CI cross-check on the same fixture.Validation
python tools/accuracy/dump_qwen38fn_moe_golden.py --out test_data/qwen38fn-moe-golden.safetensorsThe dump asserts nothing — what it stores is the reference's own output at the real geometry, so there is nothing to re-check:
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