Skip to content

Unlimited-OCR-2 on native Metal: the publisher's BF16 release, page images to structured text - #2074

Open
BobClawblaw wants to merge 24 commits into
ashhart:mainfrom
BobClawblaw:unlimited-ocr-pr
Open

BobClawblaw wants to merge 24 commits into
ashhart:mainfrom
BobClawblaw:unlimited-ocr-pr

Conversation

@BobClawblaw

@BobClawblaw BobClawblaw commented Nov 3, 2026 •

Copy link
Copy Markdown
Contributor

Fixes #2073.

This serves Baidu's BF16 baidu/Unlimited-OCR-2 (revision 0480966938291445591f69f08a13b041c073a96a) natively on Metal, as published: the engine reads the release's safetensors directly, with no converted pack and no Python or libtorch in the served engine. The image tower (SAM ViT-B, its neck and CLIP-L), the decoder (page2d rotary, DASH block-sparse attention, a 64-expert MoE) and the publisher's page-guarded greedy decoding run in Zig and Metal. Its caches stay in the model's precision: bf16 keys and values, fp32 router keys, as in the publisher's code.

Before (main) After (this PR)
baidu/Unlimited-OCR-2 (BF16) on Metal not supported served by tensorfold-native serve --vision on M5 (tensor units)
Transcribed text vs the release's own infer (torch bf16, CPU) on 6 cases (5 pages alone, all 5 as one document) — identical in 6/6; 10 of 288 layout-box coordinates differ, mostly by 1 (at most 10)
Repeats and two requests at once vs a lone request — identical replies
serve --vision MODEL — the family's chat template built in (the release ships none); a checkpoint's own template or --chat-template takes its place
--parallel above 1 — refused: "Unlimited-OCR-2 serves one stream: serve it without --parallel"
Prompt cut into chunks vs whole — bit-identical logits and replies (tf-uocr-decode --chunk, odd sizes)
Kernel contract (Metal) — 34 Unlimited-OCR-2 launches: SUMMARY 34 kernels: 34 pass gates 1-4, 0 fail, 0 known findings; 31 controls: 31 red as planted, 0 missed; 0 limit problems
Lean checker 425 problems 425 (0 new)
Existing models — unchanged: no shared kernel or path changes (see below)

Speed (M5 Ultra, another model serving on the same GPU throughout)

Pages Prompt tokens Reply tokens First token Decode
report-1 909 198 1.16 s 247.8 tok/s
report-2 909 141 0.82 s 255.0 tok/s
paper-1 909 543 0.82 s 222.1 tok/s
paper-3 909 899 0.91 s 241.0 tok/s
paper-10 909 1,959 0.82 s 241.5 tok/s
all five as one document 4,525 3,764 3.93 s 215.5 tok/s

