Skip to content

metal families: refuse a token id at or past the embedding table before the lookup - #3045

Draft
ashhart wants to merge 2 commits into
mainfrom
security/token-ids-metal
Draft

ashhart wants to merge 2 commits into
mainfrom
security/token-ids-metal

Conversation

@ashhart

@ashhart ashhart commented Nov 12, 2026

Copy link
Copy Markdown
Owner

What this changes

Before: on these paths a token id reached the embedding lookup unchecked, and no embedding kernel bounds its ids, so an id at or past the table read device memory past it. After: each path refuses such an id on the host before the ids are uploaded, and the loaders that took the embedding table as stored refuse one whose rows are not the vocabulary's. Metal paths: nemotron, gemma4, glm, glm_dsa, flashnext, the minimax_m3 drafters and step5's draft head.

  • Issue: n/a. SECURITY.md asks for private reports, so there is no public issue.
  • Stacked on: main.
  • atc overlap --json: not run.

Checks on this head

  • Head commit: e8722866bb916997142850ae0fe4f2d43a2f22ab
  • Lean checker, python3 tools/zig/lean_check.py, on the base and on this head: 345 on the base, 345 on this head, 0 new.
  • Tests run, with the command and pass, fail and skip counts:
    • zig fmt --check on 24 files: passes
    • git diff --check: passes
    • zig build native: 66/66 steps succeeded
    • zig build check-native-cuda-host compile-flashnext-mtp-parity-host compile-glm-native-open: 11/11 steps succeeded
    • zig build -Dhip=true check-native-hip-host test-native-hip-host test-environment-hip-host test-hip-family-offers-host test-qwen3-hip-host: 19/19 steps succeeded; 29/29 tests passed
    • zig build xpu-native with stand-ins for Intel's ocloc and nGEN headers: 47/47 steps succeeded
    • zig build test-step5-host test-lfm3-host test-lfm3-prefill-host test-minimax-info-host test-minimax-cuda-host test-xpu-refusal-host test-draft-xpu-host test-qwen27-config-xpu-host test-nemotron-round-depth-host test-nemotron-prefill-widths-host test-nemotron-mtp-policy-host test-nemotron-restored-head-host test-nemotron-subset-head-host test-nemotron-grants-host test-nemotron-copies-host test-gemma4-load-parts-host test-gemma4-owners-host test-gemma4-contract-host test-gemma4-prompt-head-host test-gemma4-tower-host test-glm-window-contract-host test-glm-contract-host test-glm-formats-host test-glm-copies-host test-flashnext-host test-flashnext-batch-host test-flashnext-mtp-plan-host test-flashnext-driver-host test-flashnext-single-session-host test-flashnext-scheduling-host test-flashnext-constraint-host test-flashnext-prefill-pairs-host test-flashnext-contract-host test-qwen35-rounds-host test-qwen35-verify-host test-qwen35-prefill-host test-qwen35-draft-frames-host: 194/194 steps succeeded; 977/981 tests passed (4 skipped)
    • zig build test-golden test-server-security test-pull-security: 76/76 steps succeeded; 31/31 tests passed
  • CI: runs on this push.
  • Not run, and why: No new unit test. The checks are covered by the compiles and by reading. zig build test and every step that runs on a GPU were not run: this Mac's GPU was serving a model, and no CUDA, ROCm or Intel hardware was available.

Receipt

  • Main against branch: could not be run. It needs a Metal device, and none was run.
  • Boxes it reaches: Metal serving of the families listed above.
  • New switch? Default, accepted values, profile row: none.
  • Environment, with the TensorFold commit, Zig and backend toolchain, machine and GPU, checkpoint and revision: this head on Zig 0.17.0, on an Apple silicon Mac, CPU-only steps. No GPU and no checkpoint was used.
  • Paths and platforms the change reaches, which I ran and which I could not: all four servers compile with this change. The Metal server was compiled and not started.
  • Exactness: not measured. The change adds host comparisons before an upload and changes no kernel and no arithmetic, so a valid prompt takes the same path as on main.
  • Main against this branch, same machine and checkpoint: not run, for the reason above.
  • Decode speed before and after: not measured.
  • Prompt speed before and after: not measured.
  • Licence of any new weights or drafter: n/a.

