From ab2e345a26b50be9b50440084535533d5a0262e7 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Tue, 20 Oct 2026 15:52:47 -0700 Subject: [PATCH] [executorch][muse] Verify CUDA DFlash drafts on device under any sampling Pull Request resolved: https://github.com/pytorch/executorch/pull/23984 Export dflash_verify / dflash_verify_argmax: batched on-device acceptance of every owner's drafts under that owner's temperature / top-k / top-p, with a greedy fast path. Only token IDs return to the host. Uniforms come from (seed, position, purpose) so results do not depend on batch composition. Target preamble sampling also stays on device (CudaExecutor::sample_rows). ghstack-source-id: 453099298 @exported-using-ghexport Differential Revision: [D125705716](https://our.internmc.facebook.com/intern/diff/D125705716/) --- backends/cuda/batching/cuda_executor.cpp | 70 +++- backends/cuda/batching/cuda_executor.h | 13 + .../muse-glimmer/export/export_dflash.py | 47 ++- .../models/muse-glimmer/export/export_solo.py | 30 +- .../model/dflash_token_sampler.py | 101 ++++- .../runtime/cuda_dflash_executor.cpp | 364 +++++++++++++----- .../runtime/cuda_dflash_executor.h | 61 ++- .../tests/cuda_dflash_executor_test.cpp | 54 ++- .../tests/cuda_dflash_native_test.cpp | 88 ++++- .../tests/test_cuda_dflash_packed.py | 16 + .../tests/test_dflash_token_sampler.py | 166 +++++++- 11 files changed, 858 insertions(+), 152 deletions(-) diff --git a/backends/cuda/batching/cuda_executor.cpp b/backends/cuda/batching/cuda_executor.cpp index a47197f82aa..40aad0317fc 100644 --- a/backends/cuda/batching/cuda_executor.cpp +++ b/backends/cuda/batching/cuda_executor.cpp @@ -810,21 +810,40 @@ bool CudaExecutor::execute(const BatchInput& batch, BatchOutput& out) { }); } +std::optional> CudaExecutor::sample_rows( + const ::executorch::aten::Tensor& logits, + std::size_t rows, + const std::vector& params) { + if (params.size() != rows * 4) { + return std::nullopt; + } + bool all_greedy = true; + for (std::size_t row = 0; row < rows; ++row) { + all_greedy = all_greedy && params[row * 4] == 0.0f; + } + if (all_greedy) { + return argmax_rows(logits, rows); + } + if (logits.device_type() != ::executorch::aten::DeviceType::CUDA) { + ET_LOG( + Error, + "CudaExecutor: sampled rows need device logits and %s", + kSampleMethod); + return std::nullopt; + } + if (!initialized_ || !has_sampler_ || logits.dim() < 2 || !is_dense(logits) || + logits.size(logits.dim() - 1) != vocab_size_ || + logits.scalar_type() != ::executorch::aten::ScalarType::Float || + rows == 0 || + rows > static_cast(logits.numel() / vocab_size_)) { + return std::nullopt; + } + return run_device_sampler(logits, rows, params, false); +} + std::optional> CudaExecutor::sample_on_device( const ::executorch::aten::Tensor& logits, const std::vector& sessions) { - using ::executorch::aten::ScalarType; - using ::executorch::aten::SizesType; - const ::executorch::aten::Device cuda( - ::executorch::aten::DeviceType::CUDA, logits.device_index()); - const auto rows = static_cast(sessions.size()); - // The rows that sample lead the logits; the rest are padding. - auto row_logits = make_tensor_ptr( - {rows, static_cast(vocab_size_)}, - logits.mutable_data_ptr(), - logits.scalar_type(), - cuda); - const bool all_greedy = std::all_of(sessions.begin(), sessions.end(), [this](SessionId session) { return sessions_.greedy(session); @@ -840,6 +859,25 @@ std::optional> CudaExecutor::sample_on_device( params.insert( params.end(), {row->temperature, row->top_p, row->top_k, row->coin}); } + return run_device_sampler(logits, sessions.size(), params, all_greedy); +} + +std::optional> CudaExecutor::run_device_sampler( + const ::executorch::aten::Tensor& logits, + std::size_t row_count, + const std::vector& params, + bool all_greedy) { + using ::executorch::aten::ScalarType; + using ::executorch::aten::SizesType; + const ::executorch::aten::Device cuda( + ::executorch::aten::DeviceType::CUDA, logits.device_index()); + const auto rows = static_cast(row_count); + // The rows that sample lead the logits; the rest are padding. + auto row_logits = make_tensor_ptr( + {rows, static_cast(vocab_size_)}, + logits.mutable_data_ptr(), + logits.scalar_type(), + cuda); const cudaStream_t stream = ::executorch::extension::cuda::getCallerStream().value_or( @@ -851,18 +889,18 @@ std::optional> CudaExecutor::sample_on_device( kArgmaxMethod, std::vector<::executorch::runtime::EValue>{row_logits}); } - if (sessions.size() > device_params_rows_) { + if (row_count > device_params_rows_) { if (device_params_ != nullptr) { (void)cudaFree(device_params_); device_params_ = nullptr; device_params_rows_ = 0; } ET_CHECK_OR_RETURN_ERROR( - cudaMalloc(&device_params_, sessions.size() * 4 * sizeof(float)) == + cudaMalloc(&device_params_, row_count * 4 * sizeof(float)) == cudaSuccess, MemoryAllocationFailed, "CudaExecutor: could not stage the sampler's params"); - device_params_rows_ = sessions.size(); + device_params_rows_ = row_count; } // The backend runs every method on the caller's stream if one is set, // else the per-thread stream; ordering the upload there puts it ahead of @@ -890,7 +928,7 @@ std::optional> CudaExecutor::sample_on_device( ET_LOG(Error, "CudaExecutor: the device sampler returned a bad shape"); return std::nullopt; } - std::vector host(sessions.size()); + std::vector host(row_count); const auto kind = tokens.device_type() == ::executorch::aten::DeviceType::CUDA ? cudaMemcpyDeviceToHost : cudaMemcpyHostToHost; diff --git a/backends/cuda/batching/cuda_executor.h b/backends/cuda/batching/cuda_executor.h index 8acc77c82e2..fc0094ae9ef 100644 --- a/backends/cuda/batching/cuda_executor.h +++ b/backends/cuda/batching/cuda_executor.h @@ -174,6 +174,14 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor { std::optional> argmax_rows( const ::executorch::aten::Tensor& logits, std::size_t rows); + // Samples device logits with the exported sampler. params is row-major + // [rows, 4] float32: temperature, top_p, top_k and a uniform coin, as + // extension/llm/batching/sampler.py consumes them. All-greedy rows take + // sample_argmax. + std::optional> sample_rows( + const ::executorch::aten::Tensor& logits, + std::size_t rows, + const std::vector& params); ::executorch::extension::Module& module() { return *module_; } @@ -212,6 +220,11 @@ class ET_EXPERIMENTAL CudaExecutor : public llm_batching::Executor { // Draws one token per row of `logits`, rows [0, sessions.size()), on the // device; returns them in row order. + std::optional> run_device_sampler( + const ::executorch::aten::Tensor& logits, + std::size_t rows, + const std::vector& params, + bool all_greedy); std::optional> sample_on_device( const ::executorch::aten::Tensor& logits, const std::vector& sessions); diff --git a/examples/models/muse-glimmer/export/export_dflash.py b/examples/models/muse-glimmer/export/export_dflash.py index b4478bfccb1..1d6e183cbfd 100644 --- a/examples/models/muse-glimmer/export/export_dflash.py +++ b/examples/models/muse-glimmer/export/export_dflash.py @@ -95,6 +95,50 @@ def _export_cuda_sampler_methods( } +VERIFY_METHOD = "dflash_verify" +VERIFY_ARGMAX_METHOD = "dflash_verify_argmax" + + +def cuda_owner_verifier_programs( + vocab_size: int, max_owners: int, block: int +) -> dict[str, torch.export.ExportedProgram]: + """Batched on-device verification of up to ``max_owners`` draft blocks.""" + from executorch.examples.models.muse_glimmer.model.dflash_token_sampler import ( + DFlashOwnerArgmaxVerifier, + DFlashOwnerVerifier, + ) + from torch.export import Dim, export + + if block < 3: + raise ValueError("CUDA DFlash verifiers need a block of at least 3") + owners = ( + Dim("verify_owners", min=1, max=max_owners) if max_owners > 1 else Dim.STATIC + ) + # A one-proposal width would specialize its size-1 proposal axis. + width = Dim("verify_width", min=3, max=block) + logits = torch.zeros(max_owners, block, vocab_size) + candidates = torch.zeros(max_owners, block, dtype=torch.long) + params = torch.zeros(max_owners, block + 3) + return { + VERIFY_METHOD: export( + DFlashOwnerVerifier(), + (logits, candidates, params), + dynamic_shapes=( + {0: owners, 1: width}, + {0: owners, 1: width}, + {0: owners, 1: width + 3}, + ), + strict=True, + ), + VERIFY_ARGMAX_METHOD: export( + DFlashOwnerArgmaxVerifier(), + (logits, candidates), + dynamic_shapes=({0: owners, 1: width}, {0: owners, 1: width}), + strict=True, + ), + } + + def _cuda_propagate_device_config() -> dict[str, PropagateDeviceConfig]: """Keeps the speculative chain's vocabulary-wide tensors on the device.""" from executorch.exir.passes.propagate_device_pass import PropagateDeviceConfig @@ -337,6 +381,7 @@ def prepare_cuda_packed_draft( ), strict=True, ) + programs = {"draft_forward": program} geometry = offgraph_kv_cache_geometry(draft_model) metadata = { key.replace("get_", "get_draft_", 1): value for key, value in geometry.items() @@ -357,7 +402,7 @@ def prepare_cuda_packed_draft( "get_dflash_draft_cache_key": "muse_glimmer_draft", } ) - return {"draft_forward": program}, manifest, metadata + return programs, manifest, metadata def _dflash_constant_methods( diff --git a/examples/models/muse-glimmer/export/export_solo.py b/examples/models/muse-glimmer/export/export_solo.py index 7a890ba2b6d..9fe39ce6df5 100644 --- a/examples/models/muse-glimmer/export/export_solo.py +++ b/examples/models/muse-glimmer/export/export_solo.py @@ -318,7 +318,7 @@ def sampler_programs(vocab_size: int, max_rows: int = MAX_SAMPLE_ROWS): } -def device_residency(forward_methods, device_sampling: bool): +def device_residency(forward_methods, device_sampling: bool, samplers=()): """Per-method ``PropagateDeviceConfig`` for ``to_executorch``. With device sampling the forwards keep their logits in device memory, and @@ -331,7 +331,7 @@ def device_residency(forward_methods, device_sampling: bool): for name in forward_methods } if device_sampling: - for name in (SAMPLE_METHOD, ARGMAX_METHOD): + for name in (SAMPLE_METHOD, ARGMAX_METHOD, *samplers): device[name] = PropagateDeviceConfig(skip_h2d_for_method_inputs=True) return device @@ -469,19 +469,31 @@ def export_cuda_batching( # noqa: C901 if dflash_draft is not None: programs.update(draft_programs) del draft_programs, dflash_draft - cache_methods = forward_methods + ( - ["draft_forward"] if "draft_forward" in programs else [] - ) + draft_methods = [name for name in programs if name.startswith("draft_forward")] + cache_methods = forward_methods + draft_methods + verifiers = {} if device_sampling: print(f"Exporting {SAMPLE_METHOD} and {ARGMAX_METHOD}...") programs.update(sampler_programs(config.vocab_size)) + if "draft_forward" in programs: + from executorch.examples.models.muse_glimmer.export.export_dflash import ( + cuda_owner_verifier_programs, + ) + + verifiers = cuda_owner_verifier_programs( + config.vocab_size, + max_draft_owners, + draft_metadata["get_block_size"], + ) + print(f"Exporting {', '.join(verifiers)}...") + programs.update(verifiers) def partitioner(name: str) -> CudaPartitioner: if name not in cache_methods: return CudaPartitioner( [CudaBackend.generate_method_name_compile_spec(name)] ) - is_draft = name == "draft_forward" + is_draft = name in draft_methods specs = [ CudaBackend.generate_method_name_compile_spec(name), CompileSpec("low_memory_mode", b"ON"), @@ -494,7 +506,7 @@ def partitioner(name: str) -> CudaPartitioner: b"2:0" if is_draft else step_width_spec(name), ), ] - if "draft_forward" in cache_methods: + if draft_methods: specs.append( CompileSpec( "llm_cache_registry_key", @@ -539,7 +551,9 @@ def partitioner(name: str) -> CudaPartitioner: extract_delegate_segments=True, do_quant_fusion_and_const_prop=True, memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), - propagate_device_config=device_residency(cache_methods, device_sampling), + propagate_device_config=device_residency( + cache_methods, device_sampling, tuple(verifiers) + ), emit_mutable_buffer_names=True, ), ) diff --git a/examples/models/muse-glimmer/model/dflash_token_sampler.py b/examples/models/muse-glimmer/model/dflash_token_sampler.py index b7da5db3aee..3354a1e8a83 100644 --- a/examples/models/muse-glimmer/model/dflash_token_sampler.py +++ b/examples/models/muse-glimmer/model/dflash_token_sampler.py @@ -25,6 +25,8 @@ def sampling_probabilities( temperature: torch.Tensor, top_k: torch.Tensor, top_p: torch.Tensor, + *, + batching_policy: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """Returns ``[R, V]`` filtered probabilities and ``[R]`` argmax tokens.""" vocab_size = logits.shape[-1] @@ -52,8 +54,14 @@ def sampling_probabilities( torch.cumsum(sorted_probabilities, dim=-1) - sorted_probabilities ) top_p_enabled = (top_p > 0.0) & (top_p < 1.0) + if batching_policy: + # Match BatchSampler's top-k precedence and inclusive nucleus cutoff. + top_p_enabled = top_p_enabled & ~((top_k > 0) & (top_k < vocab_size)) + nucleus = mass_before_token <= top_p + else: + nucleus = mass_before_token < top_p sorted_probabilities = torch.where( - ~top_p_enabled | (mass_before_token < top_p), + ~top_p_enabled | nucleus, sorted_probabilities, torch.zeros_like(sorted_probabilities), ) @@ -195,6 +203,97 @@ def verify_speculative( return torch.cat([decision_row + 1, correction, candidates]) +def verify_owners( + target_logits: torch.Tensor, + candidates: torch.Tensor, + params: torch.Tensor, +) -> torch.Tensor: + """Verifies each owner's greedy draft block against its target rows. + + ``target_logits`` is ``[O, W, V]`` and ``candidates`` ``[O, W]``: the anchor + then ``W - 1`` proposals. ``params`` is float32 ``[O, W + 3]``: temperature, + top_k, top_p, ``W - 1`` acceptance uniforms, then one correction uniform. + Greedy proposals have q = 1, so a proposal is accepted with its target + probability and a rejection resamples the target without that token. + Returns int64 ``[O, 2]``: accepted proposals and the correction or bonus. + """ + owners, width, vocab_size = target_logits.shape + rows = owners * width + + def per_row(column: int) -> torch.Tensor: + return params[:, column : column + 1].expand(owners, width).reshape(rows, 1) + + probabilities, greedy = sampling_probabilities( + target_logits.reshape(rows, vocab_size), + per_row(0), + per_row(1), + per_row(2), + batching_policy=True, + ) + probabilities = probabilities.view(owners, width, vocab_size) + greedy = greedy.view(owners, width) + proposals = candidates[:, 1:] + stochastic = params[:, 0] > 0.0 + proposal_probability = ( + probabilities[:, :-1].gather(-1, proposals.unsqueeze(-1)).squeeze(-1) + ) + accepted = torch.where( + stochastic.unsqueeze(-1), + params[:, 3:-1] < proposal_probability, + greedy[:, :-1] == proposals, + ) + count = torch.cumprod(accepted.to(torch.int64), dim=1).sum(dim=1) + decision = count.view(owners, 1, 1).expand(owners, 1, vocab_size) + row = probabilities.gather(1, decision).squeeze(1) + excluded_token = proposals.gather(1, count.clamp(max=width - 2).unsqueeze(1)) + excluded = row.scatter(1, excluded_token, 0.0) + excluded_sum = excluded.sum(dim=-1, keepdim=True) + excluded = excluded / torch.where( + excluded_sum > 0.0, excluded_sum, torch.ones_like(excluded_sum) + ) + coin = params[:, -1] + sampled = torch.where( + count < width - 1, + strict_categorical_sample(excluded, coin), + strict_categorical_sample(row, coin), + ) + correction = torch.where( + stochastic, sampled, greedy.gather(1, count.unsqueeze(1)).squeeze(1) + ) + return torch.stack([count, correction], dim=1) + + +def verify_owners_greedy( + target_logits: torch.Tensor, candidates: torch.Tensor +) -> torch.Tensor: + """``verify_owners`` at temperature 0, without sorting the vocabulary.""" + greedy = torch.argmax(target_logits.float(), dim=-1) + accepted = greedy[:, :-1] == candidates[:, 1:] + count = torch.cumprod(accepted.to(torch.int64), dim=1).sum(dim=1) + return torch.stack([count, greedy.gather(1, count.unsqueeze(1)).squeeze(1)], 1) + + +class DFlashOwnerVerifier(nn.Module): + """Batched draft verification; see :func:`verify_owners`.""" + + def forward( + self, + target_logits: torch.Tensor, + candidates: torch.Tensor, + params: torch.Tensor, + ) -> torch.Tensor: + return verify_owners(target_logits, candidates, params) + + +class DFlashOwnerArgmaxVerifier(nn.Module): + """Greedy batched draft verification; see :func:`verify_owners_greedy`.""" + + def forward( + self, target_logits: torch.Tensor, candidates: torch.Tensor + ) -> torch.Tensor: + return verify_owners_greedy(target_logits, candidates) + + class DFlashTokenSampler(nn.Module): """Samples ``[R, V]`` logits; returns ``[R]`` tokens and ``[R, V]`` probs.""" diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp index f6888d201a4..f4a1dd2a052 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp @@ -17,8 +17,10 @@ #include #include #include +#include #include #include +#include #include namespace executorch::extension::llm { @@ -30,26 +32,46 @@ namespace df = cuda_dflash_detail; namespace cb = ::executorch::backends::cuda::batching; namespace cuda_dflash_detail { -std::optional verify_greedy( +std::optional make_burst( const std::vector& proposals, - const std::vector& predictions) { - if (proposals.empty() || predictions.size() != proposals.size() + 1) { + std::int64_t accepted, + std::int64_t token, + std::int32_t vocab_size) { + if (proposals.empty() || accepted < 0 || + accepted > static_cast(proposals.size()) || token < 0 || + token >= vocab_size) { 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; - } - } + Burst result; + result.tokens.assign(proposals.begin(), proposals.begin() + accepted); + result.tokens.push_back(static_cast(token)); + result.accepted = static_cast(accepted); + result.bonus = result.accepted == proposals.size(); return result; } +float positional_uniform( + std::uint64_t seed, + batching::Position position, + int purpose) { + // SplitMix64 over (seed, purpose, position), as the MLX batching executor + // derives its per-prediction generators. + auto mix = [](std::uint64_t value) { + value += 0x9e3779b97f4a7c15ULL; + value = (value ^ (value >> 30)) * 0xbf58476d1ce4e5b9ULL; + value = (value ^ (value >> 27)) * 0x94d049bb133111ebULL; + return value ^ (value >> 31); + }; + const std::uint64_t value = + mix(seed ^ mix(static_cast(purpose)) ^ + mix(static_cast(position) + 1)); + std::seed_seq words{ + static_cast(value), + static_cast(value >> 32)}; + std::mt19937 generator(words); + return std::uniform_real_distribution(0.0f, 1.0f)(generator); +} + bool trim_hidden( std::vector& hidden, std::size_t row_bytes, @@ -120,6 +142,49 @@ std::optional pack_draft( } return plan; } + +DeviceStaging::~DeviceStaging() { + if (device_) { + (void)cudaFree(device_); + } +} + +std::optional> DeviceStaging::upload( + const std::vector>& parts, + void* stream) { + std::vector offsets; + std::size_t total = 0; + for (const auto& part : parts) { + total = (total + 7) & ~std::size_t{7}; + offsets.push_back(total); + total += part.second; + } + if (total > capacity_) { + if (device_ && cudaFree(device_) != cudaSuccess) { + return std::nullopt; + } + device_ = nullptr; + capacity_ = 0; + if (cudaMalloc(&device_, total) != cudaSuccess) { + return std::nullopt; + } + capacity_ = total; + } + std::vector pointers; + for (std::size_t i = 0; i < parts.size(); ++i) { + auto* destination = static_cast(device_) + offsets[i]; + if (cudaMemcpyAsync( + destination, + parts[i].first, + parts[i].second, + cudaMemcpyHostToDevice, + static_cast(stream)) != cudaSuccess) { + return std::nullopt; + } + pointers.push_back(destination); + } + return pointers; +} } // namespace cuda_dflash_detail namespace { @@ -382,11 +447,12 @@ Result MuseGlimmerCudaDFlashExecutor::create( auto& module = target->module(); ET_ASSIGN_OR_RETURN(names, module.method_names()); ET_CHECK_OR_RETURN_ERROR( - names.count(df::kForward) && + names.count(df::kForward) && names.count(df::kVerify) && + names.count(df::kVerifyArgmax) && !names.count("get_mutable_buffer_metadata") && !names.count("vision_encoder"), InvalidProgram, - "CUDA DFlash requires a tapped text-only off-graph artifact"); + "CUDA DFlash requires a tapped text-only off-graph artifact with its verifiers"); ET_ASSIGN_OR_RETURN(block, positive_int(module, df::kBlock)); ET_ASSIGN_OR_RETURN( declared_block, positive_int(module, "get_draft_dflash_block_size")); @@ -403,7 +469,7 @@ Result MuseGlimmerCudaDFlashExecutor::create( 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 >= 2 && block == declared_block && max_owners <= df::kOwnerLimit && block <= (std::numeric_limits::max() - 1) / max_owners && max_write > max_owners * block && max_cells > max_write && mask && *mask >= 0 && *mask < target->vocab_size() && @@ -417,10 +483,11 @@ Result MuseGlimmerCudaDFlashExecutor::create( "Invalid packed CUDA DFlash metadata"); const int proposals = config.dflash_n_draft == 0 ? block - 1 : config.dflash_n_draft; + // The exported verifiers take two to block - 1 proposals. ET_CHECK_OR_RETURN_ERROR( - proposals > 0 && proposals < block, + proposals >= 2 && proposals < block, InvalidArgument, - "DFlash n_draft must be in [1, trained block size - 1]"); + "DFlash n_draft must be in [2, 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)); @@ -548,12 +615,15 @@ bool MuseGlimmerCudaDFlashExecutor::initialize() { if (initialized_) { return true; } + auto& module = target_->module(); if (!target_->initialize() || batching::load_method_with_cache( - target_->module(), + module, df::kForward, ::executorch::backends::cuda::kCudaBackendId, - draft_install_) != Error::Ok) { + draft_install_) != Error::Ok || + module.load_method(df::kVerify) != Error::Ok || + module.load_method(df::kVerifyArgmax) != Error::Ok) { ET_LOG( Error, "CUDA DFlash: could not initialize distinct target/draft caches"); @@ -659,15 +729,31 @@ void MuseGlimmerCudaDFlashExecutor::set_sampling( 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"); + if (it == sessions_.end()) { + return; + } + auto& state = it->second; + state.sampling.reset(); + if (!std::isfinite(params.temperature) || params.temperature < 0 || + !std::isfinite(params.top_p) || params.top_p < 0 || params.top_p > 1) { + return; + } +#if ET_HAS_EXCEPTIONS + try { +#endif + if (seed) { + state.seed = *seed; + } else { + std::random_device device; + state.seed = (static_cast(device()) << 32) ^ device(); } + state.sampling = params; +#if ET_HAS_EXCEPTIONS + } catch (const std::exception& error) { + ET_LOG( + Error, "CUDA DFlash: could not initialize sampling: %s", error.what()); } +#endif } bool MuseGlimmerCudaDFlashExecutor::rewind(State& state, int position) { @@ -688,57 +774,51 @@ bool MuseGlimmerCudaDFlashExecutor::rewind(State& state, int position) { bool MuseGlimmerCudaDFlashExecutor::run_target( const RawStep& step, const std::vector& token_owners, - std::vector& predictions, - bool drain_context) { + std::vector& predictions) { if (token_owners.size() != step.tokens.size()) { 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; - } + ++stats_.target_preamble_calls; std::size_t predicted = 0; const bool ok = target_->execute_raw(step, [&](const cb::RawTargetOutput& output) { - if (drain_context) { - ++stats_.target_preamble_forwards; - } else { - ++stats_.target_verify_forwards; - } + ++stats_.target_preamble_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; - // 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_; - } + // The sampler or the next target slice may invalidate these taps. + std::vector copied( + 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(); + } + std::vector params; + for (const auto selected : output.selected_indices) { + const auto index = step.logit_indices[selected]; + const auto& state = sessions_.at(token_owners[index]); + const auto& sampling = *state.sampling; + params.insert( + params.end(), + {sampling.temperature, + sampling.top_p, + static_cast(sampling.top_k), + sampling.temperature > 0.0f + ? df::positional_uniform( + *state.seed, + static_cast( + step.positions[index] + 1), + 0) + : 0.0f}); } - auto tokens = - target_->argmax_rows(output.logits, output.selected_indices.size()); + auto tokens = target_->sample_rows( + output.logits, output.selected_indices.size(), params); if (!tokens) { return false; } @@ -746,9 +826,6 @@ 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; @@ -761,11 +838,9 @@ bool MuseGlimmerCudaDFlashExecutor::run_target( state.frontier = static_cast(step.positions[index]) + 1; touched.insert(owner); } - if (drain_context) { - for (const auto owner : touched) { - if (!drain(owner)) { - return false; - } + for (const auto owner : touched) { + if (!drain(owner)) { + return false; } } return true; @@ -778,6 +853,121 @@ bool MuseGlimmerCudaDFlashExecutor::run_target( })); } +std::vector MuseGlimmerCudaDFlashExecutor::verify_params( + const State& state, + batching::Position first) const { + const auto& sampling = *state.sampling; + std::vector params(proposals_ + 4, 0.0f); + params[0] = sampling.temperature; + params[1] = static_cast(sampling.top_k); + params[2] = sampling.top_p; + if (sampling.temperature > 0.0f) { + for (int j = 0; j < proposals_; ++j) { + params[3 + j] = df::positional_uniform(*state.seed, first + j, 2); + } + params.back() = df::positional_uniform(*state.seed, first, 3); + } + return params; +} + +bool MuseGlimmerCudaDFlashExecutor::verify( + const RawStep& step, + const std::vector& owners, + const std::vector& candidates, + std::vector& results) { + const int width = proposals_ + 1; + if (!verification_hidden_.begin(step.tokens.size(), row_bytes_) || + candidates.size() != owners.size() * width) { + return false; + } + results.assign(owners.size() * 2, -1); + ++stats_.target_verify_calls; + return target_->execute_raw(step, [&](const cb::RawTargetOutput& output) { + ++stats_.target_verify_forwards; + if (!output.taps || output.taps->dim() != 3 || + output.taps->scalar_type() != aten::ScalarType::BFloat16 || + output.taps->size(2) != hidden_width_ || + !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_; + } + // Every verification row samples, so a slice holds whole owner blocks. + if (output.offset % width != 0 || output.length % width != 0 || + output.selected_indices.size() != + static_cast(output.length) || + output.logits.device_type() != aten::DeviceType::CUDA || + output.logits.scalar_type() != aten::ScalarType::Float) { + return false; + } + const int first = output.offset / width; + const int count = output.length / width; + const bool greedy = std::all_of( + owners.begin() + first, + owners.begin() + first + count, + [this](SessionId owner) { + return sessions_.at(owner).sampling->temperature == 0.0f; + }); + std::vector params; + if (!greedy) { + params.reserve(count * (width + 3)); + for (int i = first; i < first + count; ++i) { + const auto row = verify_params( + sessions_.at(owners[i]), + static_cast(step.positions[i * width] + 1)); + params.insert(params.end(), row.begin(), row.end()); + } + } + const auto stream = + ::executorch::extension::cuda::getCallerStream().value_or( + cudaStreamPerThread); + std::vector> parts{ + {candidates.data() + first * width, + static_cast(count * width) * sizeof(std::int64_t)}}; + if (!greedy) { + parts.emplace_back(params.data(), params.size() * sizeof(float)); + } + auto staged = verify_inputs_.upload(parts, stream); + if (!staged) { + return false; + } + const aten::Device cuda( + aten::DeviceType::CUDA, output.logits.device_index()); + auto logits = make_tensor_ptr( + {count, width, target_->vocab_size()}, + output.logits.mutable_data_ptr(), + aten::ScalarType::Float, + cuda); + auto tokens = make_tensor_ptr( + {count, width}, (*staged)[0], aten::ScalarType::Long, cuda); + auto result = greedy + ? target_->module().execute( + df::kVerifyArgmax, std::vector{logits, tokens}) + : target_->module().execute( + df::kVerify, + {logits, + tokens, + make_tensor_ptr( + {count, width + 3}, + (*staged)[1], + aten::ScalarType::Float, + cuda)}); + if (!result.ok() || result->size() != 1 || !result->at(0).isTensor()) { + return false; + } + const auto& verdict = result->at(0).toTensor(); + if (verdict.scalar_type() != aten::ScalarType::Long || + verdict.numel() != 2 * count) { + return false; + } + return copy_host( + results.data() + 2 * first, verdict, 2 * count * sizeof(std::int64_t)); + }); +} + bool MuseGlimmerCudaDFlashExecutor::drain(SessionId owner) { auto& state = sessions_.at(owner); while (state.hidden.size() / row_bytes_ > @@ -935,7 +1125,7 @@ bool MuseGlimmerCudaDFlashExecutor::speculate( return false; } RawStep verification; - std::vector token_owners; + std::vector candidates; std::vector starts; const int width = proposals_ + 1; for (std::size_t i = 0; i < owners.size(); ++i) { @@ -947,20 +1137,20 @@ bool MuseGlimmerCudaDFlashExecutor::speculate( 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]); + candidates.push_back(verification.tokens.back()); } } - std::vector predictions; - if (!run_target(verification, token_owners, predictions, false)) { + std::vector verdicts; + if (!verify(verification, owners, candidates, verdicts)) { return false; } - std::vector bursts; + std::vector bursts; for (std::size_t i = 0; i < owners.size(); ++i) { - auto burst = df::verify_greedy( + auto burst = df::make_burst( proposals[i], - std::vector( - predictions.begin() + i * width, - predictions.begin() + (i + 1) * width)); + verdicts[2 * i], + verdicts[2 * i + 1], + target_->vocab_size()); if (!burst || !verification_hidden_.copy_rows( sessions_.at(owners[i]).hidden, i * width, burst->accepted + 1)) { @@ -1067,10 +1257,8 @@ bool MuseGlimmerCudaDFlashExecutor::execute( (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)"); + it->second.sampling->temperature < 0.0f))) { + ET_LOG(Error, "CUDA DFlash: invalid input or sampling configuration"); return reject_batch(); } const auto start = @@ -1128,7 +1316,7 @@ bool MuseGlimmerCudaDFlashExecutor::execute( } } else { std::vector predictions; - if (!run_target(step, token_owners, predictions, true)) { + if (!run_target(step, token_owners, predictions)) { return fail_batch(batch, out); } for (std::size_t i = 0; i < batch.inputs.size(); ++i) { diff --git a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h index 9316da30680..5088a165041 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.h @@ -24,16 +24,30 @@ 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"; +// On-device verification of [owners, width] candidates; see +// model/dflash_token_sampler.py. +inline constexpr char kVerify[] = "dflash_verify"; +inline constexpr char kVerifyArgmax[] = "dflash_verify_argmax"; +inline constexpr int kOwnerLimit = 4; -struct GreedyBurst { +struct Burst { std::vector tokens; std::size_t accepted = 0; bool bonus = false; }; -// Verification returns only new predictions, never the consumed anchor. -std::optional verify_greedy( +// The new tokens of one verified block: the accepted proposals, then the +// correction or bonus. Never the consumed anchor. +std::optional make_burst( const std::vector& proposals, - const std::vector& predictions); + std::int64_t accepted, + std::int64_t token, + std::int32_t vocab_size); +// A uniform in [0, 1) that depends only on the seed, position and purpose, so +// a seed reproduces a generation however its requests are batched. +float positional_uniform( + std::uint64_t seed, + batching::Position position, + int purpose); bool trim_hidden( std::vector& hidden, std::size_t row_bytes, @@ -90,6 +104,26 @@ std::optional pack_draft( int context_limit, int max_context, std::size_t row_bytes); + +// Reusable device memory for small host-built inputs. +class DeviceStaging { + public: + ~DeviceStaging(); + DeviceStaging() = default; + DeviceStaging(const DeviceStaging&) = delete; + DeviceStaging& operator=(const DeviceStaging&) = delete; + DeviceStaging(DeviceStaging&&) = delete; + DeviceStaging& operator=(DeviceStaging&&) = delete; + // Device copies of each (data, bytes) part, ordered on stream. Earlier + // uploads must be consumed first: they share one allocation. + std::optional> upload( + const std::vector>& parts, + void* stream); + + private: + void* device_ = nullptr; + std::size_t capacity_ = 0; +}; } // namespace cuda_dflash_detail struct CudaDFlashStats { @@ -114,9 +148,9 @@ struct CudaDFlashStats { 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. +// Text-only model adapter over the existing CUDA target executor. Drafts are +// greedy; dflash_verify accepts or resamples them under each session's sampling +// on the device, so only token IDs return to the host. class MuseGlimmerCudaDFlashExecutor final : public batching::Executor { public: static runtime::Result create( @@ -180,8 +214,16 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor { bool run_target( const RawStep& step, const std::vector& token_owners, - std::vector& predictions, - bool drain_context); + std::vector& predictions); + // Runs the verification step and the on-device verifier; results holds + // [accepted, correction] per owner. + bool verify( + const RawStep& step, + const std::vector& owners, + const std::vector& candidates, + std::vector& results); + std::vector verify_params(const State& state, batching::Position first) + const; bool drain(batching::SessionId owner); bool draft( const std::vector& owners, @@ -210,6 +252,7 @@ class MuseGlimmerCudaDFlashExecutor final : public batching::Executor { std::unordered_map sessions_; CudaDFlashStats stats_; cuda_dflash_detail::HiddenSnapshot verification_hidden_; + cuda_dflash_detail::DeviceStaging verify_inputs_; 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 b5fd6b17df7..73878110346 100644 --- a/examples/models/muse-glimmer/tests/cuda_dflash_executor_test.cpp +++ b/examples/models/muse-glimmer/tests/cuda_dflash_executor_test.cpp @@ -19,16 +19,16 @@ namespace aten = executorch::aten; 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}); +TEST(CudaDFlash, FirstRejectionReturnsOnlyCorrection) { + const auto burst = df::make_burst({11, 12, 13}, 0, 21, 64); 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}); +TEST(CudaDFlash, PartialAcceptanceEndsWithCorrection) { + const auto burst = df::make_burst({11, 12, 13}, 1, 22, 64); ASSERT_TRUE(burst); EXPECT_EQ(burst->tokens, std::vector({11, 22})); EXPECT_EQ(burst->accepted, 1); @@ -36,18 +36,29 @@ TEST(CudaDFlash, PartialAcceptanceStopsAtFirstMismatch) { } TEST(CudaDFlash, FullAcceptanceAddsOneBonus) { - const auto burst = df::verify_greedy({11, 12, 13}, {11, 12, 13, 14}); + const auto burst = df::make_burst({11, 12, 13}, 3, 14, 64); 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, InvalidVerificationRejected) { + EXPECT_FALSE(df::make_burst({}, 0, 1, 64)); + EXPECT_FALSE(df::make_burst({1, 2}, -1, 1, 64)); + EXPECT_FALSE(df::make_burst({1, 2}, 3, 1, 64)); + EXPECT_FALSE(df::make_burst({1, 2}, 1, 64, 64)); + EXPECT_FALSE(df::make_burst({1, 2}, 1, -1, 64)); +} + +TEST(CudaDFlash, PositionalUniformsAreSeededAndIndependentOfBatching) { + const float u = df::positional_uniform(7, 100, 2); + EXPECT_EQ(u, df::positional_uniform(7, 100, 2)); + EXPECT_GE(u, 0.0f); + EXPECT_LT(u, 1.0f); + EXPECT_NE(u, df::positional_uniform(8, 100, 2)); + EXPECT_NE(u, df::positional_uniform(7, 101, 2)); + EXPECT_NE(u, df::positional_uniform(7, 100, 3)); } TEST(CudaDFlash, NewPreambleAnchorIsRetainedOnce) { @@ -55,7 +66,15 @@ TEST(CudaDFlash, NewPreambleAnchorIsRetainedOnce) { // 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); + const auto accepted = static_cast( + std::mismatch( + predictions.begin(), + predictions.begin() + 3, + std::vector({11, 12, 13}).begin()) + .first - + predictions.begin()); + auto burst = + df::make_burst({11, 12, 13}, accepted, predictions[accepted], 64); ASSERT_TRUE(burst); std::vector output{10}; output.insert(output.end(), burst->tokens.begin(), burst->tokens.end()); @@ -69,7 +88,7 @@ TEST(CudaDFlash, NewPreambleAnchorIsRetainedOnce) { } TEST(CudaDFlash, BudgetRewindTrimsOnlyCommittedRows) { - const auto burst = df::verify_greedy({11, 12, 13}, {11, 12, 13, 14}); + const auto burst = df::make_burst({11, 12, 13}, 3, 14, 64); ASSERT_TRUE(burst); // Draft has warmed through input.end=100; target verified four rows. for (int kept = 1; kept <= 4; ++kept) { @@ -139,12 +158,11 @@ TEST(CudaDFlash, NearContextLimitRequiresWholeTrainedBlock) { 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); + const auto burst = df::make_burst( + proposals, + static_cast(rejected), + rejected == proposals.size() ? 90 : 80, + 128); ASSERT_TRUE(burst); std::vector expected( proposals.begin(), proposals.begin() + rejected); diff --git a/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp b/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp index fbd9ae8301a..a12386237ea 100644 --- a/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp +++ b/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp @@ -120,6 +120,15 @@ struct CudaDFlashTestPeer { static bool poisoned(DFlash& executor, b::SessionId sid) { return executor.sessions_.at(sid).poisoned; } + static std::optional sampling_seed( + DFlash& executor, + b::SessionId sid) { + return executor.sessions_.at(sid).seed; + } + static std::vector + verify_params(DFlash& executor, b::SessionId sid, b::Position first) { + return executor.verify_params(executor.sessions_.at(sid), first); + } }; } // namespace executorch::extension::llm using Peer = llm::CudaDFlashTestPeer; @@ -183,7 +192,9 @@ class FixtureModule final : public Module { } const auto stream = executorch::extension::cuda::getCallerStream().value_or( cudaStreamPerThread); - if (name == "sample_argmax") { + // Verification samples through the device verifier, preambles through + // sample_argmax; either is the target sampler. + if (name == "sample_argmax" || name == "dflash_verify_argmax") { if (taps_ && cudaMemsetAsync(taps_, 0, 32 * 256, stream) != cudaSuccess) { return rt::Error::Internal; } @@ -193,10 +204,11 @@ class FixtureModule final : public Module { } return actual; } - if (name != "draft_forward" && name.rfind("forward_", 0) != 0) { + const bool draft = name.rfind("draft_forward", 0) == 0; + if (!draft && name.rfind("forward_", 0) != 0) { return actual; } - target_output_ = name != "draft_forward"; + target_output_ = !draft; if (!deterministic_) { return actual; } @@ -205,7 +217,8 @@ class FixtureModule final : public Module { 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; + // Draft tokens are [owners, block]; every draft query row returns logits. + const int rows = target_output_ ? selector.numel() : tokens.numel(); 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) { @@ -328,6 +341,66 @@ class CudaDFlashNative : public ::testing::Test { std::unique_ptr executor_; }; +TEST_F( + CudaDFlashNative, + SamplingInitializesMissingSeedAndPreservesExplicitZero) { + create(); + ASSERT_TRUE(executor_); + const auto sid = open(); + executor_->set_sampling(sid, {0.8f, 0.9f, 7}, std::nullopt); + EXPECT_TRUE(Peer::sampling_seed(*executor_, sid)); + executor_->set_sampling(sid, {0.8f, 0.9f, 7}, 0); + ASSERT_TRUE(Peer::sampling_seed(*executor_, sid)); + EXPECT_EQ(*Peer::sampling_seed(*executor_, sid), 0); + const auto params = Peer::verify_params(*executor_, sid, 100); + ASSERT_EQ(params.size(), 19); + EXPECT_EQ(params[3], llm::cuda_dflash_detail::positional_uniform(0, 100, 2)); + EXPECT_EQ( + params.back(), llm::cuda_dflash_detail::positional_uniform(0, 100, 3)); + executor_->set_sampling(sid, {0.8f, 0.9f, 7}, 42); + ASSERT_TRUE(Peer::sampling_seed(*executor_, sid)); + EXPECT_EQ(*Peer::sampling_seed(*executor_, sid), 42); + executor_->set_sampling(sid, {0.8f, 0.9f, 7}, std::nullopt); + EXPECT_TRUE(Peer::sampling_seed(*executor_, sid)); +} + +TEST_F(CudaDFlashNative, MixedOwnerVerifyParamsSkipOnlyGreedyCoins) { + create(); + ASSERT_TRUE(executor_); + const std::array sampling{ + b::SamplingParams{0.0f, 1.0f, 1}, + b::SamplingParams{0.8f, 0.9f, 7}, + b::SamplingParams{0.0f, 0.5f, 3}, + b::SamplingParams{1.2f, 0.8f, 5}}; + const std::array seeds{0, 42, 19, 17}; + const std::array positions{10, 100, 50, 200}; + std::array owners; + for (int owner = 0; owner < 4; ++owner) { + owners[owner] = open(); + executor_->set_sampling(owners[owner], sampling[owner], seeds[owner]); + } + for (const int owner : {3, 0, 2, 1, 3}) { + const auto params = + Peer::verify_params(*executor_, owners[owner], positions[owner]); + ASSERT_EQ(params.size(), 19); + EXPECT_EQ(params[0], sampling[owner].temperature); + EXPECT_EQ(params[1], sampling[owner].top_k); + EXPECT_EQ(params[2], sampling[owner].top_p); + for (int proposal = 0; proposal < 15; ++proposal) { + const float expected = sampling[owner].temperature > 0.0f + ? llm::cuda_dflash_detail::positional_uniform( + seeds[owner], positions[owner] + proposal, 2) + : 0.0f; + EXPECT_EQ(params[3 + proposal], expected); + } + const float correction = sampling[owner].temperature > 0.0f + ? llm::cuda_dflash_detail::positional_uniform( + seeds[owner], positions[owner], 3) + : 0.0f; + EXPECT_EQ(params.back(), correction); + } +} + TEST_F(CudaDFlashNative, EveryRejectionPositionAndFullBonus) { create(); ASSERT_TRUE(executor_); @@ -492,15 +565,16 @@ TEST_F(CudaDFlashNative, WrongFeedbackAndNonOutputSlicesKeepPreamble) { 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); + // An invalid temperature is refused before a forward, without poisoning. + 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)); + EXPECT_EQ(executor_->stats().failed_batches, valid.failed_batches); + EXPECT_FALSE(Peer::poisoned(*executor_, sid)); } TEST_F(CudaDFlashNative, FinalWholeBlockFitsExactlyAndPreservesAllPrefillRows) { diff --git a/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py b/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py index 1509da4ea51..db7afaa2072 100644 --- a/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py +++ b/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py @@ -17,6 +17,7 @@ import torch from executorch.examples.models.muse_glimmer.export import common from executorch.examples.models.muse_glimmer.export.export_dflash import ( + cuda_owner_verifier_programs, prepare_cuda_packed_draft, ) from executorch.examples.models.muse_glimmer.export.export_solo import step_forward @@ -444,6 +445,8 @@ def owned_schema(package, name): "draft_forward", "sample", "sample_argmax", + "dflash_verify", + "dflash_verify_argmax", }, ) from executorch.backends.cuda.cuda_weight_collector import ( @@ -484,5 +487,18 @@ def owned_schema(package, name): self.assertEqual(shared_keys[0], shared_keys[2]) +class SingleOwnerVerifierExportTest(unittest.TestCase): + def test_single_owner_accepts_runtime_draft_widths(self): + methods = cuda_owner_verifier_programs(32, 1, 16) + verifier = methods["dflash_verify_argmax"].module() + for width in (3, 4, 16): + logits = torch.zeros(1, width, 32) + logits[:, :, 7] = 1 + candidates = torch.full((1, width), 7, dtype=torch.long) + self.assertEqual(verifier(logits, candidates).tolist(), [[width - 1, 7]]) + candidates[0, 2] = 5 + self.assertEqual(verifier(logits, candidates).tolist(), [[1, 7]]) + + if __name__ == "__main__": unittest.main() diff --git a/examples/models/muse-glimmer/tests/test_dflash_token_sampler.py b/examples/models/muse-glimmer/tests/test_dflash_token_sampler.py index 2d1e0a2bf56..bd3598b261b 100644 --- a/examples/models/muse-glimmer/tests/test_dflash_token_sampler.py +++ b/examples/models/muse-glimmer/tests/test_dflash_token_sampler.py @@ -12,13 +12,23 @@ import torch from executorch.examples.models.muse_glimmer.model.dflash_token_sampler import ( + DFlashOwnerArgmaxVerifier, + DFlashOwnerVerifier, DFlashTokenSampler, sample_tokens, + verify_owners, + verify_owners_greedy, ) def _reference_row( - logits: list[float], temperature: float, top_k: int, top_p: float, uniform: float + logits: list[float], + temperature: float, + top_k: int, + top_p: float, + uniform: float, + *, + batching_policy: bool = False, ) -> tuple[int, list[float]]: vocab_size = len(logits) if temperature <= 0.0: @@ -26,17 +36,22 @@ def _reference_row( return token, [float(index == token) for index in range(vocab_size)] ranked = sorted(range(vocab_size), key=lambda index: (-logits[index], index)) - if 0 < top_k < vocab_size: + use_top_k = 0 < top_k < vocab_size + if use_top_k: ranked = ranked[:top_k] maximum = max(logits[index] for index in ranked) weights = [0.0] * vocab_size for index in ranked: weights[index] = math.exp((logits[index] - maximum) / temperature) total = sum(weights) - if 0.0 < top_p < 1.0: + if 0.0 < top_p < 1.0 and not (batching_policy and use_top_k): nucleus, nucleus_sum = [], 0.0 for index in ranked: - if nucleus_sum >= top_p * total: + if ( + nucleus_sum > top_p * total + if batching_policy + else nucleus_sum >= top_p * total + ): break nucleus.append(index) nucleus_sum += weights[index] @@ -139,3 +154,146 @@ def test_strict_export_supports_dynamic_rows(self) -> None: if __name__ == "__main__": unittest.main() + + +class DFlashOwnerVerifierTest(unittest.TestCase): + """Owner verification uses the batching target's sampling distribution.""" + + def _reference(self, logits, candidates, params): + width = candidates.shape[1] + expected = [] + for owner in range(candidates.shape[0]): + row = params[owner] + distributions = [ + _reference_row( + values.tolist(), + float(row[0]), + int(row[1]), + float(row[2]), + 0.5, + batching_policy=True, + )[1] + for values in logits[owner] + ] + accepted = 0 + for index in range(width - 1): + proposal = int(candidates[owner, index + 1]) + distribution = distributions[index] + keep = ( + float(row[3 + index]) < distribution[proposal] + if row[0] > 0 + else proposal + == max(range(len(distribution)), key=distribution.__getitem__) + ) + if not keep: + break + accepted += 1 + distribution = distributions[accepted].copy() + if row[0] == 0: + correction = max(range(len(distribution)), key=distribution.__getitem__) + else: + if accepted < width - 1: + distribution[int(candidates[owner, accepted + 1])] = 0.0 + total = sum(distribution) + distribution = [value / total for value in distribution] + cumulative = 0.0 + correction = max(i for i, value in enumerate(distribution) if value > 0) + for token, value in enumerate(distribution): + cumulative += value + if cumulative > float(row[-1]): + correction = token + break + expected.append([accepted, correction]) + return torch.tensor(expected) + + def _case(self, owners, width, vocab, temperature, top_k, top_p, seed): + generator = torch.Generator().manual_seed(seed) + logits = torch.randn(owners, width, vocab, generator=generator) * 3 + greedy = logits.argmax(-1) + candidates = torch.randint(0, vocab, (owners, width), generator=generator) + # Owner o's draft agrees with the target for its first o proposals. + for owner in range(owners): + agree = min(owner, width - 1) + candidates[owner, 1 : agree + 1] = greedy[owner, :agree] + params = torch.cat( + [ + torch.tensor([[temperature, top_k, top_p]]).expand(owners, 3), + torch.rand(owners, width, generator=generator), + ], + dim=1, + ) + return logits, candidates, params + + def test_matches_independent_batch_sampling_reference(self) -> None: + for temperature, top_k, top_p in ( + (0.0, 0, 1.0), + (0.8, 0, 1.0), + (1.0, 5, 1.0), + (0.7, 0, 0.9), + (1.2, 7, 0.8), + ): + for seed in range(6): + with self.subTest(t=temperature, k=top_k, p=top_p, seed=seed): + logits, candidates, params = self._case( + 4, 16, 32, temperature, top_k, top_p, seed + ) + torch.testing.assert_close( + verify_owners(logits, candidates, params), + self._reference(logits, candidates, params), + ) + + def test_top_k_overrides_top_p_like_the_target_sampler(self) -> None: + logits = torch.full((1, 3, 6), -100.0) + logits[:, :, :4] = 0.0 + candidates = torch.tensor([[0, 0, 0]]) + params = torch.tensor([[1.0, 4.0, 0.1, 0.5, 0.5, 0.4]]) + self.assertEqual(verify_owners(logits, candidates, params).tolist(), [[0, 2]]) + + def test_top_p_includes_the_equal_cumulative_mass_boundary(self) -> None: + logits = torch.full((1, 3, 6), -100.0) + logits[:, :, :4] = 0.0 + candidates = torch.tensor([[0, 2, 2]]) + params = torch.tensor([[1.0, 0.0, 0.5, 0.2, 0.2, 0.75]]) + self.assertEqual(verify_owners(logits, candidates, params).tolist(), [[2, 2]]) + + def test_bonus_zero_coin_selects_a_supported_token(self) -> None: + logits = torch.full((1, 3, 6), -100.0) + logits[:, :, 3] = 0.0 + logits[:, :, 4] = 1.0 + candidates = torch.tensor([[0, 4, 4]]) + params = torch.tensor([[1.0, 2.0, 1.0, 0.0, 0.0, 0.0]]) + self.assertEqual(verify_owners(logits, candidates, params).tolist(), [[2, 3]]) + + def test_greedy_fast_path_matches_and_counts_prefix(self) -> None: + logits, candidates, params = self._case(4, 16, 32, 0.0, 0, 1.0, 3) + candidates[3, 1:] = logits[3, :-1].argmax(-1) + result = verify_owners_greedy(logits, candidates) + torch.testing.assert_close(result, verify_owners(logits, candidates, params)) + self.assertEqual(result[:, 0].tolist(), [0, 1, 2, 15]) + self.assertEqual(int(result[3, 1]), int(logits[3, 15].argmax())) + + def test_strict_export_supports_dynamic_owners_and_width(self) -> None: + logits, candidates, params = self._case(4, 16, 32, 0.9, 0, 1.0, 1) + owners = torch.export.Dim("owners", min=1, max=4) + width = torch.export.Dim("width", min=3, max=16) + sampled = torch.export.export( + DFlashOwnerVerifier(), + (logits, candidates, params), + dynamic_shapes=( + {0: owners, 1: width}, + {0: owners, 1: width}, + {0: owners, 1: width + 3}, + ), + strict=True, + ) + greedy = torch.export.export( + DFlashOwnerArgmaxVerifier(), + (logits, candidates), + dynamic_shapes=({0: owners, 1: width}, {0: owners, 1: width}), + strict=True, + ) + small = (logits[:2, :4], candidates[:2, :4], params[:2, :7]) + torch.testing.assert_close(sampled.module()(*small), verify_owners(*small)) + torch.testing.assert_close( + greedy.module()(*small[:2]), verify_owners_greedy(*small[:2]) + )