Skip to content

VLM inference: encode the image once per request instead of per decode step #126

Description

@camilobrownpinilla

What

During VLM generation the vision tower is re-run on every decode step. Because there is no public seam to inject precomputed visual embeddings, any autoregressive decode loop re-encodes the image (ViT + adapter) once per generated token. The current (and only) consumer is the VLM eval pipeline. For G generated
tokens this is ≈G ViT+adapter passes instead of 1; with a 16-frame video clip and a realistic
max_new_tokens, the encode cost dominates eval wall-clock.

Add a seam to compute the projected visual embeddings once per request and inject them into the
transformer without re-encoding — e.g. a prepare/forward path (or a VLMWrapper helper) that
accepts precomputed prefix_embeds / image_features in place of pixel_values. This is
independent of KV-cache decoding and changes
no training behavior and reuses the identical vision_encoder + adapter outputs.

Scope

  • kempnerforge/model/vlm.py — the four ModalityStrategy.prepare methods + VLMWrapper.forward:
    an alternate path that takes precomputed visual embeddings; factor the encode
    so the result can be computed once and
    reused. Arch-agnostic.
  • kempnerforge/eval/vlm/adapter.py — consumer: precompute the clip embedding once per request and
    reuse it across the decode loop (outputs unchanged).
  • Tests under tests/unit/ — assert the encoder runs once per request (call-count spy) and that the
    cached path's logits match the per-step path.

No new config fields (inference-time behavior).

Backward compatibility

Additive and backward compatible. The existing pixel_values forward path is unchanged (training
and the current eval path keep working); the new precomputed-embeddings path is an opt-in
alternative used by decode loops. Outputs are numerically identical (the same vision_encoder +
adapter result, reused), pinned by an equivalence test. No config changes.

Metadata

Metadata

Labels

enhancementNew feature or request

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions