diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index e69c65c213e..9a214964ce1 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -50,7 +50,7 @@ Arguments: - vr-streaming: Voxtral Realtime streaming mode - vr-offline: Voxtral Realtime offline mode - solo-text: Muse Glimmer solo text mode - - solo-text-batching: Muse Glimmer text on the CUDA batching runtime + - dflash-image-batching: Muse Glimmer DFlash text and image on the CUDA batching runtime - dflash-image: Muse Glimmer DFlash vision mode Examples: @@ -105,7 +105,7 @@ if [ -n "$MODE" ]; then exit 1 fi ;; - solo-text|solo-text-batching|dflash-image) + solo-text|dflash-image-batching|dflash-image) if [ "$HF_MODEL" != "meta-models/Muse-Glimmer-30B-GGUF" ]; then echo "Error: Mode '$MODE' can only be used with Muse Glimmer model" echo "Provided model: $HF_MODEL" @@ -114,7 +114,7 @@ if [ -n "$MODE" ]; then ;; *) echo "Error: Unsupported mode '$MODE'" - echo "Supported modes: vr-streaming, vr-offline, solo-text, solo-text-batching, dflash-image" + echo "Supported modes: vr-streaming, vr-offline, solo-text, dflash-image-batching, dflash-image" exit 1 ;; esac @@ -354,8 +354,8 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then echo "Error: Muse Glimmer requires quantization 'kquant-17gb' or 'kquant-dynamic'" exit 1 fi - if [ "$MODE" != "solo-text" ] && [ "$MODE" != "solo-text-batching" ] && [ "$MODE" != "dflash-image" ]; then - echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-batching' or 'dflash-image'" + if [ "$MODE" != "solo-text" ] && [ "$MODE" != "dflash-image-batching" ] && [ "$MODE" != "dflash-image" ]; then + echo "Error: Muse Glimmer requires mode 'solo-text', 'dflash-image-batching' or 'dflash-image'" exit 1 fi fi @@ -707,23 +707,17 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then --backend cuda \ --output-dir "${OUTPUT_DIR}" ;; - solo-text-batching) - # The off-graph cache exports the batching methods CudaExecutor drives: - # Static forward_{N}, dynamic forward_others and the device samplers. - EXPORT_START_SECONDS=$SECONDS - TMPDIR="$INDUCTOR_TMPDIR" \ - TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \ - python -m executorch.examples.models.muse_glimmer.export.export_solo \ - --gguf "$TARGET_GGUF_PATH" \ - --backend cuda \ - --use-offgraph-kv-cache \ - --output-dir "${OUTPUT_DIR}" - ;; - dflash-image) + dflash-image|dflash-image-batching) DRAFT_GGUF_FILE="dflash-Muse-Glimmer-30B-Q4_K_M.gguf" MMPROJ_GGUF_FILE="mmproj-Muse-Glimmer-30B-Q4_K_M.gguf" python -c "from huggingface_hub import hf_hub_download; hf_hub_download('${HF_MODEL}', '${DRAFT_GGUF_FILE}', local_dir='${LOCAL_MODEL_DIR}')" python -c "from huggingface_hub import hf_hub_download; hf_hub_download('${HF_MODEL}', '${MMPROJ_GGUF_FILE}', local_dir='${LOCAL_MODEL_DIR}')" + # The off-graph cache exports the batching methods the CUDA DFlash + # executor drives, plus the vision encoder and forward_from_embeddings. + OFFGRAPH_FLAG="" + if [ "$MODE" = "dflash-image-batching" ]; then + OFFGRAPH_FLAG="--use-offgraph-kv-cache" + fi EXPORT_START_SECONDS=$SECONDS TMPDIR="$INDUCTOR_TMPDIR" \ TORCHINDUCTOR_CACHE_DIR="$INDUCTOR_CACHE" \ @@ -732,10 +726,11 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then --draft-gguf "${LOCAL_MODEL_DIR}/${DRAFT_GGUF_FILE}" \ --mmproj "${LOCAL_MODEL_DIR}/${MMPROJ_GGUF_FILE}" \ --backend cuda \ + $OFFGRAPH_FLAG \ --output-dir "${OUTPUT_DIR}" ;; *) - echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-batching' or 'dflash-image'" + echo "Error: Muse Glimmer requires mode 'solo-text', 'dflash-image-batching' or 'dflash-image'" exit 1 ;; esac @@ -746,7 +741,7 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then test -f "${OUTPUT_DIR}/model.pte" test -f "${OUTPUT_DIR}/aoti_cuda_blob.ptd" - if [ "$MODE" = "dflash-image" ]; then + if [ "$MODE" = "dflash-image" ] || [ "$MODE" = "dflash-image-batching" ]; then test -f "${OUTPUT_DIR}/pos_embed.bin" fi ls -al "${OUTPUT_DIR}" diff --git a/.ci/scripts/muse_glimmer_batching_smoke.cpp b/.ci/scripts/muse_glimmer_batching_smoke.cpp index 6b9c794c706..12e821ea696 100644 --- a/.ci/scripts/muse_glimmer_batching_smoke.cpp +++ b/.ci/scripts/muse_glimmer_batching_smoke.cpp @@ -20,6 +20,7 @@ DECLARE_string(pte); DECLARE_string(tokenizer); DECLARE_string(data_path); DECLARE_string(backend); +DECLARE_bool(dflash); DECLARE_int32(max_sessions); DECLARE_int32(max_session_tokens); DECLARE_uint64(bos_id); @@ -64,6 +65,7 @@ int main(int argc, char** argv) { config.max_sessions = FLAGS_max_sessions; config.max_session_tokens = FLAGS_max_session_tokens; config.bos_id = FLAGS_bos_id; + config.use_dflash = FLAGS_dflash; auto backend = llm::create_muse_glimmer_backend(config); if (!backend.ok()) { throw std::runtime_error( diff --git a/.ci/scripts/test_model_e2e.sh b/.ci/scripts/test_model_e2e.sh index 786cd81f9e2..70b8f3855e7 100755 --- a/.ci/scripts/test_model_e2e.sh +++ b/.ci/scripts/test_model_e2e.sh @@ -51,7 +51,7 @@ Arguments: - vr-streaming: Voxtral Realtime streaming mode - vr-offline: Voxtral Realtime offline mode - solo-text: Muse Glimmer solo text mode - - solo-text-batching: Muse Glimmer text on the CUDA batching runtime + - dflash-image-batching: Muse Glimmer DFlash text and image on the CUDA batching runtime - dflash-image: Muse Glimmer DFlash vision mode Environment: @@ -126,7 +126,7 @@ if [ -n "$MODE" ]; then exit 1 fi ;; - solo-text|solo-text-batching|dflash-image) + solo-text|dflash-image-batching|dflash-image) if [ "$HF_MODEL" != "meta-models/Muse-Glimmer-30B-GGUF" ]; then echo "Error: Mode '$MODE' can only be used with Muse Glimmer model" echo "Provided model: $HF_MODEL" @@ -135,7 +135,7 @@ if [ -n "$MODE" ]; then ;; *) echo "Error: Unsupported mode '$MODE'" - echo "Supported modes: vr-streaming, vr-offline, solo-text, solo-text-batching, dflash-image" + echo "Supported modes: vr-streaming, vr-offline, solo-text, dflash-image-batching, dflash-image" exit 1 ;; esac @@ -321,12 +321,12 @@ case "$HF_MODEL" in EXPECTED_OUTPUT="Paris" IMAGE_URL="" ;; - solo-text-batching) + dflash-image-batching) # Built by `make muse-glimmer-batching-cuda` into its own directory. RUNNER_PATH="muse-glimmer-batching" RUNNER_TARGET="muse_glimmer_batching_runner" - EXPECTED_OUTPUT="Paris" - IMAGE_URL="" + EXPECTED_OUTPUT="dog" + IMAGE_URL="https://github.com/pytorch/hub/raw/master/images/dog.jpg" ;; dflash-image) RUNNER_TARGET="dflash_runner" @@ -334,7 +334,7 @@ case "$HF_MODEL" in IMAGE_URL="https://github.com/pytorch/hub/raw/master/images/dog.jpg" ;; *) - echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-batching' or 'dflash-image'" + echo "Error: Muse Glimmer requires mode 'solo-text', 'dflash-image-batching' or 'dflash-image'" exit 1 ;; esac @@ -500,12 +500,12 @@ EOF fi RUNNER_ARGS="$RUNNER_ARGS --tokenizer_path ${MODEL_DIR}/$TOKENIZER_FILE --prompt_file \"$PROMPT_FILE\" --max_new_tokens 512" RUNNER_ARGS="$RUNNER_ARGS --cuda_graph" - if [ "$MODE" = "solo-text-batching" ]; then - # The batching runner goes through ServingRuntime and takes its own flags. - printf '%s' '<|start|>user<|message|>What is the capital of France?<|eot|><|start|>assistant' > "$PROMPT_FILE" - RUNNER_ARGS="--pte ${MODEL_DIR}/model.pte --data_path ${MODEL_DIR}/aoti_cuda_blob.ptd --tokenizer ${MODEL_DIR}/$TOKENIZER_FILE --prompt \"\$(cat $PROMPT_FILE)\" --max_new_tokens 128 --temperature 0 --max_session_tokens 4096" + if [ "$MODE" = "dflash-image-batching" ]; then + # The batching runner goes through ServingRuntime and takes its own flags; + # pos_embed.bin is read from beside model.pte. + RUNNER_ARGS="--pte ${MODEL_DIR}/model.pte --data_path ${MODEL_DIR}/aoti_cuda_blob.ptd --tokenizer ${MODEL_DIR}/$TOKENIZER_FILE --prompt \"\$(cat $PROMPT_FILE)\" --max_new_tokens 512 --temperature 0 --max_session_tokens 4096 --dflash" fi - if [ "$MODE" = "dflash-image" ]; then + if [ "$MODE" = "dflash-image" ] || [ "$MODE" = "dflash-image-batching" ]; then RUNNER_ARGS="$RUNNER_ARGS --image_path ${MODEL_DIR}/test_image.jpg" fi ;; @@ -607,14 +607,15 @@ else fi echo "::endgroup::" -if [ "$DEVICE" = "cuda" ] && [ "$MODEL_NAME" = "muse_glimmer" ] && [ "$MODE" = "solo-text-batching" ]; then +if [ "$DEVICE" = "cuda" ] && [ "$MODEL_NAME" = "muse_glimmer" ] && [ "$MODE" = "dflash-image-batching" ]; then echo "::group::Run $MODEL_NAME batching serving smoke" # The CI-only correctness tool is not part of the normal serving build. bash .ci/scripts/build_muse_glimmer_batching_smoke.sh - # Concurrent staggered requests through ServingRuntime on CudaExecutor. + # Concurrent staggered requests through ServingRuntime on the DFlash executor. cmake-out/examples/models/muse-glimmer-batching/muse_glimmer_batching_smoke \ --pte "${MODEL_DIR}/model.pte" --data_path "${MODEL_DIR}/aoti_cuda_blob.ptd" \ --tokenizer "${MODEL_DIR}/$TOKENIZER_FILE" --max_sessions 8 --max_session_tokens 2048 \ + --dflash \ --mode requests --requests 5 --prompt_tokens 600 --max_new_tokens 64 --stagger_ms 200 pip install -r examples/llm_server/python/requirements.txt @@ -639,6 +640,7 @@ PY --hf-tokenizer "${MODEL_DIR}" \ --max-context 4096 \ --prefix-cache-entries 2 \ + --dflash \ --host 127.0.0.1 \ --port "$PORT" >"$SERVER_LOG" 2>&1 & SERVER_PID=$! @@ -651,13 +653,15 @@ PY } trap cleanup_mg_server EXIT - if ! python - "$PORT" "$SERVER_LOG" <<'PY' + if ! python - "$PORT" "$SERVER_LOG" "${MODEL_DIR}/test_image.jpg" <<'PY' +import base64 import json import sys import time import urllib.request +from concurrent.futures import ThreadPoolExecutor -port, log_path = sys.argv[1], sys.argv[2] +port, log_path, image_path = sys.argv[1], sys.argv[2], sys.argv[3] base = f"http://127.0.0.1:{port}" @@ -711,6 +715,38 @@ for line in stream.splitlines(): pieces.append((delta.get("content") or "") + (delta.get("reasoning_content") or "")) if len([p for p in pieces if p]) < 2 or "Paris" not in "".join(pieces): raise AssertionError(f"expected Paris streamed in several chunks, got: {pieces!r}") + +with open(image_path, "rb") as f: + image_url = "data:image/jpeg;base64," + base64.b64encode(f.read()).decode("ascii") +image_body = { + "model": "muse_glimmer", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What animal is in this image?"}, + {"type": "image_url", "image_url": {"url": image_url}}, + ], + } + ], + "max_tokens": 512, + "temperature": 0, +} + + +def answer(payload): + message = json.loads(request("/v1/chat/completions", payload))["choices"][0]["message"] + return (message.get("content") or "") + (message.get("reasoning_content") or "") + + +# Image and text requests in flight together share DFlash executor steps. +with ThreadPoolExecutor(max_workers=3) as pool: + image_a, text, image_b = pool.map(answer, (image_body, body, image_body)) +for reply in (image_a, image_b): + if "dog" not in reply.lower(): + raise AssertionError(f"expected dog in image response, got: {reply!r}") +if "Paris" not in text: + raise AssertionError(f"expected Paris beside image requests, got: {text!r}") print("Muse Glimmer CUDA batching serving smoke passed") PY then diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 296878243df..56dd6071e34 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -1220,7 +1220,7 @@ jobs: - variant: kquant-17gb mode: solo-text - variant: kquant-17gb - mode: solo-text-batching + mode: dflash-image-batching - variant: kquant-17gb mode: dflash-image - variant: kquant-dynamic diff --git a/examples/models/muse-glimmer/BATCHING_SERVING.md b/examples/models/muse-glimmer/BATCHING_SERVING.md index 90254af2722..52dfc2969a8 100644 --- a/examples/models/muse-glimmer/BATCHING_SERVING.md +++ b/examples/models/muse-glimmer/BATCHING_SERVING.md @@ -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 diff --git a/examples/models/muse-glimmer/export/export_dflash.py b/examples/models/muse-glimmer/export/export_dflash.py index 479b8a033a8..6c2eb2da01d 100644 --- a/examples/models/muse-glimmer/export/export_dflash.py +++ b/examples/models/muse-glimmer/export/export_dflash.py @@ -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, diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp index 68ae07d6f2a..b79b28be429 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp @@ -561,10 +561,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")); @@ -690,19 +689,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, @@ -713,6 +704,7 @@ Result 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 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) || @@ -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(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 reject_batch(); } @@ -1467,6 +1474,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 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)) { diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h index 481fdd2a55d..5614d43e806 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h @@ -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 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 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_;