[executorch][muse] Draft CUDA DFlash decode from pending tokens (#24041)

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

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

Author: pytorchbot

.github/workflows/cuda.yml

@@ -607,7 +607,8 @@ jobs:
         cmake -S "${PWD}/examples/models/muse-glimmer" -B "${BATCHING_BIN_DIR}" \
           -DMUSE_GLIMMER_BUILD_BATCHING_TESTS=ON
         bash .ci/scripts/build_muse_glimmer_batching_smoke.sh "${BATCHING_BIN_DIR}"
-        cmake --build "${BATCHING_BIN_DIR}" --target muse_glimmer_batching_test --parallel
+        cmake --build "${BATCHING_BIN_DIR}" \
+          --target muse_glimmer_batching_test muse_glimmer_cuda_dflash_test --parallel
         ctest --test-dir "${BATCHING_BIN_DIR}" \
           -R '^muse_glimmer_batching$' --output-on-failure --no-tests=error
         test -x "${BATCHING_BIN_DIR}/muse_glimmer_batching_runner"
@@ -619,8 +620,23 @@ jobs:
         python -m pytest -v -o "addopts=" \
           examples/models/muse-glimmer/tests/test_cuda_batching_pipeline.py \
           examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py \
+          examples/models/muse-glimmer/tests/test_dflash_token_sampler.py \
           examples/models/muse-glimmer/tests/test_checkpoint_context.py
 
+        # CUDA DFlash executor on a tiny exported target + draft. The native
+        # tests skip without the artifact, so fail if any skipped.
+        export ET_MUSE_DFLASH_TINY_DIR="${RUNNER_TEMP:-/tmp}/muse_glimmer_dflash_tiny"
+        python examples/models/muse-glimmer/tests/export_cuda_dflash_tiny.py \
+          --output-dir "${ET_MUSE_DFLASH_TINY_DIR}"
+        ctest --test-dir "${BATCHING_BIN_DIR}" -R '^muse_glimmer_cuda_dflash$' \
+          --output-on-failure --no-tests=error -V > /tmp/muse_glimmer_cuda_dflash.log 2>&1 \
+          || { cat /tmp/muse_glimmer_cuda_dflash.log; exit 1; }
+        cat /tmp/muse_glimmer_cuda_dflash.log
+        if grep -q "SKIPPED" /tmp/muse_glimmer_cuda_dflash.log; then
+          echo "::error::muse_glimmer_cuda_dflash_test skipped"
+          exit 1
+        fi
+
   unittest-cuda-quantized-gemm:
     name: unittest-cuda-quantized-gemm (${{ matrix.shard }})
     needs: [changed-files, run-decision]

examples/models/muse-glimmer/CMakeLists.txt

@@ -366,11 +366,17 @@ if(MUSE_GLIMMER_BUILD_BATCHING OR MUSE_GLIMMER_BUILD_BATCHING_TESTS)
       target_compile_definitions(gtest PUBLIC GTEST_HAS_POSIX_RE=0)
       add_executable(
         muse_glimmer_cuda_dflash_test tests/cuda_dflash_executor_test.cpp
+                                      tests/cuda_dflash_native_test.cpp
       )
       target_link_libraries(
         muse_glimmer_cuda_dflash_test PRIVATE muse_glimmer_batching_runtime
                                               GTest::gtest_main
       )
+      if(NOT APPLE AND NOT MSVC)
+        target_link_options(
+          muse_glimmer_cuda_dflash_test PRIVATE "LINKER:--export-dynamic"
+        )
+      endif()
       add_test(NAME muse_glimmer_cuda_dflash
                COMMAND muse_glimmer_cuda_dflash_test
       )

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

@@ -635,6 +635,7 @@ std::optional<SessionId> MuseGlimmerCudaDFlashExecutor::clone(
 #endif
     State copy{*target_id, *draft_id};
     copy.frontier = upto;
+    copy.pending = upto == state.frontier ? state.pending : std::nullopt;
     copy.hidden = state.hidden;
     if (!df::trim_hidden(copy.hidden, row_bytes_, draft_position, upto)) {
       target_->cache_control().seq_rm(*target_id);
@@ -680,6 +681,7 @@ bool MuseGlimmerCudaDFlashExecutor::rewind(State& state, int position) {
     return false;
   }
   state.frontier = position;
+  state.pending.reset();
   return true;
 }
 
@@ -887,7 +889,8 @@ bool MuseGlimmerCudaDFlashExecutor::draft(
 
 bool MuseGlimmerCudaDFlashExecutor::speculate(
     const batching::BatchInput& batch,
-    batching::BatchOutput& out) {
+    batching::BatchOutput& out,
+    bool pending_decode) {
   std::size_t cursor = 0;
   while (cursor < batch.inputs.size()) {
     std::vector<SessionId> owners;
@@ -899,8 +902,9 @@ bool MuseGlimmerCudaDFlashExecutor::speculate(
       auto& state = sessions_.at(input.sid);
       const auto end =
           static_cast<std::int64_t>(input.position) + input.offset + input.size;
-      if (!out.outputs[cursor] || end != state.frontier ||
-          end > max_context_ - block_) {
+      const int seed_rows = pending_decode ? 1 : 0;
+      if (!out.outputs[cursor] || end != state.frontier + seed_rows ||
+          end - seed_rows > max_context_ - block_) {
         ++cursor;
         continue;
       }
@@ -917,7 +921,10 @@ bool MuseGlimmerCudaDFlashExecutor::speculate(
       }
       context += static_cast<int>(rows);
       owners.push_back(input.sid);
-      anchors.push_back(out.outputs[cursor]->tokens.front());
+      anchors.push_back(
+          pending_decode ? (*std::get<batching::TokenInputPtr>(
+                               input.payload))[input.offset]
+                         : out.outputs[cursor]->tokens.front());
       indices.push_back(cursor++);
     }
     if (owners.empty()) {
@@ -982,9 +989,9 @@ bool MuseGlimmerCudaDFlashExecutor::speculate(
         return false;
       }
       state.frontier = committed;
-      // The regular preamble's prediction is new output, now committed by
-      // verification. Append only predictions after that anchor. The last
-      // prediction remains uncommitted: frontier = input.end + out.size() - 1.
+      // Pending decode consumes Input's already-emitted token; preamble mode
+      // consumes its new anchor. Both append only new predictions, leaving
+      // correction/bonus pending: frontier = input.end + out.size() - 1.
       auto& output_tokens = out.outputs[indices[i]]->tokens;
       output_tokens.insert(
           output_tokens.end(), burst.tokens.begin(), burst.tokens.end());
@@ -1009,6 +1016,7 @@ bool MuseGlimmerCudaDFlashExecutor::fail_batch(
     failed_ = true;
     for (auto& item : sessions_) {
       item.second.poisoned = true;
+      item.second.pending.reset();
       if (fenced) {
         item.second.hidden.clear();
       }
@@ -1045,6 +1053,7 @@ bool MuseGlimmerCudaDFlashExecutor::execute(
     std::vector<SessionId> token_owners;
     std::unordered_map<SessionId, int> cursor;
     std::vector<std::pair<SessionId, int>> rewinds;
+    bool pending_decode = !batch.inputs.empty();
     for (const auto& input : batch.inputs) {
       const auto it = sessions_.find(input.sid);
       const auto* payload =
@@ -1082,6 +1091,13 @@ bool MuseGlimmerCudaDFlashExecutor::execute(
       if (start < at->second) {
         rewinds.emplace_back(input.sid, static_cast<int>(start));
       }
+      // 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 &&
+          *it->second.pending == (**payload)[input.offset];
       at->second = static_cast<int>(end);
       for (std::size_t j = 0; j < input.size; ++j) {
         const auto token = (**payload)[input.offset + j];
@@ -1104,19 +1120,49 @@ bool MuseGlimmerCudaDFlashExecutor::execute(
         return fail_batch(batch, out);
       }
     }
-    std::vector<Token> predictions;
-    if (!run_target(step, token_owners, predictions, true)) {
+    if (pending_decode) {
+      ++stats_.pending_decode_batches;
+      stats_.pending_decode_owners += batch.inputs.size();
+      for (std::size_t i = 0; i < batch.inputs.size(); ++i) {
+        out.outputs[i] = batching::Output{batch.inputs[i].sid, {}};
+      }
+    } else {
+      std::vector<Token> predictions;
+      if (!run_target(step, token_owners, predictions, true)) {
+        return fail_batch(batch, out);
+      }
+      for (std::size_t i = 0; i < batch.inputs.size(); ++i) {
+        sessions_.at(batch.inputs[i].sid).pending.reset();
+        if (batch.inputs[i].produce_output) {
+          out.outputs[i] =
+              batching::Output{batch.inputs[i].sid, {predictions[i]}};
+        }
+      }
+    }
+    if (!speculate(batch, out, pending_decode)) {
       return fail_batch(batch, out);
     }
     for (std::size_t i = 0; i < batch.inputs.size(); ++i) {
-      if (batch.inputs[i].produce_output) {
-        out.outputs[i] =
-            batching::Output{batch.inputs[i].sid, {predictions[i]}};
+      if (out.outputs[i]) {
+        const auto& tokens = out.outputs[i]->tokens;
+        const auto& input = batch.inputs[i];
+        auto& state = sessions_.at(input.sid);
+        const auto end = static_cast<std::int64_t>(input.position) +
+            input.offset + input.size;
+        if (tokens.empty()) {
+          return fail_batch(batch, out);
+        }
+        const auto expected =
+            end + static_cast<std::int64_t>(tokens.size()) - 1;
+        if (pending_decode && expected != state.frontier) {
+          return fail_batch(batch, out);
+        }
+        // Only the final input for a repeated owner may seed the next cycle.
+        if (expected == state.frontier) {
+          state.pending = tokens.back();
+        }
       }
     }
-    if (!speculate(batch, out)) {
-      return fail_batch(batch, out);
-    }
     ++stats_.executions;
     return true;
 #if ET_HAS_EXCEPTIONS

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

@@ -97,6 +97,8 @@ struct CudaDFlashStats {
   std::uint64_t 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;
+  std::uint64_t pending_decode_owners = 0;
   std::uint64_t target_preamble_calls = 0;
   std::uint64_t target_verify_calls = 0;
   std::uint64_t target_preamble_forwards = 0;
@@ -145,6 +147,7 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor {
   }
 
  private:
+  friend struct CudaDFlashTestPeer;
   using Target = ::executorch::backends::cuda::batching::CudaExecutor;
   using RawStep = ::executorch::backends::cuda::batching::RawTargetStep;
   struct State {
@@ -153,6 +156,8 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor {
     std::int32_t target_id;
     std::int32_t draft_id;
     int frontier = 0;
+    // Last emitted correction/bonus, not yet consumed by the target cache.
+    std::optional<batching::Token> pending;
     // Host-owned committed BF16 rows [draft frontier, target frontier).
     // Verification appends only the consumed anchor and accepted proposals.
     std::vector<std::uint8_t> hidden;
@@ -183,7 +188,10 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor {
       const std::vector<batching::Token>& anchors,
       int context_limit,
       std::vector<std::vector<batching::Token>>* proposals);
-  bool speculate(const batching::BatchInput& batch, batching::BatchOutput& out);
+  bool speculate(
+      const batching::BatchInput& batch,
+      batching::BatchOutput& out,
+      bool pending_decode);
   bool fail_batch(
       const batching::BatchInput& batch,
       batching::BatchOutput& out);