Repository navigation
feat(qwen4): component kernels extracted from the brick PRs - #1242
Conversation
…egainfer-project#1153) Signed-off-by: bohaotu <bohaotu@pony.ai> Signed-off-by: Feathbow <feathbow@gmail.com>
…d from pegainfer-project#1154) Signed-off-by: bohaotu <bohaotu@pony.ai> Signed-off-by: Feathbow <feathbow@gmail.com>
…ainfer-project#1155) Signed-off-by: bohaotu <bohaotu@pony.ai> Signed-off-by: Feathbow <feathbow@gmail.com>
…egainfer-project#1151) Signed-off-by: bohaotu <bohaotu@pony.ai> Signed-off-by: Feathbow <feathbow@gmail.com>
…ect#1152) Signed-off-by: bohaotu <bohaotu@pony.ai> Signed-off-by: Feathbow <feathbow@gmail.com>
Signed-off-by: Feathbow <feathbow@gmail.com>
Signed-off-by: Feathbow <feathbow@gmail.com>
github-actions
Bot
added
the
qwen4
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
Commit attribution noticeThis pull request contains commits attributed to an account other than the PR author @FeathBow without a verified signature from that account:
If you are mentioned above and did not take part in this pull request, please say so here. Attributing a commit to someone else normally requires that person to be involved; maintainers should treat unacknowledged attribution as a review blocker. This check never fails the build. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4e72c143ef
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| - **Manifest** (`src/manifest/`, test only until slice B's loader calls it) — the tensor contract, ported from `pegainfer-gemma4/src/manifest/`: `schema.rs` derives it from `Config` and `validate.rs` classifies headers against it. It derives **exactly** the frozen text set, 1658 keys − 333 `model.visual.*` − 31 `mtp.*` = **1294 tensors**, with `lm_head.weight` the only one outside `model.language_model.`. Beyond the untied head, it differs from Gemma 4's in rank-3 shapes (the depthwise conv and two fused expert stacks per layer, which carry no `.weight` suffix) and in a mandatory skip set, because all 333 vision tensors share shard 1 with 16 text tensors. `tests/frozen_tensors.json` (148 KB, written by `tools/accuracy/dump_qwen4_frozen_tensors.py`) holds every shard header of the frozen revision, and `the_contract_is_the_frozen_checkpoint` checks all 1658 tensors against the contract. | ||
| - **Registration** — the server binary hand-mirrors every line's probe identities in a hint table that cannot be derived away. With the feature off, `feature_gate_hint` names `--features qwen4` for both `qwen4_exp` and `qwen4_exp_text`; with it on, the hint stays silent. A test under the feature binds that row to this crate's `MODEL_TYPE` / `TEXT_MODEL_TYPE`, so the mirror cannot drift from the probe. The hint is **identity-only** — it cannot run a compiled-out probe — so it also fires for the quantized siblings and the tied/YaRN variants of `qwen4_exp`, which the config refuses once the feature is on; the hint names the feature to rebuild with, not a promise of service. `pegainfer-server/tests/detect_gate.rs` then drives the shipped binary against a model directory holding this crate's frozen `config.json`, both ways: the stock build must name the feature, the enabled build must detect the line and refuse to serve. That is the only coverage of `main`'s CLI → registry → hint → `launch` wiring, and neither outcome reaches CUDA — `launch` bails before the engine is built. | ||
|
|
||
| - **Component kernels from the five bricks** (#1151 to #1155, committed under their author's name) — `csrc/qwen4/qwen4_hc_block.cu` behind `hc_mix_into` and `hc_scatter_add_into` (the gated residual's mix pass and the scatter back into the 10240-wide stream), `qwen4_moe_route.cu` behind `qwen4_moe_route_into` (fp32 softmax, canonical top-10, `norm_topk_prob`), `qwen4_moe_shared_branch_into` and `qwen4_moe_combine_into`, `qwen4_ple_block.cu` behind `ple_block_into` (the per-branch gate with its double rounding and signed sqrt, then the dilated causal conv) with the host-side hasher in `src/ngram.rs`, and the `qwen4_qsa_norm_rope.cuh` helpers the indexer now shares. Every kernel header spells out the reference's rounding points, which is the part of those PRs worth keeping. Left behind: the goldens (66 MiB, all generated from `torch.manual_seed` and `randn` weights, so they pin the port against itself rather than against the checkpoint), their replay readers, the test-only knobs that had reached production signatures (`combine_scale`, `probs_bf16`, the PLE gate variants) and the prefill-only QSA attention kernel. The five `tools/accuracy/dump_qwen38fn_*_golden.py` scripts stay as the oracle and will read the frozen checkpoint's tensors one layer at a time instead of random init. The GDN brick (#1152) contributed a finding rather than code: Qwen3.5's conv and chunkwise delta-rule kernels match the reference at this line's geometry, so this line calls the qwen35 production entries rather than a split copy. |
There was a problem hiding this comment.
The new HC/MoE/PLE entries are not exercised by any production caller or parity test: a repo-wide search finds these symbols only in their definitions, re-exports, and this documentation, while this added paragraph explicitly leaves the goldens and decoder consumer for later. When slice B begins calling them, wrong rounding, routing, or convolution behavior can therefore ship without detection; land a frozen-checkpoint decoder-layer/device parity gate (or defer the kernels) so this feature has the production E2E/model evidence required for AI-assisted changes.
AGENTS.md reference: AGENTS.md:L132-L132
Useful? React with 👍 / 👎.
Description
Carries the production halves of the five Qwen4 component PRs #1151, #1152, #1153, #1154 and #1155 onto one branch. Each PR is one commit under its author's name with the author's Signed-off-by kept, followed by two commits of ours, one dropping the knobs those kernels had grown for their tests and one for the docs.
csrc/qwen4/qwen4_hc_block.cuwithhc_mix_intoandhc_scatter_add_into, the gated residual's mix pass and the scatter back into the hyper connection stream.csrc/qwen4/qwen4_moe_route.cuwithqwen4_moe_route_into(fp32 softmax, canonical top 10,norm_topk_prob),qwen4_moe_shared_branch_intoandqwen4_moe_combine_into.csrc/qwen4/qwen4_ple_block.cuwithple_block_into(the gate for each branch, the signed sqrt, the dilated causal conv) and the host n-gram hasherpegainfer-qwen4/src/ngram.rs.qwen4_qsa_norm_rope.cuhhelpers and the indexer kernel reading them.pegainfer-qwen35/src/recurrent.rsfor its test is not taken.tools/accuracy/dump_qwen38fn_*_golden.pyscripts come along unchanged as the reference oracle.The goldens stay behind. Every one of the 66 MiB is generated from
torch.manual_seedandrandnweights, so a parity test against them pins the port against itself. Goldens from the frozen checkpoint arrive with the decoder layer that calls these kernels, one layer at a time, so each lands with its consumer. The replay tests and host readers built on those goldens stay behind too, and so do the knobs that served only them and had reached production signatures,combine_scaleon the branch combine (the launcher now uses1 / branches),probs_bf16andQwen4RouteSemanticson the router (qwen4_moe_route_intotakesrenorm: bool), and the PLE gate variants with theirgate_spanandsigned_sqrtarguments. The QSA attention kernel of #1151 is left out as well. It serves prefill only, over contiguous keys with an 8192 token cap, and up to 2051 visible tokens the sparse attention equals the dense attention the FlashInfer hd256 path already serves. The CI row,Cargo.lockand dev dependency changes that served the dropped tests are not carried either.The additions of the four kernel PRs to
pegainfer-kernels/src/ops.rs,src/ffi/qwen4.rsandsrc/ops/qwen4.rsare merged line for line. The only edits inside the authored commits are theuselines two PRs both added and one unit test two PRs had given the same name, nowonly_the_compiled_ple_geometry_passes.Nothing here has a production caller yet and
launchstill refuses the line. The default feature set compiles none of it, so serving on every existing line is unchanged.Test Env
Single GPU (sm_89, x86_64) with CUDA 13.1.
Verification
Same commands on the merge base and on this branch, same machine and target directory.
cargo clippy --release --locked -p pegainfer-kernels --features qwen4 --all-targets -- -D warningscargo clippy --release --locked -p pegainfer-qwen4 --all-targets -- -D warningscargo test --release --locked -p pegainfer-kernels --features qwen4 --lib qwen4cargo test --release --locked -p pegainfer-qwen4 --libcargo build --release --locked --features qwen4cargo fmt --all -- --checkpasses.