diff --git a/.gitignore b/.gitignore index b2a93af05..a8816af8b 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,9 @@ *.dylib *.bin *.cg +*.json +*.cactus +*.weights !cactus-engine/libs/curl/**/*.a !android/mbedtls/**/*.a @@ -29,4 +32,6 @@ python/cactus/bindings/lib/ site/ site_docs/ +TODO.md +explain.md python/cactus/code/ diff --git a/README.md b/README.md index 0b056026c..5ab4db76e 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ A hybrid edge-cloud AI engine for mobile devices & wearables. └─────────────────┘ │ ┌─────────────────┐ -│ Cactus Graph │ ←── Zero-copy computation graph +│ Cactus Graph │ ←── Zero-copy computation graph └─────────────────┘ │ ┌─────────────────┐ @@ -25,12 +25,8 @@ A hybrid edge-cloud AI engine for mobile devices & wearables. └─────────────────┘ │ ┌─────────────────┐ -│ Cactus Quants │ ←── Custom rotation-based quantization technique +│ Cactus Quants │ ←── Custom rotation-based quantization technique └─────────────────┘ - │ -┌─────────────────┐ -│Cactus Transpiler│ ←── Transpiles custom PyTorch model to Cactus. -└─────────────────┘ ``` ## Quick Demo (Mac) @@ -126,7 +122,7 @@ graph.hard_reset(); - VLM: Gemma-4-E2B-CQ4 (256px image encode time / decode) - Transcribe: Parakeet-TDT-0.6B-CQ4 (20s audio end-to-end transcribe time) - 1k-Context RAM: peak MB during the LLM benchmark -- No speculative decode or MTP, pure decode +- No speculative decode or MTP, pure decode Command: `cactus benchmark` [optional `--ios` or `--android`] @@ -141,14 +137,14 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] N/B: With 1k-context prefill and decode for 100 runs on M5 Max - `LFM2.5-VL-1.6B` = 289toks/sec -- `Qwen3-1.7B` = 155toks/sec +- `Qwen3-1.7B` = 155toks/sec - `LFM2.5-VL-450m` = 472toks/sec, image encodes in 43ms - `LFM22.5-VL-230m` = 555toks/sec -## Output Quality +## Output Quality -- Gemma-4-E2B-it accuracy across bit widths, averaged over 3 seeds. -- CQ3.26 and CQ2.54 are mixed-precision, CQ2/CQ3/CQ4 are uniformly quantized. +- Gemma-4-E2B-it accuracy across bit widths, averaged over 3 seeds. +- CQ3.26 and CQ2.54 are mixed-precision, CQ2/CQ3/CQ4 are uniformly quantized. - Full results in [docs/cactus_quants.md](/docs/cactus_quants.md): | Task | F16 (Original) | CQ4 | CQ3.26 | CQ2.54 | CQ2 | @@ -168,10 +164,10 @@ N/B: With 1k-context prefill and decode for 100 runs on M5 Max ## Supported Models -- Any HuggingFace model can be converted using `cactus convert [HF-Name]`, though experimental. -- Liquid, Gemma. whisper. parakeet and Qwen model families are especially tested. +- Any HuggingFace model can be converted using `cactus convert [HF-Name]`, though experimental. +- Liquid, Gemma. whisper. parakeet and Qwen model families are especially tested. - Some models have been pre-uploaded [here](https://huggingface.co/Cactus-Compute), just run `cactus download [HF-Name]`. -- `cactus run [HF-Name]` albeit first downloads or convert the model if not found. +- `cactus run [HF-Name]` albeit first downloads or convert the model if not found. ## Needle @@ -190,7 +186,6 @@ cactus run Cactus-Compute/needle [--tools my_tools.json] # OpenAI function-call | [Cactus Kernels](/docs/cactus_kernels.md) | C++ | ARM NEON SIMD kernels for matmul, attention, convolution, quantization, DSP, image processing | | [Cactus Quants](/docs/cactus_quants.md) | C++ | Rotation-and-codebook quantization from 4-bit to 1-bit for all weight tensors | | [Cactus Hybrid](/docs/cactus_hybrid.md) | C/Python | Route hard queries to the cloud automatically based on local model confidence | -| [Cactus Transpiler](/docs/cactus_transpiler.md) | Python | Convert any PyTorch model to a Cactus runtime graph for on-device inference | | [Python Package](/python/) | Python | Python package and CLI | ## Bindings @@ -244,13 +239,12 @@ cactus run Cactus-Compute/needle [--tools my_tools.json] # OpenAI function-call │ cactus download [model] get a bundle (prebuilt, else build) │ │ --bits 1|2|3|4|2.54|3.26 CQ quantization (default: 4) │ │ --token HuggingFace token (gated models) │ -│ --reconvert force local rebuild from source │ +│ --reconvert refresh cached bundle │ │ │ -│ cactus convert [dir] HuggingFace -> runnable cactus bundle │ -│ (CQ weights + runtime graph) │ +│ cactus convert [dir] HuggingFace -> Cactus CQ weights │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ │ --token HuggingFace token (gated models) │ -│ --reconvert force local rebuild from source │ +│ --reconvert force weight conversion from source │ │ --lora merge a LoRA adapter before converting │ │ --weights-only stop after CQ weights (skip the graph) │ │ --artifact-dir bundle output (default: weights/) │ diff --git a/apple/CMakeLists.txt b/apple/CMakeLists.txt index d2795bdd4..cdd95297d 100644 --- a/apple/CMakeLists.txt +++ b/apple/CMakeLists.txt @@ -27,7 +27,7 @@ if(BUILD_SHARED_LIBS) ) target_link_libraries(cactus_framework PUBLIC - "-framework Accelerate" "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" + "-framework Accelerate" "-framework CoreML" "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" "-framework Security" "-framework SystemConfiguration" "-framework CFNetwork" ) diff --git a/bindings/rust/cactus.rs b/bindings/rust/cactus.rs index 3a956a618..03e65c398 100644 --- a/bindings/rust/cactus.rs +++ b/bindings/rust/cactus.rs @@ -7,6 +7,7 @@ pub type CactusTokenCallback = Option; #[cfg_attr(target_os = "macos", link(name = "Accelerate", kind = "framework"))] +#[cfg_attr(target_os = "macos", link(name = "CoreML", kind = "framework"))] #[cfg_attr(target_os = "macos", link(name = "Metal", kind = "framework"))] #[cfg_attr(target_os = "macos", link(name = "MetalPerformanceShaders", kind = "framework"))] #[cfg_attr(target_os = "macos", link(name = "Foundation", kind = "framework"))] diff --git a/blog/lfm2.5_350m.md b/blog/lfm2.5_350m.md index 3c314fa2f..5b18a735c 100644 --- a/blog/lfm2.5_350m.md +++ b/blog/lfm2.5_350m.md @@ -55,7 +55,7 @@ That's two separate shader codepaths to maintain, with different memory models, And GPU inference holds a persistent GPU context, which competes with the UI rendering pipeline and drains battery even when the inference workload is light. GPUs are also energy-inefficient. On Macs and PCs this barely matters since they're plugged into a power source most of the time, and GPU inference works well there. But mobile devices run on battery, and every watt spent on GPU inference cuts directly into the user's runtime. The GPU path that works fine on a MacBook becomes a battery killer on a phone. -Dedicated AI accelerators are a more interesting target. Apple's Neural Engine and Qualcomm's Hexagon DSP both offer high-throughput inference at good power efficiency. +Dedicated AI accelerators are a more interesting target. Apple's Neural Engine and Qualcomm's Hexagon DSP both offer high-throughput inference at good power efficiency. The problem is access. Apple's Neural Engine requires conversion into a proprietary model format, which imposes its own quantization and graph constraints. You don't control the execution schedule, the memory layout, or the precision semantics. The quantisation techniques available through that toolchain severely lag the state-of-the-art methods used by GGML and Cactus. @@ -201,4 +201,4 @@ from a Vision Pro to a Raspberry Pi 5, with no special casing and no device-spec - [Cactus Engine API Reference](/docs/cactus_engine.md) - Full C API docs for completion, tool calling, and cloud handoff - [LFM2-24B-A2B](/blog/lfm2_24b_a2b.md) - Running LFM2 24B MoE locally on Mac for coding use cases - [Hybrid Transcription](/blog/hybrid_transcription.md) - On-device/cloud hybrid speech transcription with Cactus -- [Ridiculously Fast Transcription](/blog/parakeet.md) - 6M tok/sec decode speed with Parakeet CTC 1.1B \ No newline at end of file +- [Ridiculously Fast Transcription](/blog/parakeet.md) - 6M tok/sec decode speed with Parakeet CTC 1.1B diff --git a/cactus-engine/CMakeLists.txt b/cactus-engine/CMakeLists.txt index ce09c48fe..bba63c225 100644 --- a/cactus-engine/CMakeLists.txt +++ b/cactus-engine/CMakeLists.txt @@ -15,6 +15,7 @@ set(ENGINE_SOURCES src/sp.cpp src/constraints.cpp src/model.cpp + src/model_npu.cpp src/kv_compress.cpp src/engine_image.cpp src/index.cpp @@ -31,6 +32,13 @@ set(ENGINE_SOURCES src/stream.cpp src/tokenizer.cpp ) +if(APPLE) + enable_language(OBJCXX) + list(APPEND ENGINE_SOURCES src/npu_ane.mm) + set_source_files_properties(src/npu_ane.mm PROPERTIES COMPILE_FLAGS "-fobjc-arc") +else() + list(APPEND ENGINE_SOURCES src/npu.cpp) +endif() add_library(cactus_engine STATIC ${ENGINE_SOURCES}) set_target_properties(cactus_engine PROPERTIES POSITION_INDEPENDENT_CODE ON @@ -47,7 +55,7 @@ target_link_libraries(cactus_engine PUBLIC cactus_graph) if(APPLE) target_link_libraries(cactus_engine PUBLIC - "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" + "-framework CoreML" "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" "-framework Security" "-framework SystemConfiguration" "-framework CFNetwork" ) @@ -134,6 +142,7 @@ if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) add_library(cactus_engine_shared SHARED ${CMAKE_BINARY_DIR}/cactus_engine_shared_stub.cpp) set_target_properties(cactus_engine_shared PROPERTIES OUTPUT_NAME cactus_engine) add_dependencies(cactus_engine_shared ${_CACTUS_ENGINE_DEPS}) + set_property(TARGET cactus_engine_shared APPEND PROPERTY LINK_DEPENDS ${_CACTUS_ENGINE_ARCHIVES}) target_link_libraries(cactus_engine_shared PRIVATE cactus_flags) if(APPLE) @@ -147,7 +156,7 @@ if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR) target_link_libraries(cactus_engine_shared PRIVATE "${_curl}") endif() target_link_libraries(cactus_engine_shared PRIVATE - "-framework Accelerate" "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" + "-framework Accelerate" "-framework CoreML" "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" "-framework Security" "-framework SystemConfiguration" "-framework CFNetwork" ) else() diff --git a/cactus-engine/src/complete.cpp b/cactus-engine/src/complete.cpp index e8fe2e33d..d91b70272 100644 --- a/cactus-engine/src/complete.cpp +++ b/cactus-engine/src/complete.cpp @@ -372,6 +372,13 @@ std::vector> build_stop_sequences( stop_token_sequences.push_back(tokenizer->encode(stop_seq)); } + if (!tokenizer->has_chat_template()) { + stop_token_sequences.push_back(tokenizer->encode("\nUser:")); + stop_token_sequences.push_back(tokenizer->encode("\nAssistant:")); + stop_token_sequences.push_back(tokenizer->encode("\nUser")); + stop_token_sequences.push_back(tokenizer->encode("\nAssistant")); + } + if (model_type == Config::ModelType::GEMMA4) { stop_token_sequences.push_back(tokenizer->encode("")); if (has_tools) { @@ -435,6 +442,7 @@ struct PreparedPrompt { std::vector tokens; size_t context_token_count = 0; std::vector> images; + bool force_session_reset = false; std::vector> audio_features; size_t audio_num_frames = 0; @@ -575,7 +583,7 @@ PreparedPrompt prepare_prompt( } } - if (prompt.model_type == Config::ModelType::GEMMA4) { + if (handle->model->audio_preprocess_strategy() == "gemma4") { std::vector user_indices; for (size_t i = 0; i < prompt.messages.size(); i++) { if (prompt.messages[i].role == "user") user_indices.push_back(i); @@ -629,6 +637,23 @@ PreparedPrompt prepare_prompt( setup_tool_constraints(handle, prompt.tools, prompt.options.force_tools, prompt.options.temperature); } + if (tokenizer->uses_default_chat_style(prompt.messages, formatted_tools, prompt.options.enable_thinking_if_supported)) { + auto last_user = std::find_if(prompt.messages.rbegin(), prompt.messages.rend(), [](const ChatMessage& msg) { + return msg.role == "user"; + }); + if (last_user != prompt.messages.rend()) { + std::vector latest_turn; + for (const auto& msg : prompt.messages) { + if (msg.role == "system" || msg.role == "developer") { + latest_turn.push_back(msg); + } + } + latest_turn.push_back(*last_user); + prompt.messages = std::move(latest_turn); + prompt.force_session_reset = true; + } + } + prompt.rendered = tokenizer->format_chat_prompt( prompt.messages, add_generation_prompt, @@ -827,6 +852,15 @@ int cactus_complete( handle_error_response(err, response_buffer, buffer_size); return -1; } + const auto runtime_model_type = handle->model->get_config().model_type; + if (runtime_model_type == Config::ModelType::WHISPER || + runtime_model_type == Config::ModelType::PARAKEET || + runtime_model_type == Config::ModelType::PARAKEET_TDT) { + const std::string err = "Model is speech-to-text; use `cactus transcribe` instead of `cactus run`"; + CACTUS_LOG_ERROR("complete", err); + handle_error_response(err, response_buffer, buffer_size); + return -1; + } auto* tokenizer = handle->model->get_tokenizer(); auto prompt = prepare_prompt(handle, messages_json, options_json, tools_json, true, true, pcm_buffer, pcm_buffer_size); if (prompt.options.sample_seed != 0) { @@ -956,6 +990,10 @@ int cactus_complete( uint32_t next_token; size_t prompt_tokens; + if (prompt.force_session_reset && !handle->processed_tokens.empty()) { + reset_cache(handle); + } + if (has_audio && !handle->processed_tokens.empty()) { auto& cache = handle->processed_tokens; size_t common = 0; @@ -979,6 +1017,9 @@ int cactus_complete( if (!first_token_from_prefill) { auto prefill_result = do_prefill(handle, prompt, prompt.tokens); prompt_tokens = prefill_result.prefilled_count + prefill_result.remaining_tokens.size(); + if (handle->model->repetition_penalty_uses_context()) { + handle->model->seed_token_history_from_context(prompt.tokens); + } next_token = generate_first_token(handle, prefill_result, prompt, &first_token_entropy); } diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 73fc6d47a..0ee887c3d 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -17,6 +17,43 @@ class CactusGraph; namespace cactus { +namespace npu { + +struct NPUNamedInput { + std::string name; + enum class DataType { FP16, INT32 }; + const void* data; + DataType data_type = DataType::FP16; + std::vector shape; +}; + +class NPUEncoder { +public: + virtual ~NPUEncoder() = default; + virtual bool load(const std::string& model_path, const std::string& compute_units = "") = 0; + virtual bool preallocate(const std::vector& input_shape, + const std::string& input_name = "x", + const std::string& output_name = "") = 0; + virtual size_t encode(const __fp16* input, __fp16* output, + const std::vector& shape, + const std::string& input_name = "x", + const std::string& output_name = "") = 0; + virtual bool is_available() const = 0; + virtual std::vector get_input_shape() const = 0; + virtual std::vector get_output_shape() const = 0; + virtual bool has_input(const std::string&) const { return false; } + virtual std::vector get_input_shape_for(const std::string&) const { return {}; } + virtual __fp16* get_output_buffer() = 0; + virtual size_t get_output_buffer_size() const = 0; + virtual size_t encode_multimodal_input(const std::vector& inputs, + __fp16* output, + const std::string& output_name = "") = 0; +}; + +std::unique_ptr create_encoder(); +bool is_npu_available(); + +} // namespace npu namespace engine { struct Config { @@ -76,6 +113,10 @@ struct Config { uint32_t tile_size = 512; float max_pixels_tolerance = 2.0f; bool do_image_splitting = true; + bool do_resize = true; + bool do_rescale = true; + bool do_normalize = true; + bool do_convert_rgb = true; bool encoder_act_gelu = false; bool decoder_act_gelu = false; uint32_t num_encoder_layers = 0; @@ -108,7 +149,7 @@ struct Config { float rope_scaling_factor = 1.0f; float rope_mscale_all_dim = 0.0f; - enum class ModelType {QWEN = 0, GEMMA = 1, NOMIC = 3, LFM2 = 5, SIGLIP2 = 6, WHISPER = 7, MOONSHINE = 8, PARAKEET = 10, QWEN3P5 = 11, PARAKEET_TDT = 12, GEMMA3N = 13, YOUTU = 14, GEMMA4 = 15, NEEDLE = 18}; + enum class ModelType {QWEN = 0, GEMMA = 1, NOMIC = 3, LFM2 = 5, SIGLIP2 = 6, WHISPER = 7, MOONSHINE = 8, PARAKEET = 10, QWEN3P5 = 11, PARAKEET_TDT = 12, GEMMA3N = 13, YOUTU = 14, GEMMA4 = 15, NEEDLE = 18, GENERIC = 19}; uint32_t predictor_hidden_dim = 0; uint32_t predictor_num_layers = 0; uint32_t tdt_joint_dim = 0; @@ -290,6 +331,7 @@ class Tokenizer { virtual std::vector apply_chat_template(const std::vector& messages, bool add_generation_prompt = true) const; virtual std::string format_chat_prompt(const std::vector& messages, bool add_generation_prompt = true, const std::string& tools_json = "", bool enable_thinking_if_supported = false) const; + virtual bool uses_default_chat_style(const std::vector& messages, const std::string& tools_json = "", bool enable_thinking_if_supported = false) const; virtual uint32_t get_vocab_size() const = 0; virtual uint32_t get_unk_token() const = 0; @@ -312,9 +354,45 @@ class Tokenizer { size_t get_image_soft_token_count() const { return image_soft_token_count_; } void set_lfm2_vision_config(const Config& cfg) { lfm2_vision_config_ = cfg; has_lfm2_vision_config_ = true; } + void set_gemma4_raw_media_prompt(bool enabled) { gemma4_raw_media_prompt_ = enabled; } + void set_text_prompt_style(const std::string& style) { text_prompt_style_ = style; } + void set_media_prompt_style(const std::string& style) { media_prompt_style_ = style; } + void set_gemma4_turn_tokens(const std::string& turn_start_token, + const std::string& turn_end_token) { + if (!turn_start_token.empty()) gemma4_turn_start_token_ = turn_start_token; + if (!turn_end_token.empty()) gemma4_turn_end_token_ = turn_end_token; + } + void set_media_prompt_positions(const std::string& image_position, + const std::string& audio_position) { + if (!image_position.empty()) media_image_prompt_position_ = image_position; + if (!audio_position.empty()) media_audio_prompt_position_ = audio_position; + } + void set_media_prompt_order(const std::vector& order) { + if (!order.empty()) media_prompt_order_ = order; + } + void set_media_focus_policy(const std::string& policy, + const std::vector& image_keywords, + const std::vector& audio_keywords) { + media_focus_policy_ = policy; + if (!image_keywords.empty()) media_image_focus_keywords_ = image_keywords; + if (!audio_keywords.empty()) media_audio_focus_keywords_ = audio_keywords; + } + void set_media_tokens(const std::string& image_token, + const std::string& image_begin_token, + const std::string& image_end_token, + const std::string& audio_token, + const std::string& audio_begin_token, + const std::string& audio_end_token) { + if (!image_token.empty()) media_image_token_ = image_token; + if (!image_begin_token.empty()) media_image_begin_token_ = image_begin_token; + if (!image_end_token.empty()) media_image_end_token_ = image_end_token; + if (!audio_token.empty()) media_audio_token_ = audio_token; + if (!audio_begin_token.empty()) media_audio_begin_token_ = audio_begin_token; + if (!audio_end_token.empty()) media_audio_end_token_ = audio_end_token; + } protected: - enum class ModelType { UNKNOWN, GEMMA4, GEMMA, QWEN, LFM2, NEEDLE }; + enum class ModelType { UNKNOWN, GEMMA4, GEMMA, QWEN, LFM2, NEEDLE, WHISPER, PARAKEET_TDT }; ModelType model_type_ = ModelType::UNKNOWN; enum class ModelVariant { DEFAULT, VLM, EXTRACT, RAG}; ModelVariant model_variant_ = ModelVariant::DEFAULT; @@ -331,12 +409,30 @@ class Tokenizer { uint32_t vision_default_output_length_ = 280; uint32_t vision_image_size_ = 768; size_t image_soft_token_count_ = 0; + bool gemma4_raw_media_prompt_ = false; + std::string gemma4_turn_start_token_ = "<|turn>"; + std::string gemma4_turn_end_token_ = ""; + std::string text_prompt_style_; + std::string media_prompt_style_; + std::string media_image_token_ = "<|image|>"; + std::string media_image_begin_token_ = "<|image>"; + std::string media_image_end_token_ = ""; + std::string media_audio_token_ = "<|audio|>"; + std::string media_audio_begin_token_ = "<|audio>"; + std::string media_audio_end_token_ = ""; + std::string media_image_prompt_position_ = "before_text"; + std::string media_audio_prompt_position_ = "before_text"; + std::vector media_prompt_order_ = {"image", "audio"}; + std::string media_focus_policy_; + std::vector media_image_focus_keywords_; + std::vector media_audio_focus_keywords_; Config lfm2_vision_config_{}; bool has_lfm2_vision_config_ = false; TokenizerRuntimeConfig runtime_config_; void detect_model_type(const std::string& config_path); void load_chat_template(const std::string& template_file); + std::string format_default_chat_style(const std::vector& messages, bool add_generation_prompt) const; std::string format_gemma4_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported = false) const; std::string format_qwen_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported = false) const; std::string format_lfm2_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported = false) const; @@ -575,19 +671,14 @@ class ToolCallConstrainer { class Model { public: - struct DebugNode { - uint32_t layer_idx; - std::string name; - size_t node_id; - }; - Model(); explicit Model(const Config& config); ~Model(); const Config& get_config() const { return config_; } + size_t audio_soft_token_count_for_frames(size_t frame_count); + std::string audio_preprocess_strategy() const; Tokenizer* get_tokenizer() const { return tokenizer_.get(); } - const std::vector& get_debug_nodes() const; bool init(const std::string& bundle_dir, size_t context_size, const std::string& system_prompt = "", bool do_warmup = true); @@ -681,7 +772,10 @@ class Model { token_history_.erase(token_history_.begin(), token_history_.begin() + (MAX_TOKEN_HISTORY / 2)); } token_history_.push_back(token); + token_history_set_.insert(token); } + void seed_token_history_from_context(const std::vector& tokens); + bool repetition_penalty_uses_context() const { return repetition_penalty_scope_ == "context"; } double score_tokens_window_logprob(const std::vector& tokens, size_t start, size_t end, size_t context, size_t* tokens_scored); @@ -697,6 +791,12 @@ class Model { size_t last_prefill_tail_chunk_tokens() const { return last_prefill_tail_chunk_tokens_; } size_t last_prefill_tail_padding_tokens() const { return last_prefill_tail_padding_tokens_; } + bool load_npu_audio_encoder(const std::string& model_path, const std::string& compute_units = ""); + bool has_npu_audio_encoder() const { return npu_audio_encoder_ != nullptr; } + bool load_npu_vision_encoder(const std::string& model_path); + bool has_npu_vision_encoder() const { return npu_vision_encoder_ != nullptr; } + bool load_npu_source_encoder(const std::string& model_path); + void remove_thinking_tokens(const std::vector>& ranges); void compact_kv_cache() {} @@ -731,6 +831,81 @@ class Model { int value_node_id = -1; }; + struct RuntimeStateSpec { + std::string name; + std::string kind; + std::string producer; + std::vector consumers; + std::vector release_after_consumers; + std::string lifetime; + std::string transfer; + bool persist_after_component_unload = true; + bool required = true; + std::map metadata; + }; + + struct RuntimeAliasSpec { + std::string source_component; + std::string source_output; + std::string target_component; + std::string target_input; + std::string policy; + std::string lifetime; + std::string fallback; + bool required = false; + int source_node_id = -1; + int target_node_id = -1; + bool storage_stable = false; + std::map metadata; + }; + + struct RuntimeStoredTensor { + std::vector bytes; + std::shared_ptr storage; + std::vector shape; + Precision precision = Precision::FP16; + std::string producer_component; + std::string producer_output; + std::string kind; + std::string lifetime; + }; + + struct StateArena { + std::map> tensors; + + void publish(const std::string& id, std::shared_ptr storage) { + if (!id.empty() && storage && storage->valid()) tensors[id] = std::move(storage); + } + std::shared_ptr find(const std::string& id) const { + auto it = tensors.find(id); + return it == tensors.end() ? nullptr : it->second; + } + size_t byte_size() const { + size_t total = 0; + std::unordered_set seen; + for (const auto& [_, storage] : tensors) { + if (storage && seen.insert(storage.get()).second) total += storage->byte_size; + } + return total; + } + void erase(const std::string& id) { tensors.erase(id); } + void clear() { tensors.clear(); } + }; + + struct RuntimeProfileCounters { + uint64_t copied_bytes = 0; + uint64_t aliased_bytes = 0; + uint64_t media_copied_bytes = 0; + uint64_t media_aliased_bytes = 0; + uint64_t component_loads = 0; + uint64_t component_unloads = 0; + uint64_t resident_weight_bytes = 0; + uint64_t peak_resident_weight_bytes = 0; + uint64_t peak_state_bytes = 0; + uint64_t peak_runtime_store_bytes = 0; + uint64_t released_runtime_store_bytes = 0; + }; + struct Component { std::string name; std::string graph_path; @@ -743,11 +918,12 @@ class Model { std::map metadata; std::unique_ptr graph; std::vector> input_buffers; + size_t resident_weight_bytes = 0; }; - void copy_cache_state(const Component& src, Component& dst); - struct ChunkedPrefillResult { + Component* component = nullptr; + Component* encoder_component = nullptr; size_t logical_tokens = 0; size_t executed_tokens = 0; size_t padding_tokens = 0; @@ -760,6 +936,10 @@ class Model { bool load_components(const std::unordered_set& required_components); bool load_component_graph(Component& comp); void unload_component_graph(Component& comp); + size_t component_weight_bytes(const Component& comp) const; + void observe_state_arena_bytes(); + size_t runtime_state_store_byte_size() const; + void write_runtime_profile() const; bool bind_runtime_buffers(Component& comp); void run_step(uint32_t token_id, size_t position, bool read_logits, bool use_fused = true); void run_step_batch(const std::vector& token_ids, const std::vector& positions); @@ -771,7 +951,9 @@ class Model { void write_media_embeds_row(Component& comp, int embeds_idx, const uint8_t* feature_row, size_t feature_row_bytes, Precision feature_precision); void reset_encoder_cross_kv_route_state(); + void release_encoder_cross_kv_prepare_graphs(); bool finish_encoder_cross_kv_prepare(); + bool finish_encoder_cross_kv_prepare_after_source(); bool prepare_encoder_cross_kv_from_text(const std::vector& tokens); bool prepare_encoder_cross_kv_from_audio(const std::vector& audio_features); bool run_encoder_cross_kv_decoder_step(uint32_t token_id, size_t position); @@ -780,11 +962,22 @@ class Model { size_t max_tokens, const std::vector>& stop_token_sequences, const std::atomic* should_stop); - void copy_component_outputs_to_inputs(const Component& source, Component& target); - bool copy_cross_kv_outputs_to_decoder_cache_inputs(const Component& source, Component& target, size_t source_len); - void copy_encoder_outputs_to_decoder(const Component& enc); - void copy_component_outputs_to_chunk_inputs(const Component& source, Component& target, size_t token_index); - void copy_component_outputs_to_chunk_inputs_range(const Component& source, Component& target, size_t token_offset); + void copy_component_outputs_to_inputs(Component& source, Component& target); + bool copy_cross_kv_outputs_to_decoder_cache_inputs(Component& source, Component& target, size_t source_len); + void copy_component_outputs_to_chunk_inputs(Component& source, Component& target, size_t token_index); + void copy_component_outputs_to_chunk_inputs_range(Component& source, Component& target, + size_t token_offset, size_t source_tokens = 0); + void persist_component_outputs(Component& source); + void bind_runtime_state_inputs(Component& target); + void bind_component_cache_states(Component& target); + void publish_component_cache_states(Component& source); + static std::string cache_state_id(const CacheStateBinding& binding, bool key); + bool runtime_state_producer_matches(const RuntimeStateSpec& state, const Component& source) const; + bool runtime_state_consumer_matches(const RuntimeStateSpec& state, const Component& target) const; + bool output_should_persist_to_runtime_state(const Component& source, const std::string& output_name) const; + const RuntimeStoredTensor* find_runtime_state_for_input(const Component& target, const std::string& input_name) const; + void clear_runtime_states_with_lifetime(const std::string& lifetime); + void release_runtime_states_after_consumer(const Component& consumer); bool cache_states_compatible(const Component& source, const Component& target) const; void move_cache_states(Component& source, Component& target, size_t logical_current = std::numeric_limits::max()); void set_cache_current_len(Component& comp, size_t len); @@ -802,27 +995,36 @@ class Model { uint32_t sample_component_logits(Component& comp, float temperature, float top_p, size_t top_k, float min_p, bool greedy, float* out_uncertainty); uint32_t argmax_component_logits(Component& comp, size_t logit_row = std::numeric_limits::max(), - float* out_uncertainty = nullptr); - uint32_t argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_off, float* out_uncertainty); + float* out_uncertainty = nullptr, float repetition_penalty = 1.0f); + uint32_t argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_off, float* out_uncertainty, + float repetition_penalty = 1.0f); + bool is_generation_suppressed_token(uint32_t token_id) const; std::vector argmax_component_logits_batch(Component& comp, size_t batch); void write_int_input(Component& comp, const std::string& name, int64_t value); void write_int_input_at(Component& comp, const std::string& name, size_t index, int64_t value); + void fill_int_input(Component& comp, const std::string& name, int64_t value, size_t count); void write_bytes_input(Component& comp, const std::string& name, const void* data, size_t byte_size); int input_index(const Component& comp, const std::string& name) const; int output_index(const Component& comp, const std::string& name) const; - uint32_t argmax_last_logits(float* out_uncertainty = nullptr); + uint32_t argmax_last_logits(float* out_uncertainty = nullptr, float repetition_penalty = 1.0f); bool load_handoff_probe(); void maybe_capture_handoff_probe_hidden(const Component& comp, const std::string& output_name = "probe_hidden"); void run_vision_encoder(const std::string& image_path); void run_vision_encoder_lfm2_vl(const std::string& image_path); void encode_lfm2_vl_image_into_features(const std::string& image_path); bool load_lfm2_vl_position_grid(); + bool lfm2_vl_use_npu_vision() const; + bool lfm2_vl_encode_tile_npu(const float* pixel_values, const int64_t* mask, + const float* pos_embeds, size_t max_patches, int dim, + size_t patch_dim, std::vector& enc_out); void run_audio_encoder(const std::vector& audio_features); void run_audio_encoder_messages(const std::vector>& audio_features_per_message); bool run_chunk_prefill_path(const std::vector& tokens, const std::vector& image_paths, const std::vector>& audio_features_per_message); - bool build_lm_encoder_outputs_dynamic_gemma4( + size_t media_rows_to_copy(size_t output_rows, size_t frames_in_chunk) const; + std::string image_preprocess_strategy() const; + bool build_lm_encoder_outputs_dynamic_media( const std::vector& tokens, std::map>& store_bytes, std::map& store_prec, @@ -853,10 +1055,76 @@ class Model { bool prefill_tail_pad_disabled_ = false; std::string family_; + std::string runtime_plan_name_; + std::string runtime_execution_strategy_; + std::string runtime_state_owner_; + std::string runtime_cache_persistence_; + std::string runtime_output_alias_policy_; + std::string runtime_cache_transfer_policy_; + std::vector runtime_states_; + std::vector runtime_aliases_; + std::map runtime_state_store_; + StateArena state_arena_; + std::string image_preprocess_strategy_; + std::string audio_preprocess_strategy_; + std::string media_injection_strategy_; + std::string media_prefill_fallback_; + size_t media_min_new_tokens_ = 0; + size_t media_min_new_tokens_remaining_ = 0; + std::unordered_set media_min_stop_token_ids_; + std::string prompt_text_style_; + std::string prompt_media_style_; + std::string media_image_token_; + std::string media_audio_token_; + std::string media_image_begin_token_; + std::string media_image_end_token_; + std::string media_audio_begin_token_; + std::string media_audio_end_token_; + std::string media_image_prompt_position_; + std::string media_audio_prompt_position_; + std::string prompt_turn_start_token_; + std::string prompt_turn_end_token_; + std::vector media_prompt_order_; + std::string media_focus_policy_; + std::vector media_image_focus_keywords_; + std::vector media_audio_focus_keywords_; + uint32_t media_placeholder_token_id_ = 0; + size_t audio_rows_frame_divisor_ = 0; + std::unordered_set generation_suppressed_token_ids_; + std::string repetition_penalty_scope_; + std::vector media_chunk_prefill_modalities_; + std::map media_chunk_output_sources_; + std::vector image_feature_candidates_; + std::vector audio_feature_candidates_; + std::string npu_audio_encoder_mlpackage_; + std::string npu_audio_compute_units_; + std::string npu_vision_encoder_mlpackage_; + std::string npu_source_encoder_mlpackage_; + + std::unique_ptr npu_audio_encoder_; + std::unique_ptr npu_vision_encoder_; + std::unique_ptr npu_source_encoder_; + + bool audio_encode_via_npu(const std::vector& audio_features); + bool vision_encode_via_npu(const std::vector& pixel_values, + const std::vector* pixel_position_ids = nullptr); + bool source_encode_via_npu(const std::vector& tokens); std::map> media_features_; + std::map> media_feature_storages_; std::map> media_feature_shapes_; std::map media_feature_precisions_; + const uint8_t* media_feature_data(const std::string& name) const; + size_t media_feature_byte_size(const std::string& name) const; + void clear_media_feature(const std::string& name); + void publish_media_feature_storage(const std::string& name, + std::shared_ptr storage, + const std::vector& shape, + Precision precision); + void append_media_feature(const std::string& name, const void* data, size_t byte_size, + const std::vector& shape, Precision precision, + size_t rows_to_append = 0); + RuntimeProfileCounters runtime_profile_; std::vector lfm2_pos_grid_; int lfm2_pos_grid_h_ = 0; @@ -882,6 +1150,7 @@ class Model { static constexpr size_t MAX_TOKEN_HISTORY = 128; std::vector token_history_; + std::unordered_set token_history_set_; std::vector samp_recent_; std::vector samp_bias_dense_; bool samp_has_bias_ = false; @@ -913,7 +1182,6 @@ class Model { std::vector handoff_probe_head4_bias_; std::vector handoff_probe_hidden_; - mutable std::vector debug_nodes_; }; class ConvCache { diff --git a/cactus-engine/src/engine_image.cpp b/cactus-engine/src/engine_image.cpp index 3fdb92e66..b24c1ab0a 100644 --- a/cactus-engine/src/engine_image.cpp +++ b/cactus-engine/src/engine_image.cpp @@ -912,11 +912,11 @@ static Siglip2Preprocessor::Config make_lfm2_vl_siglip_config(const Config& conf sp_cfg.rescale_factor = config.rescale_factor > 0.0f ? config.rescale_factor : (1.0f / 255.0f); sp_cfg.image_mean[0] = sp_cfg.image_mean[1] = sp_cfg.image_mean[2] = config.image_mean; sp_cfg.image_std[0] = sp_cfg.image_std[1] = sp_cfg.image_std[2] = config.image_std; - sp_cfg.do_image_splitting = true; - sp_cfg.do_resize = true; - sp_cfg.do_rescale = true; - sp_cfg.do_normalize = true; - sp_cfg.do_convert_rgb = true; + sp_cfg.do_image_splitting = config.do_image_splitting; + sp_cfg.do_resize = config.do_resize; + sp_cfg.do_rescale = config.do_rescale; + sp_cfg.do_normalize = config.do_normalize; + sp_cfg.do_convert_rgb = config.do_convert_rgb; return sp_cfg; } diff --git a/cactus-engine/src/init.cpp b/cactus-engine/src/init.cpp index ca11df7fa..127e55fa9 100644 --- a/cactus-engine/src/init.cpp +++ b/cactus-engine/src/init.cpp @@ -1,6 +1,7 @@ #include "../cactus_engine.h" #include "utils.h" #include "telemetry.h" +#include "../../cactus-kernels/src/threading.h" #include #include #include @@ -19,7 +20,8 @@ static constexpr size_t RAG_MIN_CHUNK_TOKENS = 24; static constexpr size_t RAG_CHUNK_OVERLAP = 32; static void apply_no_cloud_telemetry_env() { - if (cactus::ffi::env_flag_enabled("CACTUS_NO_CLOUD_TELE")) { + if (cactus::ffi::env_flag_enabled("CACTUS_NO_CLOUD_TELE") || + cactus::ffi::env_flag_enabled("CACTUS_DISABLE_CLOUD_HANDOFF")) { cactus::telemetry::setCloudDisabled(true); } } @@ -352,10 +354,9 @@ extern "C" { int cactus_set_backend(const char* backend) { return cactus_backend_select(backend); } cactus_model_t cactus_init(const char* model_path, const char* corpus_dir, bool cache_index) { - constexpr size_t DEFAULT_CONTEXT_SIZE = 512; // matches default sliding window size + constexpr size_t DEFAULT_CONTEXT_SIZE = 512; std::string model_path_str = model_path ? std::string(model_path) : "unknown"; - std::string model_name = model_path_str; size_t last_slash = model_path_str.find_last_of("/\\"); if (last_slash != std::string::npos) { @@ -439,7 +440,10 @@ cactus_model_t cactus_init(const char* model_path, const char* corpus_dir, bool } void cactus_destroy(cactus_model_t model) { - if (model) delete static_cast(model); + if (!model) return; + CactusThreading::get_thread_pool().wait_all(); + delete static_cast(model); + CactusThreading::get_thread_pool().wait_all(); } void cactus_reset(cactus_model_t model) { diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 15788bb43..d7f1213c5 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include namespace cactus { @@ -44,6 +45,167 @@ std::vector parse_config_uint_list(const std::string& value) { return numbers; } +std::vector parse_csv_strings(const std::string& value) { + std::vector items; + std::stringstream ss(value); + std::string item; + while (std::getline(ss, item, ',')) { + size_t begin = 0; + while (begin < item.size() && std::isspace(static_cast(item[begin]))) ++begin; + size_t end = item.size(); + while (end > begin && std::isspace(static_cast(item[end - 1]))) --end; + if (end > begin) items.push_back(item.substr(begin, end - begin)); + } + return items; +} + +bool debug_media_features_enabled() { + return std::getenv("CACTUS_DEBUG_MEDIA_FEATURES") != nullptr; +} + +uint64_t byte_checksum(const std::vector& bytes) { + uint64_t sum = 0; + for (uint8_t value : bytes) sum += value; + return sum; +} + +void debug_media_feature(const std::string& stage, const std::string& name, + const std::vector& bytes, + const std::vector& shape, + Precision precision) { + if (!debug_media_features_enabled()) return; + std::cerr << "[media-debug] " << stage << " " << name + << " bytes=" << bytes.size() + << " checksum=" << byte_checksum(bytes) + << " precision=" << static_cast(precision) + << " shape="; + for (size_t i = 0; i < shape.size(); ++i) { + if (i) std::cerr << "x"; + std::cerr << shape[i]; + } + std::cerr << std::endl; + + const char* dump_dir = std::getenv("CACTUS_DEBUG_MEDIA_FEATURES_DIR"); + if (!dump_dir || bytes.empty()) return; + std::filesystem::path path = std::filesystem::path(dump_dir) / (stage + "_" + name + ".bin"); + std::ofstream out(path, std::ios::binary); + if (out.is_open()) { + out.write(reinterpret_cast(bytes.data()), static_cast(bytes.size())); + } + std::filesystem::path meta_path = std::filesystem::path(dump_dir) / (stage + "_" + name + ".txt"); + std::ofstream meta(meta_path); + if (meta.is_open()) { + meta << "precision=" << static_cast(precision) << "\nshape="; + for (size_t i = 0; i < shape.size(); ++i) { + if (i) meta << ","; + meta << shape[i]; + } + meta << "\nbytes=" << bytes.size() << "\n"; + } +} + +template +void debug_media_vector(const std::string& stage, const std::string& name, + const std::vector& values, + const std::vector& shape, + Precision precision) { + if (!debug_media_features_enabled()) return; + std::vector bytes(values.size() * sizeof(T)); + if (!values.empty()) { + std::memcpy(bytes.data(), values.data(), bytes.size()); + } + debug_media_feature(stage, name, bytes, shape, precision); +} + +std::unordered_set parse_csv_uint_set(const std::string& value) { + std::unordered_set ids; + for (const auto& item : parse_csv_strings(value)) { + try { + ids.insert(static_cast(std::stoul(item))); + } catch (...) { + } + } + return ids; +} + +std::map parse_csv_string_map(const std::string& value) { + std::map items; + for (const auto& item : parse_csv_strings(value)) { + size_t separator = item.find(':'); + if (separator == std::string::npos || separator == 0 || separator + 1 >= item.size()) { + continue; + } + + items[item.substr(0, separator)] = item.substr(separator + 1); + } + return items; +} + +std::string json_string_or_empty(const picojson::object& obj, const std::string& key) { + auto it = obj.find(key); + if (it == obj.end() || !it->second.is()) return ""; + return it->second.get(); +} + +bool json_bool_or(const picojson::object& obj, const std::string& key, bool fallback) { + auto it = obj.find(key); + if (it == obj.end() || !it->second.is()) return fallback; + return it->second.get(); +} + +std::vector json_string_array(const picojson::object& obj, const std::string& key) { + std::vector out; + auto it = obj.find(key); + if (it == obj.end() || !it->second.is()) return out; + for (const auto& value : it->second.get()) { + if (value.is()) out.push_back(value.get()); + } + return out; +} + +std::map json_string_map(const picojson::object& obj, const std::string& key) { + std::map out; + auto it = obj.find(key); + if (it == obj.end() || !it->second.is()) return out; + for (const auto& kv : it->second.get()) { + if (kv.second.is()) out[kv.first] = kv.second.get(); + } + return out; +} + +namespace { +void write_typed_buffer(std::vector& buf, Precision dst_prec, + const void* src_data, size_t src_bytes, Precision src_prec); +} // namespace + +size_t parse_ceil_divisor(const std::string& value) { + const std::string prefix = "ceil_div:"; + if (value.rfind(prefix, 0) != 0) return 0; + try { + return static_cast(std::stoul(value.substr(prefix.size()))); + } catch (...) { + return 0; + } +} + +uint32_t parse_uint32_or_zero(const std::string& value) { + if (value.empty()) return 0; + try { + return static_cast(std::stoul(value)); + } catch (...) { + return 0; + } +} + +size_t parse_size_or_zero(const std::string& value) { + if (value.empty()) return 0; + try { + return static_cast(std::stoull(value)); + } catch (...) { + return 0; + } +} + float read_scalar_value(Precision precision, const uint8_t* data, size_t index) { const uint8_t* ptr = data + PrecisionTraits::byte_offset_of(precision, index); switch (precision) { @@ -103,6 +265,77 @@ bool copy_component_tensor(CactusGraph& source_graph, return true; } +struct MediaFeatureInfo { + std::string name; + const uint8_t* data = nullptr; + size_t byte_size = 0; + size_t rows = 0; + size_t row_bytes = 0; + Precision precision = Precision::FP16; + + bool present() const { + return data != nullptr && byte_size > 0 && row_bytes > 0; + } + + const uint8_t* row(size_t index) const { + return data + index * row_bytes; + } +}; + +MediaFeatureInfo find_media_feature( + const std::map>& features, + const std::map>& storages, + const std::map>& shapes, + const std::map& precisions, + const std::vector& preferred_names, + const std::vector& fallback_names, + const std::string& excluded_name = "") { + auto load = [&](const std::string& name) -> MediaFeatureInfo { + if (name.empty() || name == excluded_name) return {}; + auto bytes_it = features.find(name); + auto storage_it = storages.find(name); + const bool have_bytes = bytes_it != features.end() && !bytes_it->second.empty(); + const bool have_storage = storage_it != storages.end() && storage_it->second && storage_it->second->valid(); + if (!have_bytes && !have_storage) return {}; + + MediaFeatureInfo info; + info.name = name; + info.data = have_storage + ? static_cast(storage_it->second->get_data()) + : bytes_it->second.data(); + info.byte_size = have_storage ? storage_it->second->byte_size : bytes_it->second.size(); + auto prec_it = precisions.find(name); + if (prec_it != precisions.end()) info.precision = prec_it->second; + + auto shape_it = shapes.find(name); + if (shape_it != shapes.end()) { + const auto& shape = shape_it->second; + if (shape.size() >= 2) { + info.rows = shape.size() >= 3 ? shape[shape.size() - 2] : shape[0]; + } else if (!shape.empty()) { + info.rows = shape[0]; + } + } + info.row_bytes = info.rows > 0 ? info.byte_size / info.rows : info.byte_size; + return info; + }; + + const auto& candidates = preferred_names.empty() ? fallback_names : preferred_names; + for (const auto& name : candidates) { + auto info = load(name); + if (info.present()) return info; + } + + if (preferred_names.empty() && fallback_names.empty()) { + for (const auto& kv : features) { + auto info = load(kv.first); + if (info.present()) return info; + } + } + + return {}; +} + struct CrossKVCacheMetadata { uint64_t current_seq_len; uint64_t max_seq_len; @@ -263,7 +496,9 @@ Model::Model() : config_() {} Model::Model(const Config& config) : config_(config) {} -Model::~Model() = default; +Model::~Model() { + write_runtime_profile(); +} namespace { @@ -451,6 +686,39 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, CACTUS_LOG_ERROR("model", "Tokenizer init failed for bundle: " << bundle_dir); return false; } + if (tokenizer_) { + tokenizer_->set_media_tokens( + media_image_token_, + media_image_begin_token_, + media_image_end_token_, + media_audio_token_, + media_audio_begin_token_, + media_audio_end_token_); + } + if (tokenizer_ && prompt_media_style_ == "raw_processor") { + tokenizer_->set_gemma4_raw_media_prompt(true); + } + if (tokenizer_) { + tokenizer_->set_text_prompt_style(prompt_text_style_); + tokenizer_->set_media_prompt_style(prompt_media_style_); + tokenizer_->set_gemma4_turn_tokens(prompt_turn_start_token_, prompt_turn_end_token_); + tokenizer_->set_media_prompt_positions(media_image_prompt_position_, media_audio_prompt_position_); + tokenizer_->set_media_prompt_order(media_prompt_order_); + tokenizer_->set_media_focus_policy( + media_focus_policy_, + media_image_focus_keywords_, + media_audio_focus_keywords_); + + media_min_stop_token_ids_.clear(); + media_min_stop_token_ids_.insert(tokenizer_->get_eos_token()); + const std::string media_stop = prompt_turn_end_token_.empty() + ? tokenizer_->get_default_stop_sequence() + : prompt_turn_end_token_; + const std::vector media_stop_tokens = tokenizer_->encode(media_stop); + if (media_stop_tokens.size() == 1) { + media_min_stop_token_ids_.insert(media_stop_tokens.front()); + } + } const bool is_text_embedding = components_.count("text_embedding") && !components_.count("decoder") @@ -497,7 +765,7 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, && components_.count("decoder_prefill_chunk"); if (has_metadata_encoder_cross_kv_route) { decode_route_ = DecodeRoute::ENCODER_CROSS_KV_STEP; - required_components = {source_encoder_name, decoder_cross_kv_name, decoder_name}; + required_components = {decoder_name}; } else if (has_chunked_prefill) { encoder_name = "lm_encoder_step"; decoder_name = "decoder_media_step"; @@ -505,8 +773,6 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, required_components = { encoder_name, decoder_name, - "lm_encoder_text_chunk", - "decoder_prefill_chunk", }; } else if (components_.count("decoder_step") && input_index(components_.at("decoder_step"), "input_ids") >= 0 @@ -519,12 +785,10 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, decoder_name = "decoder_step"; decode_route_ = DecodeRoute::CACHED_STEP; required_components = {encoder_name, decoder_name}; - if (components_.count("decoder_prefill_chunk")) { - required_components.insert("decoder_prefill_chunk"); - } - if (components_.count("lm_encoder_text_chunk")) { - required_components.insert("lm_encoder_text_chunk"); - } + } else if (components_.count("decoder_full_context")) { + decoder_name = "decoder_full_context"; + decode_route_ = DecodeRoute::FULL_CONTEXT_TEXT; + required_components = {decoder_name}; } else if (components_.count("text_lm_encoder") && components_.count("decoder")) { encoder_name = "text_lm_encoder"; decoder_name = "decoder"; @@ -534,41 +798,35 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, (components_.count("decoder") || components_.count("decoder_joint"))) { decoder_name = components_.count("decoder") ? "decoder" : "decoder_joint"; decode_route_ = DecodeRoute::DIRECT_DECODER_STEP; - required_components = {"audio_encoder", decoder_name}; + required_components = {decoder_name}; } else { - CACTUS_LOG_ERROR("model", "Bundle missing required components: need lm_encoder_step+decoder_step (LM), text_lm_encoder+decoder, audio_encoder+decoder (transcription), or source/audio_encoder+decoder_cross_kv+decoder_step"); + CACTUS_LOG_ERROR("model", "Bundle missing required components: need lm_encoder_step+decoder_step (LM), decoder_full_context, text_lm_encoder+decoder, audio_encoder+decoder (transcription), or source/audio_encoder+decoder_cross_kv+decoder_step"); return false; } - for (const auto& optional : { - "vision_encoder", - "audio_encoder", - "lm_encoder_media_step", - "decoder_prefill_chunk", - "decoder_embed_chunk", - "lm_encoder", - }) { - if (components_.count(optional)) { - required_components.insert(optional); - } - } if (!load_components(required_components)) return false; if (!encoder_name.empty()) encoder_ = &components_.at(encoder_name); if (!decoder_name.empty()) decoder_ = &components_.at(decoder_name); if (!source_encoder_name.empty()) source_encoder_ = &components_.at(source_encoder_name); if (!decoder_cross_kv_name.empty()) decoder_cross_kv_ = &components_.at(decoder_cross_kv_name); - if (components_.count("decoder_prefill_chunk") && components_.at("decoder_prefill_chunk").graph) { - decoder_prefill_ = &components_.at("decoder_prefill_chunk"); - decoder_prefill_chunk_ = decoder_prefill_; - } else if (decode_route_ != DecodeRoute::ENCODER_CROSS_KV_STEP && !components_.count("audio_encoder")) { + if (components_.count("decoder_prefill_chunk")) { + decoder_prefill_chunk_ = &components_.at("decoder_prefill_chunk"); + decoder_prefill_ = components_.count("decoder_prefill_text_chunk") + ? &components_.at("decoder_prefill_text_chunk") + : decoder_prefill_chunk_; + } else if (decode_route_ != DecodeRoute::ENCODER_CROSS_KV_STEP + && decode_route_ != DecodeRoute::FULL_CONTEXT_TEXT + && !components_.count("audio_encoder")) { CACTUS_LOG_WARN("model", "Bundle has no decoder_prefill_chunk component; prompts will prefill token-by-token (prefill speed ~= decode speed)."); } if (decoder_ && decoder_->graph) decoder_->graph->prewarm_metal_quant_weights(); if (decoder_prefill_ && decoder_prefill_->graph) decoder_prefill_->graph->prewarm_metal_quant_weights(); - if (components_.count("decoder_embed_chunk") && components_.at("decoder_embed_chunk").graph) { + if (components_.count("decoder_embed_chunk")) { decoder_embed_ = &components_.at("decoder_embed_chunk"); } - if (components_.count("lm_encoder_text_chunk") && components_.at("lm_encoder_text_chunk").graph) { - prefill_encoder_ = &components_.at("lm_encoder_text_chunk"); + if (components_.count("lm_encoder_text_chunk") || components_.count("lm_encoder_text_prefill_chunk")) { + prefill_encoder_ = components_.count("lm_encoder_text_prefill_chunk") + ? &components_.at("lm_encoder_text_prefill_chunk") + : &components_.at("lm_encoder_text_chunk"); } if (const char* env = std::getenv("CACTUS_DISABLE_PREFILL_TAIL_PAD")) { prefill_tail_pad_disabled_ = std::atoi(env) != 0; @@ -587,6 +845,7 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, decoder_, decoder_cross_kv_, decoder_prefill_, + decoder_prefill_chunk_, vision_encoder_, vision_projector_, audio_encoder_, @@ -617,16 +876,18 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, } } - if (vision_encoder_ && tokenizer_ && !vision_encoder_->output_node_ids.empty()) { + if (vision_encoder_ && tokenizer_ && !vision_encoder_->output_node_ids.empty() + && load_component_graph(*vision_encoder_)) { size_t out_node = static_cast(vision_encoder_->output_node_ids[0]); const auto& desc = vision_encoder_->graph->get_output_buffer(out_node); size_t n = 0; if (desc.shape.size() >= 3) n = desc.shape[desc.shape.size() - 2]; else if (desc.shape.size() >= 2) n = desc.shape[0]; if (n > 0) tokenizer_->set_image_soft_token_count(n); + unload_component_graph(*vision_encoder_); } - if (family_ == "lfm2_vl" && tokenizer_) { + if (image_preprocess_strategy() == "lfm2_vl" && tokenizer_) { tokenizer_->set_lfm2_vision_config(config_); } @@ -660,6 +921,117 @@ bool Model::load_manifest() { if (obj.count("family") && obj.at("family").is()) { family_ = obj.at("family").get(); } + auto read_string_key = [&](const char* key) -> std::string { + if (obj.count(key) && obj.at(key).is()) { + return obj.at(key).get(); + } + return ""; + }; + image_preprocess_strategy_ = read_string_key("media_image_preprocess_strategy"); + audio_preprocess_strategy_ = read_string_key("media_audio_preprocess_strategy"); + media_injection_strategy_ = read_string_key("media_injection_strategy"); + media_prefill_fallback_ = read_string_key("media_prefill_fallback"); + media_min_new_tokens_ = parse_size_or_zero(read_string_key("media_min_new_tokens")); + prompt_text_style_ = read_string_key("prompt_text_style"); + prompt_media_style_ = read_string_key("prompt_media_style"); + media_image_token_ = read_string_key("media_image_token"); + media_audio_token_ = read_string_key("media_audio_token"); + media_image_begin_token_ = read_string_key("media_image_begin_token"); + media_image_end_token_ = read_string_key("media_image_end_token"); + media_audio_begin_token_ = read_string_key("media_audio_begin_token"); + media_audio_end_token_ = read_string_key("media_audio_end_token"); + media_image_prompt_position_ = read_string_key("media_image_prompt_position"); + media_audio_prompt_position_ = read_string_key("media_audio_prompt_position"); + prompt_turn_start_token_ = read_string_key("prompt_turn_start_token"); + prompt_turn_end_token_ = read_string_key("prompt_turn_end_token"); + media_prompt_order_ = parse_csv_strings(read_string_key("media_order")); + media_focus_policy_ = read_string_key("media_focus_policy"); + media_image_focus_keywords_ = parse_csv_strings(read_string_key("media_image_focus_keywords")); + media_audio_focus_keywords_ = parse_csv_strings(read_string_key("media_audio_focus_keywords")); + runtime_plan_name_ = read_string_key("runtime_plan_name"); + runtime_execution_strategy_ = read_string_key("runtime_execution_strategy"); + runtime_state_owner_ = read_string_key("runtime_state_owner"); + runtime_cache_persistence_ = read_string_key("runtime_cache_persistence"); + runtime_output_alias_policy_ = read_string_key("runtime_output_alias_policy"); + runtime_cache_transfer_policy_ = read_string_key("runtime_cache_transfer_policy"); + + runtime_states_.clear(); + if (obj.count("states") && obj.at("states").is()) { + for (const auto& value : obj.at("states").get()) { + if (!value.is()) continue; + const auto& state_obj = value.get(); + RuntimeStateSpec state; + state.name = json_string_or_empty(state_obj, "name"); + state.kind = json_string_or_empty(state_obj, "kind"); + state.producer = json_string_or_empty(state_obj, "producer"); + state.consumers = json_string_array(state_obj, "consumers"); + state.release_after_consumers = json_string_array(state_obj, "release_after_consumers"); + state.lifetime = json_string_or_empty(state_obj, "lifetime"); + state.transfer = json_string_or_empty(state_obj, "transfer"); + state.persist_after_component_unload = json_bool_or(state_obj, "persist_after_component_unload", true); + state.required = json_bool_or(state_obj, "required", true); + state.metadata = json_string_map(state_obj, "metadata"); + if (!state.name.empty()) runtime_states_.push_back(std::move(state)); + } + } + + runtime_aliases_.clear(); + if (obj.count("aliases") && obj.at("aliases").is()) { + for (const auto& value : obj.at("aliases").get()) { + if (!value.is()) continue; + const auto& alias_obj = value.get(); + RuntimeAliasSpec alias; + alias.source_component = json_string_or_empty(alias_obj, "source_component"); + alias.source_output = json_string_or_empty(alias_obj, "source_output"); + alias.target_component = json_string_or_empty(alias_obj, "target_component"); + alias.target_input = json_string_or_empty(alias_obj, "target_input"); + alias.policy = json_string_or_empty(alias_obj, "policy"); + alias.lifetime = json_string_or_empty(alias_obj, "lifetime"); + alias.fallback = json_string_or_empty(alias_obj, "fallback"); + alias.required = json_bool_or(alias_obj, "required", false); + if (alias_obj.count("source_node_id") && alias_obj.at("source_node_id").is()) + alias.source_node_id = static_cast(alias_obj.at("source_node_id").get()); + if (alias_obj.count("target_node_id") && alias_obj.at("target_node_id").is()) + alias.target_node_id = static_cast(alias_obj.at("target_node_id").get()); + alias.storage_stable = json_bool_or(alias_obj, "storage_stable", false); + alias.metadata = json_string_map(alias_obj, "metadata"); + if (!alias.source_component.empty() && !alias.source_output.empty() + && !alias.target_component.empty() && !alias.target_input.empty()) { + runtime_aliases_.push_back(std::move(alias)); + } + } + } + + std::string suppress_ids = read_string_key("suppress_generation_token_ids"); + generation_suppressed_token_ids_ = parse_csv_uint_set(suppress_ids); + repetition_penalty_scope_ = read_string_key("repetition_penalty_scope"); + + media_placeholder_token_id_ = parse_uint32_or_zero(read_string_key("media_placeholder_token_id")); + + uint32_t image_token = parse_uint32_or_zero(read_string_key("media_image_token_id")); + if (image_token != 0) config_.image_token_id = image_token; + + uint32_t audio_token = parse_uint32_or_zero(read_string_key("media_audio_token_id")); + if (audio_token != 0) config_.audio_token_id = audio_token; + + audio_rows_frame_divisor_ = parse_ceil_divisor(read_string_key("media_audio_rows_per_frames")); + media_chunk_prefill_modalities_ = parse_csv_strings(read_string_key("media_chunk_prefill_modalities")); + media_chunk_output_sources_ = parse_csv_string_map(read_string_key("media_chunk_output_sources")); + image_feature_candidates_ = parse_csv_strings(read_string_key("media_image_feature_names")); + audio_feature_candidates_ = parse_csv_strings(read_string_key("media_audio_feature_names")); + + if (obj.count("npu_audio_encoder") && obj.at("npu_audio_encoder").is()) { + npu_audio_encoder_mlpackage_ = obj.at("npu_audio_encoder").get(); + } + if (obj.count("npu_audio_compute_units") && obj.at("npu_audio_compute_units").is()) { + npu_audio_compute_units_ = obj.at("npu_audio_compute_units").get(); + } + if (obj.count("npu_vision_encoder") && obj.at("npu_vision_encoder").is()) { + npu_vision_encoder_mlpackage_ = obj.at("npu_vision_encoder").get(); + } + if (obj.count("npu_source_encoder") && obj.at("npu_source_encoder").is()) { + npu_source_encoder_mlpackage_ = obj.at("npu_source_encoder").get(); + } if (!obj.count("components")) return false; for (const auto& cv : obj.at("components").get()) { const auto& c = cv.get(); @@ -741,11 +1113,16 @@ bool Model::load_components(const std::unordered_set& required_comp } bool Model::load_component_graph(Component& comp) { - if (comp.graph) return true; + if (comp.graph) { + bind_component_cache_states(comp); + bind_runtime_state_inputs(comp); + return true; + } if (comp.graph_path.empty()) return true; fs::path full = fs::path(bundle_dir_) / comp.graph_path; try { comp.graph = std::make_unique(CactusGraph::load(full.string())); + comp.graph->set_profile_label(comp.name); comp.graph->retain_outputs(comp.output_node_ids); } catch (const std::exception& e) { CACTUS_LOG_ERROR("model", "load " << comp.graph_path << ": " << e.what()); @@ -758,26 +1135,176 @@ bool Model::load_component_graph(Component& comp) { if (weight_path.is_absolute()) { fs::path local = fs::path(bundle_dir_) / weight_path.filename(); if (fs::exists(local)) weight_path = local; - } else { + } else weight_path = fs::path(bundle_dir_) / weight_path; - } + comp.graph->bind_mmap_weights(static_cast(b.node_id), weight_path.string()); } catch (const std::exception& e) { CACTUS_LOG_ERROR("model", "bind " << b.path << ": " << e.what()); return false; } } - return bind_runtime_buffers(comp); + if (!bind_runtime_buffers(comp)) return false; + bind_component_cache_states(comp); + bind_runtime_state_inputs(comp); + comp.resident_weight_bytes = component_weight_bytes(comp); + runtime_profile_.resident_weight_bytes += comp.resident_weight_bytes; + runtime_profile_.peak_resident_weight_bytes = std::max(runtime_profile_.peak_resident_weight_bytes, runtime_profile_.resident_weight_bytes); + ++runtime_profile_.component_loads; + return true; } void Model::unload_component_graph(Component& comp) { if (cactus_default_backend() == ComputeBackend::METAL) return; + const bool was_loaded = static_cast(comp.graph); if (comp.graph) { + persist_component_outputs(comp); comp.graph->release_runtime_buffers(); - comp.graph->release_all_weight_pages(); } comp.input_buffers.clear(); comp.graph.reset(); + release_runtime_states_after_consumer(comp); + runtime_profile_.resident_weight_bytes = comp.resident_weight_bytes > runtime_profile_.resident_weight_bytes + ? 0 : runtime_profile_.resident_weight_bytes - comp.resident_weight_bytes; + comp.resident_weight_bytes = 0; + if (was_loaded) ++runtime_profile_.component_unloads; +} + +size_t Model::component_weight_bytes(const Component& comp) const { + size_t total = 0; + std::unordered_set seen; + for (const auto& binding : comp.bindings) { + fs::path path(binding.path); + if (!path.is_absolute()) path = fs::path(bundle_dir_) / path; + std::error_code ec; + const std::string key = path.lexically_normal().string(); + if (seen.insert(key).second && fs::is_regular_file(path, ec)) { + total += static_cast(fs::file_size(path, ec)); + } + } + return total; +} + +void Model::observe_state_arena_bytes() { + runtime_profile_.peak_state_bytes = std::max( + runtime_profile_.peak_state_bytes, state_arena_.byte_size()); + runtime_profile_.peak_runtime_store_bytes = std::max( + runtime_profile_.peak_runtime_store_bytes, runtime_state_store_byte_size()); +} + +size_t Model::runtime_state_store_byte_size() const { + size_t total = 0; + std::unordered_set seen; + for (const auto& [_, tensor] : runtime_state_store_) { + if (tensor.storage) { + if (seen.insert(tensor.storage.get()).second) total += tensor.storage->byte_size; + } else { + total += tensor.bytes.size(); + } + } + return total; +} + +void Model::write_runtime_profile() const { + const char* destination = std::getenv("CACTUS_RUNTIME_PROFILE"); + if (!destination || !*destination) return; + std::ostringstream line; + line << "{\"copied_bytes\":" << runtime_profile_.copied_bytes + << ",\"aliased_bytes\":" << runtime_profile_.aliased_bytes + << ",\"media_copied_bytes\":" << runtime_profile_.media_copied_bytes + << ",\"media_aliased_bytes\":" << runtime_profile_.media_aliased_bytes + << ",\"component_loads\":" << runtime_profile_.component_loads + << ",\"component_unloads\":" << runtime_profile_.component_unloads + << ",\"resident_weight_bytes\":" << runtime_profile_.resident_weight_bytes + << ",\"peak_resident_weight_bytes\":" << runtime_profile_.peak_resident_weight_bytes + << ",\"state_bytes\":" << state_arena_.byte_size() + << ",\"peak_state_bytes\":" << runtime_profile_.peak_state_bytes + << ",\"runtime_store_bytes\":" << runtime_state_store_byte_size() + << ",\"peak_runtime_store_bytes\":" << runtime_profile_.peak_runtime_store_bytes + << ",\"released_runtime_store_bytes\":" << runtime_profile_.released_runtime_store_bytes << "}"; + if (std::string(destination) == "stdout" || std::string(destination) == "-") { + std::cout << "CACTUS_RUNTIME_PROFILE=" << line.str() << std::endl; + return; + } + std::ofstream out(destination, std::ios::app); + if (out.is_open()) out << line.str() << "\n"; +} + +const uint8_t* Model::media_feature_data(const std::string& name) const { + auto storage = media_feature_storages_.find(name); + if (storage != media_feature_storages_.end() && storage->second && storage->second->valid()) { + return static_cast(storage->second->get_data()); + } + auto bytes = media_features_.find(name); + return bytes == media_features_.end() || bytes->second.empty() ? nullptr : bytes->second.data(); +} + +size_t Model::media_feature_byte_size(const std::string& name) const { + auto storage = media_feature_storages_.find(name); + if (storage != media_feature_storages_.end() && storage->second && storage->second->valid()) { + return storage->second->byte_size; + } + auto bytes = media_features_.find(name); + return bytes == media_features_.end() ? 0 : bytes->second.size(); +} + +void Model::clear_media_feature(const std::string& name) { + media_features_.erase(name); + media_feature_storages_.erase(name); + media_feature_shapes_.erase(name); + media_feature_precisions_.erase(name); + runtime_state_store_.erase(name); + state_arena_.erase("media:" + name); +} + +void Model::publish_media_feature_storage(const std::string& name, + std::shared_ptr storage, + const std::vector& shape, + Precision precision) { + if (!storage || !storage->valid()) return; + clear_media_feature(name); + media_features_[name] = {}; + media_feature_storages_[name] = storage; + media_feature_shapes_[name] = shape; + media_feature_precisions_[name] = precision; + state_arena_.publish("media:" + name, std::move(storage)); + runtime_profile_.media_aliased_bytes += media_feature_byte_size(name); + runtime_profile_.aliased_bytes += media_feature_byte_size(name); + observe_state_arena_bytes(); +} + +void Model::append_media_feature(const std::string& name, const void* data, size_t byte_size, + const std::vector& shape, Precision precision, + size_t rows_to_append) { + auto& bytes = media_features_[name]; + auto storage = media_feature_storages_.find(name); + if (storage != media_feature_storages_.end() && storage->second && storage->second->valid()) { + const auto* previous = static_cast(storage->second->get_data()); + bytes.assign(previous, previous + storage->second->byte_size); + runtime_profile_.media_copied_bytes += storage->second->byte_size; + runtime_profile_.copied_bytes += storage->second->byte_size; + media_feature_storages_.erase(storage); + state_arena_.erase("media:" + name); + } + const size_t previous_bytes = bytes.size(); + bytes.resize(previous_bytes + byte_size); + if (byte_size > 0) std::memcpy(bytes.data() + previous_bytes, data, byte_size); + runtime_profile_.media_copied_bytes += byte_size; + runtime_profile_.copied_bytes += byte_size; + auto shape_it = media_feature_shapes_.find(name); + if (shape_it == media_feature_shapes_.end() || shape_it->second.empty()) { + media_feature_shapes_[name] = shape; + if (rows_to_append > 0 && !media_feature_shapes_[name].empty()) { + auto& stored_shape = media_feature_shapes_[name]; + if (stored_shape.size() >= 2) stored_shape[stored_shape.size() - 2] = rows_to_append; + else stored_shape[0] = rows_to_append; + } + } else if (rows_to_append > 0) { + auto& stored_shape = shape_it->second; + if (stored_shape.size() >= 2) stored_shape[stored_shape.size() - 2] += rows_to_append; + else if (!stored_shape.empty()) stored_shape[0] += rows_to_append; + } + media_feature_precisions_[name] = precision; } bool Model::bind_runtime_buffers(Component& comp) { @@ -827,6 +1354,26 @@ void Model::write_int_input_at(Component& comp, const std::string& name, size_t } } +void Model::fill_int_input(Component& comp, const std::string& name, int64_t value, size_t count) { + int idx = input_index(comp, name); + if (idx < 0) return; + + size_t node_id = static_cast(comp.runtime_input_node_ids[idx]); + const auto& desc = comp.graph->get_output_buffer(node_id); + auto& buf = comp.input_buffers[idx]; + size_t n = std::min(count, desc.total_size); + + for (size_t i = 0; i < n; ++i) { + size_t offset = PrecisionTraits::byte_offset_of(desc.precision, i); + write_scalar_value(desc.precision, buf.data() + offset, 0, static_cast(value)); + } + + size_t elem_size = PrecisionTraits::size_of(desc.precision); + if (elem_size > 0 && n < desc.total_size) { + std::memset(buf.data() + n * elem_size, 0, (desc.total_size - n) * elem_size); + } +} + void Model::write_bytes_input(Component& comp, const std::string& name, const void* data, size_t byte_size) { int idx = input_index(comp, name); if (idx < 0) return; @@ -845,24 +1392,11 @@ int Model::output_index(const Component& comp, const std::string& name) const { return -1; } -void Model::copy_encoder_outputs_to_decoder(const Component& enc) { - for (size_t i = 0; i < enc.output_node_ids.size() && i < enc.logical_outputs.size(); ++i) { - const std::string& out_name = enc.logical_outputs[i]; - int dst_idx = input_index(*decoder_, out_name); - if (dst_idx < 0) continue; - size_t src_node = static_cast(enc.output_node_ids[i]); - const auto& src_desc = enc.graph->get_output_buffer(src_node); - void* src_ptr = enc.graph->get_output(src_node); - auto& dst_buf = decoder_->input_buffers[dst_idx]; - size_t to_copy = std::min(src_desc.byte_size, dst_buf.size()); - std::memcpy(dst_buf.data(), src_ptr, to_copy); - } -} - void Model::run_step(uint32_t token_id, size_t position, bool /*read_logits*/, bool use_fused) { if (decode_route_ == DecodeRoute::DIRECT_DECODER_STEP) { write_int_input(*decoder_, "input_ids", static_cast(token_id)); write_int_input(*decoder_, "position_ids", static_cast(position)); + fill_int_input(*decoder_, "attention_mask", 1, position + 1); if (!use_fused || cactus_default_backend() != ComputeBackend::METAL) cactus_graph_set_prefill_consistent(true); decoder_->graph->execute(); cactus_graph_set_prefill_consistent(false); @@ -887,6 +1421,7 @@ void Model::run_step(uint32_t token_id, size_t position, bool /*read_logits*/, b } run_encoder_step(token_id, position); copy_component_outputs_to_inputs(*encoder_, *decoder_); + fill_int_input(*decoder_, "attention_mask", 1, position + 1); decoder_->graph->execute(); cactus_graph_set_prefill_consistent(false); maybe_capture_handoff_probe_hidden(*decoder_); @@ -895,6 +1430,7 @@ void Model::run_step(uint32_t token_id, size_t position, bool /*read_logits*/, b void Model::run_encoder_step(uint32_t token_id, size_t position) { write_int_input(*encoder_, "input_ids", static_cast(token_id)); write_int_input(*encoder_, "position_ids", static_cast(position)); + fill_int_input(*encoder_, "attention_mask", 1, position + 1); encoder_->graph->execute(); } @@ -937,9 +1473,11 @@ void Model::run_step_batch(const std::vector& token_ids, const std::ve for (size_t b = 0; b < token_ids.size(); ++b) { write_int_input_at(*encoder_, "input_ids", b, static_cast(token_ids[b])); write_int_input_at(*encoder_, "position_ids", b, static_cast(positions[b])); + write_int_input_at(*encoder_, "attention_mask", b, 1); } encoder_->graph->execute(); copy_component_outputs_to_inputs(*encoder_, *decoder_); + fill_int_input(*decoder_, "attention_mask", 1, positions.empty() ? 1 : (*std::max_element(positions.begin(), positions.end()) + 1)); decoder_->graph->execute(); maybe_capture_handoff_probe_hidden(*decoder_); } @@ -1078,16 +1616,246 @@ void Model::set_decode_slots(size_t num_slots) { body \ } -void Model::copy_component_outputs_to_inputs(const Component& source, Component& target) { +bool Model::runtime_state_producer_matches(const RuntimeStateSpec& state, const Component& source) const { + if (state.producer.empty()) return false; + for (const auto& producer : parse_csv_strings(state.producer)) { + if (producer == source.name) return true; + } + return false; +} + +bool Model::runtime_state_consumer_matches(const RuntimeStateSpec& state, const Component& target) const { + if (state.consumers.empty()) return false; + return std::find(state.consumers.begin(), state.consumers.end(), target.name) != state.consumers.end(); +} + +bool Model::output_should_persist_to_runtime_state(const Component& source, const std::string& output_name) const { + for (const auto& alias : runtime_aliases_) { + if (alias.source_component == source.name && alias.source_output == output_name) { + return true; + } + } + for (const auto& state : runtime_states_) { + if (!state.persist_after_component_unload || !runtime_state_producer_matches(state, source)) continue; + if (state.name == output_name) return true; + auto outputs = state.metadata.find("outputs"); + if (outputs != state.metadata.end()) { + const auto names = parse_csv_strings(outputs->second); + if (std::find(names.begin(), names.end(), output_name) != names.end()) return true; + continue; + } + // Older manifests described activation state only by producer. Keep + // their broad behavior, while new profiles can avoid retaining every + // producer output by supplying metadata.outputs. + if (state.kind == "activation") return true; + } + return false; +} + +void Model::persist_component_outputs(Component& source) { + if (!source.graph) return; + for (size_t i = 0; i < source.output_node_ids.size() && i < source.logical_outputs.size(); ++i) { + const std::string& output_name = source.logical_outputs[i]; + if (!output_should_persist_to_runtime_state(source, output_name)) continue; + const size_t node_id = static_cast(source.output_node_ids[i]); + const auto& desc = source.graph->get_output_buffer(node_id); + const void* ptr = desc.get_data(); + if (!ptr || desc.byte_size == 0) continue; + + RuntimeStoredTensor tensor; + tensor.storage = source.graph->export_tensor_storage(node_id); + if (!tensor.storage) { + tensor.bytes.assign(static_cast(ptr), static_cast(ptr) + desc.byte_size); + } + tensor.shape = desc.shape; + tensor.precision = desc.precision; + tensor.producer_component = source.name; + tensor.producer_output = output_name; + tensor.kind = "activation"; + tensor.lifetime = "request"; + for (const auto& state : runtime_states_) { + if (!runtime_state_producer_matches(state, source)) continue; + bool state_output_matches = state.name == output_name; + auto outputs = state.metadata.find("outputs"); + if (!state_output_matches && outputs != state.metadata.end()) { + const auto names = parse_csv_strings(outputs->second); + state_output_matches = std::find(names.begin(), names.end(), output_name) != names.end(); + } else if (!state_output_matches && outputs == state.metadata.end()) { + state_output_matches = state.kind == "activation"; + } + if (state_output_matches) { + tensor.kind = state.kind; + tensor.lifetime = state.lifetime.empty() ? tensor.lifetime : state.lifetime; + break; + } + } + runtime_state_store_[output_name] = std::move(tensor); + observe_state_arena_bytes(); + } +} + +const Model::RuntimeStoredTensor* Model::find_runtime_state_for_input( + const Component& target, + const std::string& input_name) const { + for (const auto& alias : runtime_aliases_) { + if (alias.target_component != target.name || alias.target_input != input_name) continue; + auto it = runtime_state_store_.find(alias.source_output); + if (it != runtime_state_store_.end()) return &it->second; + } + + for (const auto& state : runtime_states_) { + if (state.name != input_name || !runtime_state_consumer_matches(state, target)) continue; + auto it = runtime_state_store_.find(state.name); + if (it != runtime_state_store_.end()) return &it->second; + } + return nullptr; +} + +void Model::bind_runtime_state_inputs(Component& target) { + if (!target.graph) return; + for (size_t i = 0; i < target.runtime_input_node_ids.size() && i < target.logical_inputs.size(); ++i) { + const RuntimeStoredTensor* tensor = find_runtime_state_for_input(target, target.logical_inputs[i]); + if (!tensor || (!tensor->storage && tensor->bytes.empty())) continue; + size_t node_id = static_cast(target.runtime_input_node_ids[i]); + if (tensor->storage && target.graph->bind_tensor_storage(node_id, tensor->storage)) { + runtime_profile_.aliased_bytes += tensor->storage->byte_size; + continue; + } + const auto& desc = target.graph->get_output_buffer(node_id); + const void* source = tensor->storage ? tensor->storage->get_data() : tensor->bytes.data(); + const size_t source_bytes = tensor->storage ? tensor->storage->byte_size : tensor->bytes.size(); + write_typed_buffer(target.input_buffers[i], desc.precision, + source, source_bytes, tensor->precision); + runtime_profile_.copied_bytes += std::min(desc.byte_size, source_bytes); + target.graph->set_external_input(node_id, target.input_buffers[i].data(), desc.precision); + } +} + +std::string Model::cache_state_id(const CacheStateBinding& binding, bool key) { + return binding.layer_key + (key ? ":key" : ":value"); +} + +void Model::bind_component_cache_states(Component& target) { + if (std::getenv("CACTUS_DISABLE_SHARED_STATE") != nullptr) return; + if (!target.graph) return; + for (const auto& binding : target.cache_states) { + int previous = -1; + for (auto [node_id, key] : {std::pair{binding.key_node_id, true}, + std::pair{binding.value_node_id, false}}) { + if (node_id < 0 || node_id == previous) continue; + previous = node_id; + auto storage = state_arena_.find(cache_state_id(binding, key)); + if (storage) target.graph->bind_tensor_storage(static_cast(node_id), storage); + } + } +} + +void Model::publish_component_cache_states(Component& source) { + if (!source.graph || std::getenv("CACTUS_DISABLE_SHARED_STATE") != nullptr) return; + for (const auto& binding : source.cache_states) { + int previous = -1; + for (auto [node_id, key] : {std::pair{binding.key_node_id, true}, + std::pair{binding.value_node_id, false}}) { + if (node_id < 0 || node_id == previous) continue; + previous = node_id; + auto storage = source.graph->export_tensor_storage(static_cast(node_id)); + if (storage) { + state_arena_.publish(cache_state_id(binding, key), std::move(storage)); + observe_state_arena_bytes(); + } + } + } +} + +void Model::clear_runtime_states_with_lifetime(const std::string& lifetime) { + for (auto it = runtime_state_store_.begin(); it != runtime_state_store_.end();) { + if (it->second.lifetime == lifetime) it = runtime_state_store_.erase(it); + else ++it; + } + observe_state_arena_bytes(); +} + +void Model::release_runtime_states_after_consumer(const Component& consumer) { + const size_t store_bytes_before = runtime_state_store_byte_size(); + for (const auto& state : runtime_states_) { + if (std::find(state.release_after_consumers.begin(), state.release_after_consumers.end(), consumer.name) + == state.release_after_consumers.end()) { + continue; + } + + std::vector outputs; + auto output_metadata = state.metadata.find("outputs"); + if (output_metadata != state.metadata.end()) { + outputs = parse_csv_strings(output_metadata->second); + } + if (outputs.empty() && !state.name.empty()) outputs.push_back(state.name); + + for (const auto& output : outputs) { + runtime_state_store_.erase(output); + if (media_features_.count(output) || media_feature_storages_.count(output)) { + clear_media_feature(output); + } else { + state_arena_.erase("media:" + output); + } + } + } + const size_t store_bytes_after = runtime_state_store_byte_size(); + if (store_bytes_before > store_bytes_after) { + runtime_profile_.released_runtime_store_bytes += store_bytes_before - store_bytes_after; + } + observe_state_arena_bytes(); +} + +void Model::copy_component_outputs_to_inputs(Component& source, Component& target) { + persist_component_outputs(source); FOR_EACH_MATCHED_OUTPUT(source, target, { + bool alias_candidate = false; + for (const auto& alias : runtime_aliases_) { + if (alias.source_component != source.name || alias.source_output != out_name + || alias.target_component != target.name || alias.target_input != out_name) continue; + const bool nodes_match = (alias.source_node_id < 0 || alias.source_node_id == static_cast(src_node)) + && (alias.target_node_id < 0 || alias.target_node_id == static_cast(dst_node)); + const bool policy_allows = alias.policy.empty() || alias.policy == "alias_if_compatible" + || alias.policy == "alias" || alias.policy == "move"; + if (nodes_match && policy_allows + && (alias.storage_stable || alias.source_node_id < 0)) alias_candidate = true; + } + // Generic component splits can use the historical "inputs_embeds" + // logical name for an integer token-id passthrough. Such input-backed + // storage is rewritten on every scalar-prefill step and is not a + // stable embedding tensor to retain in the decoder. Keep the implicit + // zero-copy fast path for actual floating-point feature tensors; an + // explicit runtime-plan alias may still opt another representation in. + const bool input_backed_output = std::find( + source.runtime_input_node_ids.begin(), + source.runtime_input_node_ids.end(), + static_cast(src_node)) != source.runtime_input_node_ids.end(); + const bool floating_feature = !input_backed_output && ( + src_desc.precision == Precision::FP16 + || src_desc.precision == Precision::FP32); + alias_candidate = alias_candidate || (floating_feature && ( + out_name == "inputs_embeds" + || out_name == "encoder_hidden_states" + || out_name == "image_features" + || out_name == "vision_features")); + if (alias_candidate && std::getenv("CACTUS_DISABLE_OUTPUT_ALIAS") == nullptr) { + auto storage = source.graph->export_tensor_storage(src_node); + if (storage && target.graph->bind_tensor_storage(dst_node, storage)) { + runtime_profile_.aliased_bytes += storage->byte_size; + continue; + } + } std::fill(dst_buf.begin(), dst_buf.end(), 0); size_t elements = std::min(src_desc.total_size, dst_desc.total_size); if (!copy_component_tensor(*source.graph, src_desc, src_node, dst_desc, dst_buf, 0, elements, out_name)) throw std::runtime_error("component output/input precision mismatch for " + out_name); + runtime_profile_.copied_bytes += std::min(src_desc.byte_size, dst_desc.byte_size); }) + bind_runtime_state_inputs(target); } -bool Model::copy_cross_kv_outputs_to_decoder_cache_inputs(const Component& source, Component& target, size_t source_len) { +bool Model::copy_cross_kv_outputs_to_decoder_cache_inputs(Component& source, Component& target, size_t source_len) { + persist_component_outputs(source); bool copied_any = false; for (size_t i = 0; i < source.output_node_ids.size() && i < source.logical_outputs.size(); ++i) { const std::string& out_name = source.logical_outputs[i]; @@ -1106,28 +1874,21 @@ bool Model::copy_cross_kv_outputs_to_decoder_cache_inputs(const Component& sourc } copied_any = true; } + bind_runtime_state_inputs(target); return copied_any; } -void Model::copy_component_outputs_to_chunk_inputs(const Component& source, Component& target, size_t token_index) { - FOR_EACH_MATCHED_OUTPUT(source, target, { - size_t chunk_tokens = component_chunk_tokens(target, out_name); - if (chunk_tokens <= token_index || chunk_tokens == 0) - throw std::runtime_error("chunk prefill token index exceeds input capacity for " + out_name); - if (dst_desc.total_size % chunk_tokens != 0) - throw std::runtime_error("chunk prefill input shape is not token-aligned for " + out_name); - size_t elements_per_token = dst_desc.total_size / chunk_tokens; - if (src_desc.total_size != elements_per_token) - throw std::runtime_error("component output/input token shape mismatch for " + out_name); - if (!copy_component_tensor(*source.graph, src_desc, src_node, dst_desc, - dst_buf, token_index * elements_per_token, src_desc.total_size, out_name)) - throw std::runtime_error("component output/input precision mismatch for " + out_name); - }) +void Model::copy_component_outputs_to_chunk_inputs(Component& source, Component& target, size_t token_index) { + copy_component_outputs_to_chunk_inputs_range(source, target, token_index, 1); } -void Model::copy_component_outputs_to_chunk_inputs_range(const Component& source, Component& target, size_t token_offset) { +void Model::copy_component_outputs_to_chunk_inputs_range( + Component& source, + Component& target, + size_t token_offset, + size_t source_tokens_override) { FOR_EACH_MATCHED_OUTPUT(source, target, { - size_t src_tokens = component_output_tokens(source, out_name); + size_t src_tokens = source_tokens_override ? source_tokens_override : component_output_tokens(source, out_name); size_t dst_tokens = component_chunk_tokens(target, out_name); if (src_tokens == 0 || dst_tokens == 0 || token_offset + src_tokens > dst_tokens) throw std::runtime_error("chunk prefill output range exceeds input capacity for " + out_name); @@ -1137,9 +1898,19 @@ void Model::copy_component_outputs_to_chunk_inputs_range(const Component& source size_t dst_elements_per_token = dst_desc.total_size / dst_tokens; if (src_elements_per_token != dst_elements_per_token) throw std::runtime_error("component output/input token shape mismatch for " + out_name); + if (out_name == "inputs_embeds" + && std::getenv("CACTUS_DISABLE_OUTPUT_ALIAS") == nullptr + && token_offset == 0 && src_tokens == dst_tokens) { + auto storage = source.graph->export_tensor_storage(src_node); + if (storage && target.graph->bind_tensor_storage(dst_node, storage)) { + runtime_profile_.aliased_bytes += storage->byte_size; + continue; + } + } if (!copy_component_tensor(*source.graph, src_desc, src_node, dst_desc, dst_buf, token_offset * dst_elements_per_token, src_desc.total_size, out_name)) throw std::runtime_error("component output/input precision mismatch for " + out_name); + runtime_profile_.copied_bytes += src_desc.byte_size; }) } @@ -1147,33 +1918,50 @@ void Model::copy_component_outputs_to_chunk_inputs_range(const Component& source bool Model::cache_states_compatible(const Component& source, const Component& target) const { if (source.cache_states.empty() || source.cache_states.size() != target.cache_states.size()) return false; + if (!source.graph || !target.graph) return false; for (size_t i = 0; i < source.cache_states.size(); ++i) { const auto& src = source.cache_states[i]; const auto& dst = target.cache_states[i]; if (src.layer_key != dst.layer_key) return false; if (src.key_node_id < 0 || src.value_node_id < 0 || dst.key_node_id < 0 || dst.value_node_id < 0) return false; + for (auto [src_node, dst_node] : {std::pair{src.key_node_id, dst.key_node_id}, std::pair{src.value_node_id, dst.value_node_id}}) { + const size_t src_id = static_cast(src_node); + const size_t dst_id = static_cast(dst_node); + if (source.graph->get_node_op_type(src_id) != target.graph->get_node_op_type(dst_id)) return false; + } } return true; } void Model::move_cache_states(Component& source, Component& target, size_t logical_current) { - if (source.cache_states.empty() || source.cache_states.size() != target.cache_states.size()) { + if (!cache_states_compatible(source, target)) { throw std::runtime_error("prefill and step cache states are not compatible"); } for (size_t i = 0; i < source.cache_states.size(); ++i) { const auto& src = source.cache_states[i]; const auto& dst = target.cache_states[i]; - if (src.layer_key != dst.layer_key) { - throw std::runtime_error("prefill and step cache layer mismatch: " + src.layer_key + " != " + dst.layer_key); - } int moved_src = -1, moved_dst = -1; - for (auto [src_node, dst_node] : {std::pair{src.key_node_id, dst.key_node_id}, std::pair{src.value_node_id, dst.value_node_id}}) { + for (auto [src_node, dst_node, key] : { + std::tuple{src.key_node_id, dst.key_node_id, true}, + std::tuple{src.value_node_id, dst.value_node_id, false}}) { if (src_node < 0 || dst_node < 0) continue; // Conv/recurrent caches share one node for key and value; move it only once. if (src_node == moved_src && dst_node == moved_dst) continue; moved_src = src_node; moved_dst = dst_node; - target.graph->steal_cache_buffer(static_cast(dst_node), *source.graph, static_cast(src_node)); + auto storage = std::getenv("CACTUS_DISABLE_SHARED_STATE") == nullptr + ? source.graph->export_tensor_storage(static_cast(src_node)) + : nullptr; + const bool attached = storage + && target.graph->bind_tensor_storage(static_cast(dst_node), storage); + if (attached) { + state_arena_.publish(cache_state_id(src, key), storage); + runtime_profile_.aliased_bytes += storage->byte_size; + observe_state_arena_bytes(); + } else { + target.graph->steal_cache_buffer(static_cast(dst_node), *source.graph, + static_cast(src_node)); + } if (target.graph->get_node_op_type(static_cast(dst_node)) == OpType::KV_CACHE_STATE && logical_current != std::numeric_limits::max()) { auto* meta = static_cast(target.graph->get_output(static_cast(dst_node))); @@ -1252,6 +2040,77 @@ size_t Model::component_output_tokens(const Component& comp, const std::string& return 0; } +size_t Model::audio_soft_token_count_for_frames(size_t frame_count) { + if (!audio_encoder_) return 0; + if (!load_component_graph(*audio_encoder_)) return 0; + + int feature_idx = -1; + for (const auto& name : {"input_features", "audio_features"}) { + int idx = input_index(*audio_encoder_, name); + if (idx >= 0) { + feature_idx = idx; + break; + } + } + if (feature_idx < 0 || audio_encoder_->output_node_ids.empty()) return 0; + + const auto& input_desc = audio_encoder_->graph->get_output_buffer( + static_cast(audio_encoder_->runtime_input_node_ids[feature_idx])); + const auto& output_desc = audio_encoder_->graph->get_output_buffer( + static_cast(audio_encoder_->output_node_ids.front())); + + const size_t frames_per_chunk = input_desc.shape.size() >= 2 ? input_desc.shape[1] : 0; + size_t rows_per_chunk = 0; + if (output_desc.shape.size() >= 2) { + rows_per_chunk = output_desc.shape[output_desc.shape.size() - 2]; + } else if (!output_desc.shape.empty()) { + rows_per_chunk = output_desc.shape[0]; + } + if (frames_per_chunk == 0 || rows_per_chunk == 0) return 0; + + const size_t chunks = frame_count == 0 + ? 1 + : (frame_count + frames_per_chunk - 1) / frames_per_chunk; + return chunks * rows_per_chunk; +} + +size_t Model::media_rows_to_copy(size_t output_rows, size_t frames_in_chunk) const { + if (output_rows == 0) return 0; + if (audio_rows_frame_divisor_ == 0) return output_rows; + size_t expected_rows = (frames_in_chunk + audio_rows_frame_divisor_ - 1) / audio_rows_frame_divisor_; + return std::min(output_rows, expected_rows); +} + +size_t sequence_length_from_shape(const std::vector& shape) { + if (shape.empty()) return 0; + if (shape.size() >= 2 && shape[0] == 1) return shape[1]; + if (shape.size() >= 2) return shape[shape.size() - 2]; + return shape[0]; +} + +void set_sequence_length_in_shape(std::vector& shape, size_t sequence_length) { + if (shape.empty()) return; + if (shape.size() >= 2 && shape[0] == 1 && shape[1] == 1) { + shape[1] = sequence_length; + return; + } + if (shape.size() >= 2 && shape[shape.size() - 2] == 1) { + shape[shape.size() - 2] = sequence_length; + return; + } + if (shape.size() == 1 && shape[0] == 1) { + shape[0] = sequence_length; + } +} + +std::string Model::image_preprocess_strategy() const { + return image_preprocess_strategy_; +} + +std::string Model::audio_preprocess_strategy() const { + return audio_preprocess_strategy_; +} + void Model::execute_prefill_chunk(Component& chunk_comp, Component* enc_comp, size_t encoder_chunk, size_t chunk_tokens, const std::vector& tokens, size_t processed, size_t start_position) { @@ -1280,6 +2139,10 @@ void Model::execute_prefill_chunk(Component& chunk_comp, Component* enc_comp, si copy_component_outputs_to_chunk_inputs(*encoder_, chunk_comp, i); } } + const size_t real_chunk_tokens = tokens.size() > processed ? std::min(chunk_tokens, tokens.size() - processed) : 0; + for (size_t i = 0; i < real_chunk_tokens; ++i) { + write_int_input_at(chunk_comp, "attention_mask", i, 1); + } chunk_comp.graph->execute(); } @@ -1295,18 +2158,53 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector size_t { + if (!component) return 0; + auto it = component->metadata.find("prefill_chunk_tokens"); + if (it == component->metadata.end()) return 0; + try { + return static_cast(std::stoull(it->second)); + } catch (...) { + return 0; + } + }; + const size_t selected_tokens = declared_prefill_tokens(prefill_component); + const size_t shared_tokens = declared_prefill_tokens(decoder_prefill_chunk_); + const bool selected_fits = selected_tokens > 1 && selected_tokens <= tokens.size(); + const bool shared_fits = shared_tokens > 1 && shared_tokens <= tokens.size(); + const bool use_shared_component = decoder_prefill_chunk_ + && decoder_prefill_chunk_ != prefill_component + && ((shared_fits && (!selected_fits || shared_tokens > selected_tokens)) + || (!selected_fits && !shared_fits && shared_tokens > 1 + && (selected_tokens <= 1 || shared_tokens < selected_tokens))); + if (use_shared_component) { + prefill_component = decoder_prefill_chunk_; + for (Component* candidate : {lm_encoder_text_chunk_, lm_encoder_, prefill_encoder_}) { + if (candidate && declared_prefill_tokens(candidate) == shared_tokens) { + encoder_component = candidate; + break; + } + } + } + result.component = prefill_component; + result.encoder_component = encoder_component; + if (start_position != 0) { + if (std::getenv("CACTUS_DISABLE_SHARED_STATE") != nullptr || !decoder_->graph) return result; + publish_component_cache_states(*decoder_); + } + if (!load_component_graph(*prefill_component)) return result; + if (encoder_component && !load_component_graph(*encoder_component)) return result; + if (!cache_states_compatible(*prefill_component, *decoder_)) return result; + size_t component_tokens = component_chunk_tokens(*prefill_component, "inputs_embeds"); if (component_tokens <= 1) return result; size_t effective_chunk = chunk_size > 0 ? std::min(chunk_size, component_tokens) : component_tokens; if (effective_chunk != component_tokens) effective_chunk = component_tokens; size_t whole_chunks_end = (tokens.size() / effective_chunk) * effective_chunk; auto any_cache_node = [&](auto predicate) { - if (!decoder_prefill_->graph) return false; - for (const auto& state : decoder_prefill_->cache_states) { + if (!prefill_component->graph) return false; + for (const auto& state : prefill_component->cache_states) { for (int node_id : {state.key_node_id, state.value_node_id}) { if (node_id < 0) continue; if (predicate(static_cast(node_id))) return true; @@ -1315,28 +2213,28 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vectorgraph->get_node_op_type(id) == OpType::RECURRENT_CACHE_STATE; + return prefill_component->graph->get_node_op_type(id) == OpType::RECURRENT_CACHE_STATE; }); if (has_recurrent_state && whole_chunks_end > effective_chunk) { whole_chunks_end = effective_chunk; } const bool has_sliding_window_cache = any_cache_node([&](size_t id) { - return decoder_prefill_->graph->get_node_op_type(id) == OpType::KV_CACHE_STATE - && decoder_prefill_->graph->get_node_window_size(id) > 0; + return prefill_component->graph->get_node_op_type(id) == OpType::KV_CACHE_STATE + && prefill_component->graph->get_node_window_size(id) > 0; }); - const size_t tail_tokens = tokens.size() - whole_chunks_end; - const size_t padding_cutoff = std::max(1, effective_chunk / 16); const bool has_conv_state = any_cache_node([&](size_t id) { - return decoder_prefill_->graph->get_node_op_type(id) == OpType::CONV_CACHE_STATE; + return prefill_component->graph->get_node_op_type(id) == OpType::CONV_CACHE_STATE; }); + const size_t tail_tokens = tokens.size() - whole_chunks_end; + const size_t padding_cutoff = std::max(1, effective_chunk / 16); const bool pad_tail = !has_conv_state && !has_recurrent_state && !has_sliding_window_cache && tail_tokens >= padding_cutoff; const bool padded_window_too_small = any_cache_node([&](size_t id) { - if (decoder_prefill_->graph->get_node_op_type(id) != OpType::KV_CACHE_STATE) return false; - size_t window = decoder_prefill_->graph->get_node_window_size(id); - size_t sink = decoder_prefill_->graph->get_node_sink_size(id); + if (prefill_component->graph->get_node_op_type(id) != OpType::KV_CACHE_STATE) return false; + size_t window = prefill_component->graph->get_node_window_size(id); + size_t sink = prefill_component->graph->get_node_sink_size(id); return window > 0 && (window < effective_chunk * 4 || window <= effective_chunk + sink); }); const bool use_padded_tail = !pad_tail && !prefill_tail_pad_disabled_ @@ -1350,8 +2248,8 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector= 0 && input_index(*prefill_encoder_, "position_ids") >= 0) { - encoder_chunk = component_chunk_tokens(*prefill_encoder_, "input_ids"); + if (encoder_component && input_index(*encoder_component, "input_ids") >= 0 && input_index(*encoder_component, "position_ids") >= 0) { + encoder_chunk = component_chunk_tokens(*encoder_component, "input_ids"); if (encoder_chunk == 0 || effective_chunk % encoder_chunk != 0) { encoder_chunk = 0; } @@ -1359,7 +2257,7 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector>> backups; - for (const auto& state : decoder_prefill_->cache_states) { + for (const auto& state : prefill_component->cache_states) { for (int node_id : {state.key_node_id, state.value_node_id}) { if (node_id < 0) continue; size_t id = static_cast(node_id); - if (decoder_prefill_->graph->get_node_op_type(id) != OpType::KV_CACHE_STATE) continue; - backups.emplace_back(id, decoder_prefill_->graph->snapshot_cache_padded_append(id, kept_real, pads + 1)); + if (prefill_component->graph->get_node_op_type(id) != OpType::KV_CACHE_STATE) continue; + backups.emplace_back(id, prefill_component->graph->snapshot_cache_padded_append(id, kept_real, pads + 1)); } } - execute_prefill_chunk(*decoder_prefill_, prefill_encoder_, encoder_chunk, + execute_prefill_chunk(*prefill_component, encoder_component, encoder_chunk, effective_chunk, tokens, processed, start_position); for (auto& [id, backup] : backups) { - decoder_prefill_->graph->rollback_cache_padded_append(id, kept_real, pads + 1, backup); + prefill_component->graph->rollback_cache_padded_append(id, kept_real, pads + 1, backup); } processed += kept_real; tail_executed = kept_real; @@ -1405,7 +2303,7 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vectorinput_buffers[i].begin(), decoder_->input_buffers[i].end(), 0); } auto copy_start = std::chrono::high_resolution_clock::now(); - move_cache_states(*decoder_prefill_, *decoder_, start_position + result.logical_tokens); + move_cache_states(*prefill_component, *decoder_, start_position + result.logical_tokens); auto copy_end = std::chrono::high_resolution_clock::now(); last_prefill_cache_copy_ms_ = std::chrono::duration_cast(copy_end - copy_start).count() / 1000.0; } @@ -1413,54 +2311,67 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector(encoder_->runtime_input_node_ids[input_ids_idx]); - const auto& input_desc = encoder_->graph->get_output_buffer(input_node); + size_t input_node = static_cast(input_component->runtime_input_node_ids[input_ids_idx]); + const auto& input_desc = input_component->graph->get_output_buffer(input_node); if (context_tokens_.size() > input_desc.total_size) { - throw std::runtime_error("context exceeds transpiled text_lm_encoder capacity"); + throw std::runtime_error("context exceeds graph-bundle full-context capacity"); } - std::fill(encoder_->input_buffers[input_ids_idx].begin(), encoder_->input_buffers[input_ids_idx].end(), 0); - std::fill(encoder_->input_buffers[attention_mask_idx].begin(), encoder_->input_buffers[attention_mask_idx].end(), 0); + std::fill(input_component->input_buffers[input_ids_idx].begin(), input_component->input_buffers[input_ids_idx].end(), 0); + std::fill(input_component->input_buffers[attention_mask_idx].begin(), input_component->input_buffers[attention_mask_idx].end(), 0); for (size_t i = 0; i < context_tokens_.size(); ++i) { - write_int_input_at(*encoder_, "input_ids", i, static_cast(context_tokens_[i])); - write_int_input_at(*encoder_, "attention_mask", i, 1); - } - encoder_->graph->execute(); - for (size_t i = 0; i < encoder_->output_node_ids.size() && i < encoder_->logical_outputs.size(); ++i) { - const std::string& out_name = encoder_->logical_outputs[i]; - int dst_idx = input_index(*decoder_, out_name); - if (dst_idx < 0) continue; - size_t src_node = static_cast(encoder_->output_node_ids[i]); - const auto& src_desc = encoder_->graph->get_output_buffer(src_node); - void* src_ptr = encoder_->graph->get_output(src_node); - std::memcpy(decoder_->input_buffers[dst_idx].data(), src_ptr, src_desc.byte_size); + const char* token_input_name = input_index(*input_component, "input_ids") >= 0 + ? "input_ids" + : "decoder_input_ids"; + write_int_input_at(*input_component, token_input_name, i, static_cast(context_tokens_[i])); + write_int_input_at(*input_component, "attention_mask", i, 1); + } + input_component->graph->execute(); + if (encoder_) { + for (size_t i = 0; i < encoder_->output_node_ids.size() && i < encoder_->logical_outputs.size(); ++i) { + const std::string& out_name = encoder_->logical_outputs[i]; + int dst_idx = input_index(*decoder_, out_name); + if (dst_idx < 0) continue; + size_t src_node = static_cast(encoder_->output_node_ids[i]); + const auto& src_desc = encoder_->graph->get_output_buffer(src_node); + void* src_ptr = encoder_->graph->get_output(src_node); + std::memcpy(decoder_->input_buffers[dst_idx].data(), src_ptr, src_desc.byte_size); + } + decoder_->graph->execute(); } last_logit_position_ = context_tokens_.empty() ? 0 : context_tokens_.size() - 1; - decoder_->graph->execute(); } void Model::write_media_embeds_row(Component& comp, int embeds_idx, const uint8_t* feature_row, size_t feature_row_bytes, Precision feature_precision) { - auto& embeds_buf = comp.input_buffers[embeds_idx]; size_t node_id = static_cast(comp.runtime_input_node_ids[embeds_idx]); const auto& desc = comp.graph->get_output_buffer(node_id); + auto* embeds_data = static_cast(comp.graph->get_output(node_id)); + if (!embeds_data) { + throw std::runtime_error("media inputs_embeds storage is unavailable"); + } if (desc.precision == feature_precision) { - size_t to_copy = std::min(feature_row_bytes, embeds_buf.size()); - std::memcpy(embeds_buf.data(), feature_row, to_copy); - if (to_copy < embeds_buf.size()) { - std::memset(embeds_buf.data() + to_copy, 0, embeds_buf.size() - to_copy); + size_t to_copy = std::min(feature_row_bytes, desc.byte_size); + std::memcpy(embeds_data, feature_row, to_copy); + if (to_copy < desc.byte_size) { + std::memset(embeds_data + to_copy, 0, desc.byte_size - to_copy); } return; } + size_t src_elem = PrecisionTraits::size_of(feature_precision); size_t dst_elem = PrecisionTraits::size_of(desc.precision); size_t src_count = src_elem ? feature_row_bytes / src_elem : 0; - size_t dst_count = dst_elem ? embeds_buf.size() / dst_elem : 0; + size_t dst_count = dst_elem ? desc.byte_size / dst_elem : 0; size_t n = std::min(src_count, dst_count); auto load_float = [&](size_t i) -> float { if (feature_precision == Precision::FP16) return static_cast(reinterpret_cast(feature_row)[i]); @@ -1469,12 +2380,12 @@ void Model::write_media_embeds_row(Component& comp, int embeds_idx, const uint8_ }; for (size_t i = 0; i < n; ++i) { float v = load_float(i); - if (desc.precision == Precision::FP16) reinterpret_cast<__fp16*>(embeds_buf.data())[i] = static_cast<__fp16>(v); - else if (desc.precision == Precision::FP32) reinterpret_cast(embeds_buf.data())[i] = v; - else reinterpret_cast(embeds_buf.data())[i] = static_cast(v); + if (desc.precision == Precision::FP16) reinterpret_cast<__fp16*>(embeds_data)[i] = static_cast<__fp16>(v); + else if (desc.precision == Precision::FP32) reinterpret_cast(embeds_data)[i] = v; + else reinterpret_cast(embeds_data)[i] = static_cast(v); } if (n < dst_count) { - std::memset(embeds_buf.data() + n * dst_elem, 0, (dst_count - n) * dst_elem); + std::memset(embeds_data + n * dst_elem, 0, (dst_count - n) * dst_elem); } } @@ -1483,11 +2394,35 @@ void Model::run_media_step(size_t position, const uint8_t* feature_row, size_t f if (lm_encoder_media_step_) { int embeds_idx = input_index(*lm_encoder_media_step_, "inputs_embeds"); if (embeds_idx >= 0) { + if (encoder_ != nullptr) { + run_encoder_step(static_cast(media_placeholder_token_id_), position); + copy_component_outputs_to_inputs(*encoder_, *decoder_); + } write_media_embeds_row(*lm_encoder_media_step_, embeds_idx, feature_row, feature_row_bytes, feature_precision); - write_int_input(*lm_encoder_media_step_, "input_ids", 0); + write_int_input(*lm_encoder_media_step_, "input_ids", static_cast(media_placeholder_token_id_)); write_int_input(*lm_encoder_media_step_, "position_ids", static_cast(position)); lm_encoder_media_step_->graph->execute(); - copy_encoder_outputs_to_decoder(*lm_encoder_media_step_); + persist_component_outputs(*lm_encoder_media_step_); + for (size_t i = 0; i < lm_encoder_media_step_->output_node_ids.size() && i < lm_encoder_media_step_->logical_outputs.size(); ++i) { + const std::string& out_name = lm_encoder_media_step_->logical_outputs[i]; + auto source_it = media_chunk_output_sources_.find(out_name); + if (source_it != media_chunk_output_sources_.end() && source_it->second == "text_chunk") { + continue; + } + int dst_idx = input_index(*decoder_, out_name); + if (dst_idx < 0) continue; + size_t src_node = static_cast(lm_encoder_media_step_->output_node_ids[i]); + const auto& src_desc = lm_encoder_media_step_->graph->get_output_buffer(src_node); + size_t dst_node = static_cast(decoder_->runtime_input_node_ids[dst_idx]); + const auto& dst_desc = decoder_->graph->get_output_buffer(dst_node); + auto& dst_buf = decoder_->input_buffers[dst_idx]; + std::fill(dst_buf.begin(), dst_buf.end(), 0); + size_t elements = std::min(src_desc.total_size, dst_desc.total_size); + if (!copy_component_tensor(*lm_encoder_media_step_->graph, src_desc, src_node, dst_desc, dst_buf, 0, elements, out_name)) { + throw std::runtime_error("component output/input precision mismatch for " + out_name); + } + } + bind_runtime_state_inputs(*decoder_); decoder_->graph->execute(); return; } @@ -1501,6 +2436,7 @@ void Model::run_media_step(size_t position, const uint8_t* feature_row, size_t f return; } } + run_step(static_cast(config_.pad_token_id), position, false); } @@ -1511,7 +2447,8 @@ void write_typed_buffer(std::vector& buf, Precision dst_prec, void Model::run_vision_encoder(const std::string& image_path) { if (!vision_encoder_) return; - if (family_ == "lfm2_vl") { + const std::string strategy = image_preprocess_strategy(); + if (strategy == "lfm2_vl") { run_vision_encoder_lfm2_vl(image_path); return; } @@ -1541,7 +2478,7 @@ void Model::run_vision_encoder(const std::string& image_path) { if (n < cap) std::memset(buf.data() + n * elem, 0, (cap - n) * elem); }; - if (family_ == "qwen3_5" || family_ == "qwen3_vl" || config_.model_type == Config::ModelType::QWEN) { + if (strategy == "qwen3_vl") { Qwen3VlImagePreprocessed prep = preprocess_qwen3_vl_image(image_path, config_); int pv_idx = input_index(*vision_encoder_, "pixel_values"); if (pv_idx < 0) { @@ -1555,7 +2492,19 @@ void Model::run_vision_encoder(const std::string& image_path) { prep.pixel_values.size() * sizeof(float), Precision::FP32); } else { + if (!strategy.empty() && strategy != "gemma4") { + throw std::runtime_error("Unsupported image preprocess strategy: " + strategy); + } Gemma4ImagePreprocessed prep = preprocess_gemma4_image(image_path, config_); + debug_media_vector("vision_preprocess", "pixel_values", + prep.pixel_values, {prep.max_patches, prep.patch_dim}, + Precision::FP32); + debug_media_vector("vision_preprocess", "pixel_position_ids", + prep.pixel_position_ids, {prep.max_patches, 2}, + Precision::INT8); + if (has_npu_vision_encoder() && vision_encode_via_npu(prep.pixel_values, &prep.pixel_position_ids)) { + return; + } int pv_idx = input_index(*vision_encoder_, "pixel_values"); if (pv_idx >= 0) { auto& pv_buf = vision_encoder_->input_buffers[pv_idx]; @@ -1578,12 +2527,13 @@ void Model::run_vision_encoder(const std::string& image_path) { const std::string& name = vision_encoder_->logical_outputs[i]; size_t node_id = static_cast(vision_encoder_->output_node_ids[i]); const auto& desc = vision_encoder_->graph->get_output_buffer(node_id); - void* ptr = vision_encoder_->graph->get_output(node_id); - auto& slot = media_features_[name]; - slot.assign(desc.byte_size, 0); - std::memcpy(slot.data(), ptr, desc.byte_size); - media_feature_shapes_[name] = desc.shape; - media_feature_precisions_[name] = desc.precision; + auto storage = vision_encoder_->graph->export_tensor_storage(node_id); + publish_media_feature_storage(name, storage, desc.shape, desc.precision); + if (debug_media_features_enabled()) { + std::vector debug_bytes(desc.byte_size); + std::memcpy(debug_bytes.data(), media_feature_data(name), desc.byte_size); + debug_media_feature("vision_encoder", name, debug_bytes, desc.shape, desc.precision); + } } if (cactus_default_backend() != ComputeBackend::METAL) { vision_encoder_->graph->release_runtime_buffers(); @@ -1644,9 +2594,7 @@ void Model::run_vision_encoder_lfm2_vl(const std::string& image_path) { if (!load_component_graph(*vision_encoder_)) throw std::runtime_error("failed to load vision_encoder"); if (!load_component_graph(*vision_projector_)) throw std::runtime_error("failed to load vision_projector"); - media_features_.erase("image_features"); - media_feature_shapes_.erase("image_features"); - media_feature_precisions_.erase("image_features"); + clear_media_feature("image_features"); encode_lfm2_vl_image_into_features(image_path); @@ -1717,9 +2665,6 @@ void Model::encode_lfm2_vl_image_into_features(const std::string& image_path) { write_typed_buffer(buf, d.precision, src, count * sizeof(float), Precision::FP32); }; - std::vector image_features; - size_t total_tokens = 0; - std::vector pos_buf(max_patches * static_cast(dim)); std::vector enc_out_f(max_patches * static_cast(dim)); std::vector unshuf; @@ -1760,20 +2705,10 @@ void Model::encode_lfm2_vl_image_into_features(const std::string& image_path) { const void* proj_ptr = vision_projector_->graph->get_output(proj_out_node); const size_t row_bytes = text_hidden * proj_elem; const size_t append_bytes = static_cast(num_tokens) * row_bytes; - const size_t base = image_features.size(); - image_features.resize(base + append_bytes); - std::memcpy(image_features.data() + base, proj_ptr, append_bytes); - total_tokens += static_cast(num_tokens); + append_media_feature("image_features", proj_ptr, append_bytes, + {static_cast(num_tokens), text_hidden}, + proj_prec, static_cast(num_tokens)); } - - auto& slot = media_features_["image_features"]; - const size_t prev_bytes = slot.size(); - slot.resize(prev_bytes + image_features.size()); - std::memcpy(slot.data() + prev_bytes, image_features.data(), image_features.size()); - auto& shape = media_feature_shapes_["image_features"]; - if (shape.size() != 2) shape = { total_tokens, text_hidden }; - else shape[0] += total_tokens; - media_feature_precisions_["image_features"] = proj_prec; } void Model::run_audio_encoder_messages(const std::vector>& audio_features_per_message) { @@ -1783,9 +2718,7 @@ void Model::run_audio_encoder_messages(const std::vector>& au throw std::runtime_error("failed to load audio_encoder"); } for (const std::string& logical : audio_encoder_->logical_outputs) { - media_features_.erase(logical); - media_feature_shapes_.erase(logical); - media_feature_precisions_.erase(logical); + clear_media_feature(logical); } for (const auto& mel : audio_features_per_message) { if (mel.empty()) continue; @@ -1880,17 +2813,26 @@ void Model::run_audio_encoder(const std::vector& audio_features) { const size_t node_id = static_cast(audio_encoder_->output_node_ids[i]); const auto& desc = audio_encoder_->graph->get_output_buffer(node_id); void* ptr = audio_encoder_->graph->get_output(node_id); - auto& slot = media_features_[name]; - const size_t prev_bytes = slot.size(); - slot.resize(prev_bytes + desc.byte_size); - std::memcpy(slot.data() + prev_bytes, ptr, desc.byte_size); - auto shape_it = media_feature_shapes_.find(name); - if (shape_it == media_feature_shapes_.end() || shape_it->second.empty()) { - media_feature_shapes_[name] = desc.shape; - } else if (desc.shape.size() >= 2 && shape_it->second.size() == desc.shape.size()) { - shape_it->second[shape_it->second.size() - 2] += desc.shape[desc.shape.size() - 2]; + const size_t output_rows = desc.shape.size() >= 2 + ? desc.shape[desc.shape.size() - 2] + : desc.shape.empty() ? 0 : desc.shape[0]; + size_t rows_to_copy = media_rows_to_copy(output_rows, frames_in_chunk); + if (rows_to_copy == 0) continue; + const size_t row_bytes = output_rows > 0 ? desc.byte_size / output_rows : desc.byte_size; + const size_t append_bytes = rows_to_copy * row_bytes; + if (num_chunks == 1 && rows_to_copy == output_rows) { + publish_media_feature_storage( + name, audio_encoder_->graph->export_tensor_storage(node_id), + desc.shape, desc.precision); + } else { + append_media_feature(name, ptr, append_bytes, desc.shape, desc.precision, rows_to_copy); + } + if (debug_media_features_enabled()) { + std::vector debug_bytes(media_feature_byte_size(name)); + std::memcpy(debug_bytes.data(), media_feature_data(name), debug_bytes.size()); + debug_media_feature("audio_encoder", name, debug_bytes, + media_feature_shapes_[name], desc.precision); } - media_feature_precisions_[name] = desc.precision; } } } @@ -1904,9 +2846,10 @@ static float uncertainty_from_margin(float best, float second) { return std::max(0.0f, std::min(1.0f, 1.0f - confidence)); } -uint32_t Model::argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_off, float* out_uncertainty) { +uint32_t Model::argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_off, float* out_uncertainty, + float repetition_penalty) { size_t vocab = desc.shape.empty() ? 0 : desc.shape.back(); - uint32_t best = 0; + uint32_t best = std::numeric_limits::max(); float best_v = -std::numeric_limits::infinity(); float second_v = -std::numeric_limits::infinity(); const auto& tool_bias = tool_constrainer_.get_bias(); @@ -1914,6 +2857,8 @@ uint32_t Model::argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_o { uint32_t gidx; float gbest, gsecond; if (!tool_dense && tool_bias.empty() && vocab_bias_.empty() && suppressed_token_id_ < 0 && + generation_suppressed_token_ids_.empty() && media_min_new_tokens_remaining_ == 0 && + repetition_penalty == 1.0f && cactus_graph_metal_argmax(&gidx, &gbest, &gsecond)) { if (out_uncertainty) *out_uncertainty = uncertainty_from_margin(gbest, gsecond); return gidx; @@ -1925,10 +2870,15 @@ uint32_t Model::argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_o if (tool_it != tool_bias.end()) value += tool_it->second; auto vocab_it = vocab_bias_.find(static_cast(token_id)); if (vocab_it != vocab_bias_.end()) value += vocab_it->second; + const bool in_history = token_history_set_.count(static_cast(token_id)) != 0; + if (repetition_penalty > 1.0f && in_history) { + value = value >= 0.0f ? value / repetition_penalty : value * repetition_penalty; + } return value; }; auto observe_logit = [&](size_t i, float v) { if (static_cast(i) == suppressed_token_id_) return; + if (is_generation_suppressed_token(static_cast(i))) return; v = score_with_bias(i, v); if (v > best_v) { second_v = best_v; @@ -1949,10 +2899,65 @@ uint32_t Model::argmax_logits_at(const BufferDesc& desc, void* ptr, size_t row_o for (size_t i = 0; i < vocab; ++i) observe_logit(i, static_cast(p[i])); } if (out_uncertainty) *out_uncertainty = uncertainty_from_margin(best_v, second_v); + if (best == std::numeric_limits::max()) { + best = 0; + } + if (std::getenv("CACTUS_DEBUG_TOPK") != nullptr) { + std::vector> topk; + auto collect_logit = [&](size_t i, float v) { + if (static_cast(i) == suppressed_token_id_) return; + if (is_generation_suppressed_token(static_cast(i))) return; + v = score_with_bias(i, v); + topk.emplace_back(v, static_cast(i)); + }; + if (desc.precision == Precision::FP32) { + float* p = static_cast(ptr) + row_off; + for (size_t i = 0; i < vocab; ++i) collect_logit(i, p[i]); + } else if (desc.precision == Precision::FP16) { + __fp16* p = static_cast<__fp16*>(ptr) + row_off; + for (size_t i = 0; i < vocab; ++i) collect_logit(i, static_cast(p[i])); + } else { + int8_t* p = static_cast(ptr) + row_off; + for (size_t i = 0; i < vocab; ++i) collect_logit(i, static_cast(p[i])); + } + const size_t limit = std::min(10, topk.size()); + std::partial_sort(topk.begin(), topk.begin() + limit, topk.end(), std::greater<>()); + std::cerr << "[topk-debug] row_off=" << row_off << " best=" << best << " logits="; + for (size_t i = 0; i < limit; ++i) { + const uint32_t token = topk[i].second; + std::cerr << (i == 0 ? "" : " | ") << token << ":" << topk[i].first; + if (tokenizer_) { + std::cerr << ":" << std::quoted(tokenizer_->decode({token})); + } + } + std::cerr << std::endl; + } + if (debug_media_features_enabled() && media_min_new_tokens_remaining_ > 0) { + std::cerr << "[media-debug] selected token while media_min_new_tokens_remaining=" + << media_min_new_tokens_remaining_ << " token=" << best + << " best=" << best_v << " second=" << second_v << std::endl; + } return best; } -uint32_t Model::argmax_component_logits(Component& comp, size_t logit_row, float* out_uncertainty) { +bool Model::is_generation_suppressed_token(uint32_t token_id) const { + if (media_min_new_tokens_remaining_ > 0 && media_min_stop_token_ids_.count(token_id)) { + if (debug_media_features_enabled()) { + std::cerr << "[media-debug] suppress media stop token=" << token_id + << " for media_min_new_tokens_remaining=" + << media_min_new_tokens_remaining_ << std::endl; + } + return true; + } + if (generation_suppressed_token_ids_.count(token_id)) return true; + if (!generation_suppressed_token_ids_.empty()) return false; + if (config_.image_token_id != 0 && token_id == config_.image_token_id) return true; + if (config_.audio_token_id != 0 && token_id == config_.audio_token_id) return true; + return false; +} + +uint32_t Model::argmax_component_logits(Component& comp, size_t logit_row, float* out_uncertainty, + float repetition_penalty) { size_t out_node = static_cast(comp.output_node_ids.empty() ? 0 : comp.output_node_ids[0]); const auto& desc = comp.graph->get_output_buffer(out_node); void* ptr = comp.graph->get_output(out_node); @@ -1964,7 +2969,7 @@ uint32_t Model::argmax_component_logits(Component& comp, size_t logit_row, float } else if (decode_route_ == DecodeRoute::FULL_CONTEXT_TEXT) { row = std::min(last_logit_position_, seq > 0 ? seq - 1 : 0); } - return argmax_logits_at(desc, ptr, row * vocab, out_uncertainty); + return argmax_logits_at(desc, ptr, row * vocab, out_uncertainty, repetition_penalty); } std::vector Model::argmax_component_logits_batch(Component& comp, size_t batch) { @@ -1984,8 +2989,8 @@ std::vector Model::argmax_component_logits_batch(Component& comp, size return out; } -uint32_t Model::argmax_last_logits(float* out_uncertainty) { - return argmax_component_logits(*decoder_, std::numeric_limits::max(), out_uncertainty); +uint32_t Model::argmax_last_logits(float* out_uncertainty, float repetition_penalty) { + return argmax_component_logits(*decoder_, std::numeric_limits::max(), out_uncertainty, repetition_penalty); } void Model::prepare_sampling_context(float repetition_penalty) { @@ -2195,12 +3200,14 @@ bool Model::prefill_and_sample_first_token(const std::vector& tokens, if (decoder_prefill_) { chunked = run_chunked_prefill(tokens, cache_total_seq_len_, get_prefill_chunk_size(), true); if (chunked.logical_tokens == tokens.size() && chunked.padding_tokens > 0 && tokens.size() > 0) { - // Cache already moved into the step component; drop the padded last row and re-run it for logits. + // Cache already moved into the step component; drop the last row and re-run it for step-decoder logits. set_cache_current_len(*decoder_, tokens.size() - 1); cache_total_seq_len_ = tokens.size() - 1; run_step(tokens.back(), cache_total_seq_len_, true, /*use_fused=*/false); ++cache_total_seq_len_; out_token = argmax_last_logits(out_uncertainty); + if (chunked.component) unload_component_graph(*chunked.component); + if (chunked.encoder_component) unload_component_graph(*chunked.encoder_component); record_sampled_token(out_token); last_prefill_scalar_tail_tokens_ = 1; maybe_roll_compact(); @@ -2208,13 +3215,19 @@ bool Model::prefill_and_sample_first_token(const std::vector& tokens, } cache_total_seq_len_ += chunked.logical_tokens; } + if (chunked.logical_tokens < tokens.size() && chunked.logical_tokens > 0) { + if (chunked.component) unload_component_graph(*chunked.component); + if (chunked.encoder_component) unload_component_graph(*chunked.encoder_component); + } for (size_t i = chunked.logical_tokens; i < tokens.size(); ++i) { run_step(tokens[i], cache_total_seq_len_, i + 1 == tokens.size(), /*use_fused=*/false); ++cache_total_seq_len_; } last_prefill_scalar_tail_tokens_ = tokens.size() - chunked.logical_tokens; - if (chunked.logical_tokens == tokens.size() && chunked.logical_tokens > 0 && decoder_prefill_) { - out_token = argmax_component_logits(*decoder_prefill_, chunked.last_logit_row, out_uncertainty); + if (chunked.logical_tokens == tokens.size() && chunked.logical_tokens > 0 && chunked.component) { + out_token = argmax_component_logits(*chunked.component, chunked.last_logit_row, out_uncertainty); + unload_component_graph(*chunked.component); + if (chunked.encoder_component) unload_component_graph(*chunked.encoder_component); } else { out_token = argmax_last_logits(out_uncertainty); } @@ -2239,6 +3252,10 @@ void Model::prefill(const std::vector& tokens, size_t /*chunk_size*/, } ChunkedPrefillResult chunked = run_chunked_prefill(tokens, cache_total_seq_len_, get_prefill_chunk_size(), prepare_decode); cache_total_seq_len_ += chunked.logical_tokens; + if (prepare_decode && chunked.logical_tokens > 0) { + if (chunked.component) unload_component_graph(*chunked.component); + if (chunked.encoder_component) unload_component_graph(*chunked.encoder_component); + } for (size_t i = chunked.logical_tokens; i < tokens.size(); ++i) { run_step(tokens[i], cache_total_seq_len_, /*read_logits=*/false, /*use_fused=*/false); ++cache_total_seq_len_; @@ -2413,10 +3430,20 @@ void Model::reset_encoder_cross_kv_route_state() { reset_component_cache_states(*decoder_); cache_total_seq_len_ = 0; token_history_.clear(); + token_history_set_.clear(); encoder_cross_kv_ready_ = false; encoder_cross_kv_source_len_ = 0; } +void Model::release_encoder_cross_kv_prepare_graphs() { + if (source_encoder_ && source_encoder_ != decoder_) { + unload_component_graph(*source_encoder_); + } + if (decoder_cross_kv_ && decoder_cross_kv_ != decoder_) { + unload_component_graph(*decoder_cross_kv_); + } +} + bool Model::finish_encoder_cross_kv_prepare() { if (!source_encoder_ || !decoder_cross_kv_ || !decoder_) return false; source_encoder_->graph->execute(); @@ -2427,12 +3454,25 @@ bool Model::finish_encoder_cross_kv_prepare() { copy_component_outputs_to_inputs(*decoder_cross_kv_, *decoder_); } encoder_cross_kv_ready_ = true; + release_encoder_cross_kv_prepare_graphs(); + return true; +} + +bool Model::finish_encoder_cross_kv_prepare_after_source() { + if (!decoder_cross_kv_ || !decoder_) return false; + decoder_cross_kv_->graph->execute(); + if (!copy_cross_kv_outputs_to_decoder_cache_inputs(*decoder_cross_kv_, *decoder_, encoder_cross_kv_source_len_)) { + copy_component_outputs_to_inputs(*decoder_cross_kv_, *decoder_); + } + encoder_cross_kv_ready_ = true; + release_encoder_cross_kv_prepare_graphs(); return true; } bool Model::prepare_encoder_cross_kv_from_text(const std::vector& tokens) { if (!source_encoder_ || !decoder_cross_kv_ || !decoder_ || tokens.empty()) return false; if (encoder_cross_kv_source_kind_ != "text_tokens") return false; + if (!load_component_graph(*source_encoder_) || !load_component_graph(*decoder_cross_kv_)) return false; reset_encoder_cross_kv_route_state(); for (auto& buf : source_encoder_->input_buffers) { @@ -2469,6 +3509,7 @@ bool Model::prepare_encoder_cross_kv_from_text(const std::vector& toke bool Model::prepare_encoder_cross_kv_from_audio(const std::vector& audio_features) { if (!source_encoder_ || !decoder_cross_kv_ || !decoder_) return false; + if (!load_component_graph(*source_encoder_) || !load_component_graph(*decoder_cross_kv_)) return false; int feature_idx = input_index(*source_encoder_, "input_features"); if (feature_idx < 0) { feature_idx = input_index(*source_encoder_, "audio_features"); @@ -2547,7 +3588,7 @@ std::vector Model::run_encoder_cross_kv_decode_loop( return emitted; } -bool Model::build_lm_encoder_outputs_dynamic_gemma4( +bool Model::build_lm_encoder_outputs_dynamic_media( const std::vector& tokens, std::map>& store_bytes, std::map& store_prec, @@ -2556,35 +3597,48 @@ bool Model::build_lm_encoder_outputs_dynamic_gemma4( const uint32_t image_tok = config_.image_token_id; const uint32_t audio_tok = config_.audio_token_id; - - auto audio_it = media_features_.find("audio_features"); - const bool have_audio_features = audio_it != media_features_.end() && !audio_it->second.empty(); - size_t audio_rows = 0; - size_t audio_row_bytes = 0; - Precision audio_prec = Precision::FP16; - if (have_audio_features) { - const auto& shape = media_feature_shapes_["audio_features"]; - audio_rows = shape.size() >= 2 ? shape[shape.size() - 2] : 0; - audio_prec = media_feature_precisions_["audio_features"]; - audio_row_bytes = audio_rows > 0 ? audio_it->second.size() / audio_rows : audio_it->second.size(); - } - - auto image_it = media_features_.find("image_features"); - const bool have_image_features = image_it != media_features_.end() && !image_it->second.empty(); - size_t image_rows = 0; - size_t image_row_bytes = 0; - Precision image_prec = Precision::FP16; - if (have_image_features) { - const auto& shape = media_feature_shapes_["image_features"]; - image_rows = shape.size() >= 2 ? shape[shape.size() - 2] : 0; - image_prec = media_feature_precisions_["image_features"]; - image_row_bytes = image_rows > 0 ? image_it->second.size() / image_rows : image_it->second.size(); + size_t image_token_count = 0; + size_t audio_token_count = 0; + for (uint32_t token : tokens) { + if (image_tok != 0 && token == image_tok) ++image_token_count; + if (audio_tok != 0 && token == audio_tok) ++audio_token_count; + } + + const std::vector image_fallback = {"image_features", "image_embeddings", "vision_features", "inputs_embeds"}; + const std::vector audio_fallback = {"audio_features", "audio_embeddings", "encoder_hidden_states", "inputs_embeds"}; + const auto image_feature = image_token_count > 0 + ? find_media_feature(media_features_, media_feature_storages_, media_feature_shapes_, media_feature_precisions_, + image_feature_candidates_, image_fallback) + : MediaFeatureInfo{}; + const auto audio_feature = audio_token_count > 0 + ? find_media_feature(media_features_, media_feature_storages_, media_feature_shapes_, media_feature_precisions_, + audio_feature_candidates_, audio_fallback, image_feature.name) + : MediaFeatureInfo{}; + const bool have_image_features = image_feature.present(); + const bool have_audio_features = audio_feature.present(); + if (debug_media_features_enabled()) { + std::cerr << "[media-debug] inject" + << " image_tokens=" << image_token_count + << " image_feature=" << image_feature.name + << " image_rows=" << image_feature.rows + << " audio_tokens=" << audio_token_count + << " audio_feature=" << audio_feature.name + << " audio_rows=" << audio_feature.rows + << std::endl; + } + if (image_token_count > 0 && (!have_image_features || image_feature.rows < image_token_count)) { + throw std::runtime_error("Prompt image token count exceeds available image feature rows"); + } + if (audio_token_count > 0 && (!have_audio_features || audio_feature.rows < audio_token_count)) { + throw std::runtime_error("Prompt audio token count exceeds available audio feature rows"); } struct OutputInfo { std::string name; int text_idx = -1; + int text_chunk_idx = -1; int media_idx = -1; + int media_chunk_idx = -1; size_t per_token_bytes = 0; Precision precision = Precision::FP16; std::vector shape_template; @@ -2613,13 +3667,201 @@ bool Model::build_lm_encoder_outputs_dynamic_gemma4( store_bytes[info.name].assign(token_count * info.per_token_bytes, 0); store_prec[info.name] = info.precision; std::vector shape = info.shape_template; - if (shape.size() >= 2 && shape[shape.size() - 2] == 1) { - shape[shape.size() - 2] = token_count; - } else if (shape.size() == 1) { - shape[0] = token_count; - } + set_sequence_length_in_shape(shape, token_count); store_shape[info.name] = std::move(shape); } + auto debug_lm_store = [&]() { + for (const auto& kv : store_bytes) { + auto shape_it = store_shape.find(kv.first); + auto prec_it = store_prec.find(kv.first); + debug_media_feature( + "lm_encoder_store", + kv.first, + kv.second, + shape_it == store_shape.end() ? std::vector{} : shape_it->second, + prec_it == store_prec.end() ? Precision::FP16 : prec_it->second); + } + }; + + if (std::getenv("CACTUS_DISABLE_MEDIA_LM_ENCODER_CHUNK") == nullptr + && lm_encoder_text_chunk_ && lm_encoder_media_chunk_ + && load_component_graph(*lm_encoder_text_chunk_) + && load_component_graph(*lm_encoder_media_chunk_)) { + const size_t text_chunk_tokens = component_chunk_tokens(*lm_encoder_text_chunk_, "input_ids"); + const size_t media_chunk_tokens = component_chunk_tokens(*lm_encoder_media_chunk_, "input_ids"); + const int media_embeds_idx = input_index(*lm_encoder_media_chunk_, "inputs_embeds"); + bool can_chunk = text_chunk_tokens > 0 + && text_chunk_tokens == media_chunk_tokens + && media_embeds_idx >= 0 + && input_index(*lm_encoder_text_chunk_, "input_ids") >= 0 + && input_index(*lm_encoder_text_chunk_, "position_ids") >= 0 + && input_index(*lm_encoder_media_chunk_, "input_ids") >= 0 + && input_index(*lm_encoder_media_chunk_, "position_ids") >= 0; + + for (auto& info : outputs) { + info.text_chunk_idx = output_index(*lm_encoder_text_chunk_, info.name); + info.media_chunk_idx = output_index(*lm_encoder_media_chunk_, info.name); + if (info.text_chunk_idx < 0 || info.media_chunk_idx < 0) { + can_chunk = false; + break; + } + for (Component* chunk_comp : {lm_encoder_text_chunk_, lm_encoder_media_chunk_}) { + int output_idx = chunk_comp == lm_encoder_text_chunk_ ? info.text_chunk_idx : info.media_chunk_idx; + size_t node_id = static_cast(chunk_comp->output_node_ids[output_idx]); + const auto& desc = chunk_comp->graph->get_output_buffer(node_id); + if (desc.precision != info.precision || desc.total_size % media_chunk_tokens != 0) { + can_chunk = false; + break; + } + size_t per_token_elements = desc.total_size / media_chunk_tokens; + size_t per_token_bytes = PrecisionTraits::packed_size_of(desc.precision, per_token_elements); + if (per_token_bytes != info.per_token_bytes) { + can_chunk = false; + break; + } + } + if (!can_chunk) break; + } + + auto write_media_row_at = [&](size_t row, const uint8_t* feature_row, + size_t feature_row_bytes, Precision feature_precision) { + size_t node_id = static_cast(lm_encoder_media_chunk_->runtime_input_node_ids[media_embeds_idx]); + const auto& desc = lm_encoder_media_chunk_->graph->get_output_buffer(node_id); + auto* embeds_data = static_cast(lm_encoder_media_chunk_->graph->get_output(node_id)); + if (!embeds_data) { + throw std::runtime_error("lm_encoder_media_chunk inputs_embeds storage is unavailable"); + } + if (media_chunk_tokens == 0 || desc.total_size % media_chunk_tokens != 0) { + throw std::runtime_error("lm_encoder_media_chunk inputs_embeds is not token-aligned"); + } + const size_t row_elements = desc.total_size / media_chunk_tokens; + const size_t dst_offset = PrecisionTraits::byte_offset_of(desc.precision, row * row_elements); + const size_t dst_bytes = PrecisionTraits::packed_size_of(desc.precision, row_elements); + if (dst_offset + dst_bytes > desc.byte_size) { + throw std::runtime_error("lm_encoder_media_chunk media row exceeds inputs_embeds buffer"); + } + if (desc.precision == feature_precision) { + size_t to_copy = std::min(feature_row_bytes, dst_bytes); + std::memcpy(embeds_data + dst_offset, feature_row, to_copy); + if (to_copy < dst_bytes) { + std::memset(embeds_data + dst_offset + to_copy, 0, dst_bytes - to_copy); + } + return; + } + std::vector row_buf(dst_bytes, 0); + write_typed_buffer(row_buf, desc.precision, feature_row, feature_row_bytes, feature_precision); + std::memcpy(embeds_data + dst_offset, row_buf.data(), dst_bytes); + }; + + if (can_chunk) { + size_t audio_idx = 0; + size_t image_idx = 0; + for (size_t chunk_start = 0; chunk_start < token_count; chunk_start += media_chunk_tokens) { + const size_t real_tokens = std::min(media_chunk_tokens, token_count - chunk_start); + bool chunk_has_media = false; + for (size_t i = 0; i < real_tokens; ++i) { + uint32_t token = tokens[chunk_start + i]; + bool is_audio = audio_tok != 0 && token == audio_tok && have_audio_features; + bool is_image = image_tok != 0 && token == image_tok && have_image_features; + if (is_audio || is_image) { + chunk_has_media = true; + } + } + + for (auto& buf : lm_encoder_text_chunk_->input_buffers) { + std::fill(buf.begin(), buf.end(), 0); + } + for (auto& buf : lm_encoder_media_chunk_->input_buffers) { + std::fill(buf.begin(), buf.end(), 0); + } + + for (size_t i = 0; i < media_chunk_tokens; ++i) { + size_t pos = chunk_start + i; + uint32_t token = pos < token_count ? tokens[pos] : static_cast(config_.pad_token_id); + bool is_audio = pos < token_count && audio_tok != 0 && token == audio_tok && have_audio_features; + bool is_image = pos < token_count && image_tok != 0 && token == image_tok && have_image_features; + uint32_t encoder_token = (is_audio || is_image) ? media_placeholder_token_id_ : token; + write_int_input_at(*lm_encoder_text_chunk_, "input_ids", i, static_cast(encoder_token)); + write_int_input_at(*lm_encoder_text_chunk_, "position_ids", i, static_cast(pos)); + } + + lm_encoder_text_chunk_->graph->execute(); + + if (chunk_has_media) { + copy_component_outputs_to_inputs(*lm_encoder_text_chunk_, *lm_encoder_media_chunk_); + } + + for (size_t i = 0; i < media_chunk_tokens; ++i) { + size_t pos = chunk_start + i; + uint32_t token = pos < token_count ? tokens[pos] : static_cast(config_.pad_token_id); + uint32_t media_token = token; + + if (pos < token_count && audio_tok != 0 && token == audio_tok && have_audio_features) { + if (audio_idx >= audio_feature.rows) { + throw std::runtime_error("Prompt contains more audio tokens than audio feature rows"); + } + write_media_row_at(i, audio_feature.row(audio_idx), audio_feature.row_bytes, audio_feature.precision); + media_token = media_placeholder_token_id_; + ++audio_idx; + } else if (pos < token_count && image_tok != 0 && token == image_tok && have_image_features) { + if (image_idx >= image_feature.rows) { + throw std::runtime_error("Prompt contains more image tokens than image feature rows"); + } + write_media_row_at(i, image_feature.row(image_idx), image_feature.row_bytes, image_feature.precision); + media_token = media_placeholder_token_id_; + ++image_idx; + } + + write_int_input_at(*lm_encoder_media_chunk_, "input_ids", i, static_cast(media_token)); + write_int_input_at(*lm_encoder_media_chunk_, "position_ids", i, static_cast(pos)); + } + + Component* output_chunk = lm_encoder_text_chunk_; + if (chunk_has_media) { + lm_encoder_media_chunk_->graph->execute(); + output_chunk = lm_encoder_media_chunk_; + } + + for (const auto& info : outputs) { + Component* selected_chunk = output_chunk; + auto source_it = media_chunk_output_sources_.find(info.name); + const std::string source_policy = + source_it == media_chunk_output_sources_.end() ? "" : source_it->second; + if (chunk_has_media && source_it != media_chunk_output_sources_.end()) { + if (source_policy == "text_chunk") { + selected_chunk = lm_encoder_text_chunk_; + } else if (source_policy == "media_chunk") { + selected_chunk = lm_encoder_media_chunk_; + } + } + + auto copy_row_from = [&](Component* source_chunk, int output_idx, size_t row, size_t dst_pos) { + size_t node_id = static_cast(source_chunk->output_node_ids[output_idx]); + const auto& desc = source_chunk->graph->get_output_buffer(node_id); + const auto* ptr = static_cast(source_chunk->graph->get_output(node_id)); + size_t per_token_elements = desc.total_size / media_chunk_tokens; + size_t per_token_bytes = PrecisionTraits::packed_size_of(desc.precision, per_token_elements); + size_t src_offset = PrecisionTraits::byte_offset_of(desc.precision, row * per_token_elements); + std::memcpy( + store_bytes[info.name].data() + dst_pos * info.per_token_bytes, + ptr + src_offset, + per_token_bytes); + }; + + int output_idx = selected_chunk == lm_encoder_text_chunk_ ? info.text_chunk_idx : info.media_chunk_idx; + for (size_t i = 0; i < real_tokens; ++i) { + copy_row_from(selected_chunk, output_idx, i, chunk_start + i); + } + } + lm_encoder_text_chunk_->graph->release_runtime_buffers(); + lm_encoder_media_chunk_->graph->release_runtime_buffers(); + } + lm_encoder_text_chunk_->graph->release_all_weight_pages(); + lm_encoder_media_chunk_->graph->release_all_weight_pages(); + debug_lm_store(); + return true; + } + } size_t audio_idx = 0; size_t image_idx = 0; @@ -2631,22 +3873,22 @@ bool Model::build_lm_encoder_outputs_dynamic_gemma4( Precision media_prec = Precision::FP16; if (audio_tok != 0 && token == audio_tok && have_audio_features) { - if (audio_idx >= audio_rows) { - throw std::runtime_error("Gemma4 prompt contains more audio tokens than audio feature rows"); + if (audio_idx >= audio_feature.rows) { + throw std::runtime_error("Prompt contains more audio tokens than audio feature rows"); } component = lm_encoder_media_step_; - media_row = audio_it->second.data() + audio_idx * audio_row_bytes; - media_row_bytes = audio_row_bytes; - media_prec = audio_prec; + media_row = audio_feature.row(audio_idx); + media_row_bytes = audio_feature.row_bytes; + media_prec = audio_feature.precision; ++audio_idx; } else if (image_tok != 0 && token == image_tok && have_image_features) { - if (image_idx >= image_rows) { - throw std::runtime_error("Gemma4 prompt contains more image tokens than image feature rows"); + if (image_idx >= image_feature.rows) { + throw std::runtime_error("Prompt contains more image tokens than image feature rows"); } component = lm_encoder_media_step_; - media_row = image_it->second.data() + image_idx * image_row_bytes; - media_row_bytes = image_row_bytes; - media_prec = image_prec; + media_row = image_feature.row(image_idx); + media_row_bytes = image_feature.row_bytes; + media_prec = image_feature.precision; ++image_idx; } @@ -2659,7 +3901,7 @@ bool Model::build_lm_encoder_outputs_dynamic_gemma4( size_t node_id = static_cast(component->runtime_input_node_ids[embeds_idx]); const auto& desc = component->graph->get_output_buffer(node_id); write_typed_buffer(buf, desc.precision, media_row, media_row_bytes, media_prec); - write_int_input(*component, "input_ids", 0); + write_int_input(*component, "input_ids", static_cast(media_placeholder_token_id_)); write_int_input(*component, "position_ids", static_cast(pos)); } else { write_int_input(*component, "input_ids", static_cast(token)); @@ -2672,7 +3914,7 @@ bool Model::build_lm_encoder_outputs_dynamic_gemma4( size_t node_id = static_cast(component->output_node_ids[out_idx]); const auto& desc = component->graph->get_output_buffer(node_id); if (desc.byte_size != info.per_token_bytes || desc.precision != info.precision) { - throw std::runtime_error("Gemma4 dynamic output shape mismatch for " + info.name); + throw std::runtime_error("Dynamic media output shape mismatch for " + info.name); } const void* ptr = component->graph->get_output(node_id); std::memcpy(store_bytes[info.name].data() + pos * info.per_token_bytes, ptr, info.per_token_bytes); @@ -2681,18 +3923,27 @@ bool Model::build_lm_encoder_outputs_dynamic_gemma4( } encoder_->graph->release_all_weight_pages(); lm_encoder_media_step_->graph->release_all_weight_pages(); + debug_lm_store(); return true; } - bool Model::run_chunk_prefill_path(const std::vector& tokens, const std::vector& image_paths, const std::vector>& audio_features_per_message) { - if (cache_total_seq_len_ > 0) return false; + const size_t start_position = cache_total_seq_len_; const bool have_images = !image_paths.empty() && vision_encoder_ != nullptr; bool any_audio = false; for (const auto& mel : audio_features_per_message) { if (!mel.empty()) { any_audio = true; break; } } const bool have_audio = any_audio && audio_encoder_ != nullptr; + auto media_chunk_modality_allowed = [&](const std::string& modality) -> bool { + if (media_chunk_prefill_modalities_.empty()) return true; + return std::find(media_chunk_prefill_modalities_.begin(), media_chunk_prefill_modalities_.end(), modality) + != media_chunk_prefill_modalities_.end(); + }; + if ((have_images && !media_chunk_modality_allowed("vision") && !media_chunk_modality_allowed("image")) + || (have_audio && !media_chunk_modality_allowed("audio"))) { + return false; + } std::vector qwen_images; if (have_images) { @@ -2700,26 +3951,23 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, throw std::runtime_error("failed to load vision_encoder"); } for (const std::string& logical : vision_encoder_->logical_outputs) { - media_features_.erase(logical); - media_feature_shapes_.erase(logical); - media_feature_precisions_.erase(logical); + clear_media_feature(logical); } - const bool lfm2_vision = family_ == "lfm2_vl"; + const std::string strategy = image_preprocess_strategy(); + const bool lfm2_vision = strategy == "lfm2_vl"; if (lfm2_vision) { if (!vision_projector_) throw std::runtime_error("lfm2_vl requires a vision_projector component"); if (!load_lfm2_vl_position_grid()) { throw std::runtime_error("lfm2_vl vision position-embedding grid is missing from the bundle"); } if (!load_component_graph(*vision_projector_)) throw std::runtime_error("failed to load vision_projector"); - media_features_.erase("image_features"); - media_feature_shapes_.erase("image_features"); - media_feature_precisions_.erase("image_features"); + clear_media_feature("image_features"); } for (const auto& path : image_paths) { - if (family_ == "lfm2_vl") { + if (strategy == "lfm2_vl") { encode_lfm2_vl_image_into_features(path); continue; - } else if (family_ == "qwen3_5" || family_ == "qwen3_vl" || config_.model_type == Config::ModelType::QWEN) { + } else if (strategy == "qwen3_vl") { Qwen3VlImagePreprocessed prep = preprocess_qwen3_vl_image(path, config_); int pv_idx = input_index(*vision_encoder_, "pixel_values"); if (pv_idx < 0) { @@ -2734,7 +3982,19 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, Precision::FP32); qwen_images.push_back(std::move(prep)); } else { + if (!strategy.empty() && strategy != "gemma4") { + throw std::runtime_error("Unsupported image preprocess strategy: " + strategy); + } Gemma4ImagePreprocessed prep = preprocess_gemma4_image(path, config_); + debug_media_vector("vision_preprocess", "pixel_values", + prep.pixel_values, {prep.max_patches, prep.patch_dim}, + Precision::FP32); + debug_media_vector("vision_preprocess", "pixel_position_ids", + prep.pixel_position_ids, {prep.max_patches, 2}, + Precision::INT8); + if (has_npu_vision_encoder() && vision_encode_via_npu(prep.pixel_values, &prep.pixel_position_ids)) { + continue; + } int pv_idx = input_index(*vision_encoder_, "pixel_values"); if (pv_idx >= 0) { auto& pv_buf = vision_encoder_->input_buffers[pv_idx]; @@ -2774,28 +4034,30 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, const std::string& name = vision_encoder_->logical_outputs[i]; size_t node_id = static_cast(vision_encoder_->output_node_ids[i]); const auto& desc = vision_encoder_->graph->get_output_buffer(node_id); - void* ptr = vision_encoder_->graph->get_output(node_id); - auto& slot = media_features_[name]; - const size_t prev_bytes = slot.size(); - slot.resize(prev_bytes + desc.byte_size); - std::memcpy(slot.data() + prev_bytes, ptr, desc.byte_size); - auto shape_it = media_feature_shapes_.find(name); - if (shape_it == media_feature_shapes_.end() || shape_it->second.empty()) { - media_feature_shapes_[name] = desc.shape; - } else if (desc.shape.size() >= 2 && shape_it->second.size() == desc.shape.size()) { - shape_it->second[shape_it->second.size() - 2] += desc.shape[desc.shape.size() - 2]; + if (image_paths.size() == 1) { + publish_media_feature_storage( + name, vision_encoder_->graph->export_tensor_storage(node_id), + desc.shape, desc.precision); + } else { + const size_t rows = desc.shape.size() >= 2 + ? desc.shape[desc.shape.size() - 2] + : desc.shape.empty() ? 0 : desc.shape[0]; + append_media_feature(name, vision_encoder_->graph->get_output(node_id), + desc.byte_size, desc.shape, desc.precision, rows); + } + if (debug_media_features_enabled()) { + std::vector debug_bytes(desc.byte_size); + std::memcpy(debug_bytes.data(), media_feature_data(name), desc.byte_size); + debug_media_feature("vision_encoder", name, debug_bytes, desc.shape, desc.precision); } - media_feature_precisions_[name] = desc.precision; } vision_encoder_->graph->release_runtime_buffers(); - vision_encoder_->graph->release_all_weight_pages(); } if (lfm2_vision) { unload_component_graph(*vision_projector_); } unload_component_graph(*vision_encoder_); } - if (have_audio) { run_audio_encoder_messages(audio_features_per_message); } @@ -2803,86 +4065,147 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, std::map> store_bytes; std::map store_prec; std::map> store_shape; - const bool needs_dynamic_walk = family_ == "gemma4" && (have_images || have_audio) && lm_encoder_media_step_ != nullptr; + const bool needs_dynamic_walk = + media_injection_strategy_ == "token_row_replacement" + && (have_images || have_audio) + && lm_encoder_media_step_ != nullptr; if (needs_dynamic_walk) { - if (!build_lm_encoder_outputs_dynamic_gemma4(tokens, store_bytes, store_prec, store_shape)) { + if (!build_lm_encoder_outputs_dynamic_media(tokens, store_bytes, store_prec, store_shape)) { return false; } } else { if (!load_component_graph(*lm_encoder_)) { throw std::runtime_error("failed to load lm_encoder"); } - int ids_idx = input_index(*lm_encoder_, "input_ids"); - if (ids_idx >= 0) { - auto& ids_buf = lm_encoder_->input_buffers[ids_idx]; - size_t ids_node = static_cast(lm_encoder_->runtime_input_node_ids[ids_idx]); - const auto& ids_desc = lm_encoder_->graph->get_output_buffer(ids_node); - write_tokens_buffer(ids_buf, ids_desc.precision, tokens, 0); - } - int mask_idx = input_index(*lm_encoder_, "attention_mask"); - if (mask_idx >= 0) { - auto& mb = lm_encoder_->input_buffers[mask_idx]; - size_t mnode = static_cast(lm_encoder_->runtime_input_node_ids[mask_idx]); - const auto& mdesc = lm_encoder_->graph->get_output_buffer(mnode); - fill_int_buffer(mb, mdesc.precision, 1, tokens.size()); - } - - int pos_idx = input_index(*lm_encoder_, "position_ids"); - if (pos_idx >= 0) { - auto& pos_buf = lm_encoder_->input_buffers[pos_idx]; - size_t pos_node = static_cast(lm_encoder_->runtime_input_node_ids[pos_idx]); - const auto& pos_desc = lm_encoder_->graph->get_output_buffer(pos_node); - if (pos_desc.shape.size() >= 3 && pos_desc.shape[0] == 3 && !qwen_images.empty()) { - size_t capacity = pos_desc.shape[pos_desc.shape.size() - 1]; - auto positions = qwen3_vl_position_ids(tokens, capacity, qwen_images, config_.image_token_id); - write_int_vector_buffer(pos_buf, pos_desc.precision, positions); - } - } + const size_t lm_chunk_tokens = component_chunk_tokens(*lm_encoder_, "input_ids"); + const size_t encoder_tokens = lm_chunk_tokens > 0 ? lm_chunk_tokens : tokens.size(); + if (encoder_tokens == 0) return false; - for (const auto& kv : media_features_) { - const std::string& name = kv.first; - int idx = input_index(*lm_encoder_, name); - if (idx < 0) continue; - auto& dst_buf = lm_encoder_->input_buffers[idx]; - size_t node_id = static_cast(lm_encoder_->runtime_input_node_ids[idx]); - const auto& desc = lm_encoder_->graph->get_output_buffer(node_id); - Precision src_prec = media_feature_precisions_[name]; - write_typed_buffer(dst_buf, desc.precision, - kv.second.data(), kv.second.size(), src_prec); - } - lm_encoder_->graph->execute(); + struct EncoderOutputInfo { + std::string name; + int output_idx = -1; + size_t per_token_elements = 0; + size_t per_token_bytes = 0; + Precision precision = Precision::FP16; + std::vector shape; + }; + std::vector encoder_outputs; for (size_t i = 0; i < lm_encoder_->output_node_ids.size() && i < lm_encoder_->logical_outputs.size(); ++i) { - const std::string& name = lm_encoder_->logical_outputs[i]; - size_t node_id = static_cast(lm_encoder_->output_node_ids[i]); + const size_t node_id = static_cast(lm_encoder_->output_node_ids[i]); const auto& desc = lm_encoder_->graph->get_output_buffer(node_id); - void* ptr = lm_encoder_->graph->get_output(node_id); - auto& slot = store_bytes[name]; - slot.assign(desc.byte_size, 0); - std::memcpy(slot.data(), ptr, desc.byte_size); - store_prec[name] = desc.precision; - store_shape[name] = desc.shape; + if (desc.total_size == 0 || desc.total_size % encoder_tokens != 0) { + continue; + } + EncoderOutputInfo info; + info.name = lm_encoder_->logical_outputs[i]; + info.output_idx = static_cast(i); + info.per_token_elements = desc.total_size / encoder_tokens; + info.per_token_bytes = PrecisionTraits::packed_size_of(desc.precision, info.per_token_elements); + info.precision = desc.precision; + info.shape = desc.shape; + set_sequence_length_in_shape(info.shape, tokens.size()); + if (info.per_token_bytes == 0) { + continue; + } + store_bytes[info.name].assign(tokens.size() * info.per_token_bytes, 0); + store_prec[info.name] = info.precision; + store_shape[info.name] = info.shape; + encoder_outputs.push_back(std::move(info)); } + if (encoder_outputs.empty()) return false; + + for (size_t chunk_start = 0; chunk_start < tokens.size(); chunk_start += encoder_tokens) { + const size_t real_tokens = std::min(encoder_tokens, tokens.size() - chunk_start); + for (auto& input_buffer : lm_encoder_->input_buffers) { + std::fill(input_buffer.begin(), input_buffer.end(), 0); + } + + int ids_idx = input_index(*lm_encoder_, "input_ids"); + if (ids_idx >= 0) { + for (size_t i = 0; i < encoder_tokens; ++i) { + const size_t pos = chunk_start + i; + const uint32_t token = pos < tokens.size() + ? tokens[pos] + : static_cast(config_.pad_token_id); + write_int_input_at(*lm_encoder_, "input_ids", i, static_cast(token)); + } + } + + int mask_idx = input_index(*lm_encoder_, "attention_mask"); + if (mask_idx >= 0) { + for (size_t i = 0; i < real_tokens; ++i) { + write_int_input_at(*lm_encoder_, "attention_mask", i, 1); + } + } + + int pos_idx = input_index(*lm_encoder_, "position_ids"); + if (pos_idx >= 0) { + auto& pos_buf = lm_encoder_->input_buffers[pos_idx]; + size_t pos_node = static_cast(lm_encoder_->runtime_input_node_ids[pos_idx]); + const auto& pos_desc = lm_encoder_->graph->get_output_buffer(pos_node); + if (pos_desc.shape.size() >= 3 && pos_desc.shape[0] == 3 && !qwen_images.empty()) { + size_t capacity = pos_desc.shape[pos_desc.shape.size() - 1]; + auto positions = qwen3_vl_position_ids(tokens, capacity, qwen_images, config_.image_token_id); + write_int_vector_buffer(pos_buf, pos_desc.precision, positions); + } else { + for (size_t i = 0; i < encoder_tokens; ++i) { + write_int_input_at(*lm_encoder_, "position_ids", i, static_cast(start_position + chunk_start + i)); + } + } + } - if (have_images && family_ == "lfm2_vl") { - auto feat_it = media_features_.find("image_features"); + for (const auto& kv : media_features_) { + const std::string& name = kv.first; + int idx = input_index(*lm_encoder_, name); + if (idx < 0) continue; + const uint8_t* src_data = media_feature_data(name); + const size_t src_bytes = media_feature_byte_size(name); + if (!src_data || src_bytes == 0) continue; + auto& dst_buf = lm_encoder_->input_buffers[idx]; + size_t node_id = static_cast(lm_encoder_->runtime_input_node_ids[idx]); + const auto& desc = lm_encoder_->graph->get_output_buffer(node_id); + Precision src_prec = media_feature_precisions_[name]; + write_typed_buffer(dst_buf, desc.precision, + src_data, src_bytes, src_prec); + runtime_profile_.media_copied_bytes += std::min(dst_buf.size(), src_bytes); + runtime_profile_.copied_bytes += std::min(dst_buf.size(), src_bytes); + } + + lm_encoder_->graph->execute(); + + for (const auto& info : encoder_outputs) { + size_t node_id = static_cast(lm_encoder_->output_node_ids[info.output_idx]); + const auto& desc = lm_encoder_->graph->get_output_buffer(node_id); + const auto* ptr = static_cast(lm_encoder_->graph->get_output(node_id)); + for (size_t i = 0; i < real_tokens; ++i) { + const size_t src_offset = PrecisionTraits::byte_offset_of(desc.precision, i * info.per_token_elements); + const size_t dst_offset = (chunk_start + i) * info.per_token_bytes; + std::memcpy(store_bytes[info.name].data() + dst_offset, ptr + src_offset, info.per_token_bytes); + } + } + lm_encoder_->graph->release_runtime_buffers(); + } + + if (have_images && image_preprocess_strategy() == "lfm2_vl") { auto emb_it = store_bytes.find("inputs_embeds"); - if (feat_it != media_features_.end() && emb_it != store_bytes.end() && tokenizer_) { + if (media_feature_data("image_features") && emb_it != store_bytes.end() && tokenizer_) { const auto& emb_shape = store_shape["inputs_embeds"]; const auto& feat_shape = media_feature_shapes_["image_features"]; const size_t hidden = emb_shape.empty() ? 0 : emb_shape.back(); const size_t emb_seq = emb_shape.size() >= 2 ? emb_shape[emb_shape.size() - 2] : 0; - const size_t feat_rows = feat_shape.empty() ? 0 : feat_shape[0]; + const size_t feat_rows = feat_shape.size() >= 2 ? feat_shape[feat_shape.size() - 2] + : feat_shape.empty() ? 0 : feat_shape[0]; const Precision emb_prec = store_prec["inputs_embeds"]; const Precision feat_prec = media_feature_precisions_["image_features"]; const size_t emb_elem = PrecisionTraits::size_of(emb_prec); const size_t feat_elem = PrecisionTraits::size_of(feat_prec); const uint32_t image_tok = tokenizer_->get_image_token_id(); uint8_t* emb_data = emb_it->second.data(); - const uint8_t* feat_data = feat_it->second.data(); + const uint8_t* feat_data = media_feature_data("image_features"); std::vector frow(hidden); size_t fcur = 0; const size_t limit = std::min(emb_seq, tokens.size()); @@ -2906,7 +4229,12 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, lm_encoder_->graph->release_all_weight_pages(); unload_component_graph(*lm_encoder_); } - + { + std::vector consumed_media; + consumed_media.reserve(media_features_.size()); + for (const auto& [name, _] : media_features_) consumed_media.push_back(name); + for (const auto& name : consumed_media) clear_media_feature(name); + } auto embeds_shape_it = store_shape.find("inputs_embeds"); if (embeds_shape_it == store_shape.end()) { return false; @@ -2914,8 +4242,7 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, size_t full_seq = 0; { const auto& sh = embeds_shape_it->second; - if (sh.size() >= 3) full_seq = sh[sh.size() - 2]; - else if (!sh.empty()) full_seq = sh[0]; + full_seq = sequence_length_from_shape(sh); } if (full_seq == 0) return false; @@ -2967,7 +4294,13 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, if (count > 0) valid_seq = count; } valid_seq = std::min(valid_seq, full_seq); - const size_t whole_chunks_end = (valid_seq / chunk_seq) * chunk_seq; + // Media prefill graphs initialize convolution state, so they are only + // valid at the beginning of a sequence. Prefix extensions keep the + // decoder-owned cache and execute the delta through decoder_step below. + size_t whole_chunks_end = start_position == 0 ? (valid_seq / chunk_seq) * chunk_seq : 0; + if (whole_chunks_end > 0 && decoder_ != nullptr && !cache_states_compatible(*decoder_prefill_chunk_, *decoder_)) { + whole_chunks_end = 0; + } for (size_t chunk_start = 0; chunk_start < whole_chunks_end; chunk_start += chunk_seq) { for (const auto& kv : store_bytes) { const std::string& name = kv.first; @@ -2981,6 +4314,16 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, const uint8_t* src_ptr = kv.second.data() + chunk_start * src_per_pos; size_t src_slice_bytes = chunk_seq * src_per_pos; write_typed_buffer(dst_buf, desc.precision, src_ptr, src_slice_bytes, src_prec); + // Runtime-state aliases may have rebound this input to preserved + // producer storage. This path is explicitly materializing a new + // media chunk, so make the freshly written mutable buffer + // authoritative for the next execution. + decoder_prefill_chunk_->graph->set_external_input( + node_id, dst_buf.data(), desc.precision); + } + const size_t real_chunk_tokens = std::min(chunk_seq, valid_seq - chunk_start); + for (size_t i = 0; i < real_chunk_tokens; ++i) { + write_int_input_at(*decoder_prefill_chunk_, "attention_mask", i, 1); } decoder_prefill_chunk_->graph->execute(); } @@ -2989,6 +4332,9 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, decoder_prefill_chunk_->graph->release_runtime_buffers(); unload_component_graph(*decoder_prefill_chunk_); } + if (whole_chunks_end < valid_seq && (decoder_ == nullptr || !load_component_graph(*decoder_))) { + return false; + } for (size_t pos = whole_chunks_end; pos < valid_seq; ++pos) { for (const auto& kv : store_bytes) { const std::string& name = kv.first; @@ -3002,6 +4348,9 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, const uint8_t* src_ptr = kv.second.data() + pos * src_per_pos; write_typed_buffer(dst_buf, desc.precision, src_ptr, src_per_pos, src_prec); } + if (input_index(*decoder_, "position_ids") >= 0) { + write_int_input(*decoder_, "position_ids", static_cast(start_position + pos)); + } decoder_->graph->execute(); } cache_total_seq_len_ += valid_seq; @@ -3027,113 +4376,150 @@ void Model::prefill_with_media(const std::vector& tokens, prefill(tokens, get_prefill_chunk_size(), profile_file); return; } - + clear_runtime_states_with_lifetime("request"); + + const std::string fallback = media_prefill_fallback_.empty() ? "error" : media_prefill_fallback_; + const bool prefer_warm_media_step = + fallback == "warm_media_step" || fallback == "warm_media" || fallback == "media_step"; + const bool can_chunk_lm_encoder = + lm_encoder_ != nullptr + || (media_injection_strategy_ == "token_row_replacement" + && lm_encoder_text_chunk_ != nullptr + && lm_encoder_media_chunk_ != nullptr); + auto media_chunk_modality_allowed = [&](const std::string& modality) -> bool { + if (media_chunk_prefill_modalities_.empty()) return true; + return std::find(media_chunk_prefill_modalities_.begin(), media_chunk_prefill_modalities_.end(), modality) + != media_chunk_prefill_modalities_.end(); + }; + const bool media_chunk_modalities_allowed = + (!have_images || media_chunk_modality_allowed("vision") || media_chunk_modality_allowed("image")) && + (!have_audio || media_chunk_modality_allowed("audio")); const bool can_chunk_prefill = - lm_encoder_ != nullptr && decoder_prefill_chunk_ != nullptr && + std::getenv("CACTUS_DISABLE_MEDIA_CHUNK_PREFILL") == nullptr && + !prefer_warm_media_step && + media_chunk_modalities_allowed && + can_chunk_lm_encoder && decoder_prefill_chunk_ != nullptr && (vision_encoder_ != nullptr || audio_encoder_ != nullptr); if (can_chunk_prefill) { if (run_chunk_prefill_path(tokens, image_paths, audio_features_per_message)) { + media_min_new_tokens_remaining_ = media_min_new_tokens_; + if (debug_media_features_enabled()) { + std::cerr << "[media-debug] media_min_new_tokens_remaining=" + << media_min_new_tokens_remaining_ << std::endl; + } (void)profile_file; return; } } - if (!supports_warm_media_injection()) { - CACTUS_LOG_WARN("model", "Bundle supports neither chunk-prefill nor warm media injection; falling back to text-only prefill"); - prefill(tokens, get_prefill_chunk_size(), profile_file); - return; - } - if (have_images) { - for (const auto& path : image_paths) { - run_vision_encoder(path); + if (prefer_warm_media_step) { + if (!supports_warm_media_injection()) { + throw std::runtime_error("Warm media-step prefill requested, but bundle has no compatible media step"); } - } - if (have_audio) { - run_audio_encoder_messages(audio_features_per_message); - } - std::string image_feature_name; - Precision image_feature_prec = Precision::FP16; - size_t image_row_bytes = 0; - if (have_images) { - const std::vector candidates = {"image_features", "image_embeddings", "vision_features", "inputs_embeds"}; - for (const auto& name : candidates) { - if (media_features_.count(name)) { image_feature_name = name; break; } + if (have_images) { + for (const auto& path : image_paths) { + run_vision_encoder(path); + } } - if (image_feature_name.empty() && !media_features_.empty()) { - image_feature_name = media_features_.begin()->first; + if (have_audio) { + run_audio_encoder_messages(audio_features_per_message); } - if (!image_feature_name.empty()) { - const auto& shape = media_feature_shapes_[image_feature_name]; - image_feature_prec = media_feature_precisions_[image_feature_name]; - if (shape.size() >= 2) { - size_t rows = (shape.size() >= 3) ? shape[shape.size() - 2] : shape[0]; - size_t total = media_features_[image_feature_name].size(); - image_row_bytes = rows > 0 ? total / rows : total; - } else { - image_row_bytes = media_features_[image_feature_name].size(); + + std::string image_feature_name; + Precision image_feature_prec = Precision::FP16; + size_t image_row_bytes = 0; + if (have_images) { + const std::vector candidates = {"image_features", "image_embeddings", "vision_features", "inputs_embeds"}; + for (const auto& name : candidates) { + if (media_features_.count(name)) { image_feature_name = name; break; } + } + if (image_feature_name.empty() && !media_features_.empty()) { + image_feature_name = media_features_.begin()->first; + } + if (!image_feature_name.empty()) { + const auto& shape = media_feature_shapes_[image_feature_name]; + image_feature_prec = media_feature_precisions_[image_feature_name]; + if (shape.size() >= 2) { + size_t rows = (shape.size() >= 3) ? shape[shape.size() - 2] : shape[0]; + size_t total = media_feature_byte_size(image_feature_name); + image_row_bytes = rows > 0 ? total / rows : total; + } else { + image_row_bytes = media_feature_byte_size(image_feature_name); + } } } - } - std::string audio_feature_name; - Precision audio_feature_prec = Precision::FP16; - size_t audio_row_bytes = 0; - if (have_audio) { - const std::vector candidates = {"audio_features", "audio_embeddings", "encoder_hidden_states", "inputs_embeds"}; - for (const auto& name : candidates) { - if (media_features_.count(name) && name != image_feature_name) { audio_feature_name = name; break; } - } - if (audio_feature_name.empty()) { - for (const auto& kv : media_features_) { - if (kv.first != image_feature_name) { audio_feature_name = kv.first; break; } + std::string audio_feature_name; + Precision audio_feature_prec = Precision::FP16; + size_t audio_row_bytes = 0; + if (have_audio) { + const std::vector candidates = {"audio_features", "audio_embeddings", "encoder_hidden_states", "inputs_embeds"}; + for (const auto& name : candidates) { + if (media_features_.count(name) && name != image_feature_name) { audio_feature_name = name; break; } } - } - if (!audio_feature_name.empty()) { - const auto& shape = media_feature_shapes_[audio_feature_name]; - audio_feature_prec = media_feature_precisions_[audio_feature_name]; - if (shape.size() >= 2) { - size_t rows = (shape.size() >= 3) ? shape[shape.size() - 2] : shape[0]; - size_t total = media_features_[audio_feature_name].size(); - audio_row_bytes = rows > 0 ? total / rows : total; - } else { - audio_row_bytes = media_features_[audio_feature_name].size(); + if (audio_feature_name.empty()) { + for (const auto& kv : media_features_) { + if (kv.first != image_feature_name) { audio_feature_name = kv.first; break; } + } + } + if (!audio_feature_name.empty()) { + const auto& shape = media_feature_shapes_[audio_feature_name]; + audio_feature_prec = media_feature_precisions_[audio_feature_name]; + if (shape.size() >= 2) { + size_t rows = (shape.size() >= 3) ? shape[shape.size() - 2] : shape[0]; + size_t total = media_feature_byte_size(audio_feature_name); + audio_row_bytes = rows > 0 ? total / rows : total; + } else { + audio_row_bytes = media_feature_byte_size(audio_feature_name); + } } } - } - size_t image_consumed = 0; - size_t audio_consumed = 0; - const uint32_t image_tok = config_.image_token_id; - const uint32_t audio_tok = config_.audio_token_id; - const bool warm_media = supports_warm_media_injection(); - - for (size_t i = 0; i < tokens.size(); ++i) { - uint32_t t = tokens[i]; - size_t pos = cache_total_seq_len_ + i; - if (image_tok != 0 && t == image_tok && !image_feature_name.empty() && warm_media) { - const auto& feat = media_features_[image_feature_name]; - const uint8_t* row = feat.data() + image_consumed * image_row_bytes; - if (image_consumed * image_row_bytes + image_row_bytes <= feat.size()) { - run_media_step(pos, row, image_row_bytes, image_feature_prec); - ++image_consumed; - continue; + size_t image_consumed = 0; + size_t audio_consumed = 0; + const uint32_t image_tok = config_.image_token_id; + const uint32_t audio_tok = config_.audio_token_id; + + for (size_t i = 0; i < tokens.size(); ++i) { + uint32_t t = tokens[i]; + size_t pos = cache_total_seq_len_ + i; + if (image_tok != 0 && t == image_tok && !image_feature_name.empty()) { + const uint8_t* feature_data = media_feature_data(image_feature_name); + const uint8_t* row = feature_data + image_consumed * image_row_bytes; + if (image_consumed * image_row_bytes + image_row_bytes <= media_feature_byte_size(image_feature_name)) { + run_media_step(pos, row, image_row_bytes, image_feature_prec); + ++image_consumed; + continue; + } } - } - if (audio_tok != 0 && t == audio_tok && !audio_feature_name.empty() && warm_media) { - const auto& feat = media_features_[audio_feature_name]; - const uint8_t* row = feat.data() + audio_consumed * audio_row_bytes; - if (audio_consumed * audio_row_bytes + audio_row_bytes <= feat.size()) { - run_media_step(pos, row, audio_row_bytes, audio_feature_prec); - ++audio_consumed; - continue; + if (audio_tok != 0 && t == audio_tok && !audio_feature_name.empty()) { + const uint8_t* feature_data = media_feature_data(audio_feature_name); + const uint8_t* row = feature_data + audio_consumed * audio_row_bytes; + if (audio_consumed * audio_row_bytes + audio_row_bytes <= media_feature_byte_size(audio_feature_name)) { + run_media_step(pos, row, audio_row_bytes, audio_feature_prec); + ++audio_consumed; + continue; + } } + run_step(t, pos, false); } - run_step(t, pos, false); + cache_total_seq_len_ += tokens.size(); + cache_token_ids_.insert(cache_token_ids_.end(), tokens.begin(), tokens.end()); + media_min_new_tokens_remaining_ = media_min_new_tokens_; + if (debug_media_features_enabled()) { + std::cerr << "[media-debug] media_min_new_tokens_remaining=" + << media_min_new_tokens_remaining_ << std::endl; + } + return; } - cache_total_seq_len_ += tokens.size(); - cache_token_ids_.insert(cache_token_ids_.end(), tokens.begin(), tokens.end()); - (void)profile_file; + + if (fallback == "text_only") { + CACTUS_LOG_WARN("model", "Chunked media prefill unavailable; falling back to text-only prefill by profile policy"); + prefill(tokens, get_prefill_chunk_size(), profile_file); + return; + } + throw std::runtime_error("Chunked media prefill failed or is unavailable for this profile"); } uint32_t Model::decode(const std::vector& tokens, float temperature, float top_p, @@ -3171,7 +4557,8 @@ uint32_t Model::decode(const std::vector& tokens, float temperature, f } cache_total_seq_len_ += tokens.size(); if (out_entropy) *out_entropy = 0.0f; - uint32_t result = argmax_last_logits(); + uint32_t result = argmax_last_logits(nullptr, repetition_penalty); + if (media_min_new_tokens_remaining_ > 0) --media_min_new_tokens_remaining_; record_sampled_token(result); return result; } @@ -3181,8 +4568,9 @@ uint32_t Model::decode(const std::vector& tokens, float temperature, f cache_total_seq_len_ = context_tokens_.size(); cache_token_ids_ = context_tokens_; uint32_t result = (greedy && !samp_ctx_active_) - ? argmax_last_logits(out_entropy) + ? argmax_last_logits(out_entropy, repetition_penalty) : sample_component_logits(*decoder_, temp, tp, tk, min_p, greedy, out_entropy); + if (media_min_new_tokens_remaining_ > 0) --media_min_new_tokens_remaining_; record_sampled_token(result); return result; } @@ -3194,8 +4582,9 @@ uint32_t Model::decode(const std::vector& tokens, float temperature, f cache_token_ids_.insert(cache_token_ids_.end(), tokens.begin(), tokens.end()); maybe_roll_compact(); uint32_t result = (greedy && !samp_ctx_active_) - ? argmax_last_logits(out_entropy) + ? argmax_last_logits(out_entropy, repetition_penalty) : sample_component_logits(*decoder_, temp, tp, tk, min_p, greedy, out_entropy); + if (media_min_new_tokens_remaining_ > 0) --media_min_new_tokens_remaining_; record_sampled_token(result); return result; } @@ -3245,6 +4634,7 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a CACTUS_LOG_ERROR("model", "Parakeet TDT bundle missing audio_encoder or decoder component"); return emitted; } + if (!load_component_graph(*audio_enc) || !load_component_graph(*dec)) return emitted; if (!bind_runtime_buffers(*audio_enc)) return emitted; if (!bind_runtime_buffers(*dec)) return emitted; @@ -3507,16 +4897,17 @@ uint32_t Model::decode_with_images(const std::vector& tokens, const st namespace { std::vector pool_and_normalize_media_feature( - const std::vector& bytes, + const uint8_t* bytes, + size_t byte_size, const std::vector& shape, Precision precision, const std::string& source ) { const size_t elem_size = PrecisionTraits::size_of(precision); - if (elem_size == 0 || bytes.empty() || shape.empty()) { + if (elem_size == 0 || bytes == nullptr || byte_size == 0 || shape.empty()) { throw std::runtime_error(source + " produced empty feature output"); } - const size_t total_elems = bytes.size() / elem_size; + const size_t total_elems = byte_size / elem_size; const size_t hidden_dim = shape.back(); if (hidden_dim == 0 || total_elems == 0 || total_elems % hidden_dim != 0) { throw std::runtime_error(source + " feature shape inconsistent with hidden_dim"); @@ -3525,13 +4916,13 @@ std::vector pool_and_normalize_media_feature( std::vector fp32(total_elems); switch (precision) { case Precision::FP32: - std::memcpy(fp32.data(), bytes.data(), total_elems * sizeof(float)); + std::memcpy(fp32.data(), bytes, total_elems * sizeof(float)); break; case Precision::FP16: - Quantization::fp16_to_fp32(reinterpret_cast(bytes.data()), fp32.data(), total_elems); + Quantization::fp16_to_fp32(reinterpret_cast(bytes), fp32.data(), total_elems); break; case Precision::INT8: - Quantization::int8_to_fp32(reinterpret_cast(bytes.data()), fp32.data(), total_elems, 1.0f); + Quantization::int8_to_fp32(reinterpret_cast(bytes), fp32.data(), total_elems, 1.0f); break; default: throw std::runtime_error(source + " feature precision not supported for embeddings"); @@ -3573,21 +4964,19 @@ std::vector Model::get_image_embeddings(const std::string& image_path) { if (media_features_.count(name)) { output_name = name; break; } } } - auto bytes_it = media_features_.find(output_name); auto shape_it = media_feature_shapes_.find(output_name); auto prec_it = media_feature_precisions_.find(output_name); - if (bytes_it == media_features_.end() || shape_it == media_feature_shapes_.end() + if (!media_feature_data(output_name) || shape_it == media_feature_shapes_.end() || prec_it == media_feature_precisions_.end()) { throw std::runtime_error("vision_encoder produced no output for '" + output_name + "'"); } std::vector embedding = pool_and_normalize_media_feature( - bytes_it->second, shape_it->second, prec_it->second, "vision_encoder"); + media_feature_data(output_name), media_feature_byte_size(output_name), + shape_it->second, prec_it->second, "vision_encoder"); for (const std::string& name : vision_encoder_->logical_outputs) { - media_features_.erase(name); - media_feature_shapes_.erase(name); - media_feature_precisions_.erase(name); + clear_media_feature(name); } media_features_.erase(output_name); media_feature_shapes_.erase(output_name); @@ -3612,21 +5001,19 @@ std::vector Model::get_audio_embeddings(const std::vector& mel_bin run_audio_encoder_messages({mel_bins}); - auto bytes_it = media_features_.find(output_name); auto shape_it = media_feature_shapes_.find(output_name); auto prec_it = media_feature_precisions_.find(output_name); - if (bytes_it == media_features_.end() || shape_it == media_feature_shapes_.end() + if (!media_feature_data(output_name) || shape_it == media_feature_shapes_.end() || prec_it == media_feature_precisions_.end()) { throw std::runtime_error("audio_encoder produced no output for '" + output_name + "'"); } std::vector embedding = pool_and_normalize_media_feature( - bytes_it->second, shape_it->second, prec_it->second, "audio_encoder"); + media_feature_data(output_name), media_feature_byte_size(output_name), + shape_it->second, prec_it->second, "audio_encoder"); for (const std::string& name : audio_encoder_->logical_outputs) { - media_features_.erase(name); - media_feature_shapes_.erase(name); - media_feature_precisions_.erase(name); + clear_media_feature(name); } load_component_graph(*audio_encoder_); return embedding; @@ -3635,19 +5022,32 @@ std::vector Model::get_audio_embeddings(const std::vector& mel_bin void Model::reset_cache() { cache_total_seq_len_ = 0; last_logit_position_ = 0; + media_min_new_tokens_remaining_ = 0; encoder_cross_kv_ready_ = false; context_tokens_.clear(); cache_token_ids_.clear(); special_rows_.clear(); token_history_.clear(); + token_history_set_.clear(); media_features_.clear(); + media_feature_storages_.clear(); media_feature_shapes_.clear(); media_feature_precisions_.clear(); + runtime_state_store_.clear(); for (auto& kv : components_) { Component& comp = kv.second; if (!comp.graph) continue; reset_component_cache_states(comp); } + state_arena_.clear(); +} + +void Model::seed_token_history_from_context(const std::vector& tokens) { + token_history_.clear(); + token_history_set_.clear(); + const size_t start = tokens.size() > MAX_TOKEN_HISTORY ? tokens.size() - MAX_TOKEN_HISTORY : 0; + token_history_.insert(token_history_.end(), tokens.begin() + start, tokens.end()); + token_history_set_.insert(tokens.begin(), tokens.end()); } void Model::set_cache_window(size_t /*window_size*/, size_t /*sink_size*/) {} @@ -4127,6 +5527,10 @@ bool Config::from_json(const std::string& config_path) { else if (key == "tile_size") tile_size = static_cast(std::stoul(value)); else if (key == "max_pixels_tolerance") max_pixels_tolerance = std::stof(value); else if (key == "do_image_splitting") do_image_splitting = (value == "true" || value == "1"); + else if (key == "do_resize") do_resize = (value == "true" || value == "1"); + else if (key == "do_rescale") do_rescale = (value == "true" || value == "1"); + else if (key == "do_normalize") do_normalize = (value == "true" || value == "1"); + else if (key == "do_convert_rgb") do_convert_rgb = (value == "true" || value == "1"); else if (key == "precision") { if (value == "INT8") precision = Precision::INT8; else if (value == "FP16") precision = Precision::FP16; @@ -4145,6 +5549,7 @@ bool Config::from_json(const std::string& config_path) { else if (mt == "youtu") model_type = ModelType::YOUTU; else if (mt == "needle") model_type = ModelType::NEEDLE; else if (mt == "bert" || mt == "nomic") model_type = ModelType::NOMIC; + else if (mt == "generic") model_type = ModelType::GENERIC; else model_type = ModelType::GEMMA4; } else if (key == "model_variant") { @@ -4376,19 +5781,13 @@ std::unique_ptr create_model(const std::string& bundle_dir) { fs::path manifest = fs::path(bundle_dir) / "components" / "manifest.json"; if (!fs::exists(manifest)) { CACTUS_LOG_ERROR("model", - "Not a transpiled bundle (no components/manifest.json at " << bundle_dir << "). " - "Run `cactus convert ` to produce one."); + "Not a runnable Cactus bundle (no components/manifest.json at " << bundle_dir << "). " + "Run `cactus download ` to fetch one."); return nullptr; } return std::make_unique(); } -const std::vector& Model::get_debug_nodes() const { - debug_nodes_.clear(); - return debug_nodes_; -} - - double Model::score_tokens_window_logprob(const std::vector& /*tokens*/, size_t /*start*/, size_t /*end*/, size_t /*context*/, size_t* tokens_scored) { if (tokens_scored) *tokens_scored = 0; diff --git a/cactus-engine/src/model_npu.cpp b/cactus-engine/src/model_npu.cpp new file mode 100644 index 000000000..31944f5f6 --- /dev/null +++ b/cactus-engine/src/model_npu.cpp @@ -0,0 +1,359 @@ +#include "engine.h" +#include "cactus_graph.h" +#include "cactus_kernels.h" + +#include +#include +#include +#include + +namespace cactus { +namespace engine { + +bool Model::load_npu_audio_encoder(const std::string& model_path, const std::string& compute_units) { + auto encoder = npu::create_encoder(); + if (!encoder) return false; + if (!encoder->load(model_path, compute_units)) return false; + if (!encoder->is_available()) return false; + npu_audio_encoder_ = std::move(encoder); + CACTUS_LOG_INFO("model", "NPU audio encoder loaded from: " << model_path); + return true; +} + +bool Model::load_npu_vision_encoder(const std::string& model_path) { + auto encoder = npu::create_encoder(); + if (!encoder) return false; + if (!encoder->load(model_path)) return false; + if (!encoder->is_available()) return false; + npu_vision_encoder_ = std::move(encoder); + CACTUS_LOG_INFO("model", "NPU vision encoder loaded from: " << model_path); + return true; +} + +bool Model::load_npu_source_encoder(const std::string& model_path) { + auto encoder = npu::create_encoder(); + if (!encoder) return false; + if (!encoder->load(model_path)) return false; + if (!encoder->is_available()) return false; + npu_source_encoder_ = std::move(encoder); + CACTUS_LOG_INFO("model", "NPU source encoder loaded from: " << model_path); + return true; +} + +bool Model::source_encode_via_npu(const std::vector& tokens) { + if (!npu_source_encoder_ || !npu_source_encoder_->is_available() || + !source_encoder_ || !decoder_cross_kv_) { + return false; + } + + int ids_idx = input_index(*source_encoder_, "input_ids"); + if (ids_idx < 0) return false; + size_t ids_node = static_cast(source_encoder_->runtime_input_node_ids[ids_idx]); + const auto& ids_desc = source_encoder_->graph->get_output_buffer(ids_node); + if (tokens.size() > ids_desc.total_size) return false; + + std::vector input_shape; + input_shape.reserve(ids_desc.shape.size()); + for (size_t dim : ids_desc.shape) input_shape.push_back(static_cast(dim)); + + std::vector input_ids(ids_desc.total_size, 0); + std::vector attention_mask(ids_desc.total_size, 0); + for (size_t i = 0; i < tokens.size(); ++i) { + input_ids[i] = static_cast(tokens[i]); + attention_mask[i] = 1; + } + + std::vector npu_inputs; + npu_inputs.push_back({ + "input_ids", + input_ids.data(), + npu::NPUNamedInput::DataType::INT32, + input_shape, + }); + if (input_index(*source_encoder_, "attention_mask") >= 0) { + npu_inputs.push_back({ + "attention_mask", + attention_mask.data(), + npu::NPUNamedInput::DataType::INT32, + input_shape, + }); + } + + int hidden_idx = input_index(*decoder_cross_kv_, "encoder_hidden_states"); + if (hidden_idx < 0) return false; + size_t hidden_node = static_cast(decoder_cross_kv_->runtime_input_node_ids[hidden_idx]); + const auto& hidden_desc = decoder_cross_kv_->graph->get_output_buffer(hidden_node); + std::vector<__fp16> hidden(hidden_desc.total_size, __fp16(0)); + size_t written = npu_source_encoder_->encode_multimodal_input( + npu_inputs, + hidden.data(), + "encoder_hidden_states"); + if (written == 0) return false; + + auto copy_fp16_to_component_input = [&](Component& comp, int input_idx, const __fp16* src, size_t src_elems) { + if (input_idx < 0 || !src) return false; + size_t node = static_cast(comp.runtime_input_node_ids[static_cast(input_idx)]); + const auto& desc = comp.graph->get_output_buffer(node); + auto& dst = comp.input_buffers[static_cast(input_idx)]; + std::fill(dst.begin(), dst.end(), 0); + const size_t elems = std::min(src_elems, desc.total_size); + if (desc.precision == Precision::FP16) { + std::memcpy(dst.data(), src, elems * sizeof(__fp16)); + return true; + } + if (desc.precision == Precision::FP32) { + float* out = reinterpret_cast(dst.data()); + for (size_t i = 0; i < elems; ++i) out[i] = static_cast(src[i]); + return true; + } + CACTUS_LOG_WARN("model", "NPU source encoder output precision mismatch for encoder_hidden_states"); + return false; + }; + + if (!copy_fp16_to_component_input(*decoder_cross_kv_, hidden_idx, hidden.data(), written)) { + return false; + } + + auto fill_mask = [&](Component& comp) { + int idx = input_index(comp, "encoder_attention_mask"); + if (idx < 0) return; + std::fill(comp.input_buffers[static_cast(idx)].begin(), + comp.input_buffers[static_cast(idx)].end(), + 0); + for (size_t i = 0; i < tokens.size(); ++i) { + write_int_input_at(comp, "encoder_attention_mask", i, 1); + } + }; + fill_mask(*decoder_cross_kv_); + if (decoder_) fill_mask(*decoder_); + + return true; +} + +bool Model::audio_encode_via_npu(const std::vector& audio_features) { + if (!npu_audio_encoder_ || !npu_audio_encoder_->is_available() || !audio_encoder_) { + return false; + } + const std::vector input_shape = npu_audio_encoder_->get_input_shape(); + if (input_shape.empty()) return false; + + size_t expected_elems = 1; + for (int d : input_shape) { + if (d <= 0) return false; + expected_elems *= static_cast(d); + } + if (audio_features.size() > expected_elems) return false; + + std::vector<__fp16> input_fp16(expected_elems, __fp16(0)); + for (size_t i = 0; i < audio_features.size(); ++i) { + input_fp16[i] = static_cast<__fp16>(audio_features[i]); + } + + const std::vector output_shape = npu_audio_encoder_->get_output_shape(); + size_t output_elems = 1; + for (int d : output_shape) { + if (d <= 0) { output_elems = 0; break; } + output_elems *= static_cast(d); + } + if (output_elems == 0) { + output_elems = npu_audio_encoder_->get_output_buffer_size(); + } + std::vector<__fp16> output_fp16(output_elems, __fp16(0)); + + size_t written = npu_audio_encoder_->encode( + input_fp16.data(), output_fp16.data(), input_shape, "x", "encoded"); + if (written == 0) return false; + + for (size_t i = 0; i < audio_encoder_->output_node_ids.size() + && i < audio_encoder_->logical_outputs.size(); ++i) { + const std::string& name = audio_encoder_->logical_outputs[i]; + size_t node_id = static_cast(audio_encoder_->output_node_ids[i]); + const auto& desc = audio_encoder_->graph->get_output_buffer(node_id); + const size_t elem_size = (desc.precision == Precision::FP32) ? sizeof(float) : sizeof(__fp16); + const size_t copy_elems = std::min(static_cast(written), desc.byte_size / elem_size); + auto& slot = media_features_[name]; + if (desc.precision == Precision::FP32) { + const size_t prev = slot.size(); + slot.resize(prev + copy_elems * sizeof(float)); + float* dst = reinterpret_cast(slot.data() + prev); + for (size_t k = 0; k < copy_elems; ++k) dst[k] = static_cast(output_fp16[k]); + } else { + const size_t prev = slot.size(); + slot.resize(prev + copy_elems * sizeof(__fp16)); + std::memcpy(slot.data() + prev, output_fp16.data(), copy_elems * sizeof(__fp16)); + } + auto shape_it = media_feature_shapes_.find(name); + if (shape_it == media_feature_shapes_.end() || shape_it->second.empty()) { + std::vector shape; + for (int d : output_shape) shape.push_back(static_cast(d)); + media_feature_shapes_[name] = std::move(shape); + } + media_feature_precisions_[name] = desc.precision; + break; + } + return true; +} + +bool Model::vision_encode_via_npu(const std::vector& pixel_values, + const std::vector* pixel_position_ids) { + if (!npu_vision_encoder_ || !npu_vision_encoder_->is_available() || !vision_encoder_) { + return false; + } + const std::vector input_shape = npu_vision_encoder_->get_input_shape_for("x"); + if (input_shape.empty()) return false; + + size_t expected_elems = 1; + for (int d : input_shape) { + if (d <= 0) return false; + expected_elems *= static_cast(d); + } + if (pixel_values.size() > expected_elems) return false; + + const bool package_takes_positions = npu_vision_encoder_->has_input("pixel_position_ids"); + if (package_takes_positions != (pixel_position_ids != nullptr)) { + CACTUS_LOG_WARN("model", "NPU vision encoder and pixel_position_ids mismatch; " + "falling back to CPU vision encoder (rebuild the bundle with NPU vision support to fix)"); + return false; + } + + std::vector<__fp16> input_fp16(expected_elems, __fp16(0)); + for (size_t i = 0; i < pixel_values.size(); ++i) { + input_fp16[i] = static_cast<__fp16>(pixel_values[i]); + } + + const std::vector output_shape = npu_vision_encoder_->get_output_shape(); + size_t output_elems = 1; + for (int d : output_shape) { + if (d <= 0) { output_elems = 0; break; } + output_elems *= static_cast(d); + } + if (output_elems == 0) { + output_elems = npu_vision_encoder_->get_output_buffer_size(); + } + std::vector<__fp16> output_fp16(output_elems, __fp16(0)); + + size_t written = 0; + if (package_takes_positions) { + const std::vector pos_shape = + npu_vision_encoder_->get_input_shape_for("pixel_position_ids"); + if (pos_shape.empty()) return false; + size_t pos_elems = 1; + for (int d : pos_shape) { + if (d <= 0) return false; + pos_elems *= static_cast(d); + } + if (pixel_position_ids->size() > pos_elems) return false; + + std::vector positions_i32(pos_elems, -1); + for (size_t i = 0; i < pixel_position_ids->size(); ++i) { + positions_i32[i] = static_cast((*pixel_position_ids)[i]); + } + + const std::vector inputs = { + {"x", input_fp16.data(), npu::NPUNamedInput::DataType::FP16, input_shape}, + {"pixel_position_ids", positions_i32.data(), npu::NPUNamedInput::DataType::INT32, pos_shape}, + }; + written = npu_vision_encoder_->encode_multimodal_input( + inputs, output_fp16.data(), "encoded"); + } else { + written = npu_vision_encoder_->encode( + input_fp16.data(), output_fp16.data(), input_shape, "x", "encoded"); + } + if (written == 0) return false; + + for (size_t i = 0; i < vision_encoder_->output_node_ids.size() + && i < vision_encoder_->logical_outputs.size(); ++i) { + const std::string& name = vision_encoder_->logical_outputs[i]; + size_t node_id = static_cast(vision_encoder_->output_node_ids[i]); + const auto& desc = vision_encoder_->graph->get_output_buffer(node_id); + const Precision cpu_prec = desc.precision; + const size_t cpu_byte_size = desc.byte_size; + const size_t copy_elems = std::min(static_cast(written), + cpu_byte_size / sizeof(__fp16)); + auto& slot = media_features_[name]; + if (cpu_prec == Precision::FP16) { + const size_t prev = slot.size(); + slot.resize(prev + copy_elems * sizeof(__fp16)); + std::memcpy(slot.data() + prev, output_fp16.data(), copy_elems * sizeof(__fp16)); + } else if (cpu_prec == Precision::FP32) { + const size_t prev = slot.size(); + slot.resize(prev + copy_elems * sizeof(float)); + float* dst = reinterpret_cast(slot.data() + prev); + for (size_t k = 0; k < copy_elems; ++k) dst[k] = static_cast(output_fp16[k]); + } else { + const size_t prev = slot.size(); + slot.resize(prev + copy_elems * sizeof(__fp16)); + std::memcpy(slot.data() + prev, output_fp16.data(), copy_elems * sizeof(__fp16)); + } + auto shape_it = media_feature_shapes_.find(name); + std::vector npu_shape; + for (int d : output_shape) npu_shape.push_back(static_cast(d)); + if (shape_it == media_feature_shapes_.end() || shape_it->second.empty()) { + media_feature_shapes_[name] = npu_shape; + } else if (npu_shape.size() >= 2 && shape_it->second.size() == npu_shape.size()) { + shape_it->second[shape_it->second.size() - 2] += npu_shape[npu_shape.size() - 2]; + } + media_feature_precisions_[name] = cpu_prec; + break; + } + return true; +} + +bool Model::lfm2_vl_use_npu_vision() const { + return npu_vision_encoder_ != nullptr + && npu_vision_encoder_->is_available() + && npu_vision_encoder_->has_input("positional_embeddings"); +} + +bool Model::lfm2_vl_encode_tile_npu(const float* pixel_values, const int64_t* mask, + const float* pos_embeds, size_t max_patches, + int dim, size_t patch_dim, std::vector& enc_out) { + if (!npu_vision_encoder_ || !npu_vision_encoder_->is_available()) return false; + + const std::vector x_shape = npu_vision_encoder_->get_input_shape_for("x"); + const std::vector m_shape = npu_vision_encoder_->get_input_shape_for("pixel_attention_mask"); + const std::vector p_shape = npu_vision_encoder_->get_input_shape_for("positional_embeddings"); + if (x_shape.empty() || p_shape.empty()) return false; + + auto elem_count = [](const std::vector& s) -> size_t { + size_t e = 1; + for (int d : s) { if (d <= 0) return 0; e *= static_cast(d); } + return e; + }; + const size_t x_elems = elem_count(x_shape); + const size_t p_elems = elem_count(p_shape); + const size_t m_elems = m_shape.empty() ? 0 : elem_count(m_shape); + const size_t pv_count = max_patches * patch_dim; + const size_t pe_count = max_patches * static_cast(dim); + if (x_elems < pv_count || p_elems < pe_count) return false; + + std::vector<__fp16> x_fp16(x_elems, __fp16(0)); + for (size_t i = 0; i < pv_count; ++i) x_fp16[i] = static_cast<__fp16>(pixel_values[i]); + std::vector<__fp16> p_fp16(p_elems, __fp16(0)); + for (size_t i = 0; i < pe_count; ++i) p_fp16[i] = static_cast<__fp16>(pos_embeds[i]); + std::vector m_i32(m_elems ? m_elems : max_patches, 0); + for (size_t i = 0; i < max_patches && i < m_i32.size(); ++i) m_i32[i] = static_cast(mask[i]); + + const std::vector out_shape = npu_vision_encoder_->get_output_shape(); + size_t out_elems = elem_count(out_shape); + if (out_elems == 0) out_elems = npu_vision_encoder_->get_output_buffer_size(); + std::vector<__fp16> out_fp16(out_elems, __fp16(0)); + + std::vector inputs = { + {"x", x_fp16.data(), npu::NPUNamedInput::DataType::FP16, x_shape}, + {"positional_embeddings", p_fp16.data(), npu::NPUNamedInput::DataType::FP16, p_shape}, + }; + if (!m_shape.empty()) { + inputs.push_back({"pixel_attention_mask", m_i32.data(), npu::NPUNamedInput::DataType::INT32, m_shape}); + } + + size_t written = npu_vision_encoder_->encode_multimodal_input(inputs, out_fp16.data(), "encoded"); + if (written == 0) return false; + + const size_t copy = std::min(enc_out.size(), static_cast(written)); + for (size_t i = 0; i < copy; ++i) enc_out[i] = static_cast(out_fp16[i]); + return true; +} + +} +} diff --git a/cactus-engine/src/npu.cpp b/cactus-engine/src/npu.cpp new file mode 100644 index 000000000..3222b0269 --- /dev/null +++ b/cactus-engine/src/npu.cpp @@ -0,0 +1,15 @@ +#include "engine.h" + +namespace cactus { +namespace npu { + +std::unique_ptr create_encoder() { + return nullptr; +} + +bool is_npu_available() { + return false; +} + +} // namespace npu +} // namespace cactus diff --git a/cactus-engine/src/npu_ane.h b/cactus-engine/src/npu_ane.h new file mode 100644 index 000000000..7594960bf --- /dev/null +++ b/cactus-engine/src/npu_ane.h @@ -0,0 +1,108 @@ +#ifndef CACTUS_NPU_ANE_H +#define CACTUS_NPU_ANE_H + +#include "engine.h" + +#if defined(__APPLE__) +#define CACTUS_HAS_ANE 1 +#else +#define CACTUS_HAS_ANE 0 +#endif + +namespace cactus { +namespace npu { + +#if CACTUS_HAS_ANE + + +class ANEEncoder : public NPUEncoder { +public: + ANEEncoder(); + ~ANEEncoder() override; + + ANEEncoder(const ANEEncoder&) = delete; + ANEEncoder& operator=(const ANEEncoder&) = delete; + + ANEEncoder(ANEEncoder&& other) noexcept; + ANEEncoder& operator=(ANEEncoder&& other) noexcept; + + bool load(const std::string& model_path, const std::string& compute_units = "") override; + + bool preallocate(const std::vector& input_shape, + const std::string& input_name = "x", + const std::string& output_name = "") override; + + size_t encode(const __fp16* input, + __fp16* output, + const std::vector& shape, + const std::string& input_name = "x", + const std::string& output_name = "") override; + + bool is_available() const override; + + std::vector get_input_shape() const override; + + std::vector get_output_shape() const override; + + bool has_input(const std::string& name) const override; + + std::vector get_input_shape_for(const std::string& name) const override; + + __fp16* get_output_buffer() override; + + size_t get_output_buffer_size() const override; + + size_t encode_multimodal_input( + const std::vector& inputs, + __fp16* output, + const std::string& output_name = "") override; + +private: + void* impl_; +}; + +#else + +class ANEEncoder : public NPUEncoder { +public: + ANEEncoder() = default; + ~ANEEncoder() override = default; + + bool load(const std::string&, const std::string& = "") override { return false; } + + bool preallocate(const std::vector&, + const std::string& = "x", + const std::string& = "") override { return false; } + + size_t encode(const __fp16*, + __fp16*, + const std::vector&, + const std::string& = "x", + const std::string& = "") override { return 0; } + + bool is_available() const override { return false; } + + std::vector get_input_shape() const override { return {}; } + + std::vector get_output_shape() const override { return {}; } + + bool has_input(const std::string&) const override { return false; } + + std::vector get_input_shape_for(const std::string&) const override { return {}; } + + __fp16* get_output_buffer() override { return nullptr; } + + size_t get_output_buffer_size() const override { return 0; } + + size_t encode_multimodal_input( + const std::vector&, + __fp16*, + const std::string& = "") override { return 0; } +}; + +#endif // CACTUS_HAS_ANE + +} // namespace npu +} // namespace cactus + +#endif // CACTUS_NPU_ANE_H \ No newline at end of file diff --git a/cactus-engine/src/npu_ane.mm b/cactus-engine/src/npu_ane.mm new file mode 100644 index 000000000..99638e2e0 --- /dev/null +++ b/cactus-engine/src/npu_ane.mm @@ -0,0 +1,1046 @@ +#include "npu_ane.h" + +#if CACTUS_HAS_ANE + +#import +#import +#include +#include +#include +#include +#include +#include +#include +#include "engine.h" + +namespace { + +static std::string to_lower_ascii(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return value; +} + +static bool parse_compute_units(const char* raw_value, MLComputeUnits& out_units) { + if (!raw_value || raw_value[0] == '\0') { + return false; + } + + std::string value = to_lower_ascii(raw_value); + if (value == "all" || value == "default") { + out_units = MLComputeUnitsAll; + return true; + } + if (value == "cpu_and_ne" || value == "cpu-ne" || value == "cpu_ne" || + value == "ne" || value == "ane" || value == "cpuandneuralengine") { + out_units = MLComputeUnitsCPUAndNeuralEngine; + return true; + } + if (value == "cpu_and_gpu" || value == "cpu_gpu" || value == "cpu-gpu" || + value == "cpuandgpu") { + out_units = MLComputeUnitsCPUAndGPU; + return true; + } + if (value == "cpu_only" || value == "cpu-only" || value == "cpu") { + out_units = MLComputeUnitsCPUOnly; + return true; + } + + return false; +} + +static const char* compute_units_to_string(MLComputeUnits units) { + switch (units) { + case MLComputeUnitsAll: + return "ALL"; + case MLComputeUnitsCPUAndGPU: + return "CPU_AND_GPU"; + case MLComputeUnitsCPUOnly: + return "CPU_ONLY"; + case MLComputeUnitsCPUAndNeuralEngine: + return "CPU_AND_NE"; + } + return "UNKNOWN"; +} + +static bool model_path_looks_like_audio_encoder(NSString* model_path) { + if (!model_path) return false; + NSString* lower = [[model_path lastPathComponent] lowercaseString]; + return [lower containsString:@"audio_encoder"]; +} + +static bool model_path_looks_like_vision_encoder(NSString* model_path) { + if (!model_path) return false; + NSString* lower = [[model_path lastPathComponent] lowercaseString]; + return [lower containsString:@"vision_encoder"]; +} + +static void maybe_apply_compute_units_env(const char* env_name, MLComputeUnits& target) { + const char* raw = std::getenv(env_name); + if (!raw || raw[0] == '\0') return; + + MLComputeUnits parsed; + if (parse_compute_units(raw, parsed)) { + target = parsed; + return; + } + + CACTUS_LOG_WARN("npu", "Ignoring invalid " << env_name << "=" << raw); +} + +static MLComputeUnits resolve_compute_units_for_model(NSString* model_path, bool is_prefill, + const std::string& compute_units_hint) { + MLComputeUnits units = MLComputeUnitsCPUAndNeuralEngine; + + if (!is_prefill && model_path_looks_like_audio_encoder(model_path)) { + units = MLComputeUnitsCPUAndGPU; + } + if (!is_prefill && model_path_looks_like_vision_encoder(model_path)) { + units = MLComputeUnitsCPUAndGPU; + } + + if (!compute_units_hint.empty()) { + MLComputeUnits parsed; + if (parse_compute_units(compute_units_hint.c_str(), parsed)) { + units = parsed; + } else { + CACTUS_LOG_WARN("npu", "Ignoring invalid compute units hint=" << compute_units_hint); + } + } + + maybe_apply_compute_units_env("CACTUS_ANE_COMPUTE_UNITS", units); + if (is_prefill) { + maybe_apply_compute_units_env("CACTUS_ANE_PREFILL_COMPUTE_UNITS", units); + } else { + maybe_apply_compute_units_env("CACTUS_ANE_ENCODER_COMPUTE_UNITS", units); + if (model_path_looks_like_audio_encoder(model_path)) { + maybe_apply_compute_units_env("CACTUS_ANE_AUDIO_COMPUTE_UNITS", units); + } + if (model_path_looks_like_vision_encoder(model_path)) { + maybe_apply_compute_units_env("CACTUS_ANE_VISION_COMPUTE_UNITS", units); + } + } + + return units; +} + +static bool should_recompile_mlpackage(NSString* mlpackage_path, NSString* mlmodelc_path) { + NSFileManager* fm = [NSFileManager defaultManager]; + if (![fm fileExistsAtPath:mlmodelc_path]) return true; + + const char* force = std::getenv("CACTUS_ANE_FORCE_RECOMPILE"); + if (force && force[0] != '\0' && std::strcmp(force, "0") != 0) { + return true; + } + + NSError* pkg_err = nil; + NSError* cache_err = nil; + NSDictionary* pkg_attr = [fm attributesOfItemAtPath:mlpackage_path error:&pkg_err]; + NSDictionary* cache_attr = [fm attributesOfItemAtPath:mlmodelc_path error:&cache_err]; + if (!pkg_attr || !cache_attr || pkg_err || cache_err) { + return false; + } + + NSDate* pkg_mtime = pkg_attr[NSFileModificationDate]; + NSDate* cache_mtime = cache_attr[NSFileModificationDate]; + if (!pkg_mtime || !cache_mtime) { + return false; + } + return ([pkg_mtime compare:cache_mtime] == NSOrderedDescending); +} + +static NSURL* resolve_or_compile_model_url(NSString* path, NSError** error) { + NSURL* modelURL = [NSURL fileURLWithPath:path]; + if (![path hasSuffix:@".mlpackage"]) { + return modelURL; + } + + BOOL isDir = NO; + if (![[NSFileManager defaultManager] fileExistsAtPath:path isDirectory:&isDir] || !isDir) { + CACTUS_LOG_ERROR("npu", "ANE mlpackage path is not a valid directory: " << [path UTF8String]); + return modelURL; + } + + NSString* cachedPath = [[path stringByDeletingPathExtension] stringByAppendingPathExtension:@"mlmodelc"]; + NSURL* cachedURL = [NSURL fileURLWithPath:cachedPath]; + NSFileManager* fm = [NSFileManager defaultManager]; + + if (!should_recompile_mlpackage(path, cachedPath)) { + return cachedURL; + } + + if ([fm fileExistsAtPath:cachedPath]) { + NSError* rmErr = nil; + [fm removeItemAtPath:cachedPath error:&rmErr]; + if (rmErr) { + CACTUS_LOG_WARN("npu", "Failed to remove stale mlmodelc: " << [[rmErr localizedDescription] UTF8String]); + } else { + CACTUS_LOG_INFO("npu", "Removed stale mlmodelc cache, recompiling: " << [cachedPath UTF8String]); + } + } + + NSURL* compiledURL = [MLModel compileModelAtURL:modelURL error:error]; + if (*error || !compiledURL) { + CACTUS_LOG_ERROR("npu", "ANE model compilation failed: " << [[*error localizedDescription] UTF8String]); + return modelURL; + } + + NSError* moveError = nil; + [fm moveItemAtURL:compiledURL toURL:cachedURL error:&moveError]; + if (moveError) { + CACTUS_LOG_WARN("npu", "Could not persist compiled mlmodelc cache: " << [[moveError localizedDescription] UTF8String]); + return compiledURL; + } + return cachedURL; +} + +} // namespace + +@interface CactusANEImpl : NSObject + +@property (nonatomic, strong) MLModel* model; +@property (nonatomic, strong) MLModelDescription* modelDescription; +@property (nonatomic, strong) MLMultiArray* cachedInputArray; +@property (nonatomic, strong) MLMultiArray* cachedOutputArray; +@property (nonatomic, strong) NSArray* cachedShape; +@property (nonatomic, strong) NSArray* cachedOutputShape; +@property (nonatomic, strong) NSString* cachedInputName; +@property (nonatomic, strong) NSString* cachedOutputName; +@property (nonatomic, assign) NSUInteger cachedOutputSize; +@property (nonatomic, strong) MLPredictionOptions* predictionOptions; +@property (nonatomic, strong) NSMutableDictionary* cachedMultiInputArrays; +@property (nonatomic, strong) MLMultiArray* cachedMultiOutputArray; +@property (nonatomic, strong) MLPredictionOptions* multiPredictionOptions; +@property (nonatomic, assign) BOOL multiInputPreallocated; +@property (nonatomic, assign) BOOL hasOutputBackings; +@property (nonatomic, assign) BOOL hasMultiOutputBackings; + +- (instancetype)initWithModelPath:(NSString*)path computeUnits:(NSString*)cu; +- (NSArray*)getInputShape; +- (NSArray*)getOutputShape; +- (BOOL)preallocateBuffersWithInput:(NSString*)inputName + shape:(NSArray*)shape + outputName:(NSString*)outputName; +- (BOOL)canUseCachedBufferWithInput:(NSString*)inputName + shape:(NSArray*)shape + outputName:(NSString*)outputName; +- (MLMultiArray*)predictWithInput:(NSString*)inputName + data:(const __fp16*)data + shape:(NSArray*)shape + outputName:(NSString*)outputName; +- (BOOL)preallocateMultiInputBuffersWithOutputName:(NSString*)outputName; +- (size_t)predictMultiInputWithDict:(NSMutableDictionary*)inputDict + outputData:(__fp16*)output + outputName:(NSString*)outputName; + +@end + +@implementation CactusANEImpl + +- (instancetype)initWithModelPath:(NSString*)path computeUnits:(NSString*)cu { + self = [super init]; + if (self) { + NSError* error = nil; + NSURL* modelURL = resolve_or_compile_model_url(path, &error); + + MLModelConfiguration* config = [[MLModelConfiguration alloc] init]; + config.computeUnits = resolve_compute_units_for_model( + path, false, cu ? std::string([cu UTF8String]) : std::string("")); + CACTUS_LOG_INFO("npu", + "ANEEncoder compute units: " + << compute_units_to_string(config.computeUnits) + << " for " << [path UTF8String]); + + _model = [MLModel modelWithContentsOfURL:modelURL configuration:config error:&error]; + if (_model) { + _modelDescription = _model.modelDescription; + } + if (error) { + CACTUS_LOG_ERROR("npu", "ANE model load failed: " << [[error localizedDescription] UTF8String]); + } + } + return self; +} + +- (NSArray*)getInputShape { + if (!_modelDescription) return @[]; + + NSString* inputName = _cachedInputName; + if (!inputName) { + inputName = _modelDescription.inputDescriptionsByName.allKeys.firstObject; + } + + MLFeatureDescription* inputDesc = _modelDescription.inputDescriptionsByName[inputName]; + if (inputDesc && inputDesc.type == MLFeatureTypeMultiArray) { + return inputDesc.multiArrayConstraint.shape; + } + + return @[]; +} + +- (NSArray*)getOutputShape { + if (!_modelDescription) return @[]; + + NSString* outputName = _cachedOutputName; + if (!outputName) { + outputName = _modelDescription.outputDescriptionsByName.allKeys.firstObject; + } + + MLFeatureDescription* outputDesc = _modelDescription.outputDescriptionsByName[outputName]; + if (outputDesc && outputDesc.type == MLFeatureTypeMultiArray) { + return outputDesc.multiArrayConstraint.shape; + } + + return @[]; +} + +- (BOOL)preallocateBuffersWithInput:(NSString*)inputName + shape:(NSArray*)shape + outputName:(NSString*)outputName { + if (!_model) return NO; + + NSError* error = nil; + + MLMultiArrayDataType inputDataType = MLMultiArrayDataTypeFloat16; + MLFeatureDescription* inputDesc = _modelDescription.inputDescriptionsByName[inputName]; + if (inputDesc && inputDesc.multiArrayConstraint) { + inputDataType = inputDesc.multiArrayConstraint.dataType; + } + + _cachedInputArray = [[MLMultiArray alloc] + initWithShape:shape + dataType:inputDataType + error:&error]; + + if (error) { + CACTUS_LOG_ERROR("npu", "ANE preallocate input array failed: " << [[error localizedDescription] UTF8String]); + return NO; + } + + _cachedShape = [shape copy]; + _cachedInputName = [inputName copy]; + _cachedOutputName = outputName ? [outputName copy] + : _modelDescription.outputDescriptionsByName.allKeys.firstObject; + + CACTUS_LOG_DEBUG("npu", + "[CactusANE] prealloc input name=" + << [_cachedInputName UTF8String] + << " shape=" << [[_cachedInputArray.shape description] UTF8String] + << " strides=" << [[_cachedInputArray.strides description] UTF8String] + << " dtype=" << (long)_cachedInputArray.dataType); + + MLFeatureDescription* outputDesc = _modelDescription.outputDescriptionsByName[_cachedOutputName]; + if (outputDesc && outputDesc.type == MLFeatureTypeMultiArray) { + NSArray* outputShape = outputDesc.multiArrayConstraint.shape; + MLMultiArrayDataType outputDataType = outputDesc.multiArrayConstraint.dataType; + + _cachedOutputArray = [[MLMultiArray alloc] + initWithShape:outputShape + dataType:outputDataType + error:&error]; + + if (error) { + CACTUS_LOG_ERROR("npu", "ANE preallocate output array failed: " << [[error localizedDescription] UTF8String]); + return NO; + } + + _cachedOutputShape = [outputShape copy]; + _cachedOutputSize = 1; + for (NSNumber* dim in outputShape) { + _cachedOutputSize *= [dim unsignedIntegerValue]; + } + + CACTUS_LOG_DEBUG("npu", + "[CactusANE] prealloc output name=" + << [_cachedOutputName UTF8String] + << " shape=" << [[_cachedOutputArray.shape description] UTF8String] + << " strides=" << [[_cachedOutputArray.strides description] UTF8String] + << " dtype=" << (long)_cachedOutputArray.dataType); + + _predictionOptions = [[MLPredictionOptions alloc] init]; + _hasOutputBackings = NO; + if (_predictionOptions && _cachedOutputArray && _cachedOutputName) { + _predictionOptions.outputBackings = + @{ _cachedOutputName: _cachedOutputArray }; + _hasOutputBackings = YES; + } + } + + return YES; +} + +- (BOOL)canUseCachedBufferWithInput:(NSString*)inputName + shape:(NSArray*)shape + outputName:(NSString*)outputName { + if (!_cachedInputArray || !_cachedShape) return NO; + if (![_cachedInputName isEqualToString:inputName]) return NO; + if (_cachedShape.count != shape.count) return NO; + + for (NSUInteger i = 0; i < shape.count; i++) { + if (![_cachedShape[i] isEqualToNumber:shape[i]]) return NO; + } + + if (outputName && outputName.length > 0 && ![_cachedOutputName isEqualToString:outputName]) { + return NO; + } + + return YES; +} + +- (MLMultiArray*)predictWithInput:(NSString*)inputName + data:(const __fp16*)data + shape:(NSArray*)shape + outputName:(NSString*)outputName { + if (!_model) return nil; + + NSError* error = nil; + MLMultiArray* inputArray = nil; + + MLMultiArrayDataType expectedDataType = MLMultiArrayDataTypeFloat16; + MLFeatureDescription* inputDesc = _modelDescription.inputDescriptionsByName[inputName]; + if (inputDesc && inputDesc.multiArrayConstraint) { + expectedDataType = inputDesc.multiArrayConstraint.dataType; + } + + BOOL useCached = (expectedDataType == MLMultiArrayDataTypeFloat16) && + [self canUseCachedBufferWithInput:inputName shape:shape outputName:outputName]; + + if (useCached) { + inputArray = _cachedInputArray; + } else { + inputArray = [[MLMultiArray alloc] + initWithShape:shape + dataType:expectedDataType + error:&error]; + + if (error) { + CACTUS_LOG_ERROR("npu", "ANE create input array failed: " << [[error localizedDescription] UTF8String]); + return nil; + } + } + + NSUInteger totalElements = 1; + for (NSNumber* dim in shape) { + totalElements *= [dim unsignedIntegerValue]; + } + + if (expectedDataType == MLMultiArrayDataTypeFloat16) { + __fp16* inputPtr = (__fp16*)inputArray.dataPointer; + memcpy(inputPtr, data, totalElements * sizeof(__fp16)); + } else { + float* inputPtr = (float*)inputArray.dataPointer; + for (NSUInteger i = 0; i < totalElements; i++) { + inputPtr[i] = (float)data[i]; + } + } + + MLFeatureValue* inputFeature = [MLFeatureValue featureValueWithMultiArray:inputArray]; + NSDictionary* inputDict = @{inputName: inputFeature}; + id inputProvider = [[MLDictionaryFeatureProvider alloc] + initWithDictionary:inputDict + error:&error]; + + if (error) { + CACTUS_LOG_ERROR("npu", "ANE create feature provider failed: " << [[error localizedDescription] UTF8String]); + return nil; + } + + id outputProvider = nil; + + if (useCached && _hasOutputBackings && _predictionOptions) { + outputProvider = [_model predictionFromFeatures:inputProvider + options:_predictionOptions + error:&error]; + } else { + outputProvider = [_model predictionFromFeatures:inputProvider error:&error]; + } + + if (error) { + CACTUS_LOG_ERROR("npu", "ANE prediction failed: " << [[error localizedDescription] UTF8String]); + return nil; + } + + NSString* outName = outputName; + if (!outName || outName.length == 0) { + outName = useCached ? _cachedOutputName + : _modelDescription.outputDescriptionsByName.allKeys.firstObject; + } + + MLFeatureValue* outputFeature = [outputProvider featureValueForName:outName]; + static bool logged_output_feature_layout = false; + if (!logged_output_feature_layout && outputFeature && outputFeature.multiArrayValue) { + MLMultiArray* outArr = outputFeature.multiArrayValue; + CACTUS_LOG_DEBUG("npu", + "[CactusANE] output feature name=" + << [outName UTF8String] + << " shape=" << [[outArr.shape description] UTF8String] + << " strides=" << [[outArr.strides description] UTF8String] + << " dtype=" << (long)outArr.dataType); + logged_output_feature_layout = true; + } + return outputFeature.multiArrayValue; +} + +static size_t copyMLArrayToFP16(MLMultiArray* array, __fp16* output) { + size_t count = array.count; + if (!array || !output || count == 0) return 0; + + std::vector dims; + std::vector strides; + const size_t rank = array.shape.count; + bool have_layout = rank == array.strides.count && rank > 0; + if (have_layout) { + dims.resize(rank); + strides.resize(rank); + for (size_t i = 0; i < rank; ++i) { + NSInteger d = [array.shape[i] integerValue]; + NSInteger s = [array.strides[i] integerValue]; + if (d <= 0 || s < 0) { + have_layout = false; + break; + } + dims[i] = static_cast(d); + strides[i] = static_cast(s); + } + } + + if (!have_layout) { + if (array.dataType == MLMultiArrayDataTypeFloat16) { + if (output != (__fp16*)array.dataPointer) { + memcpy(output, array.dataPointer, count * sizeof(__fp16)); + } + } else { + const float* src = (const float*)array.dataPointer; + for (size_t i = 0; i < count; i++) output[i] = (__fp16)src[i]; + } + return count; + } + + auto is_contiguous = [&]() -> bool { + if (!have_layout) return false; + size_t expected = 1; + for (size_t i = rank; i-- > 0;) { + if (dims[i] > 1 && strides[i] != expected) return false; + expected *= dims[i]; + } + return true; + }; + + if (false && is_contiguous()) { + if (array.dataType == MLMultiArrayDataTypeFloat16) { + if (output != (__fp16*)array.dataPointer) { + memcpy(output, array.dataPointer, count * sizeof(__fp16)); + } + } else { + const float* src = (const float*)array.dataPointer; + for (size_t i = 0; i < count; i++) output[i] = (__fp16)src[i]; + } + return count; + } + + std::vector idx(rank, 0); + if (array.dataType == MLMultiArrayDataTypeFloat16) { + const __fp16* src = (const __fp16*)array.dataPointer; + for (size_t i = 0; i < count; ++i) { + size_t offset = 0; + for (size_t d = 0; d < rank; ++d) { + offset += idx[d] * strides[d]; + } + output[i] = src[offset]; + + for (size_t d = rank; d-- > 0;) { + idx[d]++; + if (idx[d] < dims[d]) break; + idx[d] = 0; + } + } + } else { + const float* src = (const float*)array.dataPointer; + for (size_t i = 0; i < count; ++i) { + size_t offset = 0; + for (size_t d = 0; d < rank; ++d) { + offset += idx[d] * strides[d]; + } + output[i] = (__fp16)src[offset]; + + for (size_t d = rank; d-- > 0;) { + idx[d]++; + if (idx[d] < dims[d]) break; + idx[d] = 0; + } + } + } + return count; +} + +static void copyFP16ToMLArray(const __fp16* data, size_t count, MLMultiArray* array) { + if (!array || !data || count == 0) return; + + std::vector dims; + std::vector strides; + const size_t rank = array.shape.count; + bool have_layout = rank == array.strides.count && rank > 0; + if (have_layout) { + dims.resize(rank); + strides.resize(rank); + for (size_t i = 0; i < rank; ++i) { + NSInteger d = [array.shape[i] integerValue]; + NSInteger s = [array.strides[i] integerValue]; + if (d <= 0 || s < 0) { + have_layout = false; + break; + } + dims[i] = static_cast(d); + strides[i] = static_cast(s); + } + } + + if (!have_layout) { + if (array.dataType == MLMultiArrayDataTypeFloat16) { + memcpy(array.dataPointer, data, count * sizeof(__fp16)); + } else { + float* dst = (float*)array.dataPointer; + for (size_t i = 0; i < count; i++) dst[i] = (float)data[i]; + } + return; + } + + auto is_contiguous = [&]() -> bool { + if (!have_layout) return false; + size_t expected = 1; + for (size_t i = rank; i-- > 0;) { + if (dims[i] > 1 && strides[i] != expected) return false; + expected *= dims[i]; + } + return true; + }; + + if (false && is_contiguous()) { + if (array.dataType == MLMultiArrayDataTypeFloat16) { + memcpy(array.dataPointer, data, count * sizeof(__fp16)); + } else { + float* dst = (float*)array.dataPointer; + for (size_t i = 0; i < count; i++) dst[i] = (float)data[i]; + } + return; + } + + std::vector idx(rank, 0); + if (array.dataType == MLMultiArrayDataTypeFloat16) { + __fp16* dst = (__fp16*)array.dataPointer; + for (size_t i = 0; i < count; ++i) { + size_t offset = 0; + for (size_t d = 0; d < rank; ++d) { + offset += idx[d] * strides[d]; + } + dst[offset] = data[i]; + + for (size_t d = rank; d-- > 0;) { + idx[d]++; + if (idx[d] < dims[d]) break; + idx[d] = 0; + } + } + } else { + float* dst = (float*)array.dataPointer; + for (size_t i = 0; i < count; ++i) { + size_t offset = 0; + for (size_t d = 0; d < rank; ++d) { + offset += idx[d] * strides[d]; + } + dst[offset] = (float)data[i]; + + for (size_t d = rank; d-- > 0;) { + idx[d]++; + if (idx[d] < dims[d]) break; + idx[d] = 0; + } + } + } +} + +static void copyInt32ToMLArray(const int32_t* data, size_t count, MLMultiArray* array) { + if (!array || !data || count == 0) return; + + std::vector dims; + std::vector strides; + const size_t rank = array.shape.count; + bool have_layout = rank == array.strides.count && rank > 0; + if (have_layout) { + dims.resize(rank); + strides.resize(rank); + for (size_t i = 0; i < rank; ++i) { + NSInteger d = [array.shape[i] integerValue]; + NSInteger s = [array.strides[i] integerValue]; + if (d <= 0 || s < 0) { + have_layout = false; + break; + } + dims[i] = static_cast(d); + strides[i] = static_cast(s); + } + } + + auto write_at = [&](size_t offset, size_t i) { + if (array.dataType == MLMultiArrayDataTypeInt32) { + ((int32_t*)array.dataPointer)[offset] = data[i]; + } else if (array.dataType == MLMultiArrayDataTypeFloat16) { + ((__fp16*)array.dataPointer)[offset] = static_cast<__fp16>(data[i]); + } else { + ((float*)array.dataPointer)[offset] = static_cast(data[i]); + } + }; + + if (!have_layout) { + for (size_t i = 0; i < count; ++i) write_at(i, i); + return; + } + + std::vector idx(rank, 0); + for (size_t i = 0; i < count; ++i) { + size_t offset = 0; + for (size_t d = 0; d < rank; ++d) { + offset += idx[d] * strides[d]; + } + write_at(offset, i); + + for (size_t d = rank; d-- > 0;) { + idx[d]++; + if (idx[d] < dims[d]) break; + idx[d] = 0; + } + } +} + +- (BOOL)preallocateMultiInputBuffersWithOutputName:(NSString*)outputName { + if (!_model || !_modelDescription) return NO; + + NSError* error = nil; + _cachedMultiInputArrays = [NSMutableDictionary new]; + + for (NSString* inputName in _modelDescription.inputDescriptionsByName) { + MLFeatureDescription* desc = _modelDescription.inputDescriptionsByName[inputName]; + if (desc.type != MLFeatureTypeMultiArray) continue; + + MLMultiArray* array = [[MLMultiArray alloc] + initWithShape:desc.multiArrayConstraint.shape + dataType:desc.multiArrayConstraint.dataType + error:&error]; + if (!array || error) return NO; + _cachedMultiInputArrays[inputName] = array; + } + + NSString* outName = (outputName && outputName.length > 0) + ? outputName : _modelDescription.outputDescriptionsByName.allKeys.firstObject; + + MLFeatureDescription* outputDesc = _modelDescription.outputDescriptionsByName[outName]; + if (outputDesc && outputDesc.type == MLFeatureTypeMultiArray) { + _cachedMultiOutputArray = [[MLMultiArray alloc] + initWithShape:outputDesc.multiArrayConstraint.shape + dataType:outputDesc.multiArrayConstraint.dataType + error:&error]; + if (!_cachedMultiOutputArray || error) return NO; + + _multiPredictionOptions = [[MLPredictionOptions alloc] init]; + _hasMultiOutputBackings = NO; + } + + _multiInputPreallocated = YES; + return YES; +} + +- (size_t)predictMultiInputWithDict:(NSMutableDictionary*)inputDict + outputData:(__fp16*)output + outputName:(NSString*)outputName { + if (!_model) return 0; + + NSError* error = nil; + MLDictionaryFeatureProvider* provider = + [[MLDictionaryFeatureProvider alloc] initWithDictionary:inputDict error:&error]; + if (!provider || error) return 0; + + id result = nil; + if (_hasMultiOutputBackings && _multiPredictionOptions) { + result = [_model predictionFromFeatures:provider options:_multiPredictionOptions error:&error]; + } else { + result = [_model predictionFromFeatures:provider error:&error]; + } + if (!result || error) return 0; + + MLFeatureValue* outFeature = [result featureValueForName:outputName]; + if (!outFeature || !outFeature.multiArrayValue) return 0; + return copyMLArrayToFP16(outFeature.multiArrayValue, output); +} + +@end + +namespace cactus { +namespace npu { + +static bool g_npu_enabled = true; + +ANEEncoder::ANEEncoder() : impl_(nullptr) {} + +ANEEncoder::~ANEEncoder() { + if (impl_) { + (void)(__bridge_transfer CactusANEImpl*)impl_; + impl_ = nullptr; + } +} + +ANEEncoder::ANEEncoder(ANEEncoder&& other) noexcept : impl_(other.impl_) { + other.impl_ = nullptr; +} + +ANEEncoder& ANEEncoder::operator=(ANEEncoder&& other) noexcept { + if (this != &other) { + if (impl_) { + (void)(__bridge_transfer CactusANEImpl*)impl_; + } + impl_ = other.impl_; + other.impl_ = nullptr; + } + return *this; +} + +bool ANEEncoder::load(const std::string& model_path, const std::string& compute_units) { + @autoreleasepool { + CACTUS_LOG_INFO("npu", "ANEEncoder loading model: " << model_path); + NSString* path = [NSString stringWithUTF8String:model_path.c_str()]; + + if (![[NSFileManager defaultManager] fileExistsAtPath:path]) { + CACTUS_LOG_WARN("npu", "not using NPU"); + return false; + } + + CactusANEImpl* impl = [[CactusANEImpl alloc] initWithModelPath:path computeUnits:(compute_units.empty() ? nil : [NSString stringWithUTF8String:compute_units.c_str()])]; + + if (impl && impl.model) { + impl_ = (__bridge_retained void*)impl; + CACTUS_LOG_INFO("npu", "ANEEncoder model loaded successfully: " << model_path); + return true; + } + CACTUS_LOG_ERROR("npu", "ANEEncoder model load failed: " << model_path); + return false; + } +} + +bool ANEEncoder::preallocate(const std::vector& input_shape, + const std::string& input_name, + const std::string& output_name) { + if (!impl_) return false; + + @autoreleasepool { + CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; + + NSMutableArray* shapeArray = [NSMutableArray array]; + for (int dim : input_shape) { + [shapeArray addObject:@(dim)]; + } + + NSString* inName = [NSString stringWithUTF8String:input_name.c_str()]; + NSString* outName = output_name.empty() + ? nil + : [NSString stringWithUTF8String:output_name.c_str()]; + + return [impl preallocateBuffersWithInput:inName shape:shapeArray outputName:outName]; + } +} + +size_t ANEEncoder::encode(const __fp16* input, + __fp16* output, + const std::vector& shape, + const std::string& input_name, + const std::string& output_name) { + if (!impl_ || !input || !output) return 0; + + @autoreleasepool { + CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; + + NSArray* shapeArray = impl.cachedShape; + bool shapeMatches = (shapeArray && shapeArray.count == shape.size()); + if (shapeMatches) { + for (size_t i = 0; i < shape.size(); ++i) { + if ([shapeArray[i] intValue] != shape[i]) { + shapeMatches = false; + break; + } + } + } + + if (!shapeMatches) { + NSMutableArray* newShapeArray = [NSMutableArray arrayWithCapacity:shape.size()]; + for (int dim : shape) { + [newShapeArray addObject:@(dim)]; + } + shapeArray = newShapeArray; + } + + NSString* inName = impl.cachedInputName; + if (!inName) { + inName = [NSString stringWithUTF8String:input_name.c_str()]; + } + NSString* outName = impl.cachedOutputName; + if (!outName && !output_name.empty()) { + outName = [NSString stringWithUTF8String:output_name.c_str()]; + } + + MLMultiArray* mlOutput = [impl predictWithInput:inName + data:input + shape:shapeArray + outputName:outName]; + + if (mlOutput) { + return copyMLArrayToFP16(mlOutput, output); + } + } + + return 0; +} + +bool ANEEncoder::is_available() const { + if (!impl_) return false; + CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; + return impl.model != nil; +} + +std::vector ANEEncoder::get_input_shape() const { + std::vector result; + if (!impl_) return result; + CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; + NSArray* shape = [impl getInputShape]; + for (NSNumber* dim in shape) { + result.push_back([dim intValue]); + } + return result; +} + +std::vector ANEEncoder::get_output_shape() const { + std::vector result; + if (!impl_) return result; + CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; + NSArray* shape = [impl getOutputShape]; + for (NSNumber* dim in shape) { + result.push_back([dim intValue]); + } + return result; +} + +bool ANEEncoder::has_input(const std::string& name) const { + if (!impl_) return false; + CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; + if (!impl.modelDescription) return false; + NSString* nsName = [NSString stringWithUTF8String:name.c_str()]; + return impl.modelDescription.inputDescriptionsByName[nsName] != nil; +} + +std::vector ANEEncoder::get_input_shape_for(const std::string& name) const { + std::vector result; + if (!impl_) return result; + CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; + if (!impl.modelDescription) return result; + NSString* nsName = [NSString stringWithUTF8String:name.c_str()]; + MLFeatureDescription* desc = impl.modelDescription.inputDescriptionsByName[nsName]; + if (desc && desc.type == MLFeatureTypeMultiArray) { + for (NSNumber* dim in desc.multiArrayConstraint.shape) { + result.push_back([dim intValue]); + } + } + return result; +} + +__fp16* ANEEncoder::get_output_buffer() { + if (!impl_) return nullptr; + CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; + if (!impl.hasOutputBackings || !impl.cachedOutputArray) return nullptr; + return (__fp16*)impl.cachedOutputArray.dataPointer; +} + +size_t ANEEncoder::get_output_buffer_size() const { + if (!impl_) return 0; + CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; + if (!impl.hasOutputBackings || !impl.cachedOutputArray) return 0; + return impl.cachedOutputSize; +} + +size_t ANEEncoder::encode_multimodal_input( + const std::vector& inputs, + __fp16* output, + const std::string& output_name) { + if (!impl_ || !output || inputs.empty()) return 0; + + @autoreleasepool { + CactusANEImpl* impl = (__bridge CactusANEImpl*)impl_; + + NSString* outName = output_name.empty() + ? impl.modelDescription.outputDescriptionsByName.allKeys.firstObject + : [NSString stringWithUTF8String:output_name.c_str()]; + + if (!impl.multiInputPreallocated) { + [impl preallocateMultiInputBuffersWithOutputName:outName]; + } + + NSMutableDictionary* inputDict = [NSMutableDictionary new]; + + for (const auto& input : inputs) { + NSString* nsName = [NSString stringWithUTF8String:input.name.c_str()]; + + size_t total = 1; + for (int dim : input.shape) total *= dim; + + MLMultiArray* array = impl.cachedMultiInputArrays[nsName]; + if (!array) { + MLMultiArrayDataType dataType = MLMultiArrayDataTypeFloat16; + MLFeatureDescription* desc = impl.modelDescription.inputDescriptionsByName[nsName]; + if (desc && desc.multiArrayConstraint) dataType = desc.multiArrayConstraint.dataType; + + NSMutableArray* shapeArray = [NSMutableArray arrayWithCapacity:input.shape.size()]; + for (int dim : input.shape) [shapeArray addObject:@(dim)]; + + NSError* arrayError = nil; + array = [[MLMultiArray alloc] initWithShape:shapeArray dataType:dataType error:&arrayError]; + if (!array || arrayError) return 0; + } + + if (input.data_type == NPUNamedInput::DataType::INT32) { + copyInt32ToMLArray(static_cast(input.data), total, array); + } else { + copyFP16ToMLArray(static_cast(input.data), total, array); + } + inputDict[nsName] = [MLFeatureValue featureValueWithMultiArray:array]; + } + + return [impl predictMultiInputWithDict:inputDict outputData:output outputName:outName]; + } +} + +std::unique_ptr create_encoder() { + return std::make_unique(); +} + +bool is_npu_available() { + CACTUS_LOG_INFO("npu", "is_npu_available: " << (g_npu_enabled ? "true" : "false")); + return g_npu_enabled; +} + +} // namespace npu +} // namespace cactus + +#else // !CACTUS_HAS_ANE + +namespace cactus { +namespace npu { + +std::unique_ptr create_encoder() { + return nullptr; +} + +bool is_npu_available() { + return false; +} + +} // namespace npu +} // namespace cactus + +#endif // CACTUS_HAS_ANE diff --git a/cactus-engine/src/telemetry_impl.cpp b/cactus-engine/src/telemetry_impl.cpp index 150440b53..0f5d85b1f 100644 --- a/cactus-engine/src/telemetry_impl.cpp +++ b/cactus-engine/src/telemetry_impl.cpp @@ -1302,6 +1302,11 @@ void init(const char* project_id_param, const char* project_scope_param, const c cloud_disabled_from_env = true; } } + if (const char* env = std::getenv("CACTUS_DISABLE_CLOUD_HANDOFF")) { + if (env[0] != '\0' && !(env[0] == '0' && env[1] == '\0')) { + cloud_disabled_from_env = true; + } + } const char* env_url = std::getenv("CACTUS_SUPABASE_URL"); const char* env_key = std::getenv("CACTUS_SUPABASE_KEY"); const char* env_project = std::getenv("CACTUS_PROJECT_ID"); diff --git a/cactus-engine/src/tokenizer.cpp b/cactus-engine/src/tokenizer.cpp index d792578db..370c61f3e 100644 --- a/cactus-engine/src/tokenizer.cpp +++ b/cactus-engine/src/tokenizer.cpp @@ -6,6 +6,8 @@ #include #include #include +#include +#include #include namespace cactus { @@ -57,6 +59,24 @@ std::string trim_copy(const std::string& value) { return value.substr(start, end - start + 1); } +std::string lower_copy(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return value; +} + +bool contains_keyword(const std::string& text, const std::vector& keywords) { + if (text.empty() || keywords.empty()) return false; + const std::string lowered = lower_copy(text); + for (const auto& keyword : keywords) { + if (!keyword.empty() && lowered.find(lower_copy(keyword)) != std::string::npos) { + return true; + } + } + return false; +} + TokenizerRuntimeConfig::TokenizerType parse_tokenizer_type(const std::string& value) { if (value == "bpe") return TokenizerRuntimeConfig::TokenizerType::BPE; if (value == "sentencepiece") return TokenizerRuntimeConfig::TokenizerType::SENTENCEPIECE; @@ -361,7 +381,7 @@ void Tokenizer::load_chat_template(const std::string& template_file) { } void Tokenizer::detect_model_type(const std::string& config_path) { - model_type_ = ModelType::GEMMA4; + model_type_ = ModelType::UNKNOWN; std::ifstream file(config_path); if (!file.is_open()) return; @@ -382,6 +402,11 @@ void Tokenizer::detect_model_type(const std::string& config_path) { model_type_ = ModelType::GEMMA4; } else if (lower_line.find("gemma") != std::string::npos) { model_type_ = ModelType::GEMMA; + } else if (lower_line.find("whisper") != std::string::npos) { + model_type_ = ModelType::WHISPER; + } else if (lower_line.find("parakeet_tdt") != std::string::npos || + lower_line.find("parakeet-tdt") != std::string::npos) { + model_type_ = ModelType::PARAKEET_TDT; } } else if (lower_line.find("model_variant") != std::string::npos) { if (lower_line.find("vlm") != std::string::npos) { model_variant_ = ModelVariant::VLM; } @@ -416,18 +441,69 @@ std::string Tokenizer::get_default_stop_sequence() const { if (model_type_ == ModelType::GEMMA) { return ""; } - return ""; + if (model_type_ == ModelType::GEMMA4) { + return gemma4_turn_end_token_; + } + if (model_type_ == ModelType::UNKNOWN && has_chat_template_ + && chat_template_.find("<|im_start|>") != std::string::npos + && chat_template_.find("<|im_end|>") != std::string::npos) { + return "<|im_end|>"; + } + return ""; } std::vector Tokenizer::apply_chat_template(const std::vector& messages, bool add_generation_prompt) const { return encode(format_chat_prompt(messages, add_generation_prompt)); } +bool Tokenizer::uses_default_chat_style(const std::vector& messages, const std::string& tools_json, + bool enable_thinking_if_supported) const { + const bool use_default_text_style = + model_type_ == ModelType::UNKNOWN || text_prompt_style_ == "default_chat"; + return !has_chat_template_ && tools_json.empty() && !enable_thinking_if_supported && + std::all_of(messages.begin(), messages.end(), [](const ChatMessage& msg) { + return msg.images.empty() && msg.audio.empty() && msg.tool_calls.empty() && + msg.audio_soft_token_count == 0 && msg.role != "tool"; + }) && + use_default_text_style; +} + std::string Tokenizer::format_chat_prompt(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported) const { + if (uses_default_chat_style(messages, tools_json, enable_thinking_if_supported)) { + return format_default_chat_style(messages, add_generation_prompt); + } + if (model_type_ == ModelType::QWEN) { return format_qwen_style(messages, add_generation_prompt, tools_json, enable_thinking_if_supported); } + // A common generic Hugging Face template is plain ChatML. We do not need + // to identify the model family to render it: the delimiters define the + // message contract. Preserve a literal default system message embedded + // in the template when the caller did not provide one. + if (model_type_ == ModelType::UNKNOWN && has_chat_template_ + && chat_template_.find("<|im_start|>") != std::string::npos + && chat_template_.find("<|im_end|>") != std::string::npos) { + std::vector templated_messages = messages; + const bool has_system = !messages.empty() + && (messages.front().role == "system" || messages.front().role == "developer"); + if (!has_system) { + const std::string marker = "<|im_start|>system\n"; + const size_t begin = chat_template_.find(marker); + const size_t end = begin == std::string::npos + ? std::string::npos + : chat_template_.find("<|im_end|>", begin + marker.size()); + if (begin != std::string::npos && end != std::string::npos) { + ChatMessage system_message; + system_message.role = "system"; + system_message.content = chat_template_.substr( + begin + marker.size(), end - (begin + marker.size())); + templated_messages.insert(templated_messages.begin(), std::move(system_message)); + } + } + return format_qwen_style( + templated_messages, add_generation_prompt, tools_json, enable_thinking_if_supported); + } if (model_type_ == ModelType::LFM2) { return format_lfm2_style(messages, add_generation_prompt, tools_json, enable_thinking_if_supported); } @@ -440,6 +516,30 @@ std::string Tokenizer::format_chat_prompt(const std::vector& messag return format_gemma4_style(messages, add_generation_prompt, tools_json, enable_thinking_if_supported); } +std::string Tokenizer::format_default_chat_style(const std::vector& messages, bool add_generation_prompt) const { + std::string result; + if (model_type_ == ModelType::GEMMA || model_type_ == ModelType::GEMMA4) { + result = ""; + } else if (model_type_ == ModelType::LFM2) { + result = "<|startoftext|>"; + } + + for (const auto& msg : messages) { + if (msg.role == "system" || msg.role == "developer") { + result += "System: " + msg.content + "\n"; + } else if (msg.role == "assistant" || msg.role == "model") { + result += "Assistant: " + msg.content + "\n"; + } else { + result += "User: " + msg.content + "\n"; + } + } + + if (add_generation_prompt) { + result += "Assistant:"; + } + return result; +} + std::string Tokenizer::format_gemma_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const { std::string result = ""; @@ -688,13 +788,13 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa } if (enable_thinking_if_supported || !sys_content.empty() || !tools_json.empty()) { - result += "<|turn>system\n"; + result += gemma4_turn_start_token_ + "system\n"; if (enable_thinking_if_supported) { result += "<|think|>\n"; } result += sys_content; result += tools_json; - result += "\n"; + result += gemma4_turn_end_token_ + "\n"; } auto compute_soft_tokens = [&](const std::string& image_path) -> size_t { @@ -716,10 +816,116 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa if (tw == 0) tw = side; return static_cast((th / p / k) * (tw / p / k)); }; + + bool has_media = false; + for (const auto& msg : messages) { + if (!msg.images.empty() || msg.audio_soft_token_count > 0) { + has_media = true; + break; + } + } + auto focused_modality = [&](const ChatMessage& msg) -> std::string { + if (media_focus_policy_ != "focused_modality_after_text") return ""; + if (msg.images.empty() || msg.audio_soft_token_count == 0) return ""; + + const bool image_focus = contains_keyword(msg.content, media_image_focus_keywords_); + const bool audio_focus = contains_keyword(msg.content, media_audio_focus_keywords_); + if (image_focus && !audio_focus) return "image"; + if (audio_focus && !image_focus) return "audio"; + return ""; + }; + auto effective_image_position = [&](const ChatMessage& msg) -> std::string { + const std::string focus = focused_modality(msg); + if (focus == "image") return "after_text"; + if (focus == "audio") return "before_text"; + return media_image_prompt_position_; + }; + auto effective_audio_position = [&](const ChatMessage& msg) -> std::string { + const std::string focus = focused_modality(msg); + if (focus == "audio") return "after_text"; + if (focus == "image") return "before_text"; + return media_audio_prompt_position_; + }; + auto append_separated_media = [&](std::string& out, const ChatMessage& msg, const std::string& position, bool& needs_separator) { + const std::string image_position = effective_image_position(msg); + const std::string audio_position = effective_audio_position(msg); + auto append_separator = [&]() { + if (needs_separator) out += " "; + }; + auto append_images = [&]() { + if (image_position != position) return; + for (const auto& image_path : msg.images) { + size_t n = compute_soft_tokens(image_path); + if (n > 0) { + append_separator(); + out += media_image_begin_token_; + for (size_t j = 0; j < n; j++) out += media_image_token_; + out += media_image_end_token_; + needs_separator = true; + } + } + }; + auto append_audio = [&]() { + if (audio_position != position || msg.audio_soft_token_count == 0) return; + append_separator(); + out += media_audio_begin_token_; + for (size_t j = 0; j < msg.audio_soft_token_count; j++) out += media_audio_token_; + out += media_audio_end_token_; + needs_separator = true; + }; + for (const auto& modality : media_prompt_order_) { + if (modality == "audio") append_audio(); + else if (modality == "image" || modality == "vision") append_images(); + } + }; + auto append_turn_media = [&](std::string& out, const ChatMessage& msg, const std::string& position) { + const std::string image_position = effective_image_position(msg); + const std::string audio_position = effective_audio_position(msg); + auto append_images = [&]() { + if (image_position != position) return; + for (const auto& image_path : msg.images) { + size_t n = compute_soft_tokens(image_path); + if (n > 0) { + out += media_image_begin_token_; + for (size_t j = 0; j < n; j++) out += media_image_token_; + out += media_image_end_token_; + } + } + }; + auto append_audio = [&]() { + if (audio_position != position || msg.audio_soft_token_count == 0) return; + out += media_audio_begin_token_; + for (size_t j = 0; j < msg.audio_soft_token_count; j++) out += media_audio_token_; + out += media_audio_end_token_; + }; + for (const auto& modality : media_prompt_order_) { + if (modality == "audio") append_audio(); + else if (modality == "image" || modality == "vision") append_images(); + } + }; + + if (tools_json.empty() && !enable_thinking_if_supported && has_media && + (gemma4_raw_media_prompt_ || std::getenv("CACTUS_GEMMA4_RAW_MEDIA_PROMPT") != nullptr)) { + result = ""; + for (size_t i = first_msg; i < messages.size(); i++) { + const auto& msg = messages[i]; + bool needs_separator = false; + append_separated_media(result, msg, "before_text", needs_separator); + if (!msg.content.empty()) { + if (needs_separator) result += " "; + result += msg.content; + needs_separator = true; + } + append_separated_media(result, msg, "after_text", needs_separator); + if (i + 1 < messages.size()) result += "\n"; + } + return result; + } + bool in_model_turn = false; std::vector pending_call_names; auto close_model_turn = [&]() { - if (in_model_turn) { result += "\n"; in_model_turn = false; } + if (in_model_turn) { result += gemma4_turn_end_token_ + "\n"; in_model_turn = false; } }; for (size_t i = first_msg; i < messages.size(); i++) { @@ -728,7 +934,7 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa : (msg.role == "developer") ? "system" : msg.role; if (role == "model") { - if (!in_model_turn) { result += "<|turn>model\n"; in_model_turn = true; } + if (!in_model_turn) { result += gemma4_turn_start_token_ + "model\n"; in_model_turn = true; } result += msg.content; for (const auto& tc : msg.tool_calls) { result += format_tool_call_for_prompt(tc.name, tc.arguments, true); @@ -736,7 +942,7 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa } if (msg.tool_calls.empty()) close_model_turn(); } else if (role == "tool") { - if (!in_model_turn) { result += "<|turn>model\n"; in_model_turn = true; } + if (!in_model_turn) { result += gemma4_turn_start_token_ + "model\n"; in_model_turn = true; } std::string fn = !msg.name.empty() ? msg.name : (!pending_call_names.empty() ? pending_call_names.front() : std::string("unknown")); @@ -744,29 +950,16 @@ std::string Tokenizer::format_gemma4_style(const std::vector& messa result += format_tool_response_for_prompt(fn, msg.content, true); } else { close_model_turn(); - result += "<|turn>" + role + "\n"; - for (const auto& image_path : msg.images) { - size_t n = compute_soft_tokens(image_path); - if (n > 0) { - result += "<|image>"; - for (size_t j = 0; j < n; j++) - result += "<|image|>"; - result += ""; - } - } + result += gemma4_turn_start_token_ + role + "\n"; + append_turn_media(result, msg, "before_text"); result += msg.content; - if (msg.audio_soft_token_count > 0) { - result += "<|audio>"; - for (size_t j = 0; j < msg.audio_soft_token_count; j++) - result += "<|audio|>"; - result += ""; - } - result += "\n"; + append_turn_media(result, msg, "after_text"); + result += gemma4_turn_end_token_ + "\n"; } } if (add_generation_prompt) { - if (!in_model_turn) result += "<|turn>model\n"; + if (!in_model_turn) result += gemma4_turn_start_token_ + "model\n"; } else { close_model_turn(); } diff --git a/cactus-engine/src/transcribe.cpp b/cactus-engine/src/transcribe.cpp index 9f2ebe928..90c4cf5ca 100644 --- a/cactus-engine/src/transcribe.cpp +++ b/cactus-engine/src/transcribe.cpp @@ -93,7 +93,7 @@ CACTUS_FFI_EXPORT int cactus_preprocess_audio_features( cactus::engine::Config cfg; cfg.model_type = cactus::engine::Config::ModelType::GEMMA4; cfg.audio_input_feat_size = static_cast(bins); - // The Python transpiled runtime shape-specializes Gemma4 audio to + // The graph-bundle runtime shape-specializes Gemma4 audio to // its requested capacity (currently capped at 30s). Do not apply // the native interactive model's shorter audio_soft_tokens limit // here, otherwise long-file multimodal prompts are silently diff --git a/cactus-engine/src/utils.h b/cactus-engine/src/utils.h index afc720a1b..5e84e6850 100644 --- a/cactus-engine/src/utils.h +++ b/cactus-engine/src/utils.h @@ -432,15 +432,7 @@ inline AudioPreprocessResult preprocess_audio_for_gemma4( result.num_frames = mel.size() / mel_bins; result.features = transpose_mel_to_frame_major(mel, mel_bins, result.num_frames); - size_t max_soft = model_config.audio_soft_tokens; - size_t max_frames_per_chunk = max_soft > 0 ? (max_soft * 2 - 1) * 2 - 1 : 0; - if (max_soft > 0 && result.num_frames > max_frames_per_chunk) { - size_t num_chunks = (result.num_frames + max_frames_per_chunk - 1) / max_frames_per_chunk; - result.num_soft_tokens = num_chunks * max_soft; - } else { - size_t after_stage1 = (result.num_frames + 1) / 2; - result.num_soft_tokens = (after_stage1 + 1) / 2; - } + result.num_soft_tokens = (result.num_frames + 3) / 4; return result; } diff --git a/cactus-engine/tests/ios/configure_xcode.rb b/cactus-engine/tests/ios/configure_xcode.rb index 6e8379500..a63bd43e6 100755 --- a/cactus-engine/tests/ios/configure_xcode.rb +++ b/cactus-engine/tests/ios/configure_xcode.rb @@ -232,7 +232,7 @@ def generate_app_delegate(output_path, test_files) config.build_settings['OTHER_LDFLAGS'].reject! { |flag| flag.to_s.include?('libcactus') } config.build_settings['OTHER_LDFLAGS'] << static_lib_path - ['-framework Foundation', '-framework Accelerate', '-framework Security', '-framework SystemConfiguration', '-framework CFNetwork', '-framework Metal', '-framework MetalPerformanceShaders'].each do |framework| + ['-framework Foundation', '-framework Accelerate', '-framework CoreML', '-framework Security', '-framework SystemConfiguration', '-framework CFNetwork', '-framework Metal', '-framework MetalPerformanceShaders'].each do |framework| config.build_settings['OTHER_LDFLAGS'] << framework unless config.build_settings['OTHER_LDFLAGS'].include?(framework) end if vendored_curl_lib diff --git a/cactus-engine/tests/run.cpp b/cactus-engine/tests/run.cpp index d3de3e731..68739c57c 100644 --- a/cactus-engine/tests/run.cpp +++ b/cactus-engine/tests/run.cpp @@ -1087,6 +1087,10 @@ int main(int argc, char** argv) { bool attach_media = !current_image.empty() || !current_audio.empty() || !current_pcm.empty(); if (attach_media) { cactus_reset(model); + for (auto& turn : history) { + turn.image.clear(); + turn.audio.clear(); + } } history.push_back({"user", input, current_image, current_audio}); std::string messages = build_messages(system_prompt, history); diff --git a/cactus-engine/tests/test_kv_compress.cpp b/cactus-engine/tests/test_kv_compress.cpp index 07365db95..76eaae518 100644 --- a/cactus-engine/tests/test_kv_compress.cpp +++ b/cactus-engine/tests/test_kv_compress.cpp @@ -366,6 +366,25 @@ bool test_config_parse_rolling_fields() { return cfg.kv_compress && cfg.kv_compress_trigger_len == 4096 && cfg.kv_compress_target_len == 2048; } +bool test_generic_config_skips_family_specific_validation() { + char tmpl[] = "/tmp/cactus_genericcfg_XXXXXX"; + int fd = mkstemp(tmpl); + if (fd < 0) return false; + ::close(fd); + { + std::ofstream f(tmpl); + f << "model_type=generic\n" + << "hidden_dim=2048\n" + << "num_layers=16\n" + << "attention_heads=32\n" + << "attention_kv_heads=8\n"; + } + cactus::engine::Config cfg; + bool parsed = cfg.from_json(tmpl); + std::remove(tmpl); + return parsed && cfg.model_type == cactus::engine::Config::ModelType::GENERIC; +} + bool test_trigger_zero_gates_rolling() { // The maybe_roll_compact gate fires iff kv_compress && trigger_len > 0 && seq_len >= trigger_len. const size_t n = 5000, kv_heads = 2, head_dim = 16, max_seq = 5008; @@ -993,6 +1012,7 @@ int main() { runner.run_test("rerope_zero_delta_noop", test_rerope_zero_delta_noop()); runner.run_test("rolling_bounded_compaction", test_rolling_bounded_compaction()); runner.run_test("config_parse_rolling_fields", test_config_parse_rolling_fields()); + runner.run_test("generic_config_skips_family_specific_validation", test_generic_config_skips_family_specific_validation()); runner.run_test("trigger_zero_gates_rolling", test_trigger_zero_gates_rolling()); runner.run_test("degenerate_rolling_config_disabled", test_degenerate_rolling_config_disabled()); runner.run_test("env_override_parse", test_env_override_parse()); diff --git a/cactus-engine/tests/test_llm.cpp b/cactus-engine/tests/test_llm.cpp index 8bcff7526..7f05eaa21 100644 --- a/cactus-engine/tests/test_llm.cpp +++ b/cactus-engine/tests/test_llm.cpp @@ -659,19 +659,105 @@ bool test_batch_distinct4_matches_single() { return ok; } +bool test_generate_batch_ragged() { + cactus_model_t model = load_gemma4_or_skip(); + if (!model) return true; + + auto* handle = static_cast(model); + auto* tok = handle->model->get_tokenizer(); + std::vector a = tok->encode("The capital of France is"); + std::vector b = tok->encode("Hello"); + if (a.empty() || b.empty() || a.size() == b.size()) { + std::cerr << " could not build two ragged prompts\n"; cactus_destroy(model); return false; + } + const size_t M = 16; + + auto batched = handle->model->generate_batch({a, b}, M); + if (batched.empty()) { + std::cout << " [WARN] bundle not dynamic-batch capable; skipping\n"; + cactus_destroy(model); + return true; + } + + auto ref_a = handle->model->generate_batch({a}, M); + auto ref_b = handle->model->generate_batch({b}, M); + cactus_destroy(model); + + bool ok = true; + if (batched.size() != 2 || ref_a.empty() || ref_b.empty()) { std::cerr << " unexpected stream counts\n"; return false; } + if (batched[0].size() != M || batched[1].size() != M) { std::cerr << " wrong generated length\n"; ok = false; } + if (batched[0] != ref_a[0]) { std::cerr << " ragged row 0 (len " << a.size() << ") diverged from single-stream\n"; ok = false; } + if (batched[1] != ref_b[0]) { std::cerr << " ragged row 1 (len " << b.size() << ") diverged from single-stream\n"; ok = false; } + return ok; +} + +bool test_decode_batch_throughput() { + cactus_model_t model = load_gemma4_or_skip(); + if (!model) return true; + + auto* handle = static_cast(model); + auto* tok = handle->model->get_tokenizer(); + std::vector ids = tok->encode("The"); + if (ids.empty()) { std::cerr << " empty seed\n"; cactus_destroy(model); return false; } + uint32_t seed = ids.back(); + const size_t M = 64; + + auto probe = handle->model->decode_batch(std::vector(2, seed), 2); + if (probe.empty()) { + std::cout << " [WARN] bundle not dynamic-batch capable; skipping\n"; + cactus_destroy(model); + return true; + } + + std::cout << " batched decode throughput (" << M << " tokens/stream):\n"; + std::cout << " N agg tok/s per-stream tok/s speedup\n"; + double base_agg = 0.0, best_agg = 0.0; + for (size_t N : {(size_t)1, (size_t)2, (size_t)4, (size_t)8, (size_t)16, (size_t)32, (size_t)64}) { + std::vector seeds(N, seed); + auto t0 = std::chrono::high_resolution_clock::now(); + auto streams = handle->model->decode_batch(seeds, M); + auto t1 = std::chrono::high_resolution_clock::now(); + if (streams.size() != N) break; + double secs = std::chrono::duration(t1 - t0).count(); + if (secs <= 0.0) continue; + double agg = double(N * M) / secs; + double per = double(M) / secs; + if (N == 1) base_agg = agg; + if (agg > best_agg) best_agg = agg; + std::printf(" %-4zu %9.1f %14.1f %6.2fx\n", + N, agg, per, base_agg > 0.0 ? agg / base_agg : 1.0); + } + cactus_destroy(model); + + if (base_agg <= 0.0) { std::cerr << " no N=1 baseline measured\n"; return false; } + if (best_agg <= base_agg) { + std::cerr << " batching did not improve aggregate throughput\n"; + return false; + } + return true; +} + int main() { TestUtils::TestRunner runner("LLM Tests"); + runner.run_test("1k_context", test_1k_context()); runner.run_test("streaming", test_streaming()); runner.run_test("prefill", test_prefill()); + runner.run_test("prefill_idempotent_reuse", test_prefill_idempotent_reuse()); + runner.run_test("prefill_prefix_extension_reuse", test_prefill_prefix_extension_reuse()); runner.run_test("prefill_invalidated_on_message_change", test_prefill_invalidated_on_message_change()); runner.run_test("chunked_prefill_padding", test_chunked_prefill_padding()); runner.run_test("tool_calls", test_tool_call()); runner.run_test("tool_multiple_tool_call_invocations", test_multiple_tool_call_invocations()); + runner.run_test("tool_calls_with_three_tools", test_tool_call_with_three_tools()); runner.run_test("tool_constraint_clear_releases_bias", test_tool_constraint_clear_releases_bias()); runner.run_test("partition_thinking_response", test_partition_thinking_response()); runner.run_test("prompt_retains_thinking", test_prompt_gemma4_retains_thinking()); runner.run_test("complete_thinking_api_clean", test_complete_gemma4_thinking_api_clean()); runner.run_test("multiturn_thinking_persist", test_multiturn_thinking_persist()); + runner.run_test("multiturn_turn2_distinct", test_multiturn_turn2_distinct()); + runner.run_test("decode_batch", test_decode_batch()); + runner.run_test("generate_batch_ragged", test_generate_batch_ragged()); + runner.run_test("decode_batch_throughput", test_decode_batch_throughput()); runner.run_test("batch_distinct4_matches_single", test_batch_distinct4_matches_single()); runner.print_summary(); return runner.all_passed() ? 0 : 1; diff --git a/cactus-engine/tests/transcribe.cpp b/cactus-engine/tests/transcribe.cpp index 52aac854e..3ccf558a6 100644 --- a/cactus-engine/tests/transcribe.cpp +++ b/cactus-engine/tests/transcribe.cpp @@ -92,7 +92,7 @@ static std::string whisper_prompt(const std::string& model_path, const std::stri std::string p = model_path; std::transform(p.begin(), p.end(), p.begin(), [](unsigned char c) { return (char)std::tolower(c); }); if (p.find("whisper") != std::string::npos) - return "<|startoftranscript|><|" + language + "|><|transcribe|>"; + return "<|startoftranscript|><|" + language + "|><|transcribe|><|notimestamps|>"; return ""; } diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 279d1be02..9475f4293 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -122,7 +122,32 @@ enum class OpType { SCALAR_NOT_EQUAL, RECURRENT_CACHE_STATE, RECURRENT_CACHE_WRITE, - CONV_CACHE_INITIALIZE + CONV_CACHE_INITIALIZE, + EQUAL, + LESS, + LESS_EQUAL, + GREATER, + GREATER_EQUAL, + SCALAR_EQUAL, + SCALAR_LESS, + SCALAR_LESS_EQUAL, + SCALAR_GREATER, + SCALAR_GREATER_EQUAL, + LOGICAL_NOT, + BITWISE_AND, + BITWISE_OR, + BITWISE_NOT, + WHERE, + UNFOLD, + PAD, + SCALAR_FLOOR_DIVIDE, + STRIDED_SLICE, + EXPAND, + MASKED_SELECT_PREFIX, + QKV_TQ_FUSED, + PROJECTION_PAIR_TQ_FUSED, + CONV1D_CAUSAL_CHANNEL_FIRST, + LOGITS_TQ_SOFTCAP }; struct PrecisionTraits { @@ -232,6 +257,7 @@ struct BufferDesc { size_t total_size; size_t byte_size; std::unique_ptr data; + std::shared_ptr shared_data; void* external_data; char* pooled_data; Precision precision; @@ -317,6 +343,18 @@ struct BufferDesc { void resize_from_pool(BufferPool& pool); }; + +struct TensorStorage { + std::shared_ptr data; + std::vector shape; + size_t total_size = 0; + size_t byte_size = 0; + Precision precision = Precision::FP16; + + void* get_data() const { return data.get(); } + bool valid() const { return data != nullptr && byte_size > 0; } +}; + struct OpParams { float scalar = 0.0f; float scale = 1.0f; @@ -504,6 +542,9 @@ class CactusGraph { size_t input(const std::vector& shape, Precision precision = Precision::FP16); void set_input(size_t node_id, const void* data, Precision precision); void set_external_input(size_t node_id, void* data, Precision precision); + std::shared_ptr export_tensor_storage(size_t node_id); + bool bind_tensor_storage(size_t node_id, const std::shared_ptr& storage, + bool require_exact_shape = true); void* get_output(size_t node_id); size_t add(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); @@ -512,18 +553,35 @@ class CactusGraph { size_t multiply(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); size_t divide(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); size_t not_equal(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t equal(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t less(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t less_equal(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t greater(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t greater_equal(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t bitwise_and(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t bitwise_or(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t where(size_t condition, size_t true_value, size_t false_value, ComputeBackend backend = cactus_default_backend()); + size_t masked_select_prefix(size_t input, size_t mask, ComputeBackend backend = cactus_default_backend()); size_t scalar_add(size_t input, float value, ComputeBackend backend = cactus_default_backend()); size_t scalar_subtract(size_t input, float value, ComputeBackend backend = cactus_default_backend()); size_t scalar_multiply(size_t input, float value, ComputeBackend backend = cactus_default_backend()); size_t scalar_divide(size_t input, float value, ComputeBackend backend = cactus_default_backend()); + size_t scalar_floor_divide(size_t input, float value, ComputeBackend backend = cactus_default_backend()); size_t scalar_not_equal(size_t input, float value, ComputeBackend backend = cactus_default_backend()); + size_t scalar_equal(size_t input, float value, ComputeBackend backend = cactus_default_backend()); + size_t scalar_less(size_t input, float value, ComputeBackend backend = cactus_default_backend()); + size_t scalar_less_equal(size_t input, float value, ComputeBackend backend = cactus_default_backend()); + size_t scalar_greater(size_t input, float value, ComputeBackend backend = cactus_default_backend()); + size_t scalar_greater_equal(size_t input, float value, ComputeBackend backend = cactus_default_backend()); size_t scalar_exp(size_t input, ComputeBackend backend = cactus_default_backend()); size_t scalar_sqrt(size_t input, ComputeBackend backend = cactus_default_backend()); size_t scalar_cos(size_t input, ComputeBackend backend = cactus_default_backend()); size_t scalar_sin(size_t input, ComputeBackend backend = cactus_default_backend()); size_t scalar_log(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t logical_not(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t bitwise_not(size_t input, ComputeBackend backend = cactus_default_backend()); size_t abs(size_t input, ComputeBackend backend = cactus_default_backend()); size_t pow(size_t input, float exponent, ComputeBackend backend = cactus_default_backend()); @@ -553,11 +611,15 @@ class CactusGraph { size_t reshape(size_t input, const std::vector& new_shape, ComputeBackend backend = cactus_default_backend()); size_t view(size_t input, const std::vector& new_shape, ComputeBackend backend = cactus_default_backend()); + size_t expand(size_t input, const std::vector& new_shape, ComputeBackend backend = cactus_default_backend()); size_t flatten(size_t input, int start_dim = 0, int end_dim = -1, ComputeBackend backend = cactus_default_backend()); size_t transpose(size_t input, ComputeBackend backend = cactus_default_backend()); size_t transposeN(size_t input, const std::vector& permutation, ComputeBackend backend = cactus_default_backend()); size_t slice(size_t input, int axis, size_t start, size_t length, ComputeBackend backend = cactus_default_backend()); + size_t strided_slice(size_t input, int axis, size_t start, size_t length, size_t step, ComputeBackend backend = cactus_default_backend()); size_t index(size_t input, size_t index_value, int dim, ComputeBackend backend = cactus_default_backend()); + size_t unfold(size_t input, int dimension, size_t size, size_t step, ComputeBackend backend = cactus_default_backend()); + size_t pad(size_t input, const std::vector& pads, float value = 0.0f, ComputeBackend backend = cactus_default_backend()); size_t concat(size_t input1, size_t input2, int axis = 0, ComputeBackend backend = cactus_default_backend()); size_t cat(const std::vector& inputs, int axis, ComputeBackend backend = cactus_default_backend()); @@ -625,6 +687,9 @@ class CactusGraph { size_t window_size = 0, size_t v_head_dim = 0, size_t cache_slot = 0, + size_t mask = static_cast(-1), + bool additive_mask = false, + bool is_causal = true, ComputeBackend backend = cactus_default_backend()); size_t conv_cache_state(size_t window_size, size_t hidden_dim, ComputeBackend backend = cactus_default_backend()); @@ -635,10 +700,13 @@ class CactusGraph { size_t recurrent_cache_write(size_t new_value, size_t cache_state, ComputeBackend backend = cactus_default_backend()); size_t conv1d_causal(size_t input, size_t weight, size_t kernel_size, size_t dilation = 1, ComputeBackend backend = cactus_default_backend()); + size_t conv1d_causal_channel_first(size_t input, size_t weight, size_t kernel_size, size_t dilation = 1, ComputeBackend backend = cactus_default_backend()); size_t conv1d_k3(size_t input, size_t weight, size_t stride, ComputeBackend backend = cactus_default_backend()); size_t conv1d_k7s3(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); size_t conv1d(size_t input, size_t weight, size_t stride, ComputeBackend backend = cactus_default_backend()); size_t conv1d(size_t input, size_t weight, size_t bias, size_t stride, ComputeBackend backend = cactus_default_backend()); + size_t conv1d_depthwise(size_t input, size_t weight, size_t stride = 1, ComputeBackend backend = cactus_default_backend()); + size_t conv1d_depthwise(size_t input, size_t weight, size_t bias, size_t stride, ComputeBackend backend = cactus_default_backend()); size_t conv1d_same_depthwise_k9(size_t input, size_t weight, ComputeBackend backend = cactus_default_backend()); size_t conv1d_same_depthwise_k9(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); size_t conv1d_pointwise(size_t input, size_t weight, ComputeBackend backend = cactus_default_backend()); @@ -715,7 +783,10 @@ class CactusGraph { size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, Activation activation, ComputeBackend backend = cactus_default_backend()); - size_t dense_mlp_tq_fused(size_t hidden, size_t gate_weight, size_t up_weight, size_t down_weight, float product_scale = 1.0f, ComputeBackend backend = cactus_default_backend()); + size_t dense_mlp_tq_fused(size_t hidden, size_t gate_weight, size_t up_weight, size_t down_weight, float product_scale = 1.0f, float gate_input_scale = 1.0f, ComputeBackend backend = cactus_default_backend()); + size_t qkv_tq_fused(size_t hidden, size_t query_weight, size_t key_weight, size_t value_weight, ComputeBackend backend = cactus_default_backend()); + size_t projection_pair_tq_fused(size_t hidden, size_t first_weight, size_t second_weight, ComputeBackend backend = cactus_default_backend()); + size_t logits_tq_softcap(size_t hidden, size_t weight, float cap, float projection_scale = 1.0f, ComputeBackend backend = cactus_default_backend()); size_t stats_pool(size_t input, ComputeBackend backend = cactus_default_backend()); size_t weighted_stats_pool(size_t input, size_t weights, ComputeBackend backend = cactus_default_backend()); @@ -729,7 +800,7 @@ class CactusGraph { ComputeBackend backend = cactus_default_backend()); size_t scatter_topk(size_t indices, size_t values, size_t num_classes, ComputeBackend backend = cactus_default_backend()); - size_t gather(size_t embeddings, size_t indices, ComputeBackend backend = cactus_default_backend()); + size_t gather(size_t tensor, size_t indices, int axis = 0, ComputeBackend backend = cactus_default_backend()); size_t embedding(const std::string& filename, size_t indices, ComputeBackend backend = cactus_default_backend()); size_t embedding(size_t embedding_tensor, size_t indices, ComputeBackend backend = cactus_default_backend()); size_t mmap_embeddings(const std::string& filename); @@ -749,6 +820,7 @@ class CactusGraph { void invalidate_persistent(size_t persistent_node_id); void execute(const std::string& profile_file = ""); + void set_profile_label(std::string label) { profile_label_ = std::move(label); } bool extract_ple_pathway(FusedEmbedCtx& ctx) const; void hard_reset(); void soft_reset(); @@ -764,6 +836,7 @@ class CactusGraph { const std::vector& output_shape, const OpParams& params = {}); const BufferDesc& get_output_buffer(size_t node_id) const; OpType get_node_op_type(size_t node_id) const; + const std::vector& get_node_inputs(size_t node_id) const; size_t get_node_window_size(size_t node_id) const; size_t get_node_sink_size(size_t node_id) const; size_t get_node_cache_num_slots(size_t node_id) const; @@ -803,6 +876,7 @@ class CactusGraph { std::unordered_set populated_node_ids_; std::unordered_set embedded_input_node_ids_; std::unordered_set retained_output_node_ids_; + std::string profile_label_; void build_metal_retype_plan(); void invalidate_metal_state(); std::unordered_map metal_plans_; @@ -956,17 +1030,43 @@ CACTUS_FFI_EXPORT int cactus_graph_divide(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_not_equal(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_equal(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_less(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_less_equal(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_greater(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_greater_equal(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_bitwise_and(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_bitwise_or(cactus_graph_t graph, cactus_node_t +a, cactus_node_t b, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_where(cactus_graph_t graph, cactus_node_t +condition, cactus_node_t true_value, cactus_node_t false_value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_masked_select_prefix(cactus_graph_t graph, cactus_node_t +x, cactus_node_t mask, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_add(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_subtract(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_multiply(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_divide(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_floor_divide(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_not_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_less(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_less_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_greater(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_scalar_greater_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_exp(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_sqrt(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_cos(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_sin(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scalar_log(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_logical_not(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_bitwise_not(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_abs(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); @@ -981,14 +1081,22 @@ CACTUS_FFI_EXPORT int cactus_graph_flatten( cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_reshape( cactus_graph_t graph, cactus_node_t x, const size_t* shape, size_t rank, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_expand( + cactus_graph_t graph, cactus_node_t x, const size_t* shape, size_t rank, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_transpose( cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_transpose_n( cactus_graph_t graph, cactus_node_t x, const size_t* permutation, size_t rank, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_slice( cactus_graph_t graph, cactus_node_t x, int32_t axis, size_t start, size_t length, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_strided_slice( + cactus_graph_t graph, cactus_node_t x, int32_t axis, size_t start, size_t length, size_t step, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_index( cactus_graph_t graph, cactus_node_t x, size_t index_value, int32_t dim, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_unfold( + cactus_graph_t graph, cactus_node_t x, int32_t dimension, size_t size, size_t step, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_pad( + cactus_graph_t graph, cactus_node_t x, const size_t* pads, size_t pad_count, float value, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_sum(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_mean(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_variance(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); @@ -1007,6 +1115,8 @@ CACTUS_FFI_EXPORT int cactus_graph_matmul( cactus_graph_t graph, cactus_node_t a, cactus_node_t b, bool pretransposed_rhs, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_gather( cactus_graph_t graph, cactus_node_t tensor, cactus_node_t indices, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_gather_dim( + cactus_graph_t graph, cactus_node_t tensor, cactus_node_t indices, int32_t axis, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_embedding_from_tensor( cactus_graph_t graph, cactus_node_t embedding_tensor, cactus_node_t indices, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_embedding_from_file( @@ -1033,7 +1143,7 @@ CACTUS_FFI_EXPORT int cactus_graph_glu(cactus_graph_t graph, cactus_node_t x, in CACTUS_FFI_EXPORT int cactus_graph_clamp(cactus_graph_t graph, cactus_node_t input, float lo, float hi, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_dense_mlp_tq_fused( cactus_graph_t graph, cactus_node_t hidden, cactus_node_t gate_weight, cactus_node_t up_weight, - cactus_node_t down_weight, float product_scale, cactus_node_t* out); + cactus_node_t down_weight, float product_scale, float gate_input_scale, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_layernorm( cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, cactus_node_t bias, float epsilon, bool has_bias, cactus_node_t* out); @@ -1066,6 +1176,11 @@ CACTUS_FFI_EXPORT int cactus_graph_attention_cached( cactus_graph_t graph, cactus_node_t query, cactus_node_t key_new, cactus_node_t value_new, cactus_node_t k_cache_state, cactus_node_t v_cache_state, float scale, size_t position_offset, size_t window_size, size_t v_head_dim, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_attention_cached_masked( + cactus_graph_t graph, cactus_node_t query, cactus_node_t key_new, cactus_node_t value_new, + cactus_node_t k_cache_state, cactus_node_t v_cache_state, + float scale, size_t position_offset, size_t window_size, size_t v_head_dim, + bool is_causal, cactus_node_t mask, bool additive_mask, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_conv_cache_state( cactus_graph_t graph, size_t window_size, size_t hidden_dim, cactus_node_t* out); @@ -1102,12 +1217,16 @@ CACTUS_FFI_EXPORT int cactus_graph_image_preprocess( CACTUS_FFI_EXPORT int cactus_graph_conv1d_causal( cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, size_t kernel_size, size_t dilation, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv1d_causal_channel_first( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, size_t kernel_size, size_t dilation, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_conv1d_k3( cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, size_t stride, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_conv1d_k7s3( cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, cactus_node_t bias, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_conv1d( cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, size_t stride, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_conv1d_depthwise( + cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, size_t stride, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_conv1d_same_depthwise_k9( cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_conv1d_pointwise( @@ -1146,6 +1265,9 @@ CACTUS_FFI_EXPORT int cactus_graph_moe_layer_ungated( CACTUS_FFI_EXPORT int cactus_graph_sample( cactus_graph_t graph, cactus_node_t logits, float temperature, float top_p, size_t top_k, cactus_node_t* out); +CACTUS_FFI_EXPORT int cactus_graph_logits_tq_softcap( + cactus_graph_t graph, cactus_node_t hidden, cactus_node_t weight, + float cap, float projection_scale, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_scatter_topk( cactus_graph_t graph, cactus_node_t indices, cactus_node_t values, size_t num_classes, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_persistent( @@ -1156,6 +1278,12 @@ CACTUS_FFI_EXPORT int cactus_graph_invalidate_persistent( cactus_graph_t graph, cactus_node_t persistent_node); CACTUS_FFI_EXPORT int cactus_graph_execute(cactus_graph_t graph); +CACTUS_FFI_EXPORT int cactus_graph_retain_outputs(cactus_graph_t graph, +const cactus_node_t* nodes, size_t count); +CACTUS_FFI_EXPORT int cactus_graph_get_node_op_type(cactus_graph_t graph, +cactus_node_t node, int32_t* out_op_type); +CACTUS_FFI_EXPORT int cactus_graph_get_node_inputs(cactus_graph_t graph, +cactus_node_t node, cactus_node_t* out_inputs, size_t* inout_count); CACTUS_FFI_EXPORT int cactus_graph_get_output_ptr(cactus_graph_t graph, cactus_node_t node, void** out_ptr); CACTUS_FFI_EXPORT int cactus_graph_get_output_info(cactus_graph_t graph, diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp index cc52bda29..0dd3796e4 100644 --- a/cactus-graph/src/builder.cpp +++ b/cactus-graph/src/builder.cpp @@ -13,6 +13,25 @@ bool use_fp16_kv_cache_for_builder() { return value != nullptr && std::strcmp(value, "1") == 0; } +bool op_honors_explicit_output_precision(OpType op_type) { + switch (op_type) { + case OpType::PRECISION_CAST: + case OpType::EMBEDDING: + case OpType::TOPK: + case OpType::SCATTER_TOPK: + case OpType::SAMPLE: + case OpType::KV_CACHE_APPEND: + case OpType::CONV_CACHE_STATE: + case OpType::RECURRENT_CACHE_STATE: + case OpType::IMAGE_PREPROCESS: + case OpType::MEL_FILTER_BANK: + case OpType::SPECTROGRAM: + return true; + default: + return false; + } +} + } // namespace size_t CactusGraph::attach_conv_bias(size_t node, size_t bias, size_t expected_size, const char* op_name) { @@ -343,7 +362,9 @@ size_t CactusGraph::moe_layer(size_t hidden, return tag_backend(add_node(OpType::MOE_LAYER, input_ids, hidden_buffer.shape, params), backend); } -size_t CactusGraph::dense_mlp_tq_fused(size_t hidden, size_t gate_weight, size_t up_weight, size_t down_weight, float product_scale, ComputeBackend backend) { +size_t CactusGraph::dense_mlp_tq_fused(size_t hidden, size_t gate_weight, size_t up_weight, + size_t down_weight, float product_scale, + float gate_input_scale, ComputeBackend backend) { const auto& hidden_buffer = get_output_buffer(hidden); const auto& down_buffer = get_output_buffer(down_weight); if (hidden_buffer.shape.empty()) { @@ -359,11 +380,75 @@ size_t CactusGraph::dense_mlp_tq_fused(size_t hidden, size_t gate_weight, size_t OpParams params; params.output_precision = Precision::FP16; params.scalar = product_scale; + params.scale = gate_input_scale; return tag_backend(add_node(OpType::DENSE_MLP_TQ_FUSED, {hidden, gate_weight, up_weight, down_weight}, output_shape, params), backend); } +size_t CactusGraph::qkv_tq_fused(size_t hidden, size_t query_weight, size_t key_weight, + size_t value_weight, ComputeBackend backend) { + const auto& hidden_buffer = get_output_buffer(hidden); + const auto& query_buffer = get_output_buffer(query_weight); + const auto& key_buffer = get_output_buffer(key_weight); + const auto& value_buffer = get_output_buffer(value_weight); + if (hidden_buffer.shape.empty() || query_buffer.shape.size() != 2 || + key_buffer.shape.size() != 2 || value_buffer.shape.size() != 2) { + throw std::runtime_error("qkv_tq_fused expects non-scalar hidden and rank-2 weights"); + } + const size_t hidden_dim = hidden_buffer.shape.back(); + if (query_buffer.shape[1] != hidden_dim || key_buffer.shape[1] != hidden_dim || + value_buffer.shape[1] != hidden_dim) { + throw std::runtime_error("qkv_tq_fused weight dimensions do not match hidden"); + } + std::vector output_shape = hidden_buffer.shape; + output_shape.back() = query_buffer.shape[0] + key_buffer.shape[0] + value_buffer.shape[0]; + OpParams params; + params.output_precision = Precision::FP16; + return tag_backend(add_node(OpType::QKV_TQ_FUSED, + {hidden, query_weight, key_weight, value_weight}, output_shape, params), backend); +} + +size_t CactusGraph::projection_pair_tq_fused(size_t hidden, size_t first_weight, + size_t second_weight, ComputeBackend backend) { + const auto& hidden_buffer = get_output_buffer(hidden); + const auto& first_buffer = get_output_buffer(first_weight); + const auto& second_buffer = get_output_buffer(second_weight); + if (hidden_buffer.shape.empty() || first_buffer.shape.size() != 2 || second_buffer.shape.size() != 2) { + throw std::runtime_error("projection_pair_tq_fused expects non-scalar hidden and rank-2 weights"); + } + const size_t hidden_dim = hidden_buffer.shape.back(); + if (first_buffer.shape[1] != hidden_dim || second_buffer.shape[1] != hidden_dim) { + throw std::runtime_error("projection_pair_tq_fused weight dimensions do not match hidden"); + } + std::vector output_shape = hidden_buffer.shape; + output_shape.back() = first_buffer.shape[0] + second_buffer.shape[0]; + OpParams params; + params.output_precision = Precision::FP16; + return tag_backend(add_node(OpType::PROJECTION_PAIR_TQ_FUSED, + {hidden, first_weight, second_weight}, output_shape, params), backend); +} + +size_t CactusGraph::logits_tq_softcap(size_t hidden, size_t weight, float cap, + float projection_scale, ComputeBackend backend) { + const auto& hidden_buffer = get_output_buffer(hidden); + const auto& weight_buffer = get_output_buffer(weight); + if (hidden_buffer.shape.empty() || weight_buffer.shape.size() != 2 || + !PrecisionTraits::is_cq(weight_buffer.precision) || !(cap > 0.0f)) { + throw std::runtime_error("logits_tq_softcap expects hidden, rank-2 CQ weight, and positive cap"); + } + if (weight_buffer.shape[1] != hidden_buffer.shape.back()) { + throw std::runtime_error("logits_tq_softcap weight dimensions do not match hidden"); + } + std::vector output_shape = hidden_buffer.shape; + output_shape.back() = weight_buffer.shape[0]; + OpParams params; + params.output_precision = Precision::FP16; + params.scalar = cap; + params.scale = projection_scale; + return tag_backend(add_node(OpType::LOGITS_TQ_SOFTCAP, {hidden, weight}, output_shape, params), backend); +} + size_t CactusGraph::moe_layer(size_t hidden, size_t routing_probs, size_t topk_indices, @@ -1138,7 +1223,7 @@ size_t CactusGraph::rope_gptj(size_t input, float theta, size_t position_offset, return add_node(OpType::ROPE_GPTJ, {input}, {}, params); } -size_t CactusGraph::gather(size_t tensor, size_t indices, ComputeBackend backend) { +size_t CactusGraph::gather(size_t tensor, size_t indices, int axis, ComputeBackend backend) { const auto& tensor_buffer = get_output_buffer(tensor); const auto& idx_shape = get_output_buffer(indices).shape; @@ -1146,12 +1231,25 @@ size_t CactusGraph::gather(size_t tensor, size_t indices, ComputeBackend backend throw std::runtime_error("Cannot gather from scalar tensor"); } - std::vector output_shape = idx_shape; - for (size_t i = 1; i < tensor_buffer.shape.size(); i++) { - output_shape.push_back(tensor_buffer.shape[i]); + int actual_axis = axis; + if (actual_axis < 0) actual_axis += static_cast(tensor_buffer.shape.size()); + if (actual_axis < 0 || static_cast(actual_axis) >= tensor_buffer.shape.size()) { + throw std::runtime_error("gather axis out of range"); + } + + std::vector output_shape; + + if (idx_shape.size() == tensor_buffer.shape.size()) { + output_shape = idx_shape; + } else { + output_shape = idx_shape; + for (size_t i = 1; i < tensor_buffer.shape.size(); i++) { + output_shape.push_back(tensor_buffer.shape[i]); + } } OpParams params; + params.axis = actual_axis; params.output_precision = tensor_buffer.precision; return tag_backend(add_node(OpType::GATHER, {tensor, indices}, output_shape, params), backend); @@ -1261,6 +1359,10 @@ OpType CactusGraph::get_node_op_type(size_t node_id) const { return nodes_[node_index_map_.at(node_id)]->op_type; } +const std::vector& CactusGraph::get_node_inputs(size_t node_id) const { + return nodes_[node_index_map_.at(node_id)]->input_ids; +} + size_t CactusGraph::get_node_window_size(size_t node_id) const { return nodes_[node_index_map_.at(node_id)]->params.window_size; } @@ -1656,6 +1758,7 @@ size_t CactusGraph::attention_cached(size_t query, size_t key_new, size_t value_ size_t k_cache_state, size_t v_cache_state, float scale, size_t position_offset, size_t window_size, size_t v_head_dim, size_t cache_slot, + size_t mask, bool additive_mask, bool is_causal, ComputeBackend backend) { const auto& q_shape = get_output_buffer(query).shape; size_t batch = q_shape[0]; @@ -1670,8 +1773,152 @@ size_t CactusGraph::attention_cached(size_t query, size_t key_new, size_t value_ params.window_size = window_size; params.v_head_dim = v_head_dim; params.cache_slot = cache_slot; + params.is_causal = is_causal; + params.attention_mask_is_additive = additive_mask; params.output_precision = Precision::FP16; + std::vector inputs = {query, key_new, value_new, k_cache_state, v_cache_state}; + if (mask != static_cast(-1)) inputs.push_back(mask); return tag_backend(add_node(OpType::ATTENTION_CACHED, - {query, key_new, value_new, k_cache_state, v_cache_state}, + inputs, {batch, seq_len, num_q_heads, out_v_dim}, params), backend); } + +size_t CactusGraph::equal(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::EQUAL, a, b), backend); } +size_t CactusGraph::less(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::LESS, a, b), backend); } +size_t CactusGraph::less_equal(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::LESS_EQUAL, a, b), backend); } +size_t CactusGraph::greater(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::GREATER, a, b), backend); } +size_t CactusGraph::greater_equal(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::GREATER_EQUAL, a, b), backend); } +size_t CactusGraph::bitwise_and(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::BITWISE_AND, a, b), backend); } +size_t CactusGraph::bitwise_or(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::BITWISE_OR, a, b), backend); } + +size_t CactusGraph::where(size_t condition, size_t true_value, size_t false_value, ComputeBackend backend) { + BroadcastInfo cond_true = BroadcastInfo::compute(get_output_buffer(condition).shape, + get_output_buffer(true_value).shape); + BroadcastInfo all_inputs = BroadcastInfo::compute(cond_true.output_shape, + get_output_buffer(false_value).shape); + OpParams params; + params.output_precision = get_output_buffer(true_value).precision; + return tag_backend(add_node(OpType::WHERE, {condition, true_value, false_value}, + all_inputs.output_shape, params), backend); +} + +size_t CactusGraph::masked_select_prefix(size_t input, size_t mask, ComputeBackend backend) { + const auto& in = get_output_buffer(input); + const auto& m = get_output_buffer(mask); + if ((in.shape.size() != 2 && in.shape.size() != 3) || + (m.shape.size() != 1 && m.shape.size() != 2)) + throw std::runtime_error("masked_select_prefix expects [T,D]/[1,T,D] and [T]/[1,T]"); + const size_t time = in.shape.size() == 3 ? in.shape[1] : in.shape[0]; + if (m.shape.back() != time) throw std::runtime_error("masked_select_prefix time mismatch"); + OpParams params{.output_precision = in.precision}; + return tag_backend(add_node(OpType::MASKED_SELECT_PREFIX, {input, mask}, + {time, in.shape.back()}, params), backend); +} + +size_t CactusGraph::expand(size_t input, const std::vector& new_shape, ComputeBackend backend) { + const auto& in = get_output_buffer(input); + if (new_shape.size() < in.shape.size()) throw std::runtime_error("expand target rank must be >= input rank"); + const size_t offset = new_shape.size() - in.shape.size(); + for (size_t i = 0; i < new_shape.size(); ++i) { + const size_t source = i < offset ? 1 : in.shape[i - offset]; + if (source != new_shape[i] && source != 1) + throw std::runtime_error("expand target shape is not broadcast-compatible"); + } + OpParams params{.new_shape = new_shape, .output_precision = in.precision}; + return tag_backend(add_node(OpType::EXPAND, {input}, new_shape, params), backend); +} + +size_t CactusGraph::scalar_floor_divide(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_FLOOR_DIVIDE, input, value), backend); } +size_t CactusGraph::scalar_equal(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_EQUAL, input, value), backend); } +size_t CactusGraph::scalar_less(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_LESS, input, value), backend); } +size_t CactusGraph::scalar_less_equal(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_LESS_EQUAL, input, value), backend); } +size_t CactusGraph::scalar_greater(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_GREATER, input, value), backend); } +size_t CactusGraph::scalar_greater_equal(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_GREATER_EQUAL, input, value), backend); } + +size_t CactusGraph::logical_not(size_t input, ComputeBackend backend) { + const auto& in = get_output_buffer(input); + return tag_backend(add_node(OpType::LOGICAL_NOT, {input}, in.shape, + {.output_precision = in.precision}), backend); +} + +size_t CactusGraph::bitwise_not(size_t input, ComputeBackend backend) { + const auto& in = get_output_buffer(input); + return tag_backend(add_node(OpType::BITWISE_NOT, {input}, in.shape, + {.output_precision = in.precision}), backend); +} + +size_t CactusGraph::unfold(size_t input, int dimension, size_t size, size_t step, ComputeBackend backend) { + const auto& in = get_output_buffer(input); + if (in.shape.empty() || step == 0) throw std::runtime_error("invalid unfold input or step"); + int dim = dimension < 0 ? dimension + static_cast(in.shape.size()) : dimension; + if (dim < 0 || static_cast(dim) >= in.shape.size() || size == 0 || size > in.shape[dim]) + throw std::runtime_error("unfold dimension or size out of range"); + std::vector out = in.shape; + out[dim] = ((in.shape[dim] - size) / step) + 1; + out.push_back(size); + OpParams params; + params.axis = dim; params.kernel_size = size; params.stride = step; + params.output_precision = in.precision; + return tag_backend(add_node(OpType::UNFOLD, {input}, out, params), backend); +} + +size_t CactusGraph::pad(size_t input, const std::vector& pads, float value, ComputeBackend backend) { + const auto& in = get_output_buffer(input); + if (in.shape.empty() || pads.size() % 2 || pads.size() / 2 > in.shape.size()) + throw std::runtime_error("invalid pad specification"); + std::vector out = in.shape; + for (size_t pair = 0; pair < pads.size(); pair += 2) + out[in.shape.size() - 1 - pair / 2] += pads[pair] + pads[pair + 1]; + OpParams params; + params.new_shape = pads; params.scalar = value; params.output_precision = in.precision; + return tag_backend(add_node(OpType::PAD, {input}, out, params), backend); +} + +size_t CactusGraph::strided_slice(size_t input, int axis, size_t start, size_t length, + size_t step, ComputeBackend backend) { + const auto& in = get_output_buffer(input); + if (in.shape.empty() || step == 0) throw std::runtime_error("invalid strided_slice input or step"); + int actual = axis < 0 ? axis + static_cast(in.shape.size()) : axis; + if (actual < 0 || static_cast(actual) >= in.shape.size()) + throw std::runtime_error("strided_slice axis out of range"); + const size_t extent = in.shape[actual]; + if (length == 0) length = start >= extent ? 0 : (extent - start + step - 1) / step; + if (length && start + (length - 1) * step >= extent) + throw std::runtime_error("strided_slice range out of bounds"); + std::vector out = in.shape; + out[actual] = length; + OpParams params; + params.axis = actual; params.slice_start = start; params.slice_length = length; + params.stride = step; params.output_precision = in.precision; + return tag_backend(add_node(OpType::STRIDED_SLICE, {input}, out, params), backend); +} + +size_t CactusGraph::conv1d_causal_channel_first(size_t input, size_t weight, size_t, + size_t dilation, ComputeBackend backend) { + return tag_backend(add_node(OpType::CONV1D_CAUSAL_CHANNEL_FIRST, {input, weight}, {}, + {.dilation = dilation}), backend); +} + +size_t CactusGraph::conv1d_depthwise(size_t input, size_t weight, size_t stride, ComputeBackend backend) { + const auto& in = get_output_buffer(input); const auto& w = get_output_buffer(weight); + if (in.shape.size() != 3 || w.shape.size() != 3 || w.shape[1] != 1 || w.shape[0] != in.shape[1]) + throw std::runtime_error("conv1d_depthwise shape mismatch"); + const size_t length = (in.shape[2] - w.shape[2]) / stride + 1; + OpParams params{.output_precision = in.precision, .stride = stride}; + params.num_groups = in.shape[1]; + return tag_backend(add_node(OpType::CONV1D, {input, weight}, {in.shape[0], w.shape[0], length}, params), backend); +} + +size_t CactusGraph::conv1d_depthwise(size_t input, size_t weight, size_t bias, + size_t stride, ComputeBackend backend) { + const auto& in = get_output_buffer(input); const auto& w = get_output_buffer(weight); + const auto& b = get_output_buffer(bias); + if (in.shape.size() != 3 || w.shape.size() != 3 || w.shape[1] != 1 || + w.shape[0] != in.shape[1] || b.total_size != w.shape[0]) + throw std::runtime_error("conv1d_depthwise shape or bias mismatch"); + const size_t length = (in.shape[2] - w.shape[2]) / stride + 1; + OpParams params{.output_precision = in.precision, .stride = stride}; + params.num_groups = in.shape[1]; + return tag_backend(add_node(OpType::CONV1D, {input, weight, bias}, + {in.shape[0], w.shape[0], length}, params), backend); +} diff --git a/cactus-graph/src/core.cpp b/cactus-graph/src/core.cpp index f850c99bc..2a6b33446 100644 --- a/cactus-graph/src/core.cpp +++ b/cactus-graph/src/core.cpp @@ -70,8 +70,8 @@ BufferDesc::BufferDesc(const std::vector& s, Precision prec) BufferDesc::~BufferDesc() { delete[] pooled_data; - - + + } BufferDesc::BufferDesc(BufferDesc&& other) noexcept @@ -79,6 +79,7 @@ BufferDesc::BufferDesc(BufferDesc&& other) noexcept total_size(other.total_size), byte_size(other.byte_size), data(std::move(other.data)), + shared_data(std::move(other.shared_data)), external_data(other.external_data), pooled_data(other.pooled_data), precision(other.precision), @@ -120,12 +121,13 @@ BufferDesc::BufferDesc(BufferDesc&& other) noexcept BufferDesc& BufferDesc::operator=(BufferDesc&& other) noexcept { if (this != &other) { - delete[] pooled_data; + delete[] pooled_data; shape = std::move(other.shape); total_size = other.total_size; byte_size = other.byte_size; data = std::move(other.data); + shared_data = std::move(other.shared_data); external_data = other.external_data; pooled_data = other.pooled_data; precision = other.precision; @@ -171,23 +173,25 @@ BufferDesc& BufferDesc::operator=(BufferDesc&& other) noexcept { void* BufferDesc::get_data() { if (external_data) return external_data; if (pooled_data) return pooled_data; + if (shared_data) return shared_data.get(); return data.get(); } const void* BufferDesc::get_data() const { if (external_data) return external_data; if (pooled_data) return pooled_data; + if (shared_data) return shared_data.get(); return data.get(); } void BufferDesc::allocate() { - if (!data && !external_data && !pooled_data) { + if (!data && !shared_data && !external_data && !pooled_data) { data = std::make_unique(byte_size); } } void BufferDesc::allocate_from_pool(BufferPool& pool) { - if (!data && !external_data && !pooled_data && byte_size > 0) { + if (!data && !shared_data && !external_data && !pooled_data && byte_size > 0) { pooled_data = pool.acquire(byte_size); pooled_byte_size = byte_size; } @@ -209,7 +213,7 @@ void BufferDesc::set_shape(const std::vector& new_shape) { } void BufferDesc::resize_from_pool(BufferPool& pool) { - if (data || external_data) return; + if (data || shared_data || external_data) return; if (pooled_data && byte_size == pooled_byte_size) return; if (pooled_data) release_to_pool(pool); if (byte_size > 0) { @@ -221,16 +225,19 @@ void BufferDesc::resize_from_pool(BufferPool& pool) { void BufferDesc::release_memory(BufferPool& pool) { release_to_pool(pool); data.reset(); + shared_data.reset(); external_data = nullptr; } void BufferDesc::set_external(void* ptr) { external_data = ptr; data.reset(); + shared_data.reset(); if (pooled_data) { delete[] pooled_data; } pooled_data = nullptr; + pooled_byte_size = 0; } GraphNode::GraphNode(size_t node_id, OpType type) : id(node_id), op_type(type) {} diff --git a/cactus-graph/src/execute.cpp b/cactus-graph/src/execute.cpp index fc19b2ccd..56830324e 100644 --- a/cactus-graph/src/execute.cpp +++ b/cactus-graph/src/execute.cpp @@ -40,6 +40,7 @@ using ComputeFn = void(*)(GraphNode&, const nodes_vector&, const node_index_map_ DECLARE_COMPUTE(compute_binary_op_node); DECLARE_COMPUTE(compute_unary_op_node); +DECLARE_COMPUTE(compute_where_node); DECLARE_COMPUTE(compute_activation_node); DECLARE_COMPUTE(compute_reduce_node); DECLARE_COMPUTE(compute_reshape_node); @@ -79,17 +80,22 @@ DECLARE_COMPUTE(compute_stats_pool_node); DECLARE_COMPUTE(compute_weighted_stats_pool_node); DECLARE_COMPUTE(compute_transpose_node); DECLARE_COMPUTE(compute_gather_node); +DECLARE_COMPUTE(compute_masked_select_prefix_node); DECLARE_COMPUTE(compute_slice_node); DECLARE_COMPUTE(compute_embedding_node); DECLARE_COMPUTE(compute_concat_node); DECLARE_COMPUTE(compute_cat_node); DECLARE_COMPUTE(compute_index_node); +DECLARE_COMPUTE(compute_unfold_node); DECLARE_COMPUTE(compute_bilinear_interpolation_node); DECLARE_COMPUTE(compute_sample_node); DECLARE_COMPUTE(compute_topk_node); DECLARE_COMPUTE(compute_scatter_topk_node); DECLARE_COMPUTE(compute_moe_layer_node); DECLARE_COMPUTE(compute_dense_mlp_tq_fused_node); +DECLARE_COMPUTE(compute_qkv_tq_fused_node); +DECLARE_COMPUTE(compute_projection_pair_tq_fused_node); +DECLARE_COMPUTE(compute_logits_tq_softcap_node); DECLARE_COMPUTE(compute_persistent_node); DECLARE_COMPUTE(compute_kv_cache_state_node); DECLARE_COMPUTE(compute_kv_cache_append_node); @@ -104,10 +110,14 @@ DECLARE_COMPUTE(compute_rfft_node); DECLARE_COMPUTE(compute_irfft_node); DECLARE_COMPUTE(compute_mel_filter_bank_node); DECLARE_COMPUTE(compute_spectrogram_node); +DECLARE_COMPUTE(compute_pad_node); +DECLARE_COMPUTE(compute_strided_slice_node); +DECLARE_COMPUTE(compute_expand_node); +DECLARE_COMPUTE(compute_conv1d_causal_channel_first_node); extern void shrink_thread_local_buffers(); #undef DECLARE_COMPUTE -static constexpr int OP_TYPE_COUNT = static_cast(OpType::CONV_CACHE_INITIALIZE) + 1; +static constexpr int OP_TYPE_COUNT = static_cast(OpType::LOGITS_TQ_SOFTCAP) + 1; static_assert(OP_TYPE_COUNT <= 256, "OpType dispatch table overflow"); static ComputeFn dispatch_flat[OP_TYPE_COUNT] = {}; @@ -118,11 +128,27 @@ static bool init_dispatch() { dispatch_flat[static_cast(OpType::MULTIPLY)] = compute_binary_op_node; dispatch_flat[static_cast(OpType::DIVIDE)] = compute_binary_op_node; dispatch_flat[static_cast(OpType::NOT_EQUAL)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::EQUAL)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::LESS)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::LESS_EQUAL)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::GREATER)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::GREATER_EQUAL)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::BITWISE_AND)] = compute_binary_op_node; + dispatch_flat[static_cast(OpType::BITWISE_OR)] = compute_binary_op_node; dispatch_flat[static_cast(OpType::SCALAR_ADD)] = compute_unary_op_node; dispatch_flat[static_cast(OpType::SCALAR_SUBTRACT)] = compute_unary_op_node; dispatch_flat[static_cast(OpType::SCALAR_MULTIPLY)] = compute_unary_op_node; dispatch_flat[static_cast(OpType::SCALAR_DIVIDE)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_FLOOR_DIVIDE)] = compute_unary_op_node; dispatch_flat[static_cast(OpType::SCALAR_NOT_EQUAL)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_EQUAL)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_LESS)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_LESS_EQUAL)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_GREATER)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::SCALAR_GREATER_EQUAL)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::LOGICAL_NOT)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::BITWISE_NOT)] = compute_unary_op_node; + dispatch_flat[static_cast(OpType::WHERE)] = compute_where_node; dispatch_flat[static_cast(OpType::SCALAR_EXP)] = compute_unary_op_node; dispatch_flat[static_cast(OpType::SCALAR_SQRT)] = compute_unary_op_node; dispatch_flat[static_cast(OpType::SCALAR_COS)] = compute_unary_op_node; @@ -147,6 +173,7 @@ static bool init_dispatch() { dispatch_flat[static_cast(OpType::FLATTEN)] = compute_reshape_node; dispatch_flat[static_cast(OpType::VIEW)] = compute_reshape_node; dispatch_flat[static_cast(OpType::RESHAPE)] = compute_reshape_node; + dispatch_flat[static_cast(OpType::EXPAND)] = compute_expand_node; dispatch_flat[static_cast(OpType::PRECISION_CAST)] = compute_precision_cast_node; dispatch_flat[static_cast(OpType::MATMUL)] = compute_matmul_node; dispatch_flat[static_cast(OpType::RMS_NORM)] = compute_rms_norm_node; @@ -160,6 +187,7 @@ static bool init_dispatch() { dispatch_flat[static_cast(OpType::ATTENTION_INT8_HYBRID)] = compute_attention_int8_hybrid_node; dispatch_flat[static_cast(OpType::REL_POS_BIAS)] = compute_rel_pos_bias_node; dispatch_flat[static_cast(OpType::CONV1D_CAUSAL)] = compute_conv1d_causal_node; + dispatch_flat[static_cast(OpType::CONV1D_CAUSAL_CHANNEL_FIRST)] = compute_conv1d_causal_channel_first_node; dispatch_flat[static_cast(OpType::CONV1D_K3)] = compute_conv1d_k3_node; dispatch_flat[static_cast(OpType::CONV1D_K7S3)] = compute_conv1d_k7s3_node; dispatch_flat[static_cast(OpType::CONV1D)] = compute_conv1d_node; @@ -172,17 +200,24 @@ static bool init_dispatch() { dispatch_flat[static_cast(OpType::GLU)] = compute_glu_node; dispatch_flat[static_cast(OpType::TRANSPOSE)] = compute_transpose_node; dispatch_flat[static_cast(OpType::GATHER)] = compute_gather_node; + dispatch_flat[static_cast(OpType::MASKED_SELECT_PREFIX)] = compute_masked_select_prefix_node; dispatch_flat[static_cast(OpType::SLICE)] = compute_slice_node; + dispatch_flat[static_cast(OpType::STRIDED_SLICE)] = compute_strided_slice_node; dispatch_flat[static_cast(OpType::EMBEDDING)] = compute_embedding_node; dispatch_flat[static_cast(OpType::CONCAT)] = compute_concat_node; dispatch_flat[static_cast(OpType::CAT)] = compute_cat_node; dispatch_flat[static_cast(OpType::INDEX)] = compute_index_node; + dispatch_flat[static_cast(OpType::UNFOLD)] = compute_unfold_node; + dispatch_flat[static_cast(OpType::PAD)] = compute_pad_node; dispatch_flat[static_cast(OpType::BILINEAR_INTERPOLATION)] = compute_bilinear_interpolation_node; dispatch_flat[static_cast(OpType::SAMPLE)] = compute_sample_node; dispatch_flat[static_cast(OpType::TOPK)] = compute_topk_node; dispatch_flat[static_cast(OpType::SCATTER_TOPK)] = compute_scatter_topk_node; dispatch_flat[static_cast(OpType::MOE_LAYER)] = compute_moe_layer_node; dispatch_flat[static_cast(OpType::DENSE_MLP_TQ_FUSED)] = compute_dense_mlp_tq_fused_node; + dispatch_flat[static_cast(OpType::QKV_TQ_FUSED)] = compute_qkv_tq_fused_node; + dispatch_flat[static_cast(OpType::PROJECTION_PAIR_TQ_FUSED)] = compute_projection_pair_tq_fused_node; + dispatch_flat[static_cast(OpType::LOGITS_TQ_SOFTCAP)] = compute_logits_tq_softcap_node; dispatch_flat[static_cast(OpType::PERSISTENT)] = compute_persistent_node; dispatch_flat[static_cast(OpType::LSTM_CELL)] = compute_lstm_cell_node; dispatch_flat[static_cast(OpType::GATED_DELTANET_DECODE)] = compute_gated_deltanet_decode_node; @@ -253,7 +288,30 @@ static const char* op_type_names[] = { "NOT_EQUAL", "SCALAR_NOT_EQUAL", "RECURRENT_CACHE_STATE", "RECURRENT_CACHE_WRITE", - "CONV_CACHE_INITIALIZE" + "CONV_CACHE_INITIALIZE", + "EQUAL", + "LESS", + "LESS_EQUAL", + "GREATER", + "GREATER_EQUAL", + "SCALAR_EQUAL", + "SCALAR_LESS", + "SCALAR_LESS_EQUAL", + "SCALAR_GREATER", + "SCALAR_GREATER_EQUAL", + "LOGICAL_NOT", + "BITWISE_AND", + "BITWISE_OR", + "BITWISE_NOT", + "WHERE", + "UNFOLD", + "PAD", + "SCALAR_FLOOR_DIVIDE", + "STRIDED_SLICE", + "EXPAND", + "MASKED_SELECT_PREFIX", + "QKV_TQ_FUSED", "PROJECTION_PAIR_TQ_FUSED", "CONV1D_CAUSAL_CHANNEL_FIRST", + "LOGITS_TQ_SOFTCAP" }; static const char* get_op_name(OpType op) { @@ -303,6 +361,61 @@ void CactusGraph::set_external_input(size_t node_id, void* data, Precision) { embedded_input_node_ids_.erase(node_id); } +std::shared_ptr CactusGraph::export_tensor_storage(size_t node_id) { + auto it = node_index_map_.find(node_id); + if (it == node_index_map_.end()) { + throw std::out_of_range("Unknown node id: " + std::to_string(node_id)); + } + + auto& buffer = nodes_[it->second]->output_buffer; + if (!buffer.get_data() || buffer.byte_size == 0) return {}; + + if (!buffer.shared_data) { + if (buffer.data) { + buffer.shared_data = std::shared_ptr(std::move(buffer.data)); + } else { + auto owned = std::shared_ptr(new char[buffer.byte_size], std::default_delete()); + std::memcpy(owned.get(), buffer.get_data(), buffer.byte_size); + buffer.shared_data = std::move(owned); + buffer.external_data = nullptr; + if (buffer.pooled_data) { + buffer.release_to_pool(buffer_pool_); + } + } + } + + auto storage = std::make_shared(); + storage->data = buffer.shared_data; + storage->shape = buffer.shape; + storage->total_size = buffer.total_size; + storage->byte_size = buffer.byte_size; + storage->precision = buffer.precision; + return storage; +} + +bool CactusGraph::bind_tensor_storage(size_t node_id, const std::shared_ptr& storage, bool require_exact_shape) { + auto it = node_index_map_.find(node_id); + if (it == node_index_map_.end() || !storage || !storage->valid()) return false; + + auto& node = *nodes_[it->second]; + auto& buffer = node.output_buffer; + const bool is_state = node.op_type == OpType::KV_CACHE_STATE + || node.op_type == OpType::CONV_CACHE_STATE + || node.op_type == OpType::RECURRENT_CACHE_STATE; + if (buffer.precision != storage->precision && !is_state) return false; + if (require_exact_shape && buffer.shape != storage->shape && !is_state) return false; + if (!is_state && buffer.byte_size != storage->byte_size) return false; + + buffer.release_memory(buffer_pool_); + buffer.shape = storage->shape; + buffer.total_size = storage->total_size; + buffer.byte_size = storage->byte_size; + buffer.precision = storage->precision; + buffer.shared_data = storage->data; + if (node.op_type == OpType::INPUT) embedded_input_node_ids_.erase(node_id); + return true; +} + void* CactusGraph::get_output(size_t node_id) { auto it = node_index_map_.find(node_id); if (it == node_index_map_.end()) { @@ -340,7 +453,40 @@ std::vector infer_output_shape(const GraphNode& node, const nodes_vector } case OpType::ADD: case OpType::ADD_CLIPPED: case OpType::SUBTRACT: case OpType::MULTIPLY: case OpType::DIVIDE: case OpType::NOT_EQUAL: + case OpType::EQUAL: case OpType::LESS: case OpType::LESS_EQUAL: + case OpType::GREATER: case OpType::GREATER_EQUAL: + case OpType::BITWISE_AND: case OpType::BITWISE_OR: return BroadcastInfo::compute(in(0), in(1)).output_shape; + case OpType::WHERE: { + auto first = BroadcastInfo::compute(in(0), in(1)).output_shape; + return BroadcastInfo::compute(first, in(2)).output_shape; + } + case OpType::EXPAND: + return node.params.new_shape; + case OpType::UNFOLD: { + std::vector out = in(0); + size_t axis = node.params.axis < 0 ? out.size() + static_cast(node.params.axis) + : static_cast(node.params.axis); + out[axis] = ((out[axis] - node.params.kernel_size) / node.params.stride) + 1; + out.push_back(node.params.kernel_size); + return out; + } + case OpType::STRIDED_SLICE: { + std::vector out = in(0); + size_t axis = node.params.axis < 0 ? out.size() + static_cast(node.params.axis) + : static_cast(node.params.axis); + out[axis] = node.params.slice_length; + return out; + } + case OpType::PAD: { + std::vector out = in(0); + const auto& pads = node.params.new_shape; + for (size_t pair = 0; pair + 1 < pads.size(); pair += 2) { + size_t axis = out.size() - 1 - (pair / 2); + out[axis] += pads[pair] + pads[pair + 1]; + } + return out; + } case OpType::ATTENTION: case OpType::ATTENTION_CACHED: case OpType::ATTENTION_INT8_HYBRID: { std::vector out = in(0); if (node.params.v_head_dim > 0) out.back() = node.params.v_head_dim; @@ -437,6 +583,9 @@ void CactusGraph::infer_shapes() { switch (node.op_type) { case OpType::ADD: case OpType::ADD_CLIPPED: case OpType::SUBTRACT: case OpType::MULTIPLY: case OpType::DIVIDE: case OpType::NOT_EQUAL: + case OpType::EQUAL: case OpType::LESS: case OpType::LESS_EQUAL: + case OpType::GREATER: case OpType::GREATER_EQUAL: + case OpType::BITWISE_AND: case OpType::BITWISE_OR: node.params.broadcast_info = BroadcastInfo::compute( get_input(node, 0, nodes_, node_index_map_).shape, get_input(node, 1, nodes_, node_index_map_).shape); @@ -1760,7 +1909,8 @@ void CactusGraph::execute(const std::string& profile_file) { continue; } if (preallocates_output(*node)) { - node->output_buffer.resize_from_pool(pool); + if (retained_output_node_ids_.count(node->id)) node->output_buffer.allocate(); + else node->output_buffer.resize_from_pool(pool); } run(*node); trace_nonfinite(i, *node); @@ -1818,11 +1968,14 @@ void CactusGraph::execute(const std::string& profile_file) { if (enable_profiling) { *out << "=== Graph Execution Profile ===" << std::endl; - *out << std::left << std::setw(24) << "Operation" + if (!profile_label_.empty()) *out << "Component: " << profile_label_ << std::endl; + *out << std::left << std::setw(32) << "Operation" + << std::setw(10) << "Node" << std::setw(12) << "Time (ms)" + << std::setw(14) << "Bytes" << std::setw(20) << "Output Shape" << "Backend" << std::endl; - *out << std::string(72, '-') << std::endl; + *out << std::string(96, '-') << std::endl; } for (size_t node_idx = 0; node_idx < n; ++node_idx) { @@ -1878,8 +2031,10 @@ void CactusGraph::execute(const std::string& profile_file) { } shape_str += "]"; - *out << std::left << std::setw(24) << get_op_name(node->op_type) + *out << std::left << std::setw(32) << get_op_name(node->op_type) + << std::setw(10) << node->id << std::setw(12) << std::fixed << std::setprecision(3) << ms + << std::setw(14) << node->output_buffer.byte_size << std::setw(20) << shape_str << std::endl; } else { dispatch_node(*node, nodes_, node_index_map_); @@ -2139,7 +2294,7 @@ void CactusGraph::execute(const std::string& profile_file) { auto total_duration = std::chrono::duration_cast(total_end - total_start); double total_ms = total_duration.count() / 1000.0; - *out << std::string(72, '-') << std::endl; + *out << std::string(96, '-') << std::endl; *out << "Total execution time: " << std::fixed << std::setprecision(3) << total_ms << " ms" << std::endl; *out << "================================" << std::endl; diff --git a/cactus-graph/src/graph_ffi.cpp b/cactus-graph/src/graph_ffi.cpp index 99cc069bb..223ce8c84 100644 --- a/cactus-graph/src/graph_ffi.cpp +++ b/cactus-graph/src/graph_ffi.cpp @@ -281,6 +281,110 @@ int cactus_graph_not_equal(cactus_graph_t graph, cactus_node_t a, cactus_node_t } } +CACTUS_FFI_EXPORT int cactus_graph_equal(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_equal"); + try { + *out = static_cast(as_graph(graph)->graph.equal(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_less(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_less"); + try { + *out = static_cast(as_graph(graph)->graph.less(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_less_equal(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_less_equal"); + try { + *out = static_cast(as_graph(graph)->graph.less_equal(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_greater(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_greater"); + try { + *out = static_cast(as_graph(graph)->graph.greater(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_greater_equal(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_greater_equal"); + try { + *out = static_cast(as_graph(graph)->graph.greater_equal(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_bitwise_and(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_bitwise_and"); + try { + *out = static_cast(as_graph(graph)->graph.bitwise_and(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_bitwise_or(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_bitwise_or"); + try { + *out = static_cast(as_graph(graph)->graph.bitwise_or(static_cast(a), static_cast(b))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_where(cactus_graph_t graph, cactus_node_t condition, cactus_node_t true_value, cactus_node_t false_value, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_where"); + try { + *out = static_cast(as_graph(graph)->graph.where( + static_cast(condition), + static_cast(true_value), + static_cast(false_value))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_masked_select_prefix(cactus_graph_t graph, cactus_node_t x, cactus_node_t mask, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_masked_select_prefix"); + try { + *out = static_cast(as_graph(graph)->graph.masked_select_prefix( + static_cast(x), + static_cast(mask))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_scalar_add(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { if (!graph || !out) { last_error_message = "Invalid args to cactus_graph_scalar_add"; @@ -334,6 +438,17 @@ int cactus_graph_scalar_divide(cactus_graph_t graph, cactus_node_t x, float valu } } +CACTUS_FFI_EXPORT int cactus_graph_scalar_floor_divide(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_floor_divide"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_floor_divide(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_scalar_not_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { if (!graph || !out) { last_error_message = "Invalid args to cactus_graph_scalar_not_equal"; @@ -348,6 +463,83 @@ int cactus_graph_scalar_not_equal(cactus_graph_t graph, cactus_node_t x, float v } } +CACTUS_FFI_EXPORT int cactus_graph_scalar_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_equal"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_equal(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_scalar_less(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_less"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_less(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_scalar_less_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_less_equal"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_less_equal(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_scalar_greater(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_greater"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_greater(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_scalar_greater_equal(cactus_graph_t graph, cactus_node_t x, float value, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_scalar_greater_equal"); + try { + *out = static_cast(as_graph(graph)->graph.scalar_greater_equal(static_cast(x), value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_logical_not(cactus_graph_t graph, cactus_node_t x, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_logical_not"); + try { + *out = static_cast(as_graph(graph)->graph.logical_not(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_bitwise_not(cactus_graph_t graph, cactus_node_t x, cactus_node_t *out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_bitwise_not"); + try { + *out = static_cast(as_graph(graph)->graph.bitwise_not(static_cast(x))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_scalar_exp(cactus_graph_t graph, cactus_node_t x, cactus_node_t *out) { if (!graph || !out) { last_error_message = "Invalid args to cactus_graph_scalar_exp"; @@ -501,6 +693,18 @@ int cactus_graph_reshape(cactus_graph_t graph, cactus_node_t x, const size_t* sh } } +CACTUS_FFI_EXPORT int cactus_graph_expand(cactus_graph_t graph, cactus_node_t x, const size_t* shape, size_t rank, cactus_node_t* out) { + if (!graph || !shape || rank == 0 || !out) return fail_invalid("Invalid args to cactus_graph_expand"); + try { + std::vector s(shape, shape + rank); + *out = static_cast(as_graph(graph)->graph.expand(static_cast(x), s)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_transpose(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_transpose"); try { @@ -535,6 +739,17 @@ int cactus_graph_slice(cactus_graph_t graph, cactus_node_t x, int32_t axis, size } } +CACTUS_FFI_EXPORT int cactus_graph_strided_slice(cactus_graph_t graph, cactus_node_t x, int32_t axis, size_t start, size_t length, size_t step, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_strided_slice"); + try { + *out = static_cast(as_graph(graph)->graph.strided_slice(static_cast(x), axis, start, length, step)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_index(cactus_graph_t graph, cactus_node_t x, size_t index_value, int32_t dim, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_index"); try { @@ -546,6 +761,29 @@ int cactus_graph_index(cactus_graph_t graph, cactus_node_t x, size_t index_value } } +CACTUS_FFI_EXPORT int cactus_graph_unfold(cactus_graph_t graph, cactus_node_t x, int32_t dimension, size_t size, size_t step, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_unfold"); + try { + *out = static_cast(as_graph(graph)->graph.unfold(static_cast(x), dimension, size, step)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_pad(cactus_graph_t graph, cactus_node_t x, const size_t* pads, size_t pad_count, float value, cactus_node_t* out) { + if (!graph || !pads || !out) return fail_invalid("Invalid args to cactus_graph_pad"); + try { + std::vector pad_values(pads, pads + pad_count); + *out = static_cast(as_graph(graph)->graph.pad(static_cast(x), pad_values, value)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_sum(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_sum"); try { @@ -667,6 +905,17 @@ int cactus_graph_gather(cactus_graph_t graph, cactus_node_t tensor, cactus_node_ } } +CACTUS_FFI_EXPORT int cactus_graph_gather_dim(cactus_graph_t graph, cactus_node_t tensor, cactus_node_t indices, int32_t axis, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_gather_dim"); + try { + *out = static_cast(as_graph(graph)->graph.gather(static_cast(tensor), static_cast(indices), axis)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_embedding_from_tensor(cactus_graph_t graph, cactus_node_t embedding_tensor, cactus_node_t indices, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_embedding_from_tensor"); try { @@ -1016,6 +1265,23 @@ int cactus_graph_attention_cached(cactus_graph_t graph, cactus_node_t query, cac } } +int cactus_graph_attention_cached_masked(cactus_graph_t graph, cactus_node_t query, cactus_node_t key_new, cactus_node_t value_new, + cactus_node_t k_cache_state, cactus_node_t v_cache_state, + float scale, size_t position_offset, size_t window_size, size_t v_head_dim, + bool is_causal, cactus_node_t mask, bool additive_mask, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_attention_cached_masked"); + try { + *out = static_cast(as_graph(graph)->graph.attention_cached( + static_cast(query), static_cast(key_new), static_cast(value_new), + static_cast(k_cache_state), static_cast(v_cache_state), + scale, position_offset, window_size, v_head_dim, 0, static_cast(mask), additive_mask, is_causal)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_conv_cache_state(cactus_graph_t graph, size_t window_size, size_t hidden_dim, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv_cache_state"); try { @@ -1158,6 +1424,18 @@ int cactus_graph_conv1d_causal(cactus_graph_t graph, cactus_node_t input, cactus } } +int cactus_graph_conv1d_causal_channel_first(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, size_t kernel_size, size_t dilation, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv1d_causal_channel_first"); + try { + *out = static_cast(as_graph(graph)->graph.conv1d_causal_channel_first( + static_cast(input), static_cast(weight), kernel_size, dilation)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_conv1d_k3(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, size_t stride, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv1d_k3"); try { @@ -1195,6 +1473,21 @@ int cactus_graph_conv1d(cactus_graph_t graph, cactus_node_t input, cactus_node_t } } +CACTUS_FFI_EXPORT int cactus_graph_conv1d_depthwise(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, size_t stride, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv1d_depthwise"); + try { + if (has_bias) { + *out = static_cast(as_graph(graph)->graph.conv1d_depthwise(static_cast(input), static_cast(weight), static_cast(bias), stride)); + } else { + *out = static_cast(as_graph(graph)->graph.conv1d_depthwise(static_cast(input), static_cast(weight), stride)); + } + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_conv1d_same_depthwise_k9(cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, bool has_bias, cactus_node_t bias, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_conv1d_same_depthwise_k9"); try { @@ -1381,7 +1674,7 @@ int cactus_graph_moe_layer_gated(cactus_graph_t graph, cactus_node_t hidden, cac } } -CACTUS_FFI_EXPORT int cactus_graph_dense_mlp_tq_fused(cactus_graph_t graph, cactus_node_t hidden, cactus_node_t gate_weight, cactus_node_t up_weight, cactus_node_t down_weight, float product_scale, cactus_node_t* out) { +CACTUS_FFI_EXPORT int cactus_graph_dense_mlp_tq_fused(cactus_graph_t graph, cactus_node_t hidden, cactus_node_t gate_weight, cactus_node_t up_weight, cactus_node_t down_weight, float product_scale, float gate_input_scale, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_dense_mlp_tq_fused"); try { *out = static_cast( @@ -1390,7 +1683,8 @@ CACTUS_FFI_EXPORT int cactus_graph_dense_mlp_tq_fused(cactus_graph_t graph, cact static_cast(gate_weight), static_cast(up_weight), static_cast(down_weight), - product_scale + product_scale, + gate_input_scale ) ); return 0; @@ -1400,6 +1694,43 @@ CACTUS_FFI_EXPORT int cactus_graph_dense_mlp_tq_fused(cactus_graph_t graph, cact } } +CACTUS_FFI_EXPORT int cactus_graph_qkv_tq_fused(cactus_graph_t graph, cactus_node_t hidden, cactus_node_t query_weight, cactus_node_t key_weight, cactus_node_t value_weight, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_qkv_tq_fused"); + try { + *out = static_cast(as_graph(graph)->graph.qkv_tq_fused( + static_cast(hidden), static_cast(query_weight), + static_cast(key_weight), static_cast(value_weight))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_projection_pair_tq_fused(cactus_graph_t graph, cactus_node_t hidden, cactus_node_t first_weight, cactus_node_t second_weight, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_projection_pair_tq_fused"); + try { + *out = static_cast(as_graph(graph)->graph.projection_pair_tq_fused( + static_cast(hidden), static_cast(first_weight), static_cast(second_weight))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +CACTUS_FFI_EXPORT int cactus_graph_logits_tq_softcap(cactus_graph_t graph, cactus_node_t hidden, cactus_node_t weight, float cap, float projection_scale, cactus_node_t* out) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_logits_tq_softcap"); + try { + *out = static_cast(as_graph(graph)->graph.logits_tq_softcap( + static_cast(hidden), static_cast(weight), cap, projection_scale)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_moe_layer_ungated(cactus_graph_t graph, cactus_node_t hidden, cactus_node_t routing_probs, cactus_node_t topk_indices, const cactus_node_t* w1_weights, const cactus_node_t* w2_weights, size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, int32_t activation, cactus_node_t* out) { @@ -1488,6 +1819,61 @@ int cactus_graph_execute(cactus_graph_t graph) { } } +int cactus_graph_retain_outputs(cactus_graph_t graph, const cactus_node_t* nodes, size_t count) { + if (!graph || (!nodes && count > 0)) { + last_error_message = "Invalid args to cactus_graph_retain_outputs"; + return -1; + } + try { + std::vector node_ids; + node_ids.reserve(count); + for (size_t i = 0; i < count; ++i) { + node_ids.push_back(static_cast(nodes[i])); + } + as_graph(graph)->graph.retain_outputs(node_ids); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_get_node_op_type(cactus_graph_t graph, cactus_node_t node, int32_t* out_op_type) { + if (!graph || !out_op_type) { + last_error_message = "Invalid args to cactus_graph_get_node_op_type"; + return -1; + } + try { + *out_op_type = static_cast(as_graph(graph)->graph.get_node_op_type(static_cast(node))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + +int cactus_graph_get_node_inputs(cactus_graph_t graph, cactus_node_t node, cactus_node_t* out_inputs, size_t* inout_count) { + if (!graph || !inout_count) { + last_error_message = "Invalid args to cactus_graph_get_node_inputs"; + return -1; + } + try { + const auto& inputs = as_graph(graph)->graph.get_node_inputs(static_cast(node)); + if (!out_inputs || *inout_count < inputs.size()) { + *inout_count = inputs.size(); + return out_inputs ? -2 : 0; + } + for (size_t i = 0; i < inputs.size(); ++i) { + out_inputs[i] = static_cast(inputs[i]); + } + *inout_count = inputs.size(); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_get_output_ptr(cactus_graph_t graph, cactus_node_t node, void** out_ptr) { if (!graph || !out_ptr) { last_error_message = "Invalid args to cactus_graph_get_output_ptr"; diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index 147eeb98b..2d90133a4 100644 --- a/cactus-graph/src/io.cpp +++ b/cactus-graph/src/io.cpp @@ -53,6 +53,36 @@ namespace { return filename; } + bool shape_matches_mmap_weight(const std::vector& logical_shape, + const std::vector& file_shape, + Precision precision, + const GraphFile::MappedFile& mapped_file) { + if (logical_shape == file_shape) { + return true; + } + + if (!PrecisionTraits::is_cq(precision) || mapped_file.group_size() == 0 || logical_shape.size() != file_shape.size()) { + return false; + } + + const size_t padded_k = mapped_file.group_size() * mapped_file.num_groups(); + size_t padded_dims = 0; + for (size_t i = 0; i < logical_shape.size(); ++i) { + if (logical_shape[i] == file_shape[i]) { + continue; + } + + if (file_shape[i] == padded_k && logical_shape[i] < file_shape[i]) { + ++padded_dims; + continue; + } + + return false; + } + + return padded_dims == 1; + } + inline void write_u32(std::ostream& out, uint32_t v) { out.write(reinterpret_cast(&v), sizeof(v)); } @@ -197,6 +227,13 @@ namespace { case OpType::MULTIPLY: case OpType::DIVIDE: case OpType::NOT_EQUAL: + case OpType::EQUAL: + case OpType::LESS: + case OpType::LESS_EQUAL: + case OpType::GREATER: + case OpType::GREATER_EQUAL: + case OpType::BITWISE_AND: + case OpType::BITWISE_OR: return true; default: return false; @@ -241,8 +278,10 @@ namespace { GraphFile::NodeEntry node; node.index = read_u32(in); uint32_t op_type_val = read_u32(in); - if (op_type_val > static_cast(OpType::CONV_CACHE_INITIALIZE)) { - throw std::runtime_error("Graph file corrupted: invalid op type"); + if (op_type_val > static_cast(OpType::LOGITS_TQ_SOFTCAP)) { + throw std::runtime_error( + "Graph file corrupted: invalid op type " + std::to_string(op_type_val) + + " > " + std::to_string(static_cast(OpType::LOGITS_TQ_SOFTCAP))); } node.op_type = static_cast(op_type_val); node.inputs = read_u32_vector(in); @@ -536,7 +575,7 @@ void CactusGraph::bind_mmap_weights(size_t node_id, const std::string& filename) const auto& shape = mapped_file->shape(); Precision precision = mapped_file->precision(); auto& buffer = node.output_buffer; - if (buffer.shape != shape) { + if (!shape_matches_mmap_weight(buffer.shape, shape, precision, *mapped_file)) { throw std::runtime_error("mmap weight shape mismatch for node " + std::to_string(node_id)); } if (buffer.precision != precision) { diff --git a/cactus-graph/src/ops_cache.cpp b/cactus-graph/src/ops_cache.cpp index ed9db8edd..ed628e42e 100644 --- a/cactus-graph/src/ops_cache.cpp +++ b/cactus-graph/src/ops_cache.cpp @@ -180,7 +180,9 @@ void compute_kv_cache_state_node( else max_seq = std::min(ceiling, kInitialCacheEntries); size_t kv_heads = node.params.num_kv_heads; size_t hdim = node.params.head_dim; - const bool fp16_cache = use_fp16_kv_cache(); + const bool fp16_cache = node.params.output_precision == Precision::FP16 || ( + node.params.output_precision != Precision::INT8 && use_fp16_kv_cache() + ); size_t per_slot = fp16_cache ? fp16_cache_elements(max_seq, kv_heads, hdim) : cache_buffer_size(max_seq, kv_heads, hdim); @@ -338,6 +340,7 @@ void compute_attention_cached_node( const auto& val_new_buf = get_input(node, 2, nodes, node_index_map); const auto& k_cache_buf = get_input(node, 3, nodes, node_index_map); const auto& v_cache_buf = get_input(node, 4, nodes, node_index_map); + const BufferDesc* mask_buf = node.input_ids.size() > 5 ? &get_input(node, 5, nodes, node_index_map) : nullptr; size_t slot = node.params.cache_slot; const auto* k_meta = get_meta(k_cache_buf, slot); @@ -360,6 +363,39 @@ void compute_attention_cached_node( size_t batch_size = q_shape[0]; size_t seq_len = q_shape[1]; size_t num_q_heads = q_shape[2]; + bool mask_per_head = false; + const __fp16* mask_ptr = nullptr; + size_t mask_batch_stride = 0; + + if (mask_buf != nullptr) { + if (mask_buf->precision != Precision::FP16) { + throw std::runtime_error("Cached attention mask tensor must be FP16"); + } + + if (mask_buf->shape.size() == 3) { + if (mask_buf->shape[0] != batch_size || + mask_buf->shape[1] != seq_len || + mask_buf->shape[2] != cache_len) { + throw std::runtime_error("Cached attention mask [B, T, S] shape mismatch"); + } + + mask_batch_stride = seq_len * cache_len; + } else if (mask_buf->shape.size() == 4) { + if (mask_buf->shape[0] != batch_size || + (mask_buf->shape[1] != num_q_heads && mask_buf->shape[1] != 1) || + mask_buf->shape[2] != seq_len || + mask_buf->shape[3] != cache_len) { + throw std::runtime_error("Cached attention mask [B, H|1, T, S] shape mismatch"); + } + + mask_per_head = mask_buf->shape[1] != 1; + mask_batch_stride = (mask_per_head ? num_q_heads : 1) * seq_len * cache_len; + } else { + throw std::runtime_error("Cached attention mask must be rank 3 or 4"); + } + + mask_ptr = mask_buf->data_as<__fp16>(); + } size_t new_seq_len = key_new_buf.total_size / (kv_heads * hdim); size_t history_len = (cache_len >= new_seq_len) ? cache_len - new_seq_len : 0; @@ -395,9 +431,17 @@ void compute_attention_cached_node( out_all + i * out_stride, 1, seq_len, ci, num_q_heads, kv_heads, hdim, - node.params.scale, nullptr, hist_i, node.params.window_size, - true, false, false, v_hdim); + node.params.scale, + mask_ptr ? mask_ptr + i * mask_batch_stride : nullptr, + hist_i, node.params.window_size, + node.params.is_causal, + node.params.attention_mask_is_additive, + mask_per_head, v_hdim); } else { + if (mask_ptr) { + throw std::runtime_error("Cached attention masks require FP16 KV cache"); + } + cactus_attention_hybrid_int8_fp16( q_all + i * q_stride, get_int8_data(k_cache_buf, i), @@ -425,16 +469,20 @@ void compute_attention_cached_node( batch_size, seq_len, cache_len, num_q_heads, kv_heads, hdim, node.params.scale, - nullptr, + mask_ptr, position_offset, node.params.window_size, - true, - false, - false, + node.params.is_causal, + node.params.attention_mask_is_additive, + mask_per_head, v_hdim); return; } + if (mask_ptr) { + throw std::runtime_error("Cached attention masks require FP16 KV cache"); + } + cactus_attention_hybrid_int8_fp16( query_buf.data_as<__fp16>(), cached_keys, @@ -754,4 +802,3 @@ void CactusGraph::shrink_cache_buffer(size_t node_id, size_t new_capacity) { if (target >= meta->max_seq_len) return; resize_cache_buffer(buf, target); } - diff --git a/cactus-graph/src/ops_conv.cpp b/cactus-graph/src/ops_conv.cpp index 20cf09030..8850123e7 100644 --- a/cactus-graph/src/ops_conv.cpp +++ b/cactus-graph/src/ops_conv.cpp @@ -113,6 +113,65 @@ void compute_conv1d_causal_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& X = get_input(node, 0, nodes, node_index_map); + const auto& W = get_input(node, 1, nodes, node_index_map); + auto& Y = node.output_buffer; + + if (X.shape.size() != 3) { + throw std::runtime_error("Channel-first causal conv requires 3D input [batch, channels, seq_len]"); + } + if (W.shape.size() != 3) { + throw std::runtime_error("Weight must be 3D"); + } + + const size_t N = X.shape[0]; + const size_t C = X.shape[1]; + const size_t L = X.shape[2]; + const size_t W0 = W.shape[0]; + const size_t W1 = W.shape[1]; + const size_t W2 = W.shape[2]; + const size_t dilation = node.params.dilation; + if (dilation < 1) throw std::runtime_error("dilation must be >= 1"); + + const bool standard_layout = W1 == 1; + const bool transposed_layout = W2 == 1; + if ((!standard_layout && !transposed_layout) || W0 != C) { + throw std::runtime_error("Only multiplier-1 depthwise channel-first causal convolution is supported"); + } + const size_t K = standard_layout ? W2 : W1; + Y.shape = {N, C, L}; + Y.precision = X.precision; + + auto transpose_weights = [&](const __fp16* source) { + std::vector<__fp16> transposed(C * K); + for (size_t channel = 0; channel < C; ++channel) { + for (size_t k = 0; k < K; ++k) { + transposed[channel * K + k] = source[(channel * W1 + k) * W2]; + } + } + return transposed; + }; + + if (W.precision == Precision::INT8) { + auto weights = dequantize_int8_weights_to_fp16(W, C, K, "conv1d_causal_channel_first"); + if (transposed_layout && !standard_layout) weights = transpose_weights(weights.data()); + cactus_conv1d_causal_depthwise_channel_first_f16( + X.data_as<__fp16>(), weights.data(), Y.data_as<__fp16>(), N, C, L, K, dilation); + } else if (W.precision == Precision::FP16) { + if (transposed_layout && !standard_layout) { + auto weights = transpose_weights(W.data_as<__fp16>()); + cactus_conv1d_causal_depthwise_channel_first_f16( + X.data_as<__fp16>(), weights.data(), Y.data_as<__fp16>(), N, C, L, K, dilation); + } else { + cactus_conv1d_causal_depthwise_channel_first_f16( + X.data_as<__fp16>(), W.data_as<__fp16>(), Y.data_as<__fp16>(), N, C, L, K, dilation); + } + } else { + throw std::runtime_error("Channel-first causal conv requires FP16 or INT8 weights"); + } +} + void compute_conv1d_k3_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { const auto& X = get_input(node, 0, nodes, node_index_map); const auto& W = get_input(node, 1, nodes, node_index_map); @@ -187,8 +246,9 @@ void compute_conv1d_node(GraphNode& node, const std::vector(); } else if (W.precision == Precision::INT8) { - W_fp16 = dequantize_int8_weights_to_fp16(W, C_out, C_in * K, "conv1d"); + W_fp16 = dequantize_int8_weights_to_fp16(W, C_out, (is_depthwise ? 1 : C_in) * K, "conv1d"); W_ptr = W_fp16.data(); } else { throw std::runtime_error("Conv1d only supports FP16/INT8 weights"); } + if (is_depthwise) { + __fp16* output = Y.data_as<__fp16>(); + const __fp16* input = X.data_as<__fp16>(); + const size_t L_out = Y.shape[2]; + + for (size_t n = 0; n < N; ++n) { + for (size_t c = 0; c < C_in; ++c) { + for (size_t t = 0; t < L_out; ++t) { + float acc = bias_ptr ? static_cast(bias_ptr[c]) : 0.0f; + + for (size_t k = 0; k < K; ++k) { + size_t input_idx = (n * C_in + c) * L + t * stride + k; + size_t weight_idx = c * K + k; + acc += static_cast(input[input_idx]) * static_cast(W_ptr[weight_idx]); + } + + output[(n * C_out + c) * L_out + t] = static_cast<__fp16>(acc); + } + } + } + return; + } + cactus_conv1d_f16(X.data_as<__fp16>(), W_ptr, bias_ptr, Y.data_as<__fp16>(), N, L, C_in, C_out, K, stride); } diff --git a/cactus-graph/src/ops_math.cpp b/cactus-graph/src/ops_math.cpp index e12661268..22d32a3b9 100644 --- a/cactus-graph/src/ops_math.cpp +++ b/cactus-graph/src/ops_math.cpp @@ -100,6 +100,33 @@ void dispatch_binary_op_f16(OpType op, const __fp16* lhs, const __fp16* rhs, __f } }); break; + case OpType::EQUAL: + case OpType::LESS: + case OpType::LESS_EQUAL: + case OpType::GREATER: + case OpType::GREATER_EQUAL: + case OpType::BITWISE_AND: + case OpType::BITWISE_OR: + CactusThreading::parallel_for(count, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t i = start_idx; i < end_idx; ++i) { + float a = static_cast(lhs[i]); + float b = static_cast(rhs[i]); + bool result = false; + switch (op) { + case OpType::EQUAL: result = (a == b); break; + case OpType::LESS: result = (a < b); break; + case OpType::LESS_EQUAL: result = (a <= b); break; + case OpType::GREATER: result = (a > b); break; + case OpType::GREATER_EQUAL: result = (a >= b); break; + case OpType::BITWISE_AND: result = (a != 0.0f && b != 0.0f); break; + case OpType::BITWISE_OR: result = (a != 0.0f || b != 0.0f); break; + default: break; + } + output[i] = static_cast<__fp16>(result ? 1.0f : 0.0f); + } + }); + break; default: break; } @@ -118,6 +145,20 @@ static float apply_binary_op_f32(OpType op, float lhs, float rhs) { return lhs / rhs; case OpType::NOT_EQUAL: return lhs != rhs ? 1.0f : 0.0f; + case OpType::EQUAL: + return lhs == rhs ? 1.0f : 0.0f; + case OpType::LESS: + return lhs < rhs ? 1.0f : 0.0f; + case OpType::LESS_EQUAL: + return lhs <= rhs ? 1.0f : 0.0f; + case OpType::GREATER: + return lhs > rhs ? 1.0f : 0.0f; + case OpType::GREATER_EQUAL: + return lhs >= rhs ? 1.0f : 0.0f; + case OpType::BITWISE_AND: + return (lhs != 0.0f && rhs != 0.0f) ? 1.0f : 0.0f; + case OpType::BITWISE_OR: + return (lhs != 0.0f || rhs != 0.0f) ? 1.0f : 0.0f; default: return lhs; } @@ -133,6 +174,16 @@ void dispatch_binary_op_f32(OpType op, const float* lhs, const float* rhs, float } void dispatch_unary_op_f16(OpType op, const __fp16* input, __fp16* output, size_t count, float param) { + if (op == OpType::SCALAR_FLOOR_DIVIDE) { + CactusThreading::parallel_for(count, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t i = start_idx; i < end_idx; ++i) { + output[i] = static_cast<__fp16>(std::floor(static_cast(input[i]) / param)); + } + }); + return; + } + ScalarOpType scalar_op; switch (op) { case OpType::SCALAR_ADD: scalar_op = ScalarOpType::ADD; break; @@ -158,6 +209,7 @@ static float apply_scalar_op_f32(OpType op, float value, float param) { case OpType::SCALAR_SUBTRACT: return value - param; case OpType::SCALAR_MULTIPLY: return value * param; case OpType::SCALAR_DIVIDE: return value / param; + case OpType::SCALAR_FLOOR_DIVIDE: return std::floor(value / param); case OpType::SCALAR_EXP: return std::exp(value); case OpType::SCALAR_SQRT: return std::sqrt(value); case OpType::SCALAR_COS: return std::cos(value); @@ -271,6 +323,35 @@ void compute_binary_op_node(GraphNode& node, const std::vector(); + const __fp16* rhs_data = rhs.data_as<__fp16>(); + __fp16* out_data = node.output_buffer.data_as<__fp16>(); + CactusThreading::parallel_for(total_elements, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + std::vector coords(output_shape.size()); + size_t tmp = start_idx; + for (int axis = static_cast(output_shape.size()) - 1; axis >= 0; --axis) { + coords[static_cast(axis)] = tmp % output_shape[static_cast(axis)]; + tmp /= output_shape[static_cast(axis)]; + } + for (size_t linear_idx = start_idx; linear_idx < end_idx; ++linear_idx) { + size_t lhs_idx = broadcast_linear_index(coords, lhs_strides); + size_t rhs_idx = broadcast_linear_index(coords, rhs_strides); + float a = static_cast(lhs_data[lhs_idx]); + float b = static_cast(rhs_data[rhs_idx]); + out_data[linear_idx] = static_cast<__fp16>(apply_binary_op_f32(node.op_type, a, b)); + increment_broadcast_coords(coords, output_shape); + } + }); + break; + } default: break; } } else { @@ -293,7 +374,14 @@ void compute_unary_op_node(GraphNode& node, const std::vector(); @@ -301,8 +389,22 @@ void compute_unary_op_node(GraphNode& node, const std::vector( - static_cast(input_data[i]) != scalar ? 1.0f : 0.0f); + float value = static_cast(input_data[i]); + bool result = false; + switch (node.op_type) { + case OpType::SCALAR_NOT_EQUAL: result = value != scalar; break; + case OpType::SCALAR_EQUAL: result = value == scalar; break; + case OpType::SCALAR_LESS: result = value < scalar; break; + case OpType::SCALAR_LESS_EQUAL: result = value <= scalar; break; + case OpType::SCALAR_GREATER: result = value > scalar; break; + case OpType::SCALAR_GREATER_EQUAL: result = value >= scalar; break; + case OpType::LOGICAL_NOT: + case OpType::BITWISE_NOT: + result = value == 0.0f; + break; + default: break; + } + output_data[i] = static_cast<__fp16>(result ? 1.0f : 0.0f); } }); } else { @@ -311,7 +413,22 @@ void compute_unary_op_node(GraphNode& node, const std::vector scalar; break; + case OpType::SCALAR_GREATER_EQUAL: result = value >= scalar; break; + case OpType::LOGICAL_NOT: + case OpType::BITWISE_NOT: + result = value == 0.0f; + break; + default: break; + } + output_data[i] = result ? 1.0f : 0.0f; } }); } @@ -329,11 +446,108 @@ void compute_unary_op_node(GraphNode& node, const std::vector()[index]; + if (buffer.precision == Precision::FP16) return static_cast(buffer.data_as<__fp16>()[index]); + if (buffer.precision == Precision::INT8) return static_cast(buffer.data_as()[index]); + throw std::runtime_error("where only supports FP32/FP16/INT8 condition data"); +} + +static void write_float_value(BufferDesc& buffer, size_t index, float value) { + if (buffer.precision == Precision::FP32) { + buffer.data_as()[index] = value; + return; + } + if (buffer.precision == Precision::FP16) { + buffer.data_as<__fp16>()[index] = static_cast<__fp16>(value); + return; + } + if (buffer.precision == Precision::INT8) { + buffer.data_as()[index] = static_cast(value); + return; + } + throw std::runtime_error("where only supports FP32/FP16/INT8 output data"); +} + +void compute_where_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { + const auto& condition = get_input(node, 0, nodes, node_index_map); + const auto& true_value = get_input(node, 1, nodes, node_index_map); + const auto& false_value = get_input(node, 2, nodes, node_index_map); + + if (true_value.precision != false_value.precision) { + throw std::runtime_error("where true/false branches must have matching precision"); + } + + const auto& output_shape = node.output_buffer.shape; + size_t total_elements = node.output_buffer.total_size; + std::vector condition_strides = compute_strides(condition.shape, output_shape); + std::vector true_strides = compute_strides(true_value.shape, output_shape); + std::vector false_strides = compute_strides(false_value.shape, output_shape); + + CactusThreading::parallel_for(total_elements, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + std::vector coords(output_shape.size()); + size_t tmp = start_idx; + for (int axis = static_cast(output_shape.size()) - 1; axis >= 0; --axis) { + coords[static_cast(axis)] = tmp % output_shape[static_cast(axis)]; + tmp /= output_shape[static_cast(axis)]; + } + for (size_t linear_idx = start_idx; linear_idx < end_idx; ++linear_idx) { + size_t cond_idx = broadcast_linear_index(coords, condition_strides); + size_t true_idx = broadcast_linear_index(coords, true_strides); + size_t false_idx = broadcast_linear_index(coords, false_strides); + float selected = read_value_as_float(condition, cond_idx) != 0.0f + ? read_value_as_float(true_value, true_idx) + : read_value_as_float(false_value, false_idx); + write_float_value(node.output_buffer, linear_idx, selected); + increment_broadcast_coords(coords, output_shape); + } + }); +} + void compute_activation_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { const auto& input = get_input(node, 0, nodes, node_index_map); - if (input.precision != Precision::FP16) { - throw std::runtime_error("Activation operations only support FP16 precision"); + if (input.precision != Precision::FP16 && input.precision != Precision::FP32) { + throw std::runtime_error("Activation operations only support FP16 or FP32 precision"); + } + + if (input.precision == Precision::FP32 && node.op_type != OpType::CLAMP) { + const float* in = input.data_as(); + float* out = node.output_buffer.data_as(); + + CactusThreading::parallel_for(node.output_buffer.total_size, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t i = start_idx; i < end_idx; ++i) { + const float x = in[i]; + switch (node.op_type) { + case OpType::RELU: + out[i] = std::max(x, 0.0f); + break; + case OpType::SILU: + out[i] = x / (1.0f + std::exp(-x)); + break; + case OpType::GELU: + out[i] = 0.5f * x * (1.0f + std::tanh(0.7978845608f * (x + 0.044715f * x * x * x))); + break; + case OpType::GELU_ERF: + out[i] = 0.5f * x * (1.0f + std::erf(x * 0.7071067812f)); + break; + case OpType::SIGMOID: + out[i] = 1.0f / (1.0f + std::exp(-x)); + break; + case OpType::TANH: + out[i] = std::tanh(x); + break; + case OpType::LEAKY_RELU: + out[i] = x >= 0.0f ? x : x * node.params.scalar; + break; + default: + throw std::runtime_error("Unsupported FP32 activation"); + } + } + }); + return; } switch (node.op_type) { @@ -373,10 +587,23 @@ void compute_activation_node(GraphNode& node, const std::vector(), - node.output_buffer.data_as<__fp16>(), - node.output_buffer.total_size, - node.params.scalar, node.params.scale); + if (input.precision == Precision::FP32) { + const float lo = node.params.scalar; + const float hi = node.params.scale; + const float* in = input.data_as(); + float* out = node.output_buffer.data_as(); + CactusThreading::parallel_for(node.output_buffer.total_size, CactusThreading::Thresholds::ELEMENT_WISE, + [&](size_t start_idx, size_t end_idx) { + for (size_t i = start_idx; i < end_idx; ++i) { + out[i] = std::min(std::max(in[i], lo), hi); + } + }); + } else { + cactus_clamp_f16(input.data_as<__fp16>(), + node.output_buffer.data_as<__fp16>(), + node.output_buffer.total_size, + node.params.scalar, node.params.scale); + } break; default: break; diff --git a/cactus-graph/src/ops_nn.cpp b/cactus-graph/src/ops_nn.cpp index 3d192f52e..e5eed8d89 100644 --- a/cactus-graph/src/ops_nn.cpp +++ b/cactus-graph/src/ops_nn.cpp @@ -86,6 +86,7 @@ namespace { thread_local std::vector moe_expert_offsets_buf; thread_local std::vector moe_expert_tokens_buf; thread_local std::vector moe_routing_denom_buf; + constexpr float MOE_FP16_MAX = 65504.0f; void ensure_moe_buffers(size_t max_tokens, size_t hidden_dim, size_t intermediate_dim, size_t num_experts, size_t top_k) { @@ -123,6 +124,30 @@ namespace { throw std::runtime_error("moe_layer only supports FP16 or TQ expert weights"); } + + inline __fp16 moe_finite_f16(float value) { + if (!std::isfinite(value)) { + value = std::copysign(MOE_FP16_MAX, value); + } + value = std::clamp(value, -MOE_FP16_MAX, MOE_FP16_MAX); + return static_cast<__fp16>(value); + } + + void sanitize_moe_fp16(__fp16* data, size_t count) { + for (size_t i = 0; i < count; ++i) { + float value = static_cast(data[i]); + if (!std::isfinite(value) || std::abs(value) > MOE_FP16_MAX) { + data[i] = moe_finite_f16(value); + } + } + } + + void multiply_moe_fp16_sanitized(__fp16* lhs, const __fp16* rhs, size_t count) { + for (size_t i = 0; i < count; ++i) { + float value = static_cast(lhs[i]) * static_cast(rhs[i]); + lhs[i] = moe_finite_f16(value); + } + } } void compute_moe_layer_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { @@ -146,6 +171,9 @@ void compute_moe_layer_node(GraphNode& node, const std::vector