[executorch][muse] Run the CUDA DFlash drafter in CUDA graphs (#24043)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: #23985 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/84/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/84/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/83/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/84/orig
Differential Revision:
[D125705719](https://our.internmc.facebook.com/intern/diff/D125705719/)
@diff-train-skip-merge

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

Author: pytorchbot

backends/cuda/batching/cuda_executor.cpp

@@ -198,6 +198,9 @@ Error set_backend_options(
       }
     }
   }
+  for (const std::string& name : options.additional_cuda_graph_methods) {
+    graph_methods += (graph_methods.empty() ? "" : ",") + name;
+  }
   ET_CHECK_OR_RETURN_ERROR(
       graph_methods.size() < ::executorch::runtime::kMaxOptionValueLength,
       InvalidArgument,
@@ -288,7 +291,7 @@ Result<std::unique_ptr<CudaExecutor>> CudaExecutor::create(
   config.max_session_tokens = max_session_tokens;
   config.kv_dtype = kv_dtype;
   config.initial_capacity = initial_capacity;
-  config.options = options;
+  config.options = std::move(options);
   return create(std::move(module), config);
 }
 

backends/cuda/batching/cuda_executor.h

@@ -65,6 +65,8 @@ struct CudaExecutorOptions {
   // Capture every static method into a CUDA graph. forward_others stays
   // eager: its width varies.
   bool cuda_graph_for_static_methods = true;
+  // Other static methods of the program to capture, e.g. a model adapter's.
+  std::vector<std::string> additional_cuda_graph_methods;
 };
 
 struct CudaExecutorConfig {

backends/cuda/runtime/cuda_kv_cache.cpp

@@ -451,10 +451,11 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache {
         "offgraph_kv: a %lld-token step is outside [1, %d]",
         static_cast<long long>(write_length),
         max_write_);
+    // A static draft program pads its tail with scratch rows, as a static
+    // target forward does; every query row stays a declared block row.
     const int tokens = static_cast<int>(step_seq_ids_.size());
     ET_CHECK_OR_RETURN_ERROR(
-        tokens > 0 && tokens <= width &&
-            (!dflash_block_size_ || tokens == width),
+        tokens > 0 && tokens <= width,
         InvalidState,
         "offgraph_kv: a %d-token program cannot run %d declared tokens",
         width,

backends/cuda/runtime/test/test_cuda_kv_cache.cpp

@@ -1927,7 +1927,7 @@ TEST_F(CudaCellCacheTest, DFlashUploadsKVWideCellsAndQueryWideOwnerMasks) {
 
 TEST_F(
     CudaCellCacheTest,
-    DFlashRejectsGroupingAndPaddedTokenCountsBeforeAllocation) {
+    DFlashRejectsGroupingAndNarrowProgramsBeforeAllocation) {
   auto cache_ptr = make(4, 2);
   auto* control = cache_ptr->as<cache::BatchControl>();
   auto* kv = cache_ptr->as<cu::CudaKVCache>();
@@ -1938,7 +1938,6 @@ TEST_F(
   EXPECT_FALSE(control->declare_step({a, a, b, b, a, a}));
   EXPECT_FALSE(control->declare_step({a, a, b}));
   ASSERT_TRUE(control->declare_step({b, b, a, a}));
-  EXPECT_EQ(kv->prepare_step(5, cudaStreamPerThread), Error::InvalidState);
   EXPECT_EQ(kv->prepare_step(3, cudaStreamPerThread), Error::InvalidState);
   EXPECT_EQ(kv->metrics().allocated_bytes, 0);
   EXPECT_EQ(control->pos(a), 0);
@@ -1977,6 +1976,41 @@ TEST_F(
   kv->forget_handle(&handle);
 }
 
+TEST_F(CudaCellCacheTest, DFlashStaticProgramPadsTailIntoScratch) {
+  auto cache_ptr = make(4, 2);
+  ASSERT_NE(cache_ptr, nullptr);
+  auto* control = cache_ptr->as<cache::BatchControl>();
+  auto* kv = cache_ptr->as<cu::CudaKVCache>();
+  auto fake = container();
+  auto handle = make_handle(fake);
+  ASSERT_TRUE(kv->note_handle(&handle).get());
+  const int32_t a = *control->seq_new(), b = *control->seq_new();
+  // One context row for a, none for b, padded to an eight-row program.
+  ASSERT_TRUE(control->declare_step({a, a, a, b, b}));
+  ASSERT_EQ(kv->prepare_step(8, cudaStreamPerThread), Error::Ok);
+  ASSERT_EQ(kv->rebind_for_execute(&handle), Error::Ok);
+  const auto cells = read_longs(fake.bound["cells"].data, 8);
+  EXPECT_EQ(
+      std::vector<int64_t>(cells.begin(), cells.begin() + 5),
+      std::vector<int64_t>({0, 1, 2, 3, 4}));
+  for (int row = 5; row < 8; ++row) {
+    EXPECT_EQ(cells[row], cells[7]) << row;
+    EXPECT_GE(cells[row], 5) << row;
+  }
+  using Mask = std::vector<std::vector<int>>;
+  EXPECT_EQ(
+      read_mask(fake.bound["mask0"].data, 4, 5),
+      Mask(
+          {{1, 1, 1, 0, 0},
+           {1, 1, 1, 0, 0},
+           {0, 0, 0, 1, 1},
+           {0, 0, 0, 1, 1}}));
+  ASSERT_EQ(kv->commit_step(8), Error::Ok);
+  EXPECT_EQ(control->pos(a), 1);
+  EXPECT_EQ(control->pos(b), 0);
+  kv->forget_handle(&handle);
+}
+
 TEST_F(CudaCellCacheTest, DFlashWarmupCommitsAllContextWithOneQueryPerOwner) {
   auto cache_ptr = make(4, 0);
   auto* control = cache_ptr->as<cache::BatchControl>();

examples/models/muse-glimmer/export/export_dflash.py

@@ -95,6 +95,7 @@ def _export_cuda_sampler_methods(
     }
 
 
+STATIC_DRAFT_PREFIX = "draft_forward_"
 VERIFY_METHOD = "dflash_verify"
 VERIFY_ARGMAX_METHOD = "dflash_verify_argmax"
 
@@ -382,6 +383,28 @@ def prepare_cuda_packed_draft(
             strict=True,
         )
     programs = {"draft_forward": program}
+    # Static decode drafts for CUDA graphs: N owners, each with at most one
+    # block of fresh context. The runtime pads unused context rows into the
+    # cache's scratch cell, as forward_N pads unused tokens.
+    with torch.no_grad():
+        for count in range(1, max_owners + 1):
+            width = count * block
+            # Context and queries write 2 * width K/V rows into one step.
+            if 2 * width > max_write:
+                break
+            programs[f"{STATIC_DRAFT_PREFIX}{count}"] = export(
+                combined,
+                (
+                    torch.zeros((count, block), dtype=torch.long),
+                    torch.zeros(
+                        (1, width, config.n_target_layers * config.dim),
+                        dtype=torch.bfloat16,
+                    ),
+                    torch.arange(2 * width, dtype=torch.long),
+                    torch.arange(2 * width, dtype=torch.long),
+                ),
+                strict=True,
+            )
     geometry = offgraph_kv_cache_geometry(draft_model)
     metadata = {
         key.replace("get_", "get_draft_", 1): value for key, value in geometry.items()

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

@@ -143,6 +143,41 @@ std::optional<DraftPlan> pack_draft(
   return plan;
 }
 
+bool pad_to_static(
+    DraftPlan& plan,
+    int owners,
+    int block,
+    std::size_t row_bytes) {
+  const int width = owners * block;
+  const int queries = width;
+  if (row_bytes == 0 || plan.hidden.size() % row_bytes != 0) {
+    return false;
+  }
+  const int context = static_cast<int>(plan.hidden.size() / row_bytes);
+  if (context > width ||
+      plan.positions.size() != static_cast<std::size_t>(context + queries) ||
+      plan.kv_order.size() != plan.positions.size()) {
+    return false;
+  }
+  for (const int count : plan.counts) {
+    if (count > block) {
+      return false;
+    }
+  }
+  const int pad = width - context;
+  plan.hidden.resize(static_cast<std::size_t>(width) * row_bytes, 0);
+  plan.positions.insert(plan.positions.begin() + context, pad, 0);
+  for (auto& index : plan.kv_order) {
+    if (index >= context) {
+      index += pad;
+    }
+  }
+  for (int row = context; row < width; ++row) {
+    plan.kv_order.push_back(row);
+  }
+  return true;
+}
+
 DeviceStaging::~DeviceStaging() {
   if (device_) {
     (void)cudaFree(device_);
@@ -443,6 +478,10 @@ Result<MuseGlimmerBackend> MuseGlimmerCudaDFlashExecutor::create(
   }
   target_config.max_sessions = config.max_sessions;
   target_config.max_session_tokens = config.max_session_tokens;
+  for (int n = 1; n <= df::kOwnerLimit; ++n) {
+    target_config.options.additional_cuda_graph_methods.push_back(
+        df::kStaticForwardPrefix + std::to_string(n));
+  }
   ET_ASSIGN_OR_RETURN(target, Target::create(target_config));
   auto& module = target->module();
   ET_ASSIGN_OR_RETURN(names, module.method_names());
@@ -629,6 +668,43 @@ bool MuseGlimmerCudaDFlashExecutor::initialize() {
         "CUDA DFlash: could not initialize distinct target/draft caches");
     return false;
   }
+  const auto names = module.method_names();
+  if (!names.ok()) {
+    ET_LOG(Error, "CUDA DFlash: could not list the program's methods");
+    return false;
+  }
+  for (int n = 1; n <= max_owners_; ++n) {
+    const auto name = df::kStaticForwardPrefix + std::to_string(n);
+    if (!names->count(name)) {
+      continue;
+    }
+    const auto meta = module.method_meta(name);
+    const int width = n * block_;
+    auto shaped = [&](std::size_t index, bool input, std::vector<int> sizes) {
+      auto tensor = input ? meta->input_tensor_meta(index)
+                          : meta->output_tensor_meta(index);
+      return tensor.ok() &&
+          std::equal(
+                 tensor->sizes().begin(),
+                 tensor->sizes().end(),
+                 sizes.begin(),
+                 sizes.end());
+    };
+    if (!meta.ok() || meta->num_inputs() != 4 || meta->num_outputs() != 1 ||
+        !shaped(0, true, {n, block_}) ||
+        !shaped(1, true, {1, width, hidden_width_}) ||
+        !shaped(2, true, {2 * width}) || !shaped(3, true, {2 * width}) ||
+        !shaped(0, false, {1, width, target_->vocab_size()}) ||
+        batching::load_method_with_cache(
+            module,
+            name,
+            ::executorch::backends::cuda::kCudaBackendId,
+            draft_install_) != Error::Ok) {
+      ET_LOG(Error, "CUDA DFlash: invalid static draft %s", name.c_str());
+      return false;
+    }
+    static_drafts_[n] = true;
+  }
   initialized_ = true;
   return true;
 }
@@ -1022,23 +1098,39 @@ bool MuseGlimmerCudaDFlashExecutor::draft(
   }
   const int q = static_cast<int>(owners.size()) * block_;
   const int c = static_cast<int>(plan->hidden.size() / row_bytes_);
-  const int k = c + q;
+  // Decode contexts fit one block per owner; take the static, graph-captured
+  // draft. Longer contexts drain through the dynamic one.
+  const bool use_static = proposals && static_drafts_[owners.size()] &&
+      std::all_of(plan->counts.begin(), plan->counts.end(), [this](int n) {
+                            return n <= block_;
+                          });
+  if (use_static &&
+      !df::pad_to_static(
+          *plan, static_cast<int>(owners.size()), block_, row_bytes_)) {
+    return false;
+  }
+  const int padded = static_cast<int>(plan->hidden.size() / row_bytes_);
+  const int k = padded + q;
   // Host-owned inputs are staged by the export's normal H2D boundary. Keep
   // their backing alive through the forward and its stream-fenced argmax.
   auto tokens = make_tensor_ptr(
       {static_cast<int>(owners.size()), block_}, std::move(plan->tokens));
   auto hidden = make_tensor_ptr(
-      {1, c, hidden_width_},
+      {1, padded, hidden_width_},
       std::move(plan->hidden),
       aten::ScalarType::BFloat16);
   auto positions = make_tensor_ptr({k}, std::move(plan->positions));
   auto order = make_tensor_ptr({k}, std::move(plan->kv_order));
-  auto result = target_->module().execute(
-      df::kForward, {tokens, hidden, positions, order});
+  const std::string method = use_static
+      ? df::kStaticForwardPrefix + std::to_string(owners.size())
+      : std::string(df::kForward);
+  auto result =
+      target_->module().execute(method, {tokens, hidden, positions, order});
   if (!result.ok() || result->size() != 1 || !result->at(0).isTensor()) {
     return false;
   }
   ++stats_.draft_forwards;
+  stats_.static_draft_forwards += use_static;
   const auto logits = result->at(0).toTensor();
   if (logits.dim() < 2 || logits.scalar_type() != aten::ScalarType::Float ||
       !dense(logits) ||

examples/models/muse-glimmer/runtime/cuda_dflash_executor.h

@@ -10,6 +10,7 @@
 #include <executorch/backends/cuda/batching/cuda_executor.h>
 #include <executorch/examples/models/muse-glimmer/runtime/backend_factory.h>
 
+#include <array>
 #include <unordered_map>
 
 namespace executorch::extension::llm {
@@ -24,6 +25,8 @@ inline constexpr char kTargetLayers[] = "get_n_target_layers";
 inline constexpr char kMaxOwners[] = "get_draft_max_owners";
 inline constexpr char kMaxWrite[] = "get_draft_max_step_tokens";
 inline constexpr char kMaxCells[] = "get_draft_offgraph_kv_max_cells";
+// Static draft_forward_{N}: N owners, each with at most one block of context.
+inline constexpr char kStaticForwardPrefix[] = "draft_forward_";
 // On-device verification of [owners, width] candidates; see
 // model/dflash_token_sampler.py.
 inline constexpr char kVerify[] = "dflash_verify";
@@ -104,6 +107,14 @@ std::optional<DraftPlan> pack_draft(
     int context_limit,
     int max_context,
     std::size_t row_bytes);
+// Widens a plan whose context fits one block per owner to the static
+// draft_forward_{N} shape: unused context rows go last, after every block, so
+// the cache writes them to its scratch cell and no query attends them.
+bool pad_to_static(
+    DraftPlan& plan,
+    int owners,
+    int block,
+    std::size_t row_bytes);
 
 // Reusable device memory for small host-built inputs.
 class DeviceStaging {
@@ -129,6 +140,7 @@ class DeviceStaging {
 struct CudaDFlashStats {
   std::uint64_t executions = 0;
   std::uint64_t draft_forwards = 0;
+  std::uint64_t static_draft_forwards = 0;
   std::uint64_t warmed_context_tokens = 0;
   // Calls are packed execute_raw attempts; forwards count slice callbacks.
   std::uint64_t pending_decode_batches = 0;
@@ -148,9 +160,11 @@ struct CudaDFlashStats {
   std::uint64_t failed_batches = 0;
 };
 
-// Text-only model adapter over the existing CUDA target executor. Drafts are
-// greedy; dflash_verify accepts or resamples them under each session's sampling
-// on the device, so only token IDs return to the host.
+// Text-only model adapter over the existing CUDA target executor. 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
+// sampling on the device, so only token IDs return to the host.
 class MuseGlimmerCudaDFlashExecutor final : public batching::Executor {
  public:
   static runtime::Result<MuseGlimmerBackend> create(
@@ -253,6 +267,8 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor {
   CudaDFlashStats stats_;
   cuda_dflash_detail::HiddenSnapshot verification_hidden_;
   cuda_dflash_detail::DeviceStaging verify_inputs_;
+  // static_drafts_[n]: the program has draft_forward_{n}.
+  std::array<bool, cuda_dflash_detail::kOwnerLimit + 1> static_drafts_{};
   bool initialized_ = false;
   // A failed co-batch invalidates every draft owner. Only model restart
   // recovers.

examples/models/muse-glimmer/tests/cuda_dflash_executor_test.cpp

@@ -61,6 +61,36 @@ TEST(CudaDFlash, PositionalUniformsAreSeededAndIndependentOfBatching) {
   EXPECT_NE(u, df::positional_uniform(7, 100, 3));
 }
 
+TEST(CudaDFlash, StaticPaddingPutsUnusedContextAfterEveryBlock) {
+  const std::vector<std::uint8_t> a{1, 1, 2, 2}, b{3, 3};
+  auto plan =
+      df::pack_draft({{9, 10, 70, &a, 2}, {3, 20, 80, &b, 1}}, 4, 99, 3, 64, 2);
+  ASSERT_TRUE(plan);
+  const auto ids = plan->ids;
+  ASSERT_TRUE(df::pad_to_static(*plan, 2, 4, 2));
+  EXPECT_EQ(plan->ids, ids);
+  EXPECT_EQ(
+      plan->hidden,
+      std::vector<std::uint8_t>(
+          {1, 1, 2, 2, 3, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}));
+  // Context rows 0..2, padding 3..7, queries 8..15.
+  EXPECT_EQ(
+      plan->positions,
+      std::vector<std::int64_t>(
+          {10, 11, 20, 0, 0, 0, 0, 0, 12, 13, 14, 15, 21, 22, 23, 24}));
+  EXPECT_EQ(
+      plan->kv_order,
+      std::vector<std::int64_t>(
+          {0, 1, 8, 9, 10, 11, 2, 12, 13, 14, 15, 3, 4, 5, 6, 7}));
+}
+
+TEST(CudaDFlash, StaticPaddingRejectsLongContext) {
+  const std::vector<std::uint8_t> a(2 * 5, 1);
+  auto plan = df::pack_draft({{9, 10, 70, &a, 5}}, 4, 99, 5, 64, 2);
+  ASSERT_TRUE(plan);
+  EXPECT_FALSE(df::pad_to_static(*plan, 1, 4, 2));
+}
+
 TEST(CudaDFlash, NewPreambleAnchorIsRetainedOnce) {
   // The input token (7) has already committed. The regular preamble predicts
   // 10, then verification consumes 10 and predicts the accepted/corrected tail.