1111
1212``--use-offgraph-kv-cache`` keeps the KV cache out of the program and exports
1313the batching runtime's contract. On MLX that is the selected-logits decoder. On
14- CUDA it is the text-only artifact CudaExecutor (backends/cuda/batching) drives;
15- without the flag CUDA keeps its in-graph prefill and decode methods:
14+ CUDA it is the artifact CudaExecutor (backends/cuda/batching) drives: the text
15+ forwards plus, with ``--mmproj``, ``vision_encoder``/``embed_text``/
16+ ``forward_from_embeddings``; without the flag CUDA keeps its in-graph prefill
17+ and decode methods:
1618
1719- ``forward_{N}`` for each N in ``--forward-sizes`` (default 1, 2, 4, ..., 64):
1820 T = K = N, static, each captured into a CUDA graph at runtime, its quantized
@@ -133,8 +135,8 @@ def export_and_lower(
133135 if not 1 <= max_prefill_chunk <= config .max_seq_len :
134136 raise ValueError ("max_prefill_chunk must be within [1, max_seq_len]" )
135137 if backend == "cuda" and use_offgraph_kv_cache :
136- if vision_model is not None :
137- raise ValueError ("the CUDA batching export is text-only " )
138+ if vision_fp32_mm != "none" :
139+ raise ValueError ("CUDA batching does not support vision_fp32_mm " )
138140 export_cuda_batching (
139141 model ,
140142 config ,
@@ -143,6 +145,9 @@ def export_and_lower(
143145 max_cells = max_cells ,
144146 forward_sizes = forward_sizes or DEFAULT_FORWARD_SIZES ,
145147 device_sampling = sample ,
148+ vision_model = vision_model ,
149+ pos_embed_table = pos_embed_table ,
150+ max_vision_patches = max_vision_patches ,
146151 )
147152 elif backend == "cuda" :
148153 _export_cuda (
@@ -257,6 +262,10 @@ def step_width_spec(method: str) -> bytes:
257262 return b"1:0"
258263
259264
265+ # Image prompts reach the batching cache through this method.
266+ EMBEDDINGS_METHOD = "forward_from_embeddings"
267+
268+
260269def step_forward (
261270 model : MuseGlimmerModel ,
262271 tokens : torch .Tensor ,
@@ -268,7 +277,18 @@ def step_forward(
268277 The whole step runs through the decoder so every token's K/V reaches the
269278 cache; only the rows the batch samples from go through the LM head.
270279 """
271- x = model .embed_text (tokens )
280+ return embeddings_step_forward (
281+ model , model .embed_text (tokens ), input_pos , logits_to_keep
282+ )
283+
284+
285+ def embeddings_step_forward (
286+ model : MuseGlimmerModel ,
287+ x : torch .Tensor ,
288+ input_pos : torch .Tensor ,
289+ logits_to_keep : torch .Tensor ,
290+ ) -> torch .Tensor :
291+ """``step_forward`` from decoder inputs [1, T, dim], e.g. spliced images."""
272292 taps = None
273293 if hasattr (model , "_dflash_target_layer_ids" ):
274294 sliding_mask , global_mask = model ._build_masks (input_pos )
@@ -349,9 +369,16 @@ def export_cuda_batching( # noqa: C901
349369 dflash_draft : nn .Module | None = None ,
350370 max_draft_owners : int = 4 ,
351371 draft_max_cells : int | None = None ,
372+ vision_model : nn .Module | None = None ,
373+ pos_embed_table : torch .Tensor | None = None ,
374+ max_vision_patches : int = 16384 ,
352375):
353376 """Exports the forward methods and writes the artifact.
354377
378+ With ``vision_model`` the token forwards are unchanged, and image prompts run
379+ through three more methods: ``vision_encoder``, ``embed_text`` and a dynamic
380+ ``forward_from_embeddings`` over the same cache and weights.
381+
355382 Returns the ExecutorchProgramManager so callers can inspect it.
356383 """
357384 import executorch .backends .cuda .quantize_op_dispatch # noqa: F401
@@ -379,6 +406,10 @@ def export_cuda_batching( # noqa: C901
379406 )
380407 from torch .export import Dim , export
381408
409+ has_vision = vision_model is not None and pos_embed_table is not None
410+ if (vision_model is None ) != (pos_embed_table is None ):
411+ raise ValueError ("vision_model and pos_embed_table must be given together" )
412+
382413 if dflash_draft is not None :
383414 from executorch .examples .models .muse_glimmer .export .export_dflash import (
384415 prepare_cuda_packed_draft ,
@@ -462,10 +493,52 @@ def export_cuda_batching( # noqa: C901
462493 dynamic_shapes = ({1 : width }, {0 : width }, {0 : rows }),
463494 strict = True ,
464495 )
496+ if has_vision :
497+ embeddings_width = Dim (
498+ "embeddings_width" , min = _CUDA_MIN_PREFILL_CHUNK , max = max_step_tokens
499+ )
500+ with (
501+ common .BoundMethodForward (
502+ model , types .MethodType (embeddings_step_forward , model )
503+ ),
504+ torch .no_grad (),
505+ ):
506+ print (
507+ f"Exporting { EMBEDDINGS_METHOD } (T in [{ _CUDA_MIN_PREFILL_CHUNK } , "
508+ f"{ max_step_tokens } ])..."
509+ )
510+ programs [EMBEDDINGS_METHOD ] = export (
511+ model ,
512+ (
513+ torch .zeros ((1 , max_step_tokens , config .dim ), dtype = torch .bfloat16 ),
514+ torch .arange (max_step_tokens , dtype = torch .long ),
515+ torch .zeros (max_step_tokens , dtype = torch .long ),
516+ ),
517+ dynamic_shapes = (
518+ {1 : embeddings_width },
519+ {0 : embeddings_width },
520+ {0 : rows },
521+ ),
522+ strict = True ,
523+ )
524+ with common .BoundMethodForward (model , model .embed_text ), torch .no_grad ():
525+ print ("Exporting embed_text..." )
526+ programs ["embed_text" ] = export (
527+ model ,
528+ (torch .zeros ((1 , max_step_tokens ), dtype = torch .long ),),
529+ dynamic_shapes = ({1 : Dim ("embed_width" , min = 1 , max = max_step_tokens )},),
530+ strict = True ,
531+ )
532+ programs ["vision_encoder" ] = common .export_vision_encoder (
533+ vision_model , pos_embed_table , max_vision_patches
534+ )
535+ del vision_model
465536 del model
466537 gc .collect ()
467538 torch .cuda .empty_cache ()
468- forward_methods = list (programs )
539+ forward_methods = [
540+ name for name in programs if name not in ("embed_text" , "vision_encoder" )
541+ ]
469542 if dflash_draft is not None :
470543 programs .update (draft_programs )
471544 del draft_programs , dflash_draft
@@ -520,8 +593,8 @@ def partitioner(name: str) -> CudaPartitioner:
520593 max_prefill = max_step_tokens ,
521594 activation_dtype = torch .bfloat16 ,
522595 mutable_buffer_metadata = None ,
523- has_vision = False ,
524- max_vision_patches = 0 ,
596+ has_vision = has_vision ,
597+ max_vision_patches = max_vision_patches ,
525598 )
526599 constant_methods .update (geometry )
527600 constant_methods .update (write_max_context_len (config .max_seq_len ))
@@ -560,7 +633,7 @@ def partitioner(name: str) -> CudaPartitioner:
560633 del edge
561634 gc .collect ()
562635
563- common .save_pte (et_program , output_dir , None )
636+ common .save_pte (et_program , output_dir , pos_embed_table )
564637 if torch .cuda .is_available ():
565638 peak_mb = torch .cuda .max_memory_allocated () / (1024 ** 2 )
566639 print (f"EXPORT_GPU_PEAK_MEMORY_MB: { peak_mb :.1f} " )
@@ -1040,8 +1113,9 @@ def main() -> None:
10401113 "--use-offgraph-kv-cache" ,
10411114 action = "store_true" ,
10421115 help = "Allocate the KV cache at runtime and export the batching runtime's "
1043- "contract: MLX's selected-logits decoder, or CUDA's text-only "
1044- "forward_{N} + forward_others over a cell-layout cache. Without it CUDA "
1116+ "contract: MLX's selected-logits decoder, or CUDA's "
1117+ "forward_{N} + forward_others over a cell-layout cache (plus the vision "
1118+ "methods with --mmproj). Without it CUDA "
10451119 "exports its in-graph prefill and decode methods." ,
10461120 )
10471121 parser .add_argument (
@@ -1109,10 +1183,6 @@ def main() -> None:
11091183 args = parser .parse_args ()
11101184
11111185 _validate_offgraph_prefill (parser , args )
1112- if args .backend == "cuda" and args .use_offgraph_kv_cache and args .mmproj :
1113- parser .error (
1114- "--use-offgraph-kv-cache on CUDA exports a text-only model; omit --mmproj."
1115- )
11161186
11171187 if args .backend == "cuda" and not torch .cuda .is_available ():
11181188 parser .error ("CUDA is required for the cuda backend." )
0 commit comments