Repository navigation
gpu: Qwen3.8's ubatch walk stores a delta layer's q, k and v projection as BF16, and runs β and α as one launch - #1039
Conversation
github-actions
Bot
added
gpu
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Verdict: MUST_FIX: Remove the single-use model-level projection wrapper.
Must fix
1. Inline the BF16 buffer adapter
wide38.rs:1921 has exactly one caller, wide38.rs:2614. It only selects the fallback, cuts a buffer window, and forwards to the existing site owner. The diff supplies no second use. This is an unnecessary abstraction under the fewest-elements rule (kind B).
Cost: inlining removes a function and its forwarding signature without changing launches.
Delete dense_gemm_bf16 and replace its call in gdn with:
let bf16 = if reads_q8_1(c, &gp.qkv)? {
false
} else {
let mut half =
WindowMut::<u16>::of_mut(&mut g.x, 0, m * GDN.channels())?;
site::resident_gemm_bf16(
c.gpu,
&dense_kernels(c),
(w, &gp.qkv),
&wide_in(mixed, &x.act_hid),
dense,
m,
&mut half,
)?
};Non-blocking
⚠ Complete the precision-change evidence: wide38.rs:2614. The timing table names base 904266cfc but no tested head hash; the reproduction block omits the request command. Attach both tested hashes and request commands/output. The BF16 store changes precision: also attach the error model and end-to-end count-pin result required by AGENTS.md's precision-change class. The reported component comparisons and verify-node result do not supply that evidence.
Checked
Duplicates: all three dryer pairs read; they share existing cores with different types or rules. The added model wrapper has one caller.
Complexity: gemm_into moves dispatch and adds output cases; conv_held adds packing/upload/launch/readback setup. Gate growth adds comparisons; launchers add validation/output dispatch; wide38::gdn adds the BF16 choice.
Silent failure: the packing default fills an absent halfword; unsupported BF16 outputs return false for fallback. Conv validates the projection pitch and buffer spans.
Tests: conv comparisons cover output, beta, decay and ring; BF16-store checks compare against round-to-nearest-even bits. Verify-node counts remain asserted. Static inspection only; GPU gates were not rerun.
Gates: verify-node, unsafe and PTX pins have dated reasons; no tolerance moved. The removed gemv-row clause has a dated removal reason.
Escapes: no catalog pattern matched.
Card check: needs-gpu. This pull request changes crates/gpu-gates/src/bin/gate_bf16.rs, crates/gpu-gates/src/bin/gate_gemm.rs, crates/gpu-gates/src/bin/gate_linear.rs … and carries no gate run on a card (CONTRIBUTING's evidence list); a maintainer adds gpu-verified after one.
🔨 Shingler, bloomery's review bot: approves a clean pull request and requests changes on a must-fix (how it works)
… BF16, which the conv widens as it reads The ubatch walk's q·k·v product of a delta layer (an MXFP8 site's W8A8 product, or a BF16 site's product over its held BF16 form) stores each cell as the BF16 bits of its F32 sum, rounded to nearest even (`ProductOut`, `site::resident_gemm_bf16`), into the first halfwords of the same rows. Any other site type keeps the F32 product and the F32 conv. The new conv entry, gdn_conv_prep_bf16, reads those words and widens each value: its outputs are the F32 conv's rule over the widened values, the ring included. Lane l owns the channel pairs 2l, 2l + 1 and 64 + 2l, 65 + 2l of its head, so each word of a row is one load of the lane and a warp's load is one 128-byte line, and it writes each pair's two outputs as one 8-byte store, so a warp's store is one 256-byte line. The conv and the norm are the F32 entries' helpers (`conv_channel`, `head_norm`); the body reads the channels through `ConvIn`, which the F32 entries implement with their own head and ring store. gate_linear's conv clause runs BF16 words at m=512, 2 and 1 over a random ring against the host rule over the widened values on the entry's lanes; gate_gemm's dense_mma and mxfp8_w8a8 cases hold the BF16 store to each cell's F32 bits rounded to nearest even, a NaN as a NaN.
…onv reads each row's β and α at a pitch The joined β·α weight (2·n_v rows, β's first) ran as two products past one row, β's rows and α's rows, so that each landed token-major in its own window. Both GDN conv entries (F32 and BF16 channels) now take `ba_pitch`, the values from one token's β (α) to the next, so they read the joined product's rows in place: each row's β values, then its α's. The ubatch walk's MXFP8 product and the step's and verify's gemv each run once over the whole weight, every output bit for bit the part's it replaces; a delta layer has one launch fewer past one row (GDN_ROWS_LAUNCHES goes). Qwen3.5 and the KDA entries keep their own rows at pitch n_v. gate_linear's conv clauses run the joined rows at pitch 2·n_v, the BF16 channels' among them, against the host rule and refuse a pitch below n_v by name.
0a0180e to
808f410
Compare
There was a problem hiding this comment.
Since the last review:
- Must fix 1: fixed. The single-use adapter is removed and its call is inlined at wide38.rs:2595.
- Evidence question: still open at the projection below.
Verdict: CLEAN: The interdiff resolves the wrapper finding; the precision-evidence request remains non-blocking.
Non-blocking
⚠ Complete the precision-change evidence: wide38.rs:2599. The body omits the tested head hash and benchmark request command/output. Attach those, the error model (σ against exact-forced-32.tsv), and the gate-gpu-e2e count-pin result required by AGENTS.md's precision-change class. The reported verify-node result covers graph structure.
Checked
Duplicates: no new definitions in the interdiff; the single-use adapter is deleted.
Complexity: gdn inlines the existing BF16 choice and reuses the input descriptors for both projections.
Silent failure: unsupported BF16 output returns false before launch; gdn then uses dense_gemm and the conv reads f32.
Tests: no tests changed in the interdiff. GPU gates were not rerun.
Gates: no pin, tolerance or case changes in the interdiff.
Escapes: no catalog pattern matched. No ledger lessons apply.
Card check: needs-gpu. This pull request changes crates/gpu-gates/src/bin/gate_bf16.rs, crates/gpu-gates/src/bin/gate_gemm.rs, crates/gpu-gates/src/bin/gate_linear.rs … and carries no gate run on a card (CONTRIBUTING's evidence list); a maintainer adds gpu-verified after one.
🔨 Shingler, bloomery's review bot: approves a clean pull request and requests changes on a must-fix (how it works)
|
The review is done and the contribution check passed, so I'm merging this. |
Qwen3.8 prefill에서 GDN 층의 q, k, v 투영을 BF16으로 저장하고 β와 α 투영을 한 번에 launch해서 4K prompt의 prefill 시간을 줄입니다.
Part of #565.
What and why
In a Qwen3.8 ubatch walk, each delta (GDN) layer projects q, k and v as f32 and the conv kernel reads them back as f32, and β and α each took their own launch per layer.
This PR changes two steps of that layer:
mxfp8_w8a8_gemm_bf16for the MXFP8 W8A8 site,bf16_mma_cols_gemm_bf16for the BF16 site), and the conv reads those words and widens them as it goes (gdn_conv_prep_bf16). The store rounds to nearest even. A site with no BF16 product (K-quant, for example) keeps the f32 path.Because q, k and v are rounded to BF16, a greedy continuation can differ from main's. This is a precision change of the BF16 store, as with the other BF16 stores in the walk.
Release notes: Changed: Qwen3.8 FP4-Local prompt prefill stores the GDN q, k and v projection as BF16 and runs the β and α projections as one launch; a 4096-token prefill is 0.9% faster.
Code shape
The review reads these first. A box left unticked needs one line of why.
conv_body) over aConvIntrait, and the BF16 words are read through the existingBf16Words. The BF16 store is an output form of the existing product launch owners, not a second launcher.site::resident_gemv_rows, the gemv-sideSiteWeightargument ofsite::launch, the verify-nodeVERIFY_GDN_EXTRAconstant and the gate clause that called the removed function are removed, and the pinned node counts are re-derived.n_vis a named error with a gate clause, and a BF16 store on a site that has none reports that instead of launching.tools/check-comments.shis ok.PIN(2026-10-31)lines with their reason.cargo fmtandcargo clippyare clean on what I touched; the lint count did not go up.Duplication found or removed (
path:line), or "none": the duplicate scan lists eight candidates, all different.crates/gpu/src/linear/conv.rs:284andcrates/gpu/src/linear/conv.rs:388are the f32 and BF16store_ring, which read different lane layouts.crates/gpu/src/gemm/f32tile.rs:2379andcrates/gpu/src/gemm/f32tile.rs:2459, andcrates/gpu/src/gemm/f32tile.rs:2470andcrates/gpu/src/gemm/f32tile.rs:2368, are launch-shape matches for two different MMA forms.crates/gpu/src/gemm/kernels32.rs:1368withcrates/gpu/src/gemm/kernels32.rs:1324and withcrates/gpu/src/q8f32.rs:2147,crates/gpu/src/gemm/nvfp4mma.rs:1850withcrates/gpu/src/gemm/nvfp4mma.rs:1799andcrates/gpu/src/linear/conv.rs:772withcrates/gpu/src/linear/conv.rs:713are kernel entries that call one shared body, because a#[kernel]entry cannot be generic.crates/gpu/src/site.rs:1639andcrates/gpu/src/site.rs:1622are the f32 store, which always launches, and the BF16 store, which reports whether the site has one.What I ran
Card conditions of every run below, from
nvidia-smi:Speed
One FP4-Local checkpoint (
Qwen3.8-Flash-Next-FP4-Local) on one DGX Spark, the defaultsm_121abuild. Each request is a 4096-token prompt with 128 greedy tokens. Three rounds per side in alternating order after a warm-up, medians below. The base is upstream main at904266cfc. The first two rows are the draft-off arm and the rest are the MTP draft arm.On the draft-off arm, prefill and TTFT improve by more than the same-revision resolution of the measurement (0.568% and 0.565%), and the ranges of the two sides are separate. Every other cell is inside its resolution (0.729% to 7.114%), so I claim a gain only for the draft-off prefill and TTFT.
Reproduce one round, with
baseon upstream main andheadon this branch (the MTP rounds addBLOOMERY_DRAFT=mtp BLOOMERY_MTP_DRAFT="${MTP_MODEL}" BLOOMERY_MTP_WIDTH=cost):GPU gates
gate_linearcovers the BF16 conv (reading channels as BF16 words over a random ring, β and α at pitch 64, and the refused pitch ofn_v - 1).gate_gemm --case dense_mmacovers the BF16 store of the BF16 and the MXFP8 W8A8 products. The summary lines:The BF16 store clause tells rounding from truncation. With a store that truncates, the same run fails it:
The verify node counts that β and α as one launch changes are re-derived, and the fixture tier of the end-to-end gate holds them:
The by-name site clauses after the removal of the gemv row window:
The PTX spill ratchet, with the three new entries pinned:
The A/B above ran the real
bloomery-serveof this change on the card.Hardware (if the change is for a machine)
DGX Spark (GB10), driver 595.84, CUDA 13.3, build
sm_121a.mxfp8_w8a8_gemm_bf16is ansm_121aentry.AI help
AI 코딩 에이전트가 코드를 작성한 후 설명을 썼고, 테스트를 마쳤습니다. diff 리뷰와 PR 초안 리뷰는 제가 직접 했습니다.