diff --git a/.gitignore b/.gitignore index 13529ace2a..b993f46469 100644 --- a/.gitignore +++ b/.gitignore @@ -50,6 +50,9 @@ examples/csharp/ModelChat/models !/test/models/engine/ /test/models/engine/* +!/test/models/engine/synthetic-fixed-components/ +!/test/models/engine/synthetic-fixed-components/decoder.onnx +!/test/models/engine/synthetic-fixed-components/genai_config.json !/test/models/engine/synthetic-paged-scales/ !/test/models/engine/synthetic-paged-scales/decoder.onnx !/test/models/engine/synthetic-paged-scales/genai_config.json diff --git a/CMakeLists.txt b/CMakeLists.txt index 7cc57b4fb6..73276d5794 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -205,6 +205,22 @@ endif() # The genai library itself is always embedded in the shared library list(APPEND ortgenai_embed_libs "$") +if(WIN32 AND NOT USE_WINML AND ENABLE_TESTS) + file(GLOB ort_runtime_dlls "${ORT_LIB_DIR}/*.dll") + if(ort_runtime_dlls) + add_custom_target(onnxruntime-genai-ort-runtime + COMMAND ${CMAKE_COMMAND} -E make_directory + "$" + COMMAND ${CMAKE_COMMAND} -E copy_if_different + ${ort_runtime_dlls} + "$" + DEPENDS ${ort_runtime_dlls} + COMMAND_EXPAND_LISTS + ) + add_dependencies(onnxruntime-genai onnxruntime-genai-ort-runtime) + endif() +endif() + # we keep the shared libraries disconnected on Android as they will come from separate AARs and we don't want to force # the ORT version to match in both. if(CMAKE_SYSTEM_NAME STREQUAL "Android" OR CMAKE_SYSTEM_NAME STREQUAL "Linux" OR (CMAKE_SYSTEM_NAME STREQUAL "Darwin" AND (NOT BUILD_APPLE_FRAMEWORK) AND (NOT MAC_CATALYST))) diff --git a/build.py b/build.py index cfb91e258c..41c0cec50a 100644 --- a/build.py +++ b/build.py @@ -9,6 +9,7 @@ import platform import shlex import shutil +import stat import subprocess import sys import textwrap @@ -557,12 +558,31 @@ def _run_android_tests(args: argparse.Namespace): def _get_windows_build_args(args: argparse.Namespace): + spectre_enabled = True + for definition in args.cmake_extra_defines: + name, separator, value = definition.removeprefix("-D").partition("=") + if name.split(":", 1)[0] == "ENABLE_MSVC_SPECTRE" and separator: + value = value.strip().upper() + spectre_enabled = value not in { + "", + "0", + "OFF", + "NO", + "FALSE", + "N", + "IGNORE", + "NOTFOUND", + } and not value.endswith("-NOTFOUND") + spectre_flag = " /Qspectre" if spectre_enabled else "" win_args = [ "-DCMAKE_EXE_LINKER_FLAGS_INIT=/profile /DYNAMICBASE", "-DCMAKE_MODULE_LINKER_FLAGS_INIT=/profile /DYNAMICBASE", "-DCMAKE_SHARED_LINKER_FLAGS_INIT=/profile /DYNAMICBASE", ] - cmake_c_flags = "/EHsc /Qspectre /MP /guard:cf /DWIN32 /D_WINDOWS /DWINAPI_FAMILY=100 /DWINVER=0x0A00 /D_WIN32_WINNT=0x0A00 /DNTDDI_VERSION=0x0A000000" + cmake_c_flags = ( + "/EHsc /MP /guard:cf /DWIN32 /D_WINDOWS /DWINAPI_FAMILY=100 /DWINVER=0x0A00 /D_WIN32_WINNT=0x0A00 /DNTDDI_VERSION=0x0A000000" + + spectre_flag + ) if args.config == "Release": cmake_c_flags += " /O2 /Ob2 /DNDEBUG" elif args.config == "RelWithDebInfo": @@ -575,7 +595,9 @@ def _get_windows_build_args(args: argparse.Namespace): ] if args.use_cuda: win_args += [ - '-DCMAKE_CUDA_FLAGS_INIT=/DWIN32 /D_WINDOWS /DWINAPI_FAMILY=100 /DWINVER=0x0A00 /D_WIN32_WINNT=0x0A00 /DNTDDI_VERSION=0x0A000000 -Xcompiler=" /MP /guard:cf /Qspectre " -allow-unsupported-compiler', + '-DCMAKE_CUDA_FLAGS_INIT=/DWIN32 /D_WINDOWS /DWINAPI_FAMILY=100 /DWINVER=0x0A00 /D_WIN32_WINNT=0x0A00 /DNTDDI_VERSION=0x0A000000 -Xcompiler=" /MP /guard:cf' + + spectre_flag + + ' " -allow-unsupported-compiler', ] return win_args @@ -910,7 +932,7 @@ def _build_sdk_cmake(args: argparse.Namespace, env: dict[str, str]): util.run(build_command, env=env) if args.sdk == "python" and not args.skip_wheel: - util.run(build_command + ["--target", "PyPackageBuild"], env=env) + util.run([*build_command, "--target", "PyPackageBuild"], env=env) def _build_sdk_csharp(args: argparse.Namespace, env: dict[str, str]): @@ -1003,14 +1025,20 @@ def build_examples(args: argparse.Namespace, env: dict[str, str]): if build_dir.exists(): log.info(f"Removing existing build directory: {build_dir}") - shutil.rmtree(build_dir) + + def remove_readonly(function, path, _): + os.chmod(path, stat.S_IWRITE) + function(path) + + shutil.rmtree(build_dir, onerror=remove_readonly) build_dir.mkdir() samples_to_build = ["-DMODEL_QA=ON", "-DMODEL_CHAT=ON", "-DMODEL_MM=ON", "-DWHISPER=ON", "-DMODEL_ASR=ON"] - ort_include_dir = REPO_ROOT / "ort" / "include" - ort_lib_dir = REPO_ROOT / "ort" / "lib" + ort_root = args.ort_home if args.ort_home else REPO_ROOT / "ort" + ort_include_dir = ort_root / "include" + ort_lib_dir = ort_root / "lib" oga_include_dir = REPO_ROOT / "src" oga_lib_dir = args.build_dir if util.is_windows(): diff --git a/cmake/msvc_security.cmake b/cmake/msvc_security.cmake index abb91cf19b..fb90282c48 100644 --- a/cmake/msvc_security.cmake +++ b/cmake/msvc_security.cmake @@ -80,7 +80,13 @@ unset(_msec_plat_lc) # --- /Qspectre Spectre variant 1 mitigations --- # Generator expressions restrict this to C/C++ only — CUDA nvcc does not # accept /Qspectre and will error if it is passed to it. -add_compile_options( - "$<$:/Qspectre>" - "$<$:/Qspectre>" -) +if(NOT DEFINED ENABLE_MSVC_SPECTRE) + # Standalone SDK builds do not include options.cmake. Keep mitigation on there too. + set(ENABLE_MSVC_SPECTRE ON) +endif() +if(ENABLE_MSVC_SPECTRE) + add_compile_options( + "$<$:/Qspectre>" + "$<$:/Qspectre>" + ) +endif() diff --git a/cmake/options.cmake b/cmake/options.cmake index 0ab48732da..9059abc999 100644 --- a/cmake/options.cmake +++ b/cmake/options.cmake @@ -6,6 +6,7 @@ option(USE_TRT_RTX "Build with TensorRT-RTX support" OFF) option(USE_DML "Build with DML support" OFF) option(USE_WINML "Build with WinML support" OFF) option(USE_GUIDANCE "Build with guidance support" OFF) +option(ENABLE_MSVC_SPECTRE "Enable MSVC Spectre variant 1 mitigations" ON) set(NUGET_PACKAGE_SOURCE "" CACHE STRING "Override the NuGet package source used by the build") set(NUGET_CONFIG_FILE "" CACHE FILEPATH "Use a custom NuGet configuration file") diff --git a/src/config.cpp b/src/config.cpp index 0a0589040f..2c5fbe65a7 100644 --- a/src/config.cpp +++ b/src/config.cpp @@ -75,6 +75,7 @@ void InheritSessionOptions(const Config::SessionOptions& parent, if (!child.inter_op_num_threads) child.inter_op_num_threads = parent.inter_op_num_threads; if (!child.enable_cpu_mem_arena) child.enable_cpu_mem_arena = parent.enable_cpu_mem_arena; if (!child.enable_mem_pattern) child.enable_mem_pattern = parent.enable_mem_pattern; + if (!child.use_deterministic_compute) child.use_deterministic_compute = parent.use_deterministic_compute; if (!child.log_id) child.log_id = parent.log_id; if (!child.log_severity_level) child.log_severity_level = parent.log_severity_level; if (!child.log_verbosity_level) child.log_verbosity_level = parent.log_verbosity_level; @@ -121,6 +122,16 @@ std::unique_ptr CreateMtpDecoderConfig(const Config& config) { } // An empty MTP list intentionally disables sharing the main decoder's initializers. decoder.shared_initializers = mtp.shared_initializers; + if (mtp.index_share.enabled || mtp.index_share.base_capacity != 0) { + if (mtp.index_share.indices_input.empty() || mtp.index_share.counts_input.empty() || + mtp.index_share.indices_output.empty() || mtp.index_share.counts_output.empty() || + mtp.index_share.status_output.empty() || + mtp.index_share.base_capacity <= 0 || mtp.index_share.max_draft_tokens < 1 || + mtp.index_share.max_draft_tokens > 7 || + mtp.index_share.base_capacity > std::numeric_limits::max() - mtp.index_share.max_draft_tokens + 1) { + throw std::runtime_error("MTP IndexShare requires complete selection metadata."); + } + } decoder.num_hidden_layers = mtp.num_hidden_layers; decoder.num_key_value_heads = mtp.num_key_value_heads; decoder.head_size = mtp.head_size; @@ -128,8 +139,12 @@ std::unique_ptr CreateMtpDecoderConfig(const Config& config) { decoder.inputs.input_ids = mtp.inputs.input_ids; decoder.inputs.hidden_states = mtp.inputs.hidden_states; decoder.inputs.position_ids = mtp.inputs.position_ids; + decoder.inputs.rotary_position_ids = mtp.inputs.rotary_position_ids; decoder.inputs.past_key_names = mtp.inputs.past_key_names; decoder.inputs.past_value_names = mtp.inputs.past_value_names; + decoder.inputs.past_indexer_names = mtp.inputs.past_indexer_names; + decoder.inputs.past_sequence_length = mtp.inputs.past_sequence_length; + decoder.inputs.logits_indices = mtp.inputs.logits_indices; // The projection starts from a copy of the target config, so a per-token quantized target would // otherwise leak its scale name templates into the head. The head is always an unquantized // full-attention layer and declares no scale tensors, so clear them: leaving them in place would @@ -141,20 +156,49 @@ std::unique_ptr CreateMtpDecoderConfig(const Config& config) { decoder.outputs.hidden_states = mtp.outputs.hidden_states; decoder.outputs.present_key_names = mtp.outputs.present_key_names; decoder.outputs.present_value_names = mtp.outputs.present_value_names; + decoder.outputs.present_indexer_names = mtp.outputs.present_indexer_names; decoder.outputs.present_key_scale_names.clear(); decoder.outputs.present_value_scale_names.clear(); - // The MTP graph is one full-attention layer. Paged-attention metadata and block-table names are + // The MTP graph is one decoder layer. Paged-attention metadata and block-table names are // deliberately inherited above; fixed recurrent state and a main-model sliding window are not. decoder.layer_types.assign(static_cast(mtp.num_hidden_layers), "full_attention"); decoder.conv_cache_size = 0; - decoder.state_update_capacity = 0; decoder.sliding_window.reset(); - decoder.state_groups.reset(); + if (!mtp.inputs.past_indexer_names.empty() || + !mtp.outputs.present_indexer_names.empty()) { + if (mtp.inputs.past_indexer_names.empty() || + mtp.outputs.present_indexer_names.empty() || + !decoder.state_groups) { + throw std::runtime_error( + "MTP indexer state requires matching input/output templates and a target fixed_indexer group."); + } + const auto indexer = std::find_if( + decoder.state_groups->begin(), decoder.state_groups->end(), + [](const Config::Model::Decoder::StateGroup& group) { + return group.kind == Config::Model::Decoder::StateGroupKind::FixedIndexer; + }); + if (indexer == decoder.state_groups->end() || !indexer->state_update) { + throw std::runtime_error( + "MTP indexer state requires the target fixed_indexer state_update contract."); + } + auto projected_indexer = *indexer; + projected_indexer.layer_ids = {0}; + // Chained head runs consume full present state. Compact replay is reserved for the target's + // verify transaction, which captures a whole speculative block in one invocation. + projected_indexer.state_update->enabled = false; + decoder.state_groups = std::vector{ + {Config::Model::Decoder::StateGroupKind::PagedKeyValue, {0}, std::nullopt}, + std::move(projected_indexer)}; + } else { + decoder.state_update_capacity = 0; + decoder.state_groups.reset(); + } decoder.pipeline.clear(); // A chained draft feeds every stage the previous stage's hidden states, so the head must emit its // own hidden states. Record that before clearing the MTP section the demand was inferred from. projected->engine.hidden_states_output_required = true; + projected->model.engram = {}; projected->model.mtp = {}; return projected; } @@ -356,6 +400,8 @@ struct SessionOptions_Element : JSON::Element { v_.enable_cpu_mem_arena = JSON::Get(value); } else if (name == "enable_mem_pattern") { v_.enable_mem_pattern = JSON::Get(value); + } else if (name == "use_deterministic_compute") { + v_.use_deterministic_compute = JSON::Get(value); } else if (name == "graph_optimization_level") { v_.graph_optimization_level = GetGraphOptimizationLevel(JSON::Get(value)); } else if (name == "custom_ops_library") { @@ -460,6 +506,10 @@ struct DecoderInputs_Element : JSON::Element { v_.input_ids = JSON::Get(value); } else if (name == "inputs_embeds") { v_.embeddings = JSON::Get(value); + } else if (name == "engram_embeddings") { + v_.engram_embeddings = JSON::Get(value); + } else if (name == "rotary_position_ids") { + v_.rotary_position_ids = JSON::Get(value); } else if (name == "attention_mask") { v_.attention_mask = JSON::Get(value); } else if (name == "position_ids") { @@ -468,6 +518,8 @@ struct DecoderInputs_Element : JSON::Element { v_.past_key_names = JSON::Get(value); } else if (name == "past_value_names") { v_.past_value_names = JSON::Get(value); + } else if (name == "past_indexer_names") { + v_.past_indexer_names = JSON::Get(value); } else if (name == "past_key_scale_names") { v_.past_key_scale_names = JSON::Get(value); } else if (name == "past_value_scale_names") { @@ -510,6 +562,14 @@ struct DecoderInputs_Element : JSON::Element { v_.past_conv_names = JSON::Get(value); } else if (name == "past_recurrent_names") { v_.past_recurrent_names = JSON::Get(value); + } else if (name == "past_ple_token_names") { + v_.past_ple_token_names = JSON::Get(value); + } else if (name == "past_ple_conv_names") { + v_.past_ple_conv_names = JSON::Get(value); + } else if (name == "past_indexer_kv_buffer_names") { + v_.past_indexer_kv_buffer_names = JSON::Get(value); + } else if (name == "past_indexer_state_lengths_names") { + v_.past_indexer_state_lengths_names = JSON::Get(value); } else if (name == "state_update_capture_count") { v_.state_update_capture_count = JSON::Get(value); } else if (name == "state_update_active") { @@ -545,6 +605,12 @@ struct DecoderOutputs_Element : JSON::Element { v_.present_key_names = JSON::Get(value); } else if (name == "present_value_names") { v_.present_value_names = JSON::Get(value); + } else if (name == "present_indexer_names") { + v_.present_indexer_names = JSON::Get(value); + } else if (name == "state_update_indexer_value_names") { + v_.state_update_indexer_value_names = JSON::Get(value); + } else if (name == "state_update_indexer_row_names") { + v_.state_update_indexer_row_names = JSON::Get(value); } else if (name == "present_key_scale_names") { v_.present_key_scale_names = JSON::Get(value); } else if (name == "present_value_scale_names") { @@ -559,10 +625,24 @@ struct DecoderOutputs_Element : JSON::Element { v_.present_conv_names = JSON::Get(value); } else if (name == "present_recurrent_names") { v_.present_recurrent_names = JSON::Get(value); + } else if (name == "present_ple_token_names") { + v_.present_ple_token_names = JSON::Get(value); + } else if (name == "present_ple_conv_names") { + v_.present_ple_conv_names = JSON::Get(value); + } else if (name == "present_indexer_kv_buffer_names") { + v_.present_indexer_kv_buffer_names = JSON::Get(value); + } else if (name == "present_indexer_state_lengths_names") { + v_.present_indexer_state_lengths_names = JSON::Get(value); } else if (name == "state_update_conv_value_names") { v_.state_update_conv_value_names = JSON::Get(value); } else if (name == "state_update_recurrent_capsule_names") { v_.state_update_recurrent_capsule_names = JSON::Get(value); + } else if (name == "state_update_ple_token_names") { + v_.state_update_ple_token_names = JSON::Get(value); + } else if (name == "state_update_ple_conv_value_names") { + v_.state_update_ple_conv_value_names = JSON::Get(value); + } else if (name == "state_update_indexer_names") { + v_.state_update_indexer_names = JSON::Get(value); } else if (name == "hidden_states") { v_.hidden_states = JSON::Get(value); } else if (name == "aux_hidden_states") { @@ -687,6 +767,13 @@ struct StateUpdate_Element : JSON::Element { throw std::runtime_error("Decoder state update key_head_count must be positive"); } v_.key_head_count = static_cast(count); + } else if (name == "compress_ratio") { + const auto ratio = SafeDoubleToInt64( + JSON::Get(value), "model.decoder.state_groups.state_update.compress_ratio"); + if (ratio < 1 || ratio > std::numeric_limits::max()) { + throw std::runtime_error("Decoder state update compress_ratio must be positive"); + } + v_.compress_ratio = static_cast(ratio); } else if (name == "kind" || name == "capture_count" || name == "value" || name == "active" || name == "capsule") { throw std::runtime_error( @@ -714,6 +801,10 @@ struct StateGroup_Element : JSON::Element { v_.kind = DecoderStateGroupKind::FixedConv; } else if (kind == "fixed_recurrent") { v_.kind = DecoderStateGroupKind::FixedRecurrent; + } else if (kind == "fixed_ple") { + v_.kind = DecoderStateGroupKind::FixedPle; + } else if (kind == "fixed_indexer") { + v_.kind = DecoderStateGroupKind::FixedIndexer; } else if (kind == "fixed") { throw std::runtime_error( "Decoder state group kind 'fixed' is no longer supported; use 'fixed_conv' or " @@ -977,6 +1068,14 @@ struct Decoder_Element : JSON::Element { std::to_string(Config::Model::Decoder::MaxStateUpdateCapacity)); } else if (name == "conv_cache_size") { v_.conv_cache_size = SafeDoubleToInt(JSON::Get(value), name); + } else if (name == "ple_token_pad_id") { + v_.ple_token_pad_id = SafeDoubleToInt64(JSON::Get(value), name); + if (v_.ple_token_pad_id < 0) + throw std::runtime_error("ple_token_pad_id must be nonnegative"); + } else if (name == "indexer_compress_ratio") { + v_.indexer_compress_ratio = SafeDoubleToInt(JSON::Get(value), name); + if (v_.indexer_compress_ratio <= 0) + throw std::runtime_error("indexer_compress_ratio must be positive"); } else { throw JSON::unknown_value_error{}; } @@ -1043,13 +1142,15 @@ struct Decoder_Element : JSON::Element { }; struct MtpInputs_Element : JSON::Element { - explicit MtpInputs_Element(Config::Model::Mtp::Inputs& v) : v_{v} {} + explicit MtpInputs_Element(Config::Model::Mtp& mtp) : v_{mtp.inputs}, index_share_{mtp.index_share} {} void OnValue(std::string_view name, JSON::Value value) override { if (name == "input_ids") { v_.input_ids = JSON::Get(value); } else if (name == "hidden_states") { v_.hidden_states = JSON::Get(value); + } else if (name == "rotary_position_ids") { + v_.rotary_position_ids = JSON::Get(value); } else if (name == "attention_mask") { v_.attention_mask = JSON::Get(value); } else if (name == "position_ids") { @@ -1058,6 +1159,16 @@ struct MtpInputs_Element : JSON::Element { v_.past_key_names = JSON::Get(value); } else if (name == "past_value_names") { v_.past_value_names = JSON::Get(value); + } else if (name == "past_indexer_names") { + v_.past_indexer_names = JSON::Get(value); + } else if (name == "past_sequence_length") { + v_.past_sequence_length = JSON::Get(value); + } else if (name == "past_indices") { + index_share_.indices_input = JSON::Get(value); + } else if (name == "past_counts") { + index_share_.counts_input = JSON::Get(value); + } else if (name == "logits_indices") { + v_.logits_indices = JSON::Get(value); } else { throw JSON::unknown_value_error{}; } @@ -1065,10 +1176,11 @@ struct MtpInputs_Element : JSON::Element { private: Config::Model::Mtp::Inputs& v_; + Config::Model::Mtp::IndexShare& index_share_; }; struct MtpOutputs_Element : JSON::Element { - explicit MtpOutputs_Element(Config::Model::Mtp::Outputs& v) : v_{v} {} + explicit MtpOutputs_Element(Config::Model::Mtp& mtp) : v_{mtp.outputs}, index_share_{mtp.index_share} {} void OnValue(std::string_view name, JSON::Value value) override { if (name == "logits") { @@ -1079,6 +1191,14 @@ struct MtpOutputs_Element : JSON::Element { v_.present_key_names = JSON::Get(value); } else if (name == "present_value_names") { v_.present_value_names = JSON::Get(value); + } else if (name == "present_indexer_names") { + v_.present_indexer_names = JSON::Get(value); + } else if (name == "present_indices") { + index_share_.indices_output = JSON::Get(value); + } else if (name == "present_counts") { + index_share_.counts_output = JSON::Get(value); + } else if (name == "indexer_status") { + index_share_.status_output = JSON::Get(value); } else { throw JSON::unknown_value_error{}; } @@ -1086,6 +1206,7 @@ struct MtpOutputs_Element : JSON::Element { private: Config::Model::Mtp::Outputs& v_; + Config::Model::Mtp::IndexShare& index_share_; }; struct Mtp_Element : JSON::Element { @@ -1107,6 +1228,11 @@ struct Mtp_Element : JSON::Element { if (v_.head_size <= 0) throw std::out_of_range("head_size must be > 0"); } else if (name == "main_hidden_states") { v_.main_hidden_states = JSON::Get(value); + } else if (name == "base_capacity") { + v_.index_share.base_capacity = SafeDoubleToInt(JSON::Get(value), name); + if (v_.index_share.base_capacity <= 0) throw std::out_of_range("base_capacity must be > 0"); + v_.index_share.enabled = true; + if (v_.index_share.max_draft_tokens == 0) v_.index_share.max_draft_tokens = 7; } else { throw JSON::unknown_value_error{}; } @@ -1143,8 +1269,8 @@ struct Mtp_Element : JSON::Element { Config::Model::Mtp& v_; std::unique_ptr session_options_; std::unique_ptr run_options_; - MtpInputs_Element inputs_{v_.inputs}; - MtpOutputs_Element outputs_{v_.outputs}; + MtpInputs_Element inputs_{v_}; + MtpOutputs_Element outputs_{v_}; SharedInitializers_Element shared_initializers_{v_.shared_initializers}; }; @@ -2137,6 +2263,81 @@ struct Embedding_Element : JSON::Element { EmbeddingOutputs_Element outputs_{v_.outputs}; }; +struct EngramInputs_Element : JSON::Element { + explicit EngramInputs_Element(Config::Model::Engram::Inputs& v) : v_{v} {} + + void OnValue(std::string_view name, JSON::Value value) override { + if (name == "input_ids") { + v_.input_ids = JSON::Get(value); + } else if (name == "past_ple_token_names" || name == "past_tokens") { + v_.past_tokens = JSON::Get(value); + } else { + throw JSON::unknown_value_error{}; + } + } + + private: + Config::Model::Engram::Inputs& v_; +}; + +struct EngramOutputs_Element : JSON::Element { + explicit EngramOutputs_Element(Config::Model::Engram::Outputs& v) : v_{v} {} + + void OnValue(std::string_view name, JSON::Value value) override { + if (name == "embeddings") { + v_.embeddings = JSON::Get(value); + } else if (name == "present_ple_token_names" || name == "present_tokens") { + v_.present_tokens = JSON::Get(value); + } else { + throw JSON::unknown_value_error{}; + } + } + + private: + Config::Model::Engram::Outputs& v_; +}; + +struct Engram_Element : JSON::Element { + explicit Engram_Element(Config::Model::Engram& v) : v_{v} {} + + void OnValue(std::string_view name, JSON::Value value) override { + if (name == "filename") { + v_.filename = JSON::Get(value); + } else if (name == "cache_capacity") { + const double capacity = JSON::Get(value); + if (capacity < 0 || capacity != std::trunc(capacity)) { + throw std::runtime_error("model.engram.cache_capacity must be a non-negative integer"); + } + v_.cache_capacity = static_cast(capacity); + } else { + throw JSON::unknown_value_error{}; + } + } + + Element& OnObject(std::string_view name) override { + if (name == "session_options") { + v_.session_options = Config::SessionOptions{}; + session_options_ = std::make_unique(*v_.session_options); + return *session_options_; + } + if (name == "run_options") { + v_.run_options = Config::RunOptions{}; + run_options_ = std::make_unique(*v_.run_options); + return *run_options_; + } + if (name == "inputs") return inputs_; + if (name == "outputs") return outputs_; + throw JSON::unknown_value_error{}; + } + + private: + Config::Model::Engram& v_; + std::unique_ptr session_options_; + std::unique_ptr run_options_; + EngramInputs_Element inputs_{v_.inputs}; + EngramOutputs_Element outputs_{v_.outputs}; +}; + struct Model_Element : JSON::Element { explicit Model_Element(Config::Model& v) : v_{v}, block_drafter_alias_{v_.dflash2.configured_alias_is_dspark} {} @@ -2247,6 +2448,9 @@ struct Model_Element : JSON::Element { if (name == "embedding") { return embedding_; } + if (name == "engram") { + return engram_; + } if (name == "speech") { return speech_; } @@ -2288,6 +2492,7 @@ struct Model_Element : JSON::Element { Int_Array_Element tdt_durations_{v_.tdt_durations}; Vision_Element vision_{v_.vision}; Embedding_Element embedding_{v_.embedding}; + Engram_Element engram_{v_.engram}; Speech_Element speech_{v_.speech}; AudioOutput_Element audio_output_{v_.audio_output}; Joiner_Element joiner_{v_.joiner}; @@ -3170,11 +3375,21 @@ void ParseConfig(const fs::path& filename, std::string_view json_overlay, Config ValidateRuntimeProfiles(config); } +void RefreshEngramProviders(Config& config) { + if (config.model.engram.session_options) { + auto& options = *config.model.engram.session_options; + options.providers.clear(); + for (const auto& provider : options.provider_options) + options.providers.push_back(provider.name); + } +} + void OverlayConfig(Config& config, std::string_view json) { Config candidate{config}; Root_Element root{candidate}; RootObject_Element element{root}; JSON::Parse(element, json); + RefreshEngramProviders(candidate); ValidateRuntimeProfiles(candidate); ModelStateManifest::ValidateConfig(candidate.model.decoder); std::swap(config, candidate); @@ -3371,6 +3586,12 @@ Config::Config(const fs::path& path, std::string_view json_overlay) : config_pat model.embedding.session_options->providers.push_back(provider_option.name); } } + RefreshEngramProviders(*this); + + if ((!model.engram.filename.empty() || !model.decoder.inputs.past_ple_token_names.empty()) && + model.decoder.ple_token_pad_id < 0) { + throw std::runtime_error("Engram/PLE models require model.decoder.ple_token_pad_id"); + } // Validate all config-specified filenames/paths after parsing so downstream loaders // (model/processor/adapter creation) can rely on them being safe. diff --git a/src/config.h b/src/config.h index fc255dea72..7571a19c21 100644 --- a/src/config.h +++ b/src/config.h @@ -133,6 +133,7 @@ struct Config { std::optional inter_op_num_threads; std::optional enable_cpu_mem_arena; std::optional enable_mem_pattern; + std::optional use_deterministic_compute; std::optional log_id; std::optional log_severity_level; std::optional log_verbosity_level; @@ -256,6 +257,23 @@ struct Config { } outputs; } embedding; + struct Engram { + std::string filename; + size_t cache_capacity{4096}; + std::optional session_options; + std::optional run_options; + + struct Inputs { + std::string input_ids{Defaults::InputIdsName}; + std::string past_tokens{"past_ple_tokens"}; + } inputs; + + struct Outputs { + std::string embeddings{"engram_embeddings"}; + std::string present_tokens{"present_ple_tokens"}; + } outputs; + } engram; + struct Vision { std::string filename; std::optional session_options; @@ -490,6 +508,8 @@ struct Config { // Hybrid SSM+Attention (LFM2) parameters std::vector layer_types; // Per-layer type: "conv" or "full_attention" int conv_cache_size{}; // Convolution cache width (conv_L_cache from HF config) + int64_t ple_token_pad_id{-1}; + int indexer_compress_ratio{}; // Zero means unspecified by an older export. struct SlidingWindow { // Sliding window parameters for models that process input prompt in chunks int window_size{}; // The size of the window to slide over the input prompt @@ -512,12 +532,16 @@ struct Config { PagedKeyValue, FixedConv, FixedRecurrent, + FixedPle, + FixedIndexer, }; enum class StateUpdateKind { Invalid, CausalConv, GatedDeltaNet, + Ple, + Indexer, }; static constexpr int MaxStateUpdateCapacity = 8; @@ -526,6 +550,7 @@ struct Config { int capacity{}; bool enabled{true}; int key_head_count{}; + int compress_ratio{}; }; struct StateGroup { @@ -540,8 +565,10 @@ struct Config { struct Inputs { std::string input_ids{Defaults::InputIdsName}; std::string embeddings{Defaults::InputsEmbedsName}; + std::string engram_embeddings; std::string attention_mask{Defaults::AttentionMaskName}; std::string position_ids{Defaults::PositionIdsName}; + std::string rotary_position_ids; std::string past_key_names{Defaults::PastKeyName}; std::string past_value_names{Defaults::PastValueName}; std::string past_key_scale_names; @@ -568,6 +595,11 @@ struct Config { std::string logits_indices; std::string past_conv_names{Defaults::PastConvName}; // Conv cache input name template (LFM2) std::string past_recurrent_names{Defaults::PastRecurrentName}; + std::string past_ple_token_names; + std::string past_ple_conv_names; + std::string past_indexer_names; + std::string past_indexer_kv_buffer_names; + std::string past_indexer_state_lengths_names; std::string state_update_capture_count{Defaults::StateUpdateCaptureCountName}; // Per-sequence capture count std::string state_update_active{Defaults::StateUpdateActiveName}; // Capture enable flag @@ -598,8 +630,18 @@ struct Config { std::string rnn_states{Defaults::RnnStatesName}; std::string present_conv_names{Defaults::PresentConvName}; // Conv cache output name template (LFM2) std::string present_recurrent_names{Defaults::PresentRecurrentName}; + std::string present_ple_token_names; + std::string present_ple_conv_names; + std::string present_indexer_names; + std::string present_indexer_kv_buffer_names; + std::string present_indexer_state_lengths_names; + std::string state_update_indexer_value_names; + std::string state_update_indexer_row_names; std::string state_update_conv_value_names{Defaults::StateUpdateConvValueName}; std::string state_update_recurrent_capsule_names{Defaults::StateUpdateRecurrentCapsuleName}; + std::string state_update_ple_token_names; + std::string state_update_ple_conv_value_names; + std::string state_update_indexer_names; std::string hidden_states; // Last hidden state output (when exported with include_hidden_states; e.g. fed to the MTP head) // Residual streams tapped at model.dflash2.aux_hidden_state_layers, concatenated on the // last axis. Empty unless the model was exported with aux_hidden_state_layers. @@ -644,6 +686,16 @@ struct Config { // hidden-state output and the head's feedback output. struct Mtp { bool enabled{true}; + struct IndexShare { + bool enabled{false}; + int base_capacity{}; + int max_draft_tokens{}; + std::string indices_input{"past_indices"}; + std::string counts_input{"past_counts"}; + std::string indices_output; + std::string counts_output; + std::string status_output{"indexer_status"}; + } index_share; std::string filename; // e.g. "mtp.onnx"; used by model packaging/building tools std::optional session_options; std::optional run_options; @@ -666,8 +718,12 @@ struct Config { std::string hidden_states{Defaults::HiddenStatesName}; std::string attention_mask{Defaults::AttentionMaskName}; std::string position_ids{Defaults::PositionIdsName}; + std::string rotary_position_ids; std::string past_key_names{Defaults::PastKeyName}; std::string past_value_names{Defaults::PastValueName}; + std::string past_indexer_names; + std::string past_sequence_length; + std::string logits_indices; } inputs; struct Outputs { @@ -675,6 +731,7 @@ struct Config { std::string hidden_states{"hidden_states_out"}; std::string present_key_names{Defaults::PresentKeyName}; std::string present_value_names{Defaults::PresentValueName}; + std::string present_indexer_names; } outputs; bool IsEnabled() const noexcept { return enabled && !filename.empty(); } diff --git a/src/generator/generators.cpp b/src/generator/generators.cpp index 3608babb57..c8e388b4ca 100644 --- a/src/generator/generators.cpp +++ b/src/generator/generators.cpp @@ -759,6 +759,7 @@ void Generator::AppendTokens(cpu_span input_ids) { else if (name.find("audio") != std::string::npos) used_audio = true; } + append_input_modality = (used_vision && used_audio) ? "multimodal" : used_vision ? "vision" : used_audio ? "audio" @@ -802,6 +803,9 @@ void Generator::AppendTokens(DeviceSpan input_ids) { set_extra_inputs_ = false; } + if (!model_->config_->model.engram.filename.empty()) { + input_ids.CopyDeviceToCpu(); + } search_->AppendTokens(input_ids); computed_logits_ = false; ComputeLogits(input_ids); @@ -986,6 +990,7 @@ void Generator::RewindToLength(size_t new_length) { const int64_t rewound_token_count = static_cast(current_length - new_length) * static_cast(search_->params_->BatchBeamSize()); + state_->ValidateRewindTo(new_length); search_->RewindTo(new_length); state_->RewindTo(new_length); if (guidance_logits_processor_) { @@ -1025,6 +1030,7 @@ void Generator::CropToAccepted(size_t new_length, size_t recurrent_position) { ThrowErrorIfSessionTerminated(state_->session_terminated_); if (new_length > search_->GetSequenceLength()) throw std::runtime_error("CropToAccepted: new_length exceeds current sequence length"); + state_->ValidateCropToAccepted(new_length, recurrent_position); search_->RewindTo(new_length); state_->CropToAccepted(new_length, recurrent_position); if (guidance_logits_processor_) { diff --git a/src/generator/mtp_generator.cpp b/src/generator/mtp_generator.cpp index a3d45d99c4..43902bf173 100644 --- a/src/generator/mtp_generator.cpp +++ b/src/generator/mtp_generator.cpp @@ -8,7 +8,6 @@ #include "mtp_generator.h" #include "decoding/speculative_sampling.h" -#include #include #include #include @@ -126,7 +125,17 @@ MtpGenerator::MtpGenerator(const Model& main_model, const Model& mtp_model, cons mtp_params_->use_graph_capture ? (num_speculative_tokens_ + 1) : 1; mtp_ = CreateGenerator(mtp_model_, *mtp_params_); - hidden_size_ = main_model_.config_->model.decoder.hidden_size; + const auto main_hidden_shape = main_model_.session_info_.GetOutputShape( + main_model_.config_->model.mtp.main_hidden_states); + const auto mtp_hidden_shape = mtp_model_.session_info_.GetInputShape( + mtp_model_.config_->model.decoder.inputs.hidden_states); + if (main_hidden_shape.empty() || mtp_hidden_shape.empty() || + main_hidden_shape.back() <= 0 || mtp_hidden_shape.back() <= 0 || + main_hidden_shape.back() != mtp_hidden_shape.back()) { + throw std::runtime_error( + "MtpGenerator requires matching static main-output and head-input hidden-state widths"); + } + hidden_size_ = static_cast(main_hidden_shape.back()); vocab_size_ = main_model_.config_->model.vocab_size; max_length_ = params.search.max_length; @@ -197,6 +206,8 @@ MtpGenerator::MtpGenerator(const Model& main_model, const Model& mtp_model, cons head_tokens_device_ = mtp_model_.p_device_->Allocate(num_speculative_tokens_ + 1); auto head_tokens_cpu = head_tokens_device_.CpuSpan(); std::fill(head_tokens_cpu.begin(), head_tokens_cpu.end(), head_non_pad); + head_step_token_device_ = mtp_model_.p_device_->Allocate(1); + head_step_token_device_.CpuSpan()[0] = head_non_pad; verify_tokens_device_ = main_model_.p_device_->Allocate(num_speculative_tokens_ + 1); // Logits::Update inspects the DeviceSpan's host mirror only to derive each row's unpadded @@ -246,11 +257,21 @@ void MtpGenerator::ExtractHiddenPosition(OrtValue* hidden, int position) { } void MtpGenerator::CopyHiddenRow(OrtValue* hidden, int position, Tensor& dst) { - // hidden is [1, S, H] on the main model device; copy row `position` into dst ([1,1,H]). - auto src = ByteWrapTensor(*main_model_.p_device_, *hidden); - const size_t row_bytes = dst.GetByteSpan().size(); - auto src_row = src.subspan(static_cast(position) * row_bytes, row_bytes); - dst.GetByteSpan().CopyFrom(src_row); + CopyHiddenRows(hidden, position, 1, dst.GetByteSpan()); +} + +void MtpGenerator::CopyHiddenRows( + OrtValue* hidden, int first_position, int count, DeviceSpan dst) { + const bool source_on_cpu = + hidden->GetTensorMemoryInfo().GetDeviceType() == OrtMemoryInfoDeviceType_CPU; + auto& source_device = source_on_cpu + ? *GetDeviceInterface(DeviceType::CPU) + : *main_model_.p_device_; + auto src = ByteWrapTensor(source_device, *hidden); + const size_t row_bytes = hidden_slice_->GetByteSpan().size(); + const size_t bytes = static_cast(count) * row_bytes; + dst.subspan(0, bytes).CopyFrom( + src.subspan(static_cast(first_position) * row_bytes, bytes)); } int32_t MtpGenerator::DraftHeadStep(int32_t token, bool need_draft) { @@ -258,6 +279,7 @@ int32_t MtpGenerator::DraftHeadStep(int32_t token, bool need_draft) { // `token` (the head KV grows by one), capture the head's own post-final-norm output for the next // chained step, and return the greedy draft for the token after `token`. std::array tok{token}; + mtp_->state_->SetRotaryPositionIds(main_->state_->GetRotaryPositionIds()); mtp_->AppendTokens(cpu_span(tok)); stats_.draft_forward_passes++; ++head_len_; @@ -298,6 +320,7 @@ void MtpGenerator::CaptureDraftToDevice(DeviceSpan draft) { } void MtpGenerator::DraftHeadStepToDevice(DeviceSpan token, DeviceSpan draft) { + mtp_->state_->SetRotaryPositionIds(main_->state_->GetRotaryPositionIds()); mtp_->AppendTokens(token); stats_.draft_forward_passes++; ++head_len_; @@ -345,6 +368,7 @@ int32_t MtpGenerator::DraftHeadStepMulti(const int32_t* tokens, int count) { // whose rows are MAIN-model hidden states for consecutive positions. One forward appends all // `count` tokens to the head KV; only the last row's logits are needed (the draft for the token // after tokens[count-1]) -- the earlier rows exist purely to re-materialize the accepted drafts. + mtp_->state_->SetRotaryPositionIds(main_->state_->GetRotaryPositionIds()); mtp_->AppendTokens(cpu_span(tokens, static_cast(count))); stats_.draft_forward_passes++; head_len_ += static_cast(count); @@ -363,6 +387,7 @@ int32_t MtpGenerator::DraftHeadStepMulti(const int32_t* tokens, int count) { } void MtpGenerator::DraftHeadStepMultiToDevice(DeviceSpan tokens, DeviceSpan draft) { + mtp_->state_->SetRotaryPositionIds(main_->state_->GetRotaryPositionIds()); mtp_->AppendTokens(tokens); stats_.draft_forward_passes++; head_len_ += tokens.size(); @@ -376,6 +401,7 @@ int32_t MtpGenerator::DraftHeadStepSample(int32_t token, int k) { // probs) is stored in draft_idx_[k]/draft_prob_[k] so the verify step can compute the // min(1, p(d)/q(d)) acceptance test with a cheap sparse probability lookup. std::array tok{token}; + mtp_->state_->SetRotaryPositionIds(main_->state_->GetRotaryPositionIds()); mtp_->AppendTokens(cpu_span(tok)); stats_.draft_forward_passes++; ++head_len_; @@ -399,9 +425,13 @@ int32_t MtpGenerator::DraftHeadStepSample(int32_t token, int k) { const float* MtpGenerator::MainLogitsRowsCpu(int first_row, int num_rows) { OrtValue* raw = main_->state_->GetOutput(main_model_.config_->model.decoder.outputs.logits.c_str()); - // Cast to fp32 on device (cheap kernel), then copy only the requested rows to the host. - Cast(*raw, logits_fp32_, *main_model_.p_device_, Ort::TypeToTensorType); - auto span = ByteWrapTensor(*main_model_.p_device_, *logits_fp32_); + // Cast only when needed: CPU Cast deliberately rejects identical input/output types. + OrtValue* fp32 = raw; + if (raw->GetTensorTypeAndShapeInfo()->GetElementType() != Ort::TypeToTensorType) { + Cast(*raw, logits_fp32_, *main_model_.p_device_, Ort::TypeToTensorType); + fp32 = logits_fp32_.get(); + } + auto span = ByteWrapTensor(*main_model_.p_device_, *fp32); const size_t row_bytes = static_cast(vocab_size_) * sizeof(float); auto rows = span.subspan(static_cast(first_row) * row_bytes, static_cast(num_rows) * row_bytes); @@ -542,13 +572,10 @@ void MtpGenerator::ArgmaxMainRows(int first_row, int num_rows, int32_t* out) { if (main_model_.p_device_->ArgMax(row_ptr, type, num_rows, vocab_size_, out)) return; - // Host fallback (e.g. CPU device): cast the logits to fp32, copy to the host, argmax each row. - Cast(*raw, logits_fp32_, *main_model_.p_device_, Ort::TypeToTensorType); - auto span = ByteWrapTensor(*main_model_.p_device_, *logits_fp32_); - auto cpu = span.CopyDeviceToCpu(); - const float* data = reinterpret_cast(cpu.data()); + // Host fallback (e.g. CPU device): copy the requested fp32 rows, casting only if needed. + const float* data = MainLogitsRowsCpu(first_row, num_rows); for (int r = 0; r < num_rows; ++r) - out[r] = ArgmaxRow(data + static_cast(first_row + r) * vocab_size_, vocab_size_); + out[r] = ArgmaxRow(data + static_cast(r) * vocab_size_, vocab_size_); } int32_t MtpGenerator::DraftNextToken(OrtValue* /*unused*/, int32_t token, bool need_draft) { @@ -556,6 +583,7 @@ int32_t MtpGenerator::DraftNextToken(OrtValue* /*unused*/, int32_t token, bool n // MTP head; its KV cache accumulates, so this is an O(1) incremental draft step. mtp_->SetHiddenStates(hidden_slice_); std::array tok{token}; + mtp_->state_->SetRotaryPositionIds(main_->state_->GetRotaryPositionIds()); mtp_->AppendTokens(cpu_span(tok)); stats_.draft_forward_passes++; if (!need_draft) { @@ -573,16 +601,14 @@ int32_t MtpGenerator::DraftNextToken(OrtValue* /*unused*/, int32_t token, bool n int32_t MtpGenerator::DraftTwo(OrtValue* hidden, int32_t tok0, int32_t tok1) { // Populate the [1,2,H] hidden buffer: row 0 = hidden@position L (pairs with tok0), row 1 = // hidden@position L+1 (pairs with tok1). `hidden` is the main model's [1,S,H] verify output. - auto src = ByteWrapTensor(*main_model_.p_device_, *hidden); - const size_t row_bytes = hidden_slice_->GetByteSpan().size(); // bytes of one [1,1,H] row auto dst = hidden_slice2_->GetByteSpan(); - dst.subspan(0, row_bytes).CopyFrom(src.subspan(0, row_bytes)); // row 0 <- hidden@0 - dst.subspan(row_bytes, row_bytes).CopyFrom(src.subspan(row_bytes, row_bytes)); // row 1 <- hidden@1 + CopyHiddenRows(hidden, 0, 2, dst); // One 2-token MTP forward: feeds tok0 (KV-advance) and tok1 (the next committed token); the // last-position logits give the draft for the token after tok1. mtp_->SetHiddenStates(hidden_slice2_); std::array toks{tok0, tok1}; + mtp_->state_->SetRotaryPositionIds(main_->state_->GetRotaryPositionIds()); mtp_->AppendTokens(cpu_span(toks)); stats_.draft_forward_passes++; auto logits_span = mtp_->GetLogits(); // fp32, last token, [1, V] @@ -768,15 +794,16 @@ void MtpGenerator::GenerateStepSingle(int32_t t) { has_pending_draft_ = true; length_ += 2; } else { - // 2b. Reject: roll back the speculative forward (restore recurrent state + crop KV to L), - // then re-run the single correct token t. The pipelined draft (if any) is invalid. + // 2b. Reject. Restore the pre-verify state and decode t alone, even for windowed targets: + // wide-verify row 0 can differ numerically from a single-token decode. has_pending_draft_ = false; + stats_.correction_tokens++; main_->RewindToLength(length_); std::array rerun{t}; main_->AppendTokens(cpu_span(rerun)); stats_.target_forward_passes++; - stats_.correction_tokens++; - OrtValue* hidden = main_->state_->GetOutput(main_model_.config_->model.mtp.main_hidden_states.c_str()); + OrtValue* hidden = main_->state_->GetOutput( + main_model_.config_->model.mtp.main_hidden_states.c_str()); ArgmaxMainRows(0, 1, &next_token_); ExtractHiddenPosition(hidden, 0); length_ += 1; @@ -827,14 +854,20 @@ void MtpGenerator::GenerateStepMulti(int32_t t) { current_token_device = head_tokens_device_.subspan(0, 1); current_token_device.CopyFrom(head_tokens_upload_.subspan(0, 1)); mtp_->SetHiddenStates(hidden_slice_); - DraftHeadStepToDevice(current_token_device, drafts_device_.subspan(0, 1)); + head_step_token_device_.CopyFrom(current_token_device); + DraftHeadStepToDevice( + head_step_token_device_, drafts_device_.subspan(0, 1)); } pending_refeed_count_ = -1; head_start = head_len_ - 1; // head KV length before t was appended + if (N > 1) mtp_->SnapshotState(); for (int k = 1; k < N; ++k) { mtp_->SetHiddenStates(head_out_hidden_); - DraftHeadStepToDevice(drafts_device_.subspan(static_cast(k - 1), 1), - drafts_device_.subspan(static_cast(k), 1)); + head_step_token_device_.CopyFrom( + drafts_device_.subspan(static_cast(k - 1), 1)); + DraftHeadStepToDevice( + head_step_token_device_, + drafts_device_.subspan(static_cast(k), 1)); } } else { if (pending_refeed_count_ > 0) { @@ -854,6 +887,7 @@ void MtpGenerator::GenerateStepMulti(int32_t t) { } pending_refeed_count_ = -1; head_start = head_len_ - 1; + if (N > 1) mtp_->SnapshotState(); for (int k = 1; k < N; ++k) { mtp_->SetHiddenStates(head_out_hidden_); drafts_[k] = DraftHeadStep(drafts_[k - 1]); @@ -866,12 +900,10 @@ void MtpGenerator::GenerateStepMulti(int32_t t) { // Flash attention), so a greedy argmax can occasionally differ on near-ties. This is the same // tradeoff the N=1 verify already makes. The fallback reject path re-runs the committed prefix // from its snapshot; a windowed recurrent state crops and commits directly from this verify. --- - // Snapshot the recurrent state so a rejected wide verify can replay the committed prefix with - // decode-consistent numerics. A model exported with a recurrent-state window commits directly out - // of the verify below and needs neither the replay nor the snapshot, which would otherwise cost - // 2*num_layers device copies on EVERY step. + // Legacy per-token windows commit without a baseline copy. Replay fallback and compact + // transition reconstruction still need the pre-verify state. const bool direct_commit = main_->CanCropRecurrentState(); - if (!direct_commit) main_->SnapshotState(); + if (!direct_commit || main_->state_->CropRequiresSnapshot()) main_->SnapshotState(); const bool device_verify_input = verify_tokens_device_.size() >= static_cast(N + 1); if (device_verify_input) { @@ -940,11 +972,7 @@ void MtpGenerator::GenerateStepMulti(int32_t t) { // Rows 0..a-1 are contiguous in both the verify output and the merged buffer: one D2D copy. // Row a is filled in by the finalize phase below from hidden_slice_. Tensor& hbuf = *refeed_multi_[a + 1]; - const size_t row_bytes = refeed_hidden_->GetByteSpan().size(); - auto src = ByteWrapTensor(*main_model_.p_device_, *vhidden); - hbuf.GetByteSpan() - .subspan(0, static_cast(a) * row_bytes) - .CopyFrom(src.subspan(0, static_cast(a) * row_bytes)); + CopyHiddenRows(vhidden, 0, a, hbuf.GetByteSpan()); for (int k = 0; k < a; ++k) merged_tokens_[k] = drafts_[k]; } @@ -1007,20 +1035,17 @@ void MtpGenerator::GenerateStepMultiSample(int32_t t) { const size_t head_start = head_len_; mtp_->SetHiddenStates(hidden_slice_); drafts_[0] = DraftHeadStepSample(t, 0); + if (N > 1) mtp_->SnapshotState(); for (int k = 1; k < N; ++k) { mtp_->SetHiddenStates(head_out_hidden_); drafts_[k] = DraftHeadStepSample(drafts_[k - 1], k); } // --- Verify [t, d0..d_{N-1}] in a single batched main forward. --- - // Snapshot the recurrent state ONLY for the fallback re-run rollback (models WITHOUT - // state_window). Snapshot() copies every linear-attn layer's conv+recurrent - // state (2*num_layers D2D copies + launches) on EVERY step. When the crop fast-path is - // available the reject branch uses the state window via CropToPosition and NEVER rewinds - // the recurrent state (RewindTo/RestoreSnapshot is unreachable), so the snapshot is dead - // overhead -- skip it. The predicate matches the reject-path crop-vs-fallback choice below, - // so the fallback branch still has its snapshot when it needs one. - if (!main_->CanCropRecurrentState()) main_->SnapshotState(); + // Compact crop reconstruction needs a baseline even though it supports direct commit; + // legacy per-token windows do not. + const bool direct_commit = main_->CanCropRecurrentState(); + if (!direct_commit || main_->state_->CropRequiresSnapshot()) main_->SnapshotState(); verify_tokens_[0] = t; for (int k = 0; k < N; ++k) verify_tokens_[k + 1] = drafts_[k]; main_->AppendTokens(cpu_span(verify_tokens_.data(), N + 1)); @@ -1100,13 +1125,9 @@ void MtpGenerator::GenerateStepMultiSample(int32_t t) { head_len_ = head_start + 1; if (a > 0) { Tensor& hbuf = *refeed_multi_[a]; - const size_t row_bytes = refeed_hidden_->GetByteSpan().size(); - auto dst = hbuf.GetByteSpan(); - auto src = ByteWrapTensor(*main_model_.p_device_, *vhidden); - for (int k = 0; k < a; ++k) - dst.subspan(static_cast(k) * row_bytes, row_bytes) - .CopyFrom(src.subspan(static_cast(k) * row_bytes, row_bytes)); + CopyHiddenRows(vhidden, 0, a, hbuf.GetByteSpan()); mtp_->SetHiddenStates(refeed_multi_[a]); + mtp_->state_->SetRotaryPositionIds(main_->state_->GetRotaryPositionIds()); mtp_->AppendTokens(cpu_span(drafts_.data(), a)); // d0..d_{a-1} with main hiddens stats_.draft_forward_passes++; head_len_ = head_start + 1 + static_cast(a); @@ -1119,7 +1140,7 @@ void MtpGenerator::GenerateStepMultiSample(int32_t t) { next_token_ = SampleSparseToken(target_idx_[N], target_prob_[N], rng_); CopyHiddenRow(vhidden, N, *hidden_slice_); // hidden that predicted the bonus token length_ += static_cast(N) + 1; - } else if (main_->CanCropRecurrentState()) { + } else if (direct_commit) { // Rejected at position a, LOSSLESS-CROP fast path (model exported with state_window): // skip the full re-run main forward (~25% of N=3 step time). The batched verify already advanced // the recurrent state through every token (window slot for position a = state AFTER the committed diff --git a/src/generator/mtp_generator.h b/src/generator/mtp_generator.h index d3996c4c71..146918eb84 100644 --- a/src/generator/mtp_generator.h +++ b/src/generator/mtp_generator.h @@ -70,6 +70,7 @@ struct MtpGenerator { void ExtractHiddenPosition(OrtValue* hidden, int position); // Copy one [1,1,H] row out of a [1,S,H] hidden OrtValue (on `main_model_`'s device) into `dst`. void CopyHiddenRow(OrtValue* hidden, int position, Tensor& dst); + void CopyHiddenRows(OrtValue* hidden, int first_position, int count, DeviceSpan dst); // One MTP-head forward on a single token: the head's `hidden_states` input must already be set // (via SetHiddenStates) by the caller. Appends `token` to the head KV, captures the head's own // post-final-norm output (hidden_states_out, last row) into `head_out_hidden_` for the next @@ -179,6 +180,7 @@ struct MtpGenerator { // `drafts_device_` with D2D copies. DeviceSpan head_tokens_upload_; DeviceSpan head_tokens_device_; + DeviceSpan head_step_token_device_; DeviceSpan verify_tokens_device_; std::vector verify_argmax_; // scratch: main argmax of the N+1 verify rows // CUDA scratch for the verify top-1 ids. ArgmaxMainRows launches the proven batch-1 device diff --git a/src/models/cpu_embedding.h b/src/models/cpu_embedding.h index b4389be522..b741f4e0dc 100644 --- a/src/models/cpu_embedding.h +++ b/src/models/cpu_embedding.h @@ -16,6 +16,8 @@ struct CpuEmbedding { ~Workspace(); DeviceSpan Prepare(Tensor& output); void Upload(); + std::vector engram_input_ids; + std::vector engram_past_tokens; private: void Wait(); diff --git a/src/models/decoder_only.cpp b/src/models/decoder_only.cpp index a014c8ea92..887879942b 100644 --- a/src/models/decoder_only.cpp +++ b/src/models/decoder_only.cpp @@ -1,11 +1,20 @@ #include "generator/generators.h" #include "decoder_only.h" +#include "engram.h" namespace Generators { DecoderOnly_Model::DecoderOnly_Model(std::unique_ptr config, OrtEnv& ort_env) : Model{std::move(config)} { session_decoder_ = CreateSession(ort_env, config_->model.decoder.filename, session_options_.get()); session_info_.Add(*session_decoder_); + if (!config_->model.engram.filename.empty()) { + engram_session_options_ = OrtSessionOptions::Create(); + CreateSessionOptionsFromConfig( + config_->model.engram.session_options.value_or(config_->model.decoder.session_options), + *engram_session_options_, true, /*disable_graph_capture=*/true); + session_engram_ = CreateSession(ort_env, config_->model.engram.filename, engram_session_options_.get()); + session_info_.Add(*session_engram_); + } if (!config_->model.embedding.filename.empty()) { if (!config_->engine.dynamic_batching) { throw std::runtime_error("Decoder-only CPU embedding requires the dynamic-batching Engine."); @@ -22,11 +31,56 @@ std::unique_ptr DecoderOnly_Model::CreateState(DeviceSpan sequen return std::make_unique(*this, sequence_lengths_unk, params); } +void DecoderOnly_Model::InitializeIndexShare(const Config::Model::Mtp::IndexShare& config, OrtEnv&) { + if (config.base_capacity == 0) return; + SessionInfo extend_info; + extend_info.Add(*session_decoder_); + SessionInfo decode_info; + decode_info.Add(*session_decoder_); + const auto indices_shape = extend_info.GetOutputShape(config.indices_output); + const auto counts_shape = extend_info.GetOutputShape(config.counts_output); + const int64_t output_capacity = indices_shape.size() == 2 ? indices_shape[1] : 0; + if (indices_shape.size() != 2 || output_capacity < config.base_capacity || + output_capacity > static_cast(config.base_capacity) + config.max_draft_tokens - 1 || counts_shape.size() != 1 || + extend_info.GetOutputDataType(config.indices_output) != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 || + extend_info.GetOutputDataType(config.counts_output) != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 || + decode_info.GetOutputDataType(config.status_output) != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 || + decode_info.GetOutputShape(config.status_output).size() != 1) { + throw std::runtime_error("IndexShare graph outputs do not match their configured int32 selection contract."); + } + const std::vector selection_inputs{ + config.indices_input.c_str(), config.counts_input.c_str(), "base_row_indices", + "range_starts", "range_ends"}; + for (const char* name : selection_inputs) { + if (!decode_info.HasInput(name) || decode_info.GetInputDataType(name) != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32) { + throw std::runtime_error("IndexShare decode graph is missing a configured int32 input."); + } + const auto shape = decode_info.GetInputShape(name); + const bool indices = std::string_view{name} == config.indices_input; + if (shape.size() != (indices ? 2u : 1u) || (indices && shape[1] != config.base_capacity)) { + throw std::runtime_error("IndexShare decode input shape does not match the configured selection contract."); + } + } + { + if (!decode_info.HasInput("mode") || !decode_info.HasInput("projection_rows") || + decode_info.GetInputDataType("mode") != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 || + decode_info.GetInputShape("mode") != std::vector{1} || + decode_info.GetInputDataType("projection_rows") != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 || + decode_info.GetInputShape("projection_rows").size() != 1) { + throw std::runtime_error("IndexShare requires an int32[1] mode and int64 projection rows."); + } + } + index_share_config_ = config; + index_share_config_.max_draft_tokens = static_cast(output_capacity - config.base_capacity + 1); +} + DecoderOnly_State::DecoderOnly_State(const DecoderOnly_Model& model, DeviceSpan sequence_lengths_unk, const GeneratorParams& params) : State{params, model}, model_{model}, kv_cache_(model_.p_device_kvcache_->CreateKeyValueCache(*this)), recurrent_state_(CreateRecurrentState(*this, /*graph_capture_variants_supported=*/true)), + ple_state_(CreatePleState(*this)), + indexer_cache_(CreateIndexerCache(*this)), position_inputs_{model_.p_device_inputs_->CreatePositionInputs(*this, sequence_lengths_unk, model_.config_->model.decoder.inputs.attention_mask)} { input_ids_.Add(); position_inputs_->Add(); @@ -35,6 +89,19 @@ DecoderOnly_State::DecoderOnly_State(const DecoderOnly_Model& model, DeviceSpan< kv_cache_->Add(); if (recurrent_state_) recurrent_state_->Add(); + if (ple_state_) + ple_state_->Add(); + if (indexer_cache_) + indexer_cache_->Add(); + if (model_.session_engram_) { + engram_state_ = std::make_unique(model_, *model_.session_engram_, params); + const auto& name = model_.config_->model.decoder.inputs.engram_embeddings; + const auto shape = model_.session_info_.GetInputShape(name); + if (shape.size() != 3 || shape.back() <= 0) + throw std::runtime_error("Generator Engram embeddings require [batch, sequence, static_width]"); + engram_embeddings_ = std::make_unique(*this, Embeddings::Mode::Input, name, shape.back()); + engram_embeddings_->Add(); + } // Models with a hidden_states input (e.g. the MTP self-speculative head) feed the main // model's last hidden state. Only created when the config declares the input. if (!model_.config_->model.decoder.inputs.hidden_states.empty()) { @@ -49,6 +116,8 @@ DecoderOnly_State::DecoderOnly_State(const DecoderOnly_Model& model, DeviceSpan< } } +DecoderOnly_State::~DecoderOnly_State() = default; + void DecoderOnly_State::SetExtraInputs(const std::vector& extra_inputs) { extra_inputs_.Add(extra_inputs, model_.session_decoder_->GetInputNames()); } @@ -67,6 +136,12 @@ DeviceSpan DecoderOnly_State::Run(int total_length, DeviceSpan& } UpdateInputsOutputs(next_tokens, next_indices, total_length); + if (engram_state_) { + engram_state_->UpdateInputsOutputs(next_tokens); + engram_state_->Run(total_length, next_tokens, next_indices); + engram_embeddings_->UpdateSequenceLength(num_tokens / params_->BatchBeamSize()); + engram_state_->CopyEmbeddingsTo(*engram_embeddings_); + } if (model_.config_->model.decoder.run_options.has_value()) { State::SetRunOptions(model_.config_->model.decoder.run_options.value()); } @@ -76,18 +151,30 @@ DeviceSpan DecoderOnly_State::Run(int total_length, DeviceSpan& // each captured length getting its own annotation id / static buffers. const int seq_len = static_cast(input_ids_.GetShape()[1]); const bool graph_capture_this_run = - params_->use_graph_capture && seq_len >= 1 && seq_len <= params_->max_graph_capture_length; - const int graph_capture_variant = recurrent_state_ ? recurrent_state_->GraphCaptureVariant() : 0; + params_->use_graph_capture && seq_len >= 1 && seq_len <= params_->max_graph_capture_length && + (!indexer_cache_ || !indexer_cache_->UsesHostSequenceLength()); + int graph_capture_variant = recurrent_state_ ? recurrent_state_->GraphCaptureVariant() : 0; + if (ple_state_) { + const int ple_graph_capture_variant = ple_state_->GraphCaptureVariant(); + if (recurrent_state_ && recurrent_state_->UsesGraphCaptureDoubleBuffer() && + graph_capture_variant != ple_graph_capture_variant) + throw std::runtime_error("PLE and recurrent state graph-capture buffer variants are out of sync"); + graph_capture_variant = ple_graph_capture_variant; + } // ORT captures by re-running the model inside this one Run(), which over-applies an // in-place recurrent state. Let the capture happen, then undo it and replay. - const int graph_id = seq_len * 2 + graph_capture_variant; + if (recurrent_state_ && recurrent_state_->StateUpdateCaptureEnabled()) + graph_capture_variant += 2; + const int graph_id = seq_len * 4 + graph_capture_variant; if (graph_capture_this_run && recurrent_state_ && recurrent_state_->ShouldFixUpGraphCapture(graph_id)) { recurrent_state_->SaveForGraphCapture(); State::Run(*model_.session_decoder_, true, seq_len, graph_capture_variant); recurrent_state_->RestoreAfterGraphCapture(graph_id); } State::Run(*model_.session_decoder_, graph_capture_this_run, seq_len, graph_capture_variant); + if (recurrent_state_) recurrent_state_->SetForwardLength(seq_len); + if (ple_state_) ple_state_->SetForwardLength(seq_len); return logits_.Get(); } @@ -111,6 +198,12 @@ DeviceSpan DecoderOnly_State::RunWithChunking(int total_length, DeviceSpa // Process this chunk - fills KV cache progressively UpdateInputsOutputs(chunk_tokens, next_indices, length); + if (engram_state_) { + engram_state_->UpdateInputsOutputs(chunk_tokens); + engram_state_->Run(length, chunk_tokens, next_indices); + engram_embeddings_->UpdateSequenceLength(current_chunk_size / params_->BatchBeamSize()); + engram_state_->CopyEmbeddingsTo(*engram_embeddings_); + } // Graph capture is typically disabled during context phase chunking bool graph_capture_this_run = false; // Disable graph capture during chunking @@ -123,28 +216,56 @@ DeviceSpan DecoderOnly_State::RunWithChunking(int total_length, DeviceSpa return logits_.Get(); } +void DecoderOnly_State::ValidateRewindTo(size_t index) const { + if (indexer_cache_) indexer_cache_->ValidateRewindTo(index); + if (recurrent_state_) recurrent_state_->ValidateRewindTo(index); + if (ple_state_) ple_state_->ValidateRewindTo(index); + if (engram_state_) engram_state_->ValidateRewindTo(index); +} + void DecoderOnly_State::RewindTo(size_t index) { + ValidateRewindTo(index); + if (engram_state_) engram_state_->RewindTo(index); position_inputs_->RewindTo(index); if (kv_cache_) kv_cache_->RewindTo(index); if (recurrent_state_) recurrent_state_->RewindTo(index); + if (ple_state_) + ple_state_->RewindTo(index); + if (indexer_cache_) + indexer_cache_->RewindTo(index); } void DecoderOnly_State::SnapshotState(size_t position) { + if (engram_state_) engram_state_->SnapshotState(position); if (recurrent_state_) recurrent_state_->Snapshot(position); + if (ple_state_) + ple_state_->Snapshot(position); + if (indexer_cache_) + indexer_cache_->Snapshot(position); } bool DecoderOnly_State::HasCroppableRecurrentState() const { - return recurrent_state_ && recurrent_state_->IsWindowed(); + return recurrent_state_ && recurrent_state_->IsWindowed() && !ple_state_ && !engram_state_ && + (!indexer_cache_ || indexer_cache_->HasStateUpdates()); } int64_t DecoderOnly_State::RecurrentStateWindow() const { return recurrent_state_ ? recurrent_state_->StateWindow() : 1; } +void DecoderOnly_State::ValidateCropToAccepted(size_t new_length, size_t recurrent_position) const { + (void)new_length; + if (recurrent_state_) recurrent_state_->ValidateCropPosition(recurrent_position); + if (ple_state_) ple_state_->ValidateCropPosition(recurrent_position); + if (engram_state_) engram_state_->ValidateAcceptedPrefix(recurrent_position + 1); + if (indexer_cache_) indexer_cache_->ValidateAcceptedPrefix(recurrent_position + 1); +} + void DecoderOnly_State::CropToAccepted(size_t new_length, size_t recurrent_position) { + ValidateCropToAccepted(new_length, recurrent_position); // Roll the attention KV cache + position back to the committed length, and crop the recurrent // state to the state AFTER verify token `recurrent_position` (no replay forward). Used by // lossless multi-token MTP partial accept. @@ -153,6 +274,10 @@ void DecoderOnly_State::CropToAccepted(size_t new_length, size_t recurrent_posit kv_cache_->RewindTo(new_length); if (recurrent_state_) recurrent_state_->CropToPosition(recurrent_position); + if (ple_state_) + ple_state_->CropToPosition(recurrent_position); + if (indexer_cache_) + indexer_cache_->CommitAcceptedPrefix(recurrent_position + 1); } void DecoderOnly_State::UpdateInputsOutputs(DeviceSpan& next_tokens, DeviceSpan beam_indices, int total_length) { @@ -183,6 +308,10 @@ void DecoderOnly_State::UpdateInputsOutputs(DeviceSpan& next_tokens, De kv_cache_->Update(beam_indices, kv_cache_length); if (recurrent_state_) recurrent_state_->Update(); + if (ple_state_) + ple_state_->Update(); + if (indexer_cache_) + indexer_cache_->Update(beam_indices, total_length, static_cast(new_length)); if (hidden_states_) hidden_states_->Update(static_cast(new_length)); if (hidden_states_output_) diff --git a/src/models/decoder_only.h b/src/models/decoder_only.h index 0b3597a743..9be587806a 100644 --- a/src/models/decoder_only.h +++ b/src/models/decoder_only.h @@ -1,43 +1,65 @@ #pragma once +#include #include "model.h" #include "cpu_embedding.h" #include "models/io/input_ids.h" +#include "models/io/embeddings.h" #include "models/io/logits.h" #include "io/kv_cache.h" #include "models/io/position_inputs.h" #include "models/io/extra_inputs.h" #include "models/io/hidden_states.h" #include "models/io/recurrent_state.h" +#include "models/io/ple_state.h" +#include "models/io/indexer_cache.h" namespace Generators { +struct EngramState; +struct CpuEngram; struct DecoderOnly_Model : Model { DecoderOnly_Model(std::unique_ptr config, OrtEnv& ort_env); std::unique_ptr CreateState(DeviceSpan sequence_lengths_unk, const GeneratorParams& params) const override; + void InitializeIndexShare(const Config::Model::Mtp::IndexShare& config, OrtEnv& ort_env); + std::unique_ptr session_decoder_; + Config::Model::Mtp::IndexShare index_share_config_; + std::unique_ptr session_engram_; + std::unique_ptr engram_session_options_; std::shared_ptr cpu_embedding_; + std::shared_ptr cpu_engram_; + std::once_flag cpu_engram_init_; }; struct DecoderOnly_State : State { DecoderOnly_State(const DecoderOnly_Model& model, DeviceSpan sequence_lengths_unk, const GeneratorParams& params); + ~DecoderOnly_State() override; void SetExtraInputs(const std::vector& extra_inputs) override; DeviceSpan Run(int total_length, DeviceSpan& next_tokens, DeviceSpan next_indices) override; void RewindTo(size_t index) override; + void ValidateRewindTo(size_t index) const override; void SnapshotState(size_t position) override; bool HasCroppableRecurrentState() const override; + bool CropRequiresSnapshot() const override { + return (recurrent_state_ && recurrent_state_->CropRequiresSnapshot()) || + (indexer_cache_ && indexer_cache_->HasStateUpdates()); + } int64_t RecurrentStateWindow() const override; void CropToAccepted(size_t new_length, size_t recurrent_position) override; + void ValidateCropToAccepted(size_t new_length, size_t recurrent_position) const override; // Stage the hidden_states values for the next Run (for models with a hidden_states input, // e.g. the MTP self-speculative head). No-op if the model has no hidden_states input. void SetHiddenStates(OrtValue* hidden_states) override; + OrtValue* GetRotaryPositionIds() override { return position_inputs_->GetRotaryPositionIds(); } + void SetRotaryPositionIds(OrtValue* positions) override { position_inputs_->SetRotaryPositionIds(positions); } private: DeviceSpan RunWithChunking(int total_length, DeviceSpan& next_tokens, @@ -51,9 +73,13 @@ struct DecoderOnly_State : State { Logits logits_{*this}; std::unique_ptr kv_cache_; std::unique_ptr recurrent_state_; + std::unique_ptr ple_state_; + std::unique_ptr indexer_cache_; std::unique_ptr position_inputs_; std::unique_ptr hidden_states_; // Only for models with a hidden_states input (MTP head). std::unique_ptr hidden_states_output_; // Only for models that emit a hidden_states output (CUDA-graph-safe). + std::unique_ptr engram_state_; + std::unique_ptr engram_embeddings_; ExtraInputs extra_inputs_{*this}; }; diff --git a/src/models/engram.cpp b/src/models/engram.cpp new file mode 100644 index 0000000000..a9c9de651e --- /dev/null +++ b/src/models/engram.cpp @@ -0,0 +1,201 @@ +#include "generator/generators.h" +#include "engram.h" +#include "utils.h" + +namespace Generators { + +EngramState::EngramState(const Model& model, OrtSession& session, const GeneratorParams& params) + : State{params, model}, model_{model}, session_{session} { + if (model_.config_->model.decoder.ple_token_pad_id < 0) + throw std::runtime_error("EngramState requires model.decoder.ple_token_pad_id"); + if (params.search.num_beams != 1) { + throw std::runtime_error("Engram-backed models currently support num_beams=1 only"); + } + input_ids_index_ = inputs_.size(); + inputs_.push_back(nullptr); + input_names_.push_back(model_.config_->model.engram.inputs.input_ids.c_str()); + embeddings_type_ = model_.session_info_.GetOutputDataType(model_.config_->model.engram.outputs.embeddings); + const auto embeddings_shape = model_.session_info_.GetOutputShape(model_.config_->model.engram.outputs.embeddings); + if (embeddings_shape.empty() || embeddings_shape.back() <= 0) { + throw std::runtime_error("Engram embeddings must have a static width"); + } + embeddings_width_ = embeddings_shape.back(); + embeddings_output_index_ = outputs_.size(); + outputs_.push_back(nullptr); + output_names_.push_back(model_.config_->model.engram.outputs.embeddings.c_str()); + + auto token_shape = model_.session_info_.GetInputShape(model_.config_->model.engram.inputs.past_tokens); + if (token_shape.empty()) { + throw std::runtime_error("Engram token history must have a batch dimension"); + } + if (token_shape[0] <= 0) token_shape[0] = params_->BatchBeamSize(); + for (int64_t dimension : token_shape) { + if (dimension <= 0) { + throw std::runtime_error("Engram token history must have a static shape"); + } + } + + past_tokens_ = OrtValue::CreateTensor(model_.allocator_cpu_, token_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64); + present_tokens_ = OrtValue::CreateTensor(model_.allocator_cpu_, token_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64); + InitializeTokens(*past_tokens_); + InitializeTokens(*present_tokens_); + past_input_index_ = inputs_.size(); + inputs_.push_back(past_tokens_.get()); + input_names_.push_back(model_.config_->model.engram.inputs.past_tokens.c_str()); + present_output_index_ = outputs_.size(); + outputs_.push_back(present_tokens_.get()); + output_names_.push_back(model_.config_->model.engram.outputs.present_tokens.c_str()); +} + +void EngramState::UpdateInputsOutputs(DeviceSpan& next_tokens) { + const int64_t batch_size = params_->BatchBeamSize(); + const int64_t sequence_length = static_cast(next_tokens.size()) / batch_size; + input_ids_ = OrtValue::CreateTensor(model_.allocator_cpu_, + std::array{batch_size, sequence_length}, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64); + const auto tokens_cpu = next_tokens.CpuSpan(); + std::transform(tokens_cpu.begin(), tokens_cpu.end(), input_ids_->GetTensorMutableData(), + [](int32_t token) { return static_cast(token); }); + inputs_[input_ids_index_] = input_ids_.get(); + embeddings_output_ = OrtValue::CreateTensor(model_.allocator_cpu_, + std::array{batch_size, sequence_length, embeddings_width_}, embeddings_type_); + outputs_[embeddings_output_index_] = embeddings_output_.get(); + if (!first_run_) { + std::swap(past_tokens_, present_tokens_); + inputs_[past_input_index_] = past_tokens_.get(); + outputs_[present_output_index_] = present_tokens_.get(); + } + + cache_hit_ = false; + active_cache_key_.reset(); + if (sequence_length == 1 && model_.config_->model.engram.cache_capacity > 0) { + std::vector key; + const auto past_info = past_tokens_->GetTensorTypeAndShapeInfo(); + const size_t past_count = past_info->GetElementCount(); + const auto* past_data = past_tokens_->GetTensorData(); + key.assign(past_data, past_data + past_count); + const auto* current_tokens = input_ids_->GetTensorData(); + key.insert(key.end(), current_tokens, current_tokens + batch_size); + auto cached = cache_.find(key); + if (cached != cache_.end()) { + const auto output_info = embeddings_output_->GetTensorTypeAndShapeInfo(); + const size_t output_bytes = output_info->GetElementCount() * Ort::SizeOf(output_info->GetElementType()); + if (cached->second.embeddings.size() == output_bytes && + cached->second.present_tokens.size() == present_tokens_->GetTensorTypeAndShapeInfo()->GetElementCount()) { + std::memcpy(embeddings_output_->GetTensorMutableRawData(), cached->second.embeddings.data(), output_bytes); + std::copy(cached->second.present_tokens.begin(), cached->second.present_tokens.end(), + present_tokens_->GetTensorMutableData()); + cache_hit_ = true; + } + } + active_cache_key_ = std::move(key); + } +} + +DeviceSpan EngramState::Run(int current_length, DeviceSpan& next_tokens, + DeviceSpan next_indices) { + (void)current_length; + (void)next_tokens; + (void)next_indices; + if (model_.config_->model.engram.run_options.has_value()) { + State::SetRunOptions(model_.config_->model.engram.run_options.value()); + } + if (!cache_hit_) { + State::Run(session_); + if (active_cache_key_.has_value()) { + if (cache_.size() >= model_.config_->model.engram.cache_capacity) cache_.clear(); + CacheEntry entry; + const auto output_info = embeddings_output_->GetTensorTypeAndShapeInfo(); + const size_t output_bytes = output_info->GetElementCount() * Ort::SizeOf(output_info->GetElementType()); + const auto* output_data = static_cast(embeddings_output_->GetTensorRawData()); + entry.embeddings.assign(output_data, output_data + output_bytes); + const auto* present_data = present_tokens_->GetTensorData(); + const size_t present_count = present_tokens_->GetTensorTypeAndShapeInfo()->GetElementCount(); + entry.present_tokens.assign(present_data, present_data + present_count); + cache_[*active_cache_key_] = std::move(entry); + } + } + first_run_ = false; + return {}; +} + +void EngramState::CopyEmbeddingsTo(Embeddings& destination) { + if (!destination.Get()) throw std::runtime_error("Engram embeddings are not allocated"); + CopyEmbeddingsTo(*destination.Get()); +} + +void EngramState::CopyEmbeddingsTo(OrtValue& destination) { + if (!embeddings_output_) throw std::runtime_error("Engram embeddings are not allocated"); + const auto source_info = embeddings_output_->GetTensorTypeAndShapeInfo(); + const auto destination_info = destination.GetTensorTypeAndShapeInfo(); + if (source_info->GetElementType() != destination_info->GetElementType() || + source_info->GetElementCount() != destination_info->GetElementCount()) { + throw std::runtime_error("Engram output must match the decoder Engram input type and shape"); + } + ByteWrapTensor(*model_.p_device_inputs_, destination).CopyFrom(ByteWrapTensor(*GetDeviceInterface(DeviceType::CPU), *embeddings_output_)); +} + +void EngramState::ValidateRewindTo(size_t index) const { + if (index != 0 && (!snapshot_valid_ || snapshot_position_ != index)) + throw std::runtime_error("EngramState cannot rewind to position " + std::to_string(index) + + " without a matching snapshot"); +} + +void EngramState::RewindTo(size_t index) { + ValidateRewindTo(index); + if (index == 0) { + InitializeTokens(*past_tokens_); + InitializeTokens(*present_tokens_); + snapshot_valid_ = false; + first_run_ = true; + inputs_[past_input_index_] = past_tokens_.get(); + outputs_[present_output_index_] = present_tokens_.get(); + return; + } + CopyTokens(*snapshot_tokens_, *present_tokens_); +} + +void EngramState::SnapshotState(size_t position) { + if (!snapshot_tokens_) { + auto shape = present_tokens_->GetTensorTypeAndShapeInfo()->GetShape(); + snapshot_tokens_ = OrtValue::CreateTensor(model_.allocator_cpu_, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64); + } + CopyTokens(*present_tokens_, *snapshot_tokens_); + snapshot_position_ = position; + snapshot_valid_ = true; +} + +void EngramState::ValidateAcceptedPrefix(size_t token_count) const { + if (!snapshot_valid_ || !snapshot_tokens_ || !input_ids_) { + throw std::runtime_error("EngramState accepted-prefix commit requires a snapshot and input tokens"); + } + const auto input_info = input_ids_->GetTensorTypeAndShapeInfo(); + const auto input_shape = input_info->GetShape(); + if (input_shape.size() != 2 || input_shape[0] != 1 || token_count > static_cast(input_shape[1])) { + throw std::runtime_error("EngramState accepted-prefix commit received an invalid token count"); + } +} + +void EngramState::CommitAcceptedPrefix(size_t token_count) { + ValidateAcceptedPrefix(token_count); + CopyTokens(*snapshot_tokens_, *present_tokens_); + auto* history = present_tokens_->GetTensorMutableData(); + const size_t history_size = present_tokens_->GetTensorTypeAndShapeInfo()->GetElementCount(); + const auto* tokens = input_ids_->GetTensorData(); + for (size_t token_index = 0; token_index < token_count; ++token_index) { + if (history_size > 1) std::move(history + 1, history + history_size, history); + history[history_size - 1] = tokens[token_index]; + } +} + +void EngramState::InitializeTokens(OrtValue& tokens) { + auto* data = tokens.GetTensorMutableData(); + const size_t count = tokens.GetTensorTypeAndShapeInfo()->GetElementCount(); + std::fill_n(data, count, model_.config_->model.decoder.ple_token_pad_id); +} + +void EngramState::CopyTokens(OrtValue& source, OrtValue& destination) { + const size_t bytes = source.GetTensorTypeAndShapeInfo()->GetElementCount() * sizeof(int64_t); + std::memcpy(destination.GetTensorMutableData(), source.GetTensorData(), bytes); +} + +} // namespace Generators \ No newline at end of file diff --git a/src/models/engram.h b/src/models/engram.h new file mode 100644 index 0000000000..cb2dcea50d --- /dev/null +++ b/src/models/engram.h @@ -0,0 +1,63 @@ +#pragma once + +#include "model.h" +#include "io/embeddings.h" + +namespace Generators { + +struct EngramState : State { + EngramState(const Model& model, OrtSession& session, const GeneratorParams& params); + + void UpdateInputsOutputs(DeviceSpan& next_tokens); + DeviceSpan Run(int current_length, DeviceSpan& next_tokens, + DeviceSpan next_indices = {}) override; + void CopyEmbeddingsTo(Embeddings& destination); + void CopyEmbeddingsTo(OrtValue& destination); + void RewindTo(size_t index) override; + void ValidateRewindTo(size_t index) const override; + void SnapshotState(size_t position) override; + void CommitAcceptedPrefix(size_t token_count); + void ValidateAcceptedPrefix(size_t token_count) const; + + private: + struct CacheEntry { + std::vector embeddings; + std::vector present_tokens; + }; + + struct CacheKeyHash { + size_t operator()(const std::vector& key) const noexcept { + size_t hash = 1469598103934665603ull; + for (int64_t value : key) { + hash ^= static_cast(value); + hash *= 1099511628211ull; + } + return hash; + } + }; + + void InitializeTokens(OrtValue& tokens); + void CopyTokens(OrtValue& source, OrtValue& destination); + + const Model& model_; + OrtSession& session_; + std::unique_ptr input_ids_; + std::unique_ptr embeddings_output_; + std::unique_ptr past_tokens_; + std::unique_ptr present_tokens_; + std::unique_ptr snapshot_tokens_; + size_t past_input_index_{~0U}; + size_t input_ids_index_{~0U}; + size_t embeddings_output_index_{~0U}; + size_t present_output_index_{~0U}; + ONNXTensorElementDataType embeddings_type_{}; + int64_t embeddings_width_{}; + size_t snapshot_position_{}; + bool snapshot_valid_{}; + bool first_run_{true}; + bool cache_hit_{}; + std::optional> active_cache_key_; + std::unordered_map, CacheEntry, CacheKeyHash> cache_; +}; + +} // namespace Generators \ No newline at end of file diff --git a/src/models/io/embeddings.h b/src/models/io/embeddings.h index 8297708383..f35e0f9906 100644 --- a/src/models/io/embeddings.h +++ b/src/models/io/embeddings.h @@ -3,6 +3,8 @@ #pragma once +#include + namespace Generators { struct Model; diff --git a/src/models/io/hidden_states.cpp b/src/models/io/hidden_states.cpp index b71b31f410..f4dc6b4173 100644 --- a/src/models/io/hidden_states.cpp +++ b/src/models/io/hidden_states.cpp @@ -11,7 +11,13 @@ HiddenStatesInputs::HiddenStatesInputs(State& state) : state_{state} { const std::string& name = model_.config_->model.decoder.inputs.hidden_states; type_ = model_.session_info_.GetInputDataType(name); - shape_ = {state_.params_->BatchBeamSize(), 0, model_.config_->model.decoder.hidden_size}; + const auto input_shape = model_.session_info_.GetInputShape(name); + const int64_t width = !input_shape.empty() && input_shape.back() > 0 + ? input_shape.back() + : model_.config_->model.decoder.hidden_size; + if (width <= 0) + throw std::runtime_error("hidden_states input requires a positive graph or configured hidden_size"); + shape_ = {state_.params_->BatchBeamSize(), 0, width}; } void HiddenStatesInputs::Add() { @@ -92,7 +98,13 @@ HiddenStatesOutputs::HiddenStatesOutputs(State& state) : state_{state} { const std::string& name = model_.config_->model.decoder.outputs.hidden_states; type_ = model_.session_info_.GetOutputDataType(name); - shape_ = {state_.params_->BatchBeamSize(), 0, model_.config_->model.decoder.hidden_size}; + const auto output_shape = model_.session_info_.GetOutputShape(name); + const int64_t width = !output_shape.empty() && output_shape.back() > 0 + ? output_shape.back() + : model_.config_->model.decoder.hidden_size; + if (width <= 0) + throw std::runtime_error("hidden_states output requires a positive graph or configured hidden_size"); + shape_ = {state_.params_->BatchBeamSize(), 0, width}; } void HiddenStatesOutputs::Add() { diff --git a/src/models/io/indexer_cache.cpp b/src/models/io/indexer_cache.cpp new file mode 100644 index 0000000000..1f19510fe4 --- /dev/null +++ b/src/models/io/indexer_cache.cpp @@ -0,0 +1,364 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "generator/generators.h" +#include "models/io/indexer_cache.h" +#include "models/io/static_kv_cache.h" +#include + +namespace Generators { +namespace { + +std::string ComposeIndexerName(const std::string& name_template, int layer_index) { + constexpr size_t buffer_size = 128; + char name[buffer_size]; + const int length = snprintf(name, buffer_size, name_template.c_str(), layer_index); + if (length < 0 || static_cast(length) >= buffer_size) + throw std::runtime_error("Unable to compose indexer cache name from template " + name_template); + return name; +} + +} // namespace + +IndexerCache::IndexerCache(State& state) : state_{state} { + const auto& inputs = model_.config_->model.decoder.inputs; + const auto& outputs = model_.config_->model.decoder.outputs; + const auto placeholder = inputs.past_indexer_names.find("%d"); + if (placeholder == std::string::npos) return; + + const auto prefix = inputs.past_indexer_names.substr(0, placeholder); + const auto suffix = inputs.past_indexer_names.substr(placeholder + 2); + for (const auto& name : model_.session_info_.GetInputNames()) { + if (name.size() > prefix.size() + suffix.size() && + name.compare(0, prefix.size(), prefix) == 0 && + name.compare(name.size() - suffix.size(), suffix.size(), suffix) == 0) { + layer_indices_.push_back(std::stoi(name.substr(prefix.size(), name.size() - prefix.size() - suffix.size()))); + } + } + std::sort(layer_indices_.begin(), layer_indices_.end()); + if (layer_indices_.empty()) return; + compress_ratio_ = model_.config_->model.decoder.indexer_compress_ratio; + if (model_.config_->model.decoder.state_groups) { + for (const auto& group : *model_.config_->model.decoder.state_groups) { + if (group.kind == Config::Model::Decoder::StateGroupKind::FixedIndexer && + group.state_update && group.state_update->compress_ratio > 0) { + const size_t ratio = static_cast(group.state_update->compress_ratio); + if (compress_ratio_ != 0 && compress_ratio_ != ratio) + throw std::runtime_error("IndexerCache compression metadata disagrees between decoder state groups"); + compress_ratio_ = ratio; + } + } + } + if (compress_ratio_ == 0) compress_ratio_ = 4; // Legacy dense Qwen3.8 cache contract. + if (outputs.present_indexer_names.empty()) + throw std::runtime_error("IndexerCache: present indexer name template must be configured"); + + for (int layer_index : layer_indices_) { + input_name_strings_.push_back(ComposeIndexerName(inputs.past_indexer_names, layer_index)); + output_name_strings_.push_back(ComposeIndexerName(outputs.present_indexer_names, layer_index)); + if (!model_.session_info_.HasOutput(output_name_strings_.back())) + throw std::runtime_error("IndexerCache: missing output for layer " + std::to_string(layer_index)); + } + const auto& update_value_template = + outputs.state_update_indexer_value_names; + const auto& update_row_template = + outputs.state_update_indexer_row_names; + if (update_value_template.empty() != update_row_template.empty()) { + throw std::runtime_error( + "IndexerCache: state update value and row templates must be configured together"); + } + if (!update_value_template.empty()) { + for (int layer_index : layer_indices_) { + state_update_value_name_strings_.push_back( + ComposeIndexerName(update_value_template, layer_index)); + state_update_row_name_strings_.push_back( + ComposeIndexerName(update_row_template, layer_index)); + } + state_update_values_.resize(layer_indices_.size()); + state_update_rows_values_.resize(layer_indices_.size()); + } + + type_ = model_.session_info_.GetInputDataType(input_name_strings_[0]); + shape_ = model_.session_info_.GetInputShape(input_name_strings_[0]); + if (shape_.size() != 3) + throw std::runtime_error("IndexerCache: expected rank-3 cache tensors"); + shape_[0] = state_.params_->BatchBeamSize(); + if (shape_[2] <= 0) + throw std::runtime_error("IndexerCache: head dimension must be static"); + const int64_t fixed_sequence_length = shape_[1] > 0 ? shape_[1] : 0; + share_buffer_ = ShouldUseSharedPastPresentKeyValueCache(state_); + if (share_buffer_ && fixed_sequence_length == 0 && + (inputs.past_sequence_length.empty() || !model_.session_info_.HasInput(inputs.past_sequence_length))) { + share_buffer_ = false; + } + if (fixed_sequence_length > 0) { + if (state_.params_->search.num_beams != 1) + throw std::runtime_error("IndexerCache: beam search is not supported with a fixed cache shape"); + share_buffer_ = true; + } + + for (const auto& output_name : output_name_strings_) { + const auto output_shape = model_.session_info_.GetOutputShape(output_name); + if (output_shape.size() != 3 || output_shape[2] != shape_[2]) + throw std::runtime_error("IndexerCache: input and output cache head dimensions must match"); + if (fixed_sequence_length > 0 && output_shape[1] > 0 && output_shape[1] != fixed_sequence_length) + throw std::runtime_error("IndexerCache: fixed input and output cache shapes must match"); + } + + shape_[1] = share_buffer_ + ? (fixed_sequence_length > 0 ? fixed_sequence_length : state_.params_->search.max_length) + : 0; + if (share_buffer_ && shape_[1] <= 0) + throw std::runtime_error("IndexerCache: shared caches require search.max_length > 0"); + + auto& allocator = model_.p_device_kvcache_->GetAllocator(); + pasts_.resize(layer_indices_.size()); + presents_.resize(layer_indices_.size()); + empty_pasts_.reserve(layer_indices_.size()); + for (size_t index = 0; index < layer_indices_.size(); ++index) + empty_pasts_.push_back(OrtValue::CreateTensor(allocator, shape_, type_)); + + const auto& length_name = inputs.past_sequence_length; + if (!length_name.empty() && model_.session_info_.HasInput(length_name)) { + if (model_.session_info_.GetInputDataType(length_name) != Ort::TypeToTensorType) + throw std::runtime_error("IndexerCache: past_sequence_length input must be int32"); + past_sequence_length_ = OrtValue::CreateTensor( + model_.allocator_cpu_, std::array{1}, Ort::TypeToTensorType); + } else if (share_buffer_) { + throw std::runtime_error("IndexerCache: shared caches require an int32 past_sequence_length input"); + } + if (past_sequence_length_ && state_.params_->use_graph_capture) { + Log("warning", "Dense sparse-indexer execution runs eagerly because past_sequence_length is read on the host. Use paged execution for graph capture."); + } +} + +void IndexerCache::Add() { + if (layer_indices_.empty()) return; + input_index_ = state_.inputs_.size(); + output_index_ = state_.outputs_.size(); + for (size_t index = 0; index < layer_indices_.size(); ++index) { + state_.inputs_.push_back(empty_pasts_[index].get()); + state_.input_names_.push_back(input_name_strings_[index].c_str()); + state_.outputs_.push_back(share_buffer_ ? empty_pasts_[index].get() : nullptr); + state_.output_names_.push_back(output_name_strings_[index].c_str()); + } + if (past_sequence_length_) { + state_.inputs_.push_back(past_sequence_length_.get()); + state_.input_names_.push_back(model_.config_->model.decoder.inputs.past_sequence_length.c_str()); + } + if (!state_update_value_name_strings_.empty()) { + state_update_output_index_ = state_.outputs_.size(); + for (size_t index = 0; index < layer_indices_.size(); ++index) { + state_.outputs_.push_back(nullptr); + state_.output_names_.push_back( + state_update_value_name_strings_[index].c_str()); + state_.outputs_.push_back(nullptr); + state_.output_names_.push_back( + state_update_row_name_strings_[index].c_str()); + } + } +} + +void IndexerCache::Update(DeviceSpan beam_indices, int total_length, int current_length) { + if (!beam_indices.empty()) + throw std::runtime_error("IndexerCache does not support beam reordering"); + if (past_sequence_length_) { + if (current_length < 0 || current_length > total_length) + throw std::runtime_error("IndexerCache: current length must be in [0, total length]"); + *past_sequence_length_->GetTensorMutableData() = total_length - current_length; + } + if (share_buffer_) { + if (total_length > shape_[1]) + throw std::runtime_error("IndexerCache: total length exceeds the shared cache capacity"); + if (!state_update_value_name_strings_.empty()) { + state_update_length_ = static_cast(current_length); + auto& allocator = model_.p_device_kvcache_->GetAllocator(); + for (size_t index = 0; index < layer_indices_.size(); ++index) { + state_update_values_[index] = OrtValue::CreateTensor( + allocator, + std::array{ + state_.params_->BatchBeamSize(), current_length, shape_[2]}, + type_); + state_update_rows_values_[index] = OrtValue::CreateTensor( + allocator, + std::array{ + state_.params_->BatchBeamSize(), current_length}, + ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32); + auto rows = WrapTensor( + *model_.p_device_kvcache_, *state_update_rows_values_[index]); + std::fill(rows.CpuSpan().begin(), rows.CpuSpan().end(), -1); + rows.CopyCpuToDevice(); + state_.outputs_[state_update_output_index_ + index * 2] = + state_update_values_[index].get(); + state_.outputs_[state_update_output_index_ + index * 2 + 1] = + state_update_rows_values_[index].get(); + } + } + return; + } + + if (!first_update_) { + for (size_t index = 0; index < layer_indices_.size(); ++index) { + pasts_[index] = std::move(presents_[index]); + state_.inputs_[input_index_ + index] = pasts_[index].get(); + } + } + + auto output_shape = shape_; + output_shape[1] = total_length; + auto& allocator = model_.p_device_kvcache_->GetAllocator(); + for (size_t index = 0; index < layer_indices_.size(); ++index) { + presents_[index] = OrtValue::CreateTensor(allocator, output_shape, type_); + state_.outputs_[output_index_ + index] = presents_[index].get(); + } + first_update_ = false; +} + +void IndexerCache::ValidateAcceptedPrefix(size_t token_count) const { + if (!HasStateUpdates() || snapshot_rows_.empty() || token_count > state_update_length_) { + throw std::runtime_error( + "IndexerCache accepted-prefix commit exceeds captured state updates"); + } + const size_t representative_capacity = static_cast(shape_[1]) / compress_ratio_; + for (size_t token = 0; token < token_count; ++token) { + const size_t position = snapshot_position_ + token; + const size_t row = position % compress_ratio_ == compress_ratio_ - 1 + ? position / compress_ratio_ + : representative_capacity + position % compress_ratio_; + if (row >= static_cast(shape_[1])) + throw std::runtime_error("IndexerCache accepted-prefix row exceeds the cache capacity"); + } +} + +void IndexerCache::CommitAcceptedPrefix(size_t token_count) { + ValidateAcceptedPrefix(token_count); + RewindTo(snapshot_position_); + auto& device = *model_.p_device_kvcache_; + const size_t representative_capacity = static_cast(shape_[1]) / compress_ratio_; + + for (size_t layer = 0; layer < layer_indices_.size(); ++layer) { + auto cache = ByteWrapTensor(device, *empty_pasts_[layer]); + auto updates = ByteWrapTensor(device, *state_update_values_[layer]); + for (size_t token = 0; token < token_count; ++token) { + const size_t position = snapshot_position_ + token; + const size_t row = + position % compress_ratio_ == compress_ratio_ - 1 ? position / compress_ratio_ + : representative_capacity + position % compress_ratio_; + if (row >= static_cast(shape_[1])) { + throw std::runtime_error( + "IndexerCache state update contains invalid cache row " + + std::to_string(row) + " at layer " + std::to_string(layer) + + " token " + std::to_string(token) + " with capacity " + + std::to_string(shape_[1])); + } + cache + .subspan(row * snapshot_row_bytes_, snapshot_row_bytes_) + .CopyFrom(updates.subspan(token * snapshot_row_bytes_, + snapshot_row_bytes_)); + } + } +} + +void IndexerCache::Snapshot(size_t position) { + if (!share_buffer_ || empty_pasts_.empty()) return; + + const size_t capacity = static_cast(shape_[1]); + snapshot_rows_.clear(); + const auto add_row = [&](size_t row) { + if (row < capacity && + std::find(snapshot_rows_.begin(), snapshot_rows_.end(), row) == + snapshot_rows_.end()) { + snapshot_rows_.push_back(row); + } + }; + + // Qwen's compressed indexer cache stores completed block representatives in the first quarter + // and the current four raw keys in scratch rows immediately after that prefix. + const size_t representative_capacity = capacity / compress_ratio_; + for (size_t token = 0; token <= Config::Model::Decoder::MaxStateUpdateCapacity; ++token) { + add_row((position + token) / compress_ratio_); + } + for (size_t row = 0; row < compress_ratio_; ++row) { + add_row(representative_capacity + row); + } + // Also preserve ordinary append positions for non-compressed/fallback layouts. + for (size_t row = position; row < position + 8; ++row) { + add_row(row); + } + std::sort(snapshot_rows_.begin(), snapshot_rows_.end()); + + snapshot_row_bytes_ = + static_cast(shape_[2]) * Ort::SizeOf(type_); + const size_t total_bytes = + layer_indices_.size() * snapshot_rows_.size() * snapshot_row_bytes_; + auto& device = *model_.p_device_kvcache_; + if (snapshot_data_.size() != total_bytes) { + snapshot_data_ = device.Allocate(total_bytes); + } + for (size_t layer = 0; layer < empty_pasts_.size(); ++layer) { + auto cache = ByteWrapTensor(device, *empty_pasts_[layer]); + for (size_t row_index = 0; row_index < snapshot_rows_.size(); ++row_index) { + snapshot_data_ + .subspan( + (layer * snapshot_rows_.size() + row_index) * snapshot_row_bytes_, + snapshot_row_bytes_) + .CopyFrom(cache.subspan( + snapshot_rows_[row_index] * snapshot_row_bytes_, + snapshot_row_bytes_)); + } + } + snapshot_position_ = position; +} + +void IndexerCache::ValidateRewindTo(size_t index) const { + if (layer_indices_.empty()) return; + if (share_buffer_ && index > static_cast(shape_[1])) + throw std::runtime_error("IndexerCache rewind exceeds the shared cache capacity"); + if (index != 0 && (!share_buffer_ || snapshot_rows_.empty() || index != snapshot_position_)) + throw std::runtime_error("IndexerCache can only rewind to zero or its saved snapshot position"); +} + +void IndexerCache::RewindTo(size_t index) { + ValidateRewindTo(index); + if (layer_indices_.empty()) return; + if (share_buffer_) { + if (index == 0) { + for (auto& cache : empty_pasts_) + ByteWrapTensor(*model_.p_device_kvcache_, *cache).Zero(); + snapshot_rows_.clear(); + snapshot_position_ = 0; + return; + } + if (!snapshot_rows_.empty() && index == snapshot_position_) { + auto& device = *model_.p_device_kvcache_; + for (size_t layer = 0; layer < empty_pasts_.size(); ++layer) { + auto cache = ByteWrapTensor(device, *empty_pasts_[layer]); + for (size_t row_index = 0; row_index < snapshot_rows_.size(); ++row_index) { + cache + .subspan( + snapshot_rows_[row_index] * snapshot_row_bytes_, + snapshot_row_bytes_) + .CopyFrom(snapshot_data_.subspan( + (layer * snapshot_rows_.size() + row_index) * + snapshot_row_bytes_, + snapshot_row_bytes_)); + } + } + } + return; + } + first_update_ = true; + for (size_t cache_index = 0; cache_index < layer_indices_.size(); ++cache_index) { + pasts_[cache_index].reset(); + presents_[cache_index].reset(); + state_.inputs_[input_index_ + cache_index] = empty_pasts_[cache_index].get(); + state_.outputs_[output_index_ + cache_index] = nullptr; + } +} + +std::unique_ptr CreateIndexerCache(State& state) { + auto cache = std::make_unique(state); + return cache->IsEmpty() ? nullptr : std::move(cache); +} + +} // namespace Generators \ No newline at end of file diff --git a/src/models/io/indexer_cache.h b/src/models/io/indexer_cache.h new file mode 100644 index 0000000000..645d825d47 --- /dev/null +++ b/src/models/io/indexer_cache.h @@ -0,0 +1,55 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "models/model.h" + +namespace Generators { + +struct IndexerCache { + explicit IndexerCache(State& state); + + void Add(); + void Update(DeviceSpan beam_indices, int total_length, int current_length); + void Snapshot(size_t position); + void RewindTo(size_t index); + void ValidateRewindTo(size_t index) const; + void CommitAcceptedPrefix(size_t token_count); + void ValidateAcceptedPrefix(size_t token_count) const; + bool UsesHostSequenceLength() const { return past_sequence_length_ != nullptr; } + bool HasStateUpdates() const { return !state_update_value_name_strings_.empty(); } + bool IsEmpty() const { return layer_indices_.empty(); } + + private: + State& state_; + const Model& model_{state_.model_}; + std::vector layer_indices_; + std::vector> pasts_; + std::vector> presents_; + std::vector> empty_pasts_; + std::unique_ptr past_sequence_length_; + std::vector input_name_strings_; + std::vector output_name_strings_; + std::vector state_update_value_name_strings_; + std::vector state_update_row_name_strings_; + std::vector> state_update_values_; + std::vector> state_update_rows_values_; + std::vector shape_; + ONNXTensorElementDataType type_{}; + bool share_buffer_{false}; + bool first_update_{true}; + size_t input_index_{~0U}; + size_t output_index_{~0U}; + size_t state_update_output_index_{~0U}; + size_t state_update_length_{}; + size_t snapshot_position_{}; + size_t snapshot_row_bytes_{}; + size_t compress_ratio_{}; + std::vector snapshot_rows_; + DeviceSpan snapshot_data_; +}; + +std::unique_ptr CreateIndexerCache(State& state); + +} // namespace Generators \ No newline at end of file diff --git a/src/models/io/ple_state.cpp b/src/models/io/ple_state.cpp new file mode 100644 index 0000000000..d0711ee775 --- /dev/null +++ b/src/models/io/ple_state.cpp @@ -0,0 +1,212 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "generator/generators.h" +#include "models/io/ple_state.h" +#include + +namespace Generators { +namespace { + +std::string ComposePleName(const std::string& name_template, int layer_index) { + constexpr size_t buffer_size = 128; + char name[buffer_size]; + const int length = snprintf(name, buffer_size, name_template.c_str(), layer_index); + if (length < 0 || static_cast(length) >= buffer_size) + throw std::runtime_error("Unable to compose PLE state name from template " + name_template); + return name; +} + +void FixAndValidateBatchDimension( + std::vector& shape, int batch_size, const std::string& name, + size_t batch_axis = 0) { + if (shape.empty()) + throw std::runtime_error("PleState: " + name + " must have a batch dimension"); + if (batch_axis >= shape.size()) + throw std::runtime_error("PleState: " + name + " has no batch axis"); + if (shape[batch_axis] <= 0) shape[batch_axis] = batch_size; + for (size_t axis = 0; axis < shape.size(); ++axis) { + if (shape[axis] <= 0) + throw std::runtime_error("PleState: " + name + " has unsupported dynamic dimension at axis " + + std::to_string(axis)); + } +} + +} // namespace + +PleState::PleState(State& state) : state_{state} { + const auto& inputs = model_.config_->model.decoder.inputs; + const auto& outputs = model_.config_->model.decoder.outputs; + const auto placeholder = inputs.past_ple_token_names.find("%d"); + if (placeholder == std::string::npos) return; + + const auto prefix = inputs.past_ple_token_names.substr(0, placeholder); + const auto suffix = inputs.past_ple_token_names.substr(placeholder + 2); + for (const auto& name : model_.session_info_.GetInputNames()) { + if (name.size() > prefix.size() + suffix.size() && + name.compare(0, prefix.size(), prefix) == 0 && + name.compare(name.size() - suffix.size(), suffix.size(), suffix) == 0) { + layer_indices_.push_back(std::stoi(name.substr(prefix.size(), name.size() - prefix.size() - suffix.size()))); + } + } + std::sort(layer_indices_.begin(), layer_indices_.end()); + if (layer_indices_.empty()) return; + if (model_.config_->model.decoder.ple_token_pad_id < 0) + throw std::runtime_error("PleState requires model.decoder.ple_token_pad_id"); + + if (inputs.past_ple_conv_names.empty() || outputs.present_ple_token_names.empty() || + outputs.present_ple_conv_names.empty()) + throw std::runtime_error("PleState: all PLE input and output name templates must be configured"); + + for (int layer_index : layer_indices_) { + input_name_strings_.push_back(ComposePleName(inputs.past_ple_token_names, layer_index)); + input_name_strings_.push_back(ComposePleName(inputs.past_ple_conv_names, layer_index)); + output_name_strings_.push_back(ComposePleName(outputs.present_ple_token_names, layer_index)); + output_name_strings_.push_back(ComposePleName(outputs.present_ple_conv_names, layer_index)); + if (!model_.session_info_.HasInput(input_name_strings_[input_name_strings_.size() - 1]) || + !model_.session_info_.HasOutput(output_name_strings_[output_name_strings_.size() - 2]) || + !model_.session_info_.HasOutput(output_name_strings_[output_name_strings_.size() - 1])) + throw std::runtime_error("PleState: incomplete PLE state pair for layer " + std::to_string(layer_index)); + } + + if (model_.session_info_.GetInputDataType(input_name_strings_[0]) != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64) + throw std::runtime_error("PleState: token history must use INT64"); + conv_type_ = model_.session_info_.GetInputDataType(input_name_strings_[1]); + token_shape_ = model_.session_info_.GetInputShape(input_name_strings_[0]); + conv_shape_ = model_.session_info_.GetInputShape(input_name_strings_[1]); + if (conv_shape_.size() == 4) { + state_window_ = conv_shape_[0]; + } + FixAndValidateBatchDimension(token_shape_, state_.params_->BatchBeamSize(), "token history"); + FixAndValidateBatchDimension( + conv_shape_, state_.params_->BatchBeamSize(), "convolution state", + conv_shape_.size() == 4 ? 1 : 0); + + auto& allocator = model_.p_device_kvcache_->GetAllocator(); + pasts_.reserve(layer_indices_.size() * 2); + presents_.reserve(layer_indices_.size() * 2); + for (size_t index = 0; index < layer_indices_.size(); ++index) { + pasts_.push_back(OrtValue::CreateTensor(allocator, token_shape_, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64)); + pasts_.push_back(OrtValue::CreateTensor(allocator, conv_shape_, conv_type_)); + presents_.push_back(OrtValue::CreateTensor(allocator, token_shape_, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64)); + presents_.push_back(OrtValue::CreateTensor(allocator, conv_shape_, conv_type_)); + } + InitializeStates(pasts_); + InitializeStates(presents_); +} + +void PleState::Add() { + if (layer_indices_.empty()) return; + input_index_ = state_.inputs_.size(); + output_index_ = state_.outputs_.size(); + for (size_t index = 0; index < pasts_.size(); ++index) { + state_.inputs_.push_back(pasts_[index].get()); + state_.input_names_.push_back(input_name_strings_[index].c_str()); + state_.outputs_.push_back(presents_[index].get()); + state_.output_names_.push_back(output_name_strings_[index].c_str()); + } +} + +void PleState::Update() { + for (size_t index = 0; index < pasts_.size(); ++index) { + std::swap(pasts_[index], presents_[index]); + state_.inputs_[input_index_ + index] = pasts_[index].get(); + state_.outputs_[output_index_ + index] = presents_[index].get(); + } + if (!pasts_.empty()) graph_buffer_variant_ ^= 1; +} + +void PleState::ValidateRewindTo(size_t index) const { + if (!layer_indices_.empty() && index != 0 && (!snapshot_valid_ || index != snapshot_position_)) + throw std::runtime_error("PleState cannot rewind to position " + std::to_string(index) + + " without a matching snapshot"); +} + +void PleState::RewindTo(size_t index) { + ValidateRewindTo(index); + if (layer_indices_.empty()) return; + if (index != 0) { + CopyStates(snapshot_, presents_); + return; + } + snapshot_valid_ = false; + InitializeStates(pasts_); + InitializeStates(presents_); + if (graph_buffer_variant_ != 0) { + for (size_t state_index = 0; state_index < pasts_.size(); ++state_index) + std::swap(pasts_[state_index], presents_[state_index]); + } + graph_buffer_variant_ = 0; + for (size_t state_index = 0; state_index < pasts_.size(); ++state_index) { + state_.inputs_[input_index_ + state_index] = pasts_[state_index].get(); + state_.outputs_[output_index_ + state_index] = presents_[state_index].get(); + } +} + +void PleState::Snapshot(size_t position) { + if (layer_indices_.empty()) return; + if (snapshot_.empty()) { + auto& allocator = model_.p_device_kvcache_->GetAllocator(); + snapshot_.reserve(layer_indices_.size() * 2); + for (size_t index = 0; index < layer_indices_.size(); ++index) { + snapshot_.push_back( + OrtValue::CreateTensor(allocator, token_shape_, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64)); + snapshot_.push_back(OrtValue::CreateTensor(allocator, conv_shape_, conv_type_)); + } + } + CopyStates(presents_, snapshot_); + snapshot_position_ = position; + snapshot_valid_ = true; +} + +void PleState::ValidateCropPosition(size_t position) const { + if (!IsWindowed()) { + throw std::runtime_error("PleState::CropToPosition requires windowed convolution state"); + } + const int64_t signed_slot = + static_cast(position) + state_window_ - forward_length_; + if (signed_slot < 0 || signed_slot >= state_window_) { + throw std::runtime_error("PleState crop position is outside the retained state window"); + } +} + +void PleState::CropToPosition(size_t position) { + ValidateCropPosition(position); + const int64_t signed_slot = static_cast(position) + state_window_ - forward_length_; + const size_t slot = static_cast(signed_slot); + if (slot + 1 == static_cast(state_window_)) return; + + auto& device = *model_.p_device_kvcache_; + for (size_t index = 1; index < presents_.size(); index += 2) { + auto bytes = ByteWrapTensor(device, *presents_[index]); + const size_t slot_bytes = bytes.size() / static_cast(state_window_); + bytes.subspan((state_window_ - 1) * slot_bytes, slot_bytes) + .CopyFrom(bytes.subspan(slot * slot_bytes, slot_bytes)); + } +} + +void PleState::CopyStates( + const std::vector>& source, + std::vector>& destination) { + auto& device = *model_.p_device_kvcache_; + for (size_t index = 0; index < source.size(); ++index) { + ByteWrapTensor(device, *destination[index]).CopyFrom(ByteWrapTensor(device, *source[index])); + } +} + +void PleState::InitializeStates(std::vector>& states) { + auto& device = *model_.p_device_kvcache_; + for (size_t index = 0; index < states.size(); index += 2) { + auto tokens = WrapTensor(device, *states[index]); + std::fill(tokens.CpuSpan().begin(), tokens.CpuSpan().end(), model_.config_->model.decoder.ple_token_pad_id); + tokens.CopyCpuToDevice(); + ByteWrapTensor(device, *states[index + 1]).Zero(); + } +} + +std::unique_ptr CreatePleState(State& state) { + auto ple_state = std::make_unique(state); + return ple_state->IsEmpty() ? nullptr : std::move(ple_state); +} + +} // namespace Generators \ No newline at end of file diff --git a/src/models/io/ple_state.h b/src/models/io/ple_state.h new file mode 100644 index 0000000000..a274e2611e --- /dev/null +++ b/src/models/io/ple_state.h @@ -0,0 +1,53 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "models/model.h" + +namespace Generators { + +struct PleState { + explicit PleState(State& state); + + void Add(); + void Update(); + void RewindTo(size_t index); + void ValidateRewindTo(size_t index) const; + void Snapshot(size_t position); + void SetForwardLength(int sequence_length) { forward_length_ = sequence_length; } + void CropToPosition(size_t position); + void ValidateCropPosition(size_t position) const; + bool IsWindowed() const { return state_window_ > 1; } + + bool IsEmpty() const { return layer_indices_.empty(); } + int GraphCaptureVariant() const { return graph_buffer_variant_; } + + private: + void InitializeStates(std::vector>& states); + void CopyStates(const std::vector>& source, + std::vector>& destination); + + State& state_; + const Model& model_{state_.model_}; + std::vector layer_indices_; + std::vector> pasts_; + std::vector> presents_; + std::vector> snapshot_; + std::vector input_name_strings_; + std::vector output_name_strings_; + std::vector token_shape_; + std::vector conv_shape_; + ONNXTensorElementDataType conv_type_{}; + int graph_buffer_variant_{}; + size_t snapshot_position_{}; + bool snapshot_valid_{}; + int64_t state_window_{1}; + int forward_length_{}; + size_t input_index_{~0U}; + size_t output_index_{~0U}; +}; + +std::unique_ptr CreatePleState(State& state); + +} // namespace Generators \ No newline at end of file diff --git a/src/models/io/qwen_vl_position_inputs.cpp b/src/models/io/qwen_vl_position_inputs.cpp index 4061262dde..758fc757cf 100644 --- a/src/models/io/qwen_vl_position_inputs.cpp +++ b/src/models/io/qwen_vl_position_inputs.cpp @@ -109,6 +109,24 @@ Qwen2VLPositionInputs::Qwen2VLPositionInputs(const Model& model, State& state, D attention_mask_shape_[1] = 0; // Will be set during first update attention_mask_ = std::make_unique(model_.p_device_inputs_, type_); } + const auto& history_name = model_.config_->model.decoder.inputs.rotary_position_ids; + if (!history_name.empty()) { + const auto shape = model_.session_info_.GetInputShape(history_name); + if (!has_posid_input_ || type_ != Ort::TypeToTensorType || + model_.session_info_.GetInputDataType(history_name) != Ort::TypeToTensorType || + shape.size() != 3 || shape[0] != 3 || state_.params_->search.max_length <= 0) { + throw std::runtime_error("Rotary history requires int64 [3, batch, history] and current T/H/W positions"); + } + const std::array history_shape{ + 3, state_.params_->BatchBeamSize(), state_.params_->search.max_length}; + rotary_history_ = std::make_unique( + history_shape[1], history_shape[2], model_.config_->model.context_length - 1); + rotary_history_value_ = OrtValue::CreateTensor( + model_.allocator_cpu_.GetInfo(), rotary_history_->Values().data(), + rotary_history_->Values().size_bytes(), history_shape, Ort::TypeToTensorType); + rotary_history_input_ = std::make_unique(model_.p_device_inputs_, Ort::TypeToTensorType); + rotary_history_input_->CreateTensor(history_shape, true); + } } void Qwen2VLPositionInputs::SetGridTensors(const std::shared_ptr& image_grid_thw, @@ -146,6 +164,11 @@ void Qwen2VLPositionInputs::Add() { if (has_mask_input_) { AddAttentionMask(); } + if (rotary_history_) { + rotary_input_index_ = state_.inputs_.size(); + state_.inputs_.push_back(rotary_history_input_->GetOrtTensor()); + state_.input_names_.push_back(model_.config_->model.decoder.inputs.rotary_position_ids.c_str()); + } } void Qwen2VLPositionInputs::AddPositionIDs() { @@ -347,6 +370,7 @@ void Qwen2VLPositionInputs::CreateAndInitialize3DPositionIDs(DeviceSpan rope_deltas_.push_back(max_pos_for_batch + 1 - seq_len); } + UpdateRotaryHistory(*position_ids, 0, state_.params_->search.num_beams); // Move tensor to GPU and expand by num_beams position_ids_->ort_tensor_ = model_.ExpandInputs(position_ids, state_.params_->search.num_beams); position_ids_shape_[1] *= state_.params_->search.num_beams; @@ -398,6 +422,7 @@ void Qwen2VLPositionInputs::Update3DPositionIDs(int base_pos) { DispatchOnType(type_, UpdatePositionIdsFunctor{position_ids.get(), base_pos, batch_size, seq_len, rope_deltas_}); + UpdateRotaryHistory(*position_ids, base_pos, 1); position_ids_->ort_tensor_ = model_.ExpandInputs(position_ids, 1); // No beam expansion needed, already expanded state_.inputs_[posid_input_index_] = position_ids_->GetOrtTensor(); } @@ -414,7 +439,9 @@ void Qwen2VLPositionInputs::UpdateAttentionMask() { void Qwen2VLPositionInputs::Update(DeviceSpan next_tokens, int total_length, int new_length) { if (has_posid_input_) { position_ids_shape_[2] = new_length; - if (is_first_update_) { + if (rotary_source_) { + UpdateFromRotarySource(total_length, new_length); + } else if (is_first_update_) { DispatchOnType(type_, InitPositionIdsFunctor{this, next_tokens, position_ids_shape_}); } else { Update3DPositionIDs(total_length - new_length); @@ -445,6 +472,74 @@ void Qwen2VLPositionInputs::RewindTo(size_t index) { if (has_mask_input_) { attention_mask_shape_[1] = static_cast(index); } + if (rotary_history_) { + rotary_history_->Forecast(index); + PublishRotaryHistory(); + } + if (index == 0) { + is_first_update_ = true; + position_ids_shape_[1] = state_.params_->search.batch_size; + attention_mask_shape_[0] = state_.params_->search.batch_size; + rope_deltas_.clear(); + } +} + +void Qwen2VLPositionInputs::PublishRotaryHistory() { + auto destination = rotary_history_input_->GetDeviceSpan(); + auto host = destination.CpuSpan(); + std::copy(rotary_history_->Values().begin(), rotary_history_->Values().end(), host.begin()); + destination.CopyCpuToDevice(); + state_.inputs_[rotary_input_index_] = rotary_history_input_->GetOrtTensor(); +} + +void Qwen2VLPositionInputs::UpdateRotaryHistory(OrtValue& positions, size_t start, size_t beams) { + if (!rotary_history_) return; + auto info = positions.GetTensorTypeAndShapeInfo(); + const auto shape = info->GetShape(); + rotary_history_->Update( + {positions.GetTensorData(), info->GetElementCount()}, start, shape[2], beams); + PublishRotaryHistory(); +} + +void Qwen2VLPositionInputs::UpdateFromRotarySource(int total_length, int new_length) { + if (!rotary_history_ || total_length < new_length || new_length <= 0) { + throw std::runtime_error("MTP rotary positions require a configured history and a valid step length"); + } + auto info = rotary_source_->GetTensorTypeAndShapeInfo(); + const auto shape = info->GetShape(); + if (rotary_source_->GetTensorMemoryInfo().GetDeviceType() != OrtMemoryInfoDeviceType_CPU || + info->GetElementType() != Ort::TypeToTensorType || + shape.size() != 3 || shape[0] != 3 || shape[1] != state_.params_->BatchBeamSize() || shape[2] <= 0) { + throw std::runtime_error("Main rotary position history must be a CPU int64 [3, batch, history] tensor"); + } + rotary_history_->CopyFrom({rotary_source_->GetTensorData(), info->GetElementCount()}, shape[2]); + if (total_length > state_.params_->search.max_length) { + throw std::runtime_error("MTP step exceeds rotary history capacity"); + } + position_ids_shape_[1] = shape[1]; + auto positions = OrtValue::CreateTensor(model_.allocator_cpu_, position_ids_shape_, Ort::TypeToTensorType); + auto* data = positions->GetTensorMutableData(); + const size_t capacity = state_.params_->search.max_length; + for (size_t row = 0; row < static_cast(3 * shape[1]); ++row) { + const auto source = rotary_history_->Values().subspan(row * capacity + total_length - new_length, new_length); + std::copy(source.begin(), source.end(), data + row * new_length); + } + if (state_.params_->use_graph_capture && new_length <= state_.params_->max_graph_capture_length) { + auto& buffer = captured_rotary_positions_[new_length]; + if (!buffer) { + buffer = std::make_unique(model_.p_device_inputs_, Ort::TypeToTensorType); + buffer->CreateTensor(position_ids_shape_, true); + } + auto destination = buffer->GetDeviceSpan(); + std::copy_n(positions->GetTensorData(), positions->GetTensorTypeAndShapeInfo()->GetElementCount(), + destination.CpuSpan().begin()); + destination.CopyCpuToDevice(); + state_.inputs_[posid_input_index_] = buffer->GetOrtTensor(); + } else { + position_ids_->ort_tensor_ = model_.ExpandInputs(positions, 1); + state_.inputs_[posid_input_index_] = position_ids_->GetOrtTensor(); + } + PublishRotaryHistory(); } } // namespace Generators diff --git a/src/models/io/qwen_vl_position_inputs.h b/src/models/io/qwen_vl_position_inputs.h index 08ac4bbd42..c6786d6ae9 100644 --- a/src/models/io/qwen_vl_position_inputs.h +++ b/src/models/io/qwen_vl_position_inputs.h @@ -1,6 +1,8 @@ #pragma once #include "position_inputs.h" +#include "rotary_position_history.h" +#include namespace Generators { @@ -47,6 +49,8 @@ struct Qwen2VLPositionInputs : PositionInputs { void Add() override; void Update(DeviceSpan next_tokens, int total_length, int new_length) override; void RewindTo(size_t index) override; + OrtValue* GetRotaryPositionIds() override { return rotary_history_value_.get(); } + void SetRotaryPositionIds(OrtValue* positions) override { rotary_source_ = positions; } void SetGridTensors(const std::shared_ptr& image_grid_thw, const std::shared_ptr& video_grid_thw, @@ -64,6 +68,9 @@ struct Qwen2VLPositionInputs : PositionInputs { template void CreateAndInitializeAttentionMask(DeviceSpan next_tokens, std::array shape); void UpdateAttentionMask(); + void UpdateRotaryHistory(OrtValue& positions, size_t start, size_t beams); + void PublishRotaryHistory(); + void UpdateFromRotarySource(int total_length, int new_length); const Model& model_; State& state_; @@ -76,6 +83,12 @@ struct Qwen2VLPositionInputs : PositionInputs { std::unique_ptr position_ids_; std::array attention_mask_shape_{}; std::unique_ptr attention_mask_; + std::unique_ptr rotary_history_; + std::unique_ptr rotary_history_value_; + std::unique_ptr rotary_history_input_; + size_t rotary_input_index_{}; + OrtValue* rotary_source_{}; + std::map> captured_rotary_positions_; bool is_first_update_{true}; std::shared_ptr image_grid_thw_; std::shared_ptr video_grid_thw_; diff --git a/src/models/io/recurrent_state.cpp b/src/models/io/recurrent_state.cpp index 486ca9988f..091fa55301 100644 --- a/src/models/io/recurrent_state.cpp +++ b/src/models/io/recurrent_state.cpp @@ -97,11 +97,9 @@ RecurrentState::RecurrentState(State& state, bool graph_capture_variants_support // legacy unwindowed layout (W == 1). const bool conv_windowed = conv_shape_.size() == 4; const bool recurrent_windowed = recurrent_shape_.size() == 5; - if (conv_windowed != recurrent_windowed) - throw std::runtime_error("RecurrentState: conv_state and recurrent_state must both be windowed or neither"); + if (recurrent_windowed && (!conv_windowed || conv_shape_[0] != recurrent_shape_[0])) + throw std::runtime_error("RecurrentState: convolution and recurrent windows must agree"); if (conv_windowed) { - if (conv_shape_[0] != recurrent_shape_[0]) - throw std::runtime_error("RecurrentState: conv_state and recurrent_state must use the same state window"); state_window_ = conv_shape_[0]; } @@ -146,6 +144,62 @@ RecurrentState::RecurrentState(State& state, bool graph_capture_variants_support presents_.push_back(OrtValue::CreateTensor(allocator, recurrent_shape_, recurrent_type_)); } + const auto& capsule_template = + model_.config_->model.decoder.outputs.state_update_recurrent_capsule_names; + if (model_.config_->model.decoder.state_update_capacity > 0) { + if (!conv_windowed || recurrent_shape_.size() != 4 || + recurrent_shape_[0] != 1) { + throw std::runtime_error( + "RecurrentState compact updates require batch size 1 and windowed convolution state"); + } + state_update_capacity_ = + static_cast(model_.config_->model.decoder.state_update_capacity); + if (state_update_capacity_ + 1 != static_cast(state_window_)) { + throw std::runtime_error( + "RecurrentState compact update capacity must be one less than the convolution state window"); + } + const std::array capture_shape{1}; + state_update_capture_count_ = OrtValue::CreateTensor( + allocator, capture_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32); + auto capture_count = + WrapTensor(*model_.p_device_kvcache_, *state_update_capture_count_); + capture_count.CpuSpan()[0] = static_cast(state_update_capacity_); + capture_count.CopyCpuToDevice(); + state_update_active_ = OrtValue::CreateTensor( + model_.allocator_cpu_, capture_shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32); + state_update_active_->GetTensorMutableData()[0] = 0; + + for (int layer_id : layer_indices_) { + const auto name = ComposeCacheName(capsule_template, layer_id); + output_name_strings_.push_back(name); + auto shape = model_.session_info_.GetOutputShape(name); + if (shape.size() == 2 && shape[0] <= 0) { + shape[0] = state_.params_->BatchBeamSize(); + } + if (shape.size() != 2 || shape[0] <= 0 || shape[1] <= 0) { + throw std::runtime_error("RecurrentState compact update capsule has an invalid shape"); + } + state_update_capsules_.push_back(OrtValue::CreateTensor( + allocator, shape, ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT)); + const size_t channels = static_cast(recurrent_shape_[1]); + const size_t state_width = static_cast(recurrent_shape_[2]); + const size_t key_width = static_cast(recurrent_shape_[3]); + const size_t capsule_width = static_cast(shape[1]); + if (capsule_width % state_update_capacity_ != 0) { + throw std::runtime_error( + "RecurrentState compact update capsule width is not divisible by its capacity"); + } + const size_t transition_width = capsule_width / state_update_capacity_; + const size_t fixed_width = channels + channels * state_width; + if (transition_width <= fixed_width || + (transition_width - fixed_width) % key_width != 0) { + throw std::runtime_error("RecurrentState compact update capsule width is inconsistent with state shape"); + } + state_update_key_head_counts_.push_back( + (transition_width - fixed_width) / key_width); + } + } + if (!share_buffers_) { ZeroStates(pasts_); } @@ -167,13 +221,24 @@ void RecurrentState::Add() { state_.outputs_.push_back(presents_[i].get()); state_.output_names_.push_back(output_name_strings_[i].c_str()); } + if (!state_update_capsules_.empty()) { + const auto& inputs = model_.config_->model.decoder.inputs; + state_.inputs_.push_back(state_update_capture_count_.get()); + state_.input_names_.push_back(inputs.state_update_capture_count.c_str()); + state_.inputs_.push_back(state_update_active_.get()); + state_.input_names_.push_back(inputs.state_update_active.c_str()); + for (size_t index = 0; index < state_update_capsules_.size(); ++index) { + state_.outputs_.push_back(state_update_capsules_[index].get()); + state_.output_names_.push_back(output_name_strings_[num_layers * 2 + index].c_str()); + } + } } void RecurrentState::SetForwardLength(int sequence_length) { forward_length_ = sequence_length; } -void RecurrentState::CropToPosition(size_t position) { +void RecurrentState::ValidateCropPosition(size_t position) const { if (!IsWindowed()) throw std::runtime_error( "RecurrentState::CropToPosition requires the model exported with state_window > 1"); @@ -191,21 +256,64 @@ void RecurrentState::CropToPosition(size_t position) { throw std::runtime_error( "RecurrentState::CropToPosition(" + std::to_string(position) + ") is past the last position of a forward of length " + std::to_string(forward_length_)); - const size_t slot = static_cast(signed_slot); - if (slot + 1 == static_cast(state_window_)) return; // Already the committed slot. + if (!state_update_capsules_.empty() && + (forward_length_ < 2 || position + 1 >= static_cast(forward_length_) || + position + 1 > state_update_capacity_ || !snapshot_valid_)) { + throw std::runtime_error("RecurrentState compact commit position is outside the captured verify prefix"); + } +} +void RecurrentState::CropToPosition(size_t position) { + ValidateCropPosition(position); + const int64_t signed_slot = static_cast(position) + state_window_ - forward_length_; + const size_t slot = static_cast(signed_slot); auto& device = *model_.p_device_kvcache_; // Fast path: one kernel for all 2*num_layers tensors. The per-tensor loop below issues one // cudaMemcpyAsync each, and at 30 layers that is 60 host-side driver calls on a step that is // already GPU-idle-bound. - if (TryBatchedSlotPromote(slot)) return; + if (slot + 1 != static_cast(state_window_) && + !TryBatchedSlotPromote(slot)) { + const size_t increment = recurrent_shape_.size() == 5 ? 1 : 2; + for (size_t index = 0; index < presents_.size(); index += increment) { + auto& present = presents_[index]; + auto window = ByteWrapTensor(device, *present); + const size_t slot_bytes = window.size() / static_cast(state_window_); + window.subspan((static_cast(state_window_) - 1) * slot_bytes, slot_bytes) + .CopyFrom(window.subspan(slot * slot_bytes, slot_bytes)); + } + } - for (auto& present : presents_) { - auto window = ByteWrapTensor(device, *present); - const size_t slot_bytes = window.size() / static_cast(state_window_); - window.subspan((static_cast(state_window_) - 1) * slot_bytes, slot_bytes) - .CopyFrom(window.subspan(slot * slot_bytes, slot_bytes)); + const size_t kept_count = position + 1; + if (state_update_capsules_.empty()) return; + std::vector descriptors; + descriptors.reserve(layer_indices_.size()); + for (size_t layer = 0; layer < layer_indices_.size(); ++layer) { + auto capsule = ByteWrapTensor(device, *state_update_capsules_[layer]); + const size_t channels = static_cast(recurrent_shape_[1]); + const size_t state_width = static_cast(recurrent_shape_[2]); + const size_t key_width = static_cast(recurrent_shape_[3]); + const size_t key_heads = state_update_key_head_counts_[layer]; + const size_t decay_bytes = state_update_capacity_ * channels * sizeof(float); + const size_t key_bytes = state_update_capacity_ * key_heads * key_width * sizeof(float); + const auto* capsule_data = reinterpret_cast(capsule.Span().data()); + descriptors.push_back(StateUpdateReplayDesc{ + ByteWrapTensor(device, *snapshot_[layer * 2 + 1]).Span().data(), + ByteWrapTensor(device, *presents_[layer * 2 + 1]).Span().data(), + nullptr, + capsule_data, + capsule_data + decay_bytes / sizeof(float), + capsule_data + (decay_bytes + key_bytes) / sizeof(float), + static_cast(channels), + static_cast(state_width), + static_cast(key_width), + static_cast(key_heads), + static_cast(state_update_capacity_), + static_cast(kept_count), + static_cast(sizeof(float)), + StateUpdateReplayKind::GatedDeltaNet, + }); } + device.ReplayStateUpdates(descriptors.data(), descriptors.size()); } bool RecurrentState::TryBatchedSlotPromote(size_t slot) { @@ -213,10 +321,12 @@ bool RecurrentState::TryBatchedSlotPromote(size_t slot) { // The state buffers are stable across steps when inputs alias outputs (which graph capture // requires), but re-derive the descriptors and compare so a reallocation cannot go unnoticed. - bool descriptors_changed = slot_descs_cpu_.size() != presents_.size(); - slot_descs_cpu_.resize(presents_.size()); - for (size_t i = 0; i < presents_.size(); ++i) { - auto& present = presents_[i]; + const size_t stride = recurrent_shape_.size() == 5 ? 1 : 2; + const size_t tensor_count = presents_.size() / stride; + bool descriptors_changed = slot_descs_cpu_.size() != tensor_count; + slot_descs_cpu_.resize(tensor_count); + for (size_t i = 0; i < tensor_count; ++i) { + auto& present = presents_[i * stride]; auto window = ByteWrapTensor(device, *present); auto bytes = window.Span(); const StateSlotDesc desc{reinterpret_cast(bytes.data()), @@ -248,7 +358,23 @@ void RecurrentState::Update() { if (graph_double_buffer_) graph_buffer_variant_ ^= 1; } +void RecurrentState::ValidateRewindTo(size_t index) const { + if (layer_indices_.empty() || index == 0) return; + if (snapshot_valid_ && index != snapshot_position_) { + throw std::runtime_error( + "RecurrentState::RewindTo(" + std::to_string(index) + ") cannot restore snapshot at length " + + std::to_string(snapshot_position_)); + } + if (!snapshot_valid_) { + throw std::runtime_error( + "RecurrentState::RewindTo(" + std::to_string(index) + + ") is not supported without a snapshot. Recurrent states cannot be partially rewound; " + "call Snapshot() at the target length first (e.g. for speculative decoding)."); + } +} + void RecurrentState::RewindTo(size_t index) { + ValidateRewindTo(index); if (layer_indices_.empty()) return; if (index != 0) { @@ -257,22 +383,14 @@ void RecurrentState::RewindTo(size_t index) { // snapshot captured at the target length (used by speculative decoding to // roll back a rejected draft). The caller (e.g. the MTP orchestrator) // guarantees the snapshot was taken at exactly `index`. - if (snapshot_valid_) { - if (index != snapshot_position_) { - throw std::runtime_error( - "RecurrentState::RewindTo(" + std::to_string(index) + ") cannot restore snapshot at length " + - std::to_string(snapshot_position_)); - } - RestoreSnapshot(); - return; - } - throw std::runtime_error( - "RecurrentState::RewindTo(" + std::to_string(index) + - ") is not supported without a snapshot. Recurrent states cannot be partially rewound; " - "call Snapshot() at the target length first (e.g. for speculative decoding)."); + RestoreSnapshot(); + return; } // Full reset to length 0. snapshot_valid_ = false; + if (state_update_active_) { + state_update_active_->GetTensorMutableData()[0] = 0; + } if (share_buffers_) { // Shared buffers: zero in place, addresses stay stable. ZeroStates(presents_); @@ -327,6 +445,9 @@ void RecurrentState::Snapshot(size_t position) { } } CopyStates(presents_, snapshot_); + if (state_update_active_) { + state_update_active_->GetTensorMutableData()[0] = 1; + } snapshot_position_ = position; snapshot_valid_ = true; } @@ -339,6 +460,9 @@ void RecurrentState::RestoreSnapshot() { // Copy back into the live buffers in place so their addresses stay stable // (required by CUDA-graph replay, which captures fixed buffer pointers). CopyStates(snapshot_, presents_); + if (state_update_active_) { + state_update_active_->GetTensorMutableData()[0] = 0; + } } bool RecurrentState::ShouldFixUpGraphCapture(int graph_id) const { diff --git a/src/models/io/recurrent_state.h b/src/models/io/recurrent_state.h index c20f32a152..1b9590fea6 100644 --- a/src/models/io/recurrent_state.h +++ b/src/models/io/recurrent_state.h @@ -15,6 +15,7 @@ struct RecurrentState { void Add(); void Update(); void RewindTo(size_t index); + void ValidateRewindTo(size_t index) const; // Snapshot/restore the recurrent (conv + recurrent) state buffers. Required for // speculative decoding (e.g. MTP): the recurrent state cannot be partially rewound @@ -31,13 +32,21 @@ struct RecurrentState { // the ops read back. On a partial-accept MTP step the controller crops the live state to the // accepted length by copying slot `position` into slot W-1 -- no full-cost main-model replay // forward, and no extra graph outputs to bind or keep alive across CUDA-graph capture. - bool IsWindowed() const { return state_window_ > 1; } + bool IsWindowed() const { + return state_window_ > 1 && (recurrent_shape_.size() == 5 || !state_update_capsules_.empty()); + } + bool CropRequiresSnapshot() const { return !state_update_capsules_.empty(); } int64_t StateWindow() const { return state_window_; } void SetForwardLength(int sequence_length); // Record this step's seq_len (maps position -> slot). void CropToPosition(size_t position); // Copy window slot for `position` -> slot W-1. + void ValidateCropPosition(size_t position) const; + bool StateUpdateCaptureEnabled() const { + return state_update_active_ && state_update_active_->GetTensorData()[0] != 0; + } bool IsEmpty() const { return layer_indices_.empty(); } int GraphCaptureVariant() const { return graph_buffer_variant_; } + bool UsesGraphCaptureDoubleBuffer() const { return graph_double_buffer_; } // ORT captures a CUDA graph by re-running the model inside a single user-visible Run() // until the EP reports capture complete (InferenceSession::RunImpl recursion, driven by @@ -70,8 +79,13 @@ struct RecurrentState { std::vector> pasts_; std::vector> presents_; std::vector> snapshot_; // Lazily-allocated copy of the live state for speculative rollback. - bool snapshot_valid_{false}; // Whether snapshot_ holds a valid captured state. - size_t snapshot_position_{}; // Sequence length represented by snapshot_. + std::vector> state_update_capsules_; + std::unique_ptr state_update_capture_count_; + std::unique_ptr state_update_active_; + std::vector state_update_key_head_counts_; + size_t state_update_capacity_{}; + bool snapshot_valid_{false}; // Whether snapshot_ holds a valid captured state. + size_t snapshot_position_{}; // Sequence length represented by snapshot_. // Per-position state window. When the model declares a rank-4 conv_state / rank-5 // recurrent_state, axis 1 is a static window of W per-token states (see IsWindowed above). diff --git a/src/models/io/rotary_position_history.h b/src/models/io/rotary_position_history.h new file mode 100644 index 0000000000..3ad3412255 --- /dev/null +++ b/src/models/io/rotary_position_history.h @@ -0,0 +1,95 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace Generators { + +class RotaryPositionHistory { + public: + RotaryPositionHistory(size_t batch_size, size_t capacity, + int64_t max_position = std::numeric_limits::max()) + : batch_size_{batch_size}, capacity_{capacity}, max_position_{max_position} { + if (batch_size == 0 || capacity == 0 || + batch_size > std::numeric_limits::max() / 3 / capacity || max_position < 0) { + throw std::runtime_error("Rotary position history requires a positive, representable batch and capacity"); + } + values_.resize(3 * batch_size * capacity); + } + + std::span Values() { return values_; } + + void Update(std::span positions, size_t start, size_t length, size_t beams = 1) { + if (beams == 0 || batch_size_ % beams != 0 || length == 0 || + start > capacity_ || length > capacity_ - start || + positions.size() != 3 * (batch_size_ / beams) * length) { + throw std::runtime_error("Current T/H/W positions do not fit the rotary history"); + } + const size_t source_batch = batch_size_ / beams; + if (std::any_of(positions.begin(), positions.end(), + [&](int64_t value) { return value < 0 || value > max_position_; })) { + throw std::runtime_error("Rotary coordinates exceed the available cache table"); + } + for (size_t plane = 0; plane < 3; ++plane) { + for (size_t row = 0; row < batch_size_; ++row) { + auto source = positions.subspan((plane * source_batch + row / beams) * length, length); + std::copy(source.begin(), source.end(), values_.begin() + (plane * batch_size_ + row) * capacity_ + start); + } + } + Forecast(start + length); + } + + void CopyFrom(std::span source, size_t source_capacity) { + if (source_capacity < capacity_ || source_capacity > std::numeric_limits::max() / 3 / batch_size_ || + source.size() != 3 * batch_size_ * source_capacity) { + throw std::runtime_error("Main and MTP rotary position histories have incompatible shapes"); + } + for (size_t row = 0; row < 3 * batch_size_; ++row) { + std::copy_n(source.begin() + row * source_capacity, capacity_, values_.begin() + row * capacity_); + } + } + + void Forecast(size_t length) { + if (length > capacity_) { + throw std::runtime_error("Rotary history rewind exceeds its capacity"); + } + for (size_t row = 0; row < batch_size_; ++row) { + int64_t next = 0; + if (length != 0) { + for (size_t plane = 0; plane < 3; ++plane) { + auto begin = values_.begin() + (plane * batch_size_ + row) * capacity_; + const int64_t maximum = *std::max_element(begin, begin + length); + if (maximum == std::numeric_limits::max()) { + throw std::runtime_error("Rotary coordinates exceed the supported integer range"); + } + next = std::max(next, maximum + 1); + } + } + next = std::min(next, max_position_); + for (size_t plane = 0; plane < 3; ++plane) { + auto begin = values_.begin() + (plane * batch_size_ + row) * capacity_; + for (size_t token = length; token < capacity_; ++token) { + const size_t offset = token - length; + begin[token] = offset > static_cast(max_position_ - next) + ? max_position_ + : next + static_cast(offset); + } + } + } + } + + private: + size_t batch_size_; + size_t capacity_; + int64_t max_position_; + std::vector values_; +}; + +} // namespace Generators diff --git a/src/models/io/shared_kv_cache.cpp b/src/models/io/shared_kv_cache.cpp index 7fd17a4dca..40c2eb7836 100644 --- a/src/models/io/shared_kv_cache.cpp +++ b/src/models/io/shared_kv_cache.cpp @@ -13,6 +13,7 @@ void SharedKeyValueCache::Update(DeviceSpan /*beam_indices*/, int total void SharedKeyValueCache::RewindTo(size_t index) { CheckWindowedKvCacheRewind(windowed_cache_size_, current_length_, index); + current_length_ = static_cast(index); } } // namespace Generators diff --git a/src/models/model.cpp b/src/models/model.cpp index 4fb19813a0..faf9008b87 100644 --- a/src/models/model.cpp +++ b/src/models/model.cpp @@ -95,7 +95,8 @@ State::State(const GeneratorParams& params, const Model& model) } const std::string& State::GraphIdForLength(int graph_capture_length, int graph_capture_variant) { - const int graph_key = graph_capture_length * 2 + graph_capture_variant; + const uint64_t graph_key = (static_cast(static_cast(graph_capture_length)) << 32) | + static_cast(graph_capture_variant); auto it = graph_ids_.find(graph_key); if (it != graph_ids_.end()) { return it->second.text; @@ -632,6 +633,10 @@ void Model::CreateSessionOptionsFromConfig(const Config::SessionOptions& config_ session_options.SetInterOpNumThreads(config_session_options.inter_op_num_threads.value()); } + if (config_session_options.use_deterministic_compute.has_value()) { + session_options.SetDeterministicCompute(config_session_options.use_deterministic_compute.value()); + } + if (config_session_options.enable_cpu_mem_arena.has_value()) { if (config_session_options.enable_cpu_mem_arena.value()) session_options.EnableCpuMemArena(); diff --git a/src/models/model.h b/src/models/model.h index 42c540f723..ea59876fc8 100644 --- a/src/models/model.h +++ b/src/models/model.h @@ -30,6 +30,7 @@ struct State { virtual void Finalize(int current_length) {} virtual void RewindTo(size_t index) { (void)index; }; + virtual void ValidateRewindTo(size_t index) const { (void)index; } // Snapshot/restore the model's recurrent state for speculative decoding. Default no-op // for models without recurrent state. SnapshotState() captures the conv/linear-attention @@ -44,6 +45,7 @@ struct State { // attention KV cache + position to `new_length` and crops the recurrent state to the state // AFTER verify token `recurrent_position` (that token's window slot). virtual bool HasCroppableRecurrentState() const { return false; } + virtual bool CropRequiresSnapshot() const { return false; } // Width of that window, i.e. the longest forward whose per-token states all stay addressable. // 1 when the model carries no window. virtual int64_t RecurrentStateWindow() const { return 1; } @@ -51,10 +53,16 @@ struct State { (void)new_length; (void)recurrent_position; } + virtual void ValidateCropToAccepted(size_t new_length, size_t recurrent_position) const { + (void)new_length; + (void)recurrent_position; + } // Stage a hidden_states input value for the next Run (models with a hidden_states input, // e.g. the MTP self-speculative head). Default no-op. virtual void SetHiddenStates(OrtValue* hidden_states) { (void)hidden_states; } + virtual OrtValue* GetRotaryPositionIds() { return nullptr; } + virtual void SetRotaryPositionIds(OrtValue* positions) { (void)positions; } virtual OrtValue* GetInput(const char* name); virtual OrtValue* GetOutput(const char* name); @@ -96,7 +104,7 @@ struct State { std::string text; // string form passed to the "gpu_graph_id" run option }; const std::string& GraphIdForLength(int graph_capture_length, int graph_capture_variant); - std::unordered_map graph_ids_{}; + std::unordered_map graph_ids_{}; // Session used for graph capture; not owned. Lifetime invariant: the OrtSession // outlives this State because State is owned by Generator, and Generator is // destroyed before the Model (and its session) that produced it. diff --git a/src/models/model_state_manifest.cpp b/src/models/model_state_manifest.cpp index dec62589b1..9a0cbd2fab 100644 --- a/src/models/model_state_manifest.cpp +++ b/src/models/model_state_manifest.cpp @@ -37,6 +37,17 @@ std::vector StateBindingsFor( return {{"state", inputs.past_conv_names, outputs.present_conv_names}}; case StateGroupKind::FixedRecurrent: return {{"state", inputs.past_recurrent_names, outputs.present_recurrent_names}}; + case StateGroupKind::FixedPle: + return { + {"tokens", inputs.past_ple_token_names, outputs.present_ple_token_names}, + {"conv", inputs.past_ple_conv_names, outputs.present_ple_conv_names}, + }; + case StateGroupKind::FixedIndexer: + return { + {"key", inputs.past_indexer_names, outputs.present_indexer_names}, + {"kv_buffer", inputs.past_indexer_kv_buffer_names, outputs.present_indexer_kv_buffer_names}, + {"state_lengths", inputs.past_indexer_state_lengths_names, outputs.present_indexer_state_lengths_names}, + }; case StateGroupKind::Invalid: return {}; } @@ -45,7 +56,16 @@ std::vector StateBindingsFor( bool IsFixedKind(StateGroupKind kind) { return kind == StateGroupKind::FixedConv || - kind == StateGroupKind::FixedRecurrent; + kind == StateGroupKind::FixedRecurrent || + kind == StateGroupKind::FixedPle || + kind == StateGroupKind::FixedIndexer; +} + +bool SupportsStateUpdate(StateGroupKind kind) { + return kind == StateGroupKind::FixedConv || + kind == StateGroupKind::FixedRecurrent || + kind == StateGroupKind::FixedPle || + kind == StateGroupKind::FixedIndexer; } StateUpdateKind StateUpdateKindFor(StateGroupKind kind) { @@ -55,6 +75,12 @@ StateUpdateKind StateUpdateKindFor(StateGroupKind kind) { if (kind == StateGroupKind::FixedRecurrent) { return StateUpdateKind::GatedDeltaNet; } + if (kind == StateGroupKind::FixedPle) { + return StateUpdateKind::Ple; + } + if (kind == StateGroupKind::FixedIndexer) { + return StateUpdateKind::Indexer; + } return StateUpdateKind::Invalid; } @@ -66,6 +92,10 @@ std::string_view StateGroupKindName(StateGroupKind kind) { return "fixed_conv"; case StateGroupKind::FixedRecurrent: return "fixed_recurrent"; + case StateGroupKind::FixedPle: + return "fixed_ple"; + case StateGroupKind::FixedIndexer: + return "fixed_indexer"; case StateGroupKind::Invalid: return "invalid"; } @@ -246,38 +276,62 @@ void ValidateStateUpdateSession(std::string_view group_label, output_name}; }; - const auto state_binding = StateBindingsFor(inputs, outputs, group.kind).front(); + const auto state_bindings = StateBindingsFor(inputs, outputs, group.kind); for (const int layer_id : group.layer_ids) { - const auto state_name = ExpandBinding(state_binding.output, layer_id); + const auto state_name = ExpandBinding(state_bindings.front().output, layer_id); const TensorMetadata state{ metadata.GetOutputDataType(state_name), metadata.GetOutputShape(state_name), state_name}; - ValidateRank( - group_label, "state output", state, - update_kind == StateUpdateKind::CausalConv ? 3 : 4); + const size_t state_rank = update_kind == StateUpdateKind::GatedDeltaNet + ? 4 + : update_kind == StateUpdateKind::Ple ? 2 : 3; + ValidateRank(group_label, "state output", state, state_rank); if (!DimensionsCompatible(state.shape[0], capture_count.shape[0])) { throw std::runtime_error( std::string{group_label} + " state_update capture_count input '" + capture_count_name + "' has batch dimension incompatible with state output '" + state.name + "'"); } - if (update_kind == StateUpdateKind::CausalConv) { - const auto value = get_output("value", outputs.state_update_conv_value_names, layer_id); + if (update_kind == StateUpdateKind::CausalConv || update_kind == StateUpdateKind::Ple) { + const auto& value_template = update_kind == StateUpdateKind::CausalConv + ? outputs.state_update_conv_value_names + : outputs.state_update_ple_conv_value_names; + const TensorMetadata conv_state = update_kind == StateUpdateKind::Ple + ? TensorMetadata{ + metadata.GetOutputDataType(ExpandBinding(state_bindings[1].output, layer_id)), + metadata.GetOutputShape(ExpandBinding(state_bindings[1].output, layer_id)), + ExpandBinding(state_bindings[1].output, layer_id)} + : state; + ValidateRank(group_label, "convolution state output", conv_state, 3); + const auto value = get_output("value", value_template, layer_id); ValidateRank(group_label, "state_update value output", value, 3); - if (value.data_type != state.data_type) { + if (value.data_type != conv_state.data_type) { throw std::runtime_error( std::string{group_label} + " state_update value output '" + value.name + - "' must have the same dtype as state output '" + state.name + "'"); + "' must have the same dtype as state output '" + conv_state.name + "'"); } - if (!DimensionsCompatible(state.shape[0], value.shape[0]) || + if (!DimensionsCompatible(conv_state.shape[0], value.shape[0]) || value.shape[1] != update.capacity || - !DimensionsCompatible(state.shape[1], value.shape[2])) { + !DimensionsCompatible(conv_state.shape[1], value.shape[2])) { throw std::runtime_error( std::string{group_label} + " state_update value output '" + value.name + "' has incompatible batch, capacity, or channel dimensions"); } - } else { + if (update_kind == StateUpdateKind::Ple && !outputs.state_update_ple_token_names.empty()) { + const auto tokens = get_output("tokens", outputs.state_update_ple_token_names, layer_id); + ValidateRank(group_label, "state_update tokens output", tokens, 3); + if (state.data_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 || + tokens.data_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64 || + !DimensionsCompatible(state.shape[0], tokens.shape[0]) || + tokens.shape[1] != update.capacity || + !DimensionsCompatible(state.shape[1], tokens.shape[2])) { + throw std::runtime_error( + std::string{group_label} + " state_update tokens output '" + tokens.name + + "' must be int64 [batch, capacity, token_state_width]"); + } + } + } else if (update_kind == StateUpdateKind::GatedDeltaNet) { const auto capsule = get_output( "capsule", outputs.state_update_recurrent_capsule_names, layer_id); ValidateRank(group_label, "state_update capsule output", capsule, 2); @@ -301,6 +355,25 @@ void ValidateStateUpdateSession(std::string_view group_label, std::string{group_label} + " state_update capsule output '" + capsule.name + "' width must be " + std::to_string(expected_width)); } + } else { + const auto payload = get_output("indexer", outputs.state_update_indexer_names, layer_id); + const auto buffer_name = ExpandBinding(state_bindings[1].output, layer_id); + const auto lengths_name = ExpandBinding(state_bindings[2].output, layer_id); + const TensorMetadata buffer{metadata.GetOutputDataType(buffer_name), metadata.GetOutputShape(buffer_name), buffer_name}; + const TensorMetadata lengths{metadata.GetOutputDataType(lengths_name), metadata.GetOutputShape(lengths_name), lengths_name}; + ValidateRank(group_label, "state_update indexer output", payload, 3); + ValidateRank(group_label, "indexer buffer output", buffer, 3); + ValidateRank(group_label, "indexer lengths output", lengths, 2); + if (payload.data_type != state.data_type || buffer.data_type != state.data_type || + lengths.data_type != ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32 || + payload.shape[1] != update.capacity || lengths.shape[1] != 2 || + buffer.shape[1] != 2 * update.compress_ratio - 1 || + !DimensionsCompatible(state.shape[0], payload.shape[0]) || + !DimensionsCompatible(state.shape[2], payload.shape[2]) || + !DimensionsCompatible(buffer.shape[2], payload.shape[2])) { + throw std::runtime_error( + std::string{group_label} + " has incompatible indexer state_update geometry or dtype"); + } } } } @@ -337,7 +410,9 @@ bool ModelStateManifest::HasStateGroupKind(StateGroupKind kind) const { bool ModelStateManifest::HasFixedStateGroups() const { return HasStateGroupKind(StateGroupKind::FixedConv) || - HasStateGroupKind(StateGroupKind::FixedRecurrent); + HasStateGroupKind(StateGroupKind::FixedRecurrent) || + HasStateGroupKind(StateGroupKind::FixedPle) || + HasStateGroupKind(StateGroupKind::FixedIndexer); } void ModelStateManifest::ValidateConfig(const Decoder& decoder) { @@ -372,8 +447,9 @@ void ModelStateManifest::ValidateConfig(const Decoder& decoder) { if (group.state_update) { ++fixed_state_update_group_count; const auto& update = *group.state_update; - if (!IsFixedKind(group.kind)) { - throw std::runtime_error(group_label + " only fixed state groups support state_update"); + if (!SupportsStateUpdate(group.kind)) { + throw std::runtime_error( + group_label + " does not support state_update; partial acceptance requires replay"); } if (update.capacity < 1 || update.capacity > Decoder::MaxStateUpdateCapacity) { throw std::runtime_error( @@ -398,6 +474,15 @@ void ModelStateManifest::ValidateConfig(const Decoder& decoder) { throw std::runtime_error( group_label + " fixed_recurrent state_update requires key_head_count"); } + if (group.kind != StateGroupKind::FixedRecurrent && update.key_head_count != 0) { + throw std::runtime_error(group_label + " state_update does not use key_head_count"); + } + if (group.kind == StateGroupKind::FixedIndexer && update.compress_ratio <= 0) { + throw std::runtime_error(group_label + " state_update requires compress_ratio"); + } + if (group.kind != StateGroupKind::FixedIndexer && update.compress_ratio != 0) { + throw std::runtime_error(group_label + " state_update does not use compress_ratio"); + } if (!fixed_state_update_capacity) { fixed_state_update_capacity = update.capacity; @@ -452,15 +537,27 @@ void ModelStateManifest::ValidateConfig(const Decoder& decoder) { validate_binding(binding); } if (group.state_update) { - const auto& output_template = group.kind == StateGroupKind::FixedConv - ? decoder.outputs.state_update_conv_value_names - : decoder.outputs.state_update_recurrent_capsule_names; - ValidateBindingTemplate(group_label, "state_update", "output", output_template); - for (const int layer_id : group.layer_ids) { - const auto output_name = ExpandBinding(output_template, layer_id); - if (!expanded_bindings.insert(output_name).second) { - throw std::runtime_error( - group_label + " resolves more than one binding to '" + output_name + "'"); + std::vector output_templates; + if (group.kind == StateGroupKind::FixedConv) { + output_templates = {&decoder.outputs.state_update_conv_value_names}; + } else if (group.kind == StateGroupKind::FixedRecurrent) { + output_templates = {&decoder.outputs.state_update_recurrent_capsule_names}; + } else if (group.kind == StateGroupKind::FixedPle) { + output_templates = {&decoder.outputs.state_update_ple_conv_value_names}; + if (!decoder.outputs.state_update_ple_token_names.empty()) { + output_templates.push_back(&decoder.outputs.state_update_ple_token_names); + } + } else { + output_templates = {&decoder.outputs.state_update_indexer_names}; + } + for (const auto* output_template : output_templates) { + ValidateBindingTemplate(group_label, "state_update", "output", *output_template); + for (const int layer_id : group.layer_ids) { + const auto output_name = ExpandBinding(*output_template, layer_id); + if (!expanded_bindings.insert(output_name).second) { + throw std::runtime_error( + group_label + " resolves more than one binding to '" + output_name + "'"); + } } } } @@ -471,6 +568,17 @@ void ModelStateManifest::ValidateConfig(const Decoder& decoder) { throw std::runtime_error( "All fixed state groups must declare state_update when any fixed group declares it"); } + if (decoder.state_update_capacity > 0) { + if (fixed_state_update_group_count != fixed_group_count) { + throw std::runtime_error( + "All fixed state groups must declare state_update when decoder state updates are enabled"); + } + if (!fixed_state_update_capacity || + *fixed_state_update_capacity != decoder.state_update_capacity) { + throw std::runtime_error( + "Fixed state_update capacity must match model.decoder.state_update_capacity"); + } + } } void ModelStateManifest::ValidateSession(const ModelStateMetadata& metadata) const { diff --git a/src/models/model_type.h b/src/models/model_type.h index d5e9ea07be..5fdb1f16c9 100644 --- a/src/models/model_type.h +++ b/src/models/model_type.h @@ -17,13 +17,13 @@ namespace Generators { struct ModelType { inline static bool IsLLM(const std::string& model_type) { // Large-language model (LLM) - static constexpr std::array LLM = {"chatglm", "decoder", "ernie4_5", "gemma", "gemma2", "gemma3_text", "gemma3_vl_text", "gemma4_text", "gpt2", "gptoss", "granite", "granitemoehybrid", "hunyuandensev1", "internlm2", "lfm2", "lfm2_audio_text", "lfm2_moe", "lfm2_vl_text", "llama", "mistral", "mistral3_text", "nemotron", "olmo", "phi", "phimoe", "phi3", "phi3small", "qwen2", "qwen2_5_vl_text", "qwen3", "qwen3_vl_text", "qwen3_5_moe_text", "qwen3_5_text", "smollm3"}; + static constexpr std::array LLM = {"chatglm", "decoder", "ernie4_5", "gemma", "gemma2", "gemma3_text", "gemma3_vl_text", "gemma4_text", "gpt2", "gptoss", "granite", "granitemoehybrid", "hunyuandensev1", "internlm2", "lfm2", "lfm2_audio_text", "lfm2_moe", "lfm2_vl_text", "llama", "mistral", "mistral3_text", "nemotron", "olmo", "phi", "phimoe", "phi3", "phi3small", "qwen2", "qwen2_5_vl_text", "qwen3", "qwen3_vl_text", "qwen3_5_moe_text", "qwen3_5_text", "qwen4_exp_text", "smollm3"}; return std::find(LLM.begin(), LLM.end(), model_type) != LLM.end(); } inline static bool IsVLM(const std::string& model_type) { // Vision-language model (VLM) - static constexpr std::array VLM = {"fara", "gemma3", "lfm2_vl", "mistral3", "phi3v", "qwen2_5_vl", "qwen3_vl", "qwen3_5", "qwen3_5_moe", "videochat_flash_qwen"}; + static constexpr std::array VLM = {"fara", "gemma3", "lfm2_vl", "mistral3", "phi3v", "qwen2_5_vl", "qwen3_vl", "qwen3_5", "qwen3_5_moe", "qwen4_exp", "videochat_flash_qwen"}; return std::find(VLM.begin(), VLM.end(), model_type) != VLM.end(); } @@ -75,7 +75,7 @@ struct ModelType { inline static bool IsQwenVLFamily(const std::string& model_type) { // Qwen-VL family: models requiring 3D mRoPE position IDs - static constexpr std::array QwenVL = {"fara", "qwen2_5_vl_text", "qwen2_5_vl", "qwen3_vl", "qwen3_vl_text", "qwen3_5", "qwen3_5_moe", "qwen3_5_moe_text", "qwen3_5_text", "videochat_flash_qwen"}; + static constexpr std::array QwenVL = {"fara", "qwen2_5_vl_text", "qwen2_5_vl", "qwen3_vl", "qwen3_vl_text", "qwen3_5", "qwen3_5_moe", "qwen3_5_moe_text", "qwen3_5_text", "qwen4_exp", "qwen4_exp_text", "videochat_flash_qwen"}; return std::find(QwenVL.begin(), QwenVL.end(), model_type) != QwenVL.end(); } diff --git a/src/models/multi_modal.cpp b/src/models/multi_modal.cpp index c1982fb235..96562d9c33 100644 --- a/src/models/multi_modal.cpp +++ b/src/models/multi_modal.cpp @@ -7,6 +7,7 @@ #include "models/io/qwen_vl_position_inputs.h" #include #include +#include #include namespace Generators { @@ -157,6 +158,16 @@ MultiModalLanguageModel::MultiModalLanguageModel(std::unique_ptr config, CreateSessionOptionsFromConfig(config_->model.embedding.session_options.has_value() ? config_->model.embedding.session_options.value() : config_->model.decoder.session_options, *embedding_session_options_, true, /*disable_graph_capture=*/true); embedding_session_ = CreateSession(ort_env, config_->model.embedding.filename, embedding_session_options_.get()); + if (!config_->model.engram.filename.empty()) { + engram_session_options_ = OrtSessionOptions::Create(); + CreateSessionOptionsFromConfig( + config_->model.engram.session_options.has_value() + ? config_->model.engram.session_options.value() + : config_->model.decoder.session_options, + *engram_session_options_, true, /*disable_graph_capture=*/true); + engram_session_ = CreateSession( + ort_env, config_->model.engram.filename, engram_session_options_.get()); + } decoder_session_ = CreateSession(ort_env, config_->model.decoder.filename, session_options_.get()); const auto& audio_output = config_->model.audio_output; @@ -183,6 +194,9 @@ MultiModalLanguageModel::MultiModalLanguageModel(std::unique_ptr config, session_info_.Add(*decoder_session_); session_info_.Add(*embedding_session_); + if (engram_session_) { + session_info_.Add(*engram_session_); + } if (speech) { session_info_.Add(*speech_session_); } @@ -877,8 +891,22 @@ DecoderState::DecoderState(const MultiModalLanguageModel& model, DeviceSpanCreatePositionInputs(*this, sequence_lengths, model_.config_->model.decoder.inputs.attention_mask)}, kv_cache_{model_.p_device_kvcache_->CreateKeyValueCache(*this)}, - recurrent_state_{CreateRecurrentState(*this, /*graph_capture_variants_supported=*/true)} { + recurrent_state_{CreateRecurrentState(*this, /*graph_capture_variants_supported=*/true)}, + ple_state_{CreatePleState(*this)}, + indexer_cache_{CreateIndexerCache(*this)} { inputs_embeds_.Add(); + if (!model_.config_->model.engram.filename.empty()) { + const auto engram_input_shape = model_.session_info_.GetInputShape( + model_.config_->model.decoder.inputs.engram_embeddings); + if (engram_input_shape.empty() || engram_input_shape.back() <= 0) { + throw std::runtime_error("Decoder Engram input must have a static width"); + } + engram_embeddings_ = std::make_unique( + *this, Embeddings::Mode::Input, + model_.config_->model.decoder.inputs.engram_embeddings, + engram_input_shape.back()); + engram_embeddings_->Add(); + } // Gemma4: decoder accepts per_layer_inputs from the embedding model if (!model_.config_->model.decoder.inputs.per_layer_inputs.empty()) { @@ -908,6 +936,14 @@ DecoderState::DecoderState(const MultiModalLanguageModel& model, DeviceSpanAdd(); if (recurrent_state_) recurrent_state_->Add(); + if (ple_state_) + ple_state_->Add(); + if (indexer_cache_) + indexer_cache_->Add(); + if (!model_.config_->model.decoder.outputs.hidden_states.empty()) { + hidden_states_output_ = std::make_unique(*this); + hidden_states_output_->Add(); + } } DeviceSpan DecoderState::Run(int current_length, DeviceSpan& next_tokens, DeviceSpan next_indices) { @@ -916,16 +952,30 @@ DeviceSpan DecoderState::Run(int current_length, DeviceSpan& nex } const int seq_len = static_cast(inputs_embeds_.GetShape()[1]); - const bool graph_capture_this_run = params_->use_graph_capture && seq_len == 1; - const int graph_capture_variant = recurrent_state_ ? recurrent_state_->GraphCaptureVariant() : 0; - - const int graph_id = seq_len * 2 + graph_capture_variant; + if (hidden_states_output_) + hidden_states_output_->Update(seq_len); + const bool graph_capture_this_run = params_->use_graph_capture && seq_len == 1 && + (!indexer_cache_ || !indexer_cache_->UsesHostSequenceLength()); + int graph_capture_variant = recurrent_state_ ? recurrent_state_->GraphCaptureVariant() : 0; + if (ple_state_) { + const int ple_graph_capture_variant = ple_state_->GraphCaptureVariant(); + if (recurrent_state_ && recurrent_state_->UsesGraphCaptureDoubleBuffer() && + graph_capture_variant != ple_graph_capture_variant) + throw std::runtime_error("PLE and recurrent state graph-capture buffer variants are out of sync"); + graph_capture_variant = ple_graph_capture_variant; + } + + if (recurrent_state_ && recurrent_state_->StateUpdateCaptureEnabled()) + graph_capture_variant += 2; + const int graph_id = seq_len * 4 + graph_capture_variant; if (graph_capture_this_run && recurrent_state_ && recurrent_state_->ShouldFixUpGraphCapture(graph_id)) { recurrent_state_->SaveForGraphCapture(); State::Run(*model_.decoder_session_, true, seq_len, graph_capture_variant); recurrent_state_->RestoreAfterGraphCapture(graph_id); } State::Run(*model_.decoder_session_, graph_capture_this_run, seq_len, graph_capture_variant); + if (recurrent_state_) recurrent_state_->SetForwardLength(seq_len); + if (ple_state_) ple_state_->SetForwardLength(seq_len); return logits_.Get(); } @@ -959,7 +1009,8 @@ void DecoderState::PrepareEmbeddingsForPrefill(size_t new_length) { } DeviceSpan DecoderState::RunPrefillWithChunking(int current_length, DeviceSpan& next_tokens, - DeviceSpan next_indices, size_t chunk_size) { + DeviceSpan next_indices, size_t chunk_size, + EngramState* engram_state) { if (model_.config_->model.decoder.run_options.has_value()) { State::SetRunOptions(model_.config_->model.decoder.run_options.value()); } @@ -967,23 +1018,54 @@ DeviceSpan DecoderState::RunPrefillWithChunking(int current_length, Devic const size_t num_tokens = next_tokens.size(); size_t processed_tokens = 0; int length = current_length - static_cast(num_tokens); + std::future next_engram; while (processed_tokens < num_tokens) { const size_t current_chunk_size = std::min(chunk_size, num_tokens - processed_tokens); auto chunk_tokens = next_tokens.subspan(processed_tokens, current_chunk_size); length += static_cast(current_chunk_size); + if (next_engram.valid()) { + next_engram.get(); + } else if (engram_state) { + engram_state->UpdateInputsOutputs(chunk_tokens); + engram_state->Run(length, chunk_tokens, next_indices); + } + if (decoder_input_ids_) decoder_input_ids_->Update(chunk_tokens); position_inputs_->Update(chunk_tokens, length, static_cast(current_chunk_size)); kv_cache_->Update(next_indices, length); if (recurrent_state_) recurrent_state_->Update(); + if (ple_state_) + ple_state_->Update(); + if (indexer_cache_) + indexer_cache_->Update(next_indices, length, static_cast(current_chunk_size)); logits_.Update(chunk_tokens, current_chunk_size); + if (hidden_states_output_) + hidden_states_output_->Update(static_cast(current_chunk_size)); + if (engram_state) { + engram_embeddings_->UpdateSequenceLength(current_chunk_size); + engram_state->CopyEmbeddingsTo(*engram_embeddings_); + } // Feed only this chunk's slice of the pre-computed embeddings to the decoder. inputs_embeds_.UseChunkView(processed_tokens, current_chunk_size); if (per_layer_inputs_) per_layer_inputs_->UseChunkView(processed_tokens, current_chunk_size); + const size_t next_offset = processed_tokens + current_chunk_size; + if (engram_state && next_offset < num_tokens) { + const size_t next_chunk_size = std::min(chunk_size, num_tokens - next_offset); + auto next_chunk_tokens = next_tokens.subspan(next_offset, next_chunk_size); + const int next_length = length + static_cast(next_chunk_size); + next_engram = std::async( + std::launch::async, + [engram_state, next_chunk_tokens, next_indices, next_length]() mutable { + engram_state->UpdateInputsOutputs(next_chunk_tokens); + engram_state->Run(next_length, next_chunk_tokens, next_indices); + }); + } + // Graph capture is disabled during prefill chunking. State::Run(*model_.decoder_session_, /*graph_capture_this_run=*/false); @@ -997,6 +1079,64 @@ DeviceSpan DecoderState::RunPrefillWithChunking(int current_length, Devic return logits_.Get(); } +void DecoderState::ValidateRewindTo(size_t index) const { + if (indexer_cache_) indexer_cache_->ValidateRewindTo(index); + if (recurrent_state_) recurrent_state_->ValidateRewindTo(index); + if (ple_state_) ple_state_->ValidateRewindTo(index); +} + +void DecoderState::RewindTo(size_t index) { + ValidateRewindTo(index); + position_inputs_->RewindTo(index); + if (kv_cache_) + kv_cache_->RewindTo(index); + if (recurrent_state_) + recurrent_state_->RewindTo(index); + if (ple_state_) + ple_state_->RewindTo(index); + if (indexer_cache_) + indexer_cache_->RewindTo(index); +} + +void DecoderState::SnapshotState(size_t position) { + if (recurrent_state_) + recurrent_state_->Snapshot(position); + if (ple_state_) + ple_state_->Snapshot(position); + if (indexer_cache_) + indexer_cache_->Snapshot(position); +} + +bool DecoderState::HasCroppableRecurrentState() const { + return recurrent_state_ && recurrent_state_->IsWindowed() && + !ple_state_ && + (!indexer_cache_ || indexer_cache_->HasStateUpdates()); +} + +int64_t DecoderState::RecurrentStateWindow() const { + return recurrent_state_ ? recurrent_state_->StateWindow() : 1; +} + +void DecoderState::ValidateCropToAccepted(size_t new_length, size_t recurrent_position) const { + (void)new_length; + if (recurrent_state_) recurrent_state_->ValidateCropPosition(recurrent_position); + if (ple_state_) ple_state_->ValidateCropPosition(recurrent_position); + if (indexer_cache_) indexer_cache_->ValidateAcceptedPrefix(recurrent_position + 1); +} + +void DecoderState::CropToAccepted(size_t new_length, size_t recurrent_position) { + ValidateCropToAccepted(new_length, recurrent_position); + position_inputs_->RewindTo(new_length); + if (kv_cache_) + kv_cache_->RewindTo(new_length); + if (recurrent_state_) + recurrent_state_->CropToPosition(recurrent_position); + if (ple_state_) + ple_state_->CropToPosition(recurrent_position); + if (indexer_cache_) + indexer_cache_->CommitAcceptedPrefix(recurrent_position + 1); +} + void DecoderState::UpdateInputsOutputs(DeviceSpan& next_tokens, int total_length, DeviceSpan beam_indices) { int batch_size = static_cast(inputs_embeds_.GetShape()[0]); size_t new_length = next_tokens.size() / batch_size; @@ -1006,8 +1146,13 @@ void DecoderState::UpdateInputsOutputs(DeviceSpan& next_tokens, int tot kv_cache_->Update(beam_indices, total_length); if (recurrent_state_) recurrent_state_->Update(); + if (ple_state_) + ple_state_->Update(); + if (indexer_cache_) + indexer_cache_->Update(beam_indices, total_length, static_cast(new_length)); logits_.Update(next_tokens, new_length); inputs_embeds_.UpdateSequenceLength(new_length); + if (engram_embeddings_) engram_embeddings_->UpdateSequenceLength(new_length); if (per_layer_inputs_) per_layer_inputs_->UpdateSequenceLength(new_length); } @@ -1018,8 +1163,13 @@ void DecoderState::UpdateInputsOutputs(DeviceSpan& next_tokens, int tot kv_cache_->Update(beam_indices, total_length); if (recurrent_state_) recurrent_state_->Update(); + if (ple_state_) + ple_state_->Update(); + if (indexer_cache_) + indexer_cache_->Update(beam_indices, total_length, static_cast(new_length)); logits_.Update(next_tokens, new_length); inputs_embeds_.UpdateSequenceLength(new_length); + if (engram_embeddings_) engram_embeddings_->UpdateSequenceLength(new_length); if (per_layer_inputs_) per_layer_inputs_->UpdateSequenceLength(new_length); } @@ -1034,6 +1184,9 @@ MultiModalPipelineState::MultiModalPipelineState(const MultiModalLanguageModel& speech_state_ = CreateSpeechState(model_, params); } embedding_state_ = std::make_unique(model, params); + if (model_.engram_session_) { + engram_state_ = std::make_unique(model, *model.engram_session_, params); + } decoder_state_ = std::make_unique(model_, sequence_lengths, params); if (model_.depthformer_session_) { audio_output_ = std::make_unique(model_, params); @@ -1100,16 +1253,20 @@ DeviceSpan MultiModalPipelineState::Run(int current_length, DeviceSpansearch.chunk_size; const size_t num_tokens = next_tokens.size(); + auto chunk_size = params_->search.chunk_size.value_or(0); const bool has_multimodal_content = num_image_tokens_ != 0 || num_audio_tokens_ != 0; - const bool chunk_prefill = is_prompt_ && chunk_size_opt.has_value() && chunk_size_opt.value() > 0 && - num_tokens > chunk_size_opt.value() && decoder_state_->SupportsPrefillChunking(has_multimodal_content); + const bool chunk_prefill = + is_prompt_ && chunk_size > 0 && num_tokens > chunk_size && + decoder_state_->SupportsPrefillChunking(has_multimodal_content); if (chunk_prefill) { decoder_state_->PrepareEmbeddingsForPrefill(num_tokens); } else { decoder_state_->UpdateInputsOutputs(next_tokens, current_length, next_indices); + if (engram_state_) { + engram_state_->UpdateInputsOutputs(next_tokens); + } } if (is_prompt_) { @@ -1140,9 +1297,17 @@ DeviceSpan MultiModalPipelineState::Run(int current_length, DeviceSpanper_layer_inputs_->ReuseEmbeddingsBuffer(*decoder_state_->per_layer_inputs_); } embedding_state_->Run(current_length, next_tokens, next_indices); + if (engram_state_) { + if (!chunk_prefill) { + engram_state_->Run(current_length, next_tokens, next_indices); + engram_state_->CopyEmbeddingsTo(*decoder_state_->engram_embeddings_); + } + } auto logits = chunk_prefill - ? decoder_state_->RunPrefillWithChunking(current_length, next_tokens, next_indices, chunk_size_opt.value()) + ? decoder_state_->RunPrefillWithChunking( + current_length, next_tokens, next_indices, chunk_size, + engram_state_.get()) : decoder_state_->Run(current_length, next_tokens, next_indices); is_prompt_ = false; @@ -1163,6 +1328,10 @@ DeviceSpan MultiModalPipelineState::Run(int current_length, DeviceSpanRun(current_length, next_tokens, next_indices); } + if (engram_state_) { + engram_state_->Run(current_length, next_tokens, next_indices); + engram_state_->CopyEmbeddingsTo(*decoder_state_->engram_embeddings_); + } auto logits = decoder_state_->Run(current_length, next_tokens, next_indices); return audio_output_ ? SampleAudioOrText(logits) : logits; } @@ -1181,6 +1350,49 @@ DeviceSpan MultiModalPipelineState::SampleAudioOrText(DeviceSpan l return audio_output_->SampleFrame(*hidden_states); } +void MultiModalPipelineState::ValidateRewindTo(size_t index) const { + decoder_state_->ValidateRewindTo(index); + if (engram_state_) engram_state_->ValidateRewindTo(index); +} + +void MultiModalPipelineState::RewindTo(size_t index) { + ValidateRewindTo(index); + decoder_state_->RewindTo(index); + if (engram_state_) engram_state_->RewindTo(index); + if (index == 0) { + is_prompt_ = true; + num_images_ = 0; + num_image_tokens_ = 0; + num_audio_tokens_ = 0; + } +} + +void MultiModalPipelineState::SnapshotState(size_t position) { + decoder_state_->SnapshotState(position); + if (engram_state_) engram_state_->SnapshotState(position); +} + +bool MultiModalPipelineState::HasCroppableRecurrentState() const { + return decoder_state_->HasCroppableRecurrentState(); +} + +int64_t MultiModalPipelineState::RecurrentStateWindow() const { + return decoder_state_->RecurrentStateWindow(); +} + +void MultiModalPipelineState::CropToAccepted(size_t new_length, size_t recurrent_position) { + ValidateCropToAccepted(new_length, recurrent_position); + decoder_state_->CropToAccepted(new_length, recurrent_position); + if (engram_state_) { + engram_state_->CommitAcceptedPrefix(recurrent_position + 1); + } +} + +void MultiModalPipelineState::ValidateCropToAccepted(size_t new_length, size_t recurrent_position) const { + decoder_state_->ValidateCropToAccepted(new_length, recurrent_position); + if (engram_state_) engram_state_->ValidateAcceptedPrefix(recurrent_position + 1); +} + OrtValue* MultiModalPipelineState::GetInput(const char* name) { if (vision_state_) { // Check if input name is in vision state's inputs diff --git a/src/models/multi_modal.h b/src/models/multi_modal.h index 35d0554b0c..f96b60ad27 100644 --- a/src/models/multi_modal.h +++ b/src/models/multi_modal.h @@ -5,8 +5,10 @@ #include #include +#include #include #include +#include #include #include "model.h" @@ -14,12 +16,16 @@ #include "models/io/multi_modal_features.h" #include "models/io/embeddings.h" #include "models/io/extra_inputs.h" +#include "models/io/hidden_states.h" #include "models/io/logits.h" +#include "models/io/indexer_cache.h" +#include "models/io/ple_state.h" #include "io/kv_cache.h" #include "models/io/position_inputs.h" #include "model_type.h" #include "models/io/recurrent_state.h" #include "models/lfm2_audio_output.h" +#include "engram.h" namespace Generators { @@ -40,11 +46,13 @@ struct MultiModalLanguageModel : Model { std::unique_ptr vision_session_; // pixel_values, [image_attention_mask], image_sizes -> image_features std::unique_ptr speech_session_; // audio_embeds, audio_sizes, audio_projection_mode -> audio_features std::unique_ptr embedding_session_; // input_ids, image_features, audio_features -> inputs_embeds + std::unique_ptr engram_session_; // input_ids, token history -> engram_embeddings std::unique_ptr decoder_session_; // inputs_embeds, attention_mask, kv_cache -> logits std::unique_ptr vision_session_options_; std::unique_ptr speech_session_options_; std::unique_ptr embedding_session_options_; + std::unique_ptr engram_session_options_; // LFM2-Audio speech output, present when model.audio_output names the two graphs. std::unique_ptr depthformer_session_; // hidden_states -> one audio code per run, a frame in num_codebooks runs @@ -234,6 +242,19 @@ struct DecoderState : State { DecoderState& operator=(const DecoderState&) = delete; DeviceSpan Run(int current_length, DeviceSpan& next_tokens, DeviceSpan next_indices) override; + void RewindTo(size_t index) override; + void ValidateRewindTo(size_t index) const override; + void SnapshotState(size_t position) override; + OrtValue* GetRotaryPositionIds() override { return position_inputs_->GetRotaryPositionIds(); } + void SetRotaryPositionIds(OrtValue* positions) override { position_inputs_->SetRotaryPositionIds(positions); } + bool HasCroppableRecurrentState() const override; + bool CropRequiresSnapshot() const override { + return (recurrent_state_ && recurrent_state_->CropRequiresSnapshot()) || + (indexer_cache_ && indexer_cache_->HasStateUpdates()); + } + int64_t RecurrentStateWindow() const override; + void CropToAccepted(size_t new_length, size_t recurrent_position) override; + void ValidateCropToAccepted(size_t new_length, size_t recurrent_position) const override; void UpdateInputsOutputs(DeviceSpan& next_tokens, int current_length, DeviceSpan beam_indices); // Prefill chunking (see search.chunk_size). The embedding model still runs once over the whole @@ -242,7 +263,8 @@ struct DecoderState : State { bool SupportsPrefillChunking(bool has_multimodal_content) const; void PrepareEmbeddingsForPrefill(size_t new_length); DeviceSpan RunPrefillWithChunking(int current_length, DeviceSpan& next_tokens, - DeviceSpan next_indices, size_t chunk_size); + DeviceSpan next_indices, size_t chunk_size, + EngramState* engram_state); private: friend struct MultiModalPipelineState; @@ -253,11 +275,15 @@ struct DecoderState : State { Embeddings inputs_embeds_{*this, Embeddings::Mode::Input, // Model input model_.config_->model.decoder.inputs.embeddings}; std::unique_ptr per_layer_inputs_; // Optional model input (Gemma4: per-layer conditioning) + std::unique_ptr engram_embeddings_; // Optional CPU Engram output staged on CUDA std::unique_ptr decoder_input_ids_; // Optional model input (e.g., Gemma4 decoder needs input_ids) std::unique_ptr position_inputs_; // Model input std::unique_ptr kv_cache_; // Model input std::unique_ptr recurrent_state_; // Model input (for hybrid models) - Logits logits_{*this}; // Model output + std::unique_ptr ple_state_; // Model input (Qwen4-Exp PLE) + std::unique_ptr indexer_cache_; // Model input (Qwen4-Exp QSA) + std::unique_ptr hidden_states_output_; + Logits logits_{*this}; // Model output }; struct MultiModalPipelineState : State { @@ -270,6 +296,18 @@ struct MultiModalPipelineState : State { DeviceSpan Run(int current_length, DeviceSpan& next_tokens, DeviceSpan next_indices) override; + void RewindTo(size_t index) override; + void ValidateRewindTo(size_t index) const override; + void SnapshotState(size_t position) override; + OrtValue* GetRotaryPositionIds() override { return decoder_state_->GetRotaryPositionIds(); } + void SetRotaryPositionIds(OrtValue* positions) override { decoder_state_->SetRotaryPositionIds(positions); } + bool HasCroppableRecurrentState() const override; + bool CropRequiresSnapshot() const override { + return decoder_state_->CropRequiresSnapshot() || engram_state_ != nullptr; + } + int64_t RecurrentStateWindow() const override; + void CropToAccepted(size_t new_length, size_t recurrent_position) override; + void ValidateCropToAccepted(size_t new_length, size_t recurrent_position) const override; OrtValue* GetInput(const char* name) override; @@ -288,6 +326,7 @@ struct MultiModalPipelineState : State { std::unique_ptr vision_state_; std::unique_ptr speech_state_; std::unique_ptr embedding_state_; + std::unique_ptr engram_state_; std::unique_ptr decoder_state_; std::unique_ptr audio_output_; // LFM2-Audio speech output, when the model has it std::shared_ptr adapters_; diff --git a/src/models/onnxruntime_api.h b/src/models/onnxruntime_api.h index 4b0d5cb66f..5260c0c191 100644 --- a/src/models/onnxruntime_api.h +++ b/src/models/onnxruntime_api.h @@ -687,6 +687,8 @@ struct OrtSessionOptions { OrtSessionOptions& SetInterOpNumThreads(int inter_op_num_threads); ///< Wraps OrtApi::SetInterOpNumThreads OrtSessionOptions& SetGraphOptimizationLevel(GraphOptimizationLevel graph_optimization_level); ///< Wraps OrtApi::SetSessionGraphOptimizationLevel + OrtSessionOptions& SetDeterministicCompute(bool value); + OrtSessionOptions& EnableCpuMemArena(); ///< Wraps OrtApi::EnableCpuMemArena OrtSessionOptions& DisableCpuMemArena(); ///< Wraps OrtApi::DisableCpuMemArena diff --git a/src/models/onnxruntime_inline.h b/src/models/onnxruntime_inline.h index 50bd6af4f3..f2df74f61b 100644 --- a/src/models/onnxruntime_inline.h +++ b/src/models/onnxruntime_inline.h @@ -621,6 +621,11 @@ inline OrtSessionOptions& OrtSessionOptions::SetGraphOptimizationLevel(GraphOpti return *this; } +inline OrtSessionOptions& OrtSessionOptions::SetDeterministicCompute(bool value) { + Ort::ThrowOnError(Ort::api->SetDeterministicCompute(this, value)); + return *this; +} + inline OrtSessionOptions& OrtSessionOptions::SetOptimizedModelFilePath(const ORTCHAR_T* optimized_model_filepath) { Ort::ThrowOnError(Ort::api->SetOptimizedModelFilePath(this, optimized_model_filepath)); return *this; diff --git a/src/models/preprocessing/multi_modal_processor.cpp b/src/models/preprocessing/multi_modal_processor.cpp index a0ae8443fd..bb7fa283e9 100644 --- a/src/models/preprocessing/multi_modal_processor.cpp +++ b/src/models/preprocessing/multi_modal_processor.cpp @@ -42,6 +42,7 @@ MultiModalProcessor::MultiModalProcessor(Config& config, const SessionInfo& sess {"qwen3_vl", Processor::Create}, {"qwen3_5", Processor::Create}, {"qwen3_5_moe", Processor::Create}, + {"qwen4_exp", Processor::Create}, {"videochat_flash_qwen", Processor::Create}} { auto processor = processor_factory_.find(config.model.type); if (processor != processor_factory_.end()) { diff --git a/src/models/preprocessing/qwen2_5_vl_image_processor.cpp b/src/models/preprocessing/qwen2_5_vl_image_processor.cpp index 8ec9f1f609..27f0b8d499 100644 --- a/src/models/preprocessing/qwen2_5_vl_image_processor.cpp +++ b/src/models/preprocessing/qwen2_5_vl_image_processor.cpp @@ -5,6 +5,9 @@ #include "models/model.h" #include "models/preprocessing/genai_tokenizer.h" #include "models/preprocessing/qwen2_5_vl_image_processor.h" +#include "models/threadpool.h" +#include +#include #include #include #include diff --git a/src/python/CMakeLists.txt b/src/python/CMakeLists.txt index 967c3bb242..5f58fdf5c0 100644 --- a/src/python/CMakeLists.txt +++ b/src/python/CMakeLists.txt @@ -157,7 +157,8 @@ if(BUILD_WHEEL) configure_file(${CMAKE_CURRENT_SOURCE_DIR}/__init__.py.in ${WHEEL_TARGET_NAME}/__init__.py @ONLY) # Copy over any additional python files - file(COPY "${CMAKE_CURRENT_SOURCE_DIR}/py/" DESTINATION ${WHEEL_TARGET_NAME}/) + file(COPY "${CMAKE_CURRENT_SOURCE_DIR}/py/" DESTINATION ${WHEEL_TARGET_NAME}/ + REGEX "/models/gemma[^/]+/" EXCLUDE) file(COPY "${CMAKE_CURRENT_SOURCE_DIR}/package_description.md" DESTINATION ${WHEEL_FILES_DIR}/) file(COPY "${REPO_ROOT}/ThirdPartyNotices.txt" DESTINATION ${WHEEL_TARGET_NAME}/) file(COPY "${REPO_ROOT}/LICENSE" DESTINATION ${WHEEL_TARGET_NAME}/) diff --git a/src/search.cpp b/src/search.cpp index e49a21e6a7..de5c936374 100644 --- a/src/search.cpp +++ b/src/search.cpp @@ -404,15 +404,16 @@ std::span Search_Cpu::GetScores(int batch_beam_index) { return next_token_scores_.CpuSpan().subspan(static_cast(batch_beam_index) * params_->config.model.vocab_size, params_->config.model.vocab_size); } -void Search_Cpu::ApplyMinLength(int min_length) { - if (sequences_.GetSequenceLength() >= min_length) { +void Search_Cpu::ApplyMinLength(int min_length, std::optional sequence_length) { + const int length = sequence_length.value_or(sequences_.GetSequenceLength()); + if (length >= min_length) { return; } const int batch_beam_size = params_->BatchBeamSize(); for (int i = 0; i < batch_beam_size; i++) { std::span const beam_token_scores = GetScores(i); - ApplyMinLengthToLogits(beam_token_scores, sequences_.GetSequenceLength(), min_length, + ApplyMinLengthToLogits(beam_token_scores, length, min_length, params_->config.model.eos_token_id); } } diff --git a/src/search.h b/src/search.h index 8b6099967a..dd2c81feca 100644 --- a/src/search.h +++ b/src/search.h @@ -1,4 +1,5 @@ #include "sequences.h" +#include #include #include #include "beam_search_scorer.h" @@ -119,7 +120,7 @@ struct Search : LeakChecked { virtual void SampleTopKTopP(int /*k*/, float /*p*/, float /*temperature*/, std::mt19937& /*rng*/) { assert(false); } // Scoring features - virtual void ApplyMinLength(int min_length) = 0; + virtual void ApplyMinLength(int min_length, std::optional sequence_length = std::nullopt) = 0; virtual void ApplyRepetitionPenalty(float penalty) = 0; // This is inline because the standalone CUDA shared library does not compile search.cpp. virtual void ApplyNoRepeatNgram(int ngram_size) { @@ -170,7 +171,7 @@ struct Search_Cpu : Search { DeviceSpan GetLogits() const override; void SetLogits(DeviceSpan logits) override; - void ApplyMinLength(int min_length) override; + void ApplyMinLength(int min_length, std::optional sequence_length = std::nullopt) override; void ApplyRepetitionPenalty(float penalty) override; void ApplyNoRepeatNgram(int ngram_size) override; diff --git a/src/smartptrs.h b/src/smartptrs.h index b99827397c..a95a0d7c17 100644 --- a/src/smartptrs.h +++ b/src/smartptrs.h @@ -118,6 +118,8 @@ struct PositionInputs { virtual void Add() = 0; virtual void Update(DeviceSpan next_tokens, int total_length, int new_length) = 0; virtual void RewindTo(size_t index) = 0; + virtual OrtValue* GetRotaryPositionIds() { return nullptr; } + virtual void SetRotaryPositionIds(OrtValue* positions) { (void)positions; } }; struct BatchedSamplerState { @@ -191,6 +193,8 @@ struct StateSlotDesc { enum class StateUpdateReplayKind : uint32_t { CausalConv = 1, GatedDeltaNet = 2, + Snapshot = 3, + Indexer = 4, }; struct StateUpdateReplayDesc { @@ -217,7 +221,7 @@ static_assert(std::is_trivially_copyable_v); // that boundary (Search, BatchedSampler, BatchedSamplerState, GeneratorParams, or Config). // Dynamically loaded add-ons must report this exact version before the host can safely call through // the C++ interface. -inline constexpr uint32_t kDeviceInterfaceVersion = 6; +inline constexpr uint32_t kDeviceInterfaceVersion = 13; struct DeviceInterface { virtual ~DeviceInterface() {}