Skip to content

nemotron metal: masked non-finite values stay out of prompt attention - #2322

Open
ashhart wants to merge 9 commits into
mainfrom
lead/nemotron-masked-prompt
Open

ashhart wants to merge 9 commits into
mainfrom
lead/nemotron-masked-prompt

Conversation

@ashhart

@ashhart ashhart commented Nov 6, 2026

Copy link
Copy Markdown
Owner

Part of #1009. Stacked on #2308 (the Metal masked non-finite census): the first commit here past #2308's head is "nemotron: isolate exceptional prompt values from causal…". Review the last 7 commits. This lands after #2308.

The bug: Nemotron's prompt attention on Metal could let a non-finite value (Inf or NaN) in a masked key's V reach the output. The probability there is zero, but in the tensor P×V product, 0 × Inf is NaN. #2308's census lists Nemotron's prompt cases as known findings, and this fixes them the way #2106 fixed the others.

The fix:

  • A scan: each prefill runs a small GPU scan over only the newly written keys past the first query's visible boundary.
  • A pick: a one-thread selector then chooses between the original pipeline and an exceptional one, through two mutually exclusive indirect grids, and clears its flag.
  • Finite prompts: they launch the same pipeline handles with the same grid as before. attention_nax.metal, paged.zig, frags.zig and the Metal compile options are byte-identical to main. Single-row decode (L == 1) is untouched.

Checks (Apple M5 Ultra):

  • Nemotron's contract, both probability modes: 61 kernels pass, 0 fail, 34 planted controls are caught and none missed, with 0 limit problems. One inherited expert-up NaN finding remains, outside prompt attention.
  • Contract: masked non-finite values across every Metal prompt attention, the findings registered as known (refs #1009) #2308's census: 82 entries pass, and Nemotron's Inf and NaN entries move from known to passing. 43 other-family findings stay known, and both planted controls are caught.
  • Served, main against this branch: warm base new new base, one and four streams, greedy and seeded, 150 ms stagger. The verdict is GREEN, and all 20 prompt keys are equal across the scored processes.
Prompt Main tok/s Fix tok/s Gap
2,048 9,367.1 9,389.4 +0.24%
8,192 10,232.3 10,257.7 +0.25%
16,384 9,810.2 9,836.1 +0.26%
32,768 9,008.2 9,033.9 +0.28%
65,536 7,707.5 7,727.3 +0.26%

Decode is level across all eight sweep cells. Greedy chat goes from 246.15 to 246.95 tok/s at one stream, and from 241.55 to 242.15 tok/s aggregate at four. Gaps under 3% aren't claimed as speedups.

Lean is unchanged, zig fmt --check passes, and every touched file is at 600 lines or under.

Check masked Inf and NaN separately against binary64, including visible IEEE classes, sparse padding and wrapped windows. Retain the 45 measured inherited failures as known #1009 findings and catch both planted GPU controls. Register the suite in the Metal contract and provide a source-only focused runner. Production kernels and dispatch remain unchanged.
Retain the original finite attention libraries and select a separate correction on the GPU for newly written non-finite values hidden from earlier rows. Restore visible IEEE contributions and cover both cache layouts and probability modes.
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