diff --git a/backends/cuda/batching/cuda_executor.cpp b/backends/cuda/batching/cuda_executor.cpp index b1bd377eb9a..0608b3d8466 100644 --- a/backends/cuda/batching/cuda_executor.cpp +++ b/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, diff --git a/backends/cuda/batching/cuda_executor.h b/backends/cuda/batching/cuda_executor.h index fc0094ae9ef..b117815fbd3 100644 --- a/backends/cuda/batching/cuda_executor.h +++ b/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 additional_cuda_graph_methods; }; struct CudaExecutorConfig { diff --git a/backends/cuda/runtime/cuda_kv_cache.cpp b/backends/cuda/runtime/cuda_kv_cache.cpp index 2fc4b7317ba..87800103363 100644 --- a/backends/cuda/runtime/cuda_kv_cache.cpp +++ b/backends/cuda/runtime/cuda_kv_cache.cpp @@ -449,10 +449,11 @@ class CudaCellCache final : public cache::CellCache, public CudaKVCache { "offgraph_kv: a %lld-token step is outside [1, %d]", static_cast(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(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, diff --git a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp b/backends/cuda/runtime/test/test_cuda_kv_cache.cpp index 2d996a1e535..0c32da2079f 100644 --- a/backends/cuda/runtime/test/test_cuda_kv_cache.cpp +++ b/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(); auto* kv = cache_ptr->as(); @@ -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(); + auto* kv = cache_ptr->as(); + 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(cells.begin(), cells.begin() + 5), + std::vector({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>; + 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(); diff --git a/examples/models/muse-glimmer/export/export_dflash.py b/examples/models/muse-glimmer/export/export_dflash.py index aa956842d16..ca548c08911 100644 --- a/examples/models/muse-glimmer/export/export_dflash.py +++ b/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" @@ -378,6 +379,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() diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp index 79d544f9047..55f8ee16a9d 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp @@ -141,6 +141,39 @@ std::optional 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; + const int context = static_cast(plan.hidden.size() / row_bytes); + if (row_bytes == 0 || plan.hidden.size() % row_bytes != 0 || + context > width || + plan.positions.size() != static_cast(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(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_) { cudaFree(device_); @@ -435,6 +468,10 @@ Result 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()); @@ -621,6 +658,42 @@ bool MuseGlimmerCudaDFlashExecutor::initialize() { "CUDA DFlash: could not initialize distinct target/draft caches"); return false; } + const auto names = module.method_names(); + if (!names.ok()) { + 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 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; } @@ -979,23 +1052,38 @@ bool MuseGlimmerCudaDFlashExecutor::draft( } const int q = static_cast(owners.size()) * block_; const int c = static_cast(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(owners.size()), block_, row_bytes_)) { + return false; + } + const int padded = static_cast(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(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)); + const std::string method = use_static + ? df::kStaticForwardPrefix + std::to_string(owners.size()) + : std::string(df::kForward); auto result = target_->module().execute( - df::kForward, {tokens, hidden, positions, order}); + 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) || diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h index eb9729e3c55..11dffd5f7fe 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h @@ -10,6 +10,7 @@ #include #include +#include #include 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"; @@ -102,6 +105,10 @@ std::optional 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 { @@ -125,6 +132,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; @@ -144,9 +152,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 create( @@ -242,6 +252,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 static_drafts_{}; bool initialized_ = false; // A failed co-batch invalidates every draft owner. Only model restart // recovers. diff --git a/examples/models/muse-glimmer/tests/cuda_dflash_executor_test.cpp b/examples/models/muse-glimmer/tests/cuda_dflash_executor_test.cpp index 740788d189d..2381686e456 100644 --- a/examples/models/muse-glimmer/tests/cuda_dflash_executor_test.cpp +++ b/examples/models/muse-glimmer/tests/cuda_dflash_executor_test.cpp @@ -61,6 +61,34 @@ TEST(CudaDFlash, PositionalUniformsAreSeededAndIndependentOfBatching) { EXPECT_NE(u, df::positional_uniform(7, 100, 3)); } +TEST(CudaDFlash, StaticPaddingPutsUnusedContextAfterEveryBlock) { + const std::vector 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({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( + {10, 11, 20, 0, 0, 0, 0, 0, 12, 13, 14, 15, 21, 22, 23, 24})); + EXPECT_EQ( + plan->kv_order, + std::vector( + {0, 1, 8, 9, 10, 11, 2, 12, 13, 14, 15, 3, 4, 5, 6, 7})); +} + +TEST(CudaDFlash, StaticPaddingRejectsLongContext) { + const std::vector 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. diff --git a/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py b/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py index 49528490719..3cff10f30fe 100644 --- a/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py +++ b/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py @@ -442,6 +442,9 @@ def owned_schema(package, name): "forward_8", "forward_others", "draft_forward", + # Two static owners fit max_write = 16 + 4 * 16 (2 * N * 16 rows). + "draft_forward_1", + "draft_forward_2", "sample", "sample_argmax", "dflash_verify",