smdesai/granite-docling-258M-coreai

granite-docling-258M — Core AI

IBM granite-docling-258M for Apple's Core AI runtime (iOS 27 / macOS 27). It turns an image of a document page into DocTags (text, layout boxes, tables, formulas, code, lists), entirely on device. The decoder runs on the Neural Engine; the vision encoder and the final norm + vocabulary projection run on the GPU.

The compiled .aimodelc files are built for A19 Pro (h18p) and were tested on iPhone 17 Pro. For any other chip, compile the portable sources in source/ yourself (see below).

Device	Page	Decode	Memory
iPhone 17 Pro	full arXiv page, 1,759 tokens	117 tok/s	~400 MB
iPhone 17 Pro	table page, 400 tokens	130 tok/s	~400 MB

Accuracy: on the arXiv sample page, the generated document text and formulas match an FP32 run of the base model; a few layout coordinates differ by one unit (FP16 near-ties). Output is identical across repeated runs.

Files

File	Size	Role
decoder.h18p.aimodelc	215 MB	30-layer decoder, Neural Engine, compiled for h18p
vision.h18p.aimodelc	179 MB	SigLIP2 vision tower + connector, GPU, compiled for h18p
head.h18p.aimodelc	110 MB	final norm + vocabulary projection, GPU, compiled for h18p
embed_tokens_fp16.bin	110 MB	token embedding table, [100352, 576] float16, row-major (looked up on the CPU)
tokenizer.json	6.8 MB	tokenizer (identical to the base model's)
granite_docling_coreai_bundle.json	4 KB	asset hashes, toolchain and export configuration
source/decoder.aimodel, source/vision.aimodel, source/head.aimodel	494 MB	portable sources of the three models
source/granite_docling_coreai.json	<1 KB	export configuration

Exported with coreai-torch 0.4.3 (torch 2.13) and compiled with coreai-build 3605.6.3.

Compiling for another chip

xcrun coreai-build compile source/decoder.aimodel --platform iOS --architecture <arch> --preferred-compute neural-engine --output out
xcrun coreai-build compile source/vision.aimodel  --platform iOS --architecture <arch> --preferred-compute gpu --output out
xcrun coreai-build compile source/head.aimodel    --platform iOS --architecture <arch> --preferred-compute gpu --output out

coreai-build names its output <name>.<arch>.aimodelc, the same scheme as the files here. Only the h18p build has been validated. Loading the raw .aimodel sources directly on a phone works but takes minutes to specialize, so ship compiled assets.

Model interfaces

Model	Function(s)	Inputs	Outputs
Vision	main	image_tile [1, 3, 512, 512] fp16, RGB in 0…1 (normalization is inside the model)	image_embeds [1, 64, 576] fp16
Decoder	extend_{ctx}_{q}, ctx ∈ {1024, 2048, 3072, 4096}, q ∈ {8, 64}	transformer_input [1, q, 1, 576] fp16, position_ids [1, q] uint16, in_step [1] int32, causal_mask [1, ctx, 1, q] fp16	out_logits [1, q, 1, 576] fp16 (pre-norm hidden states, despite the name)
Head	main	hidden [1, 1, 576] fp16	logits [1, 1, 100352] fp16

The decoder keeps its KV caches as runtime state: key_cache and value_cache, each [30, 1, 192, 1, ctx] fp16. Allocate both once from the 4096 function's state descriptors and pass prefix views ([..., 0..<ctx]) to the smaller buckets, so the cache carries over when the context grows. Always call the smallest bucket that holds every position written so far: the Neural Engine reads the whole padded cache each step, so smaller buckets decode faster.

Running it

1. Tile the page with Lanczos resampling: resize so the longest edge is 2048 (also when upscaling), keeping aspect ratio and rounding the short edge up to even; resize again so both edges are multiples of 512; cut 512×512 tiles row-major, then append the whole page resized to 512×512 as the global view.
2. Encode each tile with the vision model: 64 embeddings per tile.
3. Build the prompt (no BOS) and tokenize it with tokenizer.json:

   <|start_of_role|>user<|end_of_role|>
   for each tile row r, column c:  <fake_token_around_image><row_{r}_col_{c}> + <image>×64, then "\n" after each row
   \n<fake_token_around_image><global-img> + <image>×64 + <fake_token_around_image>
   Convert this page to docling.<|end_of_text|>\n<|start_of_role|>assistant<|end_of_role|>
   

   <image> is token 100270; row/column indices start at 1.
4. Embed on the CPU: each token's row from embed_tokens_fp16.bin, or, for each <image> token, the next vision embedding (tile order, 64 rows per tile). Use the raw rows; any scaling is inside the decoder.
5. Call the decoder. Each call covers the positions from the previous multiple of q up to the newest token: prefill with q = 64, decode with q = 8 (a decode call re-processes the earlier tokens of its 8-aligned window). For a call starting at position start with n real rows:
   • transformer_input rows 0…n−1 are the embeddings, the rest zero;
   • position_ids[j] = start + j; in_step = start;
   • causal_mask[0, c, 0, j] = 0 if j < n and c <= start + j, else −40000 (padding rows fully masked).
6. Head: pass the decoder output row of the newest real token to the head; argmax of logits is the next token. Stop at <|end_of_text|> (100257), <end_of_utterance> or </doctag>.

Allocate input and output buffers once from the function descriptors and reuse them as output views; runtime-allocated outputs that are not released promptly grow memory on the phone. Greedy decoding can loop on some pages; stopping when a short token cycle repeats several times is a cheap safeguard.

Conversion notes

The decoder is re-authored in the Neural Engine layout from Apple's coreai-models examples (rank-4 tensors, 1×1 convolution projections, per-head attention, static entrypoints). The first prompt token develops residual activations of ~1.2×10⁵, which overflows FP16, so the residual stream is carried at 1/32 scale (folded into the weights) and RMSNorm is expressed as a LayerNorm that stays finite in FP16.

License and attribution

Apache 2.0, the same as the base model. All model weights are IBM's granite-docling-258M; this repository only changes their format. Please cite and credit IBM's model when you use it.