diff --git a/examples/models/muse-glimmer/BATCHING_SERVING.md b/examples/models/muse-glimmer/BATCHING_SERVING.md index 1c1aaca46b0..4bed79eb1ee 100644 --- a/examples/models/muse-glimmer/BATCHING_SERVING.md +++ b/examples/models/muse-glimmer/BATCHING_SERVING.md @@ -3,7 +3,8 @@ The batching runner and multiplexed worker support MLX autoregressive generation with text only or text plus one image across the submitted conversation, and CUDA autoregressive generation with text only. They are opt-in entry points; the -legacy runners and export defaults remain unchanged. DFlash batching is MLX only. +legacy runners and export defaults remain unchanged. DFlash batching runs on MLX +and CUDA (text only). CUDA specifics are in [CUDA](#cuda) below. ## Export @@ -227,3 +228,10 @@ snapshots are enabled. Their total reservation must leave one cell for padded writes within the exported `--max-cells`. The cell pool tracks at most 64 resident sessions. Prefix snapshots clone a session's cells by reference, sharing the prompt's KV rather than copying it. + +For DFlash, export the tapped target and packed draft with +`export_dflash --backend cuda --use-offgraph-kv-cache` (add `--draft-max-cells` +and `--max-draft-owners`), then add `--dflash` to the serve command above. The +draft ABI and `--dflash-n-draft` behave as on MLX. Drafts are greedy; the target +accepts or resamples them on the device under each request's sampling. Decode +drafts of up to four owners run in CUDA graphs. diff --git a/examples/models/muse-glimmer/CMakeLists.txt b/examples/models/muse-glimmer/CMakeLists.txt index 70e24f3d61c..6129197e8a8 100644 --- a/examples/models/muse-glimmer/CMakeLists.txt +++ b/examples/models/muse-glimmer/CMakeLists.txt @@ -242,7 +242,7 @@ if(MUSE_GLIMMER_BUILD_BATCHING OR MUSE_GLIMMER_BUILD_BATCHING_TESTS) if(MUSE_GLIMMER_BUILD_BATCHING) if(_batching_backend STREQUAL "cuda") - set(_batching_executor) + set(_batching_executor runtime/cuda_dflash_executor.cpp) set(_batching_backend_library cuda_batching) else() set(_batching_executor runtime/mlx_executor.cpp) @@ -356,5 +356,24 @@ if(MUSE_GLIMMER_BUILD_BATCHING OR MUSE_GLIMMER_BUILD_BATCHING_TESTS) extension_llm_serving Threads::Threads ) add_test(NAME muse_glimmer_batching COMMAND muse_glimmer_batching_test) + if(MUSE_GLIMMER_BUILD_BATCHING AND _batching_backend STREQUAL "cuda") + # Tokenizer include paths contain regex.h; keep GoogleTest and its users + # on the same simple-regex ABI instead of the shadowed POSIX header. + add_subdirectory( + ${EXECUTORCH_ROOT}/third-party/googletest + ${CMAKE_CURRENT_BINARY_DIR}/dflash_googletest EXCLUDE_FROM_ALL + ) + target_compile_definitions(gtest PUBLIC GTEST_HAS_POSIX_RE=0) + add_executable( + muse_glimmer_cuda_dflash_test tests/cuda_dflash_executor_test.cpp + ) + target_link_libraries( + muse_glimmer_cuda_dflash_test PRIVATE muse_glimmer_batching_runtime + GTest::gtest_main + ) + add_test(NAME muse_glimmer_cuda_dflash + COMMAND muse_glimmer_cuda_dflash_test + ) + endif() endif() endif() diff --git a/examples/models/muse-glimmer/runtime/BUCK b/examples/models/muse-glimmer/runtime/BUCK new file mode 100644 index 00000000000..1cdb774e0d8 --- /dev/null +++ b/examples/models/muse-glimmer/runtime/BUCK @@ -0,0 +1,34 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +fbcode_target( + _kind = runtime.cxx_library, + name = "cuda_dflash_executor", + srcs = ["cuda_dflash_executor.cpp"], + header_namespace = "executorch/examples/models/muse-glimmer/runtime", + exported_headers = [ + "backend_factory.h", + "cuda_dflash_executor.h", + "engine/muse_glimmer_vision_runtime.h", + "prompt_preparation.h", + ], + visibility = ["PUBLIC"], + deps = [ + "//executorch/extension/cuda:caller_stream", + "//executorch/extension/llm/runner:stats", + "//executorch/runtime/backend:interface", + ], + exported_deps = [ + "fbsource//third-party/nlohmann-json:nlohmann-json", + "//executorch/backends/cuda/batching:cuda_executor", + "//executorch/extension/llm/batching:batching", + "//executorch/extension/llm/batching:executor_utils", + "//executorch/extension/llm/cache:kv_cache", + "//executorch/extension/llm/serving:serving_runtime", + "//executorch/extension/module:module", + "//executorch/extension/tensor:tensor", + ], + external_deps = [("cuda", None, "cuda-lazy"), "nlohmann_json"], +) diff --git a/examples/models/muse-glimmer/runtime/backend_factory.cpp b/examples/models/muse-glimmer/runtime/backend_factory.cpp index cb2aaa17713..0916ceb4fb0 100644 --- a/examples/models/muse-glimmer/runtime/backend_factory.cpp +++ b/examples/models/muse-glimmer/runtime/backend_factory.cpp @@ -11,6 +11,7 @@ #endif #if defined(EXECUTORCH_BUILD_CUDA) #include +#include #endif #include @@ -32,9 +33,12 @@ runtime::Result create_muse_glimmer_backend( InvalidArgument, "Muse Glimmer CUDA: a session needs room for a prompt and prediction"); ET_CHECK_OR_RETURN_ERROR( - !config.use_dflash && config.dflash_n_draft == 0, - NotSupported, - "Muse Glimmer CUDA: DFlash batching is not supported"); + config.use_dflash || config.dflash_n_draft == 0, + InvalidArgument, + "Muse Glimmer CUDA: dflash_n_draft requires DFlash"); + if (config.use_dflash) { + return MuseGlimmerCudaDFlashExecutor::create(config); + } cb::CudaExecutorConfig executor_config; executor_config.model_path = config.model_path; if (!config.data_path.empty()) { diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp new file mode 100644 index 00000000000..472bd71fe7d --- /dev/null +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp @@ -0,0 +1,954 @@ +/* + * 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 +#include +#include + +namespace executorch::extension::llm { +using batching::SessionId; +using batching::Token; +using runtime::Error; +using runtime::Result; +namespace df = cuda_dflash_detail; +namespace cb = ::executorch::backends::cuda::batching; + +namespace cuda_dflash_detail { +std::optional verify_greedy( + const std::vector& proposals, + const std::vector& predictions) { + if (proposals.empty() || predictions.size() != proposals.size() + 1) { + return std::nullopt; + } + GreedyBurst result; + for (std::size_t row = 0; row < predictions.size(); ++row) { + if (row < proposals.size() && proposals[row] == predictions[row]) { + result.tokens.push_back(proposals[row]); + ++result.accepted; + } else { + result.tokens.push_back(predictions[row]); + result.bonus = row == proposals.size(); + break; + } + } + return result; +} + +bool trim_hidden( + std::vector& hidden, + std::size_t row_bytes, + int draft_frontier, + int position) { + if (row_bytes == 0 || draft_frontier < 0 || position < 0 || + hidden.size() % row_bytes != 0) { + return false; + } + const auto rows = + static_cast(std::max(0, position - draft_frontier)); + if (rows > hidden.size() / row_bytes) { + return false; + } + hidden.resize(rows * row_bytes); + return true; +} + +std::optional pack_draft( + const std::vector& owners, + int block, + Token mask, + int context_limit, + int max_context, + std::size_t row_bytes) { + if (owners.empty() || owners.size() > 4 || block < 2 || context_limit < 0 || + max_context < block || row_bytes == 0 || + block > (std::numeric_limits::max() - context_limit) / + static_cast(owners.size())) { + return std::nullopt; + } + DraftPlan plan; + const int q = static_cast(owners.size()) * block; + plan.tokens.assign(q, static_cast(mask)); + std::unordered_set seen; + int c = 0; + for (std::size_t owner = 0; owner < owners.size(); ++owner) { + const auto& row = owners[owner]; + if (!seen.insert(row.id).second || row.id < 0 || !row.hidden || + row.hidden->size() % row_bytes != 0 || row.count < 0 || + row.position < 0 || row.count > context_limit - c || + row.position > max_context - block - row.count || + static_cast(row.count) > row.hidden->size() / row_bytes) { + return std::nullopt; + } + plan.tokens[owner * block] = static_cast(row.anchor); + plan.counts.push_back(row.count); + plan.ends.push_back(row.position + row.count); + plan.hidden.insert( + plan.hidden.end(), + row.hidden->begin(), + row.hidden->begin() + static_cast(row.count) * row_bytes); + for (int j = 0; j < row.count; ++j) { + plan.positions.push_back(row.position + j); + } + plan.ids.insert(plan.ids.end(), row.count + block, row.id); + c += row.count; + } + int offset = 0; + for (std::size_t owner = 0; owner < owners.size(); ++owner) { + for (int j = 0; j < plan.counts[owner]; ++j) { + plan.kv_order.push_back(offset++); + } + for (int j = 0; j < block; ++j) { + plan.positions.push_back(plan.ends[owner] + j); + plan.kv_order.push_back(c + owner * block + j); + } + } + return plan; +} +} // namespace cuda_dflash_detail + +namespace { +bool dense(const aten::Tensor& tensor) { + std::int64_t stride = 1; + for (auto d = tensor.dim() - 1; d >= 0; --d) { + if (tensor.size(d) > 1 && tensor.strides()[d] != stride) { + return false; + } + stride *= tensor.size(d); + } + return true; +} + +bool copy_host( + void* destination, + const aten::Tensor& tensor, + std::size_t bytes) { + if (tensor.device_type() == aten::DeviceType::CPU) { + std::memcpy(destination, tensor.const_data_ptr(), bytes); + return true; + } + if (tensor.device_type() != aten::DeviceType::CUDA) { + return false; + } + const auto stream = ::executorch::extension::cuda::getCallerStream().value_or( + cudaStreamPerThread); + return cudaMemcpyAsync( + destination, + tensor.const_data_ptr(), + bytes, + cudaMemcpyDeviceToHost, + stream) == cudaSuccess && + cudaStreamSynchronize(stream) == cudaSuccess; +} + +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( + value <= std::numeric_limits::max(), + InvalidProgram, + "DFlash metadata %s is too large", + name); + return static_cast(value); +} + +Result draft_geometry(Module& module, int block) { + ET_ASSIGN_OR_RETURN(layers, positive_int(module, "get_draft_n_caches")); + ET_ASSIGN_OR_RETURN( + heads, detail::read_int_tensor_method(module, "get_draft_kv_heads")); + ET_ASSIGN_OR_RETURN( + dims, detail::read_int_tensor_method(module, "get_draft_head_dims")); + ET_ASSIGN_OR_RETURN( + windows, detail::read_int_tensor_method(module, "get_draft_windows")); + ET_CHECK_OR_RETURN_ERROR( + heads.size() == static_cast(layers) && + dims.size() == heads.size() && windows.size() == heads.size(), + InvalidProgram, + "DFlash geometry vector sizes disagree"); + cache::CacheGeometry geometry; + geometry.dflash_block_size = block; + for (int i = 0; i < layers; ++i) { + geometry.layers.push_back( + {{windows[i] ? cache::LayerPolicy::Kind::Ring + : cache::LayerPolicy::Kind::Flat, + windows[i]}, + heads[i], + dims[i]}); + } + ET_CHECK_OR_RETURN_ERROR( + cache::valid(geometry), InvalidProgram, "Invalid DFlash draft geometry"); + return geometry; +} + +Result cache_key_metadata(Module& module, const char* name) { + ET_ASSIGN_OR_RETURN(outputs, module.execute(name)); + ET_CHECK_OR_RETURN_ERROR( + outputs.size() == 1 && outputs[0].isString() && + !outputs[0].toString().empty(), + InvalidProgram, + "DFlash %s must be a nonempty string", + name); + return std::string(outputs[0].toString()); +} + +// Nothing is mutated yet, so the batch is refused without poisoning the model. +bool reject_batch() { + ET_LOG(Error, "CUDA DFlash rejected an invalid batch"); + return false; +} +} // namespace + +MuseGlimmerCudaDFlashExecutor::MuseGlimmerCudaDFlashExecutor( + std::unique_ptr target, + std::shared_ptr draft_cache, + int block, + int proposals, + Token mask, + int max_owners, + int fresh_context, + int max_sessions, + int max_context, + int hidden_width) + : draft_cache_(std::move(draft_cache)), + draft_install_(draft_cache_), + target_(std::move(target)), + draft_ctl_(draft_cache_->as()), + block_(block), + proposals_(proposals), + max_owners_(max_owners), + fresh_context_(fresh_context), + max_sessions_(max_sessions), + max_context_(max_context), + hidden_width_(hidden_width), + mask_(mask), + row_bytes_(static_cast(hidden_width) * 2) {} + +MuseGlimmerCudaDFlashExecutor::~MuseGlimmerCudaDFlashExecutor() { + for (const auto& item : sessions_) { + target_->cache_control().seq_rm(item.second.target_id); + if (!failed_) { + draft_ctl_->seq_rm(item.second.draft_id); + } + } +} + +Result MuseGlimmerCudaDFlashExecutor::create( + const MuseGlimmerBackendConfig& config) { + ET_CHECK_OR_RETURN_ERROR( + config.backend == "cuda" && config.use_dflash && + !config.model_path.empty() && config.max_sessions > 0 && + config.max_session_tokens >= 0 && config.max_session_tokens != 1 && + config.dflash_n_draft >= 0, + InvalidArgument, + "Invalid CUDA DFlash configuration"); + cb::CudaExecutorConfig target_config; + target_config.model_path = config.model_path; + if (!config.data_path.empty()) { + target_config.data_files.push_back(config.data_path); + } + target_config.max_sessions = config.max_sessions; + target_config.max_session_tokens = config.max_session_tokens; + ET_ASSIGN_OR_RETURN(target, Target::create(target_config)); + auto& module = target->module(); + ET_ASSIGN_OR_RETURN(names, module.method_names()); + ET_CHECK_OR_RETURN_ERROR( + names.count(df::kForward) && + !names.count("get_mutable_buffer_metadata") && + !names.count("vision_encoder"), + InvalidProgram, + "CUDA DFlash requires a tapped text-only off-graph artifact"); + ET_ASSIGN_OR_RETURN(block, positive_int(module, df::kBlock)); + ET_ASSIGN_OR_RETURN( + declared_block, positive_int(module, "get_draft_dflash_block_size")); + ET_ASSIGN_OR_RETURN(max_owners, positive_int(module, df::kMaxOwners)); + ET_ASSIGN_OR_RETURN(max_write, positive_int(module, df::kMaxWrite)); + ET_ASSIGN_OR_RETURN(max_cells, positive_int(module, df::kMaxCells)); + ET_ASSIGN_OR_RETURN(tap_layers, positive_int(module, df::kTargetLayers)); + ET_ASSIGN_OR_RETURN( + layer_ids, + detail::read_int_tensor_method(module, "get_dflash_target_layers")); + ET_ASSIGN_OR_RETURN(mask, detail::read_int_method(module, df::kMask)); + ET_ASSIGN_OR_RETURN( + target_key, cache_key_metadata(module, "get_dflash_target_cache_key")); + ET_ASSIGN_OR_RETURN( + draft_key, cache_key_metadata(module, "get_dflash_draft_cache_key")); + ET_CHECK_OR_RETURN_ERROR( + block >= 2 && block == declared_block && max_owners <= 4 && + block <= (std::numeric_limits::max() - 1) / max_owners && + max_write > max_owners * block && max_cells > max_write && mask && + *mask >= 0 && *mask < target->vocab_size() && + target_key != draft_key && + layer_ids.size() == static_cast(tap_layers) && + !layer_ids.empty() && layer_ids.front() >= 0 && + std::adjacent_find( + layer_ids.begin(), layer_ids.end(), std::greater_equal()) == + layer_ids.end(), + InvalidProgram, + "Invalid packed CUDA DFlash metadata"); + const int proposals = + config.dflash_n_draft == 0 ? block - 1 : config.dflash_n_draft; + ET_CHECK_OR_RETURN_ERROR( + proposals > 0 && proposals < block, + InvalidArgument, + "DFlash n_draft must be in [1, trained block size - 1]"); + int hidden_width = 0; + for (const auto& method : target->method_calls()) { + ET_ASSIGN_OR_RETURN(meta, module.method_meta(method.name)); + ET_CHECK_OR_RETURN_ERROR( + meta.num_outputs() == 2, + InvalidProgram, + "DFlash target %s needs logits and all-token taps", + method.name.c_str()); + ET_ASSIGN_OR_RETURN(taps, meta.output_tensor_meta(1)); + ET_ASSIGN_OR_RETURN(tokens, meta.input_tensor_meta(0)); + ET_CHECK_OR_RETURN_ERROR( + taps.scalar_type() == aten::ScalarType::BFloat16 && + taps.sizes().size() == 3 && taps.sizes()[0] == 1 && + taps.sizes()[1] == tokens.sizes()[1] && taps.sizes()[2] > 0 && + taps.sizes()[2] % tap_layers == 0 && + (hidden_width == 0 || hidden_width == taps.sizes()[2]), + InvalidProgram, + "DFlash target taps have incompatible shape/dtype"); + hidden_width = taps.sizes()[2]; + } + ET_ASSIGN_OR_RETURN(draft_meta, module.method_meta(df::kForward)); + ET_CHECK_OR_RETURN_ERROR( + draft_meta.num_inputs() == 4 && draft_meta.num_outputs() == 1 && + draft_meta.num_backends() > 0, + InvalidProgram, + "DFlash draft needs four packed inputs and one logits output"); + for (std::size_t i = 0; i < draft_meta.num_backends(); ++i) { + ET_ASSIGN_OR_RETURN(backend, draft_meta.get_backend_name(i)); + ET_CHECK_OR_RETURN_ERROR( + std::string(backend) == ::executorch::backends::cuda::kCudaBackendId, + InvalidProgram, + "DFlash draft must delegate only to CUDA"); + } + ET_ASSIGN_OR_RETURN(draft_tokens, draft_meta.input_tensor_meta(0)); + ET_ASSIGN_OR_RETURN(draft_hidden, draft_meta.input_tensor_meta(1)); + ET_ASSIGN_OR_RETURN(draft_positions, draft_meta.input_tensor_meta(2)); + ET_ASSIGN_OR_RETURN(draft_order, draft_meta.input_tensor_meta(3)); + ET_ASSIGN_OR_RETURN(draft_logits, draft_meta.output_tensor_meta(0)); + ET_CHECK_OR_RETURN_ERROR( + draft_tokens.scalar_type() == aten::ScalarType::Long && + draft_tokens.sizes().size() == 2 && draft_tokens.sizes()[0] > 0 && + draft_tokens.sizes()[1] == block && + draft_hidden.scalar_type() == aten::ScalarType::BFloat16 && + draft_hidden.sizes().size() == 3 && draft_hidden.sizes()[0] == 1 && + draft_hidden.sizes()[2] == hidden_width && + draft_positions.scalar_type() == aten::ScalarType::Long && + draft_positions.sizes().size() == 1 && + draft_order.scalar_type() == aten::ScalarType::Long && + draft_order.sizes().size() == 1 && + draft_order.sizes()[0] == draft_positions.sizes()[0] && + draft_logits.scalar_type() == aten::ScalarType::Float && + draft_logits.sizes().size() >= 2 && + draft_logits.sizes()[draft_logits.sizes().size() - 1] == + target->vocab_size(), + InvalidProgram, + "DFlash draft packed tensor ABI mismatch"); + ET_ASSIGN_OR_RETURN(geometry, draft_geometry(module, block)); + const int temporary = max_owners * block; + const int draft_budget = (max_cells - 1 - temporary) / config.max_sessions; + const int max_context = std::min(target->max_session_tokens(), draft_budget); + ET_CHECK_OR_RETURN_ERROR( + max_context > block && draft_tokens.sizes()[0] == max_owners && + draft_hidden.sizes()[1] >= max_write - temporary && + draft_positions.sizes()[0] >= max_write && + (!config.max_session_tokens || + max_context == config.max_session_tokens), + InvalidArgument, + "DFlash draft pool cannot hold resident contexts plus temporary owner blocks"); + cache::CacheConfig cache_config{}; + cache_config.capacity = max_cells; + cache_config.kv_dtype = static_cast(aten::ScalarType::BFloat16); + cache_config.max_write = max_write; + ET_ASSIGN_OR_RETURN( + draft_cache, + cache::CacheFactory::global().build( + ::executorch::backends::cuda::kCudaBackendId, + cache::kind::kBatchedCell, + geometry, + cache_config)); + const auto* control = draft_cache->as(); + ET_CHECK_OR_RETURN_ERROR( + control && draft_cache->as<::executorch::backends::cuda::CudaKVCache>() && + (!control->max_seqs() || config.max_sessions <= *control->max_seqs()), + InvalidProgram, + "DFlash draft cache does not support requested resident sessions"); + ET_CHECK_OR_RETURN_ERROR( + target->vocab_size() > 200092 && + config.bos_id < static_cast(target->vocab_size()), + InvalidProgram, + "Muse Glimmer DFlash vocabulary cannot hold its special tokens"); + auto spec = std::make_shared( + MuseGlimmerPreparationSpec{ + aten::ScalarType::BFloat16, + 0, + max_context, + static_cast(target->preferred_batch_tokens()), + target->vocab_size(), + false, + 0, + config.image_limits, + config.bos_id}); + auto executor = + std::unique_ptr(new MuseGlimmerCudaDFlashExecutor( + std::move(target), + std::move(draft_cache), + block, + proposals, + static_cast(*mask), + max_owners, + max_write - temporary, + config.max_sessions, + max_context, + hidden_width)); + return MuseGlimmerBackend{std::move(executor), std::move(spec)}; +} + +std::size_t MuseGlimmerCudaDFlashExecutor::preferred_batch_tokens() const { + return target_->preferred_batch_tokens(); +} + +bool MuseGlimmerCudaDFlashExecutor::initialize() { + if (failed_) { + return false; + } + if (initialized_) { + return true; + } + if (!target_->initialize() || + batching::load_method_with_cache( + target_->module(), + df::kForward, + ::executorch::backends::cuda::kCudaBackendId, + draft_install_) != Error::Ok) { + ET_LOG( + Error, + "CUDA DFlash: could not initialize distinct target/draft caches"); + return false; + } + initialized_ = true; + return true; +} + +std::optional MuseGlimmerCudaDFlashExecutor::open_session() { + if (failed_ || sessions_.size() >= static_cast(max_sessions_) || + next_session_ == std::numeric_limits::max()) { + return std::nullopt; + } + auto target_id = target_->cache_control().seq_new(); + if (!target_id) { + return std::nullopt; + } + auto draft_id = draft_ctl_->seq_new(); + if (!draft_id) { + target_->cache_control().seq_rm(*target_id); + return std::nullopt; + } +#if ET_HAS_EXCEPTIONS + try { +#endif + const auto session = next_session_++; + sessions_.emplace(session, State{*target_id, *draft_id}); + return session; +#if ET_HAS_EXCEPTIONS + } catch (...) { + target_->cache_control().seq_rm(*target_id); + draft_ctl_->seq_rm(*draft_id); + return std::nullopt; + } +#endif +} + +void MuseGlimmerCudaDFlashExecutor::close_session(SessionId session) { + auto it = sessions_.find(session); + if (it == sessions_.end()) { + return; + } + target_->cache_control().seq_rm(it->second.target_id); + if (!failed_) { + draft_ctl_->seq_rm(it->second.draft_id); + } + sessions_.erase(it); +} + +std::optional MuseGlimmerCudaDFlashExecutor::clone( + SessionId source, + batching::Position upto) { + const auto it = sessions_.find(source); + if (failed_ || it == sessions_.end() || it->second.poisoned || upto < 0 || + upto > it->second.frontier || + sessions_.size() >= static_cast(max_sessions_) || + next_session_ == std::numeric_limits::max()) { + return std::nullopt; + } + const auto& state = it->second; + const int draft_position = draft_ctl_->pos(state.draft_id); + const int draft_upto = std::min(upto, draft_position); + auto target_id = upto == 0 + ? target_->cache_control().seq_new() + : target_->cache_control().seq_clone(state.target_id, upto); + if (!target_id) { + return std::nullopt; + } + auto draft_id = draft_upto == 0 + ? draft_ctl_->seq_new() + : draft_ctl_->seq_clone(state.draft_id, draft_upto); + if (!draft_id) { + target_->cache_control().seq_rm(*target_id); + return std::nullopt; + } +#if ET_HAS_EXCEPTIONS + try { +#endif + State copy{*target_id, *draft_id}; + copy.frontier = upto; + copy.hidden = state.hidden; + if (!df::trim_hidden(copy.hidden, row_bytes_, draft_position, upto)) { + target_->cache_control().seq_rm(*target_id); + draft_ctl_->seq_rm(*draft_id); + return std::nullopt; + } + const auto session = next_session_++; + sessions_.emplace(session, std::move(copy)); + return session; +#if ET_HAS_EXCEPTIONS + } catch (...) { + target_->cache_control().seq_rm(*target_id); + draft_ctl_->seq_rm(*draft_id); + return std::nullopt; + } +#endif +} + +void MuseGlimmerCudaDFlashExecutor::set_sampling( + SessionId session, + const batching::SamplingParams& params, + std::optional seed) { + const auto it = sessions_.find(session); + if (it != sessions_.end()) { + it->second.sampling = params; + it->second.seed = seed; + if (!std::isfinite(params.temperature) || params.temperature != 0.0f) { + ET_LOG( + Error, + "CUDA DFlash supports temperature=0 only; sampled DFlash is not implemented"); + } + } +} + +bool MuseGlimmerCudaDFlashExecutor::rewind(State& state, int position) { + const int draft_position = draft_ctl_->pos(state.draft_id); + if (position > state.frontier || + !df::trim_hidden(state.hidden, row_bytes_, draft_position, position) || + !target_->cache_control().rewind(state.target_id, position) || + (position < draft_position && + !draft_ctl_->rewind(state.draft_id, position))) { + state.poisoned = true; + return false; + } + state.frontier = position; + return true; +} + +bool MuseGlimmerCudaDFlashExecutor::run_target( + const RawStep& step, + const std::vector& token_owners, + std::vector& predictions, + bool drain_context) { + if (token_owners.size() != step.tokens.size()) { + return false; + } + predictions.resize(step.logit_indices.size()); + std::size_t predicted = 0; + const bool ok = + target_->execute_raw(step, [&](const cb::RawTargetOutput& output) { + if (!output.taps || + 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; + } + auto tokens = + target_->argmax_rows(output.logits, output.selected_indices.size()); + if (!tokens) { + return false; + } + for (std::size_t row = 0; row < tokens->size(); ++row) { + predictions[output.selected_indices[row]] = (*tokens)[row]; + ++predicted; + } + std::unordered_set touched; + for (int row = 0; row < output.length; ++row) { + const int index = output.offset + row; + const auto owner = token_owners[index]; + auto& state = sessions_.at(owner); + state.hidden.insert( + state.hidden.end(), + copied.begin() + static_cast(row) * row_bytes_, + copied.begin() + static_cast(row + 1) * row_bytes_); + state.frontier = static_cast(step.positions[index]) + 1; + touched.insert(owner); + } + if (drain_context) { + for (const auto owner : touched) { + if (!drain(owner)) { + return false; + } + } + } + return true; + }); + return ok && + predicted == + static_cast(std::count_if( + step.logit_indices.begin(), step.logit_indices.end(), [](auto index) { + return index >= 0; + })); +} + +bool MuseGlimmerCudaDFlashExecutor::drain(SessionId owner) { + auto& state = sessions_.at(owner); + while (state.hidden.size() / row_bytes_ > + static_cast(fresh_context_)) { + const int count = std::min( + fresh_context_, + max_context_ - block_ - draft_ctl_->pos(state.draft_id)); + if (count <= 0) { + break; // Near the context limit: ordinary target decode, no draft round. + } + if (!draft({owner}, {mask_}, count, nullptr)) { + return false; + } + } + return true; +} + +bool MuseGlimmerCudaDFlashExecutor::draft( + const std::vector& owners, + const std::vector& anchors, + int context_limit, + std::vector>* proposals) { + if (owners.size() != anchors.size() || owners.empty() || + owners.size() > static_cast(max_owners_)) { + return false; + } + std::vector contexts; + int remaining = context_limit; + for (std::size_t i = 0; i < owners.size(); ++i) { + auto& state = sessions_.at(owners[i]); + if (remaining < 0) { + return false; + } + const auto rows = state.hidden.size() / row_bytes_; + const int count = + static_cast(std::min(static_cast(remaining), rows)); + if (proposals && rows != static_cast(count)) { + return false; + } + contexts.push_back( + {state.draft_id, + draft_ctl_->pos(state.draft_id), + anchors[i], + &state.hidden, + count}); + remaining -= count; + } + auto plan = df::pack_draft( + contexts, block_, mask_, context_limit, max_context_, row_bytes_); + if (!plan || !draft_ctl_->declare_step(plan->ids)) { + return false; + } + const int q = static_cast(owners.size()) * block_; + const int c = static_cast(plan->hidden.size() / row_bytes_); + const int k = c + 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_}, + 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}); + if (!result.ok() || result->size() != 1 || !result->at(0).isTensor()) { + return false; + } + ++stats_.draft_forwards; + const auto logits = result->at(0).toTensor(); + if (logits.dim() < 2 || logits.scalar_type() != aten::ScalarType::Float || + !dense(logits) || + logits.numel() != static_cast(q) * target_->vocab_size() || + logits.size(logits.dim() - 1) != target_->vocab_size()) { + return false; + } + std::optional> picked; + if (proposals) { + picked = target_->argmax_rows(logits, q); + if (!picked) { + return false; + } + proposals->clear(); + } + for (std::size_t i = 0; i < owners.size(); ++i) { + auto& state = sessions_.at(owners[i]); + // The cache worker fences and removes each temporary trailing block at + // commit_step. A retained block is an ABI failure, never a manual fixup. + if (draft_ctl_->pos(state.draft_id) != plan->ends[i]) { + return false; + } + state.hidden.erase( + state.hidden.begin(), + state.hidden.begin() + + static_cast(plan->counts[i]) * row_bytes_); + if (proposals) { + std::vector prefix; + for (int j = 1; j <= proposals_; ++j) { + prefix.push_back((*picked)[i * block_ + j]); + } + proposals->push_back(std::move(prefix)); + } + } + stats_.warmed_context_tokens += c; + return true; +} + +bool MuseGlimmerCudaDFlashExecutor::speculate( + const batching::BatchInput& batch, + batching::BatchOutput& out) { + std::size_t cursor = 0; + while (cursor < batch.inputs.size()) { + std::vector owners; + std::vector anchors; + std::vector indices; + int context = 0; + while (cursor < batch.inputs.size()) { + const auto& input = batch.inputs[cursor]; + 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_) { + ++cursor; + continue; + } + if (!drain(input.sid)) { + return false; + } + const auto rows = state.hidden.size() / row_bytes_; + if (rows > static_cast(fresh_context_)) { + return false; + } + if (owners.size() == static_cast(max_owners_) || + rows > static_cast(fresh_context_ - context)) { + break; + } + context += static_cast(rows); + owners.push_back(input.sid); + anchors.push_back(out.outputs[cursor]->tokens.front()); + indices.push_back(cursor++); + } + if (owners.empty()) { + continue; + } + std::vector> proposals; + if (!draft(owners, anchors, fresh_context_, &proposals)) { + return false; + } + RawStep verification; + std::vector token_owners; + std::vector starts; + const int width = proposals_ + 1; + for (std::size_t i = 0; i < owners.size(); ++i) { + const auto& state = sessions_.at(owners[i]); + starts.push_back(state.frontier); + for (int j = 0; j < width; ++j) { + verification.tokens.push_back( + j == 0 ? anchors[i] : proposals[i][j - 1]); + verification.positions.push_back(state.frontier + j); + verification.seq_ids.push_back(state.target_id); + verification.logit_indices.push_back(verification.tokens.size() - 1); + token_owners.push_back(owners[i]); + } + } + std::vector predictions; + if (!run_target(verification, token_owners, predictions, false)) { + return false; + } + for (std::size_t i = 0; i < owners.size(); ++i) { + auto burst = df::verify_greedy( + proposals[i], + std::vector( + predictions.begin() + i * width, + predictions.begin() + (i + 1) * width)); + if (!burst || + !rewind( + sessions_.at(owners[i]), + starts[i] + static_cast(burst->tokens.size()))) { + return false; + } + // 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()); + ++stats_.cycles; + stats_.proposed += proposals_; + stats_.accepted += burst->accepted; + stats_.bonuses += burst->bonus; + stats_.rejected += !burst->bonus; + } + } + return true; +} + +bool MuseGlimmerCudaDFlashExecutor::fail_batch( + const batching::BatchInput& batch, + batching::BatchOutput& out) { + (void)batch; + if (!failed_) { + failed_ = true; + for (auto& item : sessions_) { + item.second.poisoned = true; + item.second.hidden.clear(); + } + // An abandoned declaration/forward affects the entire shared draft pool. + // Never allocate another draft SID or remove stale IDs after this clear. + draft_ctl_->clear(); + ET_LOG( + Error, + "CUDA DFlash: whole draft cache invalidated; model restart required"); + } + for (auto& output : out.outputs) { + output.reset(); + } + ++stats_.failed_batches; + return false; +} + +bool MuseGlimmerCudaDFlashExecutor::execute( + const batching::BatchInput& batch, + batching::BatchOutput& out) { +#if ET_HAS_EXCEPTIONS + try { +#endif + out.outputs.clear(); + out.outputs.resize(batch.inputs.size()); + if (!initialized_) { + return reject_batch(); + } + if (failed_) { + return fail_batch(batch, out); + } + RawStep step; + std::vector token_owners; + std::unordered_map cursor; + std::vector> rewinds; + for (const auto& input : batch.inputs) { + const auto it = sessions_.find(input.sid); + const auto* payload = + std::get_if(&input.payload); + if (it == sessions_.end() || it->second.poisoned || !payload || + !*payload || input.size == 0 || input.offset > (*payload)->size() || + input.size > (*payload)->size() - input.offset || + input.position < 0 || + input.offset > static_cast(max_context_) || + input.size > static_cast(max_context_) || + (input.produce_output && !it->second.sampling) || + (it->second.sampling && + (!std::isfinite(it->second.sampling->temperature) || + it->second.sampling->temperature != 0.0f))) { + ET_LOG( + Error, + "CUDA DFlash: invalid input or unsupported sampled configuration (temperature=0 required)"); + return reject_batch(); + } + const auto start = + static_cast(input.position) + input.offset; + const auto end = start + input.size; + auto [at, first] = cursor.emplace(input.sid, it->second.frontier); + if (target_->cache_control().pos(it->second.target_id) != + it->second.frontier) { + return fail_batch(batch, out); + } + if (start > at->second || (start < at->second && !first) || + end > max_context_ || + step.tokens.size() > + static_cast(std::numeric_limits::max()) - + input.size) { + return reject_batch(); + } + if (start < at->second) { + rewinds.emplace_back(input.sid, static_cast(start)); + } + at->second = static_cast(end); + for (std::size_t j = 0; j < input.size; ++j) { + const auto token = (**payload)[input.offset + j]; + if (token >= static_cast(target_->vocab_size())) { + return reject_batch(); + } + step.tokens.push_back(token); + step.positions.push_back(start + j); + step.seq_ids.push_back(it->second.target_id); + token_owners.push_back(input.sid); + } + step.logit_indices.push_back( + input.produce_output + ? static_cast(step.tokens.size()) - 1 + : -1); + } + // All inputs are validated before either cache is truncated. + for (const auto& item : rewinds) { + if (!rewind(sessions_.at(item.first), item.second)) { + return fail_batch(batch, out); + } + } + 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) { + if (batch.inputs[i].produce_output) { + out.outputs[i] = + batching::Output{batch.inputs[i].sid, {predictions[i]}}; + } + } + if (!speculate(batch, out)) { + return fail_batch(batch, out); + } + ++stats_.executions; + return true; +#if ET_HAS_EXCEPTIONS + } catch (...) { + ET_LOG(Error, "CUDA DFlash: batch failed; model restart required"); + return fail_batch(batch, out); + } +#endif +} +} // namespace executorch::extension::llm diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h new file mode 100644 index 00000000000..5f5f91fdf06 --- /dev/null +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h @@ -0,0 +1,172 @@ +/* + * 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. + */ +#pragma once + +#include +#include + +#include + +namespace executorch::extension::llm { + +// Exact packed-export metadata contract. Geometry tensors use the ordinary +// get_n_caches/get_kv_heads/get_head_dims/get_windows names with get_draft_. +namespace cuda_dflash_detail { +inline constexpr char kForward[] = "draft_forward"; +inline constexpr char kBlock[] = "get_block_size"; +inline constexpr char kMask[] = "get_mask_token_id"; +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"; + +struct GreedyBurst { + std::vector tokens; + std::size_t accepted = 0; + bool bonus = false; +}; +// Verification returns only new predictions, never the consumed anchor. +std::optional verify_greedy( + const std::vector& proposals, + const std::vector& predictions); +bool trim_hidden( + std::vector& hidden, + std::size_t row_bytes, + int draft_frontier, + int position); + +struct DraftContext { + std::int32_t id; + int position; + batching::Token anchor; + const std::vector* hidden; + int count; +}; +struct DraftPlan { + std::vector tokens, positions, kv_order; + std::vector ids; + std::vector hidden; + std::vector counts, ends; +}; +std::optional pack_draft( + const std::vector& owners, + int block, + batching::Token mask, + int context_limit, + int max_context, + std::size_t row_bytes); +} // namespace cuda_dflash_detail + +struct CudaDFlashStats { + std::uint64_t executions = 0; + std::uint64_t draft_forwards = 0; + std::uint64_t warmed_context_tokens = 0; + std::uint64_t cycles = 0; // per-owner verification cycles + std::uint64_t proposed = 0; + std::uint64_t accepted = 0; + std::uint64_t rejected = 0; // cycles ended by a rejection, not tokens + std::uint64_t bonuses = 0; + std::uint64_t failed_batches = 0; +}; + +// Text-only model adapter over the existing CUDA target executor. Draft runs +// eagerly; target methods retain their ordinary static CUDA graphs. Sampling +// other than exactly temperature zero is rejected, not approximated. +class MuseGlimmerCudaDFlashExecutor final : public batching::Executor { + public: + static runtime::Result create( + const MuseGlimmerBackendConfig& config); + MuseGlimmerCudaDFlashExecutor(const MuseGlimmerCudaDFlashExecutor&) = delete; + MuseGlimmerCudaDFlashExecutor& operator=( + const MuseGlimmerCudaDFlashExecutor&) = delete; + MuseGlimmerCudaDFlashExecutor(MuseGlimmerCudaDFlashExecutor&&) = delete; + MuseGlimmerCudaDFlashExecutor& operator=(MuseGlimmerCudaDFlashExecutor&&) = + delete; + ~MuseGlimmerCudaDFlashExecutor() override; + std::size_t preferred_batch_tokens() const override; + bool initialize() override; + std::optional open_session() override; + void close_session(batching::SessionId session) override; + std::optional clone( + batching::SessionId source, + batching::Position upto) override; + void set_sampling( + batching::SessionId session, + const batching::SamplingParams& params, + std::optional seed) override; + bool execute(const batching::BatchInput& batch, batching::BatchOutput& out) + override; + // Engine-thread-only, between executes. Counters measure actual adapter work. + CudaDFlashStats stats() const { + return stats_; + } + + private: + using Target = ::executorch::backends::cuda::batching::CudaExecutor; + using RawStep = ::executorch::backends::cuda::batching::RawTargetStep; + struct State { + State(std::int32_t target, std::int32_t draft) + : target_id(target), draft_id(draft) {} + 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. + std::vector hidden; + std::optional sampling; + std::optional seed; + bool poisoned = false; + }; + MuseGlimmerCudaDFlashExecutor( + std::unique_ptr target, + std::shared_ptr draft_cache, + int block, + int proposals, + batching::Token mask, + int max_owners, + int fresh_context, + int max_sessions, + int max_context, + int hidden_width); + bool rewind(State& state, int position); + bool run_target( + const RawStep& step, + const std::vector& token_owners, + std::vector& predictions, + bool drain_context); + bool drain(batching::SessionId owner); + bool draft( + const std::vector& owners, + const std::vector& anchors, + int context_limit, + std::vector>* proposals); + bool speculate(const batching::BatchInput& batch, batching::BatchOutput& out); + bool fail_batch( + const batching::BatchInput& batch, + batching::BatchOutput& out); + + // Both guards exist before either cache delegate loads. Destruction releases + // the target-owned Module/delegates before removing the draft registry entry. + std::shared_ptr draft_cache_; + cache::InstallGuard draft_install_; + std::unique_ptr target_; + cache::BatchControl* const draft_ctl_; + int block_, proposals_, max_owners_, fresh_context_, max_sessions_, + max_context_, hidden_width_; + batching::Token mask_; + std::size_t row_bytes_; + batching::SessionId next_session_ = 1; + std::unordered_map sessions_; + CudaDFlashStats stats_; + bool initialized_ = false; + // A failed co-batch invalidates every draft owner. Only model restart + // recovers. + bool failed_ = false; +}; + +} // namespace executorch::extension::llm diff --git a/examples/models/muse-glimmer/serving/serve_batching.py b/examples/models/muse-glimmer/serving/serve_batching.py index 9644c3bff73..eb02b080088 100644 --- a/examples/models/muse-glimmer/serving/serve_batching.py +++ b/examples/models/muse-glimmer/serving/serve_batching.py @@ -118,8 +118,6 @@ def _validate_backend(parser, args): parser.error("--backend cuda requires --data-path (aoti_cuda_blob.ptd)") if args.pos_embed_path is not None: parser.error("--backend cuda serves text only; omit --pos-embed-path") - if args.dflash: - parser.error("--dflash is not supported with --backend cuda") def _validate_paths(parser, args): diff --git a/examples/models/muse-glimmer/tests/BUCK b/examples/models/muse-glimmer/tests/BUCK index 809b8778b10..e07ba69892c 100644 --- a/examples/models/muse-glimmer/tests/BUCK +++ b/examples/models/muse-glimmer/tests/BUCK @@ -1,8 +1,17 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +load("@fbcode_macros//build_defs:cpp_unittest.bzl", "cpp_unittest") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") oncall("executorch") +fbcode_target( + _kind = cpp_unittest, + name = "cuda_dflash_executor_test", + srcs = ["cuda_dflash_executor_test.cpp"], + deps = ["//executorch/examples/models/muse-glimmer/runtime:cuda_dflash_executor"], + external_deps = [("cuda", None, "cuda-lazy"), "nlohmann_json"], +) + fbcode_target( _kind = python_unittest, name = "dflash_speculative_verifier_test", diff --git a/examples/models/muse-glimmer/tests/cuda_dflash_executor_test.cpp b/examples/models/muse-glimmer/tests/cuda_dflash_executor_test.cpp new file mode 100644 index 00000000000..5e8f28f1823 --- /dev/null +++ b/examples/models/muse-glimmer/tests/cuda_dflash_executor_test.cpp @@ -0,0 +1,138 @@ +/* + * 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 + +namespace df = executorch::extension::llm::cuda_dflash_detail; +using executorch::extension::llm::batching::Token; + +TEST(CudaDFlash, FirstMismatchReturnsOnlyCorrection) { + const auto burst = df::verify_greedy({11, 12, 13}, {21, 12, 13, 14}); + ASSERT_TRUE(burst); + EXPECT_EQ(burst->tokens, std::vector({21})); + EXPECT_EQ(burst->accepted, 0); + EXPECT_FALSE(burst->bonus); +} + +TEST(CudaDFlash, PartialAcceptanceStopsAtFirstMismatch) { + const auto burst = df::verify_greedy({11, 12, 13}, {11, 22, 13, 14}); + ASSERT_TRUE(burst); + EXPECT_EQ(burst->tokens, std::vector({11, 22})); + EXPECT_EQ(burst->accepted, 1); + EXPECT_FALSE(burst->bonus); +} + +TEST(CudaDFlash, FullAcceptanceAddsOneBonus) { + const auto burst = df::verify_greedy({11, 12, 13}, {11, 12, 13, 14}); + ASSERT_TRUE(burst); + EXPECT_EQ(burst->tokens, std::vector({11, 12, 13, 14})); + EXPECT_EQ(burst->accepted, 3); + EXPECT_TRUE(burst->bonus); +} + +TEST(CudaDFlash, InvalidVerificationShapeRejected) { + EXPECT_FALSE(df::verify_greedy({}, {1})); + EXPECT_FALSE(df::verify_greedy({1}, {})); + EXPECT_FALSE(df::verify_greedy({1, 2}, {1, 2})); + EXPECT_FALSE(df::verify_greedy({1}, {1, 2, 3})); +} + +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. + for (const auto& predictions : std::vector>{ + {21, 12, 13, 14}, {11, 22, 13, 14}, {11, 12, 13, 14}}) { + auto burst = df::verify_greedy({11, 12, 13}, predictions); + ASSERT_TRUE(burst); + std::vector output{10}; + output.insert(output.end(), burst->tokens.begin(), burst->tokens.end()); + EXPECT_EQ(output.front(), 10); + EXPECT_EQ(std::count(output.begin(), output.end(), 10), 1); + EXPECT_EQ(std::count(output.begin(), output.end(), 7), 0); + // Verification commits the anchor plus accepted proposals, never the + // correction/bonus. This matches the runner's end + out.size() - 1 rule. + EXPECT_EQ(output.size() - 1, burst->accepted + 1); + } +} + +TEST(CudaDFlash, BudgetRewindTrimsOnlyCommittedRows) { + const auto burst = df::verify_greedy({11, 12, 13}, {11, 12, 13, 14}); + ASSERT_TRUE(burst); + // Draft has warmed through input.end=100; target verified four rows. + for (int kept = 1; kept <= 4; ++kept) { + std::vector hidden{10, 11, 12, 13, 14, 15, 16, 17}; + const int frontier = 100 + kept - 1; + ASSERT_TRUE(df::trim_hidden(hidden, 2, 100, frontier)); + EXPECT_EQ(hidden.size(), static_cast(kept - 1) * 2); + if (kept > 1) { + EXPECT_EQ(hidden.front(), 10); + EXPECT_EQ(hidden.back(), 9 + (kept - 1) * 2); + } + } +} + +TEST(CudaDFlash, RewindBeforeDraftFrontierDropsFreshRows) { + std::vector hidden{1, 2, 3, 4}; + EXPECT_TRUE(df::trim_hidden(hidden, 2, 10, 8)); + EXPECT_TRUE(hidden.empty()); +} + +TEST(CudaDFlash, InvalidHiddenExtentRejected) { + std::vector hidden{1, 2, 3}; + EXPECT_FALSE(df::trim_hidden(hidden, 2, 0, 1)); + EXPECT_FALSE(df::trim_hidden(hidden, 0, 0, 1)); + EXPECT_FALSE(df::trim_hidden(hidden, 1, -1, 1)); + EXPECT_FALSE(df::trim_hidden(hidden, 1, 0, 4)); + EXPECT_EQ(hidden, std::vector({1, 2, 3})); +} + +TEST(CudaDFlash, PackingUsesExactKAndOriginalOwnerQueryOrder) { + std::vector a{1, 2, 3, 4}, b{5, 6}; + const auto plan = + df::pack_draft({{9, 10, 70, &a, 2}, {3, 20, 80, &b, 1}}, 4, 99, 3, 64, 2); + ASSERT_TRUE(plan); + EXPECT_EQ( + plan->tokens, + std::vector({70, 99, 99, 99, 80, 99, 99, 99})); + EXPECT_EQ( + plan->positions, + std::vector({10, 11, 20, 12, 13, 14, 15, 21, 22, 23, 24})); + EXPECT_EQ( + plan->kv_order, + std::vector({0, 1, 3, 4, 5, 6, 2, 7, 8, 9, 10})); + EXPECT_EQ( + plan->ids, std::vector({9, 9, 9, 9, 9, 9, 3, 3, 3, 3, 3})); + EXPECT_EQ(plan->hidden, std::vector({1, 2, 3, 4, 5, 6})); + EXPECT_EQ(plan->ends, std::vector({12, 21})); + EXPECT_EQ(plan->ids.size(), plan->tokens.size() + plan->hidden.size() / 2); +} + +TEST(CudaDFlash, ZeroFreshContextHasNoQueryPadding) { + std::vector empty; + const auto plan = df::pack_draft({{4, 10, 70, &empty, 0}}, 4, 99, 0, 64, 2); + ASSERT_TRUE(plan); + EXPECT_TRUE(plan->hidden.empty()); + EXPECT_EQ(plan->ids.size(), 4); + EXPECT_EQ(plan->positions, std::vector({10, 11, 12, 13})); + EXPECT_EQ(plan->kv_order, std::vector({0, 1, 2, 3})); +} + +TEST(CudaDFlash, NearContextLimitRequiresWholeTrainedBlock) { + std::vector hidden{1, 2}; + EXPECT_TRUE(df::pack_draft({{4, 59, 70, &hidden, 1}}, 4, 99, 1, 64, 2)); + EXPECT_FALSE(df::pack_draft({{4, 60, 70, &hidden, 1}}, 4, 99, 1, 64, 2)); +} + +TEST(CudaDFlash, PackingRejectsDuplicateOwnersAndOversizedContext) { + std::vector hidden{1, 2}; + EXPECT_FALSE(df::pack_draft( + {{4, 0, 70, &hidden, 1}, {4, 0, 80, &hidden, 1}}, 4, 99, 2, 64, 2)); + EXPECT_FALSE(df::pack_draft({{4, 0, 70, &hidden, 2}}, 4, 99, 2, 64, 2)); + EXPECT_FALSE(df::pack_draft({{4, 0, 70, &hidden, 1}}, 4, 99, 0, 64, 2)); +} diff --git a/examples/models/muse-glimmer/tests/test_serve_batching.py b/examples/models/muse-glimmer/tests/test_serve_batching.py index 5fc51108416..941d2ec430a 100644 --- a/examples/models/muse-glimmer/tests/test_serve_batching.py +++ b/examples/models/muse-glimmer/tests/test_serve_batching.py @@ -123,10 +123,24 @@ def test_cuda_worker_takes_the_weights_file(arguments, tmp_path): assert "--pos_embed_path" not in flags -@pytest.mark.parametrize("case", ["no_weights", "pos_embed", "dflash"]) -def test_cuda_launcher_rejects_missing_weights_vision_or_dflash( - arguments, tmp_path, case -): +def test_cuda_launcher_forwards_dflash(arguments, tmp_path): + data = tmp_path / "aoti_cuda_blob.ptd" + data.touch() + args = serve_batching._parse_args( + arguments + + ["--backend", "cuda", "--data-path", str(data), "--dflash"] + + ["--dflash-n-draft", "15"] + ) + flags = dict(arg.split("=", 1) for arg in serve_batching._worker_command(args)[1:]) + assert (flags["--backend"], flags["--dflash"], flags["--dflash_n_draft"]) == ( + "cuda", + "true", + "15", + ) + + +@pytest.mark.parametrize("case", ["no_weights", "pos_embed"]) +def test_cuda_launcher_rejects_missing_weights_or_vision(arguments, tmp_path, case): data = tmp_path / "aoti_cuda_blob.ptd" table = tmp_path / "pos_embed.bin" data.touch() @@ -136,8 +150,6 @@ def test_cuda_launcher_rejects_missing_weights_vision_or_dflash( argv += ["--data-path", str(data)] if case == "pos_embed": argv += ["--pos-embed-path", str(table)] - if case == "dflash": - argv += ["--dflash"] with pytest.raises(SystemExit) as error: serve_batching._parse_args(argv) assert error.value.code == 2