[executorch][muse] Verify CUDA DFlash drafts on device under any sampling

backends/cuda/batching/cuda_executor.cpp

@@ -810,21 +810,40 @@ bool CudaExecutor::execute(const BatchInput& batch, BatchOutput& out) {
   });
 }
 
+std::optional<std::vector<Token>> CudaExecutor::sample_rows(
+    const ::executorch::aten::Tensor& logits,
+    std::size_t rows,
+    const std::vector<float>& params) {
+  if (params.size() != rows * 4) {
+    return std::nullopt;
+  }
+  bool all_greedy = true;
+  for (std::size_t row = 0; row < rows; ++row) {
+    all_greedy = all_greedy && params[row * 4] == 0.0f;
+  }
+  if (all_greedy) {
+    return argmax_rows(logits, rows);
+  }
+  if (logits.device_type() != ::executorch::aten::DeviceType::CUDA) {
+    ET_LOG(
+        Error,
+        "CudaExecutor: sampled rows need device logits and %s",
+        kSampleMethod);
+    return std::nullopt;
+  }
+  if (!initialized_ || !has_sampler_ || logits.dim() < 2 || !is_dense(logits) ||
+      logits.size(logits.dim() - 1) != vocab_size_ ||
+      logits.scalar_type() != ::executorch::aten::ScalarType::Float ||
+      rows == 0 ||
+      rows > static_cast<std::size_t>(logits.numel() / vocab_size_)) {
+    return std::nullopt;
+  }
+  return run_device_sampler(logits, rows, params, false);
+}
+
 std::optional<std::vector<Token>> CudaExecutor::sample_on_device(
     const ::executorch::aten::Tensor& logits,
     const std::vector<SessionId>& sessions) {
-  using ::executorch::aten::ScalarType;
-  using ::executorch::aten::SizesType;
-  const ::executorch::aten::Device cuda(
-      ::executorch::aten::DeviceType::CUDA, logits.device_index());
-  const auto rows = static_cast<SizesType>(sessions.size());
-  // The rows that sample lead the logits; the rest are padding.
-  auto row_logits = make_tensor_ptr(
-      {rows, static_cast<SizesType>(vocab_size_)},
-      logits.mutable_data_ptr(),
-      logits.scalar_type(),
-      cuda);
-
   const bool all_greedy =
       std::all_of(sessions.begin(), sessions.end(), [this](SessionId session) {
         return sessions_.greedy(session);
@@ -840,6 +859,25 @@ std::optional<std::vector<Token>> CudaExecutor::sample_on_device(
     params.insert(
         params.end(), {row->temperature, row->top_p, row->top_k, row->coin});
   }
+  return run_device_sampler(logits, sessions.size(), params, all_greedy);
+}
+
+std::optional<std::vector<Token>> CudaExecutor::run_device_sampler(
+    const ::executorch::aten::Tensor& logits,
+    std::size_t row_count,
+    const std::vector<float>& params,
+    bool all_greedy) {
+  using ::executorch::aten::ScalarType;
+  using ::executorch::aten::SizesType;
+  const ::executorch::aten::Device cuda(
+      ::executorch::aten::DeviceType::CUDA, logits.device_index());
+  const auto rows = static_cast<SizesType>(row_count);
+  // The rows that sample lead the logits; the rest are padding.
+  auto row_logits = make_tensor_ptr(
+      {rows, static_cast<SizesType>(vocab_size_)},
+      logits.mutable_data_ptr(),
+      logits.scalar_type(),
+      cuda);
 
   const cudaStream_t stream =
       ::executorch::extension::cuda::getCallerStream().value_or(
@@ -851,18 +889,18 @@ std::optional<std::vector<Token>> CudaExecutor::sample_on_device(
           kArgmaxMethod,
           std::vector<::executorch::runtime::EValue>{row_logits});
     }
-    if (sessions.size() > device_params_rows_) {
+    if (row_count > device_params_rows_) {
       if (device_params_ != nullptr) {
         (void)cudaFree(device_params_);
         device_params_ = nullptr;
         device_params_rows_ = 0;
       }
       ET_CHECK_OR_RETURN_ERROR(
-          cudaMalloc(&device_params_, sessions.size() * 4 * sizeof(float)) ==
+          cudaMalloc(&device_params_, row_count * 4 * sizeof(float)) ==
               cudaSuccess,
           MemoryAllocationFailed,
           "CudaExecutor: could not stage the sampler's params");
-      device_params_rows_ = sessions.size();
+      device_params_rows_ = row_count;
     }
     // The backend runs every method on the caller's stream if one is set,
     // else the per-thread stream; ordering the upload there puts it ahead of
@@ -890,7 +928,7 @@ std::optional<std::vector<Token>> CudaExecutor::sample_on_device(
     ET_LOG(Error, "CudaExecutor: the device sampler returned a bad shape");
     return std::nullopt;
   }
-  std::vector<std::int64_t> host(sessions.size());
+  std::vector<std::int64_t> host(row_count);
   const auto kind = tokens.device_type() == ::executorch::aten::DeviceType::CUDA
       ? cudaMemcpyDeviceToHost
       : cudaMemcpyHostToHost;

backends/cuda/batching/cuda_executor.h

@@ -174,6 +174,14 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor {
   std::optional<std::vector<llm_batching::Token>> argmax_rows(
       const ::executorch::aten::Tensor& logits,
       std::size_t rows);
+  // Samples device logits with the exported sampler. params is row-major
+  // [rows, 4] float32: temperature, top_p, top_k and a uniform coin, as
+  // extension/llm/batching/sampler.py consumes them. All-greedy rows take
+  // sample_argmax.
+  std::optional<std::vector<llm_batching::Token>> sample_rows(
+      const ::executorch::aten::Tensor& logits,
+      std::size_t rows,
+      const std::vector<float>& params);
   ::executorch::extension::Module& module() {
     return *module_;
   }
@@ -212,6 +220,11 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor {
 
   // Draws one token per row of `logits`, rows [0, sessions.size()), on the
   // device; returns them in row order.
+  std::optional<std::vector<llm_batching::Token>> run_device_sampler(
+      const ::executorch::aten::Tensor& logits,
+      std::size_t rows,
+      const std::vector<float>& params,
+      bool all_greedy);
   std::optional<std::vector<llm_batching::Token>> sample_on_device(
       const ::executorch::aten::Tensor& logits,
       const std::vector<llm_batching::SessionId>& sessions);

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

@@ -95,6 +95,50 @@ def _export_cuda_sampler_methods(
     }
 
 
+VERIFY_METHOD = "dflash_verify"
+VERIFY_ARGMAX_METHOD = "dflash_verify_argmax"
+
+
+def cuda_owner_verifier_programs(
+    vocab_size: int, max_owners: int, block: int
+) -> dict[str, torch.export.ExportedProgram]:
+    """Batched on-device verification of up to ``max_owners`` draft blocks."""
+    from executorch.examples.models.muse_glimmer.model.dflash_token_sampler import (
+        DFlashOwnerArgmaxVerifier,
+        DFlashOwnerVerifier,
+    )
+    from torch.export import Dim, export
+
+    if block < 3:
+        raise ValueError("CUDA DFlash verifiers need a block of at least 3")
+    owners = (
+        Dim("verify_owners", min=1, max=max_owners) if max_owners > 1 else Dim.STATIC
+    )
+    # A one-proposal width would specialize its size-1 proposal axis.
+    width = Dim("verify_width", min=3, max=block)
+    logits = torch.zeros(max_owners, block, vocab_size)
+    candidates = torch.zeros(max_owners, block, dtype=torch.long)
+    params = torch.zeros(max_owners, block + 3)
+    return {
+        VERIFY_METHOD: export(
+            DFlashOwnerVerifier(),
+            (logits, candidates, params),
+            dynamic_shapes=(
+                {0: owners, 1: width},
+                {0: owners, 1: width},
+                {0: owners, 1: width + 3},
+            ),
+            strict=True,
+        ),
+        VERIFY_ARGMAX_METHOD: export(
+            DFlashOwnerArgmaxVerifier(),
+            (logits, candidates),
+            dynamic_shapes=({0: owners, 1: width}, {0: owners, 1: width}),
+            strict=True,
+        ),
+    }
+
+
 def _cuda_propagate_device_config() -> dict[str, PropagateDeviceConfig]:
     """Keeps the speculative chain's vocabulary-wide tensors on the device."""
     from executorch.exir.passes.propagate_device_pass import PropagateDeviceConfig
@@ -337,6 +381,7 @@ def prepare_cuda_packed_draft(
             ),
             strict=True,
         )
+    programs = {"draft_forward": program}
     geometry = offgraph_kv_cache_geometry(draft_model)
     metadata = {
         key.replace("get_", "get_draft_", 1): value for key, value in geometry.items()
@@ -357,7 +402,7 @@ def prepare_cuda_packed_draft(
             "get_dflash_draft_cache_key": "muse_glimmer_draft",
         }
     )
-    return {"draft_forward": program}, manifest, metadata
+    return programs, manifest, metadata
 
 
 def _dflash_constant_methods(

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

@@ -318,7 +318,7 @@ def sampler_programs(vocab_size: int, max_rows: int = MAX_SAMPLE_ROWS):
     }
 
 
-def device_residency(forward_methods, device_sampling: bool):
+def device_residency(forward_methods, device_sampling: bool, samplers=()):
     """Per-method ``PropagateDeviceConfig`` for ``to_executorch``.
 
     With device sampling the forwards keep their logits in device memory, and
@@ -331,7 +331,7 @@ def device_residency(forward_methods, device_sampling: bool):
         for name in forward_methods
     }
     if device_sampling:
-        for name in (SAMPLE_METHOD, ARGMAX_METHOD):
+        for name in (SAMPLE_METHOD, ARGMAX_METHOD, *samplers):
             device[name] = PropagateDeviceConfig(skip_h2d_for_method_inputs=True)
     return device
 
@@ -469,19 +469,31 @@ def export_cuda_batching(  # noqa: C901
     if dflash_draft is not None:
         programs.update(draft_programs)
         del draft_programs, dflash_draft
-    cache_methods = forward_methods + (
-        ["draft_forward"] if "draft_forward" in programs else []
-    )
+    draft_methods = [name for name in programs if name.startswith("draft_forward")]
+    cache_methods = forward_methods + draft_methods
+    verifiers = {}
     if device_sampling:
         print(f"Exporting {SAMPLE_METHOD} and {ARGMAX_METHOD}...")
         programs.update(sampler_programs(config.vocab_size))
+        if "draft_forward" in programs:
+            from executorch.examples.models.muse_glimmer.export.export_dflash import (
+                cuda_owner_verifier_programs,
+            )
+
+            verifiers = cuda_owner_verifier_programs(
+                config.vocab_size,
+                max_draft_owners,
+                draft_metadata["get_block_size"],
+            )
+            print(f"Exporting {', '.join(verifiers)}...")
+            programs.update(verifiers)
 
     def partitioner(name: str) -> CudaPartitioner:
         if name not in cache_methods:
             return CudaPartitioner(
                 [CudaBackend.generate_method_name_compile_spec(name)]
             )
-        is_draft = name == "draft_forward"
+        is_draft = name in draft_methods
         specs = [
             CudaBackend.generate_method_name_compile_spec(name),
             CompileSpec("low_memory_mode", b"ON"),
@@ -494,7 +506,7 @@ def partitioner(name: str) -> CudaPartitioner:
                 b"2:0" if is_draft else step_width_spec(name),
             ),
         ]
-        if "draft_forward" in cache_methods:
+        if draft_methods:
             specs.append(
                 CompileSpec(
                     "llm_cache_registry_key",
@@ -539,7 +551,9 @@ def partitioner(name: str) -> CudaPartitioner:
             extract_delegate_segments=True,
             do_quant_fusion_and_const_prop=True,
             memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False),
-            propagate_device_config=device_residency(cache_methods, device_sampling),
+            propagate_device_config=device_residency(
+                cache_methods, device_sampling, tuple(verifiers)
+            ),
             emit_mutable_buffer_names=True,
         ),
     )