Checklist

  • The title and description match the head commit, including the receipt and the stack.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. This branch adds no test.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths in the diff or its commits. No AI attribution lines.

For a change that reaches inference:

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. Not measured: no served run.
  • No precision traded for speed.
  • Decode and prompt processing are no slower than main. Not measured: no served run.

Notes for the reviewer

  • A request that carries such an id now fails with an error where it used to read past the table.
  • A checkpoint whose embedding table does not match its configured vocabulary is refused at load. Most of these checks want an exact match, so a table with more rows than the vocabulary is refused too.
  • Not changed here: the training paths, the dev and timing tools, and output heads that are loaded unshaped.
  • On a pair of Macs, rank 0 checks a prompt before the other Mac hears of it, so a bad id no longer makes the follower leave.

Suggested changelog line

Metal: nemotron, gemma4, glm, glm_dsa, flashnext, minimax_m3's drafters and step5's draft head refuse a token id at or past the embedding table before the lookup.

…re the lookup

The embedding kernels read the table at the id they are given and none of them bounds it, so the host check
is the only one. Each path now refuses an id at or past the vocabulary before it uploads the ids, and the loaders
that took the table as stored refuse one whose rows are not the vocabulary's.
…ssistant and the tree feed

gemma_assist_input clamps an id past the table to its last row, so the Metal assistant drafted from the wrong row
for a pending id at or past the vocabulary; its hold now refuses one before any stream's work, as the DFlash
drafter does. FZ_FEED's tree guesses become window ids and are refused when the feed is read.
test-token-ids-metal-host runs both host tests on the shader-free engine module.
@ashhart

ashhart commented Nov 12, 2026

Copy link
Copy Markdown
Owner Author

Review fixes pushed (7be24cf):

  • gemma4/assistant.zig: the Metal assistant drafter fed its pending id to the target's embedding with no bound. The kernel clamps it to the table's last row, so a bad id drafted from the wrong row rather than reading out of bounds. hold now refuses an id at or past the vocabulary before any stream's work, with a host test.
  • flashnext/tree.zig: guesses from the development feed file become window ids, so the parser now refuses one past the vocabulary, with a test.
  • New step test-token-ids-metal-host, in the macOS host CI list.

The other Metal drafter and MTP paths were already guarded: Nemotron, Flash Next, Qwen 27B DFlash, MiniMax DFlash2 and DSpark, Gemma 4 DFlash, Step-5 MTP and GLM DSA. Host steps: 10/10 with 111/111 tests, plus 19/19 compile steps.

Overlap with #3054: both PRs edit nemotron/backend.zig. #3054 moves the prompt pass into prefill_metal.zig, so taking its side of the conflict drops this PR's Nemotron prompt check without any error. Whichever lands second should put the check in Ops.open (prefill_metal.zig, right after limits.check). There it still runs once, before the cache opens or a slot is taken, for a whole pass and for a chunk a call alike.

Next: an M5 run, main against this branch, on Flash Next, Nemotron with MTP and Gemma 4.

  • Synthena (Ash's Agent)

@ashhart

ashhart commented Nov 14, 2026

Copy link
Copy Markdown
Owner Author

M5 Ultra device run, main against this head, Flash Next: REVIEW. Sampled code prompts are slower: 200.2 to 189.3 tok/s at 1 stream (-5.4%) and 231.6 to 222.6 at 2 streams (-3.9%), and the two new runs sit below both base runs. So this does not land as is. The Gemma 4 and Nemotron runs are still to come, and we will look at what costs the sampled Flash Next rounds before going further.

  • Synthena (Ash's Agent)

@ashhart

ashhart commented Nov 14, 2026

Copy link
Copy Markdown
Owner Author

Second M5 Ultra A/B confirms the loss and widens it: sampled chat 136.9 to 129.0 tok/s at 1 stream (-5.8%) and 196.2 to 189.2 at 2 (-3.5%), and greedy code 257.1 to 246.8 at 2 streams (-4.0%), with both new runs below both base runs in each. The new checks run on every shared verify round, but ids the engine produces are already in range, so we are moving the refusal to where ids enter from outside (prompts, prompt_token_ids, logit_bias, stop ids, recalled state), and then re-measuring.

  • Synthena (Ash's Agent)

This branch has not been deployed

No deployments
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