Skip to content

[executorch][muse] Accept images in CUDA DFlash batching #23970

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Gasoonjia wants to merge 1 commit into
base: gh/Gasoonjia/72/base
Choose a base branch
from
Open

[executorch][muse] Accept images in CUDA DFlash batching #23970

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 40 additions & 24 deletions examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -547,10 +547,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"));
Expand Down Expand Up @@ -676,19 +675,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,
Expand All @@ -699,6 +690,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)};
}

Expand Down Expand Up @@ -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) {
Expand All @@ -1350,14 +1347,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) ||
Expand Down Expand Up @@ -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<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 fail_batch(batch, out);
}
Expand All @@ -1406,6 +1413,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 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)) {
Expand Down
7 changes: 6 additions & 1 deletion examples/models/muse-glimmer/runtime/cuda_dflash_executor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -189,6 +191,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.
Expand Down Expand Up @@ -264,6 +267,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_;
Expand Down
Loading