diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp index 979be24b72d..8b4572d0833 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp @@ -156,6 +156,122 @@ bool copy_host( cudaStreamSynchronize(stream) == cudaSuccess; } +} // namespace + +namespace cuda_dflash_detail { +HiddenSnapshot::~HiddenSnapshot() { + fence(); + if (device_) { + cudaFree(device_); + } +} + +bool HiddenSnapshot::fence() { + if (stream_set_ && + cudaStreamSynchronize(static_cast(stream_)) != + cudaSuccess) { + return false; + } + stream_set_ = false; + return true; +} + +bool HiddenSnapshot::begin(std::size_t rows, std::size_t row_bytes) { + if (!fence() || rows == 0 || row_bytes == 0 || + rows > std::numeric_limits::max() / row_bytes) { + return false; + } + rows_ = rows; + row_bytes_ = row_bytes; + captured_ = 0; + device_source_.reset(); + stream_ = nullptr; + stream_set_ = false; + return true; +} + +bool HiddenSnapshot::capture( + const aten::Tensor& taps, + std::size_t offset, + std::size_t rows) { + const bool device = taps.device_type() == aten::DeviceType::CUDA; + if (offset != captured_ || offset > rows_ || rows == 0 || + rows > rows_ - offset || !dense(taps) || + taps.nbytes() < rows * row_bytes_ || + (!device && taps.device_type() != aten::DeviceType::CPU) || + (device_source_ && *device_source_ != device)) { + return false; + } + device_source_ = device; + const auto bytes = rows * row_bytes_; + if (device) { + const auto stream = + ::executorch::extension::cuda::getCallerStream().value_or( + cudaStreamPerThread); + if (stream_set_ && stream_ != stream) { + return false; + } + stream_ = stream; + stream_set_ = true; + if (capacity_ < rows_ * row_bytes_) { + if (device_ && cudaFree(device_) != cudaSuccess) { + return false; + } + device_ = nullptr; + capacity_ = 0; + if (cudaMalloc(&device_, rows_ * row_bytes_) != cudaSuccess) { + return false; + } + capacity_ = rows_ * row_bytes_; + } + if (cudaMemcpyAsync( + static_cast(device_) + offset * row_bytes_, + taps.const_data_ptr(), + bytes, + cudaMemcpyDeviceToDevice, + stream) != cudaSuccess) { + return false; + } + } else { + host_.resize(rows_ * row_bytes_); + std::memcpy( + host_.data() + offset * row_bytes_, taps.const_data_ptr(), bytes); + } + captured_ += rows; + return true; +} + +bool HiddenSnapshot::copy_rows( + std::vector& destination, + std::size_t offset, + std::size_t rows) { + if (offset > captured_ || rows > captured_ - offset || row_bytes_ == 0 || + destination.size() % row_bytes_ != 0 || + rows > (destination.max_size() - destination.size()) / row_bytes_) { + return false; + } + const auto start = destination.size(); + const auto bytes = rows * row_bytes_; + destination.resize(start + bytes); + if (bytes == 0) { + return true; + } + if (on_device()) { + stream_set_ = true; + return cudaMemcpyAsync( + destination.data() + start, + static_cast(device_) + offset * row_bytes_, + bytes, + cudaMemcpyDeviceToHost, + static_cast(stream_)) == cudaSuccess; + } + std::memcpy( + destination.data() + start, host_.data() + offset * row_bytes_, bytes); + return true; +} +} // namespace cuda_dflash_detail + +namespace { Result positive_int(Module& module, const char* name) { ET_ASSIGN_OR_RETURN(value, detail::read_required_positive_int(module, name)); ET_CHECK_OR_RETURN_ERROR( @@ -570,20 +686,48 @@ bool MuseGlimmerCudaDFlashExecutor::run_target( return false; } predictions.resize(step.logit_indices.size()); + if (!drain_context && + !verification_hidden_.begin(step.tokens.size(), row_bytes_)) { + return false; + } + if (drain_context) { + ++stats_.target_preamble_calls; + } else { + ++stats_.target_verify_calls; + } std::size_t predicted = 0; const bool ok = target_->execute_raw(step, [&](const cb::RawTargetOutput& output) { - if (!output.taps || + if (drain_context) { + ++stats_.target_preamble_forwards; + } else { + ++stats_.target_verify_forwards; + } + if (!output.taps || output.taps->dim() != 3 || output.taps->scalar_type() != aten::ScalarType::BFloat16 || output.taps->size(2) != hidden_width_) { return false; } - std::vector copied( - static_cast(output.length) * row_bytes_); - // Copy every real tap before invoking the sampler; no borrowed GPU - // output survives any subsequent Module::execute call. - if (!copy_host(copied.data(), *output.taps, copied.size())) { - return false; + std::vector copied; + // The next sampler or target slice may invalidate these borrowed taps. + if (drain_context) { + copied.resize(static_cast(output.length) * row_bytes_); + if (!copy_host(copied.data(), *output.taps, copied.size())) { + return false; + } + if (output.taps->device_type() == aten::DeviceType::CUDA) { + stats_.hidden_rows_d2h += output.length; + stats_.hidden_bytes_d2h += copied.size(); + } + } else { + if (!verification_hidden_.capture( + *output.taps, output.offset, output.length)) { + return false; + } + if (verification_hidden_.on_device()) { + stats_.hidden_rows_d2d += output.length; + stats_.hidden_bytes_d2d += output.length * row_bytes_; + } } auto tokens = target_->argmax_rows(output.logits, output.selected_indices.size()); @@ -594,6 +738,9 @@ bool MuseGlimmerCudaDFlashExecutor::run_target( predictions[output.selected_indices[row]] = (*tokens)[row]; ++predicted; } + if (!drain_context) { + return true; + } std::unordered_set touched; for (int row = 0; row < output.length; ++row) { const int index = output.offset + row; @@ -791,6 +938,7 @@ bool MuseGlimmerCudaDFlashExecutor::speculate( if (!run_target(verification, token_owners, predictions, false)) { return false; } + std::vector bursts; for (std::size_t i = 0; i < owners.size(); ++i) { auto burst = df::verify_greedy( proposals[i], @@ -798,22 +946,42 @@ bool MuseGlimmerCudaDFlashExecutor::speculate( predictions.begin() + i * width, predictions.begin() + (i + 1) * width)); if (!burst || - !rewind( - sessions_.at(owners[i]), - starts[i] + static_cast(burst->tokens.size()))) { + !verification_hidden_.copy_rows( + sessions_.at(owners[i]).hidden, i * width, burst->accepted + 1)) { + return false; + } + if (verification_hidden_.on_device()) { + stats_.hidden_rows_d2h += burst->accepted + 1; + stats_.hidden_bytes_d2h += (burst->accepted + 1) * row_bytes_; + } + bursts.push_back(std::move(*burst)); + } + // All host backing remains stable until this coalesced fence, even when + // verification spanned several target callbacks/static padding shapes. + if (!verification_hidden_.fence()) { + return false; + } + for (std::size_t i = 0; i < owners.size(); ++i) { + auto& state = sessions_.at(owners[i]); + const auto& burst = bursts[i]; + const int committed = starts[i] + static_cast(burst.accepted) + 1; + // Rejected KV was written speculatively, but rejected hidden was never + // appended. Do not use the ordinary rewind/hidden trimming path here. + if (!target_->cache_control().rewind(state.target_id, committed)) { 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. auto& output_tokens = out.outputs[indices[i]]->tokens; output_tokens.insert( - output_tokens.end(), burst->tokens.begin(), burst->tokens.end()); + output_tokens.end(), burst.tokens.begin(), burst.tokens.end()); ++stats_.cycles; stats_.proposed += proposals_; - stats_.accepted += burst->accepted; - stats_.bonuses += burst->bonus; - stats_.rejected += !burst->bonus; + stats_.accepted += burst.accepted; + stats_.bonuses += burst.bonus; + stats_.rejected += !burst.bonus; } } return true; @@ -823,6 +991,8 @@ bool MuseGlimmerCudaDFlashExecutor::fail_batch( const batching::BatchInput& batch, batching::BatchOutput& out) { (void)batch; + // A failed callback/copy may still have queued work touching host backing. + verification_hidden_.fence(); if (!failed_) { failed_ = true; for (auto& item : sessions_) { @@ -932,13 +1102,21 @@ bool MuseGlimmerCudaDFlashExecutor::execute( ++stats_.executions; ET_LOG( Info, - "DFLASH_STATS cycles=%llu proposed=%llu accepted=%llu rejected=%llu bonuses=%llu draft_forwards_total=%llu", + "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", 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), + 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_.hidden_rows_d2h), + static_cast(stats_.hidden_bytes_d2h), + static_cast(stats_.hidden_rows_d2d), + static_cast(stats_.hidden_bytes_d2d)); return true; #if ET_HAS_EXCEPTIONS } catch (...) { diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h index 1e110b11762..cea4c6cc9a2 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h @@ -40,6 +40,34 @@ bool trim_hidden( int draft_frontier, int position); +// Reusable owned storage for borrowed target taps. Capture precedes every +// sampler/next slice; host destinations stay alive until fence(). +class HiddenSnapshot { + public: + ~HiddenSnapshot(); + HiddenSnapshot() = default; + HiddenSnapshot(const HiddenSnapshot&) = delete; + HiddenSnapshot& operator=(const HiddenSnapshot&) = delete; + bool begin(std::size_t rows, std::size_t row_bytes); + bool capture(const aten::Tensor& taps, std::size_t offset, std::size_t rows); + bool copy_rows( + std::vector& destination, + std::size_t offset, + std::size_t rows); + bool fence(); + bool on_device() const { + return device_source_.value_or(false); + } + + private: + void* device_ = nullptr; + void* stream_ = nullptr; + std::size_t capacity_ = 0, rows_ = 0, row_bytes_ = 0, captured_ = 0; + std::vector host_; + std::optional device_source_; + bool stream_set_ = false; +}; + struct DraftContext { std::int32_t id; int position; @@ -66,6 +94,14 @@ struct CudaDFlashStats { std::uint64_t executions = 0; 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 target_preamble_calls = 0; + std::uint64_t target_verify_calls = 0; + std::uint64_t target_preamble_forwards = 0; + std::uint64_t target_verify_forwards = 0; + // Actual CUDA transfers, excluding static padding and CPU-resident taps. + std::uint64_t hidden_rows_d2h = 0, hidden_bytes_d2h = 0; + std::uint64_t hidden_rows_d2d = 0, hidden_bytes_d2d = 0; std::uint64_t cycles = 0; // per-owner verification cycles std::uint64_t proposed = 0; std::uint64_t accepted = 0; @@ -107,8 +143,8 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor { std::int32_t target_id; std::int32_t draft_id; int frontier = 0; - // Host-owned BF16 rows [draft frontier, target frontier). Copied/fenced - // before a sampler or subsequent forward can overwrite borrowed GPU taps. + // Host-owned committed BF16 rows [draft frontier, target frontier). + // Verification appends only the consumed anchor and accepted proposals. std::vector hidden; std::optional sampling; std::optional seed; @@ -155,6 +191,7 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor { batching::SessionId next_session_ = 1; std::unordered_map sessions_; CudaDFlashStats stats_; + cuda_dflash_detail::HiddenSnapshot verification_hidden_; 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 5e8f28f1823..4354af6befd 100644 --- a/examples/models/muse-glimmer/tests/cuda_dflash_executor_test.cpp +++ b/examples/models/muse-glimmer/tests/cuda_dflash_executor_test.cpp @@ -5,9 +5,16 @@ * 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 + +using executorch::extension::from_blob; +namespace aten = executorch::aten; namespace df = executorch::extension::llm::cuda_dflash_detail; using executorch::extension::llm::batching::Token; @@ -129,6 +136,101 @@ TEST(CudaDFlash, NearContextLimitRequiresWholeTrainedBlock) { EXPECT_FALSE(df::pack_draft({{4, 60, 70, &hidden, 1}}, 4, 99, 1, 64, 2)); } +TEST(CudaDFlash, EveryProposalMismatchKeepsOnlyConsumedPrefix) { + const std::vector proposals{11, 12, 13, 14}; + for (std::size_t rejected = 0; rejected <= proposals.size(); ++rejected) { + auto predictions = proposals; + predictions.push_back(90); + if (rejected < proposals.size()) { + predictions[rejected] = 80; + } + const auto burst = df::verify_greedy(proposals, predictions); + ASSERT_TRUE(burst); + std::vector expected( + proposals.begin(), proposals.begin() + rejected); + expected.push_back(rejected == proposals.size() ? 90 : 80); + EXPECT_EQ(burst->tokens, expected); + EXPECT_EQ(burst->accepted, rejected); + EXPECT_EQ(burst->bonus, rejected == proposals.size()); + } +} + +TEST(CudaDFlash, HostSnapshotPartialRowsAndInvalidExtents) { + df::HiddenSnapshot snapshot; + std::vector source{1, 2, 3, 4, 5, 6, 99, 99}; + auto taps = from_blob(source.data(), {1, 4, 2}); + ASSERT_TRUE(snapshot.begin(3, 2)); + EXPECT_FALSE(snapshot.capture(*taps, 1, 3)); + ASSERT_TRUE(snapshot.capture(*taps, 0, 3)); + std::fill(source.begin(), source.end(), 0); + std::vector kept{7, 8}; + EXPECT_FALSE(snapshot.copy_rows(kept, 2, 2)); + ASSERT_TRUE(snapshot.copy_rows(kept, 1, 2)); + ASSERT_TRUE(snapshot.fence()); + EXPECT_EQ(kept, std::vector({7, 8, 3, 4, 5, 6})); + EXPECT_FALSE(snapshot.begin(0, 2)); + EXPECT_FALSE(snapshot.begin(1, 0)); + ASSERT_TRUE(snapshot.begin(1, 2)); + EXPECT_FALSE(snapshot.copy_rows(kept, 0, 1)); +} + +TEST(CudaDFlashCuda, SnapshotSurvivesOverwritesSlicesAndFourUnequalPrefixes) { + int devices = 0; + if (cudaGetDeviceCount(&devices) != cudaSuccess || devices == 0) { + GTEST_SKIP() << "CUDA device required"; + } + void* borrowed = nullptr; + ASSERT_EQ(cudaMalloc(&borrowed, 20), cudaSuccess); + auto taps = from_blob( + borrowed, + {1, 5, 4}, + aten::ScalarType::Byte, + aten::Device(aten::DeviceType::CUDA, 0), + [](void* data) { cudaFree(data); }); + cudaStream_t stream; + ASSERT_EQ( + cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess); + { + executorch::extension::cuda::CallerStreamGuard guard(stream); + df::HiddenSnapshot snapshot; + // Grow, shrink, then reuse capacity across independent verification calls. + for (int total : {16, 8, 16}) { + ASSERT_TRUE(snapshot.begin(total, 4)); + for (int offset = 0; offset < total; offset += 3) { + const int rows = std::min(3, total - offset); + std::vector source(20, 255); // static padding sentinel + for (int byte = 0; byte < rows * 4; ++byte) { + source[byte] = offset * 4 + byte + 1; + } + ASSERT_EQ( + cudaMemcpyAsync( + borrowed, source.data(), 20, cudaMemcpyHostToDevice, stream), + cudaSuccess); + ASSERT_TRUE(snapshot.capture(*taps, offset, rows)); + // Model the sampler/next Module call overwriting graph-owned output. + ASSERT_EQ(cudaMemsetAsync(borrowed, 0, 20, stream), cudaSuccess); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + } + std::array, 4> kept; + const int owners = total / 4; + for (int owner = 0; owner < owners; ++owner) { + kept[owner] = {200, 201, 202, 203}; + ASSERT_TRUE(snapshot.copy_rows(kept[owner], owner * 4, owner + 1)); + } + ASSERT_TRUE(snapshot.fence()); + for (int owner = 0; owner < owners; ++owner) { + std::vector expected{200, 201, 202, 203}; + for (int byte = 0; byte < (owner + 1) * 4; ++byte) { + expected.push_back(owner * 16 + byte + 1); + } + EXPECT_EQ(kept[owner], expected); + } + } + } + // A successfully fenced snapshot must not retain a live-stream requirement. + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + TEST(CudaDFlash, PackingRejectsDuplicateOwnersAndOversizedContext) { std::vector hidden{1, 2}; EXPECT_FALSE(df::pack_draft(