Skip to content

lfm2 metal: masked non-finite values stay out of prompt attention - #2328

Open
ashhart wants to merge 10 commits into
mainfrom
lead/lfm2-masked-prompt
Open

ashhart wants to merge 10 commits into
mainfrom
lead/lfm2-masked-prompt

Conversation

@ashhart

@ashhart ashhart commented Nov 6, 2026

Copy link
Copy Markdown
Owner

Part of #1009. Stacked on #2322, which is stacked on #2308. Only the last commit is this PR's: "fix(lfm2): exclude masked nonfinite values from Metal prompt attention". It lands after #2322.

The fix: LFM2's prompt attention on Metal gets the same fix #2322 gives Nemotron. A non-finite value (Inf or NaN) in a masked key's V no longer reaches the output through a zero probability.

  • A GPU scan over the newly written keys, staged tails included, picks between the original pipeline and an exceptional one, through indirect grids.
  • Finite prompts run the same pipeline handles with the same grid as before.

Checks (Apple M5 Ultra):

  • LFM2's kernel contract, both probability modes: 9 kernels pass, 7 planted controls are caught, and launch limits pass.
  • Contract: masked non-finite values across every Metal prompt attention, the findings registered as known (refs #1009) #2308's census: 82 entries pass, and LFM2's two entries move from known to passing. 41 other-family findings stay known, and both controls are caught.
  • Served, main against this branch: LiquidAI/LFM2.5-24B-A2B-MLX-4bit at 667408d4, warm base new new base, fresh servers, one and four streams, distinct chat and code prompts, greedy and seeded, with a 150 ms stagger. All 200 saved keys are equal across the four measured passes.

Prefill is level. Decode, aggregate tok/s:

Prompt Temperature Streams Main Fix Change
chat 0 1 223.9 227.3 +1.5%
chat 0 4 477.1 475.1 -0.4%
chat 1 1 214.8 217.7 +1.4%
chat 1 4 420.2 421.6 +0.3%
code 0 1 247.8 252.8 +2.0%
code 0 4 498.8 506.2 +1.5%
code 1 1 238.8 239.1 +0.1%
code 1 4 419.2 421.1 +0.4%

These are within run-to-run spread, so no speedup is claimed. tests/contract/metal_main.zig goes from 600 to 547 lines. zig fmt --check passes.

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