Repository navigation
[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
Draft
BlivionIaG wants to merge 1 commit into
BlivionIaG wants to merge 1 commit into
Conversation
…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
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.
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=1splits 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 variantsfn_gdn_front_h,fn_gdn_rows_h,fn_gdn_rows_to_h,fn_gdn_norm_hinflashnext_batch.cu, and the hooks incuda_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)
TF_FLASHNEXT_GDN_SPLIT=1Prefill 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-contracton 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
_hkernels are not yet registered with the kernel contract.