Commits

  1. Image encoder: SAM ViT-B, its neck and CLIP-L on the tensor-unit GEMM, each step rounded where the publisher's torch code rounds; the weights read as stored (the release's second shard is 8 bytes off 16-byte alignment, so those tensors are copied 256-aligned).
  2. Preprocessing, bit-exact against the publisher's processor: Pillow's letterbox to the 1024 global view, the dynamic tile grid, torch's antialiased bicubic for the 640 tiles' tables.
  3. Decoder and greedy generation: page2d rotary, DASH (block summary keys, page-aware top-16 routing, entmax-1.5, 12 kept blocks, a 128-key local window), the MoE, PageCountGuard.
  4. Decode without host syncs: experts picked on the GPU through a pointer table, matvecs for one token.
  5. Attention across eight simdgroups: one GPU wait a token.
  6. Chunked prefill, bit-identical at any cut: prompt rows always take the GEMM.
  7. Serving through the lane host: one stream, greedy only, page images in chat requests, the family's chat template.
  8. Kernels in the family folder (zig/kernels/metal/uocr/) and the contract's scaffold, then five contract groups (rows, moves, attention, route, matvec), each kernel at its served shapes and edges with planted controls.
  9. A reply may fill the context: the capacity check counts the prompt and every reply token fed back.
  10. The reference script tools/zig/uocr_reference.py (pages, capture: the release's own infer, transformers 4.52.1 as the model card pins) and the recipe docs/recipes/unlimited-ocr-metal.md.
  11. The lean rule: contract groups split to at most 600 lines, one-line comments.
  12. --parallel: a named --parallel above 1 is refused by name instead of serving one stream.
  13. Other GPUs: without tensor units the refusal reads "Unlimited-OCR-2 needs an M5-class GPU (tensor units)", still before any weights are read.
  14. The chat template built in: used when the checkpoint ships none, below a checkpoint's own template and --chat-template.
  15. CHANGELOG and recipe: the line names the box coordinates that differ and M5 GPUs; the recipe's speed table is from this head.

Shared code

Outside the family's folders the PR only registers it: native/metal.zig (the unlimited-ocr dispatch and its name in two refusal lists), cli/hub.zig (title), kernels/metal/sources.zig (two embedded sources), tensorfold.zig (export and test), build/metal_targets.zig (test-uocr, tf-uocr-decode, tf-uocr-vision), and the contract wiring (registry.zig's all_metal, metal_main.zig, metal_limits.zig), and the built-in template: server/hf_text.zig uses families/uocr/chat.zig's template when the checkpoint ships none and its model_type is unlimited-ocr (a uocr_chat module added in build/server.zig and build/dsv41_draft.zig). No shared kernel or another model's path changes, so there is no all-family A/B.

Checks run on this branch

