Repository navigation
dsv41 Zig: the TFE draft head (TF_DS_DRAFTER=eagle), off by default - #42
Elevatormusic wants to merge 12 commits into
Conversation
An EAGLE-style draft head trained on the target's taps (the EAGLE study's eagle_train.py Head) drafts in place of DSpark: eagle.zig (rings, absorbs, the pass), eagle.cu (its row kernels), the converter and the reference (tools/dsv41-zig/eagle), and the self-test (TF_DS_EAGLE_CHECK). Rank 1 runs no head: rank 0 sends every draft in the verify frame.
…rnel test TF_DS_ACCEPT_LOG=N writes one JSON line on stderr each N verify rounds: the drafts asked and kept at each window position (a draft is kept when it and every earlier draft equal the round's draws), for the TFE head and DSpark alike. eagle_kernels_test.py checks eagle.cu's kernels against torch on any CUDA GPU.
…ck client An absorb of more rows than the ring wrote two rows to one ring row in one scatter, and either one could stay: the self-test's case q0 = 160 (160 rows, a 144-row ring) failed its 3% bound. The self-test now absorbs at most a ring of rows a time, and absorb refuses a larger span. Serving absorbs at most a ring of rows, as before. tfe_bench.py runs prose, code and tool-call prompts at 1 and 4 streams with drafts on and off and compares token_sha; tfe_seg.py splits the engine's TF_DS_ACCEPT_LOG lines into the drafted batches; tfe_sum.py tabulates a run.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 73084ac8f5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
…ocabulary head A head file now carries qk_norm (an RMSNorm on each query and key head before RoPE), head_type and cand_k. A two_stage head takes the top cand_k low-rank logits (fp32, a tie at the k-th value by the lower token id) and drafts the argmax of their exact logits from the target's own head, a dense bf16 copy in the head file on rank 0 (eagle_convert.py --target-head: the head dump files, their sha256 and shapes in the header). TF_DS_EAGLE_TOPK sets the candidates (1 to 512). The window goes up to 511 (512 keys), and a prompt chunk gives the head its last ring of rows. New kernels: tf_eagle_topk_kernel, tf_eagle_rescore_kernel, and the top-1 kernel through an id map. eagle_ref.py follows the same rules and checks its drafts against the trainer's two_stage logits.
The steps of a pass (each step's layer, logits and draft for every stream) are captured once a (streams, steps) and replayed; the tables they read are uploaded before, at fixed places. The extend of the real rows stays eager. Off by default.
…nears
The profile of window tfe-1792838409: a 5-step pass is GPU-bound (10.1 ms
busy of 10.3 ms); cuBLAS's few-row kernels read wo and w2 at 165 GB/s, the
attention takes 100 us at window 128 and the top-1 51 us on one block.
- eagle.cu: tf_eagle_gemv_{bf16,fp8}_kernel for up to 16 rows (out bf16,
fp32, or added to bf16); the attention with a thread a key and the
values split over the warps (window 511 in one block); the top-1 of a
vocabulary row in 64 parts and a join.
- eagle.zig: Lin linears; TF_DS_EAGLE_GEMV=wo (default: wo, w2 through the
own GEMV) | all | cublas; TF_DS_EAGLE_FP8=1 loads the e4m3 forms with a
row scale (own GEMV only); the loader uploads only the tensors in use.
- eagle_convert.py writes each linear's FP8 form (s = amax / 448, e4m3
RNE) and its sha256; eagle_ref.py --fp8 vectors and fp8cmp.
- eagle_kernels_test.py covers the new kernels; eagle_kernels_time.py.
…_stage head's candidates differ)
… qk_norm shape for the cross and self tests
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a300c333e6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| const total = @min(ch.n, eg.ring); | ||
| const cut = ch.start - start0; | ||
| var at = ch.n - total; // (the rows before: no window reaches them) |
There was a problem hiding this comment.
For a long prompt with a supported 511-token head window, this absorbs only the target's replay suffix, not the head's full context. Lanes.fillFn sets replay = len -| m.cfg.window (128 tokens), and prompt.replayCut trims ch.n before this function runs; earlier encoder-only chunks never reach this absorb. Consequently, a 4096-token prompt initializes only 128 head rows instead of the 512 keys required by its attention window, and Eagle.pass silently clips attention to that truncated history. The initial draft chains therefore compute different outputs from the full-window reference, potentially depressing acceptance for hundreds of generated tokens. Initialize the head's required prompt history, or explicitly reject unsupported window/replay combinations rather than silently truncating them.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Known limit, now explicit in 05862c8. The rows below the replay floor (prompt length - sliding_window 128) do not run the layers the taps come from (26, 37 and the final state), so the head cannot get them without a longer replay. A longer replay changes the target's prompt path and its compressed-layer snapshots, so it is a separate change. For now Eagle.open logs a warning for a head window past the model window, and tools/dsv41-zig/eagle/README.md states the limit. The deployed heads (b1f, main3) have window 128. The EAGLE long run trains with --serve-floor 128, so its heads see the serving history in training; earlier long-window heads did not. I leave this thread open for the reviewer to decide on the longer replay.
…-path tap check - FP8 is the default when the head file has the FP8 forms (main3 met the rule: position-1 acceptance within 0.005 on every tag); TF_DS_EAGLE_FP8=0 runs the bf16 forms, =1 requires the FP8 ones. A file without them runs bf16, as before. - TF_DS_EAGLE_TAPDUMP=FILE (rank 0): every prompt-chunk and round absorb writes its tap rows with the chunk's replay cut (eagle.zig tapDump). - tfe_tapcheck.py: three prompts that share a prefix and cut at three places; the prompt records must agree at their shared positions (a tap below the middle layer reads its rows at the cut).
- eagle.zig is split: eagle_file.zig (the head file and Eagle.open) and eagle_check.zig (the self-test). Each file is below 600 lines. - Every comment is one line of 120 columns or less (eagle.zig, eagle.cu, cublas.zig, lanes.zig, model.zig). The long design text is now in tools/dsv41-zig/eagle/README.md. - eagle_ref.py (rnd=True) rounds SiLU(g) to bf16 before the multiply and rounds each residual add once, as tf_eagle_silu_mul_kernel and linearDenseAdd / own GEMV mode 2 do. On main3 the step o changes by 0.3-2.1% and all 30 drafts stay the same. cross (rnd=False) is the same. - Eagle.open logs a warning for a head window past the model window (a prompt gives the head only its last sliding_window rows).
The 2026-10-24 window's dump (taps 2, 26, 37, final; cuts 1605, 1659, 1705): tap 2 differs by 0.0195 (bf16 rounding with other prefill chunk shapes; a row-offset error gives about 1), and the taps past layer 20 differ by 0.26-0.84, because rows below each prompt's replay floor run no upper layer. --gate selects the taps that must pass; the limit is 3e-2.
What this changes
Adds an EAGLE-style draft head (the "TFE head") to the DeepSeek-V4.1 Zig engine as a second drafter. It is off by
default;
TF_DS_DRAFTER=eagleon both ranks andTF_DS_EAGLE_HEAD=<file>.tfeon rank 0 turn it on. The head runson rank 0 only and sends its drafts to rank 1 in the verify frame, so a draft pass needs no collective. In eagle mode
no rank loads DSpark. Stacked on
dsv41-zig-dense(the dense-linear commits the serving image is built from).zig/src/families/dsv41/eagle.zig: head file reader, rings, prompt and round absorbs, the chained pass, and aself-test against reference vectors (
TF_DS_EAGLE_CHECK).zig/kernels/cuda/dsv41/eagle.cu: row kernels (RMSNorm, RoPE, window attention, SwiGLU, top-1 with confidence).The GEMMs are bf16 cuBLAS calls with fp32 sums (
cublas.linearDenseAddadded).model.zig,lanes.zig: the drafter switch, tap capture (taps read from the head file: layer inputs afterEngram, and tap 1000 = the final state), the draft pass in the existing rank depth policy.
TF_DS_ACCEPT_LOG=N: drafts asked and kept by window position on stderr, for either drafter.tools/dsv41-zig/eagle/:.ptto.tfeconverter, torch reference and vector writer, kernel test, fatbinextractor, and the live check client (
tfe_bench.py,tfe_seg.py,tfe_sum.py).Update (8d8691c..ce9a427): main3x heads, pass kernels, FP8 linears
qk_norm, a window up to 511 and thetwo_stagevocabulary head (topcand_klow-rank logitsrescored by a dense bf16 copy of the target's own head in the head file;
TF_DS_EAGLE_TOPKup to 512). Theconverter records the copy's source files and sha256.
TF_DS_EAGLE_GRAPH=1: a pass's steps as one CUDA graph (off by default; the gain was inside the noise).TF_DS_EAGLE_GEMV=wodefault: wo and w2;all;cublas), theattention with a thread a key (window 511 in one block), the vocabulary top-1 split over 64 blocks.
TF_DS_EAGLE_FP8=0runs bf16,=1requires FP8): the head's linears in e4m3 with an fp32 scale a row (scale = row absmax /448, round to nearest even, from the bf16 weight); the converter writes both forms. Only the draft head's numerics
change; the target's tokens do not (token_sha below).
Receipt of the update (a build of ce9a427, the same setup as below; head b1 with FP8 forms; nsys medians):
two_stagehead (window511, qk_norm) in bf16 and FP8. token_sha equal to
"draft": falsein every phase, 1 and 4 streams.(bf16), 0.76 ms (FP8; its weights read at the bandwidth roofline).
drafts, so the pilot head, not the pass, is now the limit.
Update (a300c33, 1513312): FP8 by default, 4-tap heads, the tap dump
-0.003 to +0.005 on every tag (the rule: at most 0.01). A head file without FP8 forms runs bf16 as before.
tap_norm(a layer-2 tap): already supported;eagle_rand_head.pymakes a random head of anyshape, and
eagle_ref.py crossagainst the training code passes on one (worst 2.0e-6, drafts equal).TF_DS_EAGLE_TAPDUMP=FILEandtfe_tapcheck.py: a check that a tap below the middle layer is read at theprompt chunk's replay cut. Not yet run on the pair.
Receipt
Environment: this branch (
09276c39, same tree as the measured build) cross-built with Zig 0.17.0 foraarch64-linux, the dense image's kernels plus
dsv41_eagle(nvcc 13.4, sm_121); 2x DGX Spark (GB10), TP=2;DeepSeek-V4.1-Flash-1017, EXL3 2.9 bpw pack; pool = context 262,144, eager decode, 4 streams; head b1 (271.7M
parameters, taps 26, 37, 39).
Exactness:
token_shaagainst"draft": false, 24 groups (prose, code, tool calls x 4 prompts x temperature0 and 1.0 with a seed, 1 and 4 streams): one value in every group across TFE on and off, DSpark on and off, and
the stock image's binaries.
Decode speed (aggregate tok/s of a batch): drafts off / DSpark / TFE b1.
The pilot head b1 is slower than drafts off on 8 of 9 workloads. Its acceptance is lower than DSpark's, and its
pass costs about 14-20 ms (estimate from round accounting) against DSpark's 4.3 ms. This PR adds the drafter
(off by default) for the trained heads that follow. DSpark speed on this build equals the stock image to 0.3%.
Prompt speed: n/a. Eagle mode adds tap copies and one fc GEMM per prompt chunk only when it is on; with it off
the prompt path is unchanged.
Tests run:
zig build test-cuda-host(Linux x86_64 CPU): 43/43 steps, 1184/1186 tests pass, 2 skipped. Engineself-test on the pair: every draft equals the torch reference, step-1 state error 0.4-1.2%. Kernel test against
torch (sm_86): all within bounds.
Not run: Metal and the other CUDA families (the code is DSv4.1-only and off by default). Graph decode with the
TFE head (eager only). Thinking and structured-output workloads.
Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (Drafted ="draft": falseand 4 streams = 1 stream: checked. A resumed prompt with the TFE head: not checked.)SUMMARYline is in the receipt. (Not done:dsv41_eaglefollows the other dsv41 kernels, which are not in the contract either.)