From 1bd4941d56dc5ae9d03a227edc1da22292c35fc7 Mon Sep 17 00:00:00 2001
From: camilobrownpinilla <camilobrownpinilla@gmail.com>
Date: Thu, 23 Jul 2026 15:24:09 -0400
Subject: [PATCH 1/3] Generalize VLM modality strategies; add encode_visual
 seam

Hoist vision projection into VLMWrapper.forward via a shared BaseModalityStrategy; add encode_visual for once-per-request decode.
---
 kempnerforge/model/vlm.py | 242 ++++++++++++++++++++------------------
 tests/unit/test_moma.py   |  11 +-
 tests/unit/test_vlm.py    | 101 +++++++++++++---
 3 files changed, 215 insertions(+), 139 deletions(-)

diff --git a/kempnerforge/model/vlm.py b/kempnerforge/model/vlm.py
index 660f431..6f9d8ed 100644
--- a/kempnerforge/model/vlm.py
+++ b/kempnerforge/model/vlm.py
@@ -4,7 +4,7 @@
 adapter (``MLP2LayerAdapter`` by default; ``LinearAdapter`` available
 via the ``adapter`` registry) projecting image features into the LLM
 embedding space, and the existing ``Transformer``. The arch-specific
-work (composing ``pixel_values`` + ``input_ids`` into a
+work (composing already-projected visual embeds + ``input_ids`` into a
 ``ModalityContext``) lives on a ``ModalityStrategy`` that the wrapper
 holds, so adding a new arch is one new strategy decorator on
 ``@registry.register_modality_strategy`` plus one new ``VLMConfig``
@@ -66,16 +66,20 @@ class ModalityStrategy(Protocol):
     def prepare(
         self,
         wrapper: VLMWrapper,
-        pixel_values: torch.Tensor | None,
+        visual_embeds: torch.Tensor | None,
         input_ids: torch.Tensor,
         frame_mask: torch.Tensor | None = None,
     ) -> ModalityContext:
-        """Compose a ``ModalityContext`` from raw VLM inputs.
-
-        ``pixel_values is None`` is a text-only request (no visual content):
-        the generative arches return a context that drives the pure-text
-        forward — an empty ``ModalityContext()`` for the image-prefix and
-        cross-attention arches — while a non-generative arch may reject it.
+        """Compose a ``ModalityContext`` from already-projected visual embeds.
+
+        Vision projection (encoder + adapter) is arch-independent and runs once
+        in ``VLMWrapper.forward`` (or ``VLMWrapper.encode_visual`` for cached
+        decode), so strategies receive ``visual_embeds`` and never touch pixels
+        or the vision tower. ``visual_embeds is None`` is a text-only request
+        (no visual content): the generative arches return a context that drives
+        the pure-text forward — an empty ``ModalityContext()`` for the
+        image-prefix and cross-attention arches — while a non-generative arch
+        may reject it.
         """
         ...
 
@@ -173,36 +177,44 @@ def _prefix_key_padding_mask(
     return torch.cat([vmask, text_valid], dim=1)
 
 
-@registry.register_modality_strategy("joint_decoder")
-class JointDecoderStrategy:
-    """Joint-Decoder: image embeds prepended to the text sequence.
+class BaseModalityStrategy:
+    """Shared ``ModalityStrategy`` implementation (not registered, not the Protocol).
+
+    ``prepare`` is a template method: it handles the text-only request and the
+    visual-token count ``n``, then delegates arch-specific ``ModalityContext``
+    construction to ``_build_context``. Concrete strategies override
+    ``_build_context`` (and, for the residual-free Cross-Attention arch,
+    ``num_image_tokens``; for the non-generative MoMa arch, ``_text_only_context``).
 
-    Forward path: ``feats = vision_encoder(pixel_values)``;
-    ``img_embeds = adapter(feats)``; ``ModalityContext(prefix_embeds,
-    output_slice)``. The transformer runs over the concatenated
-    ``(image, text)`` sequence and ``output_slice`` trims the image
-    positions before the LM head.
+    Vision projection is arch-independent and runs once in ``VLMWrapper.forward``,
+    so strategies receive already-projected ``visual_embeds`` and never touch
+    pixels or the vision tower.
     """
 
     def prepare(
         self,
-        wrapper: VLMWrapper,
-        pixel_values: torch.Tensor | None,
-        input_ids: torch.Tensor,  # noqa: ARG002
+        wrapper: VLMWrapper,  # noqa: ARG002
+        visual_embeds: torch.Tensor | None,
+        input_ids: torch.Tensor,
         frame_mask: torch.Tensor | None = None,
     ) -> ModalityContext:
-        if pixel_values is None:
-            # Text-only request: no image prefix, so the residual carries text
-            # alone and the LM head runs over every position (empty context ->
-            # Transformer.forward's pure-text path).
-            return ModalityContext()
-        img_embeds = _project_visual_features(wrapper, pixel_values)
-        n = img_embeds.shape[1]  # pooling-aware: the adapter's actual visual-token count
-        return ModalityContext(
-            prefix_embeds=img_embeds,
-            output_slice=slice(n, None),
-            key_padding_mask=_prefix_key_padding_mask(frame_mask, n, input_ids),
-        )
+        if visual_embeds is None:
+            return self._text_only_context()
+        n = visual_embeds.shape[1]  # pooling-/video-aware: the actual visual-token count
+        return self._build_context(visual_embeds, n, input_ids, frame_mask)
+
+    def _text_only_context(self) -> ModalityContext:
+        # Generative arches drive the pure-text forward from an empty context.
+        return ModalityContext()
+
+    def _build_context(
+        self,
+        visual_embeds: torch.Tensor,
+        n: int,
+        input_ids: torch.Tensor,
+        frame_mask: torch.Tensor | None,
+    ) -> ModalityContext:
+        raise NotImplementedError
 
     def num_image_tokens(self, wrapper: VLMWrapper) -> int:
         return wrapper.frames_per_clip * wrapper.adapter.output_num_tokens(
@@ -210,35 +222,51 @@ def num_image_tokens(self, wrapper: VLMWrapper) -> int:
         )
 
 
+@registry.register_modality_strategy("joint_decoder")
+class JointDecoderStrategy(BaseModalityStrategy):
+    """Joint-Decoder: image embeds prepended to the text sequence.
+
+    ``_build_context`` returns ``ModalityContext(prefix_embeds, output_slice)``:
+    the transformer runs over the concatenated ``(image, text)`` sequence and
+    ``output_slice`` trims the image positions before the LM head.
+    """
+
+    def _build_context(
+        self,
+        visual_embeds: torch.Tensor,
+        n: int,
+        input_ids: torch.Tensor,
+        frame_mask: torch.Tensor | None,
+    ) -> ModalityContext:
+        return ModalityContext(
+            prefix_embeds=visual_embeds,
+            output_slice=slice(n, None),
+            key_padding_mask=_prefix_key_padding_mask(frame_mask, n, input_ids),
+        )
+
+
 @registry.register_modality_strategy("cross_attention")
-class CrossAttentionStrategy:
+class CrossAttentionStrategy(BaseModalityStrategy):
     """Cross-Attention: image embeds flow as K/V into separate
     cross-attention blocks inside the transformer; the residual stream
     itself carries text only.
 
-    Forward path: ``feats = vision_encoder(pixel_values)``;
-    ``img_embeds = adapter(feats)``; ``ModalityContext(image_features,
-    image_mask)``. ``image_mask`` carries per-visual-token validity (padded
-    video frames are masked out of the image K/V); ``None`` means all image
-    tokens are valid (e.g. a single image or a full clip).
+    ``_build_context`` returns ``ModalityContext(image_features, image_mask)``.
+    ``image_mask`` carries per-visual-token validity (padded video frames are
+    masked out of the image K/V); ``None`` means all image tokens are valid
+    (e.g. a single image or a full clip).
     """
 
-    def prepare(
+    def _build_context(
         self,
-        wrapper: VLMWrapper,
-        pixel_values: torch.Tensor | None,
+        visual_embeds: torch.Tensor,
+        n: int,
         input_ids: torch.Tensor,  # noqa: ARG002
-        frame_mask: torch.Tensor | None = None,
+        frame_mask: torch.Tensor | None,
     ) -> ModalityContext:
-        if pixel_values is None:
-            # Text-only request: no image K/V. The cross-attention blocks are
-            # skipped in Transformer.forward when image_features is None, leaving
-            # the pure text backbone.
-            return ModalityContext()
-        img_embeds = _project_visual_features(wrapper, pixel_values)
         return ModalityContext(
-            image_features=img_embeds,
-            image_mask=_visual_token_mask(frame_mask, img_embeds.shape[1]),
+            image_features=visual_embeds,
+            image_mask=_visual_token_mask(frame_mask, n),
         )
 
     def num_image_tokens(self, wrapper: VLMWrapper) -> int:  # noqa: ARG002
@@ -247,12 +275,11 @@ def num_image_tokens(self, wrapper: VLMWrapper) -> int:  # noqa: ARG002
 
 
 @registry.register_modality_strategy("mot")
-class MoTStrategy:
+class MoTStrategy(BaseModalityStrategy):
     """Mixture-of-Transformers: image-then-text residual layout (same as
     Joint-Decoder) plus a per-position ``modality_ids`` tag.
 
-    Forward path: ``feats = vision_encoder(pixel_values)``;
-    ``img_embeds = adapter(feats)``;
+    ``_build_context`` returns
     ``ModalityContext(prefix_embeds, output_slice, modality_ids)``.
 
     ``modality_ids`` is built position-based: ``0`` for the first
@@ -266,88 +293,47 @@ class MoTStrategy:
     the LM head, matching ``JointDecoderStrategy``.
     """
 
-    def prepare(
+    def _build_context(
         self,
-        wrapper: VLMWrapper,
-        pixel_values: torch.Tensor | None,
+        visual_embeds: torch.Tensor,
+        n: int,
         input_ids: torch.Tensor,
-        frame_mask: torch.Tensor | None = None,
+        frame_mask: torch.Tensor | None,
     ) -> ModalityContext:
-        if pixel_values is None:
-            # Text-only request: no image prefix. The MoT forward runs with
-            # n_image=0 (an empty image stream), routing every position through
-            # the text-modality projections/FFN; no modality_ids are needed (see
-            # Transformer.forward's MoT branch).
-            return ModalityContext()
-        img_embeds = _project_visual_features(wrapper, pixel_values)
-        n = img_embeds.shape[1]  # pooling-aware: the adapter's actual visual-token count
         b, t_text = input_ids.shape
         modality_ids = torch.zeros(b, n + t_text, dtype=torch.long, device=input_ids.device)
         modality_ids[:, n:] = 1
         return ModalityContext(
-            prefix_embeds=img_embeds,
+            prefix_embeds=visual_embeds,
             output_slice=slice(n, None),
             modality_ids=modality_ids,
             key_padding_mask=_prefix_key_padding_mask(frame_mask, n, input_ids),
         )
 
-    def num_image_tokens(self, wrapper: VLMWrapper) -> int:
-        return wrapper.frames_per_clip * wrapper.adapter.output_num_tokens(
-            wrapper.vision_encoder.num_tokens
-        )
-
 
 @registry.register_modality_strategy("moma")
-class MoMaStrategy:
-    """Mixture of Modality-Aware Experts: same residual-stream layout as
-    Joint-Decoder/MoT (image embeds prepended, ``output_slice`` trims them
-    before the LM head), plus a per-position ``modality_ids`` tag the
-    MoMa FFN stack consumes for true scatter/gather dispatch (level-1
-    deterministic routing by modality).
-
-    Forward path: ``feats = vision_encoder(pixel_values)``;
-    ``img_embeds = adapter(feats)``;
-    ``ModalityContext(prefix_embeds, output_slice, modality_ids)``.
+class MoMaStrategy(MoTStrategy):
+    """Mixture of Modality-Aware Experts: same residual-stream layout and
+    ``modality_ids`` tagging as MoT (image embeds prepended, ``output_slice``
+    trims them before the LM head), so it reuses ``MoTStrategy._build_context``.
+    The MoMa FFN stack consumes the per-position tag for true scatter/gather
+    dispatch (level-1 deterministic routing by modality).
 
     Convention: ``modality_ids == 0`` for image positions and
     ``modality_ids == 1`` for text positions, matching the index order
-    of ``MoMaConfig.moma_modalities = ("image", "text")``. The MoMa
-    FFN uses these tags to dispatch tokens to per-modality expert
-    groups; positions are *not* assumed to be in any particular order,
-    so interleaved layouts work too (image-prefix is just one
-    instantiation).
-    """
+    of ``MoMaConfig.moma_modalities = ("image", "text")``.
 
-    def prepare(
-        self,
-        wrapper: VLMWrapper,
-        pixel_values: torch.Tensor | None,
-        input_ids: torch.Tensor,
-        frame_mask: torch.Tensor | None = None,
-    ) -> ModalityContext:
-        if pixel_values is None:
-            # MoMa's expert-choice routing is non-causal (is_generative=False), so it
-            # is excluded from generative/text-only evaluation; fail fast rather than
-            # emit an unusable context.
-            raise NotImplementedError(
-                "Text-only forward is not supported for the 'moma' arch "
-                "(non-causal expert-choice routing; excluded from generative evaluation)."
-            )
-        img_embeds = _project_visual_features(wrapper, pixel_values)
-        n = img_embeds.shape[1]  # pooling-aware: the adapter's actual visual-token count
-        b, t_text = input_ids.shape
-        modality_ids = torch.zeros(b, n + t_text, dtype=torch.long, device=input_ids.device)
-        modality_ids[:, n:] = 1
-        return ModalityContext(
-            prefix_embeds=img_embeds,
-            output_slice=slice(n, None),
-            modality_ids=modality_ids,
-            key_padding_mask=_prefix_key_padding_mask(frame_mask, n, input_ids),
-        )
+    Non-generative: expert-choice routing is non-causal, so the text-only path
+    is rejected (``_text_only_context`` raises) rather than driven.
+    """
 
-    def num_image_tokens(self, wrapper: VLMWrapper) -> int:
-        return wrapper.frames_per_clip * wrapper.adapter.output_num_tokens(
-            wrapper.vision_encoder.num_tokens
+    def _text_only_context(self) -> ModalityContext:
+        # MoMa's expert-choice routing is non-causal (is_generative=False), so it
+        # is excluded from generative/text-only evaluation; fail fast rather than
+        # emit an unusable context.
+        raise NotImplementedError(
+            "Text-only forward is not supported for the 'moma' arch "
+            "(non-causal expert-choice routing; excluded from generative evaluation)."
         )
 
 
@@ -401,19 +387,43 @@ def __init__(
     def num_image_tokens(self) -> int:
         return self.strategy.num_image_tokens(self)
 
+    def encode_visual(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """Project raw pixels to LLM-dim visual embeds (vision encoder + adapter).
+
+        The eval decode loop calls this once per request and feeds the result
+        back into ``forward(..., visual_embeds=...)`` at every decode step, so
+        the vision encoder + adapter run once instead of once per generated
+        token. This is not a KV cache: the transformer still re-runs over the
+        full sequence each step.
+        """
+        return _project_visual_features(self, pixel_values)
+
     def forward(
         self,
         pixel_values: torch.Tensor | None,
         input_ids: torch.Tensor,
         labels: torch.Tensor | None = None,
         frame_mask: torch.Tensor | None = None,
+        visual_embeds: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        # Vision projection happens once here, or is reused from a precomputed
+        # ``visual_embeds`` (the inference-only cached-decode path) so the vision
+        # tower + adapter do not re-run each decode step. ``pixel_values is None``
+        # with no cache is a text-only request. Strategies receive the projected
+        # embeds and never touch pixels.
+        if visual_embeds is not None:
+            assert not self.training, "visual_embeds (cached decode) is inference-only"
+            embeds = visual_embeds
+        elif pixel_values is not None:
+            embeds = _project_visual_features(self, pixel_values)
+        else:
+            embeds = None  # text-only request
         # Route the text embedding through Transformer.forward so FSDP2's
         # per-module hook intercepts the token_embedding call and
         # materializes the DTensor weight before F.embedding runs. Doing
         # the embedding externally (transformer.token_embedding(input_ids))
         # bypasses FSDP and fails with "mixed torch.Tensor and DTensor".
-        modality = self.strategy.prepare(self, pixel_values, input_ids, frame_mask=frame_mask)
+        modality = self.strategy.prepare(self, embeds, input_ids, frame_mask=frame_mask)
         logits = self.transformer(tokens=input_ids, modality=modality)
         return logits, labels
 
diff --git a/tests/unit/test_moma.py b/tests/unit/test_moma.py
index 46b300e..2d7a10a 100644
--- a/tests/unit/test_moma.py
+++ b/tests/unit/test_moma.py
@@ -240,10 +240,12 @@ class TestMoMaStrategy:
     def test_prepare_builds_modality_context(self):
         wrapper = _StubWrapper(num_tokens=4, feature_dim=8, dim=16)
         strategy = MoMaStrategy()
-        pixel_values = torch.zeros(2, 3, 8, 8)
+        # prepare() now receives already-projected visual embeds (projection is
+        # hoisted into VLMWrapper.forward); pass (B, N, dim) directly.
+        visual_embeds = torch.zeros(2, 4, 16)
         input_ids = torch.zeros(2, 6, dtype=torch.long)
 
-        ctx = strategy.prepare(wrapper, pixel_values, input_ids)
+        ctx = strategy.prepare(wrapper, visual_embeds, input_ids)
         assert isinstance(ctx, ModalityContext)
         assert ctx.prefix_embeds is not None
         assert ctx.prefix_embeds.shape == (2, 4, 16)
@@ -255,10 +257,11 @@ def test_prepare_builds_modality_context(self):
     def test_modality_ids_image_then_text(self):
         wrapper = _StubWrapper(num_tokens=3, feature_dim=8, dim=16)
         strategy = MoMaStrategy()
-        pixel_values = torch.zeros(1, 3, 8, 8)
+        # Already-projected embeds: 3 visual tokens.
+        visual_embeds = torch.zeros(1, 3, 16)
         input_ids = torch.zeros(1, 5, dtype=torch.long)
 
-        ctx = strategy.prepare(wrapper, pixel_values, input_ids)
+        ctx = strategy.prepare(wrapper, visual_embeds, input_ids)
         # First 3 positions (image) get 0; rest (text) get 1.
         assert ctx.modality_ids is not None
         ids = ctx.modality_ids[0]
diff --git a/tests/unit/test_vlm.py b/tests/unit/test_vlm.py
index 6556fa7..5a916e3 100644
--- a/tests/unit/test_vlm.py
+++ b/tests/unit/test_vlm.py
@@ -332,7 +332,7 @@ def test_joint_decoder_strategy_fills_prefix_and_slice(self):
         strategy = JointDecoderStrategy()
         pixel_values = torch.randn(1, 3, 64, 64)
         input_ids = torch.randint(0, 256, (1, 16))
-        ctx = strategy.prepare(wrapper, pixel_values, input_ids)
+        ctx = strategy.prepare(wrapper, wrapper.encode_visual(pixel_values), input_ids)
         assert ctx.prefix_embeds is not None
         assert ctx.prefix_embeds.shape == (1, 8, 64)  # (B, N, dim)
         assert ctx.output_slice == slice(8, None)
@@ -349,7 +349,7 @@ def test_mot_strategy_fills_prefix_slice_and_modality_ids(self):
         strategy = MoTStrategy()
         pixel_values = torch.randn(1, 3, 64, 64)
         input_ids = torch.randint(0, 256, (1, 16))
-        ctx = strategy.prepare(wrapper, pixel_values, input_ids)
+        ctx = strategy.prepare(wrapper, wrapper.encode_visual(pixel_values), input_ids)
         assert ctx.prefix_embeds is not None
         assert ctx.prefix_embeds.shape == (1, 8, 64)
         assert ctx.output_slice == slice(8, None)
@@ -367,7 +367,7 @@ def test_mot_modality_ids_shape_dtype_device(self, t_text: int):
         strategy = MoTStrategy()
         pixel_values = torch.randn(2, 3, 64, 64, device=DEVICE)
         input_ids = torch.randint(0, 256, (2, t_text), device=DEVICE)
-        ctx = strategy.prepare(wrapper, pixel_values, input_ids)
+        ctx = strategy.prepare(wrapper, wrapper.encode_visual(pixel_values), input_ids)
         assert ctx.modality_ids is not None
         assert ctx.modality_ids.shape == (2, 8 + t_text)
         assert ctx.modality_ids.dtype == torch.long
@@ -378,7 +378,7 @@ def test_cross_attention_strategy_fills_image_features(self):
         strategy = CrossAttentionStrategy()
         pixel_values = torch.randn(1, 3, 64, 64)
         input_ids = torch.randint(0, 256, (1, 16))
-        ctx = strategy.prepare(wrapper, pixel_values, input_ids)
+        ctx = strategy.prepare(wrapper, wrapper.encode_visual(pixel_values), input_ids)
         assert ctx.image_features is not None
         assert ctx.image_features.shape == (1, 8, 64)
         assert ctx.image_mask is None
@@ -451,7 +451,7 @@ def test_dispatch_no_isinstance_ladder(self):
 
         @registry.register_modality_strategy("dispatch_smoke_test_arch")
         class _Smoke:
-            def prepare(self, wrapper, pixel_values, input_ids):  # noqa: ARG002
+            def prepare(self, wrapper, visual_embeds, input_ids, frame_mask=None):  # noqa: ARG002
                 return ModalityContext()
 
             def num_image_tokens(self, wrapper):  # noqa: ARG002
@@ -562,7 +562,7 @@ def test_jd_prefix_length_is_pooled(self, adapter_type):
         strategy = JointDecoderStrategy()
         pixels = torch.randn(2, 3, 16, 16)
         input_ids = torch.randint(0, 256, (2, 12))
-        ctx = strategy.prepare(wrapper, pixels, input_ids)
+        ctx = strategy.prepare(wrapper, wrapper.encode_visual(pixels), input_ids)
         assert ctx.prefix_embeds is not None
         assert ctx.prefix_embeds.shape == (2, 4, 64)  # pooled prefix, model dim
         assert ctx.output_slice == slice(4, None)
@@ -585,7 +585,7 @@ def test_mot_split_uses_pooled_count(self):
         strategy = MoTStrategy()
         pixels = torch.randn(2, 3, 16, 16, device=DEVICE)
         input_ids = torch.randint(0, 256, (2, 16), device=DEVICE)
-        ctx = strategy.prepare(wrapper, pixels, input_ids)
+        ctx = strategy.prepare(wrapper, wrapper.encode_visual(pixels), input_ids)
         assert ctx.modality_ids is not None
         assert ctx.modality_ids.shape == (2, 4 + 16)  # pooled prefix + text
         assert (ctx.modality_ids[:, :4] == 0).all()
@@ -636,23 +636,18 @@ def test_num_image_tokens_is_frames_times_per_frame(self):
         assert wrapper.num_image_tokens == 16
 
     def test_projector_folds_frame_axis(self):
+        # Frame-axis folding now lives in encode_visual (VLMWrapper.forward hoists
+        # projection out of the strategies); assert on its output directly.
         wrapper = _video_wrapper(JointDecoderConfig(max_text_len=8), frames=4)
-        ctx = JointDecoderStrategy().prepare(
-            wrapper, torch.randn(2, 4, 3, 16, 16), torch.randint(0, 256, (2, 6))
-        )
-        assert ctx.prefix_embeds is not None
-        assert ctx.prefix_embeds.shape == (2, 16, 64)  # (B, F*P', dim)
-        assert ctx.output_slice == slice(16, None)
+        embeds = wrapper.encode_visual(torch.randn(2, 4, 3, 16, 16))
+        assert embeds.shape == (2, 16, 64)  # (B, F*P', dim)
 
     def test_static_count_matches_runtime_prefix(self):
         """MoT's positional split uses the build-time count; it must equal the
-        runtime prefix length (frames * per-frame)."""
+        runtime visual-token count (frames * per-frame) from encode_visual."""
         wrapper = _video_wrapper(JointDecoderConfig(max_text_len=8), frames=4)
-        ctx = JointDecoderStrategy().prepare(
-            wrapper, torch.randn(1, 4, 3, 16, 16), torch.randint(0, 256, (1, 6))
-        )
-        assert ctx.prefix_embeds is not None
-        assert ctx.prefix_embeds.shape[1] == wrapper.num_image_tokens == 16
+        embeds = wrapper.encode_visual(torch.randn(1, 4, 3, 16, 16))
+        assert embeds.shape[1] == wrapper.num_image_tokens == 16
 
     @pytest.mark.parametrize("arch", ["joint_decoder", "cross_attention", "mot", "moma"])
     def test_video_forward_all_archs(self, arch):
@@ -796,3 +791,71 @@ def test_undecodable_clip_stays_finite(self, arch):
         with torch.no_grad():
             logits, _ = w(pix, ids, frame_mask=fm)
         assert torch.isfinite(logits).all(), f"{arch}: NaN/inf with an all-padded clip"
+
+
+class TestVisualEmbedCache:
+    """The eval decode loop projects visual features once via ``encode_visual``
+    and feeds them back through ``forward(..., visual_embeds=...)``. Projection is
+    deterministic, so the cached path reproduces the uncached forward
+    bit-for-bit; the cache is inference-only (rejected in training mode)."""
+
+    def test_encode_visual_is_deterministic(self):
+        # Premise of the equivalence tests: RandomVisionEncoder seeds from the
+        # input and the adapter has fixed weights, so encode_visual is reproducible.
+        wrapper = _build_tiny_wrapper().to(DEVICE).eval()
+        pixels = torch.randn(2, 3, 16, 16, device=DEVICE)
+        with torch.no_grad():
+            assert torch.equal(wrapper.encode_visual(pixels), wrapper.encode_visual(pixels))
+
+    @pytest.mark.parametrize("arch", ["joint_decoder", "cross_attention", "mot", "moma"])
+    def test_cached_visual_embeds_equal_uncached_image(self, arch):
+        builders = {
+            "joint_decoder": _build_tiny_wrapper,
+            "cross_attention": _build_ca_tiny_wrapper,
+            "mot": _build_mot_tiny_wrapper,
+            "moma": _build_moma_tiny_wrapper,
+        }
+        wrapper = builders[arch](num_image_tokens=8).to(DEVICE).eval()
+        pixels = torch.randn(2, 3, 16, 16, device=DEVICE)
+        input_ids = torch.randint(0, 256, (2, 12), device=DEVICE)
+        with torch.no_grad():
+            ve = wrapper.encode_visual(pixels)
+            cached, _ = wrapper(pixels, input_ids, visual_embeds=ve)
+            uncached, _ = wrapper(pixels, input_ids)
+        assert torch.equal(cached, uncached), f"{arch}: cached decode diverges from projection"
+
+    @pytest.mark.parametrize("arch", ["joint_decoder", "cross_attention", "mot", "moma"])
+    def test_cached_visual_embeds_equal_uncached_video(self, arch):
+        # Video + frame_mask: masks are built from n == visual_embeds.shape[1],
+        # which caching does not change, so masking matches the uncached path.
+        ffn = 128 if arch in ("mot", "moma") else None
+        cfgs = {
+            "joint_decoder": JointDecoderConfig(max_text_len=8),
+            "cross_attention": CrossAttentionConfig(
+                max_text_len=8, cross_attention_every_n_layers=2
+            ),
+            "mot": MoTConfig(max_text_len=8),
+            "moma": MoMaConfig(max_text_len=8),
+        }
+        wrapper = _video_wrapper(cfgs[arch], frames=4, ffn_hidden_dim=ffn).to(DEVICE).eval()
+        pixels = torch.randn(2, 4, 3, 16, 16, device=DEVICE)
+        input_ids = torch.randint(0, 256, (2, 6), device=DEVICE)
+        fm = torch.tensor([[True, True, False, False], [True, True, True, True]], device=DEVICE)
+        with torch.no_grad():
+            ve = wrapper.encode_visual(pixels)
+            cached, _ = wrapper(pixels, input_ids, frame_mask=fm, visual_embeds=ve)
+            uncached, _ = wrapper(pixels, input_ids, frame_mask=fm)
+        assert torch.equal(cached, uncached), f"{arch}: cached video decode diverges"
+
+    def test_visual_embeds_in_training_mode_asserts(self):
+        # The cache path skips the vision tower, so it must never run in training
+        # (it would silently detach the encoder from the graph). Guarded by an
+        # assert, which python -O strips.
+        if not __debug__:
+            pytest.skip("assert guard is a no-op under -O")
+        wrapper = _build_tiny_wrapper().to(DEVICE).train()
+        pixels = torch.randn(1, 3, 16, 16, device=DEVICE)
+        input_ids = torch.randint(0, 256, (1, 8), device=DEVICE)
+        ve = wrapper.encode_visual(pixels)
+        with pytest.raises(AssertionError):
+            wrapper(pixels, input_ids, visual_embeds=ve)

From eed5dd0d221f13aacf1ecf61fc8fa10be39df656 Mon Sep 17 00:00:00 2001
From: camilobrownpinilla <camilobrownpinilla@gmail.com>
Date: Thu, 23 Jul 2026 15:54:00 -0400
Subject: [PATCH 2/3] Trim verbose comments and docstrings

---
 kempnerforge/model/vlm.py | 64 +++++++++++----------------------------
 tests/unit/test_moma.py   |  3 +-
 tests/unit/test_vlm.py    | 19 ++++--------
 3 files changed, 24 insertions(+), 62 deletions(-)

diff --git a/kempnerforge/model/vlm.py b/kempnerforge/model/vlm.py
index 6f9d8ed..f7ca394 100644
--- a/kempnerforge/model/vlm.py
+++ b/kempnerforge/model/vlm.py
@@ -72,14 +72,8 @@ def prepare(
     ) -> ModalityContext:
         """Compose a ``ModalityContext`` from already-projected visual embeds.
 
-        Vision projection (encoder + adapter) is arch-independent and runs once
-        in ``VLMWrapper.forward`` (or ``VLMWrapper.encode_visual`` for cached
-        decode), so strategies receive ``visual_embeds`` and never touch pixels
-        or the vision tower. ``visual_embeds is None`` is a text-only request
-        (no visual content): the generative arches return a context that drives
-        the pure-text forward — an empty ``ModalityContext()`` for the
-        image-prefix and cross-attention arches — while a non-generative arch
-        may reject it.
+        Projection is hoisted into ``VLMWrapper.forward``, so strategies get
+        ``visual_embeds`` (``None`` for text-only) and never touch pixels.
         """
         ...
 
@@ -178,17 +172,10 @@ def _prefix_key_padding_mask(
 
 
 class BaseModalityStrategy:
-    """Shared ``ModalityStrategy`` implementation (not registered, not the Protocol).
-
-    ``prepare`` is a template method: it handles the text-only request and the
-    visual-token count ``n``, then delegates arch-specific ``ModalityContext``
-    construction to ``_build_context``. Concrete strategies override
-    ``_build_context`` (and, for the residual-free Cross-Attention arch,
-    ``num_image_tokens``; for the non-generative MoMa arch, ``_text_only_context``).
-
-    Vision projection is arch-independent and runs once in ``VLMWrapper.forward``,
-    so strategies receive already-projected ``visual_embeds`` and never touch
-    pixels or the vision tower.
+    """Shared ``ModalityStrategy`` base: ``prepare`` handles the text-only case
+    and the visual-token count ``n``, then defers to ``_build_context``. Concrete
+    strategies override ``_build_context`` (Cross-Attention also ``num_image_tokens``;
+    MoMa also ``_text_only_context``).
     """
 
     def prepare(
@@ -200,11 +187,11 @@ def prepare(
     ) -> ModalityContext:
         if visual_embeds is None:
             return self._text_only_context()
-        n = visual_embeds.shape[1]  # pooling-/video-aware: the actual visual-token count
+        n = visual_embeds.shape[1]  # actual visual-token count
         return self._build_context(visual_embeds, n, input_ids, frame_mask)
 
     def _text_only_context(self) -> ModalityContext:
-        # Generative arches drive the pure-text forward from an empty context.
+        # empty context -> pure-text forward
         return ModalityContext()
 
     def _build_context(
@@ -313,24 +300,13 @@ def _build_context(
 
 @registry.register_modality_strategy("moma")
 class MoMaStrategy(MoTStrategy):
-    """Mixture of Modality-Aware Experts: same residual-stream layout and
-    ``modality_ids`` tagging as MoT (image embeds prepended, ``output_slice``
-    trims them before the LM head), so it reuses ``MoTStrategy._build_context``.
-    The MoMa FFN stack consumes the per-position tag for true scatter/gather
-    dispatch (level-1 deterministic routing by modality).
-
-    Convention: ``modality_ids == 0`` for image positions and
-    ``modality_ids == 1`` for text positions, matching the index order
-    of ``MoMaConfig.moma_modalities = ("image", "text")``.
-
-    Non-generative: expert-choice routing is non-causal, so the text-only path
-    is rejected (``_text_only_context`` raises) rather than driven.
+    """Mixture of Modality-Aware Experts: same residual layout and ``modality_ids``
+    tagging as MoT. Non-generative (expert-choice routing is non-causal), so the
+    text-only path is rejected.
     """
 
     def _text_only_context(self) -> ModalityContext:
-        # MoMa's expert-choice routing is non-causal (is_generative=False), so it
-        # is excluded from generative/text-only evaluation; fail fast rather than
-        # emit an unusable context.
+        # Non-causal expert-choice routing (is_generative=False): no text-only path.
         raise NotImplementedError(
             "Text-only forward is not supported for the 'moma' arch "
             "(non-causal expert-choice routing; excluded from generative evaluation)."
@@ -388,13 +364,10 @@ def num_image_tokens(self) -> int:
         return self.strategy.num_image_tokens(self)
 
     def encode_visual(self, pixel_values: torch.Tensor) -> torch.Tensor:
-        """Project raw pixels to LLM-dim visual embeds (vision encoder + adapter).
+        """Project pixels to LLM-dim visual embeds once, for cached decode.
 
-        The eval decode loop calls this once per request and feeds the result
-        back into ``forward(..., visual_embeds=...)`` at every decode step, so
-        the vision encoder + adapter run once instead of once per generated
-        token. This is not a KV cache: the transformer still re-runs over the
-        full sequence each step.
+        The decode loop passes the result to ``forward(..., visual_embeds=...)``
+        so the vision tower runs once per request, not once per token.
         """
         return _project_visual_features(self, pixel_values)
 
@@ -406,11 +379,8 @@ def forward(
         frame_mask: torch.Tensor | None = None,
         visual_embeds: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        # Vision projection happens once here, or is reused from a precomputed
-        # ``visual_embeds`` (the inference-only cached-decode path) so the vision
-        # tower + adapter do not re-run each decode step. ``pixel_values is None``
-        # with no cache is a text-only request. Strategies receive the projected
-        # embeds and never touch pixels.
+        # Project once here, or reuse precomputed ``visual_embeds`` (cached decode,
+        # inference-only). ``pixel_values is None`` with no cache is text-only.
         if visual_embeds is not None:
             assert not self.training, "visual_embeds (cached decode) is inference-only"
             embeds = visual_embeds
diff --git a/tests/unit/test_moma.py b/tests/unit/test_moma.py
index 2d7a10a..484d727 100644
--- a/tests/unit/test_moma.py
+++ b/tests/unit/test_moma.py
@@ -240,8 +240,7 @@ class TestMoMaStrategy:
     def test_prepare_builds_modality_context(self):
         wrapper = _StubWrapper(num_tokens=4, feature_dim=8, dim=16)
         strategy = MoMaStrategy()
-        # prepare() now receives already-projected visual embeds (projection is
-        # hoisted into VLMWrapper.forward); pass (B, N, dim) directly.
+        # prepare() takes already-projected embeds; pass (B, N, dim) directly.
         visual_embeds = torch.zeros(2, 4, 16)
         input_ids = torch.zeros(2, 6, dtype=torch.long)
 
diff --git a/tests/unit/test_vlm.py b/tests/unit/test_vlm.py
index 5a916e3..20681ee 100644
--- a/tests/unit/test_vlm.py
+++ b/tests/unit/test_vlm.py
@@ -636,8 +636,7 @@ def test_num_image_tokens_is_frames_times_per_frame(self):
         assert wrapper.num_image_tokens == 16
 
     def test_projector_folds_frame_axis(self):
-        # Frame-axis folding now lives in encode_visual (VLMWrapper.forward hoists
-        # projection out of the strategies); assert on its output directly.
+        # Frame-axis folding now lives in encode_visual.
         wrapper = _video_wrapper(JointDecoderConfig(max_text_len=8), frames=4)
         embeds = wrapper.encode_visual(torch.randn(2, 4, 3, 16, 16))
         assert embeds.shape == (2, 16, 64)  # (B, F*P', dim)
@@ -794,14 +793,11 @@ def test_undecodable_clip_stays_finite(self, arch):
 
 
 class TestVisualEmbedCache:
-    """The eval decode loop projects visual features once via ``encode_visual``
-    and feeds them back through ``forward(..., visual_embeds=...)``. Projection is
-    deterministic, so the cached path reproduces the uncached forward
-    bit-for-bit; the cache is inference-only (rejected in training mode)."""
+    """encode_visual projects once; feeding the result back via ``visual_embeds``
+    reproduces the uncached forward bit-for-bit (projection is deterministic)."""
 
     def test_encode_visual_is_deterministic(self):
-        # Premise of the equivalence tests: RandomVisionEncoder seeds from the
-        # input and the adapter has fixed weights, so encode_visual is reproducible.
+        # RandomVisionEncoder seeds from the input -> encode_visual is reproducible.
         wrapper = _build_tiny_wrapper().to(DEVICE).eval()
         pixels = torch.randn(2, 3, 16, 16, device=DEVICE)
         with torch.no_grad():
@@ -826,8 +822,7 @@ def test_cached_visual_embeds_equal_uncached_image(self, arch):
 
     @pytest.mark.parametrize("arch", ["joint_decoder", "cross_attention", "mot", "moma"])
     def test_cached_visual_embeds_equal_uncached_video(self, arch):
-        # Video + frame_mask: masks are built from n == visual_embeds.shape[1],
-        # which caching does not change, so masking matches the uncached path.
+        # Video + frame_mask: caching doesn't change n, so masking is identical.
         ffn = 128 if arch in ("mot", "moma") else None
         cfgs = {
             "joint_decoder": JointDecoderConfig(max_text_len=8),
@@ -848,9 +843,7 @@ def test_cached_visual_embeds_equal_uncached_video(self, arch):
         assert torch.equal(cached, uncached), f"{arch}: cached video decode diverges"
 
     def test_visual_embeds_in_training_mode_asserts(self):
-        # The cache path skips the vision tower, so it must never run in training
-        # (it would silently detach the encoder from the graph). Guarded by an
-        # assert, which python -O strips.
+        # Cache path is inference-only; guard is an assert (stripped under -O).
         if not __debug__:
             pytest.skip("assert guard is a no-op under -O")
         wrapper = _build_tiny_wrapper().to(DEVICE).train()

From 5478c47299fdae3b231775c090f0cc680aa372a5 Mon Sep 17 00:00:00 2001
From: camilobrownpinilla <camilobrownpinilla@gmail.com>
Date: Thu, 23 Jul 2026 17:41:21 -0400
Subject: [PATCH 3/3] Validate cached visual embeds and restore falsifiable VLM
 tests

Reject visual_embeds passed alongside pixel_values, check rank/batch/
token-count/dim, and swap the training-mode assert for a ValueError so
it survives -O.

Tests: cached calls now pass pixel_values=None, wrappers get their
identity-init projections perturbed (CA and MoT are image-blind at
construction), dispatch checks use exact types, and the MoMa strategy
tests run through the real projection chain.
---
 kempnerforge/model/vlm.py |  69 ++++++++++++++++--
 tests/unit/test_moma.py   |  13 ++--
 tests/unit/test_vlm.py    | 144 +++++++++++++++++++++++++++++++++-----
 3 files changed, 201 insertions(+), 25 deletions(-)

diff --git a/kempnerforge/model/vlm.py b/kempnerforge/model/vlm.py
index f7ca394..6f8762c 100644
--- a/kempnerforge/model/vlm.py
+++ b/kempnerforge/model/vlm.py
@@ -171,6 +171,56 @@ def _prefix_key_padding_mask(
     return torch.cat([vmask, text_valid], dim=1)
 
 
+def _expected_visual_tokens(wrapper: VLMWrapper) -> int:
+    """Static visual-token count: ``frames_per_clip`` * per-frame adapter output.
+
+    This is what ``_project_visual_features`` always produces (the frames check
+    there pins the frame count), so it is both the residual budget and the
+    expected ``n`` for a caller-supplied cache.
+    """
+    return wrapper.frames_per_clip * wrapper.adapter.output_num_tokens(
+        wrapper.vision_encoder.num_tokens
+    )
+
+
+def _validate_cached_visual_embeds(
+    wrapper: VLMWrapper, visual_embeds: torch.Tensor, input_ids: torch.Tensor
+) -> None:
+    """Re-establish, for caller-supplied embeds, the invariants the uncached path
+    gets for free from ``_project_visual_features``.
+
+    ``n = visual_embeds.shape[1]`` drives ``output_slice``, MoT's positional split
+    and the residual budget, so a malformed cache would otherwise surface as an
+    opaque shape error inside ``Transformer.forward`` (or, for a 2-D tensor, as a
+    plausible-looking ``n`` taken from the feature dim).
+    """
+    if visual_embeds.dim() != 3:
+        raise ValueError(
+            f"visual_embeds must be (B, N, dim); received a {visual_embeds.dim()}-D tensor "
+            f"of shape {tuple(visual_embeds.shape)}"
+        )
+    b, n, dim = visual_embeds.shape
+    if b != input_ids.shape[0]:
+        raise ValueError(
+            f"visual_embeds batch ({b}) does not match input_ids batch "
+            f"({input_ids.shape[0]}); a cache built for one batch cannot be reused "
+            "against a differently-sized decode batch"
+        )
+    expected_n = _expected_visual_tokens(wrapper)
+    if n != expected_n:
+        raise ValueError(
+            f"visual_embeds has {n} visual token(s) but this wrapper projects to "
+            f"{expected_n} (frames_per_clip={wrapper.frames_per_clip}); the token count "
+            "sets output_slice and the residual budget, so it must match"
+        )
+    expected_dim = wrapper.transformer.config.dim
+    if dim != expected_dim:
+        raise ValueError(
+            f"visual_embeds has feature dim {dim} but the transformer expects "
+            f"{expected_dim}; embeds must already be projected to the LLM dim"
+        )
+
+
 class BaseModalityStrategy:
     """Shared ``ModalityStrategy`` base: ``prepare`` handles the text-only case
     and the visual-token count ``n``, then defers to ``_build_context``. Concrete
@@ -204,9 +254,7 @@ def _build_context(
         raise NotImplementedError
 
     def num_image_tokens(self, wrapper: VLMWrapper) -> int:
-        return wrapper.frames_per_clip * wrapper.adapter.output_num_tokens(
-            wrapper.vision_encoder.num_tokens
-        )
+        return _expected_visual_tokens(wrapper)
 
 
 @registry.register_modality_strategy("joint_decoder")
@@ -381,8 +429,21 @@ def forward(
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         # Project once here, or reuse precomputed ``visual_embeds`` (cached decode,
         # inference-only). ``pixel_values is None`` with no cache is text-only.
+        # The cache bypasses ``_project_visual_features``, so its invariants are
+        # re-checked here rather than inherited.
         if visual_embeds is not None:
-            assert not self.training, "visual_embeds (cached decode) is inference-only"
+            if self.training:
+                raise ValueError(
+                    "visual_embeds (cached decode) is inference-only: it bypasses the "
+                    "vision encoder + adapter, which would silently receive no gradient"
+                )
+            if pixel_values is not None:
+                raise ValueError(
+                    "pass pixel_values or visual_embeds, not both; cached embeds supersede "
+                    "pixels, so the pixels would be silently ignored (and may not even "
+                    "correspond to the cached embeds)"
+                )
+            _validate_cached_visual_embeds(self, visual_embeds, input_ids)
             embeds = visual_embeds
         elif pixel_values is not None:
             embeds = _project_visual_features(self, pixel_values)
diff --git a/tests/unit/test_moma.py b/tests/unit/test_moma.py
index 484d727..f492aa7 100644
--- a/tests/unit/test_moma.py
+++ b/tests/unit/test_moma.py
@@ -42,7 +42,7 @@
     MoMaFFN,
 )
 from kempnerforge.model.transformer import Transformer
-from kempnerforge.model.vlm import MoMaStrategy
+from kempnerforge.model.vlm import MoMaStrategy, _project_visual_features
 
 DEVICE = torch.device("cpu")
 
@@ -240,8 +240,11 @@ class TestMoMaStrategy:
     def test_prepare_builds_modality_context(self):
         wrapper = _StubWrapper(num_tokens=4, feature_dim=8, dim=16)
         strategy = MoMaStrategy()
-        # prepare() takes already-projected embeds; pass (B, N, dim) directly.
-        visual_embeds = torch.zeros(2, 4, 16)
+        # prepare() takes already-projected embeds. Route pixels through the real
+        # projection helper so the encoder -> adapter chain stays under test: the
+        # (2, 4, 16) assertions below then pin that the visual-token count and the
+        # model dim survive projection, rather than restating a hand-built shape.
+        visual_embeds = _project_visual_features(wrapper, torch.zeros(2, 3, 16, 16))
         input_ids = torch.zeros(2, 6, dtype=torch.long)
 
         ctx = strategy.prepare(wrapper, visual_embeds, input_ids)
@@ -256,8 +259,8 @@ def test_prepare_builds_modality_context(self):
     def test_modality_ids_image_then_text(self):
         wrapper = _StubWrapper(num_tokens=3, feature_dim=8, dim=16)
         strategy = MoMaStrategy()
-        # Already-projected embeds: 3 visual tokens.
-        visual_embeds = torch.zeros(1, 3, 16)
+        # Projected from pixels: 3 visual tokens (stub adapter is token-count identity).
+        visual_embeds = _project_visual_features(wrapper, torch.zeros(1, 3, 16, 16))
         input_ids = torch.zeros(1, 5, dtype=torch.long)
 
         ctx = strategy.prepare(wrapper, visual_embeds, input_ids)
diff --git a/tests/unit/test_vlm.py b/tests/unit/test_vlm.py
index 20681ee..a187903 100644
--- a/tests/unit/test_vlm.py
+++ b/tests/unit/test_vlm.py
@@ -8,6 +8,7 @@
 
 import pytest
 import torch
+import torch.nn as nn
 
 from kempnerforge.config.adapter import AdapterConfig
 from kempnerforge.config.model import ModelConfig
@@ -441,7 +442,14 @@ def test_build_modality_strategy_cross_attention(self):
     def test_build_modality_strategy_mot(self):
         cfg = MoTConfig()
         strategy = build_modality_strategy(cfg)
-        assert isinstance(strategy, MoTStrategy)
+        # Exact type, not isinstance: MoMaStrategy subclasses MoTStrategy, so an
+        # isinstance check would accept a registry that maps "mot" -> MoMaStrategy.
+        assert type(strategy) is MoTStrategy
+
+    def test_build_modality_strategy_moma(self):
+        cfg = MoMaConfig()
+        strategy = build_modality_strategy(cfg)
+        assert type(strategy) is MoMaStrategy
 
     def test_dispatch_no_isinstance_ladder(self):
         """Sanity check: build_modality_strategy is a pure registry
@@ -474,7 +482,7 @@ def test_ca_wrapper_uses_ca_strategy(self):
 
     def test_mot_wrapper_uses_mot_strategy(self):
         wrapper = _build_mot_tiny_wrapper()
-        assert isinstance(wrapper.strategy, MoTStrategy)
+        assert type(wrapper.strategy) is MoTStrategy  # exact: MoMaStrategy is a MoTStrategy
 
     def test_mot_forward_logits_text_only_shape(self):
         """MoT VLMWrapper forward returns text-only logits (output_slice
@@ -642,11 +650,16 @@ def test_projector_folds_frame_axis(self):
         assert embeds.shape == (2, 16, 64)  # (B, F*P', dim)
 
     def test_static_count_matches_runtime_prefix(self):
-        """MoT's positional split uses the build-time count; it must equal the
-        runtime visual-token count (frames * per-frame) from encode_visual."""
+        """MoT's positional split and the residual budget use the build-time count;
+        it must equal what the strategy actually puts on the residual, and the
+        output_slice must start there."""
         wrapper = _video_wrapper(JointDecoderConfig(max_text_len=8), frames=4)
+        input_ids = torch.randint(0, 256, (1, 6))
         embeds = wrapper.encode_visual(torch.randn(1, 4, 3, 16, 16))
-        assert embeds.shape[1] == wrapper.num_image_tokens == 16
+        ctx = wrapper.strategy.prepare(wrapper, embeds, input_ids)
+        assert ctx.prefix_embeds is not None
+        assert ctx.prefix_embeds.shape[1] == wrapper.num_image_tokens == 16
+        assert ctx.output_slice == slice(wrapper.num_image_tokens, None)
 
     @pytest.mark.parametrize("arch", ["joint_decoder", "cross_attention", "mot", "moma"])
     def test_video_forward_all_archs(self, arch):
@@ -792,9 +805,32 @@ def test_undecodable_clip_stays_finite(self, arch):
         assert torch.isfinite(logits).all(), f"{arch}: NaN/inf with an all-padded clip"
 
 
+def _break_identity_init(wrapper: VLMWrapper) -> VLMWrapper:
+    """Randomize zero-initialized projections so the image can reach the logits.
+
+    Cross-Attention (``cross_attention.py``: ``zeros_(o_proj.weight)``,
+    ``zeros_(mlp.down_proj.weight)``) and MoT (``init.py``: per-modality residual
+    projections) are deliberately identity-at-construction. A freshly built
+    wrapper of either arch is therefore image-blind -- two completely different
+    images give bit-identical logits -- so a cached-vs-uncached logits comparison
+    on one passes no matter what forward does with visual_embeds. Perturbing the
+    zeroed weights gives those comparisons something to actually detect.
+    """
+    with torch.no_grad():
+        for module in wrapper.modules():
+            if isinstance(module, nn.Linear) and not module.weight.any():
+                module.weight.normal_(0.0, 0.02)
+    return wrapper
+
+
 class TestVisualEmbedCache:
     """encode_visual projects once; feeding the result back via ``visual_embeds``
-    reproduces the uncached forward bit-for-bit (projection is deterministic)."""
+    reproduces the uncached forward bit-for-bit (projection is deterministic).
+
+    The cached call passes ``pixel_values=None``: the wrapper rejects both
+    together, so these tests fail if forward ever stops honouring the cache.
+    Wrappers go through ``_break_identity_init`` first, without which the CA and
+    MoT cases would be blind to the image entirely."""
 
     def test_encode_visual_is_deterministic(self):
         # RandomVisionEncoder seeds from the input -> encode_visual is reproducible.
@@ -811,12 +847,12 @@ def test_cached_visual_embeds_equal_uncached_image(self, arch):
             "mot": _build_mot_tiny_wrapper,
             "moma": _build_moma_tiny_wrapper,
         }
-        wrapper = builders[arch](num_image_tokens=8).to(DEVICE).eval()
+        wrapper = _break_identity_init(builders[arch](num_image_tokens=8)).to(DEVICE).eval()
         pixels = torch.randn(2, 3, 16, 16, device=DEVICE)
         input_ids = torch.randint(0, 256, (2, 12), device=DEVICE)
         with torch.no_grad():
             ve = wrapper.encode_visual(pixels)
-            cached, _ = wrapper(pixels, input_ids, visual_embeds=ve)
+            cached, _ = wrapper(None, input_ids, visual_embeds=ve)
             uncached, _ = wrapper(pixels, input_ids)
         assert torch.equal(cached, uncached), f"{arch}: cached decode diverges from projection"
 
@@ -832,23 +868,99 @@ def test_cached_visual_embeds_equal_uncached_video(self, arch):
             "mot": MoTConfig(max_text_len=8),
             "moma": MoMaConfig(max_text_len=8),
         }
-        wrapper = _video_wrapper(cfgs[arch], frames=4, ffn_hidden_dim=ffn).to(DEVICE).eval()
+        wrapper = (
+            _break_identity_init(_video_wrapper(cfgs[arch], frames=4, ffn_hidden_dim=ffn))
+            .to(DEVICE)
+            .eval()
+        )
         pixels = torch.randn(2, 4, 3, 16, 16, device=DEVICE)
         input_ids = torch.randint(0, 256, (2, 6), device=DEVICE)
         fm = torch.tensor([[True, True, False, False], [True, True, True, True]], device=DEVICE)
         with torch.no_grad():
             ve = wrapper.encode_visual(pixels)
-            cached, _ = wrapper(pixels, input_ids, frame_mask=fm, visual_embeds=ve)
+            cached, _ = wrapper(None, input_ids, frame_mask=fm, visual_embeds=ve)
             uncached, _ = wrapper(pixels, input_ids, frame_mask=fm)
         assert torch.equal(cached, uncached), f"{arch}: cached video decode diverges"
 
-    def test_visual_embeds_in_training_mode_asserts(self):
-        # Cache path is inference-only; guard is an assert (stripped under -O).
-        if not __debug__:
-            pytest.skip("assert guard is a no-op under -O")
+    def test_cached_decode_skips_vision_encoder(self, monkeypatch):
+        """The point of the cache: the vision tower runs once per request, not
+        once per decode step. Count encoder calls rather than compare outputs."""
+        wrapper = _build_tiny_wrapper(num_image_tokens=8).to(DEVICE).eval()
+        pixels = torch.randn(2, 3, 16, 16, device=DEVICE)
+        input_ids = torch.randint(0, 256, (2, 12), device=DEVICE)
+        with torch.no_grad():
+            ve = wrapper.encode_visual(pixels)
+
+        calls = 0
+        inner = wrapper.vision_encoder.forward
+
+        def counting_forward(px):
+            nonlocal calls
+            calls += 1
+            return inner(px)
+
+        monkeypatch.setattr(wrapper.vision_encoder, "forward", counting_forward)
+        with torch.no_grad():
+            for _ in range(3):  # stand-in for decode steps
+                wrapper(None, input_ids, visual_embeds=ve)
+        assert calls == 0, "cached decode re-ran the vision encoder"
+
+        with torch.no_grad():
+            wrapper(pixels, input_ids)
+        assert calls == 1, "uncached forward should run the vision encoder exactly once"
+
+    def test_visual_embeds_in_training_mode_raises(self):
+        # Cache path is inference-only: it bypasses the encoder + adapter, which
+        # would then silently receive no gradient. ValueError, not assert, so the
+        # guard survives python -O.
         wrapper = _build_tiny_wrapper().to(DEVICE).train()
         pixels = torch.randn(1, 3, 16, 16, device=DEVICE)
         input_ids = torch.randint(0, 256, (1, 8), device=DEVICE)
         ve = wrapper.encode_visual(pixels)
-        with pytest.raises(AssertionError):
-            wrapper(pixels, input_ids, visual_embeds=ve)
+        with pytest.raises(ValueError, match="inference-only"):
+            wrapper(None, input_ids, visual_embeds=ve)
+
+    def test_pixel_values_and_visual_embeds_together_raises(self):
+        """Passing both is the wrong-image bug: the pixels are ignored and may not
+        even correspond to the cache. Reject rather than silently supersede."""
+        wrapper = _build_tiny_wrapper(num_image_tokens=8).to(DEVICE).eval()
+        pixels = torch.randn(2, 3, 16, 16, device=DEVICE)
+        input_ids = torch.randint(0, 256, (2, 12), device=DEVICE)
+        with torch.no_grad():
+            ve = wrapper.encode_visual(pixels)
+            with pytest.raises(ValueError, match="not both"):
+                wrapper(pixels, input_ids, visual_embeds=ve)
+
+    def test_rank_2_visual_embeds_raises(self):
+        # (B, dim) would otherwise be read as n = dim visual tokens.
+        wrapper = _build_tiny_wrapper(num_image_tokens=8).to(DEVICE).eval()
+        input_ids = torch.randint(0, 256, (2, 12), device=DEVICE)
+        ve = torch.randn(2, 64, device=DEVICE)
+        with torch.no_grad(), pytest.raises(ValueError, match=r"must be \(B, N, dim\)"):
+            wrapper(None, input_ids, visual_embeds=ve)
+
+    def test_batch_mismatch_visual_embeds_raises(self):
+        # A batch-1 cache reused against a batched decode (beam search, batched
+        # eval) would otherwise fail deep inside Transformer.forward.
+        wrapper = _build_tiny_wrapper(num_image_tokens=8).to(DEVICE).eval()
+        input_ids = torch.randint(0, 256, (4, 12), device=DEVICE)
+        ve = torch.randn(1, 8, 64, device=DEVICE)
+        with torch.no_grad(), pytest.raises(ValueError, match="batch"):
+            wrapper(None, input_ids, visual_embeds=ve)
+
+    def test_wrong_token_count_visual_embeds_raises(self):
+        # n sets output_slice and the residual budget, so it must match the
+        # count this wrapper projects to.
+        wrapper = _build_tiny_wrapper(num_image_tokens=8).to(DEVICE).eval()
+        input_ids = torch.randint(0, 256, (2, 12), device=DEVICE)
+        ve = torch.randn(2, 5, 64, device=DEVICE)
+        with torch.no_grad(), pytest.raises(ValueError, match="visual token"):
+            wrapper(None, input_ids, visual_embeds=ve)
+
+    def test_wrong_feature_dim_visual_embeds_raises(self):
+        # Unprojected features (adapter skipped) instead of LLM-dim embeds.
+        wrapper = _build_tiny_wrapper(num_image_tokens=8).to(DEVICE).eval()
+        input_ids = torch.randint(0, 256, (2, 12), device=DEVICE)
+        ve = torch.randn(2, 8, 96, device=DEVICE)
+        with torch.no_grad(), pytest.raises(ValueError, match="feature dim"):
+            wrapper(None, input_ids, visual_embeds=ve)