Model Demo: PaliGemma 3B INT8 Vision-Language Prefill

NOTE: The Jupyter Notebook below is included in the Chimera SDK and can be run interactively by running the following CLI command:

$ quadric sdk notebook

From the Jupyter Notebook window in your browser, select the notebook named /quadric/sdk-cli/examples/models/paligemma/paligemma_int8_demo.ipynb.

PaliGemma 3B INT8 Vision-Language Prefill on Chimera GPNPU

PaliGemma-3b-mix-448 is a vision-language model: a SigLIP-So400m image encoder turns a 448×448 image into 1024 tokens, a linear projector maps them into the 18-layer Gemma-1 decoder, and the decoder captions the image from the prompt caption en.

This notebook follows the sibling pi-0.5 example's formula. The custom ops run on the Chimera Instruction Set Simulator (ISS) as one fused Chimera Compute Language (CCL) kernel built with ChiPy: a cut of the SigLIP tower, the projector, an in-kernel bridge that concatenates the baked prompt embeddings, and a cut of the Gemma decoder — vision → language, end-to-end, in one binary. The rest of the model runs in PyTorch on the host CPU: the SigLIP patch embedding and the encoder layers before the cut (float, fp16 weights), the decoder layers after the cut (the same per-tensor INT8 torch model the kernel's scales come from) and the float final norm + lm_head that turn the last hidden state into the first caption token. The Chimera Graph Compiler (CGC) compiles the kernel and the cycle-accurate ISS runs it; the ISS output is compared against the ChiPy CPU reference of the same graph, then the first caption token and the image features are compared against the torch reference (the shipped checkpoint run in float). The whole model ships as one 3.9 GB INT8 checkpoint.

Note: we do not run the full model on the ISS here. We run 1) a single-layer test with 1 SigLIP layer + projector + 1 decoder layer and 2) a representative cut with more SigLIP and decoder layers, both at the real 1024-image-token / 1028-token prefill width. The whole model (all 27 SigLIP + 18 decoder layers in one binary) compiles and runs the same way from a shell call, but it does not clear the terminal comparison (cosine 0.935 against the ChiPy CPU model, match 26.7%); its measurement and the reason are in the README's Results section.

Why one fused kernel across two modalities?

A vision-language model has a seam where the image tokens meet the text. In a two-kernel design the vision tower's output leaves the ISS, the host CPU appends the prompt embeddings, and a second kernel starts the decoder — a round trip through host memory on every image. Here the seam is inside the kernel: paligemma::concatTokens appends the prompt embeddings, baked in at build time like the weights, to the projector output, so the kernel's only runtime input is the SigLIP hidden state at the cut and the decoder prefill starts on-chip. The same ChiPy composition that fuses the SigLIP tail with the projector fuses the projector with the Gemma decoder, and the decoder reuses pi-0.5's Gemma ops instead of a second copy of them.

Why rotate the down-projection input?

Gemma's GeGLU product, the tensor the MLP down projection consumes, is the hardest activation in the decoder to quantize per-tensor: most of its values sit below one INT8 step while a few channels carry values hundreds of times larger, so a single scale either clips the outliers or flattens the bulk to zero, and the error compounds over 18 layers until the caption describes a different room. Multiplying each token's product by an orthogonal Walsh-Hadamard matrix spreads every outlier across many channels and leaves a near-Gaussian tensor that INT8 represents well; because the rotation is orthogonal it folds into the down-projection weights, so the layer computes the same function. The kernel applies it as paligemma::hadamardMulQuantize: the rotation mixes across the 512 groups of 32 channels, one processing-element tile wide, so every butterfly stage is a tile-pair add and subtract inside on-chip memory with no data shuffled within a tile, and the host-side INT8 torch model applies the identical rotation so both paths quantize the same numbers.

The two stages in detail — the vision tower that turns the image into 1024 tokens, and the prefix-LM decoder that turns those tokens into the first caption token. The dashed region is what the fused kernel runs on the ISS; the grey boxes run in PyTorch on the host CPU.

Model stages