This head, 6f2582554, is 86857ca4b rebased onto main ac51580b2 (only CHANGELOG.md needed resolving, both lines kept) plus the THIRD_PARTY_NOTICES entry for the release's Python files. On it, Linux: the lean checker reads 425 problems, as main does; #2077 stacked on it serves the RTX 5090's seven ocr-match replies byte-identical to before the rebase. The macOS checks that follow ran on 86857ca4b, rebased on main b331731e6, and are not yet re-run on this head: zig fmt --check, zig build test, test-kernel-contract, test-golden, test-server-security test-pull-security, fuzz-cortheon, check-release-paths and dev-tools all pass on M5 Ultra, macOS, Zig 0.17.0, and tf-kernel-contract --only uocr gives the SUMMARY above. test-native passes except one test, draft sources: Metal Qwen27 copies and DFlash2 use actual slot and head facts, which crashes the same way on main b331731e6 without this branch (fixed in #2122). The recipe's speed table is from this head; its replies are byte-identical to the previous head's on all six cases. The two new tests (the --parallel refusal in test-native, the template's order in zig build test) fail when their expected text is changed. The served replies at this head, with no --chat-template, are byte-identical to the earlier head's with it, on all six cases. The contract's gate 5 records are the kernels' max_total_threads_per_threadgroup reservations; M1 and M2 builds were checked against them (--offline, or its procedure on the group's functions) when each group was registered, and no kernel has changed since. A full --offline run reports 7 limit problems in other families' records (dsv41, the MiniMax draft, Ling routing), none of them Unlimited-OCR-2's.

Known limits

  • One stream, greedy. Sampling, the publisher's optional n-gram repetition filter, a prompt cache and drafts are not offered; the model has no draft head.
  • JPEG pages need libjpeg-turbo 3.1.4.1, the repository's pinned decoder; PNG and WebP need nothing extra.
  • A request takes 4 pages by default; --vision-max-images raises that to 64. Each page takes gundam mode when its tokens fit, else base mode.
  • Layer by layer against torch bf16, one base-mode page keeps one row (row 163) outside the strict worst-row gate in layers 8-11, from a near-tie in an early MoE gate; cosine agreement passes everywhere and the reply text matches.

@ashhart

ashhart commented Nov 4, 2026

Copy link
Copy Markdown
Owner

Thanks @BobClawblaw, this is a careful port. A new family that reads the release as published, with the reference
script and the recipe beside it, is how we'd want it done. Outside its own folders it only registers itself:

  • the config check pins every geometry key;
  • the loader accounts for every tensor;
  • the 63 rotary frequencies match torch's float32 values bit for bit;
  • the 34 launches are in the contract with planted controls.

Three small things before it lands:

  1. --parallel. With --parallel 4, it quietly serves one stream. Other engines refuse a named --parallel
    they can't serve instead of serving fewer. Please refuse anything above 1 by name, for example "Unlimited-OCR-2
    serves one stream: serve it without --parallel".
  2. The CHANGELOG line. "Transcribed text identical" is true of the text, but 10 of the 288 box coordinates differ
    (mostly by 1, at most 10), and those are reply tokens. Please say that in the line, and add that it runs on M5
    GPUs.
  3. The refusal on other GPUs. On an M3 it reads "cannot load this checkpoint (NoTensorUnits)". It's good that it
    refuses before reading any weights. Could the message say what's needed, for example "Unlimited-OCR-2 needs an
    M5-class GPU (tensor units)"?

On your question: yes, we'd build the template in. Use the family's own when the checkpoint ships none, and keep --chat-template as an override. That way serve --vision MODEL just works.

Once these are in, we'll run the recipe's six cases on an M5 Ultra against your reference transcripts.

BobClawblaw added a commit to BobClawblaw/TensorFold that referenced this pull request Nov 4, 2026
@BobClawblaw

Copy link
Copy Markdown
Contributor Author

Thanks. All three are in, rebased on main df7ba29a1 (head 3fb5d5bb9):

  1. --parallel: a named --parallel above 1 now refuses with "Unlimited-OCR-2 serves one stream: serve it without --parallel", before the checkpoint is read; --parallel 1 and no flag serve as before. A test in test-native covers both.
  2. CHANGELOG: the line now says it runs on M5 GPUs and that 10 of the 288 box coordinates in the replies differ (7 by 1, at most 10).
  3. Other GPUs: NoTensorUnits now reads "Unlimited-OCR-2 needs an M5-class GPU (tensor units)", still before any weights are read. I have no pre-M5 Mac, so I could not see it on one.

The template is built in as you suggested: when the checkpoint ships none and its model_type is unlimited-ocr, the server uses families/uocr/chat_template.jinja; a checkpoint's own template or --chat-template still wins (tested in server/uocr_chat_test.zig). serve --vision MODEL needs no flag, and on the recipe's six cases the replies without the flag are byte-identical to the earlier head's with it. The recipe and the reference script no longer mention the flag.

On 02f6f41dc (the head adds only the recipe table measured on it): the CI set (fmt, test, test-native, golden, the contract tests, security, fuzz, release paths, dev tools) passes, tf-kernel-contract --only uocr gives SUMMARY 34 kernels: 34 pass gates 1-4, 0 fail, 0 known findings; 31 controls: 31 red as planted, 0 missed; 0 limit problems, and the lean checker reads 425, the same as main, 0 new. The served run: text identical to the reference on 6/6, box coordinates moved 10/288, decode 202-267 tok/s; the description's tables are updated to it.

@shuuul

shuuul commented Nov 4, 2026

Copy link
Copy Markdown
Contributor

If CUDA is not in this PR's scope, I have a measured CUDA implementation of the same family (DASH tile/step, router, MoE, SAM/CLIP tower) and can stack the backend here once the Metal family lands. Is that a welcome follow-up, and would you prefer it as a separate PR against #2074's branch or against main after merge?

@ashhart

ashhart commented Nov 4, 2026

Copy link
Copy Markdown
Owner

Thanks @BobClawblaw. All three are in: --parallel above 1 is refused by name, the CHANGELOG line says which box coordinates differ and that it runs on M5 GPUs, and the refusal on other GPUs says what's needed. The built-in template is what we hoped for, too. One thing is left: main moved this morning, and the branch now conflicts in two one-line lists, the Mac-only family list in native/metal.zig (main added Qwen3.5-MoE beside Llama) and all_metal in the contract registry. Could you rebase? Then we'll run your six cases on an M5 Ultra against your reference transcripts, and check the tensor-units refusal on an M3 Ultra, since you couldn't.

@shuuul, thanks for the CUDA offer, and yes, it's welcome. Please make it its own PR against main once this has landed, with the gate agreed on #2073 first. CUDA is something we can run on our side.

BobClawblaw added a commit to BobClawblaw/TensorFold that referenced this pull request Nov 4, 2026
@BobClawblaw

Copy link
Copy Markdown
Contributor Author

Thanks @ashhart. Rebased onto main b331731e6; the head is 86857ca4b. The two lists now keep main's entries and add ours: Unlimited-OCR-2 after Llama in both refusal chains in native/metal.zig (beside Qwen3.5-MoE), and uocr_records after affine_gather_records in all_metal. The CHANGELOG line stays at the top of Unreleased. git range-diff shows only that context changing; the commits' own lines are the same as on 3fb5d5bb9.

On this head: zig fmt --check, zig build test, test-kernel-contract, test-golden, the security tests (24/24), fuzz-cortheon, check-release-paths and dev-tools pass; tf-kernel-contract --only uocr gives SUMMARY 34 kernels: 34 pass gates 1-4, 0 fail, 0 known findings; 31 controls: 31 red as planted, 0 missed; 0 limit problems; lean is 425 against main's 425, 0 new. On the six cases, the replies are byte-identical to 3fb5d5bb9's. The text is identical to the reference in 6/6, with the same 10 of 288 box coordinates moved. I updated the recipe's speed table from this run (decode 216-255 tok/s; another model was serving on the GPU throughout).

One thing outside this PR: test-native crashes on main itself, in draft sources: Metal Qwen27 copies and DFlash2 use actual slot and head facts. That came from our #1125: Metal.facts() now reads runner.model.config.moe(), and three tests leave the runner's model undefined. test-native isn't in CI, so it didn't show there. The fix is #2122 (tests only). Everything else in test-native passes on this branch, including the --parallel refusal test.

BobClawblaw added a commit to BobClawblaw/TensorFold that referenced this pull request Nov 4, 2026
@ashhart

ashhart commented Nov 4, 2026

Copy link
Copy Markdown
Owner

Thanks @BobClawblaw, and for keeping both lists' main entries in the rebase. This is ready on our side once our runs are in: your six cases on an M5 Ultra against your reference transcripts, and the tensor-units refusal on an M3 Ultra. One small thing we'll handle: the recipe's speed row was measured with another model on the GPU, so we'll put our M5 Ultra numbers there instead. Thanks also for tracking down the test-native crash to #1125.

@ashhart

ashhart commented Nov 5, 2026

Copy link
Copy Markdown
Owner

Thanks @BobClawblaw. Main moved again: the branch now conflicts in server/hf_text.zig. Could you rebase? Our M5 Ultra and M3 Ultra runs follow on that head.

SAM ViT-B, its neck and CLIP-L to the projector's rows, for the 1024
global view and 640 tiles (position tables resized as torch does,
bit-identical). tf-uocr-vision gates each stage against float32 on a
reference dump's pixels.
…ery reply token fed back

The server clamps a reply to the room the prompt leaves; the family refused that reply by one row, so any request
without max_tokens failed once --max-tokens reached the context.
…ne-line comments, the records' duplicate test dropped
… follows, with their licences

modeling_unlimitedocr.py, deepencoder.py, dash_attention.py and modeling_deepseekv2.py (the release's MIT,
Copyright (c) 2026 Baidu) and _sparse_base_impl.py (Apache-2.0, DeepSeek-AI and The HuggingFace Inc. team), at
revision 0480966938291445591f69f08a13b041c073a96a.
@BobClawblaw

Copy link
Copy Markdown
Contributor Author

Thanks @ashhart. Rebased onto main a6e25faea; the head is 0af3eb8c8. The 24 commits are unchanged.

server/hf_text.zig keeps both changes. Kimi's native_kimi path returns before the tokenizer and template are read, so a Kimi checkpoint never reaches our built-in template. Ours applies only after readTemplates, and only when the checkpoint ships no template and its model_type is unlimited-ocr.

Checked on an M5 Ultra:

  • On a6e25faea: fmt, test and test-uocr pass.
  • On the rebase onto 4a20028aa (the same patches): the full CI set passes (fmt, lean, test, test-kernel-contract, test-golden, test-server-security/test-pull-security, fuzz-cortheon, check-release-paths, dev-tools).
  • The served build passes the recipe's six cases. The text without boxes is identical to the torch bf16 references in 6/6 cases, and 10/288 box coordinates moved (at most 10), as before. The replies are byte-identical to the previous head's run.

One thing outside this PR: zig build test-native fails one test on main itself (slide admission: Metal entry points refuse before checkpoint or device access). For kimi_k3, the test expects --slide: kimi_k3 checkpoints do not learn on Macs yet, but open now returns Kimi: learned states and on-disk prompt states are not implemented. It fails the same way on this branch. CI doesn't run that step.

@ashhart

ashhart commented Nov 5, 2026

Copy link
Copy Markdown
Owner

Thanks @BobClawblaw. To run your six cases on our M5 Ultra we need the inputs your run used: the five page images and six reference captures, with each file's size and SHA-256. They're what the text and box-coordinate comparisons read, and the PR supplies a generator rather than the files. Could you attach them to a release on your fork? We won't run the generator here. We're also running this head against main on the 27B, for the shared code it touches, and we'll post that here. Separately, we still owe the check that a pre-M5 Mac refuses by name at start.

@ashhart

ashhart commented Nov 5, 2026 •

Copy link
Copy Markdown
Owner

Our M5 Ultra run on this head (0af3eb8c8), against its merge base with main, Qwen3.8-27B: no token differences in 20 prompt rows. Prefill is level at 32K (1,669.7-1,670.0 tok/s on both sides). Decode after a prompt moves -1.9% to +1.0% across prompt lengths, and the stream sweeps -1.0% to +0.2% from 1 to 8 streams. The pattern isn't consistent across prompt lengths, so we read it as noise, though each side has only two runs. The GB10 runs are still to come.

The six OCR cases still wait on the inputs I asked for above.

@ashhart

ashhart commented Nov 5, 2026

Copy link
Copy Markdown
Owner

Our GB10 run on the same head (0af3eb8c8), Qwen3.8-27B: level. Decode after a 2048-token prompt is -0.2% to +0.0% at every start position, and the 1-to-8-stream sweeps are -0.1% to +0.1% for chat and code, greedy and sampled. The OCR cases still wait on the inputs above.

@ashhart

ashhart commented Nov 5, 2026

Copy link
Copy Markdown
Owner

Our M3 Ultra run on this head (0af3eb8c8): the native server, started with --vision on the checkpoint's six published metadata files and no weight shards, exits 1 with exactly Unlimited-OCR-2 needs an M5-class GPU (tensor units), before any weight is read. That's the refusal we wanted on chips without tensor units. The M5 Ultra and GB10 27B runs above are level, so only the six OCR cases are still to come, and they wait on the inputs I asked for.

@BobClawblaw

Copy link
Copy Markdown
Contributor Author

Thanks @ashhart, and for the 27B runs and the M3 Ultra check. The inputs are attached to a release on the fork, on a tag at this head: https://github.com/BobClawblaw/TensorFold/releases/tag/uocr-0af3eb8c8

They're the exact files every M5 Ultra ocr run on this PR read: five pages, and six references each named <case>.txt for <case>.png. document.txt is the five pages as one request, in the order report-1, report-2, paper-1, paper-3, paper-10.

file bytes sha256
report-1.png 60967 4bf458382a0a03004eca70987b504b09244bbb84a6b368cd1dac422268d96d00
report-2.png 51244 e0c165c2ea90f9f511fd3ed8dd0d1f00d39bbbdad2f41294b108f43f427e753f
paper-1.png 516919 b892374f079051cee7fc4344643fbd2d2e5772d42595a2565b4d3ebe3da311bb
paper-3.png 432135 c5867a9706c5649c37565905f705c279bb870aeedff21c00f404dfef32f31aa7
paper-10.png 356719 052d2b3578b7394787ceaabdbb649f3ee9ecb8691e853aded2f6f86e63f69086
report-1.txt 683 58c9ad432da973e4c062bde428363be0d5800e603b26fb3750ad477273beccf8
report-2.txt 502 366f37903b4eec023fb5c9a9762f88512ce5cc621b8f637546995dc6f925bdb3
paper-1.txt 2231 76c607079fe5d7b3a0beeb61785b92e203908291ae548bb77d0b57f7c3aecac8
paper-3.txt 4310 216de8962d746e0fe3003d6c724cfefc5d01392b2633abe0e3f8fb38a8e3eaf3
paper-10.txt 6197 8d377476d6ea03c6c95d7700ed8f74ce7ff62917f45cbdc4e5a6ede2a2600d82
document.txt 13998 34699873f08879229f6e66699804e963d34279734bd6eff829e064e951dc61c8

SHA256SUMS in the release has the same hashes for shasum -a 256 -c. I downloaded the assets back and all 11 check out.

The references are the release's own infer in torch 2.15.0 bf16 on an Apple arm64 CPU (transformers 4.52.1, greedy, crop_mode=True). The same code under torch on x86 CPUs gives different documents, so please compare against these files rather than regenerating them elsewhere.

Our last run on this head, from the recipe's build of 0af3eb8c8: the text without boxes is identical in 6/6 cases, and 10 of 288 box coordinates moved. Five of those are in paper-1, by at most 10; the other five, in paper-3, paper-10 and the document, move by 1 each.

@ashhart

ashhart commented Nov 6, 2026

Copy link
Copy Markdown
Owner

Thanks @BobClawblaw for the inputs and their hashes; that's what our six-case run on the M5 Ultra reads, and with
the 27B runs and the M3 Ultra refusal in, those cases are the last run on our side. Two things meanwhile:

  1. Main has moved again. core/contract/registry.zig conflicts (the test block above all_metal moved), and
    on today's main tests/contract/metal_main.zig is already at 600 lines, so the suite's three lines take it to
    603. Please rebase and keep that file at 600 or below, for example by moving a suite's setup into its own file.
  2. The notices entry. Since 5 Nov, code following a library doesn't land even with a notice; a model's own
    release code does. Your entry lists _sparse_base_impl.py (Apache-2.0, "Copyright 2023 DeepSeek-AI and The
    HuggingFace Inc. team") for the decoder's attention, page2d rotary and MoE gate. That file ships in the release
    but carries a library's copyright: by its own header it's Transformers' DeepSeek-V2 modeling code, so under the
    rule it's library code even though it ships in the release. The parts that follow it (the decoder's attention,
    page2d rotary and MoE gate) need rewriting from documented behaviour and the checkpoint's config, as LightOnOCR-3 (0.8B, 1B, 4B) on native HIP, with the model card's oracle #1965 did. A
    black-box comparison against the release is fine as the oracle, and the notice entry can then go. Could you also confirm each listed file's own header, in particular whether modeling_deepseekv2.py is
    Baidu's MIT code or carries the same DeepSeek and Hugging Face line?
    And "Thanks to @BobClawblaw." at the end of the CHANGELOG line.

@ashhart

ashhart commented Nov 6, 2026

Copy link
Copy Markdown
Owner

@BobClawblaw, the notices question is settled on our side. _sparse_base_impl.py carries "Copyright 2023 DeepSeek-AI and The HuggingFace Inc. team" and is Transformers' DeepSeek-V2 modeling code. So it counts as library code, even though it ships in the release, and library code doesn't land here even with a notice. The release's own MIT files can stay as they are, with their entry.

So the decoder's attention, the page2d rotary and the MoE gate need rewriting from documented behaviour (the paper, the config and what the release's outputs show), not from that file, and that file's notices entry goes. With that and the rebase, keeping tests/contract/metal_main.zig at 600 lines or fewer, our six OCR cases come next.

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.

Unlimited-OCR-2 (document parsing, BF16) on native Metal

3 participants