[executorch][muse] Export vision with the CUDA batching contract

examples/models/muse-glimmer/export/export_dflash.py

@@ -184,8 +184,6 @@ def export_dflash(
     mmproj loader and vision export contract to add ``vision_encoder``.
     """
     validate_dflash_export_options(backend, use_offgraph_kv_cache)
-    if use_offgraph_kv_cache and backend == "cuda" and mmproj is not None:
-        raise ValueError("the CUDA batching DFlash export is text-only")
 
     import executorch.extension.llm.export.gguf  # noqa: F401
     import executorch.extension.llm.export.int4  # noqa: F401
@@ -243,6 +241,22 @@ def export_dflash(
             activation_dtype=activation_dtype,
         )
 
+    vision_model = None
+    pos_embed_table = None
+    if mmproj is not None:
+        from executorch.examples.models.muse_glimmer.loaders.checkpoint_loader import (
+            load_mmproj_vision_model,
+        )
+
+        print("=" * 60)
+        print("Loading vision encoder...")
+        print("=" * 60)
+        vision_model, pos_embed_table, _ = load_mmproj_vision_model(
+            mmproj,
+            backend=backend,
+            activation_dtype=activation_dtype,
+        )
+
     if use_offgraph_kv_cache and backend == "cuda":
         from executorch.examples.models.muse_glimmer.export.export_solo import (
             DEFAULT_FORWARD_SIZES,
@@ -261,25 +275,12 @@ def export_dflash(
             dflash_draft=draft_model,
             max_draft_owners=max_draft_owners,
             draft_max_cells=draft_max_cells,
+            vision_model=vision_model,
+            pos_embed_table=pos_embed_table,
+            max_vision_patches=max_vision_patches,
         )
         return
 
-    vision_model = None
-    pos_embed_table = None
-    if mmproj is not None:
-        from executorch.examples.models.muse_glimmer.loaders.checkpoint_loader import (
-            load_mmproj_vision_model,
-        )
-
-        print("=" * 60)
-        print("Loading vision encoder...")
-        print("=" * 60)
-        vision_model, pos_embed_table, _ = load_mmproj_vision_model(
-            mmproj,
-            backend=backend,
-            activation_dtype=activation_dtype,
-        )
-
     backend_export = _export_dflash_mlx if backend == "mlx" else _export_dflash_cuda
     backend_export(
         target_model,

examples/models/muse-glimmer/export/export_solo.py

@@ -133,8 +133,6 @@ def export_and_lower(
         if not 1 <= max_prefill_chunk <= config.max_seq_len:
             raise ValueError("max_prefill_chunk must be within [1, max_seq_len]")
     if backend == "cuda" and use_offgraph_kv_cache:
-        if vision_model is not None:
-            raise ValueError("the CUDA batching export is text-only")
         export_cuda_batching(
             model,
             config,
@@ -143,6 +141,9 @@ def export_and_lower(
             max_cells=max_cells,
             forward_sizes=forward_sizes or DEFAULT_FORWARD_SIZES,
             device_sampling=sample,
+            vision_model=vision_model,
+            pos_embed_table=pos_embed_table,
+            max_vision_patches=max_vision_patches,
         )
     elif backend == "cuda":
         _export_cuda(
@@ -257,6 +258,10 @@ def step_width_spec(method: str) -> bytes:
     return b"1:0"
 
 
+# Image prompts reach the batching cache through this method.
+EMBEDDINGS_METHOD = "forward_from_embeddings"
+
+
 def step_forward(
     model: MuseGlimmerModel,
     tokens: torch.Tensor,
@@ -268,7 +273,18 @@ def step_forward(
     The whole step runs through the decoder so every token's K/V reaches the
     cache; only the rows the batch samples from go through the LM head.
     """
-    x = model.embed_text(tokens)
+    return embeddings_step_forward(
+        model, model.embed_text(tokens), input_pos, logits_to_keep
+    )
+
+
+def embeddings_step_forward(
+    model: MuseGlimmerModel,
+    x: torch.Tensor,
+    input_pos: torch.Tensor,
+    logits_to_keep: torch.Tensor,
+) -> torch.Tensor:
+    """``step_forward`` from decoder inputs [1, T, dim], e.g. spliced images."""
     taps = None
     if hasattr(model, "_dflash_target_layer_ids"):
         sliding_mask, global_mask = model._build_masks(input_pos)
@@ -349,9 +365,16 @@ def export_cuda_batching(  # noqa: C901
     dflash_draft: nn.Module | None = None,
     max_draft_owners: int = 4,
     draft_max_cells: int | None = None,
+    vision_model: nn.Module | None = None,
+    pos_embed_table: torch.Tensor | None = None,
+    max_vision_patches: int = 16384,
 ):
     """Exports the forward methods and writes the artifact.
 
+    With ``vision_model`` the token forwards are unchanged, and image prompts run
+    through three more methods: ``vision_encoder``, ``embed_text`` and a dynamic
+    ``forward_from_embeddings`` over the same cache and weights.
+
     Returns the ExecutorchProgramManager so callers can inspect it.
     """
     import executorch.backends.cuda.quantize_op_dispatch  # noqa: F401
@@ -462,10 +485,52 @@ def export_cuda_batching(  # noqa: C901
             dynamic_shapes=({1: width}, {0: width}, {0: rows}),
             strict=True,
         )
+    if vision_model is not None:
+        embeddings_width = Dim(
+            "embeddings_width", min=_CUDA_MIN_PREFILL_CHUNK, max=max_step_tokens
+        )
+        with (
+            common.BoundMethodForward(
+                model, types.MethodType(embeddings_step_forward, model)
+            ),
+            torch.no_grad(),
+        ):
+            print(
+                f"Exporting {EMBEDDINGS_METHOD} (T in [{_CUDA_MIN_PREFILL_CHUNK}, "
+                f"{max_step_tokens}])..."
+            )
+            programs[EMBEDDINGS_METHOD] = export(
+                model,
+                (
+                    torch.zeros(
+                        (1, max_step_tokens, config.dim), dtype=torch.bfloat16
+                    ),
+                    torch.arange(max_step_tokens, dtype=torch.long),
+                    torch.zeros(max_step_tokens, dtype=torch.long),
+                ),
+                dynamic_shapes=({1: embeddings_width}, {0: embeddings_width}, {0: rows}),
+                strict=True,
+            )
+        with common.BoundMethodForward(model, model.embed_text), torch.no_grad():
+            print("Exporting embed_text...")
+            programs["embed_text"] = export(
+                model,
+                (torch.zeros((1, max_step_tokens), dtype=torch.long),),
+                dynamic_shapes=(
+                    {1: Dim("embed_width", min=1, max=max_step_tokens)},
+                ),
+                strict=True,
+            )
+        programs["vision_encoder"] = common.export_vision_encoder(
+            vision_model, pos_embed_table, max_vision_patches
+        )
+        del vision_model
     del model
     gc.collect()
     torch.cuda.empty_cache()
-    forward_methods = list(programs)
+    forward_methods = [
+        name for name in programs if name not in ("embed_text", "vision_encoder")
+    ]
     if dflash_draft is not None:
         programs.update(draft_programs)
         del draft_programs, dflash_draft
@@ -524,8 +589,8 @@ def partitioner(name: str) -> CudaPartitioner:
         max_prefill=max_step_tokens,
         activation_dtype=torch.bfloat16,
         mutable_buffer_metadata=None,
-        has_vision=False,
-        max_vision_patches=0,
+        has_vision=pos_embed_table is not None,
+        max_vision_patches=max_vision_patches,
     )
     constant_methods.update(geometry)
     constant_methods.update(write_max_context_len(config.max_seq_len))
