[executorch][ci] Test Muse Glimmer batching on CUDA in OSS CI

.ci/scripts/BUCK

@@ -0,0 +1,13 @@
+load("@fbsource//tools/build_defs:fb_native_wrapper.bzl", "fb_native")
+
+oncall("executorch")
+
+# CMake builds this CI-only tool; group its sources without a product target.
+fb_native.filegroup(
+    name = "muse_glimmer_batching_smoke_sources",
+    srcs = [
+        "muse_glimmer_batching_smoke.cpp",
+        "muse_glimmer_batching_workload.h",
+        "tests/muse_glimmer_batching_workload_test.cpp",
+    ],
+)

.ci/scripts/build_muse_glimmer_batching_smoke.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Run after the normal batching build; reuse its backend and package settings.
+set -euo pipefail
+ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)
+BUILD_DIR="${1:-${ROOT}/cmake-out/examples/models/muse-glimmer-batching}"
+cmake -S "${ROOT}/examples/models/muse-glimmer" -B "$BUILD_DIR" \
+  -DMUSE_GLIMMER_BUILD_BATCHING=ON -DMUSE_GLIMMER_BUILD_CI_SMOKE=ON
+cmake --build "$BUILD_DIR" --target muse_glimmer_batching_smoke --parallel

.ci/scripts/export_model_artifact.sh

@@ -44,6 +44,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: Muse Glimmer DFlash vision mode
 
 Examples:
@@ -95,7 +96,7 @@ if [ -n "$MODE" ]; then
         exit 1
       fi
       ;;
-    solo-text|dflash-image)
+    solo-text|solo-text-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"
@@ -104,7 +105,7 @@ if [ -n "$MODE" ]; then
       ;;
     *)
       echo "Error: Unsupported mode '$MODE'"
-      echo "Supported modes: vr-streaming, vr-offline, solo-text, dflash-image"
+      echo "Supported modes: vr-streaming, vr-offline, solo-text, solo-text-batching, dflash-image"
       exit 1
       ;;
   esac
@@ -288,8 +289,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" != "dflash-image" ]; then
-    echo "Error: Muse Glimmer requires mode 'solo-text' or 'dflash-image'"
+  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'"
     exit 1
   fi
 fi
@@ -565,6 +566,18 @@ 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)
       DRAFT_GGUF_FILE="dflash-Muse-Glimmer-30B-Q4_K_M.gguf"
       MMPROJ_GGUF_FILE="mmproj-Muse-Glimmer-30B-Q4_K_M.gguf"
@@ -581,7 +594,7 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then
           --output-dir "${OUTPUT_DIR}"
       ;;
     *)
-      echo "Error: Muse Glimmer requires mode 'solo-text' or 'dflash-image'"
+      echo "Error: Muse Glimmer requires mode 'solo-text', 'solo-text-batching' or 'dflash-image'"
       exit 1
       ;;
   esac

.ci/scripts/muse_glimmer_batching_smoke.cpp

@@ -0,0 +1,101 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#include "muse_glimmer_batching_workload.h"
+
+#include <executorch/examples/models/muse-glimmer/runtime/backend_factory.h>
+#include <gflags/gflags.h>
+#include <pytorch/tokenizers/hf_tokenizer.h>
+
+#include <iostream>
+#include <stdexcept>
+#include <string>
+#include <utility>
+
+DECLARE_string(pte);
+DECLARE_string(tokenizer);
+DECLARE_string(data_path);
+DECLARE_string(backend);
+DECLARE_int32(max_sessions);
+DECLARE_int32(max_session_tokens);
+DECLARE_uint64(bos_id);
+
+DEFINE_string(mode, "requests", "Correctness workload: requests");
+DEFINE_int32(requests, 8, "Number of concurrent requests");
+DEFINE_int32(prompt_tokens, 512, "Prompt positions including BOS per request");
+DEFINE_int32(max_new_tokens, 128, "Exact completion budget per request");
+DEFINE_int32(stagger_ms, 0, "Delay between request submissions");
+DEFINE_int32(max_decode, 32, "Maximum decodes admitted per step");
+DEFINE_int32(timeout_s, 600, "Request workload deadline");
+
+int main(int argc, char** argv) {
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
+  try {
+    namespace llm = ::executorch::extension::llm;
+    namespace ci = muse_glimmer_ci;
+    if (FLAGS_mode != "requests" || FLAGS_pte.empty() ||
+        FLAGS_tokenizer.empty() || FLAGS_max_sessions <= 0 ||
+        FLAGS_max_session_tokens < 0) {
+      throw std::invalid_argument(
+          "invalid model or correctness workload flags");
+    }
+    const ci::RequestOptions options{
+        FLAGS_requests,
+        FLAGS_prompt_tokens,
+        FLAGS_max_new_tokens,
+        FLAGS_stagger_ms,
+        FLAGS_max_decode,
+        FLAGS_timeout_s};
+    // Reject impossible workloads before loading tokenizer or model weights.
+    ci::validate_requests(
+        options, FLAGS_max_sessions, FLAGS_max_session_tokens);
+    tokenizers::HFTokenizer tokenizer;
+    if (tokenizer.load(FLAGS_tokenizer) != tokenizers::Error::Ok) {
+      throw std::runtime_error("could not load the tokenizer");
+    }
+    llm::MuseGlimmerBackendConfig config;
+    config.backend = FLAGS_backend;
+    config.model_path = FLAGS_pte;
+    config.data_path = FLAGS_data_path;
+    config.max_sessions = FLAGS_max_sessions;
+    config.max_session_tokens = FLAGS_max_session_tokens;
+    config.bos_id = FLAGS_bos_id;
+    auto backend = llm::create_muse_glimmer_backend(config);
+    if (!backend.ok()) {
+      throw std::runtime_error(
+          "could not create the executor: error " +
+          std::to_string(static_cast<int>(backend.error())));
+    }
+    ci::EvidenceExecutor executor(*backend->executor);
+    llm::serving::ServingRuntimeConfig serving;
+    serving.max_sessions = options.requests;
+    serving.max_requests = options.requests;
+    serving.max_pending_operations = options.requests;
+    serving.max_context_length = backend->preparation->max_context_length;
+    auto result = ci::run_requests(
+        options,
+        executor,
+        tokenizer,
+        serving,
+        [spec = backend->preparation](
+            const llm::serving::PromptPreparationContext& context,
+            llm::serving::PromptInput input) {
+          return llm::prepare_muse_glimmer_input(
+              context, std::move(input), spec);
+        });
+    const auto evidence = executor.evidence();
+    ci::check_batching(evidence);
+    std::cout << "Batching correctness passed: " << result.accepted
+              << " requests completed with exact token budgets; "
+              << evidence.multi_session_steps
+              << " successful multi-session executor steps\n";
+    return 0;
+  } catch (const std::exception& error) {
+    std::cerr << "batching smoke failed: " << error.what() << '\n';
+    return 1;
+  }
+}