From 10e44973241a84fc39e268d32658ecc3a3630bc0 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sun, 18 Oct 2026 16:45:48 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- examples/models/muse-glimmer/CMakeLists.txt | 6 + .../runtime/cuda_dflash_executor.cpp | 97 ++- .../runtime/cuda_dflash_executor.h | 10 +- .../tests/cuda_dflash_native_test.cpp | 646 ++++++++++++++++++ 4 files changed, 739 insertions(+), 20 deletions(-) create mode 100644 examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp diff --git a/examples/models/muse-glimmer/CMakeLists.txt b/examples/models/muse-glimmer/CMakeLists.txt index 6129197e8a8..abcd30994d9 100644 --- a/examples/models/muse-glimmer/CMakeLists.txt +++ b/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 ) diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp index 8b4572d0833..2e12d017618 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp @@ -629,6 +629,7 @@ std::optional 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); @@ -674,6 +675,7 @@ bool MuseGlimmerCudaDFlashExecutor::rewind(State& state, int position) { return false; } state.frontier = position; + state.pending.reset(); return true; } @@ -878,7 +880,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 owners; @@ -890,8 +893,9 @@ bool MuseGlimmerCudaDFlashExecutor::speculate( auto& state = sessions_.at(input.sid); const auto end = static_cast(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; } @@ -908,7 +912,10 @@ bool MuseGlimmerCudaDFlashExecutor::speculate( } context += static_cast(rows); owners.push_back(input.sid); - anchors.push_back(out.outputs[cursor]->tokens.front()); + anchors.push_back( + pending_decode ? (*std::get( + input.payload))[input.offset] + : out.outputs[cursor]->tokens.front()); indices.push_back(cursor++); } if (owners.empty()) { @@ -971,9 +978,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()); @@ -997,6 +1004,7 @@ bool MuseGlimmerCudaDFlashExecutor::fail_batch( failed_ = true; for (auto& item : sessions_) { item.second.poisoned = true; + item.second.pending.reset(); item.second.hidden.clear(); } // An abandoned declaration/forward affects the entire shared draft pool. @@ -1028,6 +1036,7 @@ bool MuseGlimmerCudaDFlashExecutor::execute( std::vector token_owners; std::unordered_map cursor; std::vector> rewinds; + bool pending_decode = !batch.inputs.empty(); for (const auto& input : batch.inputs) { const auto it = sessions_.find(input.sid); const auto* payload = @@ -1063,6 +1072,13 @@ bool MuseGlimmerCudaDFlashExecutor::execute( if (start < at->second) { rewinds.emplace_back(input.sid, static_cast(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(end); for (std::size_t j = 0; j < input.size; ++j) { const auto token = (**payload)[input.offset + j]; @@ -1085,34 +1101,77 @@ bool MuseGlimmerCudaDFlashExecutor::execute( return fail_batch(batch, out); } } - std::vector predictions; - if (!run_target(step, token_owners, predictions, true)) { + const auto before = stats_; + 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 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(input.position) + + input.offset + input.size; + if (tokens.empty()) { + return fail_batch(batch, out); + } + const auto expected = + end + static_cast(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(); + } } } - const auto before = stats_; - if (!speculate(batch, out)) { - return fail_batch(batch, out); - } ++stats_.executions; ET_LOG( Info, - "DFLASH_STATS cycles=%llu proposed=%llu accepted=%llu rejected=%llu bonuses=%llu draft_forwards_total=%llu target_preamble_calls_total=%llu target_verify_calls_total=%llu target_preamble_forwards_total=%llu target_verify_forwards_total=%llu hidden_rows_d2h_total=%llu hidden_bytes_d2h_total=%llu hidden_rows_d2d_total=%llu hidden_bytes_d2d_total=%llu", + "DFLASH_STATS cycles=%llu proposed=%llu accepted=%llu rejected=%llu bonuses=%llu draft_forwards_total=%llu", static_cast(stats_.cycles - before.cycles), static_cast(stats_.proposed - before.proposed), static_cast(stats_.accepted - before.accepted), static_cast(stats_.rejected - before.rejected), static_cast(stats_.bonuses - before.bonuses), - static_cast(stats_.draft_forwards), + static_cast(stats_.draft_forwards)); + // Keep each record below the runtime's 256-byte log-message limit. + ET_LOG( + Info, + "DFLASH_PENDING_STATS executions_total=%llu pending_decode_batches_total=%llu pending_decode_owners_total=%llu", + static_cast(stats_.executions), + static_cast(stats_.pending_decode_batches), + static_cast(stats_.pending_decode_owners)); + ET_LOG( + Info, + "DFLASH_TARGET_STATS target_preamble_calls_total=%llu target_verify_calls_total=%llu target_preamble_forwards_total=%llu target_verify_forwards_total=%llu", static_cast(stats_.target_preamble_calls), static_cast(stats_.target_verify_calls), static_cast(stats_.target_preamble_forwards), - static_cast(stats_.target_verify_forwards), + static_cast(stats_.target_verify_forwards)); + ET_LOG( + Info, + "DFLASH_HIDDEN_STATS hidden_rows_d2h_total=%llu hidden_bytes_d2h_total=%llu hidden_rows_d2d_total=%llu hidden_bytes_d2d_total=%llu", static_cast(stats_.hidden_rows_d2h), static_cast(stats_.hidden_bytes_d2h), static_cast(stats_.hidden_rows_d2d), diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h index cea4c6cc9a2..f3ee876d4ac 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h @@ -95,6 +95,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; @@ -137,12 +139,15 @@ 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 { 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 pending; // Host-owned committed BF16 rows [draft frontier, target frontier). // Verification appends only the consumed anchor and accepted proposals. std::vector hidden; @@ -173,7 +178,10 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor { const std::vector& anchors, int context_limit, std::vector>* 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); diff --git a/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp b/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp new file mode 100644 index 00000000000..9774a983631 --- /dev/null +++ b/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp @@ -0,0 +1,646 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace llm = executorch::extension::llm; +namespace b = llm::batching; +namespace cb = executorch::backends::cuda::batching; +namespace c = llm::cache; +namespace rt = executorch::runtime; +namespace aten = executorch::aten; +using DFlash = llm::MuseGlimmerCudaDFlashExecutor; +using Tokens = std::vector; +using executorch::extension::Module; +using executorch::extension::TensorPtr; + +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) { + auto target = cb::CudaExecutor::create( + std::move(module), 4, 96, static_cast(aten::ScalarType::BFloat16)); + EXPECT_EQ(target.error(), rt::Error::Ok); + if (!target.ok()) { + return nullptr; + } + EXPECT_EQ(target.get()->vocab_size(), 256); + EXPECT_EQ(target.get()->preferred_batch_tokens(), 32); + auto& program = target.get()->module(); + auto mask = detail::read_int_method(program, "get_mask_token_id"); + auto layers = detail::read_int_method(program, "get_draft_n_caches"); + auto heads = detail::read_int_tensor_method(program, "get_draft_kv_heads"); + auto dims = detail::read_int_tensor_method(program, "get_draft_head_dims"); + auto windows = detail::read_int_tensor_method(program, "get_draft_windows"); + if (!mask.ok() || !*mask || !layers.ok() || !*layers || !heads.ok() || + !dims.ok() || !windows.ok()) { + return nullptr; + } + EXPECT_EQ(**layers, heads->size()); + c::CacheGeometry geometry; + geometry.dflash_block_size = 16; + for (std::size_t i = 0; i < heads->size(); ++i) { + geometry.layers.push_back( + {{(*windows)[i] ? c::LayerPolicy::Kind::Ring + : c::LayerPolicy::Kind::Flat, + (*windows)[i]}, + (*heads)[i], + (*dims)[i]}); + } + c::CacheConfig config{}; + config.capacity = 512; + config.max_write = 96; + config.kv_dtype = static_cast(aten::ScalarType::BFloat16); + auto draft = c::CacheFactory::global().build( + executorch::backends::cuda::kCudaBackendId, + c::kind::kBatchedCell, + geometry, + config); + EXPECT_EQ(draft.error(), rt::Error::Ok); + if (!draft.ok()) { + return nullptr; + } + return std::unique_ptr(new DFlash( + std::move(target.get()), + std::move(draft.get()), + 16, + proposals, + static_cast(**mask), + 4, + 32, + 4, + 96, + 128)); + } + static int frontier(DFlash& executor, b::SessionId sid) { + return executor.sessions_.at(sid).frontier; + } + static int draft_position(DFlash& executor, b::SessionId sid) { + return executor.draft_ctl_->pos(executor.sessions_.at(sid).draft_id); + } + static const std::vector& hidden( + DFlash& executor, + b::SessionId sid) { + return executor.sessions_.at(sid).hidden; + } + static std::optional pending(DFlash& executor, b::SessionId sid) { + return executor.sessions_.at(sid).pending; + } + static void check(DFlash& executor, b::SessionId sid, int expected) { + const auto& state = executor.sessions_.at(sid); + EXPECT_FALSE(state.poisoned); + EXPECT_EQ(state.frontier, expected); + EXPECT_EQ(executor.target_->cache_control().pos(state.target_id), expected); + const int draft = executor.draft_ctl_->pos(state.draft_id); + ASSERT_GE(draft, 0); + ASSERT_LE(draft, expected); + EXPECT_EQ(state.hidden.size(), (expected - draft) * 256u); + } + static bool rewind(DFlash& executor, b::SessionId sid, int position) { + return executor.rewind(executor.sessions_.at(sid), position); + } + static bool poisoned(DFlash& executor, b::SessionId sid) { + return executor.sessions_.at(sid).poisoned; + } +}; +} // namespace executorch::extension::llm +using Peer = llm::CudaDFlashTestPeer; + +namespace { +// Independent greedy fixture: prediction at absolute position p is 20+p. +// Hidden row p has finite BF16 values with bits 0x3f00 + 8*(p%16) + column%8. +// Real tiny CUDA target/draft methods still run and commit their actual caches; +// deterministic replacement outputs isolate adapter acceptance/lifetime logic +// from the tiny model's known numerical parity caveat. +b::Token prediction(int position) { + return 20 + position; +} +std::vector hidden_rows(int start, int rows) { + std::vector bytes(rows * 256); + for (int row = 0; row < rows; ++row) { + for (int column = 0; column < 128; ++column) { + const std::uint16_t bits = 0x3f00 + 8 * ((start + row) % 16) + column % 8; + std::memcpy(bytes.data() + row * 256 + column * 2, &bits, 2); + } + } + return bytes; +} +b::Input +input(b::SessionId sid, int position, Tokens tokens, bool output = true) { + const auto size = tokens.size(); + return { + sid, + output, + 0, + size, + std::make_shared(std::move(tokens)), + position}; +} +Tokens predictions(int start, int count) { + Tokens result(count); + std::iota(result.begin(), result.end(), prediction(start)); + return result; +} + +class FixtureModule final : public Module { + public: + FixtureModule(const std::string& dir, bool deterministic) + : Module( + dir + "/model.pte", + std::vector{dir + "/aoti_cuda_blob.ptd"}), + deterministic_(deterministic) {} + ~FixtureModule() override { + if (taps_) { + cudaFree(taps_); + } + } + std::array rejected{15, 15, 15, 15}; + bool fail_target_sampler = false; + rt::Result> execute( + const std::string& name, + const std::vector& inputs) override { + auto actual = Module::execute(name, inputs); + if (!actual.ok()) { + return actual.error(); + } + const auto stream = executorch::extension::cuda::getCallerStream().value_or( + cudaStreamPerThread); + if (name == "sample_argmax") { + if (taps_ && cudaMemsetAsync(taps_, 0, 32 * 256, stream) != cudaSuccess) { + return rt::Error::Internal; + } + if (target_output_ && fail_target_sampler) { + fail_target_sampler = false; + return rt::Error::Internal; + } + return actual; + } + if (name != "draft_forward" && name.rfind("forward_", 0) != 0) { + return actual; + } + target_output_ = name != "draft_forward"; + if (!deterministic_) { + return actual; + } + const auto& positions = inputs[target_output_ ? 1 : 2].toTensor(); + const auto* p = positions.const_data_ptr(); + const auto& tokens = inputs[0].toTensor(); + const int width = tokens.size(1); + const auto& selector = inputs[target_output_ ? 2 : 0].toTensor(); + const int rows = target_output_ ? selector.numel() : width; + std::vector logits(rows * 256, -100.0f); + const int context = target_output_ ? 0 : inputs[1].toTensor().size(1); + for (int row = 0; row < rows; ++row) { + const int position = target_output_ + ? p[selector.const_data_ptr()[row]] + : p[context + row] - 1; + int token = static_cast(prediction(position)); + if (!target_output_ && row % 16 == rejected[row / 16] + 1) { + token += 100; + } + logits[row * 256 + token] = 100.0f; + } + logits_ = device_tensor( + logits.data(), + logits.size() * sizeof(float), + {1, rows, 256}, + aten::ScalarType::Float); + if (!logits_) { + return rt::Error::Internal; + } + if (!target_output_) { + return std::vector{logits_}; + } + if (!taps_ && cudaMalloc(&taps_, 32 * 256) != cudaSuccess) { + return rt::Error::MemoryAllocationFailed; + } + std::vector bytes(width * 256, 255); + for (int row = 0; row < width; ++row) { + auto value = hidden_rows(p[row], 1); + std::copy(value.begin(), value.end(), bytes.begin() + row * 256); + } + if (cudaMemcpyAsync( + taps_, + bytes.data(), + bytes.size(), + cudaMemcpyHostToDevice, + stream) != cudaSuccess || + cudaStreamSynchronize(stream) != cudaSuccess) { + return rt::Error::Internal; + } + tap_view_ = executorch::extension::from_blob( + taps_, + {1, width, 128}, + aten::ScalarType::BFloat16, + aten::Device(aten::DeviceType::CUDA, 0)); + return std::vector{logits_, tap_view_}; + } + + private: + static TensorPtr device_tensor( + const void* data, + std::size_t bytes, + std::vector sizes, + aten::ScalarType type) { + void* device = nullptr; + if (cudaMalloc(&device, bytes) != cudaSuccess) { + return nullptr; + } + const auto stream = executorch::extension::cuda::getCallerStream().value_or( + cudaStreamPerThread); + if (cudaMemcpyAsync(device, data, bytes, cudaMemcpyHostToDevice, stream) != + cudaSuccess || + cudaStreamSynchronize(stream) != cudaSuccess) { + cudaFree(device); + return nullptr; + } + return executorch::extension::from_blob( + device, + std::move(sizes), + type, + aten::Device(aten::DeviceType::CUDA, 0), + [](void* pointer) { cudaFree(pointer); }); + } + bool deterministic_, target_output_ = false; + void* taps_ = nullptr; + TensorPtr logits_, tap_view_; +}; + +class CudaDFlashNative : public ::testing::Test { + protected: + void SetUp() override { + const char* dir = std::getenv("ET_MUSE_DFLASH_TINY_DIR"); + if (!dir || !*dir) { + GTEST_SKIP() + << "Set ET_MUSE_DFLASH_TINY_DIR to the existing tiny packed artifact"; + } + int devices = 0; + if (cudaGetDeviceCount(&devices) != cudaSuccess || devices == 0) { + GTEST_SKIP() << "CUDA device required"; + } + rt::runtime_init(); + dir_ = dir; + } + void create(bool deterministic = true, int proposals = 15) { + auto module = std::make_unique(dir_, deterministic); + fixture_ = module.get(); + executor_ = Peer::create(std::move(module), proposals); + ASSERT_TRUE(executor_); + ASSERT_TRUE(executor_->initialize()); + } + b::SessionId open() { + auto sid = executor_->open_session(); + EXPECT_TRUE(sid); + if (!sid) { + return -1; + } + executor_->set_sampling(*sid, {}, 0); + return *sid; + } + void seed(b::SessionId sid, int rows) { + b::BatchOutput out; + ASSERT_TRUE( + executor_->execute({{input(sid, 0, Tokens(rows, 1), false)}}, out)); + Peer::check(*executor_, sid, rows); + EXPECT_FALSE(Peer::pending(*executor_, sid)); + EXPECT_EQ(Peer::hidden(*executor_, sid), hidden_rows(0, rows)); + } + std::string dir_; + FixtureModule* fixture_ = nullptr; + std::unique_ptr executor_; +}; + +TEST_F(CudaDFlashNative, EveryRejectionPositionAndFullBonus) { + create(); + ASSERT_TRUE(executor_); + for (int rejected = 0; rejected <= 15; ++rejected) { + SCOPED_TRACE(rejected); + const auto sid = open(); + seed(sid, 3); + fixture_->rejected.fill(rejected); + b::BatchOutput preamble; + ASSERT_TRUE(executor_->execute({{input(sid, 3, {1})}}, preamble)); + ASSERT_TRUE(preamble.outputs[0]); + EXPECT_EQ(preamble.outputs[0]->tokens, predictions(3, rejected + 2)); + const int start = 5 + rejected; + Peer::check(*executor_, sid, start); + EXPECT_EQ(Peer::hidden(*executor_, sid), hidden_rows(4, rejected + 1)); + const auto before = executor_->stats(); + b::BatchOutput decode; + ASSERT_TRUE(executor_->execute( + {{input(sid, start, {prediction(start - 1)})}}, decode)); + ASSERT_TRUE(decode.outputs[0]); + EXPECT_EQ(decode.outputs[0]->tokens, predictions(start, rejected + 1)); + Peer::check(*executor_, sid, start + rejected + 1); + EXPECT_EQ(Peer::draft_position(*executor_, sid), start); + EXPECT_EQ(Peer::hidden(*executor_, sid), hidden_rows(start, rejected + 1)); + EXPECT_EQ(Peer::pending(*executor_, sid), prediction(start + rejected)); + const auto after = executor_->stats(); + 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.accepted - before.accepted, rejected); + EXPECT_EQ(after.bonuses - before.bonuses, rejected == 15); + EXPECT_EQ(after.hidden_rows_d2d - before.hidden_rows_d2d, 16); + EXPECT_EQ(after.hidden_rows_d2h - before.hidden_rows_d2h, rejected + 1); + EXPECT_EQ( + after.hidden_bytes_d2h - before.hidden_bytes_d2h, + (rejected + 1) * 256u); + executor_->close_session(sid); + } +} + +TEST_F(CudaDFlashNative, FourUnequalOwnersSplitCallbacksAndMixedFallback) { + create(); + ASSERT_TRUE(executor_); + std::array owners; + std::array starts; + const std::array rejected{0, 2, 7, 15}; + for (int owner = 0; owner < 4; ++owner) { + owners[owner] = open(); + seed(owners[owner], owner + 2); + fixture_->rejected.fill(rejected[owner]); + b::BatchOutput out; + ASSERT_TRUE( + executor_->execute({{input(owners[owner], owner + 2, {1})}}, out)); + ASSERT_TRUE(out.outputs[0]); + EXPECT_EQ( + out.outputs[0]->tokens, predictions(owner + 2, rejected[owner] + 2)); + starts[owner] = owner + 4 + rejected[owner]; + } + fixture_->rejected = rejected; + const auto before = executor_->stats(); + b::BatchInput batch; + for (int owner = 0; owner < 4; ++owner) { + batch.inputs.push_back( + input(owners[owner], starts[owner], {prediction(starts[owner] - 1)})); + } + b::BatchOutput out; + ASSERT_TRUE(executor_->execute(batch, out)); + for (int owner = 0; owner < 4; ++owner) { + ASSERT_TRUE(out.outputs[owner]); + EXPECT_EQ(out.outputs[owner]->sid, owners[owner]); + EXPECT_EQ( + out.outputs[owner]->tokens, + predictions(starts[owner], rejected[owner] + 1)); + Peer::check(*executor_, owners[owner], starts[owner] + rejected[owner] + 1); + EXPECT_EQ( + Peer::hidden(*executor_, owners[owner]), + hidden_rows(starts[owner], rejected[owner] + 1)); + } + const auto after = executor_->stats(); + EXPECT_EQ(after.target_preamble_calls, before.target_preamble_calls); + EXPECT_EQ(after.target_verify_calls - before.target_verify_calls, 1); + EXPECT_EQ(after.target_verify_forwards - before.target_verify_forwards, 2); + EXPECT_EQ(after.hidden_rows_d2d - before.hidden_rows_d2d, 64); + EXPECT_EQ(after.hidden_rows_d2h - before.hidden_rows_d2h, 28); + // A decode owner sharing a batch with new prompt rows retains the preamble. + const int position0 = starts[0] + 1; + const int position1 = starts[1] + 3; + const auto idle_hidden = Peer::hidden(*executor_, owners[3]); + const int idle_position = Peer::frontier(*executor_, owners[3]); + fixture_->rejected.fill(0); + ASSERT_TRUE(executor_->execute( + {{input(owners[0], position0, {prediction(position0 - 1)}), + input(owners[1], position1, {1, 2})}}, + out)); + EXPECT_EQ( + executor_->stats().target_preamble_calls - after.target_preamble_calls, + 1); + EXPECT_EQ( + executor_->stats().pending_decode_batches, after.pending_decode_batches); + EXPECT_EQ(out.outputs[0]->tokens, predictions(position0, 2)); + EXPECT_EQ(out.outputs[1]->tokens, predictions(position1 + 1, 2)); + EXPECT_EQ(Peer::hidden(*executor_, owners[3]), idle_hidden); + Peer::check(*executor_, owners[3], idle_position); +} + +TEST_F(CudaDFlashNative, PartialVerificationWidthExcludesStaticPadding) { + create(true, 5); + ASSERT_TRUE(executor_); + for (int rejected : {2, 5}) { + const auto sid = open(); + seed(sid, 3); + fixture_->rejected.fill(rejected); + b::BatchOutput out; + ASSERT_TRUE(executor_->execute({{input(sid, 3, {1})}}, out)); + ASSERT_TRUE(out.outputs[0]); + EXPECT_EQ(out.outputs[0]->tokens, predictions(3, rejected + 2)); + const int start = rejected + 5; + const auto before = executor_->stats(); + ASSERT_TRUE(executor_->execute( + {{input(sid, start, {prediction(start - 1)})}}, out)); + ASSERT_TRUE(out.outputs[0]); + EXPECT_EQ(out.outputs[0]->tokens, predictions(start, rejected + 1)); + EXPECT_EQ(Peer::hidden(*executor_, sid), hidden_rows(start, rejected + 1)); + Peer::check(*executor_, sid, start + rejected + 1); + const auto after = executor_->stats(); + EXPECT_EQ(after.target_verify_forwards - before.target_verify_forwards, 1); + EXPECT_EQ(after.hidden_rows_d2d - before.hidden_rows_d2d, 6); + EXPECT_EQ(after.hidden_rows_d2h - before.hidden_rows_d2h, rejected + 1); + executor_->close_session(sid); + } +} + +TEST_F(CudaDFlashNative, WrongFeedbackAndNonOutputSlicesKeepPreamble) { + create(); + ASSERT_TRUE(executor_); + const auto sid = open(); + seed(sid, 3); + fixture_->rejected.fill(0); + b::BatchOutput out; + ASSERT_TRUE(executor_->execute({{input(sid, 3, {1})}}, out)); + Peer::check(*executor_, sid, 5); + const auto before = executor_->stats(); + ASSERT_TRUE(executor_->execute({{input(sid, 5, {1})}}, out)); + ASSERT_TRUE(out.outputs[0]); + EXPECT_EQ(out.outputs[0]->tokens, predictions(5, 2)); + EXPECT_EQ( + executor_->stats().pending_decode_batches, before.pending_decode_batches); + EXPECT_EQ( + executor_->stats().target_preamble_calls - before.target_preamble_calls, + 1); + ASSERT_TRUE( + executor_->execute({{input(sid, 7, {prediction(6)}, false)}}, out)); + EXPECT_FALSE(out.outputs[0]); + EXPECT_FALSE(Peer::pending(*executor_, sid)); + Peer::check(*executor_, sid, 8); + ASSERT_TRUE(executor_->execute({{input(sid, 8, {1})}}, out)); + EXPECT_EQ( + executor_->stats().pending_decode_batches, before.pending_decode_batches); + EXPECT_EQ( + executor_->stats().target_preamble_calls - before.target_preamble_calls, + 3); + // Unsupported temperature fails before a target/draft forward. + executor_->set_sampling(sid, {0.5f, 1.0f, 0}, 0); + const auto valid = executor_->stats(); + EXPECT_FALSE(executor_->execute({{input(sid, 10, {prediction(9)})}}, out)); + EXPECT_FALSE(out.outputs[0]); + EXPECT_EQ( + executor_->stats().target_preamble_calls, valid.target_preamble_calls); + EXPECT_EQ(executor_->stats().target_verify_calls, valid.target_verify_calls); + EXPECT_TRUE(Peer::poisoned(*executor_, sid)); +} + +TEST_F(CudaDFlashNative, FinalWholeBlockFitsExactlyAndPreservesAllPrefillRows) { + create(); + ASSERT_TRUE(executor_); + const auto sid = open(); + b::BatchOutput out; + ASSERT_TRUE(executor_->execute({{input(sid, 0, Tokens(78, 1), false)}}, out)); + Peer::check(*executor_, sid, 78); + EXPECT_EQ(executor_->stats().hidden_rows_d2h, 78); + fixture_->rejected.fill(0); + ASSERT_TRUE(executor_->execute({{input(sid, 78, {1})}}, out)); + EXPECT_EQ(out.outputs[0]->tokens, predictions(78, 2)); + Peer::check(*executor_, sid, 80); + fixture_->rejected.fill(15); + const auto before = executor_->stats(); + ASSERT_TRUE(executor_->execute({{input(sid, 80, {prediction(79)})}}, out)); + EXPECT_EQ(out.outputs[0]->tokens, predictions(80, 16)); + EXPECT_EQ( + executor_->stats().target_preamble_calls, before.target_preamble_calls); + Peer::check(*executor_, sid, 96); + EXPECT_EQ(Peer::hidden(*executor_, sid), hidden_rows(80, 16)); +} + +TEST_F(CudaDFlashNative, CloneBudgetRewindCapacityAndFailureAreClosed) { + create(); + ASSERT_TRUE(executor_); + const auto sid = open(); + seed(sid, 3); + b::BatchOutput out; + ASSERT_TRUE(executor_->execute({{input(sid, 3, {1})}}, out)); + Peer::check(*executor_, sid, 20); + auto full = executor_->clone(sid, 20); + auto clipped = executor_->clone(sid, 7); + ASSERT_TRUE(full); + ASSERT_TRUE(clipped); + Peer::check(*executor_, *full, 20); + Peer::check(*executor_, *clipped, 7); + EXPECT_EQ(Peer::hidden(*executor_, *clipped), hidden_rows(4, 3)); + EXPECT_EQ(Peer::pending(*executor_, *full), prediction(19)); + EXPECT_FALSE(Peer::pending(*executor_, *clipped)); + EXPECT_FALSE(executor_->clone(sid, 21)); + executor_->set_sampling(*full, {}, 0); + executor_->set_sampling(*clipped, {}, 0); + fixture_->rejected.fill(0); + const auto before = executor_->stats(); + ASSERT_TRUE(executor_->execute({{input(*full, 20, {prediction(19)})}}, out)); + EXPECT_EQ(out.outputs[0]->tokens, Tokens({prediction(20)})); + EXPECT_EQ( + executor_->stats().target_preamble_calls, before.target_preamble_calls); + Peer::check(*executor_, sid, 20); + ASSERT_TRUE(Peer::rewind(*executor_, sid, 7)); + EXPECT_FALSE(Peer::pending(*executor_, sid)); + EXPECT_EQ(Peer::hidden(*executor_, sid), hidden_rows(4, 3)); + ASSERT_TRUE(executor_->execute({{input(sid, 7, {prediction(6)})}}, out)); + EXPECT_EQ( + executor_->stats().target_preamble_calls - before.target_preamble_calls, + 1); + EXPECT_EQ(out.outputs[0]->tokens, predictions(7, 2)); + // Refill without output to the final trained-block boundary, then decode. + ASSERT_TRUE(executor_->execute({{input(sid, 9, Tokens(72, 1), false)}}, out)); + Peer::check(*executor_, sid, 81); + const auto capacity_before = executor_->stats(); + ASSERT_TRUE(executor_->execute({{input(sid, 81, {1})}}, out)); + EXPECT_EQ(out.outputs[0]->tokens, Tokens({prediction(81)})); + Peer::check(*executor_, sid, 82); + EXPECT_EQ( + executor_->stats().target_verify_calls, + capacity_before.target_verify_calls); + EXPECT_EQ( + executor_->stats().pending_decode_batches, + capacity_before.pending_decode_batches); + // Fail the target sampler after its forward and snapshot have already run. + fixture_->fail_target_sampler = true; + EXPECT_FALSE(executor_->execute({{input(*full, 21, {prediction(20)})}}, out)); + EXPECT_FALSE(out.outputs[0]); + EXPECT_EQ(executor_->stats().failed_batches, 1); + for (const auto owner : {sid, *full, *clipped}) { + EXPECT_TRUE(Peer::poisoned(*executor_, owner)); + EXPECT_TRUE(Peer::hidden(*executor_, owner).empty()); + EXPECT_FALSE(Peer::pending(*executor_, owner)); + } + EXPECT_FALSE(executor_->open_session()); + EXPECT_FALSE(executor_->clone(*full, 0)); + EXPECT_FALSE(executor_->initialize()); +} + +TEST_F(CudaDFlashNative, UnmodifiedOldArtifactRunsSoloAndFourOwnerDecode) { + create(false); + ASSERT_TRUE(executor_); + std::array owners; + b::BatchInput batch; + for (int owner = 0; owner < 4; ++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; + ASSERT_TRUE(executor_->execute(batch, out)); + // Frozen eager first-token goldens, not outputs from this implementation. + const std::array first{7, 194, 84, 217}; + for (int owner = 0; owner < 4; ++owner) { + ASSERT_TRUE(out.outputs[owner]); + ASSERT_FALSE(out.outputs[owner]->tokens.empty()); + EXPECT_EQ(out.outputs[owner]->tokens.front(), first[owner]); + } + for (int round = 0; round < 4; ++round) { + batch.inputs.clear(); + for (int owner = 0; owner < 4; ++owner) { + ASSERT_TRUE(out.outputs[owner]); + ASSERT_FALSE(out.outputs[owner]->tokens.empty()); + const int start = Peer::frontier(*executor_, owners[owner]); + Peer::check(*executor_, owners[owner], start); + batch.inputs.push_back( + input(owners[owner], start, {out.outputs[owner]->tokens.back()})); + } + const auto before = executor_->stats(); + ASSERT_TRUE(executor_->execute(batch, out)); + EXPECT_EQ( + executor_->stats().target_preamble_calls, before.target_preamble_calls); + EXPECT_EQ( + executor_->stats().pending_decode_owners - before.pending_decode_owners, + 4); + for (int owner = 0; owner < 4; ++owner) { + ASSERT_TRUE(out.outputs[owner]); + Peer::check( + *executor_, + owners[owner], + batch.inputs[owner].position + out.outputs[owner]->tokens.size()); + } + } + const auto before = executor_->stats(); + const int start = Peer::frontier(*executor_, owners[0]); + ASSERT_TRUE(executor_->execute( + {{input(owners[0], start, {out.outputs[0]->tokens.back()})}}, out)); + ASSERT_TRUE(out.outputs[0]); + Peer::check(*executor_, owners[0], start + out.outputs[0]->tokens.size()); + EXPECT_EQ( + executor_->stats().pending_decode_owners - before.pending_decode_owners, + 1); + EXPECT_EQ( + executor_->stats().target_preamble_calls, before.target_preamble_calls); + EXPECT_EQ(executor_->stats().failed_batches, 0); +} +} // namespace