Repository navigation
Nemotron 3.5 Lightning CUDA prefill: one row departs from float64 after a near-tie routing pick in layer 6 #1489
Description
On main e9b69ed70 (GB10, sm_121, CUDA 13.0), one row of a 384-token CUDA prefill is far from a float64 reference from layer 7 on. At the input of layer 18 (the norm output that feeds its Mamba in_proj), column 2206 reads 52.5 where float64 gives -1.18. The row's RMS is 0.63.
How it was found
The float64 comparison for #1353 and #1354 (#548). The reference runs every op in float64 on the CPU from the MLX weights dequantized exactly (q * scale + bias); Mamba dt is not clamped. It saves every layer's norm output. The CUDA side is tensorfold prefill's --dump. The worst column per projection class was 4.42 (in units of the layer's RMS) on main and 0.45 on a placebo build that reverses one fold in inv_rms_v. All of the difference comes from this row.
- Model:
nemotron-3.5-lightning-mlx4(MLX 4-bit). - Prompt: token 1 (BOS), then the first 383 tokens of
CONTRIBUTING.md(identical at2e33e297dande9b69ed70), tokenized with the model's tokenizer, no special tokens. - Run:
tensorfold prefill MODEL pf.json pass --dump DIR, where pf.json is{"pass": [those 384 ids]}. One prefill of 384 rows. - Position: row 162 (0-based), token 2741 (
ribut, inside[README workflow](README.md#contributing-with-). - Value:
018_y.bin, row 162, column 2206.
What happens at that row
Layer 6 is a MoE layer. For row 162, the 6th routed pick is a near tie:
| 6th pick | sigmoid + bias: expert 38 | expert 46 | |
|---|---|---|---|
| float64, on float64's layer-6 input | 38 | 43.197611 | 43.197528 |
| main's layer-6 input, scored in float64 | 46 | 43.197337 | 43.197661 |
| main's GPU picks | 46 | ||
| placebo's GPU picks | 38 |
The float64 margin is 8.3e-5, about 22 fp32 ulps at 43. Main's layer-6 input for this row is no worse than other rows': relative L2 error 0.0117, against a median of 0.0115 over all 384 rows. Its error is just enough to flip the pick.
In float64, swapping expert 38 for 46 (slot weight 0.316) moves the layer-6 output by 0.525 (L2). On the GPU, main and the placebo's residuals after layer 6 differ by 0.502, against a residual norm of 4.2. The swap accounts for the whole departure.
After that the row's error keeps growing, well past the other rows:
| layer input | 6 | 7 | 9 | 11 | 14 | 18 | 30 | 40 | final norm |
|---|---|---|---|---|---|---|---|---|---|
| main, row 162 rel-L2 | 0.0117 | 0.112 | 0.282 | 1.12 | 1.44 | 1.89 | 1.26 | 0.64 | 0.32 |
| placebo, row 162 | 0.0127 | 0.010 | 0.013 | 0.017 | 0.069 | 0.112 | 0.136 | 0.126 | 0.089 |
| main, median row | 0.0115 | 0.0081 | 0.0099 | 0.0116 | 0.022 | 0.031 |
The decode path doesn't show it. Teacher-forced one-row steps over the same tokens put row 162's layer-18 input at 0.11 rel-L2 on both main and the placebo, column 2206 at -1.16 against -1.18.
Open questions:
- Should one routing pick that flips at a near tie move a row this far (1.9 rel-L2 by layer 18)?
- Or is something after layer 6 amplifying it for this row?
- The
klrun over the same passage (64-token prompt, then teacher-forced decode) puts this position's logits at KL 0.0003 against float64, so the outlier stays in the prefill path.
Files (reference arrays and GPU dumps for main and the placebo) are available on request.
Activity
jschmied commentedon Oct 27, 2026
It also reproduces on a newer main,
b30d5b9d3(GB10, CUDA 13.0), with the same command and prompt. The layer-18 input dump (018_y.bin) is byte-identical toe9b69ed70's:ashhart commentedon Oct 28, 2026
Thanks @jschmied, this is exactly what we hoped for, and the placebo makes it easy to read.
The pick itself is a near tie (8.3e-5, about 22 ulps at 43). Main's layer-6 input for that row is as accurate as any other row's, so the flip alone isn't a defect. What we can't tell yet is where the growth after it comes from: the row reaches 1.9 rel-L2 by layer 18, while the median row stays near 0.03. Two checks would separate the model from the engine:
Could you attach main's dumps for row 162 from layer 6 to layer 18, with the float64 arrays? Nothing in 1.1.8 changes Nemotron's CUDA prefill, so they still apply.
jschmied commentedon Oct 28, 2026
I ran both checks. The model amplifies the swap. The CUDA prefill is right for the pick it made, so this is not a bug.
All runs on main
d0c9dc9fb(1.1.8 + 1, GB10, CUDA 13.0,zig build). Its prefill and teacher dumps are byte-identical toe9b69ed70's, so the earlier numbers still apply.1. Float64 with row 162's sixth pick at layer 6 forced to expert 46. The float64 run also reaches 52.68 at column 2206 of layer 18's input, against 52.5 on the GPU and -1.18 unforced. Row 162's relative L2 error at each layer's input:
Against the forced reference the prefill stays between 0.009 and 0.043. The median row against the unforced reference was 0.008–0.031 over the same layers. The swap changes the routing after it. In float64, the forced run picks different experts for row 162 in every later MoE layer: layers 8, 10, 13, 15 and 17 each change 1 to 6 of their 6 picks. The prefill picks the forced run's experts in all of them except one near tie at layer 15: expert 84 over 4, a float64 margin of 2.1e-4 at 35.9. So expert 38 → 46 at layer 6 sends the row down a different route, and the prefill follows that route as closely as any other row follows its own.
2. Teacher-forced decode at row 162, layer 6 picks expert 38, the same as float64 (23, 20, 81, 125, 103, 38). Scored in float64, the decode's layer-6 input puts 38 ahead of 46 by 2.1e-3, against 8.3e-5 in float64 and -3.2e-4 for the prefill's input. So the decode never takes the swap, and its 0.11 at layer 18 is not "46 and still near 0.11". Decode picks were taken with two dump-only lines (main doesn't dump them on the decode path); the decode's dumped activations are byte-identical to main's. The decode's own 0.11 at layers 16–18 comes with later pick differences on the 38 route. It matches float64's picks through layer 10, then differs in one pick each at layers 13, 15 and 17. I didn't force those to test them.
Files (row 162 only, ~0.7 MB): https://gist.github.com/jschmied/02d317a6f6553642fffd05c38ffedca4.
r162.npzholds main's prefill and decode inputs for layers 6–18 (bf16 widened to float32), the float64 arrays for the normal and the forced run (float64), and the routed picks of all four at the MoE layers up to 17. The README gives shapes, dtypes and the exact commands, andref64f.pyis the reference with the forcing option.Should I close the issue as expected behavior? One routing near tie (22 ulps) moves this row to a different set of experts for the rest of the stack.
ashhart commentedon Oct 28, 2026
Thanks @jschmied, that settles it, and it's a model behaviour rather than an engine bug. Forcing the sixth pick to 46 in float64 reaches 52.7 at column 2206 of layer 18's input, where the GPU reads 52.5. Against that forced reference the prefill stays between 0.009 and 0.043 at every layer, as close as the median row follows its own route. One 22-ulp tie at layer 6 sends row 162 to a different set of experts for the rest of the stack, and the prefill follows that route correctly. The decode never takes the swap (38 leads by 2.1e-3 on its input), which explains why it stays near float64.
Closing as expected behaviour. Thanks for the forced reference and the gist; they make this easy to check later.