From c3beec834036d73c1eff0bb8960f0e8b829f401b Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Tue, 20 Oct 2026 15:52:47 -0700 Subject: [PATCH] [executorch][muse] Run the CUDA DFlash drafter in CUDA graphs Pull Request resolved: https://github.com/pytorch/executorch/pull/23985 Export static draft_forward_{N} (N owners, at most one block of new context each) and capture them in CUDA graphs through the existing CudaExecutor graph option; longer contexts still drain through the dynamic draft_forward. Decode contexts are padded to the static shape with scratch rows the cache never commits. Pool growth recaptures like the target's static methods. ghstack-source-id: 453099299 @exported-using-ghexport Differential Revision: [D125705719](https://our.internmc.facebook.com/intern/diff/D125705719/) --- backends/cuda/batching/cuda_executor.cpp | 5 +- backends/cuda/batching/cuda_executor.h | 2 + backends/cuda/runtime/cuda_kv_cache.cpp | 5 +- .../cuda/runtime/test/test_cuda_kv_cache.cpp | 38 ++++++- .../muse-glimmer/export/export_dflash.py | 23 ++++ .../runtime/cuda_dflash_executor.cpp | 100 +++++++++++++++++- .../runtime/cuda_dflash_executor.h | 22 +++- .../tests/cuda_dflash_executor_test.cpp | 30 ++++++ .../tests/cuda_dflash_native_test.cpp | 86 +++++++++++++-- .../tests/test_cuda_dflash_packed.py | 3 + 10 files changed, 296 insertions(+), 18 deletions(-) diff --git a/backends/cuda/batching/cuda_executor.cpp b/backends/cuda/batching/cuda_executor.cpp index 40aad0317fc..1e4567fc7b7 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, @@ -288,7 +291,7 @@ Result> 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); } 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 91a63798c7e..ddafe60d366 100644 --- a/backends/cuda/runtime/cuda_kv_cache.cpp +++ b/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(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 1d6e183cbfd..f28af1b6114 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" @@ -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() diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp index f4a1dd2a052..0389cd5a2f8 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp @@ -143,6 +143,41 @@ 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; + if (row_bytes == 0 || plan.hidden.size() % row_bytes != 0) { + return false; + } + const int context = static_cast(plan.hidden.size() / row_bytes); + if (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_) { (void)cudaFree(device_); @@ -443,6 +478,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()); @@ -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 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(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)); - 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) || diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h index 5088a165041..4bda7961f39 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"; @@ -104,6 +107,14 @@ 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 { @@ -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 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 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 73878110346..bf7ce17a9a4 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,36 @@ 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/cuda_dflash_native_test.cpp b/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp index a12386237ea..80ed5e774e0 100644 --- a/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp +++ b/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp @@ -17,6 +17,7 @@ #include #include #include +#include namespace llm = executorch::extension::llm; namespace b = llm::batching; @@ -33,11 +34,21 @@ namespace executorch::extension::llm { // Tiny artifacts cannot pass Muse's special-token vocabulary check. Keep the // production factory unchanged and construct only the validated tiny geometry. struct CudaDFlashTestPeer { - static std::unique_ptr create( - std::unique_ptr module, - int proposals) { + static std::unique_ptr + create(std::unique_ptr module, int proposals, bool draft_graphs) { + // Capture the static drafts in CUDA graphs, as the production create does. + cb::CudaExecutorOptions options; + for (int n = 1; draft_graphs && n <= cuda_dflash_detail::kOwnerLimit; ++n) { + options.additional_cuda_graph_methods.push_back( + cuda_dflash_detail::kStaticForwardPrefix + std::to_string(n)); + } auto target = cb::CudaExecutor::create( - std::move(module), 4, 96, static_cast(aten::ScalarType::BFloat16)); + std::move(module), + 4, + 96, + static_cast(aten::ScalarType::BFloat16), + /*initial_capacity=*/-1, + std::move(options)); EXPECT_EQ(target.error(), rt::Error::Ok); if (!target.ok()) { return nullptr; @@ -312,10 +323,13 @@ class CudaDFlashNative : public ::testing::Test { rt::runtime_init(); dir_ = dir; } - void create(bool deterministic = true, int proposals = 15) { + void create( + bool deterministic = true, + int proposals = 15, + bool draft_graphs = true) { auto module = std::make_unique(dir_, deterministic); fixture_ = module.get(); - executor_ = Peer::create(std::move(module), proposals); + executor_ = Peer::create(std::move(module), proposals, draft_graphs); ASSERT_TRUE(executor_); ASSERT_TRUE(executor_->initialize()); } @@ -430,6 +444,8 @@ TEST_F(CudaDFlashNative, EveryRejectionPositionAndFullBonus) { EXPECT_EQ(after.target_preamble_calls, before.target_preamble_calls); EXPECT_EQ(after.target_verify_calls - before.target_verify_calls, 1); EXPECT_EQ(after.pending_decode_owners - before.pending_decode_owners, 1); + EXPECT_EQ(after.draft_forwards - before.draft_forwards, 1); + EXPECT_EQ(after.static_draft_forwards - before.static_draft_forwards, 1); EXPECT_EQ(after.cycles - before.cycles, 1); EXPECT_EQ(after.proposed - before.proposed, 15); EXPECT_EQ(after.accepted - before.accepted, rejected); @@ -586,9 +602,16 @@ TEST_F(CudaDFlashNative, FinalWholeBlockFitsExactlyAndPreservesAllPrefillRows) { Peer::check(*executor_, sid, 78); EXPECT_EQ(executor_->stats().hidden_rows_d2h, 78); fixture_->rejected.fill(0); + const auto drain = executor_->stats(); ASSERT_TRUE(executor_->execute({{input(sid, 78, {1})}}, out)); EXPECT_EQ(out.outputs[0]->tokens, predictions(78, 2)); Peer::check(*executor_, sid, 80); + // 79 rows exceed one static block: drain dynamically, then draft statically. + const auto drained = executor_->stats(); + EXPECT_EQ(drained.static_draft_forwards - drain.static_draft_forwards, 1); + EXPECT_GT( + drained.draft_forwards - drain.draft_forwards, + drained.static_draft_forwards - drain.static_draft_forwards); fixture_->rejected.fill(15); const auto before = executor_->stats(); ASSERT_TRUE(executor_->execute({{input(sid, 80, {prediction(79)})}}, out)); @@ -720,4 +743,55 @@ TEST_F(CudaDFlashNative, UnmodifiedOldArtifactRunsSoloAndFourOwnerDecode) { executor_->stats().target_preamble_calls, before.target_preamble_calls); EXPECT_EQ(executor_->stats().failed_batches, 0); } + +// The production adapter captures draft_forward_{N} in CUDA graphs; on real +// outputs those replays must decode exactly as the eager static drafts do. +TEST_F(CudaDFlashNative, CapturedStaticDraftsMatchEager) { + auto run = [&](bool draft_graphs) { + std::vector emitted; + // The graph method list is a process-wide backend option set at load. + executor_.reset(); + create(/*deterministic=*/false, /*proposals=*/15, draft_graphs); + if (!executor_) { + return std::pair{emitted, llm::CudaDFlashStats{}}; + } + // The tiny artifact exports static drafts for up to three owners. + std::array owners; + b::BatchInput batch; + for (int owner = 0; owner < 3; ++owner) { + owners[owner] = open(); + Tokens prompt(owner + 3); + std::iota(prompt.begin(), prompt.end(), 1 + 17 * owner); + batch.inputs.push_back(input(owners[owner], 0, std::move(prompt))); + } + b::BatchOutput out; + for (int round = 0; round < 4; ++round) { + EXPECT_TRUE(executor_->execute(batch, out)); + batch.inputs.clear(); + for (int owner = 0; owner < 3; ++owner) { + if (!out.outputs[owner] || out.outputs[owner]->tokens.empty()) { + ADD_FAILURE() << "no output for owner " << owner; + return std::pair{emitted, executor_->stats()}; + } + emitted.push_back(out.outputs[owner]->tokens); + batch.inputs.push_back(input( + owners[owner], + Peer::frontier(*executor_, owners[owner]), + {out.outputs[owner]->tokens.back()})); + } + } + return std::pair{emitted, executor_->stats()}; + }; + const auto [graph_tokens, graph_stats] = run(true); + const auto [eager_tokens, eager_stats] = run(false); + EXPECT_EQ(graph_tokens, eager_tokens); + EXPECT_GT(graph_stats.static_draft_forwards, 0); + EXPECT_EQ( + graph_stats.static_draft_forwards, eager_stats.static_draft_forwards); + EXPECT_EQ(graph_stats.draft_forwards, eager_stats.draft_forwards); + EXPECT_EQ(graph_stats.accepted, eager_stats.accepted); + EXPECT_EQ(graph_stats.bonuses, eager_stats.bonuses); + EXPECT_EQ(graph_stats.failed_batches, 0); + EXPECT_EQ(eager_stats.failed_batches, 0); +} } // namespace 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 db7afaa2072..9ace17699b7 100644 --- a/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py +++ b/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py @@ -443,6 +443,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",