From 71401093156992003ea2158168f9e8c98cfee31a Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Sun, 18 Oct 2026 21:59:01 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/cuda/batching/cuda_executor.cpp | 64 +++- backends/cuda/batching/cuda_executor.h | 13 + .../muse-glimmer/export/export_dflash.py | 43 ++- .../models/muse-glimmer/export/export_solo.py | 30 +- .../model/dflash_token_sampler.py | 87 +++++ .../runtime/cuda_dflash_executor.cpp | 326 +++++++++++++----- .../runtime/cuda_dflash_executor.h | 60 +++- .../tests/cuda_dflash_executor_test.cpp | 52 ++- .../tests/cuda_dflash_native_test.cpp | 16 +- .../tests/test_cuda_dflash_packed.py | 2 + .../tests/test_dflash_token_sampler.py | 100 ++++++ 11 files changed, 649 insertions(+), 144 deletions(-) diff --git a/backends/cuda/batching/cuda_executor.cpp b/backends/cuda/batching/cuda_executor.cpp index 236eb43a455..b1bd377eb9a 100644 --- a/backends/cuda/batching/cuda_executor.cpp +++ b/backends/cuda/batching/cuda_executor.cpp @@ -791,21 +791,34 @@ 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 || logits.device_type() != ::executorch::aten::DeviceType::CUDA) { + // Host logits are sampled greedily only; the exported sampler is on CUDA. + return all_greedy ? argmax_rows(logits, rows) : std::nullopt; + } + if (!initialized_ || !has_sampler_ || logits.dim() < 2 || + 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); @@ -821,6 +834,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( @@ -832,18 +864,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 @@ -871,7 +903,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..aa956842d16 100644 --- a/examples/models/muse-glimmer/export/export_dflash.py +++ b/examples/models/muse-glimmer/export/export_dflash.py @@ -95,6 +95,46 @@ 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 + + owners = Dim("verify_owners", min=1, max=max_owners) + # 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 +377,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 +398,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 d7a0015b798..f3d80a418a1 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,12 +469,24 @@ 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: @@ -484,7 +496,7 @@ def partitioner(name: str) -> CudaPartitioner: CompileSpec("autotune_at_compile_time", b"OFF"), ] ) - 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"), @@ -498,7 +510,7 @@ def partitioner(name: str) -> CudaPartitioner: ), CompileSpec("autotune_at_compile_time", b"OFF"), ] - if "draft_forward" in cache_methods: + if draft_methods: specs.append( CompileSpec( "llm_cache_registry_key", @@ -543,7 +555,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..17d12c45fc4 100644 --- a/examples/models/muse-glimmer/model/dflash_token_sampler.py +++ b/examples/models/muse-glimmer/model/dflash_token_sampler.py @@ -195,6 +195,93 @@ 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) + ) + 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), + 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 2e12d017618..79d544f9047 100644 --- a/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp +++ b/examples/models/muse-glimmer/runtime/cuda_dflash_executor.cpp @@ -19,6 +19,7 @@ #include #include #include +#include #include namespace executorch::extension::llm { @@ -30,26 +31,45 @@ 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 +140,49 @@ std::optional pack_draft( } return plan; } + +DeviceStaging::~DeviceStaging() { + if (device_) { + 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 { @@ -376,11 +439,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")); @@ -397,7 +461,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() && @@ -411,10 +475,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)); @@ -542,12 +607,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"); @@ -656,11 +724,6 @@ void MuseGlimmerCudaDFlashExecutor::set_sampling( 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"); - } } } @@ -682,57 +745,48 @@ 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(); } - auto tokens = - target_->argmax_rows(output.logits, output.selected_indices.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), + df::positional_uniform( + state.seed.value_or(0), + static_cast(step.positions[index] + 1), + 0)}); + } + auto tokens = target_->sample_rows( + output.logits, output.selected_indices.size(), params); if (!tokens) { return false; } @@ -740,9 +794,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; @@ -755,11 +806,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; @@ -772,6 +821,113 @@ bool MuseGlimmerCudaDFlashExecutor::run_target( })); } +std::vector MuseGlimmerCudaDFlashExecutor::verify_params( + const State& state, + batching::Position first) const { + const auto& sampling = *state.sampling; + const auto seed = state.seed.value_or(0); + std::vector params{ + sampling.temperature, + static_cast(sampling.top_k), + sampling.top_p}; + for (int j = 0; j < proposals_; ++j) { + params.push_back(df::positional_uniform(seed, first + j, 2)); + } + params.push_back(df::positional_uniform(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; + bool greedy = true; + std::vector params; + for (int i = first; i < first + count; ++i) { + const auto& state = sessions_.at(owners[i]); + greedy = greedy && state.sampling->temperature == 0.0f; + const auto row = verify_params( + state, 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_ > @@ -926,7 +1082,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) { @@ -938,20 +1094,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)) { @@ -1050,10 +1206,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 fail_batch(batch, out); } const auto start = @@ -1110,7 +1264,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 f3ee876d4ac..eb9729e3c55 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, @@ -88,6 +102,24 @@ 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; + // 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 { @@ -112,9 +144,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( @@ -170,8 +202,17 @@ 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, @@ -200,6 +241,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 4354af6befd..740788d189d 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,13 @@ 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 +86,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 +156,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 9774a983631..4a43d08107e 100644 --- a/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp +++ b/examples/models/muse-glimmer/tests/cuda_dflash_native_test.cpp @@ -183,7 +183,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 +195,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 +208,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) { @@ -489,8 +493,8 @@ 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 fails before a target/draft forward. + executor_->set_sampling(sid, {-0.5f, 1.0f, 0}, 0); const auto valid = executor_->stats(); EXPECT_FALSE(executor_->execute({{input(sid, 10, {prediction(9)})}}, out)); EXPECT_FALSE(out.outputs[0]); 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 dbeb804f7a8..49528490719 100644 --- a/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py +++ b/examples/models/muse-glimmer/tests/test_cuda_dflash_packed.py @@ -444,6 +444,8 @@ def owned_schema(package, name): "draft_forward", "sample", "sample_argmax", + "dflash_verify", + "dflash_verify_argmax", }, ) from executorch.backends.cuda.cuda_weight_collector import ( 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..7292a750c56 100644 --- a/examples/models/muse-glimmer/tests/test_dflash_token_sampler.py +++ b/examples/models/muse-glimmer/tests/test_dflash_token_sampler.py @@ -12,8 +12,13 @@ import torch from executorch.examples.models.muse_glimmer.model.dflash_token_sampler import ( + DFlashOwnerArgmaxVerifier, + DFlashOwnerVerifier, DFlashTokenSampler, sample_tokens, + verify_owners, + verify_owners_greedy, + verify_speculative, ) @@ -139,3 +144,98 @@ def test_strict_export_supports_dynamic_rows(self) -> None: if __name__ == "__main__": unittest.main() + + +class DFlashOwnerVerifierTest(unittest.TestCase): + """The batched verifier equals per-owner ``verify_speculative`` on greedy drafts.""" + + def _reference(self, logits, candidates, params): + width = candidates.shape[1] + expected = [] + for owner in range(candidates.shape[0]): + row = params[owner] + result = verify_speculative( + logits[owner], + torch.zeros(width - 1, logits.shape[-1]), + candidates[owner], + row[0:1], + row[1:2].long(), + row[2:3], + torch.tensor([True]), + row[3 : width + 2], + row[width + 2 : width + 3], + ) + expected.append([int(result[0]) - 1, int(result[1])]) + 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_single_owner_verifier(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_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]) + )