@@ -4,7 +4,7 @@ The batching runner and multiplexed worker support MLX and CUDA autoregressive
 generation with text only or text plus one image across the submitted
 conversation. They are opt-in entry points; the
 legacy runners and export defaults remain unchanged. DFlash batching runs on MLX
-and CUDA (text only).
+(text only) and CUDA, which also serves images when exported with `--mmproj`.
 CUDA specifics are in [CUDA](#cuda) below.
 
 ## Export
 
@@ -261,13 +261,6 @@ def export_dflash(
             activation_dtype=activation_dtype,
         )
 
-    if (
-        draft_model is not None
-        and use_offgraph_kv_cache
-        and backend == "cuda"
-        and vision_model is not None
-    ):
-        raise ValueError("the CUDA batching DFlash export does not support vision yet")
     if use_offgraph_kv_cache and backend == "cuda":
         from executorch.examples.models.muse_glimmer.export.export_solo import (
             DEFAULT_FORWARD_SIZES,
 
@@ -561,10 +561,9 @@ Result<MuseGlimmerBackend> MuseGlimmerCudaDFlashExecutor::create(
   ET_CHECK_OR_RETURN_ERROR(
       names.count(df::kForward) && names.count(df::kVerify) &&
           names.count(df::kVerifyArgmax) &&
-          !names.count("get_mutable_buffer_metadata") &&
-          !names.count("vision_encoder"),
+          !names.count("get_mutable_buffer_metadata"),
       InvalidProgram,
-      "CUDA DFlash requires a tapped text-only off-graph artifact with its verifiers");
+      "CUDA DFlash requires a tapped off-graph artifact with its verifiers");
   ET_ASSIGN_OR_RETURN(block, positive_int(module, df::kBlock));
   ET_ASSIGN_OR_RETURN(
       declared_block, positive_int(module, "get_draft_dflash_block_size"));
@@ -690,19 +689,11 @@ Result<MuseGlimmerBackend> MuseGlimmerCudaDFlashExecutor::create(
           config.bos_id < static_cast<std::uint64_t>(target->vocab_size()),
       InvalidProgram,
       "Muse Glimmer DFlash vocabulary cannot hold its special tokens");
-  auto spec = std::make_shared<const MuseGlimmerPreparationSpec>(
-      MuseGlimmerPreparationSpec{
-          aten::ScalarType::BFloat16,
-          0,
-          max_context,
-          static_cast<std::int32_t>(target->preferred_batch_tokens()),
-          target->vocab_size(),
-          false,
-          0,
-          config.image_limits,
-          config.bos_id});
-  auto executor =
-      std::unique_ptr<batching::Executor>(new MuseGlimmerCudaDFlashExecutor(
+  std::shared_ptr<MuseGlimmerModelEmbeddings> embeddings;
+  ET_ASSIGN_OR_RETURN(
+      spec, make_cuda_preparation(*target, config, max_context, &embeddings));
+  auto executor = std::unique_ptr<MuseGlimmerCudaDFlashExecutor>(
+      new MuseGlimmerCudaDFlashExecutor(
           std::move(target),
           std::move(draft_cache),
           block,
@@ -713,6 +704,7 @@ Result<MuseGlimmerBackend> MuseGlimmerCudaDFlashExecutor::create(
           config.max_sessions,
           max_context,
           hidden_width));
+  executor->embeddings_ = std::move(embeddings);
   return MuseGlimmerBackend{std::move(executor), std::move(spec)};
 }
 
@@ -1390,6 +1382,11 @@ bool MuseGlimmerCudaDFlashExecutor::fail_batch(
   return false;
 }
 
+bool MuseGlimmerCudaDFlashExecutor::accepts(
+    const batching::PreparedInput& input) const {
+  return embeddings_ && embeddings_->accepts(input);
+}
+
 bool MuseGlimmerCudaDFlashExecutor::execute(
     const batching::BatchInput& batch,
     batching::BatchOutput& out) {
@@ -1409,14 +1406,21 @@ bool MuseGlimmerCudaDFlashExecutor::execute(
     std::unordered_map<SessionId, int> cursor;
     std::vector<std::pair<SessionId, int>> rewinds;
     bool pending_decode = !batch.inputs.empty();
+    bool prepared_step = false;
     for (const auto& input : batch.inputs) {
       const auto it = sessions_.find(input.sid);
       const auto* payload =
           std::get_if<batching::TokenInputPtr>(&input.payload);
-      if (it == sessions_.end() || it->second.poisoned || !payload ||
-          !*payload || input.size == 0 || input.offset > (*payload)->size() ||
-          input.size > (*payload)->size() - input.offset ||
-          input.position < 0 ||
+      const auto* prepared =
+          std::get_if<batching::PreparedInputPtr>(&input.payload);
+      const bool tokens = payload && *payload;
+      const bool image = prepared && *prepared && accepts(**prepared);
+      const std::size_t available = tokens ? (*payload)->size()
+          : image                          ? (*prepared)->size()
+                                           : 0;
+      if (it == sessions_.end() || it->second.poisoned || (!tokens && !image) ||
+          input.size == 0 || input.offset > available ||
+          input.size > available - input.offset || input.position < 0 ||
           input.offset > static_cast<std::size_t>(max_context_) ||
           input.size > static_cast<std::size_t>(max_context_) ||
           (input.produce_output && !it->second.sampling) ||
@@ -1447,13 +1451,16 @@ bool MuseGlimmerCudaDFlashExecutor::execute(
       // A one-token prompt/replay is not necessarily emitted decode feedback.
       // Stay on the legacy path unless every owner has an exact pending seed,
       // one input, and room for the complete trained draft block.
-      pending_decode = pending_decode && first && input.produce_output &&
-          input.size == 1 && start == it->second.frontier &&
-          start <= max_context_ - block_ && it->second.pending &&
+      pending_decode = pending_decode && tokens && first &&
+          input.produce_output && input.size == 1 &&
+          start == it->second.frontier && start <= max_context_ - block_ &&
+          it->second.pending &&
           *it->second.pending == (**payload)[input.offset];
+      prepared_step = prepared_step || image;
       at->second = static_cast<int>(end);
       for (std::size_t j = 0; j < input.size; ++j) {
-        const auto token = (**payload)[input.offset + j];
+        // Prepared rows reach the target as embeddings, not tokens.
+        const Token token = tokens ? (**payload)[input.offset + j] : 0;
         if (token >= static_cast<std::uint64_t>(target_->vocab_size())) {
           return reject_batch();
         }
@@ -1467,6 +1474,15 @@ bool MuseGlimmerCudaDFlashExecutor::execute(
               ? static_cast<std::int64_t>(step.tokens.size()) - 1
               : -1);
     }
+    if (prepared_step) {
+      auto embeddings =
+          embeddings_->materialize(target_->module(), batch.inputs);
+      if (!embeddings.ok()) {
+        ET_LOG(Error, "CUDA DFlash: could not embed the prepared inputs");
+        return reject_batch();
+      }
+      step.embeddings = std::move(*embeddings);
+    }
     // All inputs are validated before either cache is truncated.
     for (const auto& item : rewinds) {
       if (!rewind(sessions_.at(item.first), item.second)) {
 
@@ -172,7 +172,9 @@ struct CudaDFlashStats {
   std::uint64_t failed_batches = 0;
 };
 
-// Text-only model adapter over the existing CUDA target executor. Decode drafts
+// Model adapter over the existing CUDA target executor. Image prompts prefill
+// through the target's embeddings forward; their rows' taps then feed the
+// draft like any other context. Decode drafts
 // run the static draft_forward_{N} in CUDA graphs, as the target's static
 // methods do; longer contexts drain through the dynamic draft_forward. Drafts
 // are greedy; dflash_verify accepts or resamples them under each session's
@@ -199,6 +201,7 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor {
       batching::SessionId session,
       const batching::SamplingParams& params,
       std::optional<std::uint64_t> seed) override;
+  bool accepts(const batching::PreparedInput& input) const override;
   bool execute(const batching::BatchInput& batch, batching::BatchOutput& out)
       override;
   // Engine-thread-only, between executes. Counters measure actual adapter work.
@@ -276,6 +279,8 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor {
   std::size_t row_bytes_;
   batching::SessionId next_session_ = 1;
   std::unordered_map<batching::SessionId, State> sessions_;
+  // Null for a text-only program.
+  std::shared_ptr<MuseGlimmerModelEmbeddings> embeddings_;
   CudaDFlashStats stats_;
   cuda_dflash_detail::HiddenSnapshot verification_hidden_;
   cuda_dflash_detail::DeviceStaging verify_inputs_;