@@ -564,7 +629,7 @@ def partitioner(name: str) -> CudaPartitioner:
     del edge
     gc.collect()
 
-    common.save_pte(et_program, output_dir, None)
+    common.save_pte(et_program, output_dir, pos_embed_table)
     if torch.cuda.is_available():
         peak_mb = torch.cuda.max_memory_allocated() / (1024**2)
         print(f"EXPORT_GPU_PEAK_MEMORY_MB: {peak_mb:.1f}")
@@ -1113,10 +1178,6 @@ def main() -> None:
     args = parser.parse_args()
 
     _validate_offgraph_prefill(parser, args)
-    if args.backend == "cuda" and args.use_offgraph_kv_cache and args.mmproj:
-        parser.error(
-            "--use-offgraph-kv-cache on CUDA exports a text-only model; omit --mmproj."
-        )
 
     if args.backend == "cuda" and not torch.cuda.is_available():
         parser.error("CUDA is required for the cuda backend.")

examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py

@@ -184,6 +184,96 @@ def setUp(self) -> None:
         if not torch.cuda.is_available():
             self.skipTest("CUDA required")
 
+    def test_vision_adds_embedding_methods_over_the_same_cache(self) -> None:
+        """Images add vision, embed_text and one dynamic embeddings forward."""
+        try:
+            import gguf  # noqa: F401
+        except ImportError:
+            self.skipTest("gguf package required")
+        from executorch.backends.cuda.passes.lower_offgraph_kv import (
+            OFFGRAPH_KV_COMPILE_SPEC,
+            OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC,
+        )
+        from executorch.examples.models.muse_glimmer.loaders.checkpoint_loader import (
+            load_mmproj_vision_model,
+        )
+        from executorch.examples.models.muse_glimmer.tests.test_pipeline import (
+            _mmproj_vision_config,
+            build_mmproj_gguf,
+        )
+        from executorch.exir.lowered_backend_module import get_lowered_submodules
+
+        with (
+            tempfile.TemporaryDirectory() as ckpt_dir,
+            tempfile.TemporaryDirectory() as out_dir,
+        ):
+            save_checkpoint(ckpt_dir)
+            model, config = load_prequantized_model(
+                ckpt_dir, max_seq_len=TINY_CONFIG.max_seq_len
+            )
+            mmproj = os.path.join(ckpt_dir, "mmproj.gguf")
+            build_mmproj_gguf(mmproj)
+            vision, pos_table, _ = load_mmproj_vision_model(
+                mmproj, config=_mmproj_vision_config(), backend="cuda"
+            )
+            program = export_cuda_batching(
+                model,
+                config,
+                out_dir,
+                max_step_tokens=self.MAX_STEP,
+                forward_sizes=self.FORWARD_SIZES,
+                max_cells=self.MAX_CELLS,
+                vision_model=vision,
+                pos_embed_table=pos_table,
+                max_vision_patches=64,
+            )
+            self.assertTrue(os.path.exists(os.path.join(out_dir, "pos_embed.bin")))
+            methods = set(program.methods)
+            self.assertLessEqual(
+                {"forward_from_embeddings", "embed_text", "vision_encoder"}, methods
+            )
+            embeddings = program.exported_program("forward_from_embeddings")
+            lowered = get_lowered_submodules(embeddings.graph_module)
+            self.assertEqual(len(lowered), 1)
+            specs = {spec.key: spec.value for spec in lowered[0][1].compile_specs}
+            # Same cache and step-width contract as the token forwards.
+            others = get_lowered_submodules(
+                program.exported_program("forward_others").graph_module
+            )[0][1].compile_specs
+            self.assertEqual(
+                specs[OFFGRAPH_KV_COMPILE_SPEC],
+                {s.key: s.value for s in others}[OFFGRAPH_KV_COMPILE_SPEC],
+            )
+            self.assertEqual(specs[OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC], b"1:0")
+            for name in ("embed_text", "vision_encoder"):
+                sub = get_lowered_submodules(
+                    program.exported_program(name).graph_module
+                )
+                self.assertEqual(len(sub), 1, name)
+                self.assertNotIn(
+                    OFFGRAPH_KV_COMPILE_SPEC,
+                    {spec.key for spec in sub[0][1].compile_specs},
+                    name,
+                )
+
+            from executorch.runtime import Runtime, Verification
+
+            loaded = Runtime.get().load_program(
+                os.path.join(out_dir, "model.pte"), verification=Verification.Minimal
+            )
+            meta = loaded.metadata("forward_from_embeddings")
+            self.assertEqual(meta.num_inputs(), 3)
+            self.assertEqual(
+                meta.input_tensor_meta(0).sizes(), (1, self.MAX_STEP, TINY_CONFIG.dim)
+            )
+            self.assertEqual(
+                loaded.load_method("has_vision_encoder").execute([])[0], True
+            )
+            self.assertEqual(
+                loaded.load_method("get_vision_hidden_size").execute([])[0],
+                TINY_CONFIG.dim,
+            )
+
     def test_exports_forward_methods_over_a_cell_cache(self) -> None:
         from executorch.backends.cuda.passes.lower_offgraph_kv import (
             OFFGRAPH_KV_COMPILE_SPEC,