mtmd: support multi-row batching for deepseek-ocr

[ngxson]

Overview

Ref:

• mtmd: add batching support for internvl #24775
• mtmd: deepseek-ocr v1 multi-tile #24717

Allow multi-row batching for deepseek-ocr:

• Only rows with same size will be batched together
• Overview image won't be batched, as it have different size

CC @sfallah for review

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: yes, the featured is already planned out in the mentioned PR; AI only used to implement that plan

[ngxson]

/bot review

[ggml-gh-bot]

Automated code review

I've traced the graph construction (deepseekocr.cpp, deepseekocr2.cpp), the encode/batch glue in clip.cpp/mtmd.cpp (chunk merging, n_tokens() sizing, mtmd_batch_get_output_embd slicing), the preprocessor (fuse_row is true only for v1), and the can_batch_with / support_batch gating. The change is self-contained to the mtmd vision encoder and the logic is consistent. Review below.

---

Review: mtmd multi-row batching for deepseek-ocr

Overall the patch is correct and reasonably scoped. The batch-dimension bookkeeping (n_rows_batch captured before n_batch is repurposed, the per-row SAM-input concat, the per-row tile weave, the n_dim, per-item, n_rows_batch output) is internally consistent, and the mtmd side (mtmd_batch_encode_impl entry merging, n_tokens() using batch_f32.entries.size(), and mtmd_batch_get_output_embd's sequential offset split) lands each chunk's slice on exactly the right flat region of the 3D output tensor. I verified the equivalent single-row path is bit-for-bit unchanged when only one chunk is present, so no regression for the unbatched case. No blocking issues found.

Will slow the review

(point 1) No test covering the new batched path. The existing tools/mtmd/tests/test-deepseek-ocr.py runs mtmd-cli, which now silently routes multi-row images through mtmd_batch because support_batch() is true. OCR-output equality vs ground truth is a weak check here: because the batched matmuls are mathematically per-item-independent, embeddings should be identical to the previous per-row path, so a CER pass proves almost nothing about correctness of the new concat/weave ordering. Add a deterministic comparison: encode the same multi-row image with batching forced on vs off (or compare the MTMD_DEBUG_EMBEDDINGS dump of a batched run to per-row runs concatenated in order). This is the one test that would actually catch a future concat-order or reshape-axis bug. The deepseek-ocr2 (v2) tile-batching path (support_batch newly enabled for v2) should be exercised too — confirm the test image produces multiple tiles with v2, not just v1 rows.

(point 2) v2 overview-batch relies on img.add_viewsep of the first entry only. deepseekocr2.cpp decides whether to append the view separator with if (img.add_viewsep), and the v1 overview weave similarly trusts the first entry. can_batch_with only compares nx/ny/pos, not add_viewsep or actual spatial size. In practice overview chunks (nx includes the separator) and tile/row chunks (nx without it) have different token counts so they never co-batch, so this is currently safe — but it is a structural invariant that the graph trusts implicitly. A one-line assertion at the top of each build (loop over img_batch->entries and assert entry.add_viewsep == img.add_viewsep, like the v1 non-overview loop already does for the other fields) would make the invariant explicit and guard against a future preprocessor change.

Nits

(point 3) deepseekocr2.cpp:36 — ggml_repeat_4d(ctx0, query_embed, query_embed->ne[0], num_queries, inp->ne[2], 1) passes the local num_queries instead of query_embed->ne[1] for the second dim. They are equal today by construction, but using the tensor's own dim (or asserting equality) would make the intent self-documenting and catch a mismatch if the query tensors ever load with an unexpected shape.

(point 4) deepseekocr.cpp:257 — for the overview branch n_rows_batch is captured but the weave below still uses n_batch (equal to it, since the overview branch does not reassign n_batch). Harmless, but a reader has to reason about why one name is captured and another is used for the same value in the same branch; consider using one name consistently.

(point 5) The SAM-input combine in deepseekocr.cpp builds a ggml_concat chain of depth n_rows_batch. With max_tiles small this is fine, but it is more copy/alloc operations than necessary. Not worth changing unless perf data shows it matters; mentioning only so it is a conscious choice.

Style check

ASCII-only, 4-space indent, bracket placement, and naming all match the surrounding code. The added comments explain why (the 4-axis limitation for the weave, the n_batch repurpose) and are not verbose restatements — consistent with AGENTS.md.

This review was generated automatically by pi coding agent using zai-org/GLM-5.2. It may contain mistakes. Maintainers make the final call.

[sfallah]

Opened #26615 on top of this branch: both per-row loops fold into a single permute each. n_dim and grid_x are adjacent and in memory order, so they merge into one axis and free the 4th for n_rows_batch.

Not a speedup — graph nodes just stay flat at 1634 instead of growing ~15 per row.

[sfallah]

@ngxson
thanks for advancing the work on the DSOCR model.
FYI: I ran my OCR regression tests on Metal and CUDA and all the numbers came out identical to the unbatched path.

[ngxson]

perfect, thanks for the confirmation. I'm merging this now