• Host CPU front end (float torch) — image preprocessing, SigLIP patch + position embedding and the encoder layers before the cut (fp16 weights from the INT8 checkpoint) produce vision_hidden[1,1024,1152], the kernel's only runtime input.
• SigLIP tail (kernel) — the last N encoder layers: siglip::layerNormQuantize → biased INT8 q/k/v projections → siglip::sdpaBidir (bidirectional multi-head attention, FixedPoint32 Q/K/V quantized inside the op) → out_proj → GELU MLP. The residual stream stays FixedPoint32.
• Projector (kernel) — post_layernorm + the 1152→2048 multi_modal_projector as SigLIP ops → image_features[1,1024,2048].
• Bridge (kernel) — paligemma::concatTokens appends the 4 baked prompt embeddings → the [1,1028,2048] prefill sequence.
• Gemma decoder (kernel) — the first M decoder layers, prefix-LM (bidirectional over the 1028 tokens), with pi-0.5's ops: pi05::rmsNormQuantize → INT8 q/k/v → applyRoPE → attention → o proj → GeGLU MLP → residual adds. Two ops are PaliGemma's own: paligemma::sdpaPersistOffsetFp16V keeps Q and K int8 but runs the softmax and P·V in fp16 with an fp16 V cache (the int8 probability grid deleted most of the diffuse attention mass over 1028 keys), and paligemma::hadamardMulQuantize rotates the GeGLU product with a tile Hadamard transform before quantizing it to INT8 for the down projection (see below). The norm → q/k/v and norm → gate/up boundaries are smoothed per channel in memory (helpers/smooth_quant.py): the norm weight is divided and the matmul weights multiplied by the same per-channel factors, so the INT8 grid sees a flat tensor and the layer is unchanged in exact arithmetic.
• Host CPU tail (torch) — decoder layers M..17 (INT8) → float final norm + lm_head → the first caption token.

Model specs

Note — split pipeline.

• On the Chimera GPNPU: the last N SigLIP encoder layers, the projector, the prompt bridge and the first M Gemma decoder layers, as one fused CCL kernel.
• On the host CPU: the SigLIP patch embedding and the encoder layers before the cut, the decoder layers after the cut, and the final norm + lm_head read-out, in PyTorch.

The full pipeline can run end-to-end on the GPNPU — see yolox_e2e_chipy.ipynb or retina_net_e2e_chipy.ipynb for examples that compose CCL custom ops with ChiPy to keep every stage on-chip.

1. Setup

The SDK's openpi quantization helpers (quadric_pyquant.models.openpi) need transformers >= 4.53. A kernel whose transformers already meets that floor (the SDK images do) needs nothing; a kernel built from examples/requirements.txt carries 4.51.0, so for it the next cell side-installs transformers 4.55.0 into a notebook-local directory (the pi-0.5 mechanism; 4.55 runs on Python 3.9 and 3.10) and puts it on the path for the kernel and the helper subprocesses. Then the helpers/ package is made importable and the verbose CGC logs are silenced.

import os
import sys
from importlib.metadata import version
from pathlib import Path

from packaging.version import Version

SIDE_INSTALL = Version(version("transformers")) < Version("4.53")
if SIDE_INSTALL:
    if not (Path("paligemma_t455") / "transformers").exists():
        %pip install --quiet --no-deps --target=paligemma_t455 transformers==4.55.0 "tokenizers>=0.21,<0.22"
    T455 = os.path.abspath("paligemma_t455")
    sys.path.insert(0, T455)
    os.environ["PYTHONPATH"] = T455 + os.pathsep + os.environ.get("PYTHONPATH", "")
os.environ["MPLBACKEND"] = "Agg"
print("transformers side-install:", "4.55.0 (notebook-local)" if SIDE_INSTALL else "not needed")

WARNING: Retrying (Retry(total=4, connect=None, read=None, redirect=None, status=None)) after connection broken by 'NewConnectionError('<pip._vendor.urllib3.connection.HTTPSConnection object at 0x767552c585b0>: Failed to establish a new connection: [Errno -2] Name or service not known')': /transformers/
WARNING: Retrying (Retry(total=3, connect=None, read=None, redirect=None, status=None)) after connection broken by 'NewConnectionError('<pip._vendor.urllib3.connection.HTTPSConnection object at 0x767552c587c0>: Failed to establish a new connection: [Errno -2] Name or service not known')': /transformers/
WARNING: Retrying (Retry(total=2, connect=None, read=None, redirect=None, status=None)) after connection broken by 'NewConnectionError('<pip._vendor.urllib3.connection.HTTPSConnection object at 0x767552c589a0>: Failed to establish a new connection: [Errno -2] Name or service not known')': /transformers/
WARNING: Retrying (Retry(total=1, connect=None, read=None, redirect=None, status=None)) after connection broken by 'NewConnectionError('<pip._vendor.urllib3.connection.HTTPSConnection object at 0x767552c58cd0>: Failed to establish a new connection: [Errno -2] Name or service not known')': /transformers/
WARNING: Retrying (Retry(total=0, connect=None, read=None, redirect=None, status=None)) after connection broken by 'NewConnectionError('<pip._vendor.urllib3.connection.HTTPSConnection object at 0x767552c58eb0>: Failed to establish a new connection: [Errno -2] Name or service not known')': /transformers/
WARNING: Retrying (Retry(total=4, connect=None, read=None, redirect=None, status=None)) after connection broken by 'NewConnectionError('<pip._vendor.urllib3.connection.HTTPSConnection object at 0x76755264ca30>: Failed to establish a new connection: [Errno -2] Name or service not known')': /tokenizers/
WARNING: Retrying (Retry(total=3, connect=None, read=None, redirect=None, status=None)) after connection broken by 'NewConnectionError('<pip._vendor.urllib3.connection.HTTPSConnection object at 0x76755264ccd0>: Failed to establish a new connection: [Errno -2] Name or service not known')': /tokenizers/
WARNING: Retrying (Retry(total=2, connect=None, read=None, redirect=None, status=None)) after connection broken by 'NewConnectionError('<pip._vendor.urllib3.connection.HTTPSConnection object at 0x76755264ce80>: Failed to establish a new connection: [Errno -2] Name or service not known')': /tokenizers/
WARNING: Retrying (Retry(total=1, connect=None, read=None, redirect=None, status=None)) after connection broken by 'NewConnectionError('<pip._vendor.urllib3.connection.HTTPSConnection object at 0x76755264d060>: Failed to establish a new connection: [Errno -2] Name or service not known')': /tokenizers/
WARNING: Retrying (Retry(total=0, connect=None, read=None, redirect=None, status=None)) after connection broken by 'NewConnectionError('<pip._vendor.urllib3.connection.HTTPSConnection object at 0x76755264d240>: Failed to establish a new connection: [Errno -2] Name or service not known')': /tokenizers/
WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
Note: you may need to restart the kernel to use updated packages.
transformers side-install: 4.55.0 (notebook-local)

import logging
from pathlib import Path

assert (
    Path.cwd() / "helpers" / "paligemma_env.py"
).exists(), "run this notebook from examples/models/paligemma"

## examples/models on sys.path makes `paligemma.*` importable as a namespace package
## (the pi0_5 example does the same for `pi0_5.*`).
PKG_PARENT = str(Path.cwd().parent.resolve())
if PKG_PARENT not in sys.path:
    sys.path.insert(0, PKG_PARENT)

from paligemma.helpers import paligemma_env as env

logging.getLogger("epu").setLevel(logging.ERROR)
logging.getLogger("epu").handlers = [logging.NullHandler()]

print(env.summary())
print("\nimages:", ", ".join(p.name for p in sorted(env.IMAGES.glob("*.jpg"))))

example dir          : examples/models/paligemma   OK
model files (hf/)    : examples/models/paligemma/model/hf/config.json   MISSING
model checkpoint     : examples/models/paligemma/model/paligemma_quant_pertensor_params.pth   OK
input images         : examples/models/paligemma/data/images   OK
reports (npz + logs) : examples/models/paligemma/data/reports   MISSING
ISS builds           : examples/models/paligemma/data/builds   MISSING
pi-0.5 ops           : examples/models/pi0_5   OK
pi05_ops.hpp         : examples/models/pi0_5/pi05_ops.hpp   OK
tvm ($TVM_PYTHON/$TVM_HOME) : importable in this env
clang ($QLLVM_ROOT_PATH) : set, found

images: 00_flowers.jpg, 01_ramen.jpg, 02_coffee.jpg, 03_coffee_machine.jpg, 04_dog.jpg, 05_cat.jpg, 06_backpack.jpg, 07_bay_bridge.jpg, 08_sheep.jpg, 09_mug.jpg, 10_keyboard.jpg

2. Download model

The model is hosted on S3 as one INT8 checkpoint, paligemma_int8.pth (3.9 GB): the frozen per-tensor int8 codes of the 126 Gemma decoder linears (exactly what the kernel bakes), fp16 SigLIP / projector / embedding tensors for the parts that run in torch on the host CPU, the calibrated scales, and the config/tokenizer/processor files (unpacked next to it). The next cell populates model/hf/ once, local-first. The FP32 weights are not hosted; the optional fp32 backend in the Interactive section needs them placed locally. The checkpoint is a quantized derivative of google/paligemma-3b-mix-448 and is distributed subject to the Gemma Terms of Use.

License Notice

This example runs google/paligemma-3b-mix-448, which the Appendix to the Gemma Terms of Use lists as Gemma, as an INT8 model derivative hosted on Quadric's model bucket. The Terms govern Gemma and "Model Derivatives", which they define to include "modifications to Gemma" and "works based on Gemma": https://ai.google.dev/gemma/terms

Use is subject to the restrictions in section 3.2, which incorporate the Gemma Prohibited Use Policy: https://ai.google.dev/gemma/prohibited_use_policy Section 3.1 sets the conditions for distributing Gemma or Model Derivatives: a copy of the agreement, the required notice, and a statement of any modification — see ATTRIBUTION.md, LICENSE-gemma.txt and NOTICE in this directory. Downloading the checkpoint below means accepting the Terms.

from IPython.display import Image, display

S3_BASE = "https://sdk-cli-models.s3.us-east-2.amazonaws.com/paligemma"
S3_MODELS = {"paligemma_int8.pth": f"{S3_BASE}/paligemma_int8.pth"}

model_dir = env.fetch_model(s3_models=S3_MODELS)
print("model dir:", model_dir.relative_to(env.PALIGEMMA_DIR))

display(Image(filename=str(sorted(env.IMAGES.glob("*.jpg"))[0]), width=336))
print("prompt:", repr("caption en"))

  paligemma_int8.pth                     3892.7 MB — downloaded
  config.json                               0.0 MB — unpacked
  generation_config.json                    0.0 MB — unpacked
  added_tokens.json                         0.0 MB — unpacked
  special_tokens_map.json                   0.0 MB — unpacked
  tokenizer.json                           17.5 MB — unpacked
  tokenizer.model                           4.3 MB — unpacked
  tokenizer_config.json                     0.0 MB — unpacked
  preprocessor_config.json                  0.0 MB — unpacked
model dir: model/hf

prompt: 'caption en'

3. Compile & Run

run_fused_kernel handles the full flow in one call. On the first call it captures the torch reference in a stock subprocess — the shipped checkpoint run in float (fp16 SigLIP / embeddings, decoder weights on their int8 grid): every SigLIP and decoder layer boundary, the RoPE tables, the logits and the greedy caption — then builds the INT8 model (the decoder's norm → matmul boundaries are smoothed per channel on that input, the attention's probability and V grids widened to mirror the kernel's fp16 P·V, and the weights frozen onto per-tensor INT8 grids), keeps float copies of the final norm and lm_head, and runs the INT8 torch decoder on the reference inputs for its own references and a host-only INT8 caption. Then, per configuration, it builds the fused @chipy.func() (SigLIP tail → projector → concatTokens → decoder cut, whose attention keeps Q and K int8 but runs the softmax and P·V in fp16 with an fp16 V cache, and whose MLP quantizes the Hadamard-rotated GeGLU product for the down projection; the SigLIP tail and the decoder quantize on symmetric INT8 windows, matching the CCL quantizers, which carry no zero point, and the SigLIP layernorm writes the normalized rows to a fixed-point scratch buffer before quantizing, and its attention runs in fp16 end-to-end because an INT8 probability grid deletes most of the diffuse SigLIP attention mass), runs one CPU forward through the ChiPy chain to harvest frac_bits for every FixedPoint32 intermediate, compiles the kernel with CGC, runs it on the ISS and compares the terminal output against the ChiPy CPU reference. Finally the host CPU finishes the remaining decoder layers + the float read-out for the first token, greedy-decodes a caption from the kernel's image features, and prints both next to the torch reference, plus the ISS cycle count.

ISS vs ChiPy CPU. cosine ≥ 0.96 OR match% ≥ 96% between the ISS and the ChiPy CPU output of the kernel's terminal tensor (the last decoder layer in the cut): the CCL kernels against the ChiPy model of the same graph. The comparison against the torch reference follows in section 4.

Cost. The single-layer test is the fast end-to-end run; the representative run is the heavier one. Each decoder layer at seq=1028 carries ~105 MB of INT8 weights into the kernel; the host caption decode (12 greedy steps through the 18-layer INT8 torch decoder) adds about half a minute per run. Both cuts build into the same ccl_build/paligemma_vlm_kernel directory.

3.1 Single-layer test (1 SigLIP layer + projector + 1 decoder layer)

The last SigLIP encoder layer, the projector, the prompt bridge and decoder layer 0, at the real prefill width. Small enough to compile + run on the ISS quickly while exercising the entire vision → bridge → language path and the host CPU hand-off on both ends. The cell prints the ISS-vs-ChiPy-CPU comparison, the first token from the ISS-backed chain next to the torch reference's, the caption decoded from the kernel's image features, and the ISS cycle count.

from paligemma.helpers.paligemma_e2e import run_fused_kernel

single_layer = run_fused_kernel(num_vision_layers=1, num_decoder_layers=1, num_cores=1)
assert single_layer["passed"]
print(f"single layer: {single_layer['cycles']:,} cycles ({single_layer['cycles'] / 1e6:.2f}M)")

Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. `use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`.


[patch] attention + activations + rmsnorm
[tag] tag_all_modules
[config] generate_quant_config -> defaults.w_per_channel=False
[quantize] quantize_model


WARNING:root:Detected 1 tied-parameter group(s); quantization will break the tie (each site quantized independently):
  - model.language_model.embed_tokens.weight is lm_head.weight


  quantize_s=16.1
[load_quant_params] paligemma_quant_pertensor_params_nested_vision.pth (calibrated scales)
[siglip] 217 activation windows made symmetric INT8 (the CCL quantizers carry no zero point)
[decoder] 144 activation windows made symmetric INT8 (the pi05 quantizers carry no zero point)
[decoder] 36 softmax-probability and V grids widened to 16 bits (the kernel's P.V runs in fp16)
[smooth] 36 norm -> matmul boundaries folded (alpha 0.5, statistics from the calibration input) in 10s
[hadamard] rotated 18 down_proj inputs (tile Hadamard, group 32, symmetric INT8; windows calibrated on the run input)
[freeze] freeze_model_weights
  freeze_s=3.9
  set int8_matmul=False on 90 QMatMul
[int8] model loaded + quantized in place in 60s
[ref] torch reference for 00_flowers.jpg / 'caption en' in 0s: caption 'A vase filled with flowers on a table. The vase is', first token id=235280 'A'
[int8] INT8 torch decoder on the clean inputs in 92s: layer 0 cosine vs torch reference = 0.9999, host-only first token id=235280 'A', host-only caption 'A vase filled with water and flowers on a table.'
========================================================================
ISS Validation: fused PaliGemma kernel (1 SigLIP + projector + 1 decoder layers, 1 core(s), image 00_flowers.jpg, prompt 'caption en')
========================================================================
  [chain] AttentionConfig causal=False (prefix-LM bidirectional)
[chipy-cpu] calibrated 29 ops and built the CPU reference in 1.5s
[cgc] compiled paligemma_vlm_kernel in 43s
[cgc] >2 GB const blob patched, rebuilding the host binary
[iss] ran in 673s
  decoder layer 0 out [TERMINAL]: ISS vs ChiPy-CPU cosine=0.999842 max_diff=1.7899 mean_diff=0.0543 match=100.00%  [PASS]
  decoder layer 0 out [informational]: ChiPy-CPU vs INT8 torch cosine=0.999300, ISS vs INT8 torch cosine=0.999297, ISS vs torch reference cosine=0.999282
  image features [informational]: ISS vs ChiPy-CPU cosine=0.999764, ISS vs torch reference cosine=0.998478

  first token  torch reference           : id=235280 (top-5: 'A', 'a', 'In', 'The', 'This')
  first token  INT8 torch (host only)    : id=235280 (top-5: 'A', 'a', 'In', 'The', 'This')
  first token  ChiPy-CPU kernel + host   : id=235280 (top-5: 'A', 'a', 'In', 'The', 'This')
  first token  ISS kernel + host         : id=235280 (top-5: 'A', 'a', 'In', 'The', 'This')
  ISS first token == torch reference     : True  (top-5 shared 5/5, logit Pearson r=0.9803)
  second token after the reference's first ('A') [informational]:
    torch reference                      : ' vase' (top-5: ' vase', ' white', ' small', ' clear', ' glass')
    ISS image features + host decode     : ' vase' (top-5: ' vase', ' clear', ' white', ' glass', ' small')  match=True

  caption  torch reference               : 'A vase filled with flowers on a table. The vase is'
  caption  INT8 torch (host only)        : 'A vase filled with water and flowers on a table.'
  caption  ISS image features + host     : 'A vase of flowers on a table in a living room.'  (342s host decode)

ISS profile: 42,397,629 cycles (42.40M) = 24.9 ms @ 1.7 GHz on 1 core(s)

PASS
single layer: 42,397,629 cycles (42.40M)

3.2 Representative run (3 SigLIP layers + projector + 2 decoder layers)

A larger cut of the same graph: the last three SigLIP layers feed the projector and bridge, then two decoder layers write two layers' worth of K/V into the shared INT8 cache. Still the real 1028-token prefill; the host CPU finishes decoder layers 2–17.

representative = run_fused_kernel(num_vision_layers=3, num_decoder_layers=2, num_cores=1)
assert representative["passed"]
cycles = representative["cycles"]
print(f"representative: {cycles:,} cycles ({cycles / 1e6:.2f}M)")

========================================================================
ISS Validation: fused PaliGemma kernel (3 SigLIP + projector + 2 decoder layers, 1 core(s), image 00_flowers.jpg, prompt 'caption en')
========================================================================
  [chain] AttentionConfig causal=False (prefix-LM bidirectional)
[chipy-cpu] calibrated 66 ops and built the CPU reference in 7.0s
[cgc] compiled paligemma_vlm_kernel in 153s
[cgc] >2 GB const blob patched, rebuilding the host binary
[iss] ran in 1023s
  decoder layer 1 out [TERMINAL]: ISS vs ChiPy-CPU cosine=0.999624 max_diff=4.0240 mean_diff=0.0829 match=100.00%  [PASS]
  decoder layer 1 out [informational]: ChiPy-CPU vs INT8 torch cosine=0.999331, ISS vs INT8 torch cosine=0.999334, ISS vs torch reference cosine=0.999328
  image features [informational]: ISS vs ChiPy-CPU cosine=0.999056, ISS vs torch reference cosine=0.998132

  first token  torch reference           : id=235280 (top-5: 'A', 'a', 'In', 'The', 'This')
  first token  INT8 torch (host only)    : id=235280 (top-5: 'A', 'a', 'In', 'The', 'This')
  first token  ChiPy-CPU kernel + host   : id=235280 (top-5: 'A', 'a', 'In', 'The', 'This')
  first token  ISS kernel + host         : id=235280 (top-5: 'A', 'a', 'In', 'The', 'This')
  ISS first token == torch reference     : True  (top-5 shared 5/5, logit Pearson r=0.9823)
  second token after the reference's first ('A') [informational]:
    torch reference                      : ' vase' (top-5: ' vase', ' white', ' small', ' clear', ' glass')
    ISS image features + host decode     : ' vase' (top-5: ' vase', ' white', ' clear', ' glass', ' small')  match=True

  caption  torch reference               : 'A vase filled with flowers on a table. The vase is'
  caption  INT8 torch (host only)        : 'A vase filled with water and flowers on a table.'
  caption  ISS image features + host     : 'A vase of flowers on a table in a living room.'  (53s host decode)

ISS profile: 94,259,771 cycles (94.26M) = 55.4 ms @ 1.7 GHz on 1 core(s)

PASS
representative: 94,259,771 cycles (94.26M)

4. Compare the Fused Kernel Against the Torch Reference

The comparison in section 3 is ISS against the ChiPy CPU model of the same custom-op graph, so it checks CGC and the ISS rather than the kernel's numerics. This section compares both runs against the torch reference — the shipped checkpoint run in float, captured in a stock process that never imports the SDK's patched modeling — with compare_with_reference:

• Decoder output. The kernel's terminal hidden state (the last decoder layer in the cut) must reach cosine ≥ 0.95 against the reference output of the same layer.
• Image features. The kernel's image_features[1,1024,2048] must reach cosine ≥ 0.95 against the reference inputs_embeds for the 1024 image tokens.
• First token. After the host CPU tail, the ISS-backed chain's first token must match the reference's: the same argmax, or, when the reference's two leading candidates are near-tied ('A' and 'a' for this image, the same word in two cases and less than two logits apart), the same pair of leading candidates. At least 3 of the top-5 must be shared. The logit Pearson correlation is printed for information.

Any failed comparison raises. For a sanity check beyond the first token, each run also prints the second caption token: the reference's own second greedy token next to the token the host CPU decodes from the kernel's image features after the reference's first token. It is printed for information, not asserted.

from paligemma.helpers.paligemma_e2e import compare_with_reference

compare_with_reference(single_layer)
compare_with_reference(representative)

Torch reference comparison: 1 SigLIP + projector + 1 decoder layer(s)
  decoder layer 0 out  ISS vs torch reference cosine=0.9993 (min 0.95)  [PASS]
  image features  ISS vs torch reference cosine=0.9985 (min 0.95)  [PASS]
  first token     ISS kernel + host tail 'A' vs torch reference 'A' (argmax match=True), top-2 [235250, 235280] vs [235250, 235280], top-5 shared 5/5 (rule: argmax match or same top-2 set, and >= 3/5)  [PASS]
  logit Pearson r=0.9803  [informational]
PASS
Torch reference comparison: 3 SigLIP + projector + 2 decoder layer(s)
  decoder layer 1 out  ISS vs torch reference cosine=0.9993 (min 0.95)  [PASS]
  image features  ISS vs torch reference cosine=0.9981 (min 0.95)  [PASS]
  first token     ISS kernel + host tail 'A' vs torch reference 'A' (argmax match=True), top-2 [235250, 235280] vs [235250, 235280], top-5 shared 5/5 (rule: argmax match or same top-2 set, and >= 3/5)  [PASS]
  logit Pearson r=0.9823  [informational]
PASS





{'passed': True,
 'cosine': 0.9996240966588946,
 'match': 99.99981000729572,
 'cycles': 94259771,
 'first_id': 235280,
 'first_token': 'A',
 'reference_first_id': 235280,
 'reference_first_token': 'A',
 'iss_top5': [235280, 235250, 886, 651, 1596],
 'reference_top5': [235280, 235250, 886, 651, 1596],
 'logits_pearson': 0.9823447234885938,
 'second_token': ' vase',
 'reference_second_token': ' vase',
 'reference_cosine': 0.9993276658562168,
 'feature_cosine': 0.9981316200554299,
 'caption': 'A vase of flowers on a table in a living room.',
 'fp32_caption': 'A vase filled with flowers on a table. The vase is',
 'int8_caption': 'A vase filled with water and flowers on a table.',
 'num_vision_layers': 3,
 'num_decoder_layers': 2,
 'compile_s': 152.77938121557236,
 'iss_s': 1023.3831305727363}

5. Interactive — pick an image, a prompt and a backend

Set the three variables and run. Any image works (drop a .jpg into data/images/ and point IMAGE at its index, or pass a path); the prompt is a PaliGemma task prefix such as caption en or answer en what color is the table?.

• int8-cpu — the per-tensor INT8 torch decoder on the host CPU on the reference image features, float read-out (about a minute for a new image/prompt pair, then cached).
• iss — the fused kernel on the ISS for the chosen cut (num_vision_layers, num_decoder_layers), compared against the ChiPy CPU reference; the caption is decoded on the host CPU from the kernel's image features (about 10 minutes for the 1+1 cut).
• fp32 — the stock FP32 model, greedy-decoded in a clean subprocess; needs the FP32 checkpoint in model/hf/ (not hosted: full_model.pth, or the Hugging Face shards).

gradio_app.py is the same thing as a standalone chat UI (see the Summary).

from IPython.display import Image, display

from paligemma.helpers.paligemma_e2e import caption

## ---- set these, then run ---------------------------------------------------------------------
IMAGE = 0  # index into sorted(data/images/*.jpg), or a path to any .jpg
QUESTION = "answer en what is on the table?"
BACKEND = (
    "int8-cpu"  # "fp32" (reference) | "int8-cpu" (INT8 decoder on host CPU) | "iss" (fused kernel)
)
## ----------------------------------------------------------------------------------------------

images = sorted(env.IMAGES.glob("*.jpg"))
display(Image(filename=str(images[IMAGE] if isinstance(IMAGE, int) else IMAGE), width=336))
result = caption(image=IMAGE, prompt=QUESTION, backend=BACKEND)
print(
    f"\n{BACKEND}: {result['caption']!r}  (first token {result['first_token']!r}, {result['seconds']:.0f}s)"
)

[ref] torch reference for 00_flowers.jpg / 'answer en what is on the table?' in 35s: caption 'vase', first token id=123002 'vase'
[int8] INT8 torch decoder on the clean inputs in 13s: layer 0 cosine vs torch reference = 0.9999, host-only first token id=123002 'vase', host-only caption 'vase'
[int8-cpu] 00_flowers.jpg / 'answer en what is on the table?' -> 'vase'

int8-cpu: 'vase'  (first token 'vase', 49s)

Summary

Key takeaways

1. Custom ops on the ISS, the rest in torch on the host CPU. Only the ops that need custom CCL — the SigLIP and Gemma attention, the fused norm-quantize and INT8 matmul variants, the bridge — go into the kernel. The rest of the model runs unchanged in torch on either side of it, and the kernel drops back into that chain at a float32 layer boundary.
2. One fused binary across the modalities. Vision tail, projector, prompt bridge and decoder compile into a single CCL kernel; the only runtime input is the SigLIP hidden state at the cut, and the prompt embeddings are a build-time constant like the weights.
3. Same ops as pi-0.5, plus one rotation. The Gemma decoder reuses pi05_ops.hpp and the pi-0.5 ChiPy op builders (examples/models/pi0_5); the one PaliGemma-specific decoder op, paligemma::hadamardMulQuantize, rotates the GeGLU product before its INT8 quantization, which is what brings the INT8 decoder's first token and caption back onto the reference's. The SigLIP ops keep the residual stream in FixedPoint32, which is what holds the vision cosine.
4. Results, not just numerics. Besides the two comparisons and the cycle count, every run prints the first caption token from the ISS-backed chain next to the torch reference's, and the caption decoded from the kernel's image features.

Gradio chat interface (standalone)

gradio_app.py serves the demo as a chat: attach an image or pick one of the tracked example images from the dropdown, type a prompt, and choose the backend (INT8 CPU, INT8 fused kernel on the ISS with a selectable cut, or FP32 torch when the FP32 checkpoint is present). gradio pins websockets<13 while this repo pins 13.1, so it is installed into a side directory:

cd examples/models/paligemma
pip install --quiet --no-warn-conflicts --target=paligemma_gradio gradio==4.44.1 "numpy<2" "huggingface_hub<1"
PYTHONPATH=paligemma_t455:paligemma_gradio python gradio_app.py      # http://127.0.0.1:7861

Citation

@article{beyer2024paligemma,
  title   = {PaliGemma: A versatile 3B VLM for transfer},
  author  = {Beyer, Lucas and Steiner, Andreas and Pinto, Andr{\'e} Susano and Kolesnikov, Alexander and Wang, Xiao and Salz, Daniel and Neumann, Maxim and Alabdulmohsin, Ibrahim and Tschannen, Michael and Bugliarello, Emanuele and Unterthiner, Thomas and Keysers, Daniel and Koppula, Skanda and Liu, Fangyu and Grycner, Adam and Gritsenko, Alexey and Houlsby, Neil and Kumar, Manoj and Rong, Keran and Eisenschlos, Julian and Kabra, Rishabh and Bauer, Matthias and Bo{\v{s}}njak, Matko and Chen, Xi and Minderer, Matthias and Voigtlaender, Paul and Bica, Ioana and Balazevic, Ivana and Puigcerver, Joan and Papalampidi, Pinelopi and Henaff, Olivier and Xiong, Xi and Soricut, Radu and Harmsen, Jeremiah and Zhai, Xiaohua},
  journal = {arXiv preprint arXiv:2407.07726},
  year    = {2024}
}