diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp index f0dc4013bb8..f053a7784f8 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp @@ -547,10 +547,9 @@ Result 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")); @@ -676,19 +675,11 @@ Result MuseGlimmerCudaDFlashExecutor::create( config.bos_id < static_cast(target->vocab_size()), InvalidProgram, "Muse Glimmer DFlash vocabulary cannot hold its special tokens"); - auto spec = std::make_shared( - MuseGlimmerPreparationSpec{ - aten::ScalarType::BFloat16, - 0, - max_context, - static_cast(target->preferred_batch_tokens()), - target->vocab_size(), - false, - 0, - config.image_limits, - config.bos_id}); - auto executor = - std::unique_ptr(new MuseGlimmerCudaDFlashExecutor( + std::shared_ptr embeddings; + ET_ASSIGN_OR_RETURN( + spec, make_cuda_preparation(*target, config, max_context, &embeddings)); + auto executor = std::unique_ptr( + new MuseGlimmerCudaDFlashExecutor( std::move(target), std::move(draft_cache), block, @@ -699,6 +690,7 @@ Result MuseGlimmerCudaDFlashExecutor::create( config.max_sessions, max_context, hidden_width)); + executor->embeddings_ = std::move(embeddings); return MuseGlimmerBackend{std::move(executor), std::move(spec)}; } @@ -1334,6 +1326,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) { @@ -1350,14 +1347,21 @@ bool MuseGlimmerCudaDFlashExecutor::execute( std::unordered_map cursor; std::vector> 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(&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(&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(max_context_) || input.size > static_cast(max_context_) || (input.produce_output && !it->second.sampling) || @@ -1386,13 +1390,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(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(target_->vocab_size())) { return fail_batch(batch, out); } @@ -1406,6 +1413,15 @@ bool MuseGlimmerCudaDFlashExecutor::execute( ? static_cast(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 fail_batch(batch, out); + } + 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)) { diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h index 2a023aa57cd..96b6bdbad62 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h @@ -168,7 +168,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 @@ -189,6 +191,7 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor { batching::SessionId session, const batching::SamplingParams& params, std::optional 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. @@ -264,6 +267,8 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor { std::size_t row_bytes_; batching::SessionId next_session_ = 1; std::unordered_map sessions_; + // Null for a text-only program. + std::shared_ptr embeddings_; CudaDFlashStats stats_; cuda_dflash_detail::HiddenSnapshot verification_hidden_; cuda_dflash_detail::DeviceStaging verify_inputs_;