Skip to content

qwen3_5 cuda: a dense MLP's NVFP4 prompt gate, up and act in one lean launch, the same bytes - #3252

Open
jschmied wants to merge 1 commit into
ashhart:mainfrom
jschmied:fusedact
Open

jschmied wants to merge 1 commit into
ashhart:mainfrom
jschmied:fusedact

Conversation

@jschmied

Copy link
Copy Markdown
Contributor

Head 19cdfdfbd on main 223da7659b. 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_lean products (gate, up: 17408 x 5120 each), writes
both 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.

GB10, Qwen3.8-27B NVFP4 (unsloth, nvfp4 head), DFlash2, cold prompt, A-B-B-A main this PR paired median slower
8K (8175-8194 tokens) x3 a arm 6.455-6.622 s 6.249-6.402 s -3.28 % 0 of 6
32K (32666-32689 tokens) x2 a arm 27.908-28.030 s 27.101-27.159 s -2.99 % 0 of 4
decode, 1 stream, 256 tokens 53.80 / 53.90 tok/s 54.08 / 54.12 tok/s +0.52 / +0.40 % 0 of 2

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 at
2048 / 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.cu includes it, 397 lines): prompt_lean moved here and
    templated 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_SWIGLU a tile is 128 features: gate's two
    64-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 wn takes gate subtiles 2wn, 2wn+1 and up's 16+2wn, 17+2wn, so one lane holds a feature's
    gate 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 takes tf_q35_mlp_act's steps
    (g / (1 + expf(-g)) * u, rounded), skipping rows past P[2] as that kernel does.
  • zig/src/cuda/nvfp4.zig: Prompt.swiglu and Prompt.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.zig swiglu + cuda_forward.zig: a dense layer whose gate and up are NVFP4
    takes 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.
  • Registry: the renamed lean symbols and a record for the pair (255/255/255/247/247 registers, sm_80/86/89/120/121).

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 and
HIP are untouched (HIP has its own epilogue form, #2344).

Same bits

  • Kernel contract (tf-cuda-test prompt-swiglu-contract, also run after the gates in kernel-contract; file
    zig/tests/cuda/contract_prompt_swiglu.zig): the pair against main's path (two prompt_lean launches, then
    tf_q35_mlp_act) byte for byte, SUMMARY prompt_swiglu: 3 pass, 0 fail; 2 controls red, 0 missed:
    • served shapes: 6 bit-equal (2048 and 4096 x 17408 x 5120, 129 rows, 1 x 128 x 64, 17 x 256 x 1088, 300 x 384 x 576),
      non-finite inputs and a NaN block scale included, five repeats byte-identical, guards intact;
    • the device row count (P[2] = rows - 37): 3 shapes bit-equal, rows past it left as the act kernel leaves them;
    • every bf16 gate value: 65536 rows whose gate is exactly row r's bit pattern (a 1.0 code at input 0) against 128
      up values each, through the real kernel: bit-equal, NaN, Inf and denormals included;
    • planted controls, red at shapes: gate and up swapped (contract_prompt_swiglu.zig .swapped), and the gate's
      scale one fp32 ulp up (.gate_ulp, a pre-rounding change the rule must see).
  • Off == main: the two existing lean instances keep their SASS: cuobjdump -sass of prompt16 from main and this
    PR, 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=121 set: 113 of 114 images byte-identical (sha256); prompt16 differs only by the added
    function. Filtered kernel-contract prompt16: 7 kernels pass gates 1-4, live limits PASS (247 registers recorded).
  • Replies equal to main (27B, DFlash2 drafting and --no-drafts): 49 of 49 identical: greedy and seeded (temp
    0.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 == fresh: a 3-turn chat over an 8K document, greedy and seeded, with the prompt cache (5 turns
    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

  • x86 (-Dfixed-cpu test test-golden test-kernel-contract test-server-security test-pull-security fuzz-cortheon check-release-paths, tools/pr_check.py lean 341 vs 341 (0 new), zig fmt --check, git diff --check): pass.
  • GB10 (nvcc 13.0, sm_121): 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-host 400/400 steps, 6332 tests pass
    (30 skipped); prompt-swiglu-contract and kernel-contract prompt16 as above.
  • Not run: the full kernel-contract (40-60 min; its other kernels' images are byte-identical to main's).

Notes

  • Overlaps qwen35 cuda: fuse the dense MLP's NVFP4 prompt gate and up projections into one launch #2410 (GustavBlack, gate+up in one launch, act kernel kept; level on GB10): both touch prompt16 and the
    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.
  • No new refusal or refusal test: Prompt.swiglu returns FormatMismatch for a non-pair, but the forward asks
    Prompt.pairs() first and runs today's three launches otherwise, so nothing refuses on a served path.
  • Bandwidth-bound here; on a discrete card (sm_120 5090) the saved write/read and act kernel should matter at least
    as much. A tester run (Maxrio9 sm_120) of the 8K/32K cells would be welcome.
  • Proposed CHANGELOG line (Unreleased): "Qwen3.8 on CUDA runs a dense layer's NVFP4 prompt gate, up and SwiGLU in one
    launch; replies are byte-identical and cold prefill on GB10 is about 3 % quicker at 8K and 32K."

… 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.
@ashhart

ashhart commented Nov 14, 2026

Copy link
Copy Markdown
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.

  • Synthena (Ash's Agent)

This branch has not been deployed

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants