Skip to content

Flash Next CUDA: the prompt DeltaNet block in one kernel a value head - #1217

Closed
jschmied wants to merge 2 commits into
ashhart:mainfrom
jschmied:fn-gdn-fused
Closed

Flash Next CUDA: the prompt DeltaNet block in one kernel a value head#1217
jschmied wants to merge 2 commits into
ashhart:mainfrom
jschmied:fn-gdn-fused

Conversation

@jschmied

@jschmied jschmied commented Oct 23, 2026

Copy link
Copy Markdown
Contributor

A lone stream's prompt pass runs the GDN block as three launches: fn_gdn_front writes every row's q, k, v, g and beta as fp32 (41 KB a row, 63 MB for a 1,536-row chunk), fn_gdn_chain_tpr reads them back for the delta rule, and fn_gdn_norm reads the fp32 readouts again for the gated norm. On a GB10 the front is bound by DRAM (~240 GB/s), so overlapping it with the chain on a second stream hides almost nothing (measured: -0.5 % at 1,536 rows).

fn_gdn_fused keeps those rows in shared memory:

  • 256 threads walk the chain with fn_gdn_chain_tpr's code;
  • 128 producer threads make the next 16 rows' inputs with fn_gdn_front's op sequence;
  • the same producers run the gated norm of the previous 16 rows with fn_gdn_norm's.

Every row's arithmetic is unchanged, so the bits are too. The recurrence is row by row, so where the rows are cut (chunk size, resume point) does not change them either. Passes of 64 rows or more on one stream use it; TF_FLASHNEXT_GDN_FUSED=0 keeps the three launches.

GDN block, GB10 three launches fused
1,024 rows 1.277 ms 1.040 ms
1,536 rows 1.907 ms 1.547 ms
4,096 rows 5.008 ms 4.078 ms

Bytes (state, readouts, gated rows, sums) are equal from 64 to 4,096 rows.

Checks

  • Kernel contract: nonfinite, repeat and 6 shapes bit-equal to front -> rows_to -> norm. The gate-5 record covers sm_80/86/89/120/121 (nvcc 13.0.88; no 13.2/13.3 override).
  • Full zig build test -Dnvcc=/usr/local/cuda/bin/nvcc: exit 0 (278/278 steps; 3,491 of 3,506 tests passed, 15 skipped). lean_check: 0 new against the merge base. zig fmt is clean.

Served Flash Next FP4-Local (MTP and prompt cache on, chunk 1,536)

Token hashes: 136 of 136 equal to main's. This covers greedy and seeded, c=1 and c=4, and cold, warm and resumed requests.

cold TTFT main fused
~10k tokens 7.81 / 7.85 s 7.72 / 7.67 s
~38k tokens 30.98 / 30.24 s 30.02 / 29.93 s

That is about -1.5 to -2 %.

Gotchas

@jschmied
Flash Next CUDA: the prompt DeltaNet block in one kernel a value head
d241180 
A lone stream's prompt pass ran the GDN block as three launches: fn_gdn_front
wrote every row's q, k, v, g and beta as fp32 (41 KB a row, 63 MB for a
1536-row chunk), fn_gdn_chain_tpr read them back for the delta rule, and
fn_gdn_norm read the fp32 readouts again for the gated norm. On a GB10 the
front is DRAM-bound (~240 GB/s), so running it beside the chain on a second
stream hides almost nothing.

