Skip to content

[Draft] EXL3 Flash Next: DeltaNet heads split across the two ranks for decode rows, opt-in (slice 7, #1419) - #52

Draft
BlivionIaG wants to merge 1 commit into
chain/s6-drafts-guardfrom
chain/s7-gdn-split
Draft

BlivionIaG wants to merge 1 commit into
chain/s6-drafts-guardfrom
chain/s7-gdn-split

Conversation

@BlivionIaG

Copy link
Copy Markdown
Owner

Draft, slice 7 of the chain in ashhart#1419 (opt-in, off by default). Stacked on slice 6; review only the last commit.

What

TF_FLASHNEXT_GDN_SPLIT=1 splits the 48 DeltaNet value heads (and 16 key heads) of every linear-attention layer across the two TP=2 ranks for decode rows (rounds and windows under 64 rows). Rank r runs key heads 8r..8r+7 and value heads 24r..24r+23: its q, k, v and z projections read half of the replicated EXL3 weights through column views (a pointer offset and a shorter extent, no copy), the recurrence updates half of the state in the whole-layer layouts, and the out projection, cut by input rows, gives a partial sum that one fp32 all-reduce adds. Prompt passes keep running every head, so weights, conv windows and states keep their sizes and layouts.

Files: cuda_exl3_gdn_split.zig (views), head-subset variants fn_gdn_front_h, fn_gdn_rows_h, fn_gdn_rows_to_h, fn_gdn_norm_h in flashnext_batch.cu, and the hooks in cuda_batch.zig / cuda_prompt_load.zig.

Why it is opt-in

Measurements (2x RTX 5060 Ti sm_120, TP=2, one user, 192 tokens, temperature 0, second request)

prompt slice 6 (replicated) decode tok/s slice 7 with TF_FLASHNEXT_GDN_SPLIT=1
plain decode, 88 lines short 48.9, zig 36.2, docs 38.2 short 51.4, zig 37.4, docs 40.0
drafts on, 80 lines short 46.5, zig 104.7, docs 57.3 short 46.2, zig 104.8, docs 57.5

Prefill is unchanged by design. Extending the split to prompt passes was tried and rejected (TTFT +18%).

Tests run

Build, 294 EXL3 host tests, compile-cuda-tests, test-kernel-contract on the slice below; this slice's own checks are still to be rerun on the final stack. Not run: the on-GPU kernel-contract gates, the FP64 references, other architectures (sm_120 only).

Not done in this PR

  • The four new _h kernels are not yet registered with the kernel contract.
  • No FP64 reference or paired-error receipts (see above).

…code rows (opt-in TF_FLASHNEXT_GDN_SPLIT=1)

Rank r runs key heads 8r..8r+7 and value heads 24r..24r+23 of every linear-attention layer for rounds and windows under 64
rows: its q, k, v and z projections read half of the replicated EXL3 weights through column views (no copy), the recurrence
updates half of the state in the whole-layer layouts, and the out projection, cut by input rows, gives a partial sum that one
fp32 all-reduce adds. Prompt passes still run every head. The fp32 summation order differs from the replicated path, so the
text is not byte-identical to it (24 prompts: 20 identical, 4 diverge after 84-423 characters, all checkable answers correct).

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.

1 participant