feat(models): Support PaddleOCR-VL vision model (rework of #2320) - #2356
Open
subin9 wants to merge 42 commits into
Open
feat(models): Support PaddleOCR-VL vision model (rework of #2320)#2356subin9 wants to merge 42 commits into
subin9 wants to merge 42 commits into
Conversation
…sion model Adds PaddleOCR-VL, a document-OCR VLM: a SigLIP/NaViT variable-resolution vision encoder, an Adaptive-MLP (mlp_AR) connector, and an ERNIE-4.5-0.3B language model. Loads via --arch paddleocr_vl (HF PaddleOCRVLForConditionalGeneration). Follows the Apache-2.0 PaddleOCR-VL modeling code (transformers modeling_paddleocr_vl.py) and its GGUF conventions. mistral.rs is MIT; this credits the source project.
…tConfig Add an optional QuantizedConfig to TextConfig and top-level Config (serde default) so the text decoder load path can build quantizable layers. Prep for wiring q/k/v/o through Column/RowParallelLayer.
Build ERNIE decoder q/k/v/o via ColumnParallelLayer/RowParallelLayer so they are Arc<dyn QuantMethod> and ISQ-swappable. Thread the loading context (device mapper, loading_isq, per-layer comm) through Attention::load -> DecoderLayer::load -> ErnieTextModel::load -> PaddleOcrVlModel::new -> the loader, which no longer drops normal_loading_metadata.
Drop the per-call causal_mask_offset (Tensor::from_vec on the decode hot path). The MultimodalModel forward now builds the mask once with CausalMasker::make_causal_mask (None for single-token decode) and threads &AttentionMask through the text decoder into Sdpa/paged attention.
Rewrite MultimodalModel::forward to handle [batch, seq] and the full seqlen_offsets slice instead of assuming batch 1. Positions come from get_rope_index_batched + mrope_position_ids_for_input; text/decode embeds use a new on-device Merger::embed_tokens, dropping the per-token from_vec position and embed paths.
Behavior-preserving rewrites clippy suggests: inline the write! arg, iterate map values, use ? in block_pool lookup, drop the wildcard field pattern covered by `..`.
…r ISQ Under ISQ the LM projections quantize onto the GPU while the vision tower, connector, token embedding, and RMSNorm weights stayed on the CPU staging device, crashing a GPU ISQ run (input must live on CUDA / device mismatch in mul). Load tower/connector/merger on real_device and route the three RMSNorm loads through the device mapper, matching qwen2vl/qwen3/llama.
…example The example defaulted to CPU/f32 for deterministic parity. Add opt-in env toggles (PADDLEOCR_VL_GPU/_ISQ/_PAGED) to exercise the accelerator paths without a recompile. Default behavior unchanged.
…cking The vision-specific args carry a single image_grid_thw, so a batch>1 forward with an image has no per-sequence grid and get_rope_index_batched panicked on grids[b] (engine worker died with "channel closed"). Return a clear model error for that case so the engine reports it per-request and stays alive; text/decode batching and single-image (batch=1) prefill are unaffected.
Replace the batch=1 guard with real per-row batching. The vision-specific args now carry one image grid per batch row (was a single tuple), the inputs processor passes all rows' grids (was grid_accum.first()), and the forward splits the concatenated pixel_values by each grid's t*h*w, runs vision+connector+merge per row, and stacks to [batch, seq, D]. get_rope_index_batched already handled per-sequence grids. Concurrent image requests now batch through one LM forward instead of erroring; verified byte-identical to sequential over a 9-crop batch, and the weight-free unit tests pass.
`cargo fmt --all -- --check` and `cargo clippy --tests -- -D warnings` both failed on this branch, so CI could not go green on any later commit. Pure mechanical changes: rustfmt over llg.rs, the model's mod.rs and the recognize example, plus one cloned_ref_to_slice_refs lint in a text.rs test helper.
…ping them preprocess() took the first image and discarded the rest, but the chat template still emitted one placeholder per image, so expand_placeholders indexed grids[1] on the second one and panicked with index-out-of-bounds (the engine worker dies with "channel closed"). Two image_url parts in one OpenAI message reach this. Reject at the root instead: preprocess bails when it is not handed exactly one image, and the caller propagates that with `?` rather than .expect(), so the request gets a clear error and the engine stays up. One image per request is the documented scope; supporting N needs per-row grid lists in the forward.
process_inputs gated everything on `input_seqs.iter().all(|seq| seq.has_images())`, so one text-only request in the batch dropped pixel_values AND the image grids for every row, and the OCR rows silently came back positioned as if they had no image. The forward had the matching assumption: it walked image_grid_thw with the grid index as the batch index, so it only worked when every row had exactly one image. Split the two things that were conflated. image_grid_thw is now one Option per row and is kept whenever that row's prompt has an image at all, because mrope is recomputed from the whole prompt on every pass; vision_rows separately names the rows whose image tokens land in this pass, in pixel_values concat order. The forward embeds all rows as text and overwrites only vision_rows with the merged result, so a text-only row and an OCR row batch together correctly. Sequence::has_images is already window-scoped, so decode steps and prompt chunks that hold no image tokens keep the grid and skip the tower instead of re-running it. Also drops the duplicate detok/re-encode that built input_ids_full: it is taken from seq.get_toks() directly now, which is what set_toks_and_reallocate stored.
…ontent supports_prefix_cacher was never overridden, so it took the trait default of false and the engine disabled prefix caching for this model entirely: cache_blocks was never reached and paged attention shipped with half of what it advertises. Turning it on is only safe once a block hash can tell two images apart. Every expanded placeholder is the same token id, so two requests that share a prompt and a grid shape have byte-identical token streams and hash to identical blocks. Pages of a scanned PDF share a size, so this is the normal case, not a corner: with the cacher on and no image span registered, three different crops of the same size all came back with the first crop's text. The inputs processor now registers the IMAGE_START..IMAGE_END span as an mm feature carrying the image content hash, which is what generate_mm_extra_keys folds into the block hash. That also clamps a hit off the middle of the span, where input_ids would hold fewer image slots than the connector emits rows and Merger::forward, which counts slots from the start of input_ids, would scatter them at the wrong offsets. Registering the span also lets the prompt chunk planner see the image, so chunked prefill cuts at the span edges instead of through it. That splits an OCR prompt into a text prefix, the span, and a text suffix, so a row's grid is attached only once its token window actually holds the image tokens: get_rope_index emits a position block for every grid it is handed, and the first chunk stops before the span, so a grid there emits positions for tokens that are not in the window. No is_first_prompt_chunk mask override is added: a later chunk carries prefix_cache_len = chunk.start, so paged takes the prefix gather path, and that path reads causality off the mask. Forcing AttentionMask::None there falls back to flash_params.causal and attends non-causally within the chunk, which measurably corrupts output. Recorded as a comment so it does not get "fixed" later. Verified on CUDA against a local checkpoint over a 9-page PDF (86 layout regions, including 4 OTSL tables): prefix-cache-only is byte-identical to the plain bf16 run, paged and paged+prefix-cache differ by 2 characters in 9220, and CPU f32 by one. tests/paged_serve.rs covers the image leak and the mixed text/image batch; it enables the cacher explicitly because the builder defaults it off while the server does not.
…ble sequences The completion branch drains self.running with `while !self.running.is_empty()` and, for a sequence whose has_images() disagreed with the batch being formed, pushed it back onto self.running. That is the same queue the loop pops from, so once every remaining sequence disagrees with running[0] the loop rotates them forever: one engine thread pins a core and no request in flight ever completes. Reaching it only takes a text-only request arriving while an image request is running, which any OpenAI-compatible client can do against a multimodal model. Reproduced with PaddleOCR-VL on CUDA: two image requests plus one text-only request in flight with paged attention on, and the engine stops making progress (GPU idle, one core at 100%). It reproduces the same way before this branch, so it is not specific to PaddleOCR-VL, but the model advertises paged attention and this is on that path. Collect the incompatible ones in a separate queue so the loop always drains, then preempt them after it so they are retried in a later step, which is what the length bucketing right below already does with its non-selected buckets. tests/paged_serve.rs::paged_mixed_text_and_image_batch covers this shape.
…mplate The checkpoint's template reads a message's content as a list of typed parts, because that is how it finds the image part. A plain text message arrives as a bare string, and MessagesAction::Keep hands that to jinja untouched, which iterates it per character: `content["type"]` never matches, so the text is dropped and the model is asked to continue an empty user turn. It answers with byte-fallback garbage. Wrap bare string content in a single text part before rendering. The text-only prompt now decodes token for token the same as transformers 5.13 on the same checkpoint, and image requests are unchanged (the image path already sent typed parts). The body of the default Processor::process moves to a free function so this override can reuse the tokenizer handling instead of restating it.
…uite parity_serve.rs and paged_serve.rs became one tests/paddleocr_vl.rs, and the coverage grew to the paths whose failure mode on real weights is a plausible wrong answer rather than an error: - greedy parity against the transformers golden (CPU f32) - text-only parity against the transformers golden, which fails as byte-fallback garbage if the template never receives the text - extra images rejected, engine still serving afterwards - paged prefix cache does not serve one page's image KV for another - a text-only request in flight next to an image one still makes progress Each skips unless PADDLEOCR_VL_WEIGHTS points at a checkpoint, and the paged pair additionally needs cuda/metal, so a default `cargo test` run stays weight-free.
Lifts the one-image restriction added earlier in this branch. That commit stopped the crash but kept the scope at a single image, because the forward mapped one grid to one batch row; the per-row rework since then is what makes N workable. Most of the machinery already handled it: get_rope_index walks a list of grids per sequence, expand_placeholders spends one grid per placeholder, Merger::forward scatters connector rows in order, and find_placeholder_delimited_ranges returns a range per image so each one gets its own prefix-cache content hash. What was left was preprocess returning a single grid and the forward reading one grid per row. preprocess now emits one grid row per image with their patches concatenated in message order, and the forward runs the tower per image and concatenates the connector output before the scatter. Single-image output is unchanged over the 86-region corpus.
A sequence that gets preempted re-prefills its whole prompt, and for this model that meant re-running the SigLIP tower over an image whose output cannot have changed. On 20 concurrent mixed-size crops that put serving them together at 0.16x of serving them one at a time; text-only requests in the same shape do not regress like that, because their re-prefill is cheap. Uses the existing EncoderCacheManager the way idefics2 and qwen3_vl do: the inputs processor already had the per-image content hashes, so the forward looks up vision+connector output per image and only runs the tower on a miss. Mixed-size concurrency goes to 0.64x measured cold on both sides, and a repeated page skips the tower entirely. Output over the 86-region corpus is unchanged. The remaining gap is the scheduler keeping one length bucket per step and preempting the rest, which shows up on text-only requests too (5x for same-length batches against 1x for mixed) and needs it to separate its running set from the batch it schedules.
Rebasing onto master picked up four signature changes this branch predates: - `embedding()` returns `Arc<dyn QuantMethod>`, so `Merger` holds one and calls `embedding_forward` with the activation dtype instead of `Module::forward`. - `ModelInputs` gained `adapter_leases`, filled the same way the other vision inputs processors do. - `DeviceMappedModelLoader::non_mapped_size_in_bytes` gained a `quantization` parameter. - `IsqModelLoader::promoted_isq_predicates` is now required; the embedding and lm_head are promoted, matching qwen3_vl. The manual `Display`/`Error` impls for `MistralRsError` are dropped: master derives them through `thiserror` now and the hand-written ones collide. Kept as its own commit rather than squashed back: the embedding change lands in code whose shape several later commits on this branch rewrote, so folding it into the commit that first added the file would mean rewriting the adaptation once per intermediate shape.
… prompt The example had grown into a test harness: a layout-manifest reader with four environment toggles for device, ISQ, paged attention and the prefix cacher. No other example in the repo reads environment variables, and none of that belongs in a usage sample. It now takes an image path and an optional task prompt, which is the whole API surface a reader needs. The configuration those toggles covered is exercised by tests/paddleocr_vl.rs and by the server flags instead. Docs page regenerated with docs/scripts/render_examples.py.
Code Metrics Report━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ Language Files Lines Code Comments Blanks ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ C Header 23 4454 3116 790 548 CSS 3 282 252 6 24 CUDA 121 25439 20835 1707 2897 Dockerfile 1 35 19 9 7 HTML 2 27 27 0 0 JavaScript 3 578 563 12 3 Jinja2 7 694 656 5 33 JSON 27 17989 17986 0 3 Makefile 1 18 16 0 2 MDX 36 6285 0 4616 1669 Metal Shading Lan| 37 14409 11401 1136 1872 PowerShell 1 657 571 31 55 Python 151 12539 10418 480 1641 Shell 3 1071 852 115 104 Plain Text 53 10687 0 9209 1478 TOML 28 1389 1209 39 141 TypeScript 11 1649 1410 66 173 YAML 3 25 23 2 0 ───────────────────────────────────────────────────────────────────────────────── Jupyter Notebooks 3 122 83 23 16 |- Markdown 1 60 30 22 8 |- Python 1 122 113 1 8 (Total) 304 226 46 32 ───────────────────────────────────────────────────────────────────────────────── Markdown 273 12055 0 8981 3074 |- BASH 24 300 221 47 32 |- Dockerfile 2 14 12 0 2 |- JSON 6 289 289 0 0 |- PowerShell 1 1 1 0 0 |- Python 135 7349 6120 306 923 |- Rust 62 3840 2849 394 597 |- TOML 7 92 78 0 14 (Total) 23940 9570 9728 4642 ───────────────────────────────────────────────────────────────────────────────── Rust 701 331485 296772 5817 28896 |- Markdown 418 9947 452 8304 1191 (Total) 341432 297224 14121 30087 ───────────────────────────────────────────────────────────────────────────────── Svelte 19 1969 1826 51 92 |- CSS 1 4 4 0 0 |- JavaScript 19 921 767 25 129 (Total) 2894 2597 76 221 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ Total 1507 466797 378971 42194 45632 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Hi @EricLBuehler, this reworks #2320 along the lines you asked for. Closes #2128.
Adds PaddleOCR-VL-1.5, a document-OCR VLM: SigLIP/NaViT variable-resolution vision encoder -> Adaptive-MLP (
mlp_AR) connector -> ERNIE-4.5-0.3B LM. Loads via--arch paddleocr_vl(HFPaddleOCRVLForConditionalGeneration).Addressing the review
ColumnParallelLayer/RowParallelLayer/ReplicatedLayerwithisq_layer_regexeson the loader. Q4K, Q8_0, Q5K and HQQ4 all run, and a UQFF round-trip reproduces the direct-ISQ output.all(has_images), so a text-only request in the batch no longer costs the image rows theirpixel_valuesand mrope grids.Multi-image requests and text-only requests both work now; the connector output is cached by image content hash so a re-prefill skips the vision tower.
One fix landed outside the model: the completion scheduler pushed image-incompatible sequences back onto the queue it was draining, so a text-only request arriving next to an image one spun without either completing. It is in
paged_attention/scheduler.rs, reproduces onmaster, and is a standalone commit if you would rather take it separately.Results
RTX 4070 Ti SUPER, a 9-page PDF through the layout stage, 86 regions. Reference is GPU bf16 without paged attention:
The paged differences are kernel precision and predate this branch. Also driven over the HTTP server with
--paged-attn on: streaming, sampling, multi-turn, and concurrent mixed text/image requests.tests/paddleocr_vl.rscovers greedy and text-only parity against the transformers golden, multi-image, the prefix-cache leak and the mixed batch. They skip unlessPADDLEOCR_VL_WEIGHTSis set (the paged pair also needscuda), so a defaultcargo teststays weight-free:Notes
<fcel>/<nl>tokens (special=false) that the completion detok drops without it. Text and formula regions are unaffected.