fn_gdn_fused keeps those rows in shared memory: 256 threads walk the chain
(fn_gdn_chain_tpr's code) while 128 producer threads make the next 16 rows'
inputs with fn_gdn_front's op sequence and the gated norm of the previous 16
with fn_gdn_norm's. Each row's arithmetic is unchanged, so the bits are too.
The recurrence is row by row, so no cut of the rows changes them either.
Passes of 64 rows or more with one stream use it; TF_FLASHNEXT_GDN_FUSED=0
keeps the three launches.

- Kernel contract: nonfinite, repeat and 6 shapes bit-equal to
  front -> rows_to -> norm; gate-5 record (sm_80/86/89/120/121).
- GDN block alone, GB10: 1.907 -> 1.547 ms at 1536 rows, 5.008 -> 4.078 ms
  at 4096 (-19 %); bytes equal from 64 to 4096 rows.
- Served FP4-Local, MTP and prompt cache on: 136 of 136 token hashes equal
  to main's (greedy and seeded, c=1 and c=4, cold, warm and resumed);
  cold TTFT at ~10k 7.81/7.85 -> 7.72/7.67 s, at ~38k 30.98/30.24 ->
  30.02/29.93 s.
@jschmied
Flash Next CUDA: fused GDN kernel comments within the lean check
c184e94
ashhart added a commit that referenced this pull request Oct 23, 2026
@ashhart
Maintainer fix-up for #1217: the prompt contract serves the admitted …
e03144b 
…rows and the state alias

The registered contract covered 64 and 300 rows, so the shapes gate never ran
the default chunk (1,536) nor the most a chunk may ask for (16,384). The served
shapes now carry the production chunk constants, and the gate derives the tile
seams from them.

Production binds the fused block's state as both its input and its output, but
the contract allocated the two apart and declared no alias. The spec now
declares the in-place state, the launch passes one buffer for it, and every
launch restores the starting state, so the repeats compare the same pass.

The new host test in contract_flashnext_glue_test.zig covers the served range
and the alias.
ashhart added a commit that referenced this pull request Oct 23, 2026
@ashhart
Maintainer fix-up for #1217: the contract glue's callers use the buff…
c0b0afc 
…er index, and the Linux build compiles the glue

The M45 fix-up changed Glue.out to return the buffer index (!u64), but the chain,
scores, fused and read-out cases still called it as a statement, so Zig 0.17
rejected the discarded value:
zig/tests/cuda/contract_flashnext_glue.zig:166:17: error: value of type 'u64'
ignored. The callers discard the index where nothing needs it, and the fused
case keeps it, where the state is bound in place as its own output.

The gate that would have caught it: the runner now compiles the contract glue's
device path (main.zig's compileContractGlue), so the Linux build of the CUDA
test binary sees the same body the device runs. Without the fix, `zig build
compile-cuda-tests -Dtarget=aarch64-linux-gnu` fails on that error; with it,
the step succeeds (10/10 steps, the tf-cuda-test executable for
aarch64-linux-gnu). The host contract test still passes (2/2).

Nothing of the contributor's is amended.
@ashhart

ashhart commented Oct 23, 2026

Copy link
Copy Markdown
Owner

Thanks @jschmied, this is on main now as c0b0afc: your two commits with your name and dates, plus two maintainer fix-ups. The prompt contract serves the admitted rows and the state alias, and the contract glue's callers use the buffer index, with the glue compiled in the Linux build.

What we checked on a DGX Spark (GB10):

  • Kernel contract: every gate passes, and every planted control goes red.
  • Qwen3.8 Flash Next MLX 4-bit: served replies identical to main's. One early 64k cell looked 19.7% slower, but alternated reruns put it at +1.3%; the gap was page-cache noise.
  • nvidia/Qwen3.8-Flash-Next-FP4-Local: tokens identical to main in every cell, prefill 0.5 to 1.6% faster from 2k to 64k, and the 1/2/4/8-stream sweep within 0.6%.

Closing as landed. It ships in the next release.

- Synthema, Ash's AI Assistant.

@ashhart ashhart closed this Oct 23, 2026
ashhart added a commit that referenced this pull request Oct 24, 2026
@ashhart
Carried onto main: #1155's shared GDN pass on #1217's fused block
7faf345 
#1155's first record and main's #1217 both change gdnBlock in
zig/src/families/flashnext/cuda_batch.zig: the record keeps each lane's rows
on its own conv window and state, and #1217 runs a lone pass's prompt DeltaNet
in one kernel. The record carries the shared pass; this commit carries #1217's
block onto it.

zig/src/families/flashnext/cuda_batch.zig:272-276: a lone pass keeps its
projection and conv window before the fused branch, so the one-block path fills
the rollback's keep as main does.
zig/src/families/flashnext/cuda_batch.zig:278-283: the fused launch for a lone
pass of 64 rows or more.
zig/src/families/flashnext/cuda_batch.zig:305: the gated norm through
gdn_fused.norm, the same kernel and grid as the inline launch it replaces.

Refs #1155
@ashhart ashhart mentioned this pull request Oct 24, 2026
Closed
9 tasks done
ashhart added a commit that referenced this pull request Oct 24, 2026
@ashhart
Maintainer fix-up for #1154: complete scores on worker exit and refus…
1a4b6bf 
…e ranked scoring

Carried onto main: the CUDA build wiring keeps main's own registrations
beside the new one. Main holds the architecture contract test inline
(zig/build/cuda.zig:285) and #1217's glue host beside it (:286); the
scoring host takes the proposal host's place (:287) because
zig/build/flashnext_scoring_host.zig:4 calls the proposal host's own
steps, so the proposal checks and the score step both run. The file
keeps its 600 lines. The worker's exit still closes the scoring queue
(zig/src/core/lane_host.zig:522 and zig/src/core/lane_host_score.zig:66),
and the ranked vtable still advertises the score hook for one rank only
(zig/src/families/flashnext/cuda_lanes.zig:214, :217 and :218).
@ashhart

ashhart commented Oct 24, 2026

Copy link
Copy Markdown
Owner

This shipped in TensorFold 1.1.4: https://github.com/ashhart/TensorFold/releases/tag/v1.1.4

Thank you @jschmied for the contribution.

- Synthema, Ash's AI Assistant.

ashhart added a commit that referenced this pull request Oct 25, 2026
@ashhart
Keep the fused GDN launch record beside the paging measurements
d2d624c 
Carried onto main: zig/src/core/contract/flashnext_tile_records.zig:13 retains #1217's fn_gdn_fused record (384 threads and 82496 shared bytes), while :14 keeps #1169's measured qsa_tile register counts of 218 on sm_120 and sm_121. The paging entry-point records in flashnext_pages_pending.zig stay as authored.

Both registrations describe kernels present on this line; neither record replaces the other.
@ashhart ashhart added the landed Landed on main by the maintainers; the commits keep the contributor's authorship label Oct 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

landed Landed on main by the maintainers; the commits keep the contributor's authorship

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@jschmied @ashhart