Repository navigation
Conversation
… pair launch, the same bytes prompt_lean gains a SwiGLU pair form: a 128 x 256 tile stages gate's two 64-column blocks and up's at the same features, each warp holds a feature's gate and up sums in one lane, and the epilogue rounds both to bf16 and takes tf_q35_mlp_act's steps. The two existing instances keep their SASS on sm_80/86/89/120/121 (renamed: int out kind, four trailing arguments). Qwen3.5 dense layers with NVFP4 gate and up take it for prompt rows past 16.
Owner
|
Thanks @jschmied. The host tests pass, and the contract compares the fused result with the replaced operations without loosening byte equality. We will run the device check on our side for the new pair and the existing BF16 and FP32 lean launches, including partial tiles and shortened active row counts. Outputs must be bit identical, both planted controls must fail, and launch limits must pass for the compiled CUDA targets. We will also check served tokens and prompt and decode speed at one and four streams.
|
This branch has not been deployed
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
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.
Head
19cdfdfbdon main223da7659b. GB10 only measured; exact, on by default for NVFP4 gate/up prompt rows past 16.On main a Qwen3.8-27B dense layer's prompt MLP runs two
prompt_leanproducts (gate, up: 17408 x 5120 each), writesboth as bf16, reads them back in
tf_q35_mlp_act, and writes the act. bound1110 measured that act kernel at 3.1-3.4 %of cold TTFT at the bandwidth floor. This PR computes gate and up of the same 128 features in one lean tile and takes
the act in the epilogue, so neither bf16 product is written or read again and the act kernel does not run.
Fresh server per arm, page cache dropped before each start, 90 s between arms; SM clock per arm 2279 / 2278 / 2283 /
2284 MHz (the box's 2300 MHz cap; throttle flags are that cap's, same in every arm). Decode texts identical in all arms.
These rows are on this head's binary; an earlier run on 3b21a9a68e (same kernels, comments since) gave -3.26 % / -2.87 %.
Kernel level (
tf-cuda-test prompt-swiglu-bench, same process, A-B-A): two lean products + act vs the pair launch at2048 / 4096 / 8192 rows: 9.16 vs 7.96 ms (1.15x), 18.44 vs 15.44 ms (1.19x), 36.17 vs 30.81 ms (1.17x).
What changes
zig/kernels/cuda/prompt16_lean.cuh(new;prompt16.cuincludes it, 397 lines):prompt_leanmoved here andtemplated on an output kind (
LEAN_BF16,LEAN_F32,LEAN_SWIGLU) with four trailing arguments (up's words,scales and scale, and the act kernel's device row count). For
LEAN_SWIGLUa tile is 128 features: gate's two64-column blocks and up's two at the same features are staged as the usual 256-column tile (no repack, no second
weight copy), and warp
wntakes gate subtiles2wn, 2wn+1and up's16+2wn, 17+2wn, so one lane holds a feature'sgate and up sums. Each sum is its own column's mma chain over the same K order as a plain lean product. The epilogue
rounds both to bf16 (
acc * scale, as the bf16 lean epilogue) and then takestf_q35_mlp_act's steps(
g / (1 + expf(-g)) * u, rounded), skipping rows pastP[2]as that kernel does.zig/src/cuda/nvfp4.zig:Prompt.swigluandPrompt.pairs(both FP4, same n/k/npad, whole 128-feature tiles);the lean launch passes the four trailing arguments (zeros for one product).
zig/src/families/qwen3_5/cuda_modelopt.zigswiglu+cuda_forward.zig: a dense layer whose gate and up are NVFP4takes the pair for prompt rows past 16 with no learning hook and no group sums; everything else (FP8/fp8g layers,
EXL3, affine-4, BF16, decode/verify rows,
--learn) runs today's three launches.Shared code: the kernel is the CUDA core prompt tile (
prompt16_lean.cuh), the launcher core (cuda/nvfp4.zig);only the wiring is qwen3_5's. Any family with NVFP4 gate/up prompt rows and SiLU can call
Prompt.swiglu. Metal andHIP are untouched (HIP has its own epilogue form, #2344).
Same bits
tf-cuda-test prompt-swiglu-contract, also run after the gates inkernel-contract; filezig/tests/cuda/contract_prompt_swiglu.zig): the pair against main's path (twoprompt_leanlaunches, thentf_q35_mlp_act) byte for byte,SUMMARY prompt_swiglu: 3 pass, 0 fail; 2 controls red, 0 missed:non-finite inputs and a NaN block scale included, five repeats byte-identical, guards intact;
P[2]= rows - 37): 3 shapes bit-equal, rows past it left as the act kernel leaves them;up values each, through the real kernel: bit-equal, NaN, Inf and denormals included;
contract_prompt_swiglu.zig.swapped), and the gate'sscale one fp32 ulp up (
.gate_ulp, a pre-rounding change the rule must see).cuobjdump -sassofprompt16from main and thisPR, per function, sm_80/86/89/120/121 (nvcc 13.0): 7 of 7 functions identical, 1 added (the pair). The built
zig build fatbins -Dsm=121set: 113 of 114 images byte-identical (sha256); prompt16 differs only by the addedfunction. Filtered
kernel-contract prompt16: 7 kernels pass gates 1-4, live limits PASS (247 registers recorded).--no-drafts): 49 of 49 identical: greedy and seeded (temp0.8, top-p 0.95, seed 1234) one at a time, 4 and 8 at once, greedy over 2K and 8K prompts (2048/4096-row pair
launches); drafted == plain greedy on both binaries.
resumed at row 8192, their 24-125 new rows on the pair launch) against the same turns on a cache-off server: equal on
both binaries, and main == this PR.
Checks on
19cdfdfbd-Dfixed-cpu test test-golden test-kernel-contract test-server-security test-pull-security fuzz-cortheon check-release-paths,tools/pr_check.pylean 341 vs 341 (0 new),zig fmt --check,git diff --check): pass.test-cuda-host test-qwen35-launch-args-host test-qwen35-modelopt-admission-host test-qwen35-stack-symbols-host test-qwen35-formats-host test-qwen35-prefill-host400/400 steps, 6332 tests pass(30 skipped);
prompt-swiglu-contractandkernel-contract prompt16as above.kernel-contract(40-60 min; its other kernels' images are byte-identical to main's).Notes
qwen3_5 dense MLP site. This PR does the act in the epilogue, which is where the GB10 gain is. Either can land first;
I'll rebase onto the other when you say.
Prompt.swiglureturnsFormatMismatchfor a non-pair, but the forward asksPrompt.pairs()first and runs today's three launches otherwise, so nothing refuses on a served path.as much. A tester run (Maxrio9 sm_120) of the 8K/32K cells would be welcome.
launch; replies are byte-identical and cold prefill on GB10 is about 3 % quicker at 8K and 32K."