Repository navigation
Conversation
…d lerp residual Dory's residual stream lives on the unit sphere, so its rows need two kernels norms.cu did not have, both a block a row (256 threads, 16-byte loads, rows to 8,192 in 64s), so a row's bits never depend on the row count: - tf_l2_rows: y = bf16(s / sqrt(sum(s^2) + eps)), s = x, or bf16(x + r) (a loop's entry, justnorm(state + input_h); with r null, a loop's exit). - tf_slerp_rows: dory.py's slerp_norm, h = bf16(l2(a + |alpha scale| (b - a))) with a = bf16(l2(res)) and b = bf16(l2(out)), the lerp in fp32 steps. Norms.l2 and Norms.slerp launch them; no existing kernel or launch changes. Both are registered with the kernel contract. Their check compares every output with float64's value at the bf16 rounding points as a range over the fp32 run's spread (core/l2_ref.zig); three controls each drop one rounding point (the sum before the norm, the slerp's a, its b) and go red at gate 3. Gate 4's integer rows make an fp16 sum overflow to a zero output (85x and 51x the bound); a bf16 sum's scale cancels in the norm, so that margin is stated, not proven, as for the RMS tails. Launch records from nvcc 13.0's sm_80/86/89/120/121 images.
…2 norm and q's scale tf_qk_rows_l2 takes dory.py's order: rotate-half RoPE on the projection's bf16 rows (rounded to bf16), then each head's L2 norm with eps on the sum, then q times an fp32 scale vector ((sqk f)^2, formed on the host as torch forms it) after the normalized row's rounding; k is its normalized row, v as is, and k and v go into the caches as tf_qk_rows writes them. Shape.l2 and Shape.q_scale select it; every other call keeps its kernel. It is registered with the kernel contract: every output against float64's range at its rounding points (cosf's ulps on the fp32 angle pos * inv, the inverse's fp32 spread), at head dim 256, 16 query heads over 4, a third of the rows past position 199,800. Two controls each drop one rounding point (the rotated head, q's normalized head) and go red at gate 3. Gate 4's integer heads (no rotation) overflow an fp16 sum to a zero output. core/l2_ref.zig's intervals now keep infinite ends and widen by fp32's least subnormal step, so a planted infinity leaves only the outputs float64 makes non-finite.
…then L2-normed Two activation rows for Dory, a block a row, 16-byte loads, fp32 scale vectors formed on the host as torch forms them: - tf_gate_l2_rows: out = bf16(l2(bf16(o * sigmoid(g * sg)))) over all heads of a row (dory.py's output gate, then justnorm over num_heads * head_dim); - tf_swiglu_l2_rows: out = bf16(l2(bf16(silu(bf16(g * sug)) * bf16(u * suu)))) (its MLP). sigmoid and silu are torch's fp32 forms, 1 / (1 + expf(-x)) and x / (1 + expf(-x)) with IEEE division. ActRows.gateL2 and ActRows.swigluL2 launch them; no existing launch changes. Both are registered with the kernel contract: every output against float64's range at its rounding points (expf's ulps, fp32's exp range, the inverse's spread) at Dory's widths, 4,096 and 7,168; four controls each drop one rounding point (the gated row, the scaled gate, the scaled up, the product) and go red at gate 3. Gate 4: the gate's rows are halves of integers (sigmoid(0) = 1/2), exact in fp32; SwiGLU's squares are no integers, so its fp32 sums are checked against the shapes gate's bound, and an fp16 sum overflows either way.
Dory's logits are the head's bf16 output times logit_scale / std per vocabulary column (dory.py: logits.float() * (logit_scale.float() / initializer_range)). tf_dense_rows_scaled_* are the same tiles with that epilogue, out = fp32(bf16(x . w[n])) * scale[n] in one fp32 product; DenseRows.rowsScaled launches them. tf_dense_rows_* keep their code: their records from nvcc 13.0 are unchanged. Registered with the kernel contract as replacements: on the same inputs every output is bit-equal to its unscaled tile's fp32 outputs, rounded to bf16 and scaled on the host, at every served shape and edge (4,103 columns, 2,560 inputs, the row counts the wrapper sends each tile); gates 1, 2 and 4 against float64 of the same steps, with a NaN and an infinite scale planted.
…ry (host only) Nemotron-Labs Dory 1.0 is a looped nGPT stack: input sublayers once, the recurrent block n_recurrent_loops times, output sublayers once, attention and MLP sublayers alternating. - config.zig parses config.json and generation_config.json's end tokens: the block counts, sliding/full/mlp kinds, both RoPE thetas, the nGPT constants and their fp32 factors, the loop count (and withLoops, 1 to 4, for a load option). It refuses last_loop KV (chunk- dependent by construction), quantizations other than ModelOpt W4A16_NVFP4 at group 16 or plain bf16, scaled or partial RoPE, tied heads, biases and shapes the kernels do not take. - names.zig maps every tensor of the NVFP4 (or bf16) checkpoint with dtype and shape checks and refuses leftovers; the pinned checkpoint (TF_DORY_DIR) maps its 1,083 tensors. - schedule.zig gives the virtual sublayers for L loops with a per_loop cache per attention step: 8 + 20 L + 8 caches, matching llama.cpp #30306's physical[] and cache_layer[] and vLLM #60149's loop * layers + layer indices. - memory.zig counts weight bytes as stored (equal to the pinned file's tensor data) and KV bytes by loops, context and sliding ring: 36, 56 and 96 KiB a token at 1, 2 and 4 loops.
ref.zig runs Dory's forward in float64 over prompt chunks with a stream's KV caches, following vLLM #60149's dory.py (ce9dd0ed): the nGPT slerp residual, q and k L2-normed after RoPE with q scaled by (sqk f)^2, the scaled sigmoid output gate and the row L2 over all heads, the scaled SwiGLU and its row L2, the loop entry l2(state + input_h) and exit l2(h), and the logit scale. Options give dory.py's bf16 rounding points with torch's fp32 scale vectors, per-step taps, and last_loop KV. ref_test.zig checks, on a tiny random config (2 + 4 + 2 sublayers, window 3): - per_loop: one chunk, several chunkings and one-token decode give the same logits to the bit, at 1 to 4 loops, in float64 and at the bf16 points; - last_loop: equal to per_loop at one loop and within one chunk, but at two loops the rows after a cut differ, the chunk dependence that keeps it out of this family; - every logit within 1e-12 of dory.py itself run in float64 (its fp32 casts as float64) on the same weights at 1, 2 and 4 loops (ref_fixture.zig; the oracle is a dev-time script, not part of the tree).
nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-MXFP8 and -BF16 next to the NVFP4 map. - config.zig: Quant gains mxfp8. quantOf takes compressed-tensors' mxfp8-quantized (8-bit float codes, group 32) besides ModelOpt's W4A16_NVFP4 and plain bf16, and refuses other compressed-tensors groups by name. - names.zig: an MXFP8 projection is its E4M3 codes [n, k] and UE8M0 scales u8 [n, k/32]; the tiny checkpoint test runs all three quants, and the pinned tests (TF_DORY_DIR, TF_DORY_MXFP8_DIR, TF_DORY_BF16_DIR) share one body. - memory.zig: the published MXFP8 and BF16 files' tensor data (4,750,663,680 and 7,950,458,880 bytes) equal weightBytes; parameters are the same for every quant.
…dows bit-equal
cuda_weights.zig loads the ModelOpt W4A16 NVFP4 checkpoint: every projection packed on the
host for the lane matmul with its own weight_scale_2 (q, k, v and g launch separately), the
nGPT scale vectors formed in fp32 as torch forms them ((sqk f)^2, (s_gate f)^2, suv
sqrt(hidden), logit_scale / std), embeddings and head bf16.
cuda_model.zig runs the sublayer schedule over per_loop KV caches (8 + 20 L + 8: full
caches over the context, sliding ones over a ring of a chunk and the window): embed, the
loop entry and exit L2 rows, the projections (prompt GEMM for prompt chunks, the lane for
decode rows), tf_qk_rows_l2 into the cache, window attention at head dim 256 (window 1024,
or none over the linear cache), the gate and SwiGLU L2 rows, the slerp residual, and the
scaled head. Prompt chunks to 2,048 rows, decode windows to 16.
tf-cuda-test dory MODEL LOOPS N TEXT, on nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-NVFP4
(41f02335), GB10:
- greedy text is sane at 1, 2 and 4 loops ("The capital of France is" -> " Paris. The
capital of the United States is Washington, D.C. ...");
- drafted == plain: verify windows of 2, 3, 5, 8 and 16 rows copied from the plain run give
every row's logits bit-equal to the plain one-row decode (35 and 54 windows a run);
- a prompt cut in two chunks gives the last row's logits bit-equal to one chunk's, and the
prompt fed as decode rows gives the same first token.
Eager decode, no graphs: 14.7, 20.8 and 33.7 ms a token at 1, 2 and 4 loops.
7 of 9 tasks
Owner
|
Landed on main in 6528fbe..81b476d: the whole Dory stack (#1256-#1261), with your records under your authorship plus five maintainer children. The children refuse host-only formats and refund failed upload ownership, pin release launch arguments, keep the unmeasured Dory launches pending in the contract, and own queued host inputs through stream completion. Every existing family file, contract entry and recorded launch is byte-identical. Thank you, Jürgen. On a GB10 with nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-NVFP4 @ 41f02335, served at 2 loops with CUDA graphs:
- Synthema, Ash's AI Assistant. |
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.
What this changes
Steps 1, 2 and 4 of the Dory plan (#1233): Nemotron-Labs Dory 1.0 4B as a native CUDA family, one stream, NVFP4 (ModelOpt W4A16). Four commits:
dory: family skeleton, config, weight map, sublayer schedule and memory (host only)dory: the whole forward in float64, per_loop chunk-invariant to the bitdory: MXFP8 and BF16 checkpoint maps (host only)dory: the CUDA forward for one stream on NVFP4, greedy and verify windows bit-equalThe forward is built on the four kernels of #1256-#1259. It is not yet reachable from
tensorfold runor the server; that is #1261.This PR closes nothing. The issue is #1233.
Stacking. Stacked on #1259 (dense_rows), on #1258, #1257, #1256; the #1245 contract include fix is #1256's base commit. Retarget to
mainas the kernel PRs land.What each commit does
config.zigparsesconfig.jsonand the end tokens fromgeneration_config.json, withwithLoops(1 to 4) for a load option. It refuseslast_loopKV, tied heads, biases, scaled or partial RoPE and quantizations other than ModelOpt W4A16_NVFP4 (group 16) or plain bf16.names.zigmaps every tensor with dtype and shape checks and refuses leftovers: the pinned checkpoint maps all 1,083 tensors.schedule.ziggives the virtual sublayers for L loops with a per_loop cache per attention step, 8 + 20 L + 8 caches, matching llama.cpp #30306'sphysical[]/cache_layer[]and vLLM #60149'sloop * layers + layer.memory.zigcounts weight bytes as stored (3,202,376,832, equal to the file's tensor data; 3,975,229,440 parameters) and KV bytes: 36, 56 and 96 KiB a token at 1, 2 and 4 loops.ref.zigruns the whole forward in float64 over prompt chunks with a stream's KV, following vLLM #60149'sdory.py(ce9dd0ed). Options give dory.py's bf16 rounding points, per-step taps andlast_loopKV.weightBytes.cuda_weights.zigpacks the projections for the lane matmul, with q, k, v and g launching separately, and forms the nGPT scale vectors in fp32 as torch does.cuda_model.zigruns the sublayer schedule over per_loop caches (full caches over the context, sliding ones a ring of a chunk plus the window): embed, loop entry and exit L2 rows, projections (prompt GEMM for chunks, the lane for decode rows),tf_qk_rows_l2, window attention at head dim 256 (window 1024, or none), the gate and SwiGLU L2 rows, the slerp residual, the scaled head. Prompt chunks to 2,048 rows, decode windows to 16.Results
f64 reference against vLLM. Run on a tiny config (2 + 4 + 2 sublayers, window 3), weights from splitmix64 codes (every value exact in bf16). The Zig reference matches
dory.pyfrom vLLM #60149 @ ce9dd0ed, with its.float()casts turned into.double(), within 1e-12 relative on every logit, at 1, 2 and 4 loops. A 1e-10 relative change of eps fails it. The oracle is a dev-time script, not in the tree.Chunk invariance. In per_loop, the chunkings {12}, {5,4,3}, {1x12} and {3,9} give bit-identical logits at 1-4 loops, in f64 and at the bf16 points.
last_loopequals per_loop at L=1 and within one chunk, and differs after a cut at L=2: the chunk dependence that keeps it out of this family.CUDA, pinned checkpoint (
nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-NVFP4@ 41f02335, GB10):tf-cuda-test dory MODEL LOOPS N TEXT: greedy text is sane at L=1, 2, 4 ("The capital of France is" -> " Paris. The capital of the United States is Washington, D.C. ...").Against vLLM #60149 on the same GB10, same NVFP4 checkpoint. vLLM eager,
TRITON_F32_DEFAULT=ieee,n_recurrent_loopsset by override. Three prompts ("The capital of France is", "def fibonacci(n):", "Water boils at a temperature of"), 48 greedy tokens each, top-5 logprobs. "Top-1" and "max |d logprob|" are over the positions with identical context (up to and including the first divergence).Six of nine runs agree on all 48 tokens. The three divergences are near ties at the diverging position:
On vLLM's own chosen token, the logprob differs by at most 0.125 (mean 0.006-0.014) in every run. The per-sublayer hidden state against vLLM's differs by a relative L2 of 0.01 on average (max 0.016 at the last call), which is what bf16 steps of different order give; the first sublayers are at 0.004. We read this as bf16 rounding, not a difference in the model, but do not claim more than the table shows: the other pins, llama.cpp #30306 and mlx-lm #1988, were not run.
Not in this PR
Serving (#1261), CUDA graphs, a lanes backend, snapshots, paged KV, MXFP8/BF16 on the GPU.
Receipt
b1a7f2386(the top of this PR); NVIDIA GB10 (sm_121), CUDA 13.0 nvcc, Zig 0.17.0.zig build test-dory-hostpasses; withTF_DORY_DIR, the pinned checkpoint maps all 1,083 tensors with none left over.zig build test -Dnvcc=...: exit 0 at this branch head (zig build test -Dnvcc=/usr/local/cuda/bin/nvcc, the full step).main.Choices ashhart can change
per_looponly;last_loopis refused at config parse because it is chunk-dependent by construction and cannot pass drafted == plain.--loopsas a load option, hereConfig.withLoops; a per-request loop count would cost a graph set per L.kv_poollater.