[executorch][muse] Add cuda dflash executor (#24039)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: #23981 by
@Gasoonjia
^ Please use this as the source of truth for the PR details, comments,
and reviews
ghstack PR base:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/80/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/80/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/79/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/80/orig
Differential Revision:
[D125705713](https://our.internmc.facebook.com/intern/diff/D125705713/)
@diff-train-skip-merge

Co-authored-by: gasoonjia <gasoonjia@icloud.com>

Author: pytorchbot

examples/models/muse-glimmer/BATCHING_SERVING.md

@@ -3,7 +3,8 @@
 The batching runner and multiplexed worker support MLX autoregressive generation
 with text only or text plus one image across the submitted conversation, and
 CUDA autoregressive generation with text only. They are opt-in entry points; the
-legacy runners and export defaults remain unchanged. DFlash batching is MLX only.
+legacy runners and export defaults remain unchanged. DFlash batching runs on MLX
+and CUDA (text only).
 CUDA specifics are in [CUDA](#cuda) below.
 
 ## Export
@@ -227,3 +228,10 @@ snapshots are enabled. Their total reservation must leave one cell for padded
 writes within the exported `--max-cells`.
 The cell pool tracks at most 64 resident sessions. Prefix snapshots clone a
 session's cells by reference, sharing the prompt's KV rather than copying it.
+
+For DFlash, export the tapped target and packed draft with
+`export_dflash --backend cuda --use-offgraph-kv-cache` (add `--draft-max-cells`
+and `--max-draft-owners`), then add `--dflash` to the serve command above. The
+draft ABI and `--dflash-n-draft` behave as on MLX. Drafts are greedy; the target
+accepts or resamples them on the device under each request's sampling. Decode
+drafts of up to four owners run in CUDA graphs.

examples/models/muse-glimmer/CMakeLists.txt

@@ -242,7 +242,7 @@ if(MUSE_GLIMMER_BUILD_BATCHING OR MUSE_GLIMMER_BUILD_BATCHING_TESTS)
 
   if(MUSE_GLIMMER_BUILD_BATCHING)
     if(_batching_backend STREQUAL "cuda")
-      set(_batching_executor)
+      set(_batching_executor runtime/cuda_dflash_executor.cpp)
       set(_batching_backend_library cuda_batching)
     else()
       set(_batching_executor runtime/mlx_executor.cpp)
@@ -356,5 +356,24 @@ if(MUSE_GLIMMER_BUILD_BATCHING OR MUSE_GLIMMER_BUILD_BATCHING_TESTS)
                                          extension_llm_serving Threads::Threads
     )
     add_test(NAME muse_glimmer_batching COMMAND muse_glimmer_batching_test)
+    if(MUSE_GLIMMER_BUILD_BATCHING AND _batching_backend STREQUAL "cuda")
+      # Tokenizer include paths contain regex.h; keep GoogleTest and its users
+      # on the same simple-regex ABI instead of the shadowed POSIX header.
+      add_subdirectory(
+        ${EXECUTORCH_ROOT}/third-party/googletest
+        ${CMAKE_CURRENT_BINARY_DIR}/dflash_googletest EXCLUDE_FROM_ALL
+      )
+      target_compile_definitions(gtest PUBLIC GTEST_HAS_POSIX_RE=0)
+      add_executable(
+        muse_glimmer_cuda_dflash_test tests/cuda_dflash_executor_test.cpp
+      )
+      target_link_libraries(
+        muse_glimmer_cuda_dflash_test PRIVATE muse_glimmer_batching_runtime
+                                              GTest::gtest_main
+      )
+      add_test(NAME muse_glimmer_cuda_dflash
+               COMMAND muse_glimmer_cuda_dflash_test
+      )
+    endif()
   endif()
 endif()

examples/models/muse-glimmer/runtime/BUCK

@@ -0,0 +1,34 @@
+load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target")
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+oncall("executorch")
+
+fbcode_target(
+    _kind = runtime.cxx_library,
+    name = "cuda_dflash_executor",
+    srcs = ["cuda_dflash_executor.cpp"],
+    header_namespace = "executorch/examples/models/muse-glimmer/runtime",
+    exported_headers = [
+        "backend_factory.h",
+        "cuda_dflash_executor.h",
+        "engine/muse_glimmer_vision_runtime.h",
+        "prompt_preparation.h",
+    ],
+    visibility = ["PUBLIC"],
+    deps = [
+        "//executorch/extension/cuda:caller_stream",
+        "//executorch/extension/llm/runner:stats",
+        "//executorch/runtime/backend:interface",
+    ],
+    exported_deps = [
+        "fbsource//third-party/nlohmann-json:nlohmann-json",
+        "//executorch/backends/cuda/batching:cuda_executor",
+        "//executorch/extension/llm/batching:batching",
+        "//executorch/extension/llm/batching:executor_utils",
+        "//executorch/extension/llm/cache:kv_cache",
+        "//executorch/extension/llm/serving:serving_runtime",
+        "//executorch/extension/module:module",
+        "//executorch/extension/tensor:tensor",
+    ],
+    external_deps = [("cuda", None, "cuda-lazy"), "nlohmann_json"],
+)

examples/models/muse-glimmer/runtime/backend_factory.cpp

@@ -11,6 +11,7 @@
 #endif
 #if defined(EXECUTORCH_BUILD_CUDA)
 #include <executorch/backends/cuda/batching/cuda_executor.h>
+#include <executorch/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h>
 #endif
 
 #include <memory>
@@ -32,9 +33,12 @@ runtime::Result<MuseGlimmerBackend> create_muse_glimmer_backend(
         InvalidArgument,
         "Muse Glimmer CUDA: a session needs room for a prompt and prediction");
     ET_CHECK_OR_RETURN_ERROR(
-        !config.use_dflash && config.dflash_n_draft == 0,
-        NotSupported,
-        "Muse Glimmer CUDA: DFlash batching is not supported");
+        config.use_dflash || config.dflash_n_draft == 0,
+        InvalidArgument,
+        "Muse Glimmer CUDA: dflash_n_draft requires DFlash");
+    if (config.use_dflash) {
+      return MuseGlimmerCudaDFlashExecutor::create(config);
+    }
     cb::CudaExecutorConfig executor_config;
     executor_config.model_path = config.model_path;
     if (!config.data_path.empty()) {