PaDoc: Layout-Grounded Parallel Decoding for Document Parsing

PaDoc is an open-source implementation of tree-supervised training and fork-token decoding for document parsing. A single image-text model emits a compact main sequence and starts independent content branches at learned fork tokens. No draft model or additional prediction head is required.

main:    layout 1 [fork] -> layout 2 [fork] -> layout 3 [fork] -> EOS
                    |                 |                 |
branch:             metadata/text 1  metadata/text 2   metadata/text 3

This repository intentionally exposes one method and three supported paths:

Packed Transformers decoding, private WebUI code, legacy vLLM adapters, native-AR benchmarks, experimental vLLM patches, and cluster-specific configuration are not part of this package.

The default fork map is:

{"</SP_LAYOUT>": "<SP_META>"}

Training expands each tree into one flat sequence:

[prompt | main | shadow+branch_1 | ... | shadow+branch_n]

The shadow token repeats the main-stream trigger. It is label-masked, reuses the trigger position, and lets the first branch transition match inference. Fork-target labels are also masked because the inference engines inject those tokens. Each branch sees the main prefix before its trigger and its own causal history, but no later main tokens or sibling branches.

The two training backends implement that same visibility rule:

• padoc_tree_fa2 is the default experiment path. It stores compact O(S) segment metadata and executes the tree pattern as a virtual FlashAttention-2 variable-length batch.
• sdpa materializes a boolean (S, S) tree mask. It is slower and uses quadratic mask memory, but is useful for inspection and correctness checks.

See docs/method.md for the exact layout, labels, position IDs, and attention semantics.

Python 3.10-3.12 is supported; Python 3.12 is used for the tested environments. Transformers training/inference and vLLM deployment use different Torch stacks, so install them in separate environments. Each setup script selects its backend-specific lock file automatically; do not combine the two environments.

From the PaDoc repository root, create the environment and run the setup script. The tested stack uses Torch 2.8.0 with CUDA 12.8 and FlashAttention 2.8.3:

conda create -n padoc-transformers python=3.12 -y
conda activate padoc-transformers
bash scripts/setup_transformers_env.sh

The setup script installs the locked CUDA 12.8 dependencies, FlashAttention, and PaDoc in the required order.

The tested deployment stack uses vllm==0.19.1 and its separate Torch 2.10 CUDA 12.8 environment:

conda create -n padoc-vllm python=3.12 -y
conda activate padoc-vllm
bash scripts/setup_vllm_env.sh

Register the four PaDoc tokens and store the fork map in config.json:

padoc-preprocess \
  --base-model /path/to/base-image-text-model \
  --output /path/to/padoc-base

Fine-tuned checkpoints derived from that output retain the metadata. Both inference backends read it automatically.

Training accepts JSON, JSONL, a directory of those files, or a YAML data mix. Paths in a YAML mix are resolved relative to that YAML file. A sanitized, text-only record is included at examples/train.jsonl:

{
  "query": "Extract the document regions.",
  "images": [],
  "response": {
    "prefix": "",
    "children": [
      {
        "prefix": "<SP_LAYOUT>40 20 620 80</SP_LAYOUT>",
        "children": [
          {
            "prefix": "<SP_META>{\"category\":\"title\"}</SP_META>Example title",
            "children": []
          }
        ]
      }
    ]
  }
}

A child whose prefix starts with a fork-target token is a branch. Other children continue the main stream. The flat layout accepts branches from the main stream and rejects nested branch-of-branch forks explicitly.

The single training entry defaults to the reproducible padoc_tree_fa2 backend:

padoc-train \
  --model /path/to/padoc-base \
  --data-config configs/data.example.yaml \
  --output runs/padoc \
  --freeze-vit \
  --gradient-checkpointing

For local multi-GPU training, use the same entry through torchrun:

PADOC_NPROC_PER_NODE=8 bash scripts/train_torchrun.sh \
  /path/to/padoc-base runs/padoc configs/data.example.yaml

Run the explicit two-dimensional tree-mask reference with:

padoc-train \
  --model /path/to/padoc-base \
  --data-config configs/data.example.yaml \
  --output runs/padoc-sdpa \
  --attention-backend sdpa

Portable single-device and torchrun wrappers are available in scripts/train_single.sh and scripts/train_torchrun.sh.

Activate the backend-specific environment before running its commands. Both backends return the same main-sequence and branch result structure.

Transformers inference has two selectable execution modes:

• parallel (default) prefills the prompt once, snapshots the parent KV cache at each fork, and advances the main stream plus all active branches together in one lockstep GPU batch. --max-concurrent-branches bounds that batch; later branches wait in a queue without losing their fork-time cache state.
• sequential is the batch=1 reference. It finishes a branch after each fork before resuming the main stream, which is useful for tracing and comparison.

Both modes use the same fork semantics and greedy token selection. The parallel path is ordinary Transformers decoding with independently padded KV states; it does not depend on packed decoding or a custom attention patch.

CUDA_VISIBLE_DEVICES=0 padoc-infer \
  --model /path/to/checkpoint \
  --image /path/to/page.jpg \
  --query "Parse this document." \
  --device cuda:0 \
  --execution-mode parallel \
  --max-concurrent-branches 8 \
  --max-new-tokens 512 \
  --max-branch-tokens 512 \
  --json

Use --execution-mode sequential for the reference path. Setting CUDA_VISIBLE_DEVICES=0 exposes exactly one physical GPU to the process; inside that process the selected device is still cuda:0. Replace the value before the command to select another physical GPU. PaDoc requires an explicit device and does not use automatic model sharding for this inference path.

sdpa is the inference default. flash_attention_2 can be selected with --attn-implementation when the optional kernel is installed.

Start the local Transformers dashboard:

CUDA_VISIBLE_DEVICES=0 padoc-transformers-serve \
  --model /path/to/checkpoint \
  --device cuda:0 \
  --execution-mode parallel \
  --max-concurrent-branches 8 \
  --host 127.0.0.1 \
  --port 8000

Open http://127.0.0.1:8000. The dashboard exposes Parallel and Sequential as a segmented control and reports the live main state, active and queued branch counts, and current scheduler batch size. It reads newline-delimited events from POST /v1/stream, draws parsed [0, 1000] layout boxes directly over the uploaded page, and updates the corresponding region content as branch tokens arrive. The server keeps one model loaded and serializes documents; parallelism is within each document's main/branch decode batch.

The deployment engine submits the main sequence and every branch as independent requests to one AsyncLLMEngine. Branch prompts reuse the exact parent prefix, multimodal UUIDs, and cache salt. Prefix caching is always enabled.

CUDA_VISIBLE_DEVICES=0 env -u LD_LIBRARY_PATH padoc-vllm infer \
  --model /path/to/checkpoint \
  --image /path/to/page.jpg \
  --query "Parse this document." \
  --tensor-parallel-size 1 \
  --max-model-len 8192 \
  --max-num-seqs 16 \
  --max-num-batched-tokens 4096 \
  --max-new-tokens 512 \
  --max-branch-tokens 512 \
  --max-concurrent-branches 8 \
  --max-total-branches 64 \
  --max-total-generated-tokens 8192 \
  --document-timeout-s 120

--max-concurrent-branches controls live child requests; excess branches wait in a queue. --max-total-branches is the cumulative per-document cap. Token and wall-clock limits provide independent safeguards. The total generated-token budget counts model-emitted tokens; injected fork targets are structural inputs and are not charged to it.

Start the local vLLM dashboard and HTTP API:

CUDA_VISIBLE_DEVICES=0 env -u LD_LIBRARY_PATH padoc-vllm serve \
  --model /path/to/checkpoint \
  --tensor-parallel-size 1 \
  --host 0.0.0.0 \
  --port 8000

Open http://127.0.0.1:8000. The same bbox-and-region dashboard consumes newline-delimited streaming from POST /v1/stream; non-streaming clients can use POST /v1/generate. GET /health reports backend readiness and prefix-caching support. vLLM always uses parallel external child requests; its continuous scheduler performs the batching, so the dashboard does not expose a sequential toggle for this backend.

pytest -q

The tests check tree expansion, label masking, logical positions, equivalence between compact metadata and the explicit mask, sequential and lockstep-batched cache forking, and the external-fork scheduler's two branch limits. They also cover both dashboard HTTP surfaces, browser image payloads, static assets, and streaming on both backends.

PaDoc is released under the MIT License.