Skip to content

cuda: NVFP4 experts' prompt rows take the warp's row tile count as a constant - #805

Merged
ashhart merged 1 commit into
ashhart:mainfrom
Bizuayeu:pr/nvfp4-experts-prompt-tiles
Oct 14, 2026
Merged

ashhart merged 1 commit into
ashhart:mainfrom
Bizuayeu:pr/nvfp4-experts-prompt-tiles

Conversation

@Bizuayeu

@Bizuayeu Bizuayeu commented Oct 13, 2026 •

Copy link
Copy Markdown
Contributor

What this changes

The staged NVFP4 expert prompt kernel reads its row-tile count at run time inside the mma chain; this makes it a
template constant with one instantiation per count, so the branch leaves the chain. The sums keep their order and
every output row keeps its bytes. #764 (skip on the same prompt entries) touches the same three call sites; whichever
lands second takes the other's lines with a small rebase.

Receipt

  • Environment: 1.0.9 (c36b536d) with this commit for the tests, on one NVIDIA GB10 (sm_121, CUDA 13.3). The
    exactness and timings are the same commit on 1.0.8 (0efbd447): the kernels on that GB10, the prefill through
    the GLM-5.3-Flash CUDA family (glm: GLM-5.3-Flash (NVIDIA's NVFP4 checkpoint) as a CUDA family at TP over hosts #804, nvidia/GLM-5.3-Flash-NVFP4 at 423acf37) at TP=2 on two (driver
    580.173.02, CUDA 13.0, Linux 6.17.0-1032-nvidia).
  • Exactness: the old and new kernels' outputs (swiglu, down, relu²) hash equal at 1 to 2,048 rows in 11 row counts on
    random inputs at GLM-5.3-Flash's per-rank widths (1,024 x 4,096, 17 experts, 9 slots); tf-cuda-test quant-ref
    4 passed. Through the family at TP=2 the 8k and 32k prompt checks gave the same ids with and without it. On
    1.0.9, tf-cuda-test nvfp4-experts against the Python engine's recorded bytes (packing, plan, up and down, at its
    default shape and a GLM-like one) passes with and without it.
  • Prompt speed: not from tools/prefill_cold.py (1.0.0 removed it; the family is not upstream yet). Kernels at
    2,048 rows (ncu, median of 10): swiglu 5,565 → 3,912 us (29.7% shorter), down 2,500 → 2,116 (15.4%), relu²
    4,673 → 3,810 (18.5%). The family's prefill at TP=2 (median of 2): about 2k tokens 1.540 → 1.480 s, 8,192
    tokens 5.610 → 5.290 s (5.7% shorter), 32,768 tokens 21.855 → 20.705 s (5.3% shorter).
  • Decode speed: unchanged by construction (the decode kernels are not touched).
  • Tests run (empty cache, this commit alone on 1.0.9): zig build test-cuda-host 429/429, zig build test 644/649
    (5 skipped), tf-cuda-test quant-ref 4 passed, zig fmt.
  • Not run, and why: Nemotron's NVFP4 prompt path (no checkpoint here); Metal (not touched; no Mac here).

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
  • No precision traded for speed. If the bits change, the description says which sums change and why.
  • Prompt processing is no slower than the last release.
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. (No new test: every
    output keeps its bytes, which quant-ref, nvfp4-experts and the comparison above check.)
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

@Bizuayeu

Copy link
Copy Markdown
Contributor Author

Updated the description: on 1.0.9, tf-cuda-test nvfp4-experts against the Python engine's recorded bytes passes with
and without this commit, and the checklist's test item is ticked with its reason (no new test; the outputs keep their
bytes).

@ashhart
ashhart merged commit f243307 into ashhart:main Oct 14, 2026
4 checks passed
ThinkOffApp pushed a commit to ThinkOffApp/TensorFold that referenced this pull request Oct 14, 2026
ThinkOffApp pushed a commit to ThinkOffApp/TensorFold that referenced this pull request Oct 14, 2026
@ashhart

ashhart commented Oct 14, 2026

Copy link
Copy Markdown
Owner

Thank you, @Bizuayeu. Nemotron's NVFP4 expert prompt rows take the warp's row tile count as a compile-time constant, so one mma body runs each count. The store mask is still m0 < cnt.

On a GB10, teacher-forced rows are bit-equal to the base, 16,109 of 16,109. Served replies are equal, 17 of 17. Drafted replies equal "draft": false, 17 of 17. Cold prefill is 4.1% faster at 2k tokens, 7.5% at 8k, 6.9% at 16k, 6.6% at 32k, and 5.3% at 64k.

Your commit is in 1.0.10 under your name (4e63c6528). The patch matches your head (4ed879dc3). The CHANGELOG line is 343e64398. Your head 4ed879dc3 is merged at f2433079b with no change to the tree. We wrapped prompt_body's parameter lines (46ebd305a).

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