Cache visual features in the VLM eval decode loop

examples/vlm-evaluation/adapter.py

@@ -15,8 +15,8 @@
 v1 scope and deliberate choices (see README.md in this directory):
 
 - **Generation: no transformer KV cache, batched, data-parallel.** The decode loop
-  re-runs the transformer (including the vision encoder + adapter) over the
-  growing sequence each step. There is no transformer KV cache
+  re-runs the transformer over the growing sequence each step; the vision encoder +
+  adapter run **once** per request (``encode_visual``). There is no transformer KV cache
   (``Transformer.forward`` forbids combining ``kv_caches`` with any
   image-conditioning route), and KempnerForge has no image-conditioned KV-cache
   decode path. Requests are decoded in batches
@@ -453,16 +453,16 @@ def _generate_batch(
     visual tokens from attention as in training). ``pixel_values is None`` is a
     text-only batch: the vision tower is skipped, ``num_image_tokens`` is 0, and
     ``model(None, ...)`` runs the pure-text forward. There is no transformer KV
-    cache and no vision cache: ``model(...)`` re-runs over the growing
-    **right-padded** batch each step, re-encoding the vision tower each time.
-    Right-padding matches the training
+    cache; the vision tower is encoded **once** per request and reused each step,
+    while ``model(...)`` re-runs the transformer over the growing **right-padded**
+    batch. Right-padding matches the training
     layout: the image prefix stays at positions ``0..n-1`` and text is contiguous
     from ``n`` for every row (so image/text RoPE distances are consistent across
     rows), and the trailing pads are causally masked, so a batched forward gives
     each row the same real-position logits as decoding it alone. Each row's next
     token is read at its own last real position; EOS / ``max_new_tokens`` / first
-    ``until`` match are tracked per row. ``B == 1`` reproduces the single-request
-    path exactly.
+    ``until`` match are tracked per row, and a row that finishes is dropped from
+    every later forward. ``B == 1`` reproduces the single-request path exactly.
     """
     until: list[str] = resolved["until"]
     max_new_tokens: int = resolved["max_new_tokens"]
@@ -498,43 +498,76 @@ def _generate_batch(
         prompts.append(ids)
 
     generated: list[list[int]] = [[] for _ in range(batch_size)]
-    done = [False] * batch_size
-    row_index = torch.arange(batch_size, device=device)
+    # No decode steps to run: return before touching the vision tower, whose output
+    # would be discarded by the zero-iteration loop below.
+    if max_new_tokens == 0:
+        return [""] * batch_size, [0] * batch_size
+
+    # Encode the clip once and reuse it every step; text-only batches skip this.
+    visual_embeds = None if pixel_values is None else model.encode_visual(pixel_values)
+
+    # Token buffer, allocated once: prompt + room for the whole generation, pad-filled.
+    # A step writes its sampled token in place, so no per-row tensor is rebuilt or
+    # copied to the device each iteration. ``lengths`` stays a Python list: it indexes
+    # the buffer and bounds the slice, and reading it off a device tensor would force
+    # a sync every step.
+    max_prompt_len = max(p.shape[0] for p in prompts)
+    token_buf = torch.full(
+        (batch_size, max_prompt_len + max_new_tokens), pad_id, dtype=torch.long, device=device
+    )
+    lengths = [p.shape[0] for p in prompts]
+    for i, p in enumerate(prompts):
+        token_buf[i, : lengths[i]] = p
+
+    # Rows still generating. Finished rows are dropped from the forward batch: their
+    # logits were discarded anyway, and with right-padding + causal attention a row's
+    # logits at its own real positions depend on neither the other rows nor the batch
+    # width, so survivors decode bit-identically to the uncompacted loop.
+    active = list(range(batch_size))
 
     for _ in range(max_new_tokens):
-        # Rebuild the right-padded batch from prompt + tokens generated so far.
-        seqs = [
-            torch.cat([prompts[i], torch.tensor(generated[i], dtype=torch.long, device=device)])
-            for i in range(batch_size)
-        ]
-        real_len = torch.tensor([s.shape[0] for s in seqs], device=device)
-        cur_max = int(real_len.max().item())
-        input_ids = torch.full((batch_size, cur_max), pad_id, dtype=torch.long, device=device)
-        for i, s in enumerate(seqs):
-            input_ids[i, : s.shape[0]] = s
-
-        logits, _ = model(pixel_values, input_ids, frame_mask=frame_mask)
+        cur_max = max(lengths[i] for i in active)
+        if len(active) == batch_size:
+            # Full batch: slice the buffer in place and pass the visual tensors
+            # through untouched (no index_select, no copy).
+            input_ids = token_buf[:, :cur_max]
+            step_embeds, step_frame_mask = visual_embeds, frame_mask
+        else:
+            # Compacted: the visual tensors are row-aligned with input_ids, and
+            # VLMWrapper.forward rejects a cache whose batch differs from input_ids.
+            active_index = torch.tensor(active, device=device)
+            input_ids = token_buf[active_index, :cur_max]
+            step_embeds = None if visual_embeds is None else visual_embeds[active_index]
+            step_frame_mask = None if frame_mask is None else frame_mask[active_index]
+        real_len = torch.tensor([lengths[i] for i in active], device=device)
+
+        # pixel_values is never passed: the clip is already encoded into
+        # visual_embeds, and forward rejects both together (they could disagree).
+        # Text-only batches have no cache and pass None for both.
+        logits, _ = model(None, input_ids, frame_mask=step_frame_mask, visual_embeds=step_embeds)
         # Each row's next-token logits sit at its own last real position (the
         # output is already trimmed to text positions for JD/MoT; CA has no
         # image prefix), not at [-1] (a pad for shorter rows).
-        next_logits = logits[row_index, real_len - 1]
+        next_logits = logits[torch.arange(len(active), device=device), real_len - 1]
         next_tokens = sample(next_logits, temperature, top_k, top_p)
 
-        for i in range(batch_size):
-            if done[i]:
-                continue
-            token_id = int(next_tokens[i].item())
+        still_active: list[int] = []
+        for j, i in enumerate(active):
+            token_id = int(next_tokens[j].item())
             if eos_id is not None and token_id == eos_id:
-                done[i] = True
                 continue
             generated[i].append(token_id)
+            token_buf[i, lengths[i]] = token_id
+            lengths[i] += 1
             if len(generated[i]) >= max_new_tokens:
-                done[i] = True
-            elif until:
+                continue
+            if until:
                 text = tokenizer.decode(generated[i], skip_special_tokens=True)
                 if _first_stop(text, until) is not None:
-                    done[i] = True
-        if all(done):
+                    continue
+            still_active.append(i)
+        active = still_active
+        if not active:
             break
 
     outputs: list[str] = []

examples/vlm-evaluation/tests/unit/test_adapter.py

@@ -9,6 +9,7 @@
 
 from __future__ import annotations
 
+import inspect
 import json
 import sys
 import types
@@ -590,18 +591,50 @@ def test_per_row_eos_independent(self, arch_wrapper):
 
 
 class _CaptureModel:
-    """Minimal ``VLMWrapper`` stand-in: records the ``frame_mask`` each forward
-    receives and returns deterministic zero logits (greedy -> token 0), so the
-    decode-loop plumbing can be asserted without a real transformer."""
+    """Minimal ``VLMWrapper`` stand-in: records what each forward receives
+    (``pixel_values``, ``visual_embeds``, ``frame_mask``, batch width), counts
+    ``encode_visual`` / forward calls, and returns deterministic zero logits
+    (greedy -> token 0), so the decode-loop plumbing can be asserted without a
+    real transformer.
+
+    ``__call__`` mirrors ``VLMWrapper.forward``'s signature exactly, including
+    ``labels`` and parameter order; ``test_capture_model_matches_forward_signature``
+    pins that so a positional call site cannot bind arguments differently here
+    than in production.
+    """
 
     def __init__(self, num_image_tokens: int, vocab_size: int = 256) -> None:
         self.num_image_tokens = num_image_tokens
         self._vocab = vocab_size
         self.seen_frame_masks: list[torch.Tensor | None] = []
+        self.seen_pixel_values: list[torch.Tensor | None] = []
+        self.seen_visual_embeds: list[torch.Tensor | None] = []
+        self.seen_batch_sizes: list[int] = []
+        self.encode_calls = 0
+        self.forward_calls = 0
+        self.returned_embeds: torch.Tensor | None = None
+
+    def encode_visual(self, pixel_values):
+        self.encode_calls += 1
+        # Identifiable object whose rows are distinguishable (row i is filled with i),
+        # so a forward can be asserted to have received *this* cache, and a compacted
+        # forward to have received the surviving rows of it.
+        batch = pixel_values.shape[0]
+        self.returned_embeds = (
+            torch.arange(batch, dtype=torch.float32)
+            .view(batch, 1, 1)
+            .expand(-1, self.num_image_tokens, 1)
+            .clone()
+        )
+        return self.returned_embeds
 
-    def __call__(self, pixel_values, input_ids, frame_mask=None):
-        del pixel_values
+    def __call__(self, pixel_values, input_ids, labels=None, frame_mask=None, visual_embeds=None):
+        del labels
+        self.forward_calls += 1
         self.seen_frame_masks.append(frame_mask)
+        self.seen_pixel_values.append(pixel_values)
+        self.seen_visual_embeds.append(visual_embeds)
+        self.seen_batch_sizes.append(input_ids.shape[0])
         b, t = input_ids.shape
         return torch.zeros(b, t, self._vocab), None
 
@@ -653,6 +686,153 @@ def test_generate_batch_text_only_budget_excludes_image_tokens():
     assert len(out) == 1
 
 
+def test_encode_visual_called_once_per_generate_batch():
+    """Vision encoded once per request, not once per decode step -- and the cache that
+    encode_visual returned is what every forward actually conditions on.
+
+    Counting encode calls alone is not enough: dropping ``visual_embeds=`` from the
+    model call would leave the count at 1 while the real wrapper re-projected the clip
+    on every step. The identity assertion is what makes that revert fail.
+    """
+    model = _CaptureModel(num_image_tokens=8)
+    pixel_values = torch.randn(2, 3, 16, 16)
+    prompt_ids = [torch.tensor([5, 9], dtype=torch.long), torch.tensor([7, 3], dtype=torch.long)]
+    max_new = 5
+    r = _resolve_gen_kwargs({"max_new_tokens": max_new}, 128)
+    _gen_texts(model, _MockTokenizer(), pixel_values, prompt_ids, r, 64)
+    assert model.encode_calls == 1
+    assert model.forward_calls == max_new
+    # No row finishes early here, so every step takes the uncompacted fast path and
+    # forwards the cache object itself.
+    assert model.returned_embeds is not None
+    assert all(ve is model.returned_embeds for ve in model.seen_visual_embeds)
+
+
+def test_generate_batch_never_passes_pixels_with_cache():
+    """VLMWrapper.forward rejects pixel_values and visual_embeds together, so the decode
+    loop must pass only the cache once the clip is encoded."""
+    model = _CaptureModel(num_image_tokens=8)
+    prompt_ids = [torch.tensor([5, 9], dtype=torch.long)]
+    r = _resolve_gen_kwargs({"max_new_tokens": 3}, 128)
+    _gen_texts(model, _MockTokenizer(), torch.randn(1, 3, 16, 16), prompt_ids, r, 64)
+    assert model.forward_calls == 3
+    assert all(pv is None for pv in model.seen_pixel_values)
+    assert all(ve is not None for ve in model.seen_visual_embeds)
+
+
+def test_capture_model_matches_forward_signature():
+    """_CaptureModel must accept arguments exactly as VLMWrapper.forward does.
+
+    A drifting mock (missing ``labels``, or ``visual_embeds`` ahead of ``frame_mask``)
+    would keep every decode-loop test green while a positional call site in production
+    bound the cache to ``labels`` and dropped image conditioning.
+    """
+    real = list(inspect.signature(VLMWrapper.forward).parameters)
+    mock = list(inspect.signature(_CaptureModel.__call__).parameters)
+    assert real[0] == "self" and mock[0] == "self"
+    assert real[1:] == mock[1:]
+
+
+def test_encode_visual_not_called_for_text_only():
+    """A text-only batch (pixel_values=None) never touches the vision tower."""
+    model = _CaptureModel(num_image_tokens=8)
+    prompt_ids = [torch.tensor([5, 9, 12], dtype=torch.long)]
+    r = _resolve_gen_kwargs({"max_new_tokens": 3}, 128)
+    _gen_texts(model, _MockTokenizer(), None, prompt_ids, r, 64)
+    assert model.encode_calls == 0
+    assert model.forward_calls == 3
+    assert all(ve is None for ve in model.seen_visual_embeds)
+
+
+def test_zero_max_new_tokens_skips_vision_and_decode():
+    """max_new_tokens=0 is a supported task value; with no decode steps to run, the
+    vision tower must not be run only for its output to be discarded."""
+    model = _CaptureModel(num_image_tokens=8)
+    prompt_ids = [torch.tensor([5, 9], dtype=torch.long), torch.tensor([7], dtype=torch.long)]
+    r = _resolve_gen_kwargs({"max_new_tokens": 0}, 128)
+    texts, counts = _gb(model, _MockTokenizer(), torch.randn(2, 3, 16, 16), prompt_ids, r, 64)
+    assert texts == ["", ""] and counts == [0, 0]
+    assert model.encode_calls == 0 and model.forward_calls == 0
+
+
+class _RowTokenModel(_CaptureModel):
+    """_CaptureModel that makes greedy decoding row-dependent: row j's argmax is its
+    own first prompt token, which is stable across steps and survives compaction (the
+    row is identified by its content, not its position). Lets one row hit EOS while
+    another keeps generating."""
+
+    def __call__(self, pixel_values, input_ids, labels=None, frame_mask=None, visual_embeds=None):
+        logits, _ = super().__call__(pixel_values, input_ids, labels, frame_mask, visual_embeds)
+        logits = logits.clone()
+        for j in range(input_ids.shape[0]):
+            logits[j, :, int(input_ids[j, 0])] = 1.0
+        return logits, None
+
+
+def test_finished_rows_dropped_from_forward_batch():
+    """A row that hits EOS stops being forwarded: its logits were discarded anyway, so
+    keeping it in the batch is pure waste (and skews the reported tok/s)."""
+    model = _RowTokenModel(num_image_tokens=8)
+    # Row 0 generates token 5 forever; row 1 generates token 7 == eos and finishes
+    # on the first step.
+    prompt_ids = [torch.tensor([5, 9], dtype=torch.long), torch.tensor([7, 3], dtype=torch.long)]
+    max_new = 4
+    r = _resolve_gen_kwargs({"max_new_tokens": max_new}, 128)
+    texts, counts = _gb(
+        model, _MockTokenizer(eos_token_id=7), torch.randn(2, 3, 16, 16), prompt_ids, r, 64
+    )
+    assert texts == ["5 5 5 5", ""] and counts == [4, 0]
+    # Step 1 sees both rows; every later step sees only the survivor.
+    assert model.seen_batch_sizes == [2, 1, 1, 1]
+    # The compacted cache is row 0 of the original (encode_visual fills row i with i),
+    # so the survivor stays conditioned on its own image.
+    assert model.seen_visual_embeds[-1].shape[0] == 1
+    assert torch.equal(model.seen_visual_embeds[-1][0], model.returned_embeds[0])
+
+
+def test_compaction_preserves_survivor_output():
+    """Dropping finished rows must not change what the remaining rows generate."""
+    prompt_ids = [torch.tensor([5, 9], dtype=torch.long), torch.tensor([7, 3], dtype=torch.long)]
+    r = _resolve_gen_kwargs({"max_new_tokens": 4}, 128)
+    batched = _gen_texts(
+        _RowTokenModel(num_image_tokens=8),
+        _MockTokenizer(eos_token_id=7),
+        torch.randn(2, 3, 16, 16),
+        prompt_ids,
+        r,
+        64,
+    )
+    alone = _gen_texts(
+        _RowTokenModel(num_image_tokens=8),
+        _MockTokenizer(eos_token_id=7),
+        torch.randn(1, 3, 16, 16),
+        [prompt_ids[0]],
+        r,
+        64,
+    )
+    assert batched[0] == alone[0]
+
+
+def test_frame_mask_compacted_with_finished_rows():
+    """frame_mask is row-aligned with the forward batch, so it must be compacted too --
+    otherwise the survivor would be masked with another row's padded-frame pattern."""
+    model = _RowTokenModel(num_image_tokens=8)
+    frame_mask = torch.tensor([[True, True], [True, False]])
+    prompt_ids = [torch.tensor([5, 9], dtype=torch.long), torch.tensor([7, 3], dtype=torch.long)]
+    r = _resolve_gen_kwargs({"max_new_tokens": 3}, 128)
+    _gen_texts(
+        model,
+        _MockTokenizer(eos_token_id=7),
+        torch.randn(2, 2, 3, 16, 16),
+        prompt_ids,
+        r,
+        64,
+        frame_mask=frame_mask,
+    )
+    assert model.seen_frame_masks[0] is frame_mask  # full batch: passed through as-is
+    assert torch.equal(model.seen_frame_masks[-1], frame_mask[:1])
+
+
 # ---------------------------------------------------------------------------
 # Guards: arch + not-implemented methods
 # ---------------------------------------------------------------------------