From 22c03179b0f84eae424f926258d198aa71d7db67 Mon Sep 17 00:00:00 2001 From: jakmro Date: Wed, 8 Jul 2026 17:33:03 +0200 Subject: [PATCH 1/4] remove coreml --- README.md | 10 +- apple/CMakeLists.txt | 2 +- bindings/rust/cactus.rs | 1 - blog/gemma4.md | 6 +- blog/hybrid_transcription.md | 2 +- blog/lfm2.5_350m.md | 14 +- cactus-engine/CMakeLists.txt | 13 +- cactus-engine/src/engine.h | 68 -- cactus-engine/src/model.cpp | 182 +-- cactus-engine/src/model_npu.cpp | 359 ------ cactus-engine/src/npu.cpp | 15 - cactus-engine/src/npu_ane.h | 107 -- cactus-engine/src/npu_ane.mm | 1046 ----------------- cactus-engine/tests/ios/configure_xcode.rb | 2 +- cactus-graph/cactus_graph.h | 2 +- cactus-graph/src/ops_conv.cpp | 28 - cactus-graph/src/ops_nn.cpp | 14 - cactus-graph/src/ops_tensor.cpp | 4 - cactus-graph/src/param_io.cpp | 2 +- docs/cactus_engine.md | 4 +- docs/cactus_graph.md | 6 +- docs/cactus_transpiler.md | 5 - docs/choose-bindings.md | 4 +- llms.txt | 4 +- python/README.md | 8 +- python/cactus/bindings/cactus.py | 3 +- python/cactus/cli/__init__.py | 34 +- python/cactus/cli/code.py | 2 - python/cactus/cli/common.py | 5 - python/cactus/cli/compile.py | 1 - python/cactus/cli/convert.py | 17 +- python/cactus/cli/download.py | 24 +- python/cactus/cli/model.py | 42 +- python/cactus/cli/runtime.py | 1 - python/cactus/cli/upload.py | 5 +- python/cactus/cli/utils.py | 35 +- python/cactus/server.py | 2 +- python/cactus/transpile/README.md | 4 +- python/cactus/transpile/component_pipeline.py | 7 +- python/cactus/transpile/hf_model.py | 55 +- python/cactus/transpile/model_adapters.py | 39 +- python/cactus/transpile/npu/README.md | 86 -- python/cactus/transpile/npu/__init__.py | 3 - python/cactus/transpile/npu/audio.py | 115 -- .../transpile/npu/coremltools_patches.py | 434 ------- python/cactus/transpile/npu/pipeline.py | 94 -- python/cactus/transpile/npu/source.py | 118 -- python/cactus/transpile/npu/vision.py | 126 -- python/cactus/transpile/tdt_runtime.py | 135 --- python/pyproject.toml | 1 - python/tests/bundles.py | 4 +- python/tests/test_bindings.py | 10 +- python/tests/test_cli_run.py | 1 - python/tests/test_download.py | 27 +- python/tests/test_model.py | 8 +- python/tests/test_server_live.py | 2 +- 56 files changed, 110 insertions(+), 3238 deletions(-) delete mode 100644 cactus-engine/src/model_npu.cpp delete mode 100644 cactus-engine/src/npu.cpp delete mode 100644 cactus-engine/src/npu_ane.h delete mode 100644 cactus-engine/src/npu_ane.mm delete mode 100644 python/cactus/transpile/npu/README.md delete mode 100644 python/cactus/transpile/npu/__init__.py delete mode 100644 python/cactus/transpile/npu/audio.py delete mode 100644 python/cactus/transpile/npu/coremltools_patches.py delete mode 100644 python/cactus/transpile/npu/pipeline.py delete mode 100644 python/cactus/transpile/npu/source.py delete mode 100644 python/cactus/transpile/npu/vision.py diff --git a/README.md b/README.md index ea6d56f81..1c5a4c621 100644 --- a/README.md +++ b/README.md @@ -124,7 +124,7 @@ graph.hard_reset(); - LLM: Gemma-4-E2B-CQ4 (1k-context prefill / decode for 100 tokens) - VLM: Gemma-4-E2B-CQ4 (256px image encode time / decode) -- Transcribe: Parakeet-TDT-0.6B-CQ4 (20s audio ens-to-end transcribe time) +- Transcribe: Parakeet-TDT-0.6B-CQ4 (20s audio end-to-end transcribe time) - 1k-Context RAM: peak MB during the LLM benchmark Command: `cactus benchmark` [optional `--ios` or `--android`] @@ -208,7 +208,6 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ │ │ cactus run [model|path] run a model (downloads if needed) │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --weights general|apple weights bundle variant (default: general) │ │ --backend auto|cpu|metal inference backend (default: auto) │ │ --image image file for VLM inference │ │ --audio audio file for audio chat │ @@ -222,20 +221,17 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ --file audio file to transcribe (WAV) │ │ --language language code (default: en) │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --weights general|apple weights bundle variant (default: general) │ │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild from source │ │ │ │ cactus download [model] get a bundle (prebuilt, else build) │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --weights general|apple weights bundle variant (default: general) │ │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild from source │ │ │ │ cactus convert [dir] HuggingFace -> runnable cactus bundle │ │ (CQ weights + runtime graph) │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --weights general|apple weights bundle variant (default: general) │ │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild from source │ │ --lora merge a LoRA adapter before converting │ @@ -246,7 +242,6 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ --host bind address (default: 127.0.0.1) │ │ --port port (default: 8080) │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --weights general|apple weights bundle variant (default: general) │ │ --backend auto|cpu|metal inference backend (default: auto) │ │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild from source │ @@ -257,7 +252,6 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ cactus code run the AI coding agent (TUI / print) │ │ --serve-model auto-start a server with this model │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --weights general|apple weights bundle variant (default: general) │ │ --backend auto|cpu|metal inference backend (default: auto) │ │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild from source │ @@ -282,7 +276,6 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ --model default: google/gemma-4-E2B-it │ │ --transcription-model default: nvidia/parakeet-tdt-0.6b-v3 │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --weights general|apple weights bundle variant (default: general) │ │ --backend auto|cpu|metal inference backend (default: auto) │ │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild of test models │ @@ -298,7 +291,6 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ --model default: google/gemma-4-E2B-it │ │ --transcription-model default: nvidia/parakeet-tdt-0.6b-v3 │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --weights general|apple weights bundle variant (default: general) │ │ --backend auto|cpu|metal inference backend (default: auto) │ │ --ios run on connected iPhone │ │ --android run on connected Android │ diff --git a/apple/CMakeLists.txt b/apple/CMakeLists.txt index cdd95297d..d2795bdd4 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 CoreML" "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" + "-framework Accelerate" "-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 e126d6c00..f00ecaded 100644 --- a/bindings/rust/cactus.rs +++ b/bindings/rust/cactus.rs @@ -7,7 +7,6 @@ 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 = "Foundation", kind = "framework"))] #[cfg_attr(target_os = "macos", link(name = "Security", kind = "framework"))] #[cfg_attr(target_os = "macos", link(name = "SystemConfiguration", kind = "framework"))] diff --git a/blog/gemma4.md b/blog/gemma4.md index 2ad4e9fb5..cd07ff058 100644 --- a/blog/gemma4.md +++ b/blog/gemma4.md @@ -19,14 +19,14 @@ This matters because latency compounds. A pipeline that transcribes audio, then ## Performance -Cactus targets ARM across platforms: Apple Silicon Macs, iPhones, iPads, Vision Pro, and Android devices with ARM64 chipsets (Snapdragon, Dimensity, Tensor). On Apple hardware we additionally leverage the Neural Engine for vision and audio encoders, while on Android and Linux the runtime uses NEON, i8mm, and dot-product intrinsics. +Cactus targets ARM across platforms: Apple Silicon Macs, iPhones, iPads, Vision Pro, and Android devices with ARM64 chipsets (Snapdragon, Dimensity, Tensor). On Apple hardware we additionally leverage Metal GPU acceleration, while on Android and Linux the runtime uses NEON, i8mm, and dot-product intrinsics. | Metric (M5 Mac/iPad/Vision Pro) | E2B | |---|---| | 4096-token prefill | 660 tok/s | | 1024-token decode | 40 tok/s | | 30s audio end-to-end | 0.3s | -| Image encode (ANE) | 0.7s | +| Image encode | 0.7s | 40 tok/s decode is faster than most people read. You're not waiting for the model. The model is waiting for you. @@ -128,7 +128,7 @@ Cactus supports React Native, Flutter, Swift, Kotlin, Python, Rust, and C++. Pic |---|---|---| | React Native | [React Native bindings](/bindings/react-native/) | Native bridge over the C API | | Flutter | [cactus-flutter](/bindings/flutter/) | Dart FFI bindings | -| Swift (iOS/macOS) | [cactus-swift](/bindings/swift/) | XCFramework with NPU support | +| Swift (iOS/macOS) | [cactus-swift](/bindings/swift/) | XCFramework with Metal support | | Kotlin (Android) | [cactus-kotlin](/bindings/kotlin/) | JNI + Kotlin Multiplatform | | Python | [cactus-compute](/python/) | `pip install cactus-compute` | | Rust | [Rust bindings](/bindings/rust/) | Copy `cactus.rs`, link `libcactus_engine.a` | diff --git a/blog/hybrid_transcription.md b/blog/hybrid_transcription.md index 973bf22ac..e92ce139c 100644 --- a/blog/hybrid_transcription.md +++ b/blog/hybrid_transcription.md @@ -53,7 +53,7 @@ cactus transcribe --file path/to/audio.wav ## Built for the Edge, from the Ground Up Achieving this required relying on our industry-leading on-device engine. We built Cactus from the ground up exactly for this. -Because the engine sits directly on the metal, it is optimized for edge constraints. It targets ARM-based CPUs and NPUs natively, resulting in super low RAM usage (e.g. Moonshine transcription for example, runs with sub-10MB RAM utilization). This is a critical factor for mobile and wearable devices where memory is strictly rationed by the OS. +Because the engine sits directly on the metal, it is optimized for edge constraints. It targets ARM-based CPUs natively, resulting in super low RAM usage (e.g. Moonshine transcription for example, runs with sub-10MB RAM utilization). This is a critical factor for mobile and wearable devices where memory is strictly rationed by the OS. By keeping the footprint minimal, the engine operates silently in the background without draining battery or causing OS-level memory warnings, making it viable for continuous, ambient listening. diff --git a/blog/lfm2.5_350m.md b/blog/lfm2.5_350m.md index 9f2e40750..3c314fa2f 100644 --- a/blog/lfm2.5_350m.md +++ b/blog/lfm2.5_350m.md @@ -41,7 +41,7 @@ most people read at roughly 4-5 words per second, so even the Raspberry Pi's 30 Second, the RAM gap between Apple and Android: 56-85 MB versus 300-330 MB for the same 355 MB model. That gap is not a bug. It's a direct consequence of how the two platforms handle memory-mapped files, which we'll get into in the Cactus format section below. -## Why CPU, Not GPU or NPU +## Why CPU, Not GPU or Dedicated Accelerators It is tempting to reach for the GPU. On paper, mobile GPUs have 10-50x the FLOPS of a single CPU core. But inference on a 350M-parameter model is not a FLOPS problem, it's a memory bandwidth problem. @@ -55,17 +55,17 @@ 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. -NPU is 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 CoreML model conversion, which imposes its own quantization and graph constraints. +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 CoreML severely lag the state-of-the-art methods used by GGML and Cactus. +The quantisation techniques available through that toolchain severely lag the state-of-the-art methods used by GGML and Cactus. Qualcomm's QNN SDK has similar restrictions (limited quantisation scheme support) and not all Android devices ship Qualcomm chips. -Both lock you into a vendor-specific format that may not support your model architecture at all, and neither works on budget devices that lack the NPU entirely. +Both lock you into a vendor-specific format that may not support your model architecture at all, and neither works on budget devices that lack the accelerator entirely. CPU is the universal target. Every ARM device, flagship, budget, wearable, Raspberry Pi, and most recent ones support DOTPROD or I8MM extensions. -A single well-optimized CPU codepath works everywhere, the main challenge is energy-inefficiency compared to NPU. +A single well-optimized CPU codepath works everywhere, the main challenge is energy-inefficiency compared to dedicated accelerators. -ARM is also closing this gap. The progression from NEON to DOTPROD to I8MM to SME2 reflects a deliberate push to bring matrix-processing efficiency onto the CPU itself. Each generation narrows the energy and throughput gap with dedicated NPUs. SME2, landing on upcoming ARMv9.2-A cores, adds streaming matrix operations that approach NPU-class efficiency for INT8 workloads while remaining fully programmable. The bet is that general-purpose CPU silicon will continue absorbing the capabilities that once justified a separate accelerator. +ARM is also closing this gap. The progression from NEON to DOTPROD to I8MM to SME2 reflects a deliberate push to bring matrix-processing efficiency onto the CPU itself. Each generation narrows the energy and throughput gap with dedicated accelerators. SME2, landing on upcoming ARMv9.2-A cores, adds streaming matrix operations that approach accelerator-class efficiency for INT8 workloads while remaining fully programmable. The bet is that general-purpose CPU silicon will continue absorbing the capabilities that once justified a separate accelerator. ## Why Single Core for Decode diff --git a/cactus-engine/CMakeLists.txt b/cactus-engine/CMakeLists.txt index 7e057285e..ce09c48fe 100644 --- a/cactus-engine/CMakeLists.txt +++ b/cactus-engine/CMakeLists.txt @@ -16,7 +16,6 @@ set(ENGINE_SOURCES src/constraints.cpp src/model.cpp src/kv_compress.cpp - src/model_npu.cpp src/engine_image.cpp src/index.cpp src/index_ffi.cpp @@ -32,14 +31,6 @@ 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 @@ -56,7 +47,7 @@ target_link_libraries(cactus_engine PUBLIC cactus_graph) if(APPLE) target_link_libraries(cactus_engine PUBLIC - "-framework CoreML" "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" + "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" "-framework Security" "-framework SystemConfiguration" "-framework CFNetwork" ) @@ -156,7 +147,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 CoreML" "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" + "-framework Accelerate" "-framework Foundation" "-framework Metal" "-framework MetalPerformanceShaders" "-framework Security" "-framework SystemConfiguration" "-framework CFNetwork" ) else() diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index d0cffb294..73fc6d47a 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -17,46 +17,6 @@ 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 { @@ -164,9 +124,6 @@ struct Config { enum class Activation {GELU = 0, SILU = 1}; Activation activation = Activation::SILU; - enum class Backend {CPU = 0, NPU = 1}; - Backend default_backend = Backend::CPU; - enum class Precision {INT8 = 0, FP16 = 1, FP32 = 2}; Precision precision = Precision::FP32; @@ -740,14 +697,6 @@ 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() {} @@ -823,7 +772,6 @@ class Model { size_t feature_row_bytes, Precision feature_precision); void reset_encoder_cross_kv_route_state(); 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); @@ -869,9 +817,6 @@ class Model { 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, @@ -908,19 +853,6 @@ class Model { bool prefill_tail_pad_disabled_ = false; std::string family_; - 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_shapes_; diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index c8a5cb6c9..6f9eb1c7e 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -632,30 +632,6 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, cache_max_seq_len_ = context_size; - if (!npu_audio_encoder_mlpackage_.empty()) { - std::string full_path = bundle_dir + "/" + npu_audio_encoder_mlpackage_; - if (!load_npu_audio_encoder(full_path, npu_audio_compute_units_)) { - CACTUS_LOG_WARN("model", "NPU audio encoder load failed for " << full_path << "; falling back to CPU"); - } - } - if (!npu_vision_encoder_mlpackage_.empty()) { - std::string full_path = bundle_dir + "/" + npu_vision_encoder_mlpackage_; - if (!load_npu_vision_encoder(full_path)) { - CACTUS_LOG_WARN("model", "NPU vision encoder load failed for " << full_path << "; falling back to CPU"); - } else if (tokenizer_) { - const auto& npu_out = npu_vision_encoder_->get_output_shape(); - size_t npu_rows = 0; - if (npu_out.size() >= 3) npu_rows = static_cast(npu_out[npu_out.size() - 2]); - else if (npu_out.size() >= 2) npu_rows = static_cast(npu_out[0]); - if (npu_rows > 0) tokenizer_->set_image_soft_token_count(npu_rows); - } - } - if (!npu_source_encoder_mlpackage_.empty()) { - std::string full_path = bundle_dir + "/" + npu_source_encoder_mlpackage_; - if (!load_npu_source_encoder(full_path)) { - CACTUS_LOG_WARN("model", "NPU source encoder load failed for " << full_path << "; falling back to CPU"); - } - } if (load_handoff_probe()) { const bool is_parakeet = config_.model_type == Config::ModelType::PARAKEET_TDT; const Component* probe_comp = is_parakeet ? audio_encoder_ : decoder_; @@ -684,18 +660,6 @@ bool Model::load_manifest() { if (obj.count("family") && obj.at("family").is()) { family_ = obj.at("family").get(); } - 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(); @@ -1592,9 +1556,6 @@ void Model::run_vision_encoder(const std::string& image_path) { Precision::FP32); } else { Gemma4ImagePreprocessed prep = preprocess_gemma4_image(image_path, config_); - 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]; @@ -1695,7 +1656,6 @@ void Model::run_vision_encoder_lfm2_vl(const std::string& image_path) { void Model::encode_lfm2_vl_image_into_features(const std::string& image_path) { Lfm2VlImagePreprocessed prep = preprocess_lfm2_vl_image(image_path, config_); - const bool use_npu = lfm2_vl_use_npu_vision(); const int dim = lfm2_pos_grid_dim_; const int factor = config_.downsample_factor ? static_cast(config_.downsample_factor) : 2; const size_t max_patches = prep.max_num_patches; @@ -1721,13 +1681,11 @@ void Model::encode_lfm2_vl_image_into_features(const std::string& image_path) { if (proj_rows == 0 || text_hidden == 0) { throw std::runtime_error("lfm2_vl vision_projector has an invalid output shape"); } - if (!use_npu) { - const auto& enc_out_desc0 = vision_encoder_->graph->get_output_buffer(enc_out_node); - const size_t enc_elem = PrecisionTraits::size_of(enc_out_desc0.precision); - const size_t enc_out_elems = enc_elem ? enc_out_desc0.byte_size / enc_elem : 0; - if (enc_out_elems < max_patches * static_cast(dim)) { - throw std::runtime_error("lfm2_vl vision_encoder output is smaller than the preprocessed patch grid"); - } + const auto& enc_out_desc0 = vision_encoder_->graph->get_output_buffer(enc_out_node); + const size_t enc_elem = PrecisionTraits::size_of(enc_out_desc0.precision); + const size_t enc_out_elems = enc_elem ? enc_out_desc0.byte_size / enc_elem : 0; + if (enc_out_elems < max_patches * static_cast(dim)) { + throw std::runtime_error("lfm2_vl vision_encoder output is smaller than the preprocessed patch grid"); } auto write_mask = [&](int comp_idx, const int64_t* src, size_t count) { @@ -1779,25 +1737,17 @@ void Model::encode_lfm2_vl_image_into_features(const std::string& image_path) { interpolate_position_embeddings( lfm2_pos_grid_.data(), lfm2_pos_grid_h_, lfm2_pos_grid_w_, dim, h, w, pos_buf.data()); - if (use_npu) { - const float* pv_src = prep.pixel_values.data() + t * max_patches * patch_dim; - const int64_t* m_src = prep.pixel_attention_mask.data() + t * max_patches; - if (!lfm2_vl_encode_tile_npu(pv_src, m_src, pos_buf.data(), max_patches, dim, patch_dim, enc_out_f)) { - throw std::runtime_error("lfm2_vl NPU vision encode failed"); - } - } else { - write_float_input(*vision_encoder_, pv_idx, - prep.pixel_values.data() + t * max_patches * patch_dim, max_patches * patch_dim); - if (pm_idx >= 0) { - write_mask(pm_idx, prep.pixel_attention_mask.data() + t * max_patches, max_patches); - } - write_float_input(*vision_encoder_, pe_idx, pos_buf.data(), max_patches * static_cast(dim)); - vision_encoder_->graph->execute(); - const auto& enc_desc = vision_encoder_->graph->get_output_buffer(enc_out_node); - const void* enc_ptr = vision_encoder_->graph->get_output(enc_out_node); - typed_buffer_to_float(enc_ptr, enc_desc.precision, enc_out_f.data(), - max_patches * static_cast(dim)); + write_float_input(*vision_encoder_, pv_idx, + prep.pixel_values.data() + t * max_patches * patch_dim, max_patches * patch_dim); + if (pm_idx >= 0) { + write_mask(pm_idx, prep.pixel_attention_mask.data() + t * max_patches, max_patches); } + write_float_input(*vision_encoder_, pe_idx, pos_buf.data(), max_patches * static_cast(dim)); + vision_encoder_->graph->execute(); + const auto& enc_desc = vision_encoder_->graph->get_output_buffer(enc_out_node); + const void* enc_ptr = vision_encoder_->graph->get_output(enc_out_node); + typed_buffer_to_float(enc_ptr, enc_desc.precision, enc_out_f.data(), + max_patches * static_cast(dim)); unshuf.assign(static_cast(num_tokens) * cff, 0.0f); pixel_unshuffle(enc_out_f.data(), h, w, dim, factor, unshuf.data()); @@ -1850,9 +1800,6 @@ void Model::run_audio_encoder_messages(const std::vector>& au void Model::run_audio_encoder(const std::vector& audio_features) { if (!audio_encoder_) return; - if (has_npu_audio_encoder() && audio_encode_via_npu(audio_features)) { - return; - } const std::vector candidate_input_names = {"input_features", "audio_features"}; int feature_idx = -1; for (const auto& name : candidate_input_names) { @@ -2483,16 +2430,6 @@ bool Model::finish_encoder_cross_kv_prepare() { 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; - 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; @@ -2525,10 +2462,6 @@ bool Model::prepare_encoder_cross_kv_from_text(const std::vector& toke } } - if (source_encode_via_npu(tokens)) { - return finish_encoder_cross_kv_prepare_after_source(); - } - return finish_encoder_cross_kv_prepare(); } @@ -2800,9 +2733,6 @@ bool Model::run_chunk_prefill_path(const std::vector& tokens, qwen_images.push_back(std::move(prep)); } else { Gemma4ImagePreprocessed prep = preprocess_gemma4_image(path, config_); - 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]; @@ -3341,72 +3271,9 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a write_typed_buffer(feat_buf, feat_desc.precision, transposed.data(), transposed.size() * sizeof(float), Precision::FP32); - std::vector<__fp16> npu_hidden_storage; - bool used_npu = false; - size_t npu_hidden_T = 0; - if (has_npu_audio_encoder()) { - const std::vector in_shape = npu_audio_encoder_->get_input_shape(); - const std::vector out_shape = npu_audio_encoder_->get_output_shape(); - if (in_shape.size() >= 3 && out_shape.size() >= 3 && - in_shape[1] > 0 && in_shape[2] > 0 && out_shape[1] > 0 && out_shape[2] > 0 && - static_cast(in_shape[2]) == expected_mels && - copy_frames <= static_cast(in_shape[1])) { - const size_t window_frames = static_cast(in_shape[1]); - const size_t window_hidden = static_cast(out_shape[1]); - const size_t hidden_dim_npu = static_cast(out_shape[2]); - const size_t chunk_input_elems = window_frames * expected_mels; - const size_t chunk_output_elems = window_hidden * hidden_dim_npu; - const size_t num_chunks = (copy_frames + window_frames - 1) / window_frames; - const size_t total_hidden_T = num_chunks * window_hidden; - npu_hidden_storage.assign(total_hidden_T * hidden_dim_npu, __fp16(0)); - const bool want_probe = handoff_probe_loaded_ && - static_cast(handoff_probe_feat_dim_) == hidden_dim_npu; - std::vector<__fp16> input_fp16(chunk_input_elems); - bool all_ok = num_chunks > 0; - for (size_t c = 0; c < num_chunks && all_ok; ++c) { - if (should_stop && should_stop->load()) return emitted; - const size_t frame_start = c * window_frames; - const size_t frame_end = std::min(frame_start + window_frames, copy_frames); - std::fill(input_fp16.begin(), input_fp16.end(), __fp16(0)); - for (size_t t = frame_start; t < frame_end; ++t) { - const size_t local = (t - frame_start) * expected_mels; - const size_t src = t * expected_mels; - for (size_t m = 0; m < expected_mels; ++m) { - input_fp16[local + m] = static_cast<__fp16>(transposed[src + m]); - } - } - __fp16* out_ptr = npu_hidden_storage.data() + c * chunk_output_elems; - size_t written = npu_audio_encoder_->encode( - input_fp16.data(), out_ptr, in_shape, "x", "encoded"); - if (written == 0) { all_ok = false; break; } - } - if (all_ok) { - used_npu = true; - const size_t valid_input = copy_frames; - npu_hidden_T = (valid_input * window_hidden + window_frames - 1) / window_frames; - if (npu_hidden_T > total_hidden_T) npu_hidden_T = total_hidden_T; - CACTUS_LOG_INFO("model", "Parakeet audio encoder ran on NPU (" - << num_chunks << " chunks, " << npu_hidden_T << " valid hidden frames)"); - if (want_probe) { - const size_t probe_elems = npu_hidden_T * hidden_dim_npu; - handoff_probe_hidden_.reserve(handoff_probe_hidden_.size() + probe_elems); - const __fp16* pp = npu_hidden_storage.data(); - for (size_t i = 0; i < probe_elems; ++i) { - handoff_probe_hidden_.push_back(static_cast(pp[i])); - } - CACTUS_LOG_INFO("cloud_handoff", "Captured " << npu_hidden_T - << " NPU probe frames for the handoff probe"); - } - } else { - CACTUS_LOG_WARN("model", "NPU audio encoder chunk failed; falling back to CPU graph"); - } - } - } - if (!used_npu) { - if (should_stop && should_stop->load()) return emitted; - audio_enc->graph->execute(); - maybe_capture_handoff_probe_hidden(*audio_enc, "encoder_hidden_states"); - } + if (should_stop && should_stop->load()) return emitted; + audio_enc->graph->execute(); + maybe_capture_handoff_probe_hidden(*audio_enc, "encoder_hidden_states"); int hidden_idx = output_index(*audio_enc, "encoder_hidden_states"); if (hidden_idx < 0) { @@ -3415,19 +3282,14 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a } size_t hidden_node = static_cast(audio_enc->output_node_ids[hidden_idx]); const auto& hidden_desc = audio_enc->graph->get_output_buffer(hidden_node); - const uint8_t* hidden_ptr; - if (used_npu) { - hidden_ptr = reinterpret_cast(npu_hidden_storage.data()); - } else { - hidden_ptr = static_cast(audio_enc->graph->get_output(hidden_node)); - } + const uint8_t* hidden_ptr = static_cast(audio_enc->graph->get_output(hidden_node)); if (hidden_desc.shape.size() < 3 || hidden_ptr == nullptr) { CACTUS_LOG_ERROR("model", "encoder_hidden_states must be 3D [B, T, D]"); return emitted; } - const size_t T = used_npu ? npu_hidden_T : hidden_desc.shape[1]; + const size_t T = hidden_desc.shape[1]; const size_t D = hidden_desc.shape[2]; - const Precision hidden_precision = used_npu ? Precision::FP16 : hidden_desc.precision; + const Precision hidden_precision = hidden_desc.precision; const size_t hidden_elem = PrecisionTraits::size_of(hidden_precision); const size_t frame_bytes = D * hidden_elem; @@ -3506,7 +3368,7 @@ std::vector Model::transcribe_parakeet_tdt(const std::vector& a size_t commit_to = T; if (stream) { size_t valid_hidden = T; - if (!used_npu && expected_frames > 0) + if (expected_frames > 0) valid_hidden = std::min(T, (copy_frames * T) / expected_frames); commit_to = (end_frame > 0) ? std::min(end_frame, valid_hidden) : valid_hidden; } diff --git a/cactus-engine/src/model_npu.cpp b/cactus-engine/src/model_npu.cpp deleted file mode 100644 index 8c421dee1..000000000 --- a/cactus-engine/src/model_npu.cpp +++ /dev/null @@ -1,359 +0,0 @@ -#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 (re-transpile with --npu 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 deleted file mode 100644 index 3222b0269..000000000 --- a/cactus-engine/src/npu.cpp +++ /dev/null @@ -1,15 +0,0 @@ -#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 deleted file mode 100644 index 07b678687..000000000 --- a/cactus-engine/src/npu_ane.h +++ /dev/null @@ -1,107 +0,0 @@ -#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 deleted file mode 100644 index 3bcf2b51c..000000000 --- a/cactus-engine/src/npu_ane.mm +++ /dev/null @@ -1,1046 +0,0 @@ -#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/tests/ios/configure_xcode.rb b/cactus-engine/tests/ios/configure_xcode.rb index 7a8196a41..6e8379500 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 CoreML', '-framework Foundation', '-framework Accelerate', '-framework Security', '-framework SystemConfiguration', '-framework CFNetwork', '-framework Metal', '-framework MetalPerformanceShaders'].each do |framework| + ['-framework Foundation', '-framework Accelerate', '-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-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 4d08a1ef9..e1f59b38b 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -84,7 +84,7 @@ class Logger { #define CACTUS_LOG_WARN(component, msg) CACTUS_LOG(cactus::LogLevel::WARN, component, msg) #define CACTUS_LOG_ERROR(component, msg) CACTUS_LOG(cactus::LogLevel::ERROR, component, msg) -enum class ComputeBackend { CPU, NPU, METAL }; +enum class ComputeBackend { CPU, METAL }; enum class Activation { SILU, GELU, GELU_ERF, RELU, SIGMOID, TANH }; diff --git a/cactus-graph/src/ops_conv.cpp b/cactus-graph/src/ops_conv.cpp index d17f1340d..20cf09030 100644 --- a/cactus-graph/src/ops_conv.cpp +++ b/cactus-graph/src/ops_conv.cpp @@ -41,10 +41,6 @@ std::vector<__fp16> dequantize_int8_weights_to_fp16( } void compute_conv1d_causal_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU causal convolution operation not yet implemented"); - } - 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; @@ -118,10 +114,6 @@ void compute_conv1d_causal_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU causal convolution operation not yet implemented"); - } - 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; @@ -238,10 +230,6 @@ void compute_conv1d_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU conv1d_same_depthwise_k9 operation not yet implemented"); - } - const auto& X = get_input(node, 0, nodes, node_index_map); const auto& W = get_input(node, 1, nodes, node_index_map); const BufferDesc* B = nullptr; @@ -322,10 +310,6 @@ void compute_conv1d_same_depthwise_k9_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU conv2d_k3s2p1 operation not yet implemented"); - } - const auto& X = get_input(node, 0, nodes, node_index_map); const auto& W = get_input(node, 1, nodes, node_index_map); const BufferDesc* B = nullptr; @@ -399,10 +383,6 @@ void compute_conv2d_k3s2p1_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU conv2d_depthwise_k3s2p1 operation not yet implemented"); - } - const auto& X = get_input(node, 0, nodes, node_index_map); const auto& W = get_input(node, 1, nodes, node_index_map); const BufferDesc* B = nullptr; @@ -488,10 +468,6 @@ void compute_conv2d_depthwise_k3s2p1_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU conv2d_pointwise_1x1 operation not yet implemented"); - } - const auto& X = get_input(node, 0, nodes, node_index_map); const auto& W = get_input(node, 1, nodes, node_index_map); const BufferDesc* B = nullptr; @@ -578,10 +554,6 @@ void compute_conv2d_pointwise_1x1_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU conv1d_pointwise operation not yet implemented"); - } - const auto& X = get_input(node, 0, nodes, node_index_map); const auto& W = get_input(node, 1, nodes, node_index_map); const BufferDesc* B = nullptr; diff --git a/cactus-graph/src/ops_nn.cpp b/cactus-graph/src/ops_nn.cpp index e0ad4de91..b8f0eb033 100644 --- a/cactus-graph/src/ops_nn.cpp +++ b/cactus-graph/src/ops_nn.cpp @@ -40,12 +40,6 @@ void compute_matmul_node(GraphNode& node, const std::vector 0) { if (lhs_buffer.precision != Precision::FP16) { throw std::runtime_error("TQ matmul requires FP16 activations"); @@ -485,10 +479,6 @@ void compute_rms_norm_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU RoPE operation not yet implemented"); - } - const auto& input_buffer = get_input(node, 0, nodes, node_index_map); const auto& shape = input_buffer.shape; @@ -612,10 +602,6 @@ void compute_rel_pos_bias_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU attention operation not yet implemented"); - } - if (node.input_ids.size() < 3 || node.input_ids.size() > 4) { throw std::runtime_error("Attention operation requires 3 or 4 inputs (query, key, value[, mask]), got " + std::to_string(node.input_ids.size()) + " inputs"); diff --git a/cactus-graph/src/ops_tensor.cpp b/cactus-graph/src/ops_tensor.cpp index c33a59996..1afb75563 100644 --- a/cactus-graph/src/ops_tensor.cpp +++ b/cactus-graph/src/ops_tensor.cpp @@ -7,10 +7,6 @@ #include void compute_transpose_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { - if (node.params.backend == ComputeBackend::NPU) { - throw std::runtime_error("NPU transpose operation not yet implemented"); - } - const auto& input_buffer = get_input(node, 0, nodes, node_index_map); if (input_buffer.precision != Precision::FP16) { diff --git a/cactus-graph/src/param_io.cpp b/cactus-graph/src/param_io.cpp index 2e3fdadd4..b173c3395 100644 --- a/cactus-graph/src/param_io.cpp +++ b/cactus-graph/src/param_io.cpp @@ -314,7 +314,7 @@ void read_field(std::istream& in, ParamField field, OpParams& params) { case ParamField::PretransposedRhs: params.pretransposed_rhs = (read_u32(in) != 0); break; case ParamField::Backend: { uint32_t backend_val = read_u32(in); - if (backend_val > static_cast(ComputeBackend::NPU)) { + if (backend_val > static_cast(ComputeBackend::METAL)) { throw std::runtime_error("Graph file corrupted: invalid backend"); } params.backend = static_cast(backend_val); diff --git a/docs/cactus_engine.md b/docs/cactus_engine.md index 00d8dc6e4..3e16d8ec0 100644 --- a/docs/cactus_engine.md +++ b/docs/cactus_engine.md @@ -25,9 +25,7 @@ cactus auth `cactus download` fetches a **pre-built runtime bundle** (CQ weights + serialized graph + manifest) from [huggingface.co/Cactus-Compute](https://huggingface.co/Cactus-Compute) into -`weights/-cq[-]/`. Defaults to `--weights general` (the -portable build, used on every platform); pass `--weights apple` for the -Apple-specific (CoreML/NPU) variant. The result can be loaded directly via +`weights/-cq/`. The result can be loaded directly via `cactus_init()`. For models not on Cactus-Compute, build a bundle from source with diff --git a/docs/cactus_graph.md b/docs/cactus_graph.md index 2c56e0f50..8c0f4d858 100644 --- a/docs/cactus_graph.md +++ b/docs/cactus_graph.md @@ -489,9 +489,9 @@ size_t similarity = fixture.graph().divide(dot_product, fixture.graph().multiply 1. **Batch operations**: Process multiple samples together 2. **Pre-transpose weights**: Use `pretransposed_rhs=true` for matmul when possible 3. **Fused operations**: The framework automatically fuses compatible operations -4. **Backend selection**: Use NPU backend for supported operations: +4. **Backend selection**: Use the Metal backend for supported operations: ```cpp - size_t result = graph.matmul(a, b, false, ComputeBackend::NPU); + size_t result = graph.matmul(a, b, false, ComputeBackend::METAL); ``` ### Graph Construction @@ -522,7 +522,7 @@ size_t CactusGraph::relu(size_t input) { ``` 3. **Implement the op in the execution engine** Implement the kernel or graph op code in the relevant file, usually in `cactus-kernels/src/` -Register the new op in the dispatch table in `cactus-graph/src/execute.cpp` for the supported backends (CPU, NPU) +Register the new op in the dispatch table in `cactus-graph/src/execute.cpp` for the supported backends (CPU, Metal) 4. **Export op in FFI bindings** - header: `cactus-graph/cactus_graph.h` (in the `extern "C"` block) diff --git a/docs/cactus_transpiler.md b/docs/cactus_transpiler.md index 3be3c77ef..3ba4ebb1d 100644 --- a/docs/cactus_transpiler.md +++ b/docs/cactus_transpiler.md @@ -438,10 +438,6 @@ slot, so single-stream decode stays lean while fixed-batch decode (`decode_batch | `--no-fuse-rope` | Disable RoPE fusion | | `--no-fuse-attention` | Disable attention fusion | -NPU emission (CoreML `.mlpackage`s for Apple Silicon audio and vision encoders) -is available through the `--npu`, `--npu-quantize`, `--npu-audio-quantize`, and -`--npu-vision-quantize` flags on `cactus convert`. - ### `cactus run` ```bash @@ -454,7 +450,6 @@ huggingface.co/Cactus-Compute) or a local path to a bundle directory. | Option | Description | |--------|-------------| | `--bits 1\|2\|3\|4` | CQ quantization bits when downloading (default: 4) | -| `--weights general\|apple` | Weights bundle variant (default: general — portable on every platform) | | `--token ` | HuggingFace token (gated models) | | `--prompt ` | Input prompt | | `--input-ids ` | Comma-separated token IDs | diff --git a/docs/choose-bindings.md b/docs/choose-bindings.md index b37769035..df3a90eb7 100644 --- a/docs/choose-bindings.md +++ b/docs/choose-bindings.md @@ -23,7 +23,7 @@ Pick the right one for your platform and use case: - **React Native** -- already have a React Native project? Just `npm install` and go - **Flutter** -- cross-platform mobile + Mac with full native bindings - **Kotlin** -- native Android apps or Kotlin Multiplatform - - **Swift** -- native iOS/macOS apps with the Apple Neural Engine + - **Swift** -- native iOS/macOS apps with Metal acceleration !!! tip "Server-side or scripting?" - **Python** -- server-side inference, batch processing, or rapid prototyping @@ -37,7 +37,7 @@ Pick the right one for your platform and use case: - **[React Native](/bindings/react-native/)** -- Native bridge modules over the C API for iOS and Android - **[Python](/python/)** -- Module-level FFI bindings, mirrors the C API -- **[Swift](/bindings/swift/)** -- XCFramework for iOS/macOS with NPU support +- **[Swift](/bindings/swift/)** -- XCFramework for iOS/macOS with Metal support - **[Kotlin](/bindings/kotlin/)** -- JNI bindings + Kotlin Multiplatform support - **[Flutter](/bindings/flutter/)** -- Dart FFI bindings for Android, iOS, and macOS - **[Rust](/bindings/rust/)** -- Raw `extern "C"` FFI declarations diff --git a/llms.txt b/llms.txt index 2be6be79e..420d17e95 100644 --- a/llms.txt +++ b/llms.txt @@ -4,7 +4,7 @@ Cactus has three layers: Cactus Kernels (ARM SIMD kernels for Apple, Snapdragon, Google Tensor, Exynos, MediaTek, Raspberry Pi), Cactus Graph (zero-copy computation graph, like PyTorch for mobile), and Cactus Engine (high-level inference API with OpenAI-compatible chat completion, tool calling, auto-RAG, vision, audio, transcription, embeddings, and cloud handoff). -Supported models include Gemma 4 (E2B, E4B with vision + audio), Gemma 3/3n, Qwen 3/3.5, LiquidAI LFM2/LFM2.5 (incl. VL and MoE), Whisper, Moonshine, Parakeet (CTC/TDT), and Nomic Embed. Models run at INT4/INT8 precision with Apple NPU acceleration on supported hardware. Pre-converted weights at https://huggingface.co/Cactus-Compute. +Supported models include Gemma 4 (E2B, E4B with vision + audio), Gemma 3/3n, Qwen 3/3.5, LiquidAI LFM2/LFM2.5 (incl. VL and MoE), Whisper, Moonshine, Parakeet (CTC/TDT), and Nomic Embed. Models run at INT4/INT8 precision with Metal GPU acceleration on supported hardware. Pre-converted weights at https://huggingface.co/Cactus-Compute. ## Docs @@ -44,7 +44,7 @@ Supported models include Gemma 4 (E2B, E4B with vision + audio), Gemma 3/3n, Qwe - **Thinking**: Models that support chain-of-thought (Gemma 4) return a `thinking` field when `enable_thinking_if_supported` is set - **TurboQuant-H**: Hadamard rotation + Lloyd-Max codebook quantization for 2-bit per-layer embeddings, reducing embedding-dominated models by 40% - **INT4 Quantization**: Lossy weight quantization for ~50% memory reduction with minimal quality loss -- **NPU Acceleration**: Apple Neural Engine support for vision, audio, and transcription model encoders +- **Metal Acceleration**: Apple GPU support for faster inference on Apple silicon ## Links diff --git a/python/README.md b/python/README.md index 6bf7e0dff..616f4344c 100644 --- a/python/README.md +++ b/python/README.md @@ -20,9 +20,9 @@ cactus build --python ```bash -# Download pre-built bundles (defaults to --weights general, the portable build) +# Download pre-built bundles cactus download LiquidAI/LFM2-VL-450M -cactus download openai/whisper-small --weights apple # CoreML/NPU variant +cactus download openai/whisper-small # Optional: set your Cactus Cloud API key for automatic cloud fallback cactus auth @@ -61,8 +61,8 @@ from cactus import ensure_model, get_bundle_dir bundle = ensure_model("openai/whisper-tiny") # Or resolve the expected on-disk location explicitly -bundle_dir = get_bundle_dir("openai/whisper-tiny", bits=4, platform=None) -# -> Path("weights/whisper-tiny-cq4") (or `-cq4-apple` with platform="apple") +bundle_dir = get_bundle_dir("openai/whisper-tiny", bits=4) +# -> Path("weights/whisper-tiny-cq4") ``` ### Init / Lifecycle diff --git a/python/cactus/bindings/cactus.py b/python/cactus/bindings/cactus.py index 5eae1feee..638ff4070 100644 --- a/python/cactus/bindings/cactus.py +++ b/python/cactus/bindings/cactus.py @@ -1488,8 +1488,7 @@ class Graph: CQ3 = 5 CQ4 = 6 CPU = 0 - NPU = 1 - METAL = 2 + METAL = 1 ACT_SILU = 0 ACT_GELU = 1 ACT_GELU_ERF = 2 diff --git a/python/cactus/cli/__init__.py b/python/cactus/cli/__init__.py index 960b3c095..f65aeaf18 100644 --- a/python/cactus/cli/__init__.py +++ b/python/cactus/cli/__init__.py @@ -7,16 +7,8 @@ DEFAULT_TRANSCRIPTION_MODEL_ID, DEFAULT_TEST_MODEL_ID, DEFAULT_TEST_TRANSCRIPTION_MODEL_ID, - SUPPORTED_WEIGHTS_VARIANTS, ) from .download import cmd_download - -_WEIGHTS_CHOICES = ("general", *SUPPORTED_WEIGHTS_VARIANTS) -_WEIGHTS_HELP = ( - f"weights bundle variant: general = portable bundle for every platform (default); " - f"or a vendor variant: {', '.join(SUPPORTED_WEIGHTS_VARIANTS) if SUPPORTED_WEIGHTS_VARIANTS else '(none yet)'}" -) -_WEIGHTS_PIPE = "|".join(_WEIGHTS_CHOICES) from .compile import cmd_build from .serve import cmd_serve from .transcribe import cmd_transcribe @@ -44,8 +36,6 @@ def _build_parent(): p = argparse.ArgumentParser(add_help=False) p.add_argument("--bits", type=int, choices=[1, 2, 3, 4], default=4, help="CQ quantization (default: 4)") - p.add_argument("--weights", choices=_WEIGHTS_CHOICES, default="general", - help=_WEIGHTS_HELP) p.add_argument("--token", help="HuggingFace token") p.add_argument("--reconvert", action="store_true", help="Force local rebuild from source") @@ -128,7 +118,6 @@ def create_parser(): cactus run [model|path] run a model (default: {DEFAULT_MODEL_ID}) --bits 1|2|3|4 CQ quantization (default: 4) - --weights {_WEIGHTS_PIPE:<23} weights bundle variant (default: general) --image image file for VLM inference --audio audio file for audio chat --system system prompt @@ -141,19 +130,16 @@ def create_parser(): --file audio file to transcribe (WAV) --language language code (default: en) --bits 1|2|3|4 CQ quantization (default: 4) - --weights {_WEIGHTS_PIPE:<23} weights bundle variant (default: general) --token HuggingFace token (gated models) --reconvert force local rebuild from source cactus download [model] fetch a prebuilt bundle, else build locally (default: {DEFAULT_MODEL_ID}) --bits 1|2|3|4 CQ quantization (default: 4) - --weights {_WEIGHTS_PIPE:<23} weights bundle variant (default: general) --token HuggingFace token (gated models) --reconvert force local rebuild from source cactus convert [dir] build a runnable bundle locally (skips prebuilt fetch) --bits 1|2|3|4 CQ quantization (default: 4) - --weights {_WEIGHTS_PIPE:<23} weights bundle variant (default: general) --token HuggingFace token (gated models) --reconvert force local rebuild from source --lora merge a LoRA adapter before converting @@ -162,7 +148,6 @@ def create_parser(): --host bind address (default: 127.0.0.1) --port port (default: 8080) --bits 1|2|3|4 CQ quantization (default: 4) - --weights {_WEIGHTS_PIPE:<23} weights bundle variant (default: general) --token HuggingFace token (gated models) --reconvert force local rebuild from source --no-cloud-handoff disable automatic cloud handoff @@ -182,7 +167,6 @@ def create_parser(): --model default: {DEFAULT_TEST_MODEL_ID} --transcription-model default: {DEFAULT_TEST_TRANSCRIPTION_MODEL_ID} --bits 1|2|3|4 CQ quantization (default: 4) - --weights {_WEIGHTS_PIPE:<23} weights bundle variant (default: general) --token HuggingFace token (gated models) --reconvert force local rebuild of test models --suite run a single test suite from any @@ -196,7 +180,6 @@ def create_parser(): --model default: {DEFAULT_TEST_MODEL_ID} --transcription-model default: {DEFAULT_TEST_TRANSCRIPTION_MODEL_ID} --bits 1|2|3|4 CQ quantization (default: 4) - --weights {_WEIGHTS_PIPE:<23} weights bundle variant (default: general) --backend auto|cpu|metal inference backend (default: auto) --ios run on connected iPhone --android run on connected Android @@ -209,13 +192,12 @@ def create_parser(): Advanced / build pipeline — convert runs automatically inside run, serve, transcribe and download; reach for it only to control the - build (custom flags, LoRA, NPU, debugging): + build (custom flags, LoRA, debugging): cactus convert [dir] HuggingFace -> runnable cactus bundle (quantizes weights to CQ, then builds the runtime graph) --bits 1|2|3|4 CQ quantization (default: 4) - --weights {_WEIGHTS_PIPE:<23} weights bundle variant (default: general) --token HuggingFace token (defaults to $HF_TOKEN) --reconvert force weight conversion from source --lora merge a LoRA adapter before converting @@ -246,10 +228,6 @@ def create_parser(): --no-fuse-attention-block disable attention-block fusion --no-fuse-add-clipped disable add-clipped fusion --no-fuse-gated-deltanet disable gated DeltaNet fusion - --npu also emit CoreML .mlpackage(s) for NPU - --npu-quantize 0|4|8 force both NPU encoders to this quant - --npu-audio-quantize 0|4|8 audio encoder quant (default int8) - --npu-vision-quantize 0|4|8 vision encoder quant (default fp16) ----------------------------------------------------------------- """ @@ -396,8 +374,6 @@ def create_parser(): help="Output directory (default: weights/)") convert_parser.add_argument("--bits", type=int, choices=[1, 2, 3, 4], default=4, help="CQ quantization bits (default: 4)") - convert_parser.add_argument("--weights", choices=_WEIGHTS_CHOICES, default="general", - help=_WEIGHTS_HELP) convert_parser.add_argument("--token", help="HuggingFace token") convert_parser.add_argument("--lora", help="Path or HF id of a LoRA adapter to merge before converting (requires `peft`)") @@ -463,14 +439,6 @@ def create_parser(): help="Disable add-clipped fusion") convert_parser.add_argument("--no-fuse-gated-deltanet", action="store_true", help="Disable gated DeltaNet fusion") - convert_parser.add_argument("--npu", action="store_true", - help="Also emit CoreML .mlpackage(s) for Apple NPU encoders") - convert_parser.add_argument("--npu-quantize", type=int, choices=[0, 4, 8], default=None, - help="Legacy: force both NPU encoders to same quant (0=fp16, 4=int4, 8=int8)") - convert_parser.add_argument("--npu-audio-quantize", type=int, choices=[0, 4, 8], default=None, - help="NPU audio encoder quant: 0=fp16, 4=int4, 8=int8 (default: 8)") - convert_parser.add_argument("--npu-vision-quantize", type=int, choices=[0, 4, 8], default=None, - help="NPU vision encoder quant: 0=fp16, 4=int4, 8=int8 (default: 0; int4 degrades Gemma4 vision)") convert_parser.add_argument("--cache-context-length", default=None, help="KV cache context length for cached decode graphs (default: model config)") diff --git a/python/cactus/cli/code.py b/python/cactus/cli/code.py index b8362ebf2..7d38b7077 100644 --- a/python/cactus/cli/code.py +++ b/python/cactus/cli/code.py @@ -130,7 +130,6 @@ def cmd_code(args) -> int: model_flags_given = ( getattr(args, "reconvert", False) or getattr(args, "bits", 4) != 4 - or getattr(args, "weights", "general") != "general" or bool(getattr(args, "token", None)) ) if args.serve_model: @@ -146,7 +145,6 @@ def cmd_code(args) -> int: sys.executable, "-m", "cactus", "serve", serve_model, "--host", args.host, "--port", str(args.port), "--no-access-log", "--bits", str(getattr(args, "bits", 4)), - "--weights", getattr(args, "weights", "general"), "--backend", getattr(args, "backend", "auto"), ] if getattr(args, "token", None): diff --git a/python/cactus/cli/common.py b/python/cactus/cli/common.py index 345dc8e81..dd2afe714 100644 --- a/python/cactus/cli/common.py +++ b/python/cactus/cli/common.py @@ -39,11 +39,6 @@ def transpiled_root() -> Path: DEFAULT_TEST_TRANSCRIPTION_MODEL_ID = DEFAULT_TRANSCRIPTION_MODEL_ID -# Vendor-specific weight-bundle variants; add a new vendor by appending its -# name here. "general" (the portable bundle) is always available and default. -SUPPORTED_WEIGHTS_VARIANTS: tuple[str, ...] = ("apple",) - - RED = '\033[0;31m' GREEN = '\033[0;32m' YELLOW = '\033[1;33m' diff --git a/python/cactus/cli/compile.py b/python/cactus/cli/compile.py index 429cbdb39..c49a0648a 100644 --- a/python/cactus/cli/compile.py +++ b/python/cactus/cli/compile.py @@ -97,7 +97,6 @@ def build_binary( "-o", name, *curl_link, "-framework", "Accelerate", - "-framework", "CoreML", "-framework", "Foundation", "-framework", "Metal", "-framework", "MetalPerformanceShaders", diff --git a/python/cactus/cli/convert.py b/python/cactus/cli/convert.py index 6b9191d26..f6da1b1db 100644 --- a/python/cactus/cli/convert.py +++ b/python/cactus/cli/convert.py @@ -3,7 +3,7 @@ from pathlib import Path from .common import GREEN, RED, YELLOW, print_color -from .download import get_bundle_dir, get_weights_dir, resolve_weights_variant +from .download import get_bundle_dir, get_weights_dir def _merge_lora_adapter(base_model_id, lora_path, token=None): @@ -56,19 +56,14 @@ def cmd_convert(args): return 1 source_model_id = merged_dir - platform = resolve_weights_variant(getattr(args, "weights", "general")) output_dir = args.output_dir or str( - get_bundle_dir(args.model_id, bits=args.bits, platform=platform) + get_bundle_dir(args.model_id, bits=args.bits) ) - if platform == "apple" and not getattr(args, "npu", False): - args.npu = True - try: ensure_weights( source_model_id, bits=args.bits, - platform=platform, token=args.token, reconvert=args.reconvert, output_dir=output_dir, @@ -164,14 +159,6 @@ def cmd_transpile(args): extra_args.append("--no-fuse-add-clipped") if args.no_fuse_gated_deltanet: extra_args.append("--no-fuse-gated-deltanet") - if args.npu: - extra_args.append("--npu") - if args.npu_quantize is not None: - extra_args.extend(["--npu-quantize", str(args.npu_quantize)]) - if args.npu_audio_quantize is not None: - extra_args.extend(["--npu-audio-quantize", str(args.npu_audio_quantize)]) - if args.npu_vision_quantize is not None: - extra_args.extend(["--npu-vision-quantize", str(args.npu_vision_quantize)]) if args.cache_context_length is not None: extra_args.extend(["--cache-context-length", str(args.cache_context_length)]) diff --git a/python/cactus/cli/download.py b/python/cactus/cli/download.py index 14192e34b..be4a8f13c 100644 --- a/python/cactus/cli/download.py +++ b/python/cactus/cli/download.py @@ -3,7 +3,6 @@ from .common import ( BLUE, GREEN, - SUPPORTED_WEIGHTS_VARIANTS, print_color, weights_root, ) @@ -16,27 +15,16 @@ def get_weights_dir(model_id: str) -> Path: return weights_root() / get_model_dir_name(model_id) -def get_bundle_dir(model_id: str, *, bits: int = 4, platform: str | None = None) -> Path: +def get_bundle_dir(model_id: str, *, bits: int = 4) -> Path: from .utils import variant_suffix - return weights_root() / f"{get_model_dir_name(model_id)}-{variant_suffix(bits, platform)}" - - -def resolve_weights_variant(choice: str | None) -> str | None: - """Map the --weights choice to the bundle variant suffix: general (the - default on every platform) selects the portable bundle (no suffix).""" - if choice in (None, "general"): - return None - if choice in SUPPORTED_WEIGHTS_VARIANTS: - return choice - raise ValueError(f"unknown weights variant {choice!r}; " - f"supported: general, {', '.join(SUPPORTED_WEIGHTS_VARIANTS)}") + return weights_root() / f"{get_model_dir_name(model_id)}-{variant_suffix(bits)}" def ensure_model(model_id: str) -> Path: return download_bundle(model_id) -def download_bundle(model_id: str, *, bits: int = 4, platform: str | None = None, +def download_bundle(model_id: str, *, bits: int = 4, token: str | None = None, output_dir: Path | None = None) -> Path: from .utils import ( download_cq_archive, @@ -49,10 +37,10 @@ def download_bundle(model_id: str, *, bits: int = 4, platform: str | None = None repo_id = suggested_cq_repo(model_id) local_name = get_model_dir_name(model_id) - bundle_dir = Path(output_dir) if output_dir else get_bundle_dir(model_id, bits=bits, platform=platform) + bundle_dir = Path(output_dir) if output_dir else get_bundle_dir(model_id, bits=bits) revision = resolve_weight_revision(repo_id, token=token) - label = variant_suffix(bits, platform) + label = variant_suffix(bits) print() print_color(BLUE, f"Fetching {repo_id} [{label}] @ {revision or 'main'}") @@ -60,7 +48,7 @@ def download_bundle(model_id: str, *, bits: int = 4, platform: str | None = None if not archives: raise RuntimeError(f"no bundles published at {repo_id}") - resolution = resolve_archive(repo_id, local_name, archives, bits, platform=platform) + resolution = resolve_archive(repo_id, local_name, archives, bits) download_cq_archive(resolution, bundle_dir, token=token, revision=revision) print_color(GREEN, f"Ready at {bundle_dir}") return bundle_dir diff --git a/python/cactus/cli/model.py b/python/cactus/cli/model.py index 483067597..63bfc5cbc 100644 --- a/python/cactus/cli/model.py +++ b/python/cactus/cli/model.py @@ -3,7 +3,7 @@ import os import shutil -from dataclasses import dataclass, replace +from dataclasses import dataclass from pathlib import Path from .common import GREEN, PROJECT_ROOT, RED, YELLOW, print_color @@ -32,10 +32,10 @@ def _convert_from_source(model_id, *, bits, token, weights_dir, skip_model_load= return weights_dir -def ensure_weights(model_id, *, bits=4, platform=None, token=None, reconvert=False, output_dir=None, skip_model_load=False): +def ensure_weights(model_id, *, bits=4, token=None, reconvert=False, output_dir=None, skip_model_load=False): from .download import get_bundle_dir - weights_dir = Path(output_dir) if output_dir else get_bundle_dir(model_id, bits=bits, platform=platform) + weights_dir = Path(output_dir) if output_dir else get_bundle_dir(model_id, bits=bits) if reconvert and weights_dir.exists(): print_color(YELLOW, "Removing cached weights for reconversion...") @@ -191,14 +191,10 @@ class TranspileOptions: system_prompt: str | None = None trust_remote_code: bool = False local_files_only: bool = False - npu: bool = False - npu_quantize: int | None = None - npu_audio_quantize: int | None = None - npu_vision_quantize: int | None = None cache_context_length: str | int | None = None -def ensure_runnable_bundle(model_id, *, bits=4, platform=None, token=None, +def ensure_runnable_bundle(model_id, *, bits=4, token=None, reconvert=False, prebuilt=True, output_dir=None, transpile=None): """Resolve a runnable bundle via the full fallback ladder, building if needed. @@ -218,35 +214,31 @@ def ensure_runnable_bundle(model_id, *, bits=4, platform=None, token=None, if str(model_id).startswith(("/", "./", "../", "~")) and not Path(model_id).expanduser().exists(): raise RuntimeError(f"path not found: {model_id}") - cached = Path(output_dir) if output_dir else get_bundle_dir(model_id, bits=bits, platform=platform) + cached = Path(output_dir) if output_dir else get_bundle_dir(model_id, bits=bits) if _has_transpiled_bundle(cached) and not reconvert: return cached if prebuilt and not reconvert: try: - return download_bundle(model_id, bits=bits, platform=platform, + return download_bundle(model_id, bits=bits, token=token, output_dir=cached) except (RuntimeError, OSError) as exc: print_color(YELLOW, f"No prebuilt bundle ({exc}); building locally") opts = transpile or TranspileOptions() - if platform == "apple" and not opts.npu: - opts = replace(opts, npu=True) - return ensure_bundle(model_id, bits=bits, platform=platform, token=token, + return ensure_bundle(model_id, bits=bits, token=token, reconvert=reconvert, output_dir=cached, transpile=opts) def prepare_bundle(args, *, model_id=None, transpile=None, prebuilt=True, output_dir=None, fail_prefix="Model setup failed"): - """Resolve the weights variant from args and return a runnable bundle, with uniform - error handling shared by every model command. Returns the bundle Path, or - None (after printing the error) on failure.""" - from .download import resolve_weights_variant + """Return a runnable bundle, with uniform error handling shared by every + model command. Returns the bundle Path, or None (after printing the error) + on failure.""" try: return ensure_runnable_bundle( args.model_id if model_id is None else model_id, bits=getattr(args, "bits", 4), - platform=resolve_weights_variant(getattr(args, "weights", "general")), token=getattr(args, "token", None), reconvert=getattr(args, "reconvert", False), prebuilt=prebuilt, @@ -258,7 +250,7 @@ def prepare_bundle(args, *, model_id=None, transpile=None, prebuilt=True, return None -def ensure_bundle(model_id, *, bits=4, platform=None, token=None, +def ensure_bundle(model_id, *, bits=4, token=None, reconvert=False, output_dir=None, transpile=None, skip_model_load=False): from .download import get_bundle_dir @@ -270,10 +262,10 @@ def ensure_bundle(model_id, *, bits=4, platform=None, token=None, if output_dir is not None: output_dir = Path(output_dir).expanduser().resolve() else: - output_dir = get_bundle_dir(model_id, bits=bits, platform=platform) + output_dir = get_bundle_dir(model_id, bits=bits) ensure_weights( - model_id, bits=bits, platform=platform, token=token, + model_id, bits=bits, token=token, reconvert=reconvert, output_dir=output_dir, skip_model_load=skip_model_load, ) @@ -362,14 +354,6 @@ def ensure_bundle(model_id, *, bits=4, platform=None, token=None, extra_args.append("--trust-remote-code") if opts.local_files_only: extra_args.append("--local-files-only") - if opts.npu: - extra_args.append("--npu") - if opts.npu_quantize is not None: - extra_args.extend(["--npu-quantize", str(int(opts.npu_quantize))]) - if opts.npu_audio_quantize is not None: - extra_args.extend(["--npu-audio-quantize", str(int(opts.npu_audio_quantize))]) - if opts.npu_vision_quantize is not None: - extra_args.extend(["--npu-vision-quantize", str(int(opts.npu_vision_quantize))]) if opts.cache_context_length is not None: extra_args.extend(["--cache-context-length", str(opts.cache_context_length)]) diff --git a/python/cactus/cli/runtime.py b/python/cactus/cli/runtime.py index 54956a7c4..95cd3f2c3 100644 --- a/python/cactus/cli/runtime.py +++ b/python/cactus/cli/runtime.py @@ -75,7 +75,6 @@ def _link_python_runtime_library(*, static_lib, library_path): "-Wl,-install_name,@rpath/libcactus_engine.dylib", "-lcurl", "-framework", "Accelerate", - "-framework", "CoreML", "-framework", "Foundation", "-framework", "Metal", "-framework", "MetalPerformanceShaders", diff --git a/python/cactus/cli/upload.py b/python/cactus/cli/upload.py index 2a9e9486f..f38339659 100644 --- a/python/cactus/cli/upload.py +++ b/python/cactus/cli/upload.py @@ -4,7 +4,7 @@ from pathlib import Path from .common import BLUE, GREEN, PROJECT_ROOT, print_color -from .download import get_model_dir_name, resolve_weights_variant +from .download import get_model_dir_name from .model import prepare_bundle from .utils import suggested_cq_repo, variant_suffix @@ -15,9 +15,8 @@ def cmd_upload(args): if not token: raise SystemExit("missing HuggingFace token: pass --token or set HF_TOKEN") - platform = resolve_weights_variant(getattr(args, "weights", "general")) repo_id = suggested_cq_repo(args.model_id) - stem = f"{get_model_dir_name(args.model_id)}-{variant_suffix(args.bits, platform)}" + stem = f"{get_model_dir_name(args.model_id)}-{variant_suffix(args.bits)}" archive_name = f"{stem}.zip" tag = (PROJECT_ROOT / "CACTUS_VERSION").read_text(encoding="utf-8").strip() diff --git a/python/cactus/cli/utils.py b/python/cactus/cli/utils.py index 71882d016..b05f6c608 100644 --- a/python/cactus/cli/utils.py +++ b/python/cactus/cli/utils.py @@ -15,25 +15,24 @@ from typing import Iterable -CQ_VARIANT_RE = re.compile(r"-cq([1-4])(?:-([a-z]+))?$", re.IGNORECASE) +CQ_VARIANT_RE = re.compile(r"-cq([1-4])$", re.IGNORECASE) ARCHIVE_SUFFIXES = (".zip", ".tar", ".tar.gz", ".tgz") -def variant_suffix(bits: int, platform: str | None) -> str: - return f"cq{bits}-{platform.lower()}" if platform else f"cq{bits}" +def variant_suffix(bits: int) -> str: + return f"cq{bits}" @dataclass(frozen=True) class CqArchive: filename: str bits: int - platform: str | None = None # None = generic CPU bundle size: int | None = None sha256: str | None = None @property def suffix(self) -> str: - return variant_suffix(self.bits, self.platform) + return variant_suffix(self.bits) @dataclass(frozen=True) @@ -46,7 +45,7 @@ class CqResolution: def suggested_cq_repo(model_id: str) -> str: name = str(model_id).strip().replace("_", "-").split("/")[-1] - name = re.sub(r"-cq\d*(?:-[a-z]+)?$", "", name, flags=re.IGNORECASE) + name = re.sub(r"-cq\d*$", "", name, flags=re.IGNORECASE) return f"Cactus-Compute/{name}" @@ -57,13 +56,12 @@ def archive_stem(filename: str) -> str: return filename -def parse_cq_variant(filename: str) -> tuple[int, str | None] | None: +def parse_cq_variant(filename: str) -> int | None: stem = archive_stem(Path(filename).name) m = CQ_VARIANT_RE.search(stem) if not m: return None - platform = m.group(2).lower() if m.group(2) else None - return int(m.group(1)), platform + return int(m.group(1)) def is_supported_archive(filename: str) -> bool: @@ -71,10 +69,6 @@ def is_supported_archive(filename: str) -> bool: return any(lower.endswith(suffix) for suffix in ARCHIVE_SUFFIXES) -def _platform_sort_key(platform: str | None) -> tuple[int, str]: - return (0, "") if platform is None else (1, platform) - - def archives_from_repo_files(repo_files: Iterable[str], sizes: dict[str, int] | None = None, sha256s: dict[str, str] | None = None) -> tuple[CqArchive, ...]: sizes = sizes or {} @@ -83,13 +77,12 @@ def archives_from_repo_files(repo_files: Iterable[str], sizes: dict[str, int] | for filename in repo_files: if not is_supported_archive(filename): continue - variant = parse_cq_variant(filename) - if variant is None: + bits = parse_cq_variant(filename) + if bits is None: continue - bits, platform = variant - archives.append(CqArchive(filename=filename, bits=bits, platform=platform, + archives.append(CqArchive(filename=filename, bits=bits, size=sizes.get(filename), sha256=sha256s.get(filename))) - return tuple(sorted(archives, key=lambda a: (a.bits, _platform_sort_key(a.platform)))) + return tuple(sorted(archives, key=lambda a: a.bits)) VERSION_TAG_RE = re.compile(r"^v?(\d+(?:\.\d+){0,2})$", re.IGNORECASE) @@ -147,13 +140,13 @@ def resolve_weight_revision(repo_id: str, runtime_version: str | None = None, def resolve_archive(repo_id: str, local_name: str, archives: Iterable[CqArchive], - bits: int, platform: str | None = None) -> CqResolution: + bits: int) -> CqResolution: available = tuple(archives) if not available: raise RuntimeError(f"No CQ archives found in {repo_id}") - wanted = variant_suffix(bits, platform) - match = next((a for a in available if a.bits == bits and a.platform == platform), None) + wanted = variant_suffix(bits) + match = next((a for a in available if a.bits == bits), None) if not match: choices = ", ".join(a.suffix for a in available) raise RuntimeError(f"{wanted} not found in {repo_id}. Available: {choices}") diff --git a/python/cactus/server.py b/python/cactus/server.py index 0192c0e0b..7b38344b1 100644 --- a/python/cactus/server.py +++ b/python/cactus/server.py @@ -118,7 +118,7 @@ def _discover(self, root: Path) -> None: def _resolve(self, model_id: str) -> ModelInfo | None: """Match an exact bundle id, or fall back to the HuggingFace id / bare - stem (e.g. 'google/gemma-4-E2B-it' -> 'gemma-4-e2b-it-cq4-apple') when it + stem (e.g. 'google/gemma-4-E2B-it' -> 'gemma-4-e2b-it-cq4') when it is unambiguous.""" info = self.models.get(model_id) if info is not None: diff --git a/python/cactus/transpile/README.md b/python/cactus/transpile/README.md index 650e0686b..978a57e4b 100644 --- a/python/cactus/transpile/README.md +++ b/python/cactus/transpile/README.md @@ -667,9 +667,7 @@ cactus run ``` The runtime reads `components/manifest.json` from the bundle directory, mmaps -the `.weights` files, and loads the serialized `.cactus` graphs. Per-platform -`.mlpackage` files (Apple CoreML / future vendor platforms) are dispatched -to the matching hardware when present. +the `.weights` files, and loads the serialized `.cactus` graphs. ## Major Situations And How They Are Handled diff --git a/python/cactus/transpile/component_pipeline.py b/python/cactus/transpile/component_pipeline.py index 5cc9af49b..f161319cc 100644 --- a/python/cactus/transpile/component_pipeline.py +++ b/python/cactus/transpile/component_pipeline.py @@ -1,11 +1,10 @@ from __future__ import annotations import copy -from contextlib import AbstractContextManager from dataclasses import dataclass from dataclasses import field import re -from typing import Any, Callable +from typing import Any import numpy as np import torch @@ -30,10 +29,6 @@ class ComponentModuleSpec: graph_meta: dict[str, object] = field(default_factory=dict) metadata: dict[str, object] = field(default_factory=dict) dynamic_batch_axis: int | None = None - npu_module: torch.nn.Module | None = None - npu_runtime_input_count: int = 1 - npu_reparam: Callable[[torch.nn.Module], AbstractContextManager] | None = None - npu_example_inputs: tuple[torch.Tensor, ...] | None = None @dataclass diff --git a/python/cactus/transpile/hf_model.py b/python/cactus/transpile/hf_model.py index d3014c8cd..5a8a2d282 100644 --- a/python/cactus/transpile/hf_model.py +++ b/python/cactus/transpile/hf_model.py @@ -376,7 +376,6 @@ def _write_component_bundle( component_io_signatures: dict[str, dict[str, tuple[str, ...]]] | None = None, component_metadata: dict[str, dict[str, object]] | None = None, graph_filename: str = "graph.cactus", - npu_encoder_mlpackages: dict[str, str] | None = None, ) -> Path: bundle_dir = artifact_dir / "components" component_order = [ @@ -502,11 +501,6 @@ def _write_component_bundle( "inputs": _serialize_json_compatible(inputs_metadata), "components": manifest_components, } - for manifest_key, mlpackage_path in (npu_encoder_mlpackages or {}).items(): - if mlpackage_path: - manifest_payload[manifest_key] = mlpackage_path - if family in ("parakeet_tdt", "gemma4") and "npu_audio_encoder" in manifest_payload: - manifest_payload["npu_audio_compute_units"] = "CPU_AND_NE" _write_json(manifest_path, manifest_payload) return manifest_path @@ -633,24 +627,6 @@ def _run_component_pipeline_transpile( if weights_dir: print(f"weights_dir={weights_dir}") - npu_enabled = bool(getattr(args, "npu", False)) - npu_quantize = getattr(args, "npu_quantize", None) - npu_audio_quantize = getattr(args, "npu_audio_quantize", None) - npu_vision_quantize = getattr(args, "npu_vision_quantize", None) - npu_encoder_mlpackages: dict[str, str] = {} - if npu_enabled and artifact_dir is not None: - from .npu import run_encoder_pipeline - npu_encoder_mlpackages = run_encoder_pipeline( - component_specs, - artifact_dir, - enabled=True, - quantize_bits=npu_quantize, - audio_quantize_bits=npu_audio_quantize, - vision_quantize_bits=npu_vision_quantize, - ) - import gc as _gc - _gc.collect() - print("capture_begin=true", flush=True) captured_components = {} for spec in component_specs: @@ -807,7 +783,6 @@ def _run_component_pipeline_transpile( component_io_signatures=component_io_signatures, component_metadata=component_metadata, graph_filename=args.graph_filename, - npu_encoder_mlpackages=npu_encoder_mlpackages, ) print(f"saved_component_bundle_manifest={component_manifest_path}") for component in transpiled_component_graphs: @@ -2914,7 +2889,7 @@ def main() -> int: help=( "Instantiate the HF model on the meta device for graph capture instead of " "loading checkpoint tensors into RAM. Requires converted CQ weights and " - "only supports bundle generation, not execution or NPU emission." + "only supports bundle generation, not execution." ), ) parser.add_argument( @@ -2951,32 +2926,6 @@ def main() -> int: "(for example: vision_encoder,audio_encoder,lm_encoder,decoder)." ), ) - parser.add_argument( - "--npu", - action="store_true", - help="Also emit CoreML .mlpackage(s) for Apple Neural Engine audio + vision encoders.", - ) - parser.add_argument( - "--npu-quantize", - type=int, - default=None, - choices=[0, 4, 8], - help="Legacy override that forces BOTH audio and vision encoders to the same weight quant (0=fp16, 4=int4, 8=int8). When unset, per-component defaults apply: audio=int8, vision=fp16.", - ) - parser.add_argument( - "--npu-audio-quantize", - type=int, - default=None, - choices=[0, 4, 8], - help="Quantization for the audio encoder .mlpackage (0=fp16, 4=int4, 8=int8). Default int8: Conformer-style encoders absorb int8 quant noise well.", - ) - parser.add_argument( - "--npu-vision-quantize", - type=int, - default=None, - choices=[0, 4, 8], - help="Quantization for the vision encoder .mlpackage (0=fp16, 4=int4, 8=int8). Default fp16: ViT-style towers are small enough that fp16 is the safer correctness default — int4 visibly degrades on Gemma 4 vision.", - ) parser.add_argument("--no-fuse-gated-deltanet", action="store_true") parser.add_argument("--no-fuse-rms-norm", action="store_true") parser.add_argument("--no-fuse-rope", action="store_true") @@ -3008,8 +2957,6 @@ def main() -> int: if args.low_memory_load: if validated_weights_dir is None: raise RuntimeError("--low-memory-load requires --weights-dir with converted Cactus CQ weights") - if args.npu: - raise RuntimeError("--low-memory-load cannot emit NPU .mlpackage files because CoreML export needs real weights") if not args.skip_execute: print("note=low_memory_load_forces_skip_execute=true") args.skip_execute = True diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index fb6f09988..ba9e3bf04 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -16,7 +16,6 @@ from cactus.transpile.component_pipeline import ComponentModuleSpec from cactus.convert.cactus_adapters.tensor_io import CACTUS_MAGIC from cactus.convert.cactus_adapters.tensor_io import align_offset -from cactus.convert.model_adapters.naming import gemma4_scale_factor _GEMMA4_SAFE_TEXT_MLP_PRODUCT_SCALE = 1.0 / 64.0 @@ -1707,30 +1706,6 @@ def _gemma4_compute_native_like_image_features( return projected -_GEMMA4_NPU_PROJ_WEIGHTS = ( - ("embed_audio", "embed_audio_proj.weights"), - ("embed_vision", "embed_vision_proj.weights"), -) - - -@contextmanager -def _gemma4_npu_projection_reparam(module: torch.nn.Module): - backbone = getattr(module, "multimodal_backbone", None) - scaled: list[tuple[torch.nn.Linear, float]] = [] - for attr, cactus_name in _GEMMA4_NPU_PROJ_WEIGHTS: - projection = getattr(getattr(backbone, attr, None), "embedding_projection", None) - factor = gemma4_scale_factor(cactus_name) - if isinstance(projection, torch.nn.Linear) and factor != 1.0: - scaled.append((projection, factor)) - for projection, factor in scaled: - projection.weight.data.mul_(factor) - try: - yield - finally: - for projection, factor in scaled: - projection.weight.data.div_(factor) - - def _gemma4_compute_native_like_audio_features( multimodal_backbone: torch.nn.Module, input_features: torch.Tensor, @@ -4524,7 +4499,6 @@ def _require(component: str) -> None: raise RuntimeError("Gemma4 audio component requested but input_features/input_features_mask were not prepared") vision_encoder: Gemma4VisionEncoderAdapter | None = None - vision_encoder_npu: Gemma4VisionEncoderAdapter | None = None audio_encoder: Gemma4AudioEncoderAdapter | None = None native_image_soft_token_counts: tuple[int, ...] | None = None native_image_pool_shapes: tuple[tuple[int, int, int], ...] | None = None @@ -4545,12 +4519,6 @@ def _require(component: str) -> None: native_image_soft_token_counts=native_image_soft_token_counts, native_image_pool_shapes=native_image_pool_shapes, ).eval() - vision_encoder_npu = Gemma4VisionEncoderAdapter( - model, - weights_dir=weights_dir, - native_image_soft_token_counts=native_image_soft_token_counts, - native_image_pool_shapes=None, - ).eval() if "audio_encoder" in expanded_components: _gemma4_make_audio_mask_fp16_safe(model) audio_encoder = Gemma4AudioEncoderAdapter(model, weights_dir=weights_dir).eval() @@ -4703,7 +4671,7 @@ def _require(component: str) -> None: } specs: list[ComponentModuleSpec] = [] if "vision_encoder" in expanded_components: - if vision_encoder is None or vision_encoder_npu is None or pixel_values is None: + if vision_encoder is None or pixel_values is None: raise RuntimeError("Gemma4 vision_encoder spec requires prepared pixel_values") specs.append(ComponentModuleSpec( component="vision_encoder", @@ -4713,9 +4681,6 @@ def _require(component: str) -> None: output_keys=("image_features",), graph_meta={**common_graph_meta, "component": "vision_encoder"}, metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - npu_module=vision_encoder_npu, - npu_runtime_input_count=2, - npu_reparam=_gemma4_npu_projection_reparam, )) if "audio_encoder" in expanded_components: if audio_encoder is None or input_features is None or input_features_mask is None: @@ -4728,7 +4693,6 @@ def _require(component: str) -> None: output_keys=("audio_features",), graph_meta={**common_graph_meta, "component": "audio_encoder"}, metadata={"family": "gemma4", "task": "multimodal_causal_lm_logits"}, - npu_reparam=_gemma4_npu_projection_reparam, )) if "lm_encoder" in expanded_components: if image_features is None: @@ -7006,7 +6970,6 @@ def _ensure_decoder_inputs() -> tuple[torch.Tensor, ...]: output_keys=("last_hidden_state",), graph_meta={**common_graph_meta, "component": "vision_encoder"}, metadata=metadata, - npu_runtime_input_count=3, )) if "vision_projector" in expanded_components: specs.append(ComponentModuleSpec( diff --git a/python/cactus/transpile/npu/README.md b/python/cactus/transpile/npu/README.md deleted file mode 100644 index 8a51e1013..000000000 --- a/python/cactus/transpile/npu/README.md +++ /dev/null @@ -1,86 +0,0 @@ -# NPU Transpiler Pipeline - -Emits CoreML `.mlpackage`s so the runtime engine can dispatch through the -Apple Neural Engine. Triggered by `--npu` on `cactus convert` (forwarded to -`python -m cactus.transpile.hf_model --npu …`), or programmatically via -`TranspileOptions(npu=True)` in `cli/model.py`. Runtime side: -`cactus-engine/src/model_npu.cpp`. - -**Scope: audio + vision encoders only.** Text-decoder prefill is -intentionally *not* on NPU — CPU prefill is the supported path. - -| Entry point | Emits | manifest key | Runtime dispatch | -|---|---|---|---| -| `run_encoder_pipeline` → `emit_audio_encoder_mlpackage` (audio.py) | audio encoder | `npu_audio_encoder` | `audio_encode_via_npu` | -| `run_encoder_pipeline` → `emit_vision_encoder_mlpackage` (vision.py) | vision encoder | `npu_vision_encoder` | `vision_encode_via_npu` | - -`run_encoder_pipeline` reuses the exact `ComponentModuleSpec` adapter -modules + example inputs the graph transpiler captures, so the NPU encoder -stays in lockstep with the CPU component graph. Auxiliary adapter inputs -(masks, position ids) are baked into the exported model as constants — -the ANE requires a static single-input signature. - -## Status - -| Target | Works | Notes | -|---|---|---| -| Parakeet TDT audio encoder | **yes — verified end-to-end** | 4.5× faster warm than CPU on M-series (~155ms vs ~693ms TTFT). Use `CACTUS_ANE_AUDIO_COMPUTE_UNITS=ALL`. | -| LFM-VL vision encoder | **yes — verified end-to-end** | ~2× faster on image-specific TTFT overhead, ~4× faster on the vision encoder graph alone (the rest of TTFT is preprocessing + LM prefill, CPU-bound). | -| Gemma 4 audio + vision encoder emit | **emit verified (int8); mlpackage output matches HF reference** | Vision int8: cos_sim 0.9995 vs HF reference (MAE 0.018). Audio int8: cos_sim 0.9963 (MAE 0.0055 after applying the cactus `1/16` audio rescale, see `_gemma4_compute_native_like_audio_features`). Vision: 162 MB; audio: 282 MB. C++ runtime dispatch path is the same code that ships Parakeet ASR + LFM-VL chat — exercised end-to-end on those models in this session. Full Gemma 4 cactus runtime end-to-end requires a v2 bundle (`cactus convert`, ~15–20 GB peak), which doesn't fit on the 16 GB host where these mlpackages were produced. | -| `coremltools.optimize.linear_quantize_weights` | works | Per-component defaults: **audio int8, vision fp16**. ANE compute is fp16 internally either way; quant only affects weight storage + memory bandwidth. Audio encoders (Conformer-style) absorb int8 quant noise well — verified on Parakeet + Gemma 4. Vision encoders (ViT-style: SigLIP2, Gemma 4 vision) are more sensitive — int4 visibly corrupts Gemma 4 vision output; fp16 is the rock-solid default given the towers are small (one-shot per request, bandwidth not critical). Override via `--npu-quantize` (forces both) or `--npu-audio-quantize` / `--npu-vision-quantize` (per-component). | - -## Why not NPU text prefill? - -Earlier iterations also emitted a text-decoder prefill `.mlpackage` -(via the now-removed `prefill.py`). Removed because: - -1. Wrapping the full text decoder (~2B params for Gemma 4 E2B) needs - ~15–20 GB peak during `coremltools.convert`, which OOMs 16 GB hosts — - regardless of whether the encoder emit succeeds. -2. CPU chunked prefill (post-PR #659) is fast enough that ANE prefill - isn't a clear win for the v2 runtime. -3. `origin/main`'s NPU prefill mlpackages were pre-built externally, - not regenerated by a transpiler — there's no precedent for v2 to - regenerate them reliably. - -The C++ runtime side (`NPUPrefill` class, `prefill_via_npu`, `load_npu_prefill`, -the `ANEPrefill` Obj-C wrapper) has been removed in lockstep. - -## Known coremltools 9.0 patches - -`coremltools_patches.py` carries op patches for gaps the v2 encoders hit -during conversion: - -- `__and__` / `__or__` / `__xor__` — bitwise ops on bool tensors -- `new_ones` — `tensor.new_ones(shape)` -- `one_hot` — cast indices to int32 (Gemma 4 vision feeds fp32 ids) -- `unfold` — translates to `mb.sliding_windows` + a transpose so the - window-size dim lands at the *end* (matching PyTorch semantics; MIL - inserts it at `axis+1` by default, which breaks Gemma 4 audio's - chunked-local attention permute chain) -- `layer_norm` epsilon dtype unification (iOS17+ enforces gamma/beta/eps - same dtype; HF fp16-traced models hit this) -- `add_fp16_cast` MIL pass skipping `layer_norm` / `batch_norm` / etc. -- mb binop scalar-cast wrappers (sub/mul/add/div SDPA mask path). - *Only Var operands get a `mb.cast` injected* — numpy arrays/scalars - are retyped via `numpy.astype` so we don't accidentally materialize a - Var in the wrong block scope when downstream passes (e.g. - `divide_to_multiply`) call `mb.mul(..., before_op=op)` from inside a - nested block. Mis-scoping trips the "var not visible in the block" - validator and aborts conversion. - -Add more here as new architectures expose new gaps. - -## Where future NPU perf gains can come from - -1. **Weight quantization (int4 / int8)** — biggest single perf knob on - ANE. Smaller weights ⇒ less memory bandwidth ⇒ faster matmuls. - Wired via `--npu-quantize`. -2. **SDPA preservation** — `torch.nn.functional.scaled_dot_product_attention` - maps to a single MIL op when coremltools sees it intact. HF defaults - to `attn_implementation="sdpa"` for compatible models, so we benefit - automatically. -3. **Per-family adapter optimizations** — e.g., the Gemma 4 vision tower - could probably be split (CPU patch embed + NPU transformer + CPU - projector, matching main's `model_siglip2.cpp` pattern) for cleaner - memory usage. Not done; whole-encoder mlpackage works today. diff --git a/python/cactus/transpile/npu/__init__.py b/python/cactus/transpile/npu/__init__.py deleted file mode 100644 index e2fc7da08..000000000 --- a/python/cactus/transpile/npu/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .pipeline import run_encoder_pipeline - -__all__ = ["run_encoder_pipeline"] diff --git a/python/cactus/transpile/npu/audio.py b/python/cactus/transpile/npu/audio.py deleted file mode 100644 index 6d72bedee..000000000 --- a/python/cactus/transpile/npu/audio.py +++ /dev/null @@ -1,115 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Any - -import torch - - -class AudioEncoderWrapper(torch.nn.Module): - def __init__( - self, - audio_module: torch.nn.Module, - baked_inputs: tuple[torch.Tensor, ...] = (), - ): - super().__init__() - self.audio = audio_module - self._n_baked = len(baked_inputs) - for idx, tensor in enumerate(baked_inputs): - self.register_buffer(f"_baked_{idx}", tensor, persistent=False) - - def forward(self, input_features: torch.Tensor) -> torch.Tensor: - extra = tuple(getattr(self, f"_baked_{i}") for i in range(self._n_baked)) - return self.audio(input_features, *extra) - - -def _import_coremltools() -> Any: - import coremltools as ct - from .coremltools_patches import apply_all_coremltools_patches - apply_all_coremltools_patches() - return ct - - -def _apply_weight_quantization(mlmodel: Any, bits: int) -> Any: - try: - from coremltools.optimize.coreml import ( - linear_quantize_weights, - OpLinearQuantizerConfig, - OptimizationConfig, - ) - op_config = OpLinearQuantizerConfig( - mode="linear_symmetric", - dtype=f"int{bits}", - granularity="per_channel", - ) - config = OptimizationConfig(global_config=op_config) - return linear_quantize_weights(mlmodel, config) - except Exception as exc: - print(f"npu.audio: weight quantization to int{bits} failed ({type(exc).__name__}: {exc}); keeping FP16 weights") - return mlmodel - - -def emit_audio_encoder_mlpackage( - audio_module: torch.nn.Module, - bundle_dir: Path, - *, - example_input: torch.Tensor, - baked_inputs: tuple[torch.Tensor, ...] = (), - filename: str = "audio_encoder.mlpackage", - input_name: str = "x", - output_name: str = "encoded", - minimum_deployment_target: str = "iOS17", - quantize_bits: int | None = None, -) -> str | None: - ct = _import_coremltools() - - wrapper = AudioEncoderWrapper(audio_module, baked_inputs) - wrapper.eval() - - try: - with torch.no_grad(): - exported = torch.export.export(wrapper, (example_input,)) - exported = exported.run_decompositions({}) - except Exception as exc: - print(f"npu.audio: torch.export failed ({type(exc).__name__}: {exc}); skipping mlpackage emit") - return None - - del wrapper - import gc as _gc - _gc.collect() - - if quantize_bits == 4: - minimum_deployment_target = "iOS18" - target_attr = getattr(ct.target, minimum_deployment_target, None) or ct.target.iOS17 - - from .coremltools_patches import build_cactus_pass_pipeline - try: - mlmodel = ct.convert( - exported, - inputs=[ct.TensorType(name=input_name, shape=tuple(example_input.shape))], - outputs=[ct.TensorType(name=output_name)], - compute_precision=ct.precision.FLOAT16, - convert_to="mlprogram", - minimum_deployment_target=target_attr, - pass_pipeline=build_cactus_pass_pipeline(), - ) - except Exception as exc: - print(f"npu.audio: coremltools.convert failed ({type(exc).__name__}: {exc})") - return None - - if quantize_bits is not None: - before_id = id(mlmodel) - mlmodel = _apply_weight_quantization(mlmodel, quantize_bits) - if id(mlmodel) != before_id: - print(f"npu.audio: applied int{quantize_bits} weight quantization") - - bundle_dir.mkdir(parents=True, exist_ok=True) - out_path = bundle_dir / filename - try: - mlmodel.save(str(out_path)) - except Exception as exc: - print(f"npu.audio: mlpackage save failed ({type(exc).__name__}: {exc})") - return None - - print(f"npu.audio: wrote {out_path} (input_shape={tuple(example_input.shape)})") - return filename diff --git a/python/cactus/transpile/npu/coremltools_patches.py b/python/cactus/transpile/npu/coremltools_patches.py deleted file mode 100644 index 9bb0de1cd..000000000 --- a/python/cactus/transpile/npu/coremltools_patches.py +++ /dev/null @@ -1,434 +0,0 @@ -from __future__ import annotations - -_APPLIED = False - - -def apply_all_coremltools_patches() -> None: - global _APPLIED - if _APPLIED: - return - try: - _patch_register_func_allow_dunder() - _register_new_ones_op() - _register_logical_and_op() - _override_layer_norm_translator() - _override_one_hot_translator() - _register_unfold_op() - _register_noop_assert_ops() - _patch_pipeline_remove_fuse_prelu() - _patch_mb_binops_scalar_cast() - _patch_fp16_cast_skip_layer_norm() - _APPLIED = True - print("npu.coremltools_patches: applied") - except Exception as exc: - print(f"npu.coremltools_patches: failed to apply ({type(exc).__name__}: {exc})") - - -def _patch_register_func_allow_dunder() -> None: - from coremltools.converters.mil.frontend.torch import torch_op_registry - - cls = torch_op_registry.TorchOpsRegistry - if getattr(cls.register_func, "_cactus_dunder_patched", False): - return - - - def register_func(self, func=None, torch_alias=None, override=False): - f_name = func.__name__ - all_f_names = [f_name] - if torch_alias is not None: - all_f_names.extend(torch_alias) - for name in all_f_names: - is_dunder = name.startswith("__") and name.endswith("__") - if name.endswith("_") and not is_dunder: - raise Exception( - f'Attempting to register "{name}" op. Do not register inplace ops.' - ) - if not override and name in self.name_to_func_mapping: - raise ValueError(f"Torch op {name} already registered.") - self.set_func_by_name(func, name) - - register_func._cactus_dunder_patched = True - cls.register_func = register_func - - -def _register_new_ones_op() -> None: - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( - _TORCH_OPS_REGISTRY, - ) - from coremltools.converters.mil import Builder as mb - from coremltools.converters.mil.frontend.torch.ops import _get_inputs - - if "new_ones" in _TORCH_OPS_REGISTRY.name_to_func_mapping: - return - - def _to_int32(v): - if isinstance(v, list): - casted = [_to_int32(x) for x in v] - return mb.concat(values=casted, axis=0) - dtype = getattr(v, "dtype", None) - name = getattr(dtype, "__name__", str(dtype) if dtype else "") - if "int32" in name: - return v - return mb.cast(x=v, dtype="int32") - - def new_ones(context, node): - inputs = _get_inputs(context, node, min_expected=2) - size = _to_int32(inputs[1]) - res = mb.fill(shape=size, value=1.0, name=node.name) - context.add(res, node.name) - - _TORCH_OPS_REGISTRY.set_func_by_name(new_ones, "new_ones") - - -def _register_logical_and_op() -> None: - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( - _TORCH_OPS_REGISTRY, - ) - from coremltools.converters.mil import Builder as mb - - def _force_bool(v): - dtype = getattr(v, "dtype", None) - name = getattr(dtype, "__name__", str(dtype) if dtype else "") - if "bool" in name: - return v - if "fp16" in name or "fp32" in name or "fp64" in name or "double" in name or "float" in name: - non_zero = mb.not_equal(x=v, y=0.0) - return non_zero - if "int" in name: - non_zero = mb.not_equal(x=v, y=0) - return non_zero - return mb.cast(x=v, dtype="bool") - - def make_logical(op_kind): - def _impl(context, node): - inputs = [context[i] for i in node.inputs] - x = _force_bool(inputs[0]) - y = _force_bool(inputs[1]) - if op_kind == "and": - out = mb.logical_and(x=x, y=y, name=node.name) - elif op_kind == "or": - out = mb.logical_or(x=x, y=y, name=node.name) - else: - out = mb.logical_xor(x=x, y=y, name=node.name) - context.add(out, node.name) - return _impl - - for tag, op_kind in [("__and_", "and"), ("__or_", "or"), ("__xor_", "xor")]: - if tag not in _TORCH_OPS_REGISTRY.name_to_func_mapping: - _TORCH_OPS_REGISTRY.set_func_by_name(make_logical(op_kind), tag) - for tag, op_kind in [("__and__", "and"), ("__or__", "or"), ("__xor__", "xor")]: - if tag not in _TORCH_OPS_REGISTRY.name_to_func_mapping: - _TORCH_OPS_REGISTRY.set_func_by_name(make_logical(op_kind), tag) - for tag in ("bitwise_and", "and"): - _TORCH_OPS_REGISTRY.set_func_by_name(make_logical("and"), tag) - for tag in ("bitwise_or", "or"): - _TORCH_OPS_REGISTRY.set_func_by_name(make_logical("or"), tag) - - -def _override_layer_norm_translator() -> None: - import numpy as np - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( - _TORCH_OPS_REGISTRY, - ) - from coremltools.converters.mil.frontend.torch.ops import _get_inputs - from coremltools.converters.mil import Builder as mb - - def _dtype_name(v): - d = getattr(v, "dtype", None) - return getattr(d, "__name__", str(d) if d else "") - - def _to_typed_const(v, target): - if v is None: - return v - np_dtype = np.float16 if target == "fp16" else np.float32 - if isinstance(v, (int, float)) and not isinstance(v, bool): - return mb.const(val=np_dtype(v)) - if hasattr(v, "val") and v.val is not None and hasattr(v, "dtype"): - try: - return mb.const(val=np.asarray(v.val, dtype=np_dtype)) - except Exception: - pass - if hasattr(v, "dtype") and _dtype_name(v) != target: - return mb.cast(x=v, dtype=target) - return v - - def layer_norm(context, node): - inputs = _get_inputs(context, node, min_expected=2) - nargs = len(inputs) - x = inputs[0] - normalized_shape = inputs[1] - weight = inputs[2] if nargs > 2 else None - bias = inputs[3] if nargs > 3 else None - eps = inputs[4] if nargs > 4 else None - if eps is None: - eps = 1e-5 - - ref = weight if weight is not None else x - ref_name = _dtype_name(ref) - target = "fp16" if ("fp16" in ref_name or "float16" in ref_name) else ( - "fp32" if ("fp32" in ref_name or "float32" in ref_name) else "fp16" - ) - - if x is not None and hasattr(x, "dtype") and _dtype_name(x) != target: - x = mb.cast(x=x, dtype=target) - weight = _to_typed_const(weight, target) - bias = _to_typed_const(bias, target) - eps = _to_typed_const(eps, target) - - out = mb.layer_norm( - x=x, - axes=list(range(-len(normalized_shape.val), 0)), - gamma=weight, - beta=bias, - epsilon=eps, - name=node.name, - ) - - if node.kind == "native_layer_norm": - context.add((out, None, None), torch_name=node.name) - else: - context.add(out) - - _TORCH_OPS_REGISTRY.set_func_by_name(layer_norm, "layer_norm") - _TORCH_OPS_REGISTRY.set_func_by_name(layer_norm, "native_layer_norm") - - -def _patch_pipeline_remove_fuse_prelu() -> None: - return - - -_PASSES_TO_DROP = [ - "common::fuse_prelu", -] - - -def build_cactus_pass_pipeline(): - from coremltools.converters.mil.mil.passes.pass_pipeline import PassPipeline - pipeline = PassPipeline.DEFAULT - for pass_name in _PASSES_TO_DROP: - try: - pipeline.remove_passes([pass_name]) - except Exception: - pass - return pipeline - - -def _patch_mb_binops_scalar_cast() -> None: - import numpy as np - from coremltools.converters.mil import Builder as mb - - def _dtype_name(v): - d = getattr(v, "dtype", None) - return getattr(d, "__name__", str(d) if d else "") - - def _is_fp(name): - return "fp" in name or "float" in name or "double" in name - - def _np_for(dtype_name): - if "fp16" in dtype_name or "float16" in dtype_name: - return np.float16 - if "fp32" in dtype_name or "float32" in dtype_name: - return np.float32 - if "fp64" in dtype_name or "float64" in dtype_name or "double" in dtype_name: - return np.float64 - return np.float32 - - def _is_var(v): - return hasattr(v, "op") and hasattr(v, "sym_type") - - def _retype_in_place(v, target_dtype_name): - np_dtype = _np_for(target_dtype_name) - if isinstance(v, np.ndarray): - return v.astype(np_dtype) - if isinstance(v, (int, float)) and not isinstance(v, bool): - return np_dtype(v) - try: - return np.asarray(v, dtype=np_dtype) - except Exception: - return v - - def _scalar_const(s, dtype_name): - return mb.const(val=_np_for(dtype_name)(s)) - - for op_name in ("sub", "mul", "add", "div", "real_div"): - original = getattr(mb, op_name, None) - if original is None or getattr(original, "_cactus_scalar_cast_patched", False): - continue - - def make_wrapper(orig): - def wrapper(**kwargs): - x = kwargs.get("x") - y = kwargs.get("y") - xn = _dtype_name(x) - yn = _dtype_name(y) - if xn and not yn and _is_fp(xn) and isinstance(y, (int, float)) and not isinstance(y, bool): - kwargs["y"] = _scalar_const(y, xn) - elif yn and not xn and _is_fp(yn) and isinstance(x, (int, float)) and not isinstance(x, bool): - kwargs["x"] = _scalar_const(x, yn) - elif xn and yn and xn != yn: - target = "fp16" if (xn == "fp16" or yn == "fp16") else xn - if xn != target: - kwargs["x"] = mb.cast(x=x, dtype=target) if _is_var(x) else _retype_in_place(x, target) - if yn != target: - kwargs["y"] = mb.cast(x=y, dtype=target) if _is_var(y) else _retype_in_place(y, target) - return orig(**kwargs) - wrapper._cactus_scalar_cast_patched = True - wrapper.__name__ = orig.__name__ - return wrapper - - setattr(mb, op_name, make_wrapper(original)) - - select_op = getattr(mb, "select", None) - if select_op is not None and not getattr(select_op, "_cactus_scalar_cast_patched", False): - def _select_wrap(orig): - def wrapper(**kwargs): - a = kwargs.get("a") - b = kwargs.get("b") - an = _dtype_name(a) - bn = _dtype_name(b) - if an and not bn and _is_fp(an) and isinstance(b, (int, float)) and not isinstance(b, bool): - kwargs["b"] = _scalar_const(b, an) - elif bn and not an and _is_fp(bn) and isinstance(a, (int, float)) and not isinstance(a, bool): - kwargs["a"] = _scalar_const(a, bn) - elif an and bn and an != bn: - target = "fp16" if (an == "fp16" or bn == "fp16") else an - if an != target: - kwargs["a"] = mb.cast(x=a, dtype=target) if _is_var(a) else _retype_in_place(a, target) - if bn != target: - kwargs["b"] = mb.cast(x=b, dtype=target) if _is_var(b) else _retype_in_place(b, target) - return orig(**kwargs) - wrapper._cactus_scalar_cast_patched = True - wrapper.__name__ = orig.__name__ - return wrapper - setattr(mb, "select", _select_wrap(select_op)) - - orig_ln = getattr(mb, "layer_norm", None) - if orig_ln is not None and not getattr(orig_ln, "_cactus_scalar_cast_patched", False): - def _ln_wrap(orig): - def wrapper(**kwargs): - gamma = kwargs.get("gamma") - gn = _dtype_name(gamma) - if _is_fp(gn): - eps = kwargs.get("epsilon") - if isinstance(eps, (int, float)) and not isinstance(eps, bool): - kwargs["epsilon"] = _scalar_const(eps, gn) - elif hasattr(eps, "dtype") and _dtype_name(eps) != gn: - kwargs["epsilon"] = mb.cast(x=eps, dtype=gn) - return orig(**kwargs) - wrapper._cactus_scalar_cast_patched = True - wrapper.__name__ = orig.__name__ - return wrapper - setattr(mb, "layer_norm", _ln_wrap(orig_ln)) - - for op_name in ("batch_norm", "instance_norm", "rms_norm"): - orig_op = getattr(mb, op_name, None) - if orig_op is None or getattr(orig_op, "_cactus_scalar_cast_patched", False): - continue - def _bn_wrap(orig): - def wrapper(**kwargs): - ref = kwargs.get("gamma") or kwargs.get("x") - rn = _dtype_name(ref) - if _is_fp(rn): - eps = kwargs.get("epsilon") - if isinstance(eps, (int, float)) and not isinstance(eps, bool): - kwargs["epsilon"] = _scalar_const(eps, rn) - elif hasattr(eps, "dtype") and _dtype_name(eps) != rn: - kwargs["epsilon"] = mb.cast(x=eps, dtype=rn) - return orig(**kwargs) - wrapper._cactus_scalar_cast_patched = True - wrapper.__name__ = orig.__name__ - return wrapper - setattr(mb, op_name, _bn_wrap(orig_op)) - - -def _patch_fp16_cast_skip_layer_norm() -> None: - try: - from coremltools.converters.mil.mil.passes.defs.quantization import ( - FP16ComputePrecision, - ) - except Exception as exc: - print(f"npu.coremltools_patches: skip_layer_norm patch unavailable ({exc})") - return - base = set(FP16ComputePrecision._UNSUPPORTED_FP16_OPS) - for op in ("layer_norm", "batch_norm", "instance_norm", "rms_norm"): - base.add(op) - FP16ComputePrecision._UNSUPPORTED_FP16_OPS = base - - -def _override_one_hot_translator() -> None: - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( - _TORCH_OPS_REGISTRY, - ) - from coremltools.converters.mil.frontend.torch.ops import _get_inputs, _get_kwinputs - from coremltools.converters.mil import Builder as mb - - def _dtype_name(v): - d = getattr(v, "dtype", None) - return getattr(d, "__name__", str(d) if d else "") - - def one_hot(context, node): - inputs = _get_inputs(context, node, expected=(1, 2)) - labels = inputs[0] - num_classes = inputs[1] if len(inputs) > 1 else -1 - num_classes = _get_kwinputs(context, node, "num_classes", default=[num_classes])[0] - if hasattr(num_classes, "val") and num_classes.val is not None: - num_classes = num_classes.val - - if hasattr(labels, "dtype") and _dtype_name(labels) != "int32": - labels = mb.cast(x=labels, dtype="int32") - - res = mb.one_hot(indices=labels, one_hot_vector_size=num_classes, name=node.name) - context.add(res) - - _TORCH_OPS_REGISTRY.set_func_by_name(one_hot, "one_hot") - - -def _register_noop_assert_ops() -> None: - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( - _TORCH_OPS_REGISTRY, - ) - - def _noop(context, node): - return - - _TORCH_OPS_REGISTRY.set_func_by_name(_noop, "_assert_tensor_metadata") - - -def _register_unfold_op() -> None: - from coremltools.converters.mil.frontend.torch.torch_op_registry import ( - _TORCH_OPS_REGISTRY, - ) - from coremltools.converters.mil.frontend.torch.ops import _get_inputs - from coremltools.converters.mil import Builder as mb - - def unfold(context, node): - inputs = _get_inputs(context, node, min_expected=4) - x = inputs[0] - dimension = inputs[1].val if hasattr(inputs[1], "val") else int(inputs[1]) - size = inputs[2].val if hasattr(inputs[2], "val") else int(inputs[2]) - step = inputs[3].val if hasattr(inputs[3], "val") else int(inputs[3]) - - rank = len(x.shape) if hasattr(x, "shape") else x.rank - dim = int(dimension) - if dim < 0: - dim += rank - - windowed = mb.sliding_windows( - x=x, - axis=dim, - size=int(size), - stride=int(step), - ) - - out_rank = rank + 1 - perm = list(range(out_rank)) - size_axis = dim + 1 - perm.pop(size_axis) - perm.append(size_axis) - if perm == list(range(out_rank)): - out = mb.identity(x=windowed, name=node.name) - else: - out = mb.transpose(x=windowed, perm=perm, name=node.name) - context.add(out) - - _TORCH_OPS_REGISTRY.set_func_by_name(unfold, "unfold") diff --git a/python/cactus/transpile/npu/pipeline.py b/python/cactus/transpile/npu/pipeline.py deleted file mode 100644 index e15b25cce..000000000 --- a/python/cactus/transpile/npu/pipeline.py +++ /dev/null @@ -1,94 +0,0 @@ -from __future__ import annotations - -from contextlib import nullcontext -from pathlib import Path - -from .audio import emit_audio_encoder_mlpackage -from .source import emit_source_encoder_mlpackage -from .vision import emit_vision_encoder_mlpackage - - -_ENCODER_COMPONENTS = { - "audio_encoder": (emit_audio_encoder_mlpackage, "audio_encoder.mlpackage"), - "source_encoder": (emit_source_encoder_mlpackage, "source_encoder.mlpackage"), - "vision_encoder": (emit_vision_encoder_mlpackage, "vision_encoder.mlpackage"), -} - - -def run_encoder_pipeline( - component_specs, - artifact_dir: Path, - *, - enabled: bool = True, - quantize_bits: int | None = None, - audio_quantize_bits: int | None = None, - vision_quantize_bits: int | None = None, -) -> dict[str, str]: - results: dict[str, str] = {} - if not enabled or not component_specs: - return results - - bundle_root = artifact_dir / "components" - bundle_root.mkdir(parents=True, exist_ok=True) - - def _resolve(per_component: int | None, default: int | None) -> int | None: - chosen = per_component if per_component is not None else ( - quantize_bits if quantize_bits is not None else default - ) - return None if chosen == 0 else chosen - - component_quants = { - "audio_encoder": _resolve(audio_quantize_bits, default=8), - "source_encoder": _resolve(None, default=None), - "vision_encoder": _resolve(vision_quantize_bits, default=None), - } - - for spec in component_specs: - component = getattr(spec, "component", None) - if component not in _ENCODER_COMPONENTS: - continue - emit_fn, filename = _ENCODER_COMPONENTS[component] - example_inputs = tuple( - getattr(spec, "npu_example_inputs", None) or getattr(spec, "example_inputs", ()) or () - ) - if not example_inputs: - print(f"npu.pipeline: {component} spec has no example inputs; skipping") - continue - module = spec.npu_module or spec.module - n_runtime = min(max(1, spec.npu_runtime_input_count), len(example_inputs)) - qbits = component_quants[component] - qdesc = f"int{qbits}" if qbits else "fp16" - print(f"npu.pipeline: emitting {component} quant={qdesc}") - reparam = getattr(spec, "npu_reparam", None) - with (reparam(module) if reparam else nullcontext()): - if component == "source_encoder": - emitted = emit_fn( - module, - bundle_root, - example_inputs=example_inputs, - input_names=tuple(getattr(spec, "input_keys", ()) or ()), - filename=filename, - quantize_bits=qbits, - ) - elif component == "vision_encoder": - names = ("x",) + tuple(spec.input_keys[1:n_runtime]) - emitted = emit_fn( - module, - bundle_root, - runtime_inputs=tuple(zip(names, example_inputs[:n_runtime])), - baked_inputs=example_inputs[n_runtime:], - filename=filename, - quantize_bits=qbits, - ) - else: - emitted = emit_fn( - module, - bundle_root, - example_input=example_inputs[0], - baked_inputs=example_inputs[1:], - filename=filename, - quantize_bits=qbits, - ) - if emitted: - results[f"npu_{component}"] = f"components/{emitted}" - return results diff --git a/python/cactus/transpile/npu/source.py b/python/cactus/transpile/npu/source.py deleted file mode 100644 index 8dad63283..000000000 --- a/python/cactus/transpile/npu/source.py +++ /dev/null @@ -1,118 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import torch -import torch.nn.functional as F - -from .audio import _apply_weight_quantization, _import_coremltools - - -class SourceEncoderWrapper(torch.nn.Module): - def __init__( - self, - source_module: torch.nn.Module, - input_dtypes: tuple[torch.dtype, ...], - ): - super().__init__() - self.source = source_module - self.input_dtypes = input_dtypes - - def forward(self, *args: torch.Tensor) -> torch.Tensor: - typed_args = [] - for arg, dtype in zip(args, self.input_dtypes, strict=True): - typed_args.append(arg.to(dtype=dtype) if arg.dtype != dtype else arg) - out = self.source(*typed_args) - if isinstance(out, tuple): - return out[0] - return out - - -def emit_source_encoder_mlpackage( - source_module: torch.nn.Module, - bundle_dir: Path, - *, - example_inputs: tuple[torch.Tensor, ...], - input_names: tuple[str, ...], - filename: str = "source_encoder.mlpackage", - output_name: str = "encoder_hidden_states", - minimum_deployment_target: str = "iOS18", - quantize_bits: int | None = None, -) -> str | None: - ct = _import_coremltools() - if not example_inputs: - print("npu.source: source_encoder spec has no example inputs; skipping") - return None - - input_dtypes = tuple(t.dtype for t in example_inputs) - exported_inputs = tuple( - t.to(dtype=torch.int32) if t.dtype in {torch.int64, torch.long} else t - for t in example_inputs - ) - wrapper = SourceEncoderWrapper(source_module, input_dtypes) - wrapper.eval() - - original_sdpa = F.scaled_dot_product_attention - - def _sdpa_without_gqa(query, key, value, *args, **kwargs): - enable_gqa = bool(kwargs.pop("enable_gqa", False)) - if enable_gqa and query.dim() >= 3 and key.dim() >= 3 and query.shape[-3] != key.shape[-3]: - repeats = int(query.shape[-3]) // max(int(key.shape[-3]), 1) - key = key.repeat_interleave(repeats, dim=-3) - value = value.repeat_interleave(repeats, dim=-3) - return original_sdpa(query, key, value, *args, **kwargs) - - try: - F.scaled_dot_product_attention = _sdpa_without_gqa - with torch.no_grad(): - exported = torch.export.export(wrapper, exported_inputs) - exported = exported.run_decompositions({}) - except Exception as exc: - print(f"npu.source: torch.export failed ({type(exc).__name__}: {exc}); skipping mlpackage emit") - return None - finally: - F.scaled_dot_product_attention = original_sdpa - - del wrapper - import gc as _gc - _gc.collect() - - target_attr = getattr(ct.target, minimum_deployment_target, None) or ct.target.iOS17 - input_types = [] - for name, tensor in zip(input_names, exported_inputs, strict=True): - dtype = np.int32 if tensor.dtype in {torch.int32, torch.int64, torch.long} else None - input_types.append(ct.TensorType(name=name, shape=tuple(tensor.shape), dtype=dtype)) - - from .coremltools_patches import build_cactus_pass_pipeline - try: - mlmodel = ct.convert( - exported, - inputs=input_types, - outputs=[ct.TensorType(name=output_name)], - compute_precision=ct.precision.FLOAT16, - convert_to="mlprogram", - minimum_deployment_target=target_attr, - pass_pipeline=build_cactus_pass_pipeline(), - ) - except Exception as exc: - print(f"npu.source: coremltools.convert failed ({type(exc).__name__}: {exc})") - return None - - if quantize_bits is not None: - before_id = id(mlmodel) - mlmodel = _apply_weight_quantization(mlmodel, quantize_bits) - if id(mlmodel) != before_id: - print(f"npu.source: applied int{quantize_bits} weight quantization") - - bundle_dir.mkdir(parents=True, exist_ok=True) - out_path = bundle_dir / filename - try: - mlmodel.save(str(out_path)) - except Exception as exc: - print(f"npu.source: mlpackage save failed ({type(exc).__name__}: {exc})") - return None - - shapes = ",".join(f"{name}={tuple(t.shape)}" for name, t in zip(input_names, exported_inputs, strict=True)) - print(f"npu.source: wrote {out_path} ({shapes})") - return filename diff --git a/python/cactus/transpile/npu/vision.py b/python/cactus/transpile/npu/vision.py deleted file mode 100644 index 97087efdd..000000000 --- a/python/cactus/transpile/npu/vision.py +++ /dev/null @@ -1,126 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Any - -import torch - - -class VisionEncoderWrapper(torch.nn.Module): - def __init__( - self, - vision_module: torch.nn.Module, - baked_inputs: tuple[torch.Tensor, ...] = (), - ): - super().__init__() - self.vision = vision_module - self._n_baked = len(baked_inputs) - for idx, tensor in enumerate(baked_inputs): - self.register_buffer(f"_baked_{idx}", tensor, persistent=False) - - def forward(self, *runtime_inputs: torch.Tensor) -> torch.Tensor: - coerced = tuple( - t if torch.is_floating_point(t) else t.to(torch.long) - for t in runtime_inputs - ) - extra = tuple(getattr(self, f"_baked_{i}") for i in range(self._n_baked)) - return self.vision(*coerced, *extra) - - -def _import_coremltools() -> Any: - import coremltools as ct - from .coremltools_patches import apply_all_coremltools_patches - apply_all_coremltools_patches() - return ct - - -def _apply_weight_quantization(mlmodel: Any, bits: int) -> Any: - try: - from coremltools.optimize.coreml import ( - linear_quantize_weights, - OpLinearQuantizerConfig, - OptimizationConfig, - ) - op_config = OpLinearQuantizerConfig( - mode="linear_symmetric", - dtype=f"int{bits}", - granularity="per_channel", - ) - config = OptimizationConfig(global_config=op_config) - return linear_quantize_weights(mlmodel, config) - except Exception as exc: - print(f"npu.vision: weight quantization to int{bits} failed ({type(exc).__name__}: {exc}); keeping FP16 weights") - return mlmodel - - -def emit_vision_encoder_mlpackage( - vision_module: torch.nn.Module, - bundle_dir: Path, - *, - runtime_inputs: tuple[tuple[str, torch.Tensor], ...], - baked_inputs: tuple[torch.Tensor, ...] = (), - filename: str = "vision_encoder.mlpackage", - output_name: str = "encoded", - minimum_deployment_target: str = "iOS18", - quantize_bits: int | None = None, -) -> str | None: - ct = _import_coremltools() - - example_tensors = tuple(tensor for _, tensor in runtime_inputs) - wrapper = VisionEncoderWrapper(vision_module, baked_inputs=baked_inputs) - wrapper.eval() - - try: - with torch.no_grad(): - exported = torch.export.export(wrapper, example_tensors) - exported = exported.run_decompositions({}) - except Exception as exc: - print(f"npu.vision: torch.export failed ({type(exc).__name__}: {exc}); skipping mlpackage emit") - return None - - del wrapper - import gc as _gc - _gc.collect() - - target_attr = getattr(ct.target, minimum_deployment_target, None) or ct.target.iOS17 - - import numpy as np - ct_inputs = [ - ct.TensorType(name=name, shape=tuple(tensor.shape)) - if torch.is_floating_point(tensor) - else ct.TensorType(name=name, shape=tuple(tensor.shape), dtype=np.int32) - for name, tensor in runtime_inputs - ] - - from .coremltools_patches import build_cactus_pass_pipeline - try: - mlmodel = ct.convert( - exported, - inputs=ct_inputs, - outputs=[ct.TensorType(name=output_name)], - compute_precision=ct.precision.FLOAT16, - convert_to="mlprogram", - minimum_deployment_target=target_attr, - pass_pipeline=build_cactus_pass_pipeline(), - ) - except Exception as exc: - print(f"npu.vision: coremltools.convert failed ({type(exc).__name__}: {exc})") - return None - - if quantize_bits is not None: - before_id = id(mlmodel) - mlmodel = _apply_weight_quantization(mlmodel, quantize_bits) - if id(mlmodel) != before_id: - print(f"npu.vision: applied int{quantize_bits} weight quantization") - - bundle_dir.mkdir(parents=True, exist_ok=True) - out_path = bundle_dir / filename - try: - mlmodel.save(str(out_path)) - except Exception as exc: - print(f"npu.vision: mlpackage save failed ({type(exc).__name__}: {exc})") - return None - - shapes = ", ".join(f"{name}={tuple(t.shape)}" for name, t in runtime_inputs) - print(f"npu.vision: wrote {out_path} (inputs: {shapes})") - return filename diff --git a/python/cactus/transpile/tdt_runtime.py b/python/cactus/transpile/tdt_runtime.py index 01275d7ec..9f0b25dfb 100644 --- a/python/cactus/transpile/tdt_runtime.py +++ b/python/cactus/transpile/tdt_runtime.py @@ -3,7 +3,6 @@ from dataclasses import dataclass from dataclasses import replace import json -import os from pathlib import Path import re from typing import Any @@ -416,17 +415,6 @@ def _relative_position_bias(query: torch.Tensor, relative_key: torch.Tensor, *, return gathered * float(scale) -def _relative_position_bias_static(query: torch.Tensor, relative_key: torch.Tensor, *, scale: float) -> torch.Tensor: - batch, heads, seq_len, _ = query.shape - pos_len = 2 * seq_len - 1 - scores = torch.matmul(query, relative_key.transpose(-1, -2)) - x = F.pad(scores, (1, 0)) - x = x.reshape(batch, heads, pos_len + 1, seq_len) - x = x[:, :, 1:, :] - x = x.reshape(batch, heads, seq_len, pos_len) - return x[:, :, :, :seq_len] * float(scale) - - class ParakeetTDTFeedForward(nn.Module): def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, torch.Tensor]): super().__init__() @@ -658,111 +646,6 @@ def forward(self, input_features: torch.Tensor) -> torch.Tensor: return x -class ParakeetTDTANESelfAttention(nn.Module): - def __init__(self, attn: "ParakeetTDTSelfAttention", seq_len: int): - super().__init__() - self.linear_q = attn.linear_q - self.linear_k = attn.linear_k - self.linear_v = attn.linear_v - self.linear_out = attn.linear_out - self.linear_pos = attn.linear_pos - self.pos_bias_u = attn.pos_bias_u - self.pos_bias_v = attn.pos_bias_v - self.config = attn.config - weight = attn.linear_pos.weight - rel_pos = _relative_position_embeddings( - seq_len=int(seq_len), - hidden_dim=self.config.hidden_dim, - device=weight.device, - dtype=weight.dtype, - ) - with torch.no_grad(): - rel_k = attn.linear_pos(rel_pos).view( - 1, 2 * int(seq_len) - 1, self.config.attention_heads, self.config.attention_head_dim - ).flip(1) - self.register_buffer("_rel_k", rel_k.detach(), persistent=False) - - def forward(self, x: torch.Tensor, key_bias: torch.Tensor | None = None) -> torch.Tensor: - batch, seq_len, _ = x.shape - num_heads = self.config.attention_heads - head_dim = self.config.attention_head_dim - - q = self.linear_q(x).view(batch, seq_len, num_heads, head_dim) - k = self.linear_k(x).view(batch, seq_len, num_heads, head_dim) - v = self.linear_v(x).view(batch, seq_len, num_heads, head_dim) - - q_u = q + self.pos_bias_u.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device) - q_v = q + self.pos_bias_v.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device) - - rel_k_heads = self._rel_k.to(dtype=x.dtype, device=x.device).permute(0, 2, 1, 3) - rel_bias = _relative_position_bias_static( - q_v.permute(0, 2, 1, 3), - rel_k_heads, - scale=self.config.attention_scale, - ) - if key_bias is not None: - rel_bias = rel_bias + key_bias - attn = F.scaled_dot_product_attention( - q_u.permute(0, 2, 1, 3), - k.permute(0, 2, 1, 3), - v.permute(0, 2, 1, 3), - attn_mask=rel_bias, - dropout_p=0.0, - is_causal=False, - ) - attn = attn.permute(0, 2, 1, 3).reshape(batch, seq_len, self.config.hidden_dim) - return self.linear_out(attn) - - -class ParakeetTDTANEEncoderLayer(nn.Module): - def __init__(self, layer: "ParakeetTDTEncoderLayer", seq_len: int): - super().__init__() - self.feed_forward1 = layer.feed_forward1 - self.feed_forward2 = layer.feed_forward2 - self.self_attn = ParakeetTDTANESelfAttention(layer.self_attn, seq_len) - self.conv = layer.conv - self.norm_feed_forward1 = layer.norm_feed_forward1 - self.norm_self_att = layer.norm_self_att - self.norm_conv = layer.norm_conv - self.norm_feed_forward2 = layer.norm_feed_forward2 - self.norm_out = layer.norm_out - self.config = layer.config - - def forward(self, x: torch.Tensor, conv_mask: torch.Tensor, key_bias: torch.Tensor) -> torch.Tensor: - x = x + 0.5 * self.feed_forward1(self.norm_feed_forward1(x), activation=self.config.encoder_hidden_act) - x = x + self.self_attn(self.norm_self_att(x), key_bias) - x = x + self.conv(self.norm_conv(x), conv_mask) - x = x + 0.5 * self.feed_forward2(self.norm_feed_forward2(x), activation=self.config.encoder_hidden_act) - return self.norm_out(x) - - -class ParakeetTDTANEEncoder(nn.Module): - def __init__(self, encoder: "ParakeetTDTEncoder", seq_len: int): - super().__init__() - self.pre_encode = encoder.pre_encode - self.layers = nn.ModuleList( - [ParakeetTDTANEEncoderLayer(layer, seq_len) for layer in encoder.layers] - ) - factor = max(1, int(encoder.layers[0].config.subsampling_factor)) - self._subsample_stages = max(0, factor.bit_length() - 1) - - def _frame_masks(self, input_features: torch.Tensor, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]: - valid = (input_features.abs().sum(-1) > 0).to(input_features.dtype).sum(-1, keepdim=True) - for _ in range(self._subsample_stages): - valid = torch.floor((valid - 1.0) / 2.0) + 1.0 - positions = torch.arange(seq_len, device=input_features.device, dtype=input_features.dtype) - enc_valid = (positions.view(1, 1, seq_len) < valid.view(-1, 1, 1)).to(input_features.dtype) - key_bias = (1.0 - enc_valid).unsqueeze(1) * -1e4 - return enc_valid, key_bias - - def forward(self, input_features: torch.Tensor) -> torch.Tensor: - x = self.pre_encode(input_features) - conv_mask, key_bias = self._frame_masks(input_features, x.shape[1]) - for layer in self.layers: - x = layer(x, conv_mask, key_bias) - return x - - class ParakeetTDTDecoderCell(nn.Module): def __init__(self, hidden_dim: int, prefix: str, state_dict: dict[str, torch.Tensor]): super().__init__() @@ -988,21 +871,6 @@ def build_parakeet_tdt_component_specs( "adapter_family": "parakeet_tdt", } - factor = max(1, int(model.config.subsampling_factor)) - npu_frames = int(os.environ.get("CACTUS_NPU_PARAKEET_FRAMES", "3000")) - if npu_frames <= 0 or npu_frames % factor != 0: - raise ValueError(f"CACTUS_NPU_PARAKEET_FRAMES must be a positive multiple of {factor}, got {npu_frames}") - npu_input_features = torch.zeros( - (int(input_features.shape[0]), npu_frames, int(input_features.shape[-1])), - device=input_features.device, - dtype=input_features.dtype, - ) - valid_frames = min(int(input_features.shape[1]), npu_frames) - npu_input_features[:, :valid_frames, :] = input_features[:, :valid_frames, :] - with torch.no_grad(): - ane_seq_len = int(model.encoder.pre_encode(npu_input_features).shape[1]) - ane_encoder = ParakeetTDTANEEncoder(model.encoder, ane_seq_len) - return [ ComponentModuleSpec( component="audio_encoder", @@ -1012,9 +880,6 @@ def build_parakeet_tdt_component_specs( output_keys=("encoder_hidden_states",), graph_meta={**common_graph_meta, "component": "audio_encoder"}, metadata={"family": "parakeet_tdt", "task": "tdt_transcription"}, - npu_module=ane_encoder, - npu_example_inputs=(npu_input_features,), - npu_runtime_input_count=1, ), ComponentModuleSpec( component="decoder", diff --git a/python/pyproject.toml b/python/pyproject.toml index f5c09791e..d42b37994 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -36,7 +36,6 @@ dependencies = [ "Pillow>=11.0.0,<12", "torchvision>=0.23.0,<1", "sentencepiece>=0.2.0", - "coremltools>=9.0; sys_platform == 'darwin'", ] [project.optional-dependencies] diff --git a/python/tests/bundles.py b/python/tests/bundles.py index 00f1548c9..5a5aaf950 100644 --- a/python/tests/bundles.py +++ b/python/tests/bundles.py @@ -21,8 +21,8 @@ def _valid_bundle(path: Path) -> bool: def _iter_bundle_candidates(name: str): """Yield the bare bundle dir for `name` plus any suffixed `name-cq*` variants - (e.g. `gemma-4-e2b-it-cq4`, `gemma-4-e2b-it-cq4-apple`) so callers that know a - model's bare stem still find bundles built under the suffixed convention.""" + (e.g. `gemma-4-e2b-it-cq4`) so callers that know a model's bare stem still + find bundles built under the suffixed convention.""" bare = WEIGHTS / name if bare.exists(): yield bare diff --git a/python/tests/test_bindings.py b/python/tests/test_bindings.py index 79ad255dd..a84f8d725 100644 --- a/python/tests/test_bindings.py +++ b/python/tests/test_bindings.py @@ -336,14 +336,8 @@ def test_run_rejects_bare_name(self): with pytest.raises(SystemExit): self.parser.parse_args(["run", "whisper-base", "--prompt", "hi"]) - def test_run_weights_flag(self): - args = self.parser.parse_args(["run", "Foo/Bar", "--weights", "apple"]) - assert args.weights == "apple" - assert args.bits == 4 - - def test_download_weights_flag(self): - args = self.parser.parse_args(["download", "Foo/Bar", "--weights", "general", "--bits", "2"]) - assert args.weights == "general" + def test_download_bits_flag(self): + args = self.parser.parse_args(["download", "Foo/Bar", "--bits", "2"]) assert args.bits == 2 diff --git a/python/tests/test_cli_run.py b/python/tests/test_cli_run.py index 97aa6a59c..15b976b40 100644 --- a/python/tests/test_cli_run.py +++ b/python/tests/test_cli_run.py @@ -39,7 +39,6 @@ def fake_subprocess_run(cmd): no_cloud_tele=False, model_id="org/model", bits=4, - weights="general", token=None, reconvert=False, system=None, diff --git a/python/tests/test_download.py b/python/tests/test_download.py index cd9d3bcc4..4591d94a8 100644 --- a/python/tests/test_download.py +++ b/python/tests/test_download.py @@ -45,13 +45,9 @@ def test_suggested_cq_repo_strips_existing_cq_suffix(self): ) def test_parse_cq_variant(self): - self.assertEqual(parse_cq_variant("gemma-4-E2B-it-cq1.zip"), (1, None)) - self.assertEqual(parse_cq_variant("gemma-4-E2B-it-cq4.zip"), (4, None)) - self.assertEqual(parse_cq_variant("gemma-4-E2B-it-CQ3.tar.gz"), (3, None)) - self.assertEqual(parse_cq_variant("gemma-4-E2B-it-cq4-apple.tar.gz"), (4, "apple")) - self.assertEqual(parse_cq_variant("model-cq2-APPLE.zip"), (2, "apple")) - # Forward-compatible: any vendor name matching [a-z]+ parses. - self.assertEqual(parse_cq_variant("model-cq4-qualcomm.tar.gz"), (4, "qualcomm")) + self.assertEqual(parse_cq_variant("gemma-4-E2B-it-cq1.zip"), 1) + self.assertEqual(parse_cq_variant("gemma-4-E2B-it-cq4.zip"), 4) + self.assertEqual(parse_cq_variant("gemma-4-E2B-it-CQ3.tar.gz"), 3) self.assertIsNone(parse_cq_variant("README.md")) self.assertIsNone(parse_cq_variant("L4V4A4.zip")) @@ -61,39 +57,28 @@ def test_archives_from_repo_files(self): "gemma-4-E2B-it-cq2.zip", "gemma-4-E2B-it-cq3.zip", "gemma-4-E2B-it-cq4.zip", - "gemma-4-E2B-it-cq4-apple.zip", "README.md", "config.json", ] sizes = {"gemma-4-E2B-it-cq4.zip": 500_000_000} archives = archives_from_repo_files(files, sizes=sizes) - self.assertEqual(len(archives), 5) + self.assertEqual(len(archives), 4) self.assertEqual(archives[0].bits, 1) - self.assertIsNone(archives[0].platform) - # Generic CPU bundle sorts BEFORE vendor-specific for the same bits. - self.assertEqual((archives[3].bits, archives[3].platform), (4, None)) - self.assertEqual((archives[4].bits, archives[4].platform), (4, "apple")) + self.assertEqual(archives[3].bits, 4) self.assertEqual(archives[3].size, 500_000_000) def test_resolve_archive_finds_match(self): - files = ["model-cq1.zip", "model-cq2.zip", "model-cq3.zip", "model-cq4.zip", "model-cq4-apple.zip"] + files = ["model-cq1.zip", "model-cq2.zip", "model-cq3.zip", "model-cq4.zip"] archives = archives_from_repo_files(files) resolution = resolve_archive("Cactus-Compute/model", "model", archives, 4) self.assertEqual(resolution.archive.filename, "model-cq4.zip") self.assertEqual(resolution.archive.bits, 4) - self.assertIsNone(resolution.archive.platform) - - apple = resolve_archive("Cactus-Compute/model", "model", archives, 4, platform="apple") - self.assertEqual(apple.archive.filename, "model-cq4-apple.zip") - self.assertEqual(apple.archive.platform, "apple") def test_resolve_archive_missing_variant_errors(self): files = ["model-cq1.zip", "model-cq2.zip"] archives = archives_from_repo_files(files) with self.assertRaisesRegex(RuntimeError, "cq4 not found"): resolve_archive("Cactus-Compute/model", "model", archives, 4) - with self.assertRaisesRegex(RuntimeError, "cq2-apple not found"): - resolve_archive("Cactus-Compute/model", "model", archives, 2, platform="apple") def test_resolve_archive_empty_errors(self): with self.assertRaisesRegex(RuntimeError, "No CQ archives"): diff --git a/python/tests/test_model.py b/python/tests/test_model.py index 3addbffd7..543243ae7 100644 --- a/python/tests/test_model.py +++ b/python/tests/test_model.py @@ -17,10 +17,8 @@ cactus_transcribe, ) from cactus.cli.model import ensure_bundle -from cactus.cli.download import resolve_weights_variant -_PLATFORM = resolve_weights_variant("general") _ASSETS_DIR = PROJECT_ROOT / "cactus-engine" / "tests" / "assets" _TEST_IMAGE = _ASSETS_DIR / "test_monkey.png" _TEST_AUDIO = _ASSETS_DIR / "test.wav" @@ -29,7 +27,7 @@ class TestVLMModel(unittest.TestCase): @classmethod def setUpClass(cls): - cls.weights_dir = ensure_bundle("LiquidAI/LFM2-VL-450M", platform=_PLATFORM) + cls.weights_dir = ensure_bundle("LiquidAI/LFM2-VL-450M") cls.model = cactus_init(str(cls.weights_dir), None, False) @classmethod @@ -66,7 +64,7 @@ class TestWhisperModel(unittest.TestCase): @classmethod def setUpClass(cls): - cls.weights_dir = ensure_bundle("openai/whisper-small", platform=_PLATFORM) + cls.weights_dir = ensure_bundle("openai/whisper-small") cls.model = cactus_init(str(cls.weights_dir), None, False) @classmethod @@ -107,7 +105,7 @@ class TestNomicEmbedding(unittest.TestCase): @classmethod def setUpClass(cls): - cls.weights_dir = ensure_bundle(cls.MODEL_ID, platform=_PLATFORM) + cls.weights_dir = ensure_bundle(cls.MODEL_ID) cls.model = cactus_init(str(cls.weights_dir), None, False) @classmethod diff --git a/python/tests/test_server_live.py b/python/tests/test_server_live.py index d75f2b738..e4091c334 100644 --- a/python/tests/test_server_live.py +++ b/python/tests/test_server_live.py @@ -25,7 +25,7 @@ def _default_llm_bundle() -> Path: """Locate the canonical gemma-4-e2b-it LLM bundle under whichever convention - it was built with (bare `gemma-4-e2b-it` or suffixed `...-cq4[-apple]`).""" + it was built with (bare `gemma-4-e2b-it` or suffixed `...-cq4`).""" for candidate in _iter_bundle_candidates("gemma-4-e2b-it"): if _valid_bundle(candidate) and _read_model_type(candidate) in LLM_TYPES: return candidate From 02df848428c57d6004e4590777edf9649a37eacb Mon Sep 17 00:00:00 2001 From: jakmro Date: Thu, 9 Jul 2026 01:10:21 +0200 Subject: [PATCH 2/4] fixes + cleanups --- bindings/rust/cactus.rs | 2 ++ cactus-code/package-lock.json | 34 ---------------------------------- docs/cactus_transpiler.md | 12 ++++++------ python/cactus/cli/__init__.py | 6 ++++++ python/cactus/cli/list.py | 15 ++++++++++----- 5 files changed, 24 insertions(+), 45 deletions(-) diff --git a/bindings/rust/cactus.rs b/bindings/rust/cactus.rs index f00ecaded..3a956a618 100644 --- a/bindings/rust/cactus.rs +++ b/bindings/rust/cactus.rs @@ -7,6 +7,8 @@ pub type CactusTokenCallback = Option; #[cfg_attr(target_os = "macos", link(name = "Accelerate", 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"))] #[cfg_attr(target_os = "macos", link(name = "Security", kind = "framework"))] #[cfg_attr(target_os = "macos", link(name = "SystemConfiguration", kind = "framework"))] diff --git a/cactus-code/package-lock.json b/cactus-code/package-lock.json index 87f1343bb..681a5975c 100644 --- a/cactus-code/package-lock.json +++ b/cactus-code/package-lock.json @@ -1336,29 +1336,6 @@ "dev": true, "license": "ISC" }, - "node_modules/ws": { - "version": "8.21.0", - "resolved": "https://registry.npmjs.org/ws/-/ws-8.21.0.tgz", - "integrity": "sha512-Vsp28b7DRcimFQvrqu2Wek3z1iYxDCWqHYB8Qsnk/S4RfaCQzPGPyBNuVjJV3cd6UiKtUtp6sNM77gWvzcCH+g==", - "license": "MIT", - "optional": true, - "peer": true, - "engines": { - "node": ">=10.0.0" - }, - "peerDependencies": { - "bufferutil": "^4.0.1", - "utf-8-validate": ">=5.0.2" - }, - "peerDependenciesMeta": { - "bufferutil": { - "optional": true - }, - "utf-8-validate": { - "optional": true - } - } - }, "node_modules/yaml": { "version": "2.9.0", "resolved": "https://registry.npmjs.org/yaml/-/yaml-2.9.0.tgz", @@ -1374,17 +1351,6 @@ "url": "https://github.com/sponsors/eemeli" } }, - "node_modules/zod": { - "version": "3.25.76", - "resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz", - "integrity": "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==", - "license": "MIT", - "optional": true, - "peer": true, - "funding": { - "url": "https://github.com/sponsors/colinhacks" - } - }, "packages/agent": { "name": "@earendil-works/pi-agent-core", "version": "0.80.2", diff --git a/docs/cactus_transpiler.md b/docs/cactus_transpiler.md index 3ba4ebb1d..784def696 100644 --- a/docs/cactus_transpiler.md +++ b/docs/cactus_transpiler.md @@ -227,21 +227,21 @@ cactus convert nvidia/parakeet-tdt-0.6b-v3 ./parakeet-weights --bits 4 \ ## Running Saved Bundles -Transpiled bundles are saved to `./weights//` by default (alongside their CQ -weights). Run them later without re-transpiling: +Transpiled bundles are saved to `./weights/-cq/` by default (alongside +their CQ weights). Run them later without re-transpiling: ```bash # Text -cactus run ./weights/qwen3-0.6b --prompt "Write a haiku" +cactus run ./weights/qwen3-0.6b-cq4 --prompt "Write a haiku" # Multimodal -cactus run ./weights/gemma-4-e2b-it \ +cactus run ./weights/gemma-4-e2b-it-cq4 \ --image photo.jpg \ --audio speech.wav \ --prompt "What do you see?" # Audio -cactus run ./weights/parakeet-tdt-0.6b-v3 \ +cactus run ./weights/parakeet-tdt-0.6b-v3-cq4 \ --audio meeting.wav ``` @@ -374,7 +374,7 @@ component. Use `loaded.reset()` between independent generations. A transpiled bundle looks like this: ``` -weights// +weights/-cq/ raw_ir.json # IR before optimization (debugging) optimized_ir.json # IR after fusion passes (debugging) graph.cactus # serialized runtime graph diff --git a/python/cactus/cli/__init__.py b/python/cactus/cli/__init__.py index f65aeaf18..6cf394f02 100644 --- a/python/cactus/cli/__init__.py +++ b/python/cactus/cli/__init__.py @@ -442,6 +442,12 @@ def create_parser(): convert_parser.add_argument("--cache-context-length", default=None, help="KV cache context length for cached decode graphs (default: model config)") + upload_parser = subparsers.add_parser("upload", + parents=[_build_parent()], + help="Build a runnable bundle locally and upload it to Cactus-Compute on HuggingFace") + upload_parser.add_argument("model_id", type=_hf_id_or_path, + help="HuggingFace model id (e.g. openai/whisper-base)") + return parser diff --git a/python/cactus/cli/list.py b/python/cactus/cli/list.py index fafebbe76..24cac392a 100644 --- a/python/cactus/cli/list.py +++ b/python/cactus/cli/list.py @@ -87,9 +87,14 @@ def _collect(roots): def cmd_list(_args): - roots = (weights_root(), transpiled_root()) - converted, transpiled = _categorize(roots) - _print_section("Converted weights (cactus convert)", BLUE, converted) - print() - _print_section("Runnable bundles (cactus convert / download)", CYAN, transpiled) + models = _collect((weights_root(), transpiled_root())) + print_color(CYAN, "Available models") + if not models: + print(" (none)") + return 0 + name_w = max(len(p.name) for p, _, _, _ in models) + type_w = max(len("type"), max(len(t) for _, t, _, _ in models)) + print(f" {'name':<{name_w}} {'type':<{type_w}} {'quant':<5} {'size':>10} location") + for p, model_type, quant, size in models: + print(f" {p.name:<{name_w}} {model_type:<{type_w}} {quant:<5} {_human_size(size):>10} {p.parent}") return 0 From e7fa115546b8489ec4169d3b86073b8a9326d9f5 Mon Sep 17 00:00:00 2001 From: jakmro Date: Thu, 9 Jul 2026 03:31:00 +0200 Subject: [PATCH 3/4] cleanups --- README.md | 10 +- cactus-engine/cactus_engine.h | 11 + cactus-engine/src/model.cpp | 12 +- cactus-engine/test.sh | 2 +- cactus-engine/tests/run.cpp | 4 +- cactus-engine/tests/test_utils.cpp | 2 +- cactus-graph/cactus_graph.h | 273 ++++++------ cactus-graph/src/builder.cpp | 412 +++++++++--------- cactus-graph/src/execute.cpp | 40 +- cactus-graph/src/graph_ffi.cpp | 27 ++ cactus-graph/src/io.cpp | 4 +- cactus-graph/src/metal_plan.cpp | 6 + cactus-graph/src/ops_cache.cpp | 2 +- cactus-graph/src/ops_nn.cpp | 2 +- cactus-graph/src/param_io.cpp | 15 +- cactus-graph/tests/test_metal_parity.cpp | 47 +- docs/cactus_graph.md | 3 +- python/README.md | 6 + python/cactus/__init__.py | 2 - python/cactus/bindings/cactus.py | 508 ++++++++++------------ python/cactus/cli/__init__.py | 18 +- python/cactus/cli/code.py | 3 +- python/cactus/cli/serve.py | 4 +- python/cactus/cli/test.py | 3 +- python/cactus/transpile/runtime_compat.py | 109 ++--- 25 files changed, 812 insertions(+), 713 deletions(-) diff --git a/README.md b/README.md index 1c5a4c621..c064037a8 100644 --- a/README.md +++ b/README.md @@ -208,7 +208,7 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ │ │ cactus run [model|path] run a model (downloads if needed) │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --backend auto|cpu|metal inference backend (default: auto) │ +│ --backend cpu|metal inference backend (default: auto) │ │ --image image file for VLM inference │ │ --audio audio file for audio chat │ │ --system system prompt │ @@ -242,7 +242,7 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ --host bind address (default: 127.0.0.1) │ │ --port port (default: 8080) │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --backend auto|cpu|metal inference backend (default: auto) │ +│ --backend cpu|metal inference backend (default: auto) │ │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild from source │ │ --no-cloud-handoff disable automatic cloud handoff │ @@ -252,7 +252,7 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ cactus code run the AI coding agent (TUI / print) │ │ --serve-model auto-start a server with this model │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --backend auto|cpu|metal inference backend (default: auto) │ +│ --backend cpu|metal inference backend (default: auto) │ │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild from source │ │ --host server address (default: 127.0.0.1) │ @@ -276,7 +276,7 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ --model default: google/gemma-4-E2B-it │ │ --transcription-model default: nvidia/parakeet-tdt-0.6b-v3 │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --backend auto|cpu|metal inference backend (default: auto) │ +│ --backend cpu|metal inference backend (default: auto) │ │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild of test models │ │ --suite run a single test suite by name │ @@ -291,7 +291,7 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ --model default: google/gemma-4-E2B-it │ │ --transcription-model default: nvidia/parakeet-tdt-0.6b-v3 │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --backend auto|cpu|metal inference backend (default: auto) │ +│ --backend cpu|metal inference backend (default: auto) │ │ --ios run on connected iPhone │ │ --android run on connected Android │ │ │ diff --git a/cactus-engine/cactus_engine.h b/cactus-engine/cactus_engine.h index 5dacb0bf6..a60e4aec2 100644 --- a/cactus-engine/cactus_engine.h +++ b/cactus-engine/cactus_engine.h @@ -110,6 +110,17 @@ CACTUS_FFI_EXPORT int cactus_transcribe( size_t pcm_buffer_size ); +CACTUS_FFI_EXPORT int cactus_preprocess_audio_features( + const char* audio_file_path, + const char* model_type, + size_t mel_bins, + float* features_buffer, + size_t buffer_size, + size_t* feature_count, + size_t* out_mel_bins, + size_t* out_frames +); + CACTUS_FFI_EXPORT cactus_stream_transcribe_t cactus_stream_transcribe_start( cactus_model_t model, const char* options_json diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 6f9eb1c7e..8cbb7c1dc 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -771,7 +771,7 @@ bool Model::load_component_graph(Component& comp) { } void Model::unload_component_graph(Component& comp) { - if (cactus_backend_metal()) return; + if (cactus_default_backend() == ComputeBackend::METAL) return; if (comp.graph) { comp.graph->release_runtime_buffers(); comp.graph->release_all_weight_pages(); @@ -863,14 +863,14 @@ void Model::run_step(uint32_t token_id, size_t position, bool /*read_logits*/, b 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)); - if (!use_fused || !cactus_backend_metal()) cactus_graph_set_prefill_consistent(true); + if (!use_fused || cactus_default_backend() != ComputeBackend::METAL) cactus_graph_set_prefill_consistent(true); decoder_->graph->execute(); cactus_graph_set_prefill_consistent(false); maybe_capture_handoff_probe_hidden(*decoder_); return; } if (!use_fused) cactus_graph_set_prefill_consistent(true); - if (use_fused && cactus_backend_metal()) { + if (use_fused && cactus_default_backend() == ComputeBackend::METAL) { if (ple_probe_state_ == 0) ple_probe_state_ = encoder_->graph->extract_ple_pathway(fused_embed_ctx_) ? 1 : 2; if (ple_probe_state_ == 1) { @@ -985,7 +985,7 @@ std::vector> Model::generate_batch(const std::vector 1 || cactus_backend_metal()) { + if (batch > 1 || cactus_default_backend() == ComputeBackend::METAL) { exec_batch = std::max(batch, decoder_cache_num_slots()); } reset_component_cache_states(*decoder_); @@ -1585,7 +1585,7 @@ void Model::run_vision_encoder(const std::string& image_path) { media_feature_shapes_[name] = desc.shape; media_feature_precisions_[name] = desc.precision; } - if (!cactus_backend_metal()) { + if (cactus_default_backend() != ComputeBackend::METAL) { vision_encoder_->graph->release_runtime_buffers(); vision_encoder_->graph->release_all_weight_pages(); } @@ -1791,7 +1791,7 @@ void Model::run_audio_encoder_messages(const std::vector>& au if (mel.empty()) continue; run_audio_encoder(mel); } - if (!cactus_backend_metal()) { + if (cactus_default_backend() != ComputeBackend::METAL) { audio_encoder_->graph->release_runtime_buffers(); audio_encoder_->graph->release_all_weight_pages(); } diff --git a/cactus-engine/test.sh b/cactus-engine/test.sh index 7d581662c..aa793087a 100755 --- a/cactus-engine/test.sh +++ b/cactus-engine/test.sh @@ -8,7 +8,7 @@ ASSETS_DIR="$(pwd)/tests/assets" IOS_MODE=false ANDROID_MODE=false SUITE="" -BACKEND="auto" +BACKEND="" while [[ $# -gt 0 ]]; do case $1 in diff --git a/cactus-engine/tests/run.cpp b/cactus-engine/tests/run.cpp index 58d47ce67..6e8f520b2 100644 --- a/cactus-engine/tests/run.cpp +++ b/cactus-engine/tests/run.cpp @@ -860,8 +860,8 @@ int main(int argc, char** argv) { std::cout << "Backend: " << (backend == "metal" ? "Metal GPU" : "CPU") << "\n"; else std::cout << "Metal not available; using CPU\n"; - } else if (!backend.empty() && backend != "auto") { - std::cerr << "Error: unknown backend '" << backend << "' (expected 'cpu', 'metal', or 'auto')\n"; + } else if (!backend.empty()) { + std::cerr << "Error: unknown backend '" << backend << "' (expected 'cpu' or 'metal')\n"; return 1; } diff --git a/cactus-engine/tests/test_utils.cpp b/cactus-engine/tests/test_utils.cpp index a18076bca..76c3ddd97 100644 --- a/cactus-engine/tests/test_utils.cpp +++ b/cactus-engine/tests/test_utils.cpp @@ -10,7 +10,7 @@ void apply_backend() { if (applied) return; applied = true; const char* b = std::getenv("CACTUS_TEST_BACKEND"); - if (!b || !*b || std::strcmp(b, "auto") == 0) return; + if (!b || !*b) return; if (cactus_set_backend(b) == 0) std::cout << "Backend: " << (std::strcmp(b, "metal") == 0 ? "Metal GPU" : "CPU") << "\n"; else diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index e1f59b38b..d43d11532 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -18,7 +18,6 @@ #include int cactus_backend_select(const char* backend); -bool cactus_backend_metal(); namespace cactus { @@ -84,7 +83,9 @@ class Logger { #define CACTUS_LOG_WARN(component, msg) CACTUS_LOG(cactus::LogLevel::WARN, component, msg) #define CACTUS_LOG_ERROR(component, msg) CACTUS_LOG(cactus::LogLevel::ERROR, component, msg) -enum class ComputeBackend { CPU, METAL }; +enum class ComputeBackend { CPU = 0, METAL = 1 }; + +ComputeBackend cactus_default_backend(); enum class Activation { SILU, GELU, GELU_ERF, RELU, SIGMOID, TANH }; @@ -334,7 +335,7 @@ struct OpParams { std::vector permutation; Precision output_precision = Precision::FP16; BroadcastInfo broadcast_info; - ComputeBackend backend = ComputeBackend::CPU; + ComputeBackend backend = cactus_default_backend(); size_t dilation = 1; size_t stride = 1; @@ -505,94 +506,96 @@ class CactusGraph { void set_external_input(size_t node_id, void* data, Precision precision); void* get_output(size_t node_id); - size_t add(size_t input1, size_t input2); - size_t add_clipped(size_t input1, size_t input2); - size_t subtract(size_t input1, size_t input2); - size_t multiply(size_t input1, size_t input2); - size_t divide(size_t input1, size_t input2); - size_t not_equal(size_t input1, size_t input2); - - - size_t scalar_add(size_t input, float value); - size_t scalar_subtract(size_t input, float value); - size_t scalar_multiply(size_t input, float value); - size_t scalar_divide(size_t input, float value); - size_t scalar_not_equal(size_t input, float value); - size_t scalar_exp(size_t input); - size_t scalar_sqrt(size_t input); - size_t scalar_cos(size_t input); - size_t scalar_sin(size_t input); - size_t scalar_log(size_t input); - - size_t abs(size_t input); - size_t pow(size_t input, float exponent); - size_t precision_cast(size_t input, Precision target_precision); - - size_t relu(size_t input); - size_t leaky_relu(size_t input, float negative_slope = 0.01f); - size_t clamp(size_t input, float lo, float hi); - size_t silu(size_t input); - size_t gelu(size_t input); - size_t gelu_erf(size_t input); - size_t sigmoid(size_t input); - size_t tanh(size_t input); - size_t glu(size_t input, int axis = -1); - - - size_t sum(size_t input, int axis); - size_t mean(size_t input, int axis); - size_t variance(size_t input, int axis); - size_t min(size_t input, int axis); - size_t max(size_t input, int axis); - size_t cumsum(size_t input, int axis); - size_t softmax(size_t input, int axis = -1); - size_t topk(size_t input, size_t k); - - size_t reshape(size_t input, const std::vector& new_shape); - size_t view(size_t input, const std::vector& new_shape); - size_t flatten(size_t input, int start_dim = 0, int end_dim = -1); - size_t transpose(size_t input, ComputeBackend backend = ComputeBackend::CPU); - size_t transposeN(size_t input, const std::vector& permutation, ComputeBackend backend = ComputeBackend::CPU); - size_t slice(size_t input, int axis, size_t start, size_t length); - size_t index(size_t input, size_t index_value, int dim); - size_t concat(size_t input1, size_t input2, int axis = 0); - size_t cat(const std::vector& inputs, int axis); + size_t add(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t add_clipped(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + size_t subtract(size_t input1, size_t input2, ComputeBackend backend = cactus_default_backend()); + 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 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_not_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 abs(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t pow(size_t input, float exponent, ComputeBackend backend = cactus_default_backend()); + size_t precision_cast(size_t input, Precision target_precision, ComputeBackend backend = cactus_default_backend()); + + size_t relu(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t leaky_relu(size_t input, float negative_slope = 0.01f, ComputeBackend backend = cactus_default_backend()); + size_t clamp(size_t input, float lo, float hi, ComputeBackend backend = cactus_default_backend()); + size_t silu(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t gelu(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t gelu_erf(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t sigmoid(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t tanh(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t glu(size_t input, int axis = -1, ComputeBackend backend = cactus_default_backend()); + + + void set_node_backend(size_t node_id, ComputeBackend backend); + + size_t sum(size_t input, int axis, ComputeBackend backend = cactus_default_backend()); + size_t mean(size_t input, int axis, ComputeBackend backend = cactus_default_backend()); + size_t variance(size_t input, int axis, ComputeBackend backend = cactus_default_backend()); + size_t min(size_t input, int axis, ComputeBackend backend = cactus_default_backend()); + size_t max(size_t input, int axis, ComputeBackend backend = cactus_default_backend()); + size_t cumsum(size_t input, int axis, ComputeBackend backend = cactus_default_backend()); + size_t softmax(size_t input, int axis = -1, ComputeBackend backend = cactus_default_backend()); + size_t topk(size_t input, size_t k, ComputeBackend backend = cactus_default_backend()); + + 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 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 index(size_t input, size_t index_value, int dim, 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()); size_t matmul( size_t input1, size_t input2, bool pretransposed_rhs = false, - ComputeBackend backend = ComputeBackend::CPU); + ComputeBackend backend = cactus_default_backend()); - size_t rms_norm(size_t input, size_t weight, float epsilon = 1e-5f); - size_t layernorm(size_t input, size_t weight, size_t bias, float epsilon = 1e-5f); - size_t layernorm(size_t input, size_t weight, float epsilon = 1e-5f); - size_t groupnorm(size_t input, size_t weight, size_t bias, size_t num_groups = 32, float epsilon = 1e-5f); - size_t batchnorm(size_t input, size_t weight, size_t bias, size_t running_mean, size_t running_var, int axis = 1, float epsilon = 1e-5f); + size_t rms_norm(size_t input, size_t weight, float epsilon = 1e-5f, ComputeBackend backend = cactus_default_backend()); + size_t layernorm(size_t input, size_t weight, size_t bias, float epsilon = 1e-5f, ComputeBackend backend = cactus_default_backend()); + size_t layernorm(size_t input, size_t weight, float epsilon = 1e-5f, ComputeBackend backend = cactus_default_backend()); + size_t groupnorm(size_t input, size_t weight, size_t bias, size_t num_groups = 32, float epsilon = 1e-5f, ComputeBackend backend = cactus_default_backend()); + size_t batchnorm(size_t input, size_t weight, size_t bias, size_t running_mean, size_t running_var, int axis = 1, float epsilon = 1e-5f, ComputeBackend backend = cactus_default_backend()); - size_t rope(size_t input, float theta, size_t position_offset = 0, ComputeBackend backend = ComputeBackend::CPU); - size_t rope_gptj(size_t input, float theta, size_t position_offset = 0, size_t rot_dim = 0, ComputeBackend backend = ComputeBackend::CPU); + size_t rope(size_t input, float theta, size_t position_offset = 0, ComputeBackend backend = cactus_default_backend()); + size_t rope_gptj(size_t input, float theta, size_t position_offset = 0, size_t rot_dim = 0, ComputeBackend backend = cactus_default_backend()); size_t attention(size_t query, size_t key, size_t value, float scale, - bool is_causal = true, ComputeBackend backend = ComputeBackend::CPU); + bool is_causal = true, ComputeBackend backend = cactus_default_backend()); size_t attention(size_t query, size_t key, size_t value, float scale, - size_t position_offset, ComputeBackend backend = ComputeBackend::CPU); + size_t position_offset, ComputeBackend backend = cactus_default_backend()); size_t attention(size_t query, size_t key, size_t value, float scale, - size_t position_offset, size_t window_size, ComputeBackend backend = ComputeBackend::CPU); + size_t position_offset, size_t window_size, ComputeBackend backend = cactus_default_backend()); size_t attention_masked( size_t query, size_t key, size_t value, size_t mask, float scale, - bool is_causal = true, ComputeBackend backend = ComputeBackend::CPU, + bool is_causal = true, ComputeBackend backend = cactus_default_backend(), bool additive_mask = false, size_t position_offset = 0, size_t window_size = 0, float logit_cap = 0.0f); - size_t rel_pos_bias(size_t query, size_t relative_key, float scale); + size_t rel_pos_bias(size_t query, size_t relative_key, float scale, ComputeBackend backend = cactus_default_backend()); size_t attention_int8_hybrid( size_t query, size_t key_new, size_t value_new, float scale, size_t position_offset, const int8_t* cached_keys, const int8_t* cached_values, const float* k_scales, const float* v_scales, size_t cache_len, size_t num_kv_heads, size_t head_dim, - size_t window_size = 0, size_t v_head_dim = 0); + size_t window_size = 0, size_t v_head_dim = 0, ComputeBackend backend = cactus_default_backend()); size_t kv_cache_state( size_t max_seq_len, @@ -600,14 +603,16 @@ class CactusGraph { size_t head_dim, size_t window_size = 0, size_t sink_size = 4, - size_t num_slots = 1); + size_t num_slots = 1, + ComputeBackend backend = cactus_default_backend()); size_t kv_cache_append( size_t new_kv, size_t cache_state_node, size_t window_size = 0, size_t sink_size = 4, - size_t cache_slot = 0); + size_t cache_slot = 0, + ComputeBackend backend = cactus_default_backend()); size_t attention_cached( size_t query, @@ -619,47 +624,48 @@ class CactusGraph { size_t position_offset = 0, size_t window_size = 0, size_t v_head_dim = 0, - size_t cache_slot = 0); - - size_t conv_cache_state(size_t window_size, size_t hidden_dim); - size_t conv_cache_append(size_t new_data, size_t cache_state_node); - size_t conv_cache_initialize(size_t rows, size_t cache_state_node); - - size_t recurrent_cache_state(const std::vector& shape, Precision precision); - size_t recurrent_cache_write(size_t new_value, size_t cache_state); - - size_t conv1d_causal(size_t input, size_t weight, size_t kernel_size, size_t dilation = 1); - size_t conv1d_k3(size_t input, size_t weight, size_t stride); - size_t conv1d_k7s3(size_t input, size_t weight, size_t bias); - size_t conv1d(size_t input, size_t weight, size_t stride); - size_t conv1d(size_t input, size_t weight, size_t bias, size_t stride); - size_t conv1d_same_depthwise_k9(size_t input, size_t weight); - size_t conv1d_same_depthwise_k9(size_t input, size_t weight, size_t bias); - size_t conv1d_pointwise(size_t input, size_t weight); - size_t conv1d_pointwise(size_t input, size_t weight, size_t bias); - size_t conv2d_k3s2p1(size_t input, size_t weight); - size_t conv2d_k3s2p1(size_t input, size_t weight, size_t bias); - size_t conv2d_depthwise_k3s2p1(size_t input, size_t weight); - size_t conv2d_depthwise_k3s2p1(size_t input, size_t weight, size_t bias); - size_t conv2d_pointwise_1x1(size_t input, size_t weight); - size_t conv2d_pointwise_1x1(size_t input, size_t weight, size_t bias); - size_t conv2d_k3s1p1(size_t input, size_t weight); - size_t conv2d_k3s1p1(size_t input, size_t weight, size_t bias); - size_t stft(size_t input, size_t weight, size_t stride, size_t num_fft_bins); - - size_t rfft(size_t input); - size_t irfft(size_t input, size_t output_length); + size_t cache_slot = 0, + ComputeBackend backend = cactus_default_backend()); + + size_t conv_cache_state(size_t window_size, size_t hidden_dim, ComputeBackend backend = cactus_default_backend()); + size_t conv_cache_append(size_t new_data, size_t cache_state_node, ComputeBackend backend = cactus_default_backend()); + size_t conv_cache_initialize(size_t rows, size_t cache_state_node, ComputeBackend backend = cactus_default_backend()); + + size_t recurrent_cache_state(const std::vector& shape, Precision precision, ComputeBackend backend = cactus_default_backend()); + 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_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_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()); + size_t conv1d_pointwise(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_k3s2p1(size_t input, size_t weight, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_k3s2p1(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_depthwise_k3s2p1(size_t input, size_t weight, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_depthwise_k3s2p1(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_pointwise_1x1(size_t input, size_t weight, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_pointwise_1x1(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_k3s1p1(size_t input, size_t weight, ComputeBackend backend = cactus_default_backend()); + size_t conv2d_k3s1p1(size_t input, size_t weight, size_t bias, ComputeBackend backend = cactus_default_backend()); + size_t stft(size_t input, size_t weight, size_t stride, size_t num_fft_bins, ComputeBackend backend = cactus_default_backend()); + + size_t rfft(size_t input, ComputeBackend backend = cactus_default_backend()); + size_t irfft(size_t input, size_t output_length, ComputeBackend backend = cactus_default_backend()); size_t mel_filter_bank( size_t num_frequency_bins, size_t num_mel_filters, float min_frequency, float max_frequency, size_t sampling_rate, - int norm_type = 1, int scale_type = 2); + int norm_type = 1, int scale_type = 2, ComputeBackend backend = cactus_default_backend()); size_t spectrogram( size_t waveform, size_t mel_filters_node, size_t frame_length, size_t hop_length, size_t fft_length, float power = 2.0f, bool center = true, int pad_mode = 0, float mel_floor = 1e-10f, int log_mel_mode = 0, float dither = 0.0f, float preemphasis = 0.0f, - bool remove_dc_offset = false); + bool remove_dc_offset = false, ComputeBackend backend = cactus_default_backend()); size_t image_preprocess( size_t pixel_input, @@ -667,57 +673,65 @@ class CactusGraph { int target_width, int target_height, int patch_size, int channels = 3, float rescale_factor = 1.0f / 255.0f, - const float* mean = nullptr, const float* std_dev = nullptr); + const float* mean = nullptr, const float* std_dev = nullptr, + ComputeBackend backend = cactus_default_backend()); - size_t bilinear_interpolation(size_t pos_embeds, size_t dst_height, size_t dst_width, bool align_corners = true); - size_t maxpool1d(size_t input, size_t kernel_size, size_t stride); + size_t bilinear_interpolation(size_t pos_embeds, size_t dst_height, size_t dst_width, bool align_corners = true, ComputeBackend backend = cactus_default_backend()); + size_t maxpool1d(size_t input, size_t kernel_size, size_t stride, ComputeBackend backend = cactus_default_backend()); size_t lstm_cell( size_t input, size_t h_prev, size_t c_prev, - size_t weight_ih, size_t weight_hh, size_t bias_ih, size_t bias_hh); + size_t weight_ih, size_t weight_hh, size_t bias_ih, size_t bias_hh, + ComputeBackend backend = cactus_default_backend()); size_t bilstm_sequence( size_t input, size_t w_ih_fwd, size_t w_hh_fwd, size_t b_ih_fwd, size_t b_hh_fwd, - size_t w_ih_bwd, size_t w_hh_bwd, size_t b_ih_bwd, size_t b_hh_bwd); + size_t w_ih_bwd, size_t w_hh_bwd, size_t b_ih_bwd, size_t b_hh_bwd, + ComputeBackend backend = cactus_default_backend()); size_t gated_deltanet_decode( size_t query, size_t key, size_t value, - size_t gate_log, size_t beta, size_t initial_state, float scale = 0.0f); + size_t gate_log, size_t beta, size_t initial_state, float scale = 0.0f, + ComputeBackend backend = cactus_default_backend()); size_t gated_deltanet_prefill( size_t query, size_t key, size_t value, size_t gate_log, size_t beta, size_t initial_state, - size_t chunk_size = 64, float scale = 0.0f); + size_t chunk_size = 64, float scale = 0.0f, + ComputeBackend backend = cactus_default_backend()); - size_t altup_predict(size_t coefs, const size_t* streams, size_t num_streams); - size_t altup_correct(size_t coefs, size_t innovation, const size_t* predictions, size_t num_predictions); - size_t gaussian_topk(size_t input, float ppf); + size_t altup_predict(size_t coefs, const size_t* streams, size_t num_streams, ComputeBackend backend = cactus_default_backend()); + size_t altup_correct(size_t coefs, size_t innovation, const size_t* predictions, size_t num_predictions, ComputeBackend backend = cactus_default_backend()); + size_t gaussian_topk(size_t input, float ppf, ComputeBackend backend = cactus_default_backend()); size_t moe_layer( size_t hidden, size_t routing_probs, size_t topk_indices, const std::vector& w1_weights, const std::vector& w3_weights, const std::vector& w2_weights, size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, - Activation activation = Activation::SILU, size_t per_expert_scale = 0); + Activation activation = Activation::SILU, size_t per_expert_scale = 0, + ComputeBackend backend = cactus_default_backend()); size_t moe_layer( size_t hidden, size_t routing_probs, size_t topk_indices, const std::vector& w1_weights, const std::vector& w2_weights, size_t num_experts, size_t num_experts_per_tok, bool normalize_routing, float epsilon, float routed_scaling_factor, - Activation activation); - 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); - size_t stats_pool(size_t input); - size_t weighted_stats_pool(size_t input, size_t weights); + 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 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()); size_t sample( size_t logits, float temperature = 0.6f, float top_p = 0.95f, size_t top_k = 20, - const std::unordered_map& logit_bias = {}); + const std::unordered_map& logit_bias = {}, + ComputeBackend backend = cactus_default_backend()); size_t sample_with_options( size_t logits, float temperature, float top_p, float min_p, float repetition_penalty, - size_t top_k, const std::unordered_map& logit_bias = {}); - size_t scatter_topk(size_t indices, size_t values, size_t num_classes); + size_t top_k, const std::unordered_map& logit_bias = {}, + 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); - size_t embedding(const std::string& filename, size_t indices); - size_t embedding(size_t embedding_tensor, size_t indices); + size_t gather(size_t embeddings, size_t indices, 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); size_t mmap_weights(const std::string& filename); void bind_mmap_weights(size_t node_id, const std::string& filename); @@ -730,7 +744,7 @@ class CactusGraph { void clear_buffer_pool(); void retain_outputs(const std::vector& node_ids); - size_t persistent(size_t source_node); + size_t persistent(size_t source_node, ComputeBackend backend = cactus_default_backend()); bool is_populated(size_t persistent_node_id) const; void invalidate_persistent(size_t persistent_node_id); @@ -771,6 +785,7 @@ class CactusGraph { private: size_t binary_broadcast_op(OpType op, size_t input1, size_t input2); + size_t tag_backend(size_t node_id, ComputeBackend backend); void infer_shapes(); size_t reduction_op(OpType op, size_t input, int axis); size_t attach_conv_bias(size_t node, size_t bias, size_t expected_size, const char* op_name); @@ -904,6 +919,10 @@ CACTUS_FFI_EXPORT cactus_graph_t cactus_graph_create(void); CACTUS_FFI_EXPORT void cactus_graph_destroy(cactus_graph_t graph); CACTUS_FFI_EXPORT int cactus_graph_hard_reset(cactus_graph_t graph); +CACTUS_FFI_EXPORT int32_t cactus_graph_default_backend(void); +CACTUS_FFI_EXPORT int cactus_graph_set_node_backend( + cactus_graph_t graph, cactus_node_t node, int32_t backend); + CACTUS_FFI_EXPORT int cactus_graph_save(cactus_graph_t graph, const char* filename); CACTUS_FFI_EXPORT cactus_graph_t cactus_graph_load(const char* filename); @@ -1012,6 +1031,10 @@ CACTUS_FFI_EXPORT int cactus_graph_gelu_erf(cactus_graph_t graph, cactus_node_t CACTUS_FFI_EXPORT int cactus_graph_sigmoid(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_tanh(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_glu(cactus_graph_t graph, cactus_node_t x, int32_t axis, cactus_node_t* out); +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_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); diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp index da2ed3ce5..cc52bda29 100644 --- a/cactus-graph/src/builder.cpp +++ b/cactus-graph/src/builder.cpp @@ -27,32 +27,43 @@ size_t CactusGraph::input(const std::vector& shape, Precision precision) return add_node(OpType::INPUT, {}, shape, {.output_precision = precision}); } +void CactusGraph::set_node_backend(size_t node_id, ComputeBackend backend) { + auto it = node_index_map_.find(node_id); + if (it == node_index_map_.end()) throw std::runtime_error("set_node_backend: unknown node id"); + nodes_[it->second]->params.backend = backend; +} + +size_t CactusGraph::tag_backend(size_t node_id, ComputeBackend backend) { + nodes_[node_index_map_[node_id]]->params.backend = backend; + return node_id; +} + size_t CactusGraph::binary_broadcast_op(OpType op, size_t input1, size_t input2) { BroadcastInfo bi = BroadcastInfo::compute( get_output_buffer(input1).shape, get_output_buffer(input2).shape); return add_node(op, {input1, input2}, bi.output_shape, {.broadcast_info = bi}); } -size_t CactusGraph::add(size_t a, size_t b) { return binary_broadcast_op(OpType::ADD, a, b); } -size_t CactusGraph::add_clipped(size_t a, size_t b) { return binary_broadcast_op(OpType::ADD_CLIPPED, a, b); } -size_t CactusGraph::subtract(size_t a, size_t b) { return binary_broadcast_op(OpType::SUBTRACT, a, b); } -size_t CactusGraph::multiply(size_t a, size_t b) { return binary_broadcast_op(OpType::MULTIPLY, a, b); } -size_t CactusGraph::divide(size_t a, size_t b) { return binary_broadcast_op(OpType::DIVIDE, a, b); } -size_t CactusGraph::not_equal(size_t a, size_t b) { return binary_broadcast_op(OpType::NOT_EQUAL, a, b); } +size_t CactusGraph::add(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::ADD, a, b), backend); } +size_t CactusGraph::add_clipped(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::ADD_CLIPPED, a, b), backend); } +size_t CactusGraph::subtract(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::SUBTRACT, a, b), backend); } +size_t CactusGraph::multiply(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::MULTIPLY, a, b), backend); } +size_t CactusGraph::divide(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::DIVIDE, a, b), backend); } +size_t CactusGraph::not_equal(size_t a, size_t b, ComputeBackend backend) { return tag_backend(binary_broadcast_op(OpType::NOT_EQUAL, a, b), backend); } -size_t CactusGraph::abs(size_t input) { +size_t CactusGraph::abs(size_t input, ComputeBackend backend) { const auto& input_buffer = get_output_buffer(input); OpParams params{.output_precision = input_buffer.precision}; - return add_node(OpType::ABS, {input}, input_buffer.shape, params); + return tag_backend(add_node(OpType::ABS, {input}, input_buffer.shape, params), backend); } -size_t CactusGraph::pow(size_t input, float exponent) { +size_t CactusGraph::pow(size_t input, float exponent, ComputeBackend backend) { const auto& input_buffer = get_output_buffer(input); OpParams params{.scalar = exponent, .output_precision = input_buffer.precision}; - return add_node(OpType::POW, {input}, input_buffer.shape, params); + return tag_backend(add_node(OpType::POW, {input}, input_buffer.shape, params), backend); } -size_t CactusGraph::view(size_t input, const std::vector& new_shape) { +size_t CactusGraph::view(size_t input, const std::vector& new_shape, ComputeBackend backend) { const auto& input_buffer = get_output_buffer(input); size_t input_elements = 1; @@ -70,10 +81,10 @@ size_t CactusGraph::view(size_t input, const std::vector& new_shape) { } OpParams params{.new_shape = new_shape}; - return add_node(OpType::VIEW, {input}, new_shape, params); + return tag_backend(add_node(OpType::VIEW, {input}, new_shape, params), backend); } -size_t CactusGraph::flatten(size_t input, int start_dim, int end_dim) { +size_t CactusGraph::flatten(size_t input, int start_dim, int end_dim, ComputeBackend backend) { const auto& input_buffer = get_output_buffer(input); const auto& shape = input_buffer.shape; size_t rank = shape.size(); @@ -104,7 +115,7 @@ size_t CactusGraph::flatten(size_t input, int start_dim, int end_dim) { } OpParams params{.new_shape = output_shape}; - return add_node(OpType::FLATTEN, {input}, output_shape, params); + return tag_backend(add_node(OpType::FLATTEN, {input}, output_shape, params), backend); } size_t CactusGraph::matmul(size_t input1, size_t input2, bool pretransposed_rhs, ComputeBackend backend) { @@ -171,12 +182,12 @@ size_t CactusGraph::transposeN(size_t input, const std::vector& permutat return add_node(OpType::TRANSPOSE, {input}, output_shape, params); } -size_t CactusGraph::reshape(size_t input, const std::vector& new_shape) { +size_t CactusGraph::reshape(size_t input, const std::vector& new_shape, ComputeBackend backend) { OpParams params{.new_shape = new_shape}; - return add_node(OpType::RESHAPE, {input}, new_shape, params); + return tag_backend(add_node(OpType::RESHAPE, {input}, new_shape, params), backend); } -size_t CactusGraph::index(size_t input, size_t index_value, int dim) { +size_t CactusGraph::index(size_t input, size_t index_value, int dim, ComputeBackend backend) { const auto& input_buffer = get_output_buffer(input); const auto& shape = input_buffer.shape; @@ -211,7 +222,7 @@ size_t CactusGraph::index(size_t input, size_t index_value, int dim) { } OpParams params{.axis = actual_dim, .output_precision = input_buffer.precision, .index_value = index_value}; - return add_node(OpType::INDEX, {input}, output_shape, params); + return tag_backend(add_node(OpType::INDEX, {input}, output_shape, params), backend); } size_t CactusGraph::reduction_op(OpType op, size_t input, int axis) { @@ -227,13 +238,13 @@ size_t CactusGraph::reduction_op(OpType op, size_t input, int axis) { return add_node(op, {input}, out, {.axis = axis, .output_precision = buf.precision}); } -size_t CactusGraph::sum(size_t input, int axis) { return reduction_op(OpType::SUM, input, axis); } -size_t CactusGraph::mean(size_t input, int axis) { return reduction_op(OpType::MEAN, input, axis); } -size_t CactusGraph::variance(size_t input, int axis) { return reduction_op(OpType::VARIANCE, input, axis); } -size_t CactusGraph::min(size_t input, int axis) { return reduction_op(OpType::MIN, input, axis); } -size_t CactusGraph::max(size_t input, int axis) { return reduction_op(OpType::MAX, input, axis); } +size_t CactusGraph::sum(size_t input, int axis, ComputeBackend backend) { return tag_backend(reduction_op(OpType::SUM, input, axis), backend); } +size_t CactusGraph::mean(size_t input, int axis, ComputeBackend backend) { return tag_backend(reduction_op(OpType::MEAN, input, axis), backend); } +size_t CactusGraph::variance(size_t input, int axis, ComputeBackend backend) { return tag_backend(reduction_op(OpType::VARIANCE, input, axis), backend); } +size_t CactusGraph::min(size_t input, int axis, ComputeBackend backend) { return tag_backend(reduction_op(OpType::MIN, input, axis), backend); } +size_t CactusGraph::max(size_t input, int axis, ComputeBackend backend) { return tag_backend(reduction_op(OpType::MAX, input, axis), backend); } -size_t CactusGraph::cumsum(size_t input, int axis) { +size_t CactusGraph::cumsum(size_t input, int axis, ComputeBackend backend) { const auto& input_buffer = get_output_buffer(input); if (input_buffer.shape.empty()) { throw std::runtime_error("Cumsum requires at least one dimension"); @@ -247,12 +258,12 @@ size_t CactusGraph::cumsum(size_t input, int axis) { throw std::runtime_error("Invalid axis for cumsum operation"); } - return add_node(OpType::CUMSUM, {input}, input_buffer.shape, {.axis = actual_axis, .output_precision = input_buffer.precision}); + return tag_backend(add_node(OpType::CUMSUM, {input}, input_buffer.shape, {.axis = actual_axis, .output_precision = input_buffer.precision}), backend); } -size_t CactusGraph::rms_norm(size_t input, size_t weight, float epsilon) { +size_t CactusGraph::rms_norm(size_t input, size_t weight, float epsilon, ComputeBackend backend) { OpParams params{.epsilon = epsilon}; - return add_node(OpType::RMS_NORM, {input, weight}, {}, params); + return tag_backend(add_node(OpType::RMS_NORM, {input, weight}, {}, params), backend); } size_t CactusGraph::rope(size_t input, float theta, size_t position_offset, ComputeBackend backend) { @@ -260,12 +271,12 @@ size_t CactusGraph::rope(size_t input, float theta, size_t position_offset, Comp return add_node(OpType::ROPE, {input}, {}, params); } -size_t CactusGraph::softmax(size_t input, int axis) { +size_t CactusGraph::softmax(size_t input, int axis, ComputeBackend backend) { OpParams params{.axis = axis}; - return add_node(OpType::SOFTMAX, {input}, {}, params); + return tag_backend(add_node(OpType::SOFTMAX, {input}, {}, params), backend); } -size_t CactusGraph::topk(size_t input, size_t k) { +size_t CactusGraph::topk(size_t input, size_t k, ComputeBackend backend) { const auto& input_buffer = get_output_buffer(input); if (input_buffer.shape.empty()) { @@ -275,7 +286,7 @@ size_t CactusGraph::topk(size_t input, size_t k) { std::vector output_shape = {2, input_buffer.shape[0], k}; OpParams params{.output_precision = Precision::FP32, .top_k = k}; - return add_node(OpType::TOPK, {input}, output_shape, params); + return tag_backend(add_node(OpType::TOPK, {input}, output_shape, params), backend); } size_t CactusGraph::moe_layer(size_t hidden, @@ -290,7 +301,8 @@ size_t CactusGraph::moe_layer(size_t hidden, float epsilon, float routed_scaling_factor, Activation activation, - size_t per_expert_scale) { + size_t per_expert_scale, + ComputeBackend backend) { const auto& hidden_buffer = get_output_buffer(hidden); const auto& routing_buffer = get_output_buffer(routing_probs); const auto& topk_buffer = get_output_buffer(topk_indices); @@ -328,10 +340,10 @@ size_t CactusGraph::moe_layer(size_t hidden, params.activation = activation; params.moe_gated = true; - return add_node(OpType::MOE_LAYER, input_ids, hidden_buffer.shape, params); + 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) { +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) { const auto& hidden_buffer = get_output_buffer(hidden); const auto& down_buffer = get_output_buffer(down_weight); if (hidden_buffer.shape.empty()) { @@ -347,9 +359,9 @@ 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; - return add_node(OpType::DENSE_MLP_TQ_FUSED, - {hidden, gate_weight, up_weight, down_weight}, - output_shape, params); + return tag_backend(add_node(OpType::DENSE_MLP_TQ_FUSED, + {hidden, gate_weight, up_weight, down_weight}, + output_shape, params), backend); } size_t CactusGraph::moe_layer(size_t hidden, @@ -362,7 +374,8 @@ size_t CactusGraph::moe_layer(size_t hidden, bool normalize_routing, float epsilon, float routed_scaling_factor, - Activation activation) { + Activation activation, + ComputeBackend backend) { const auto& hidden_buffer = get_output_buffer(hidden); const auto& routing_buffer = get_output_buffer(routing_probs); const auto& topk_buffer = get_output_buffer(topk_indices); @@ -398,25 +411,25 @@ size_t CactusGraph::moe_layer(size_t hidden, params.moe_gated = false; params.activation = activation; - return add_node(OpType::MOE_LAYER, input_ids, hidden_buffer.shape, params); + return tag_backend(add_node(OpType::MOE_LAYER, input_ids, hidden_buffer.shape, params), backend); } -size_t CactusGraph::layernorm(size_t input, size_t weight, size_t bias, float epsilon) { +size_t CactusGraph::layernorm(size_t input, size_t weight, size_t bias, float epsilon, ComputeBackend backend) { OpParams params{.epsilon = epsilon}; - return add_node(OpType::LAYERNORM, {input, weight, bias}, {}, params); + return tag_backend(add_node(OpType::LAYERNORM, {input, weight, bias}, {}, params), backend); } -size_t CactusGraph::layernorm(size_t input, size_t weight, float epsilon) { +size_t CactusGraph::layernorm(size_t input, size_t weight, float epsilon, ComputeBackend backend) { OpParams params{.epsilon = epsilon}; - return add_node(OpType::LAYERNORM, {input, weight}, {}, params); + return tag_backend(add_node(OpType::LAYERNORM, {input, weight}, {}, params), backend); } -size_t CactusGraph::groupnorm(size_t input, size_t weight, size_t bias, size_t num_groups, float epsilon) { +size_t CactusGraph::groupnorm(size_t input, size_t weight, size_t bias, size_t num_groups, float epsilon, ComputeBackend backend) { OpParams params{.epsilon = epsilon, .num_groups = num_groups}; - return add_node(OpType::GROUPNORM, {input, weight, bias}, {}, params); + return tag_backend(add_node(OpType::GROUPNORM, {input, weight, bias}, {}, params), backend); } -size_t CactusGraph::batchnorm(size_t input, size_t weight, size_t bias, size_t running_mean, size_t running_var, int axis, float epsilon) { +size_t CactusGraph::batchnorm(size_t input, size_t weight, size_t bias, size_t running_mean, size_t running_var, int axis, float epsilon, ComputeBackend backend) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); const auto& b = get_output_buffer(bias); @@ -442,7 +455,7 @@ size_t CactusGraph::batchnorm(size_t input, size_t weight, size_t bias, size_t r OpParams params{.epsilon = epsilon, .axis = actual_axis}; params.output_precision = xin.precision; - return add_node(OpType::BATCHNORM, {input, weight, bias, running_mean, running_var}, xin.shape, params); + return tag_backend(add_node(OpType::BATCHNORM, {input, weight, bias, running_mean, running_var}, xin.shape, params), backend); } size_t CactusGraph::attention(size_t query, size_t key, size_t value, float scale, bool is_causal, ComputeBackend backend) { @@ -483,7 +496,7 @@ size_t CactusGraph::attention_masked(size_t query, size_t key, size_t value, siz return add_node(OpType::ATTENTION, {query, key, value, mask}, {qs[0], qs[1], qs[2], vs[3]}, params); } -size_t CactusGraph::rel_pos_bias(size_t query, size_t relative_key, float scale) { +size_t CactusGraph::rel_pos_bias(size_t query, size_t relative_key, float scale, ComputeBackend backend) { const auto& q = get_output_buffer(query); const auto& r = get_output_buffer(relative_key); @@ -511,14 +524,14 @@ size_t CactusGraph::rel_pos_bias(size_t query, size_t relative_key, float scale) } OpParams params{.scale = scale, .output_precision = q.precision}; - return add_node(OpType::REL_POS_BIAS, {query, relative_key}, {B, H, T, T}, params); + return tag_backend(add_node(OpType::REL_POS_BIAS, {query, relative_key}, {B, H, T, T}, params), backend); } size_t CactusGraph::attention_int8_hybrid(size_t query, size_t key_new, size_t value_new, float scale, size_t position_offset, const int8_t* cached_keys, const int8_t* cached_values, const float* k_scales, const float* v_scales, size_t cache_len, size_t num_kv_heads, size_t head_dim, - size_t window_size, size_t v_head_dim) { + size_t window_size, size_t v_head_dim, ComputeBackend backend) { OpParams params; params.scale = scale; params.position_offset = position_offset; @@ -536,15 +549,15 @@ size_t CactusGraph::attention_int8_hybrid(size_t query, size_t key_new, size_t v const auto& q_buf = get_output_buffer(query); out_shape = {q_buf.shape[0], q_buf.shape[1], q_buf.shape[2], v_head_dim}; } - return add_node(OpType::ATTENTION_INT8_HYBRID, {query, key_new, value_new}, out_shape, params); + return tag_backend(add_node(OpType::ATTENTION_INT8_HYBRID, {query, key_new, value_new}, out_shape, params), backend); } -size_t CactusGraph::conv1d_causal(size_t input, size_t weight, size_t, size_t dilation) { +size_t CactusGraph::conv1d_causal(size_t input, size_t weight, size_t, size_t dilation, ComputeBackend backend) { OpParams params{.dilation = dilation}; - return add_node(OpType::CONV1D_CAUSAL, {input, weight}, {}, params); + return tag_backend(add_node(OpType::CONV1D_CAUSAL, {input, weight}, {}, params), backend); } -size_t CactusGraph::conv1d_k3(size_t input, size_t weight, size_t stride) { +size_t CactusGraph::conv1d_k3(size_t input, size_t weight, size_t stride, ComputeBackend backend) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); @@ -566,10 +579,10 @@ size_t CactusGraph::conv1d_k3(size_t input, size_t weight, size_t stride) { params.output_precision = xin.precision; std::vector out_shape{N, C_out, L_out}; - return add_node(OpType::CONV1D_K3, {input, weight}, out_shape, params); + return tag_backend(add_node(OpType::CONV1D_K3, {input, weight}, out_shape, params), backend); } -size_t CactusGraph::conv1d_k7s3(size_t input, size_t weight, size_t bias) { +size_t CactusGraph::conv1d_k7s3(size_t input, size_t weight, size_t bias, ComputeBackend backend) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); const auto& b = get_output_buffer(bias); @@ -594,20 +607,20 @@ size_t CactusGraph::conv1d_k7s3(size_t input, size_t weight, size_t bias) { params.output_precision = xin.precision; std::vector out_shape{N, C_out, L_out}; - return add_node(OpType::CONV1D_K7S3, {input, weight, bias}, out_shape, params); + return tag_backend(add_node(OpType::CONV1D_K7S3, {input, weight, bias}, out_shape, params), backend); } -size_t CactusGraph::conv1d(size_t input, size_t weight, size_t stride) { +size_t CactusGraph::conv1d(size_t input, size_t weight, size_t stride, ComputeBackend backend) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); if (xin.shape.size() != 3) throw std::runtime_error("conv1d expects N,C,L"); if (w.shape.size() != 3) throw std::runtime_error("conv1d weight expects [C_out, C_in, K]"); size_t L_out = (xin.shape[2] - w.shape[2]) / stride + 1; OpParams params{.stride = stride}; - return add_node(OpType::CONV1D, {input, weight}, {xin.shape[0], w.shape[0], L_out}, params); + return tag_backend(add_node(OpType::CONV1D, {input, weight}, {xin.shape[0], w.shape[0], L_out}, params), backend); } -size_t CactusGraph::conv1d(size_t input, size_t weight, size_t bias, size_t stride) { +size_t CactusGraph::conv1d(size_t input, size_t weight, size_t bias, size_t stride, ComputeBackend backend) { const auto& b = get_output_buffer(bias); const auto& w = get_output_buffer(weight); if (b.total_size != w.shape[0]) throw std::runtime_error("conv1d bias size mismatch"); @@ -617,10 +630,10 @@ size_t CactusGraph::conv1d(size_t input, size_t weight, size_t bias, size_t stri if (w.shape.size() != 3) throw std::runtime_error("conv1d weight expects [C_out, C_in, K]"); size_t L_out = (xin.shape[2] - w.shape[2]) / stride + 1; OpParams params{.output_precision = xin.precision, .stride = stride}; - return add_node(OpType::CONV1D, {input, weight, bias}, {xin.shape[0], w.shape[0], L_out}, params); + return tag_backend(add_node(OpType::CONV1D, {input, weight, bias}, {xin.shape[0], w.shape[0], L_out}, params), backend); } -size_t CactusGraph::conv1d_same_depthwise_k9(size_t input, size_t weight) { +size_t CactusGraph::conv1d_same_depthwise_k9(size_t input, size_t weight, ComputeBackend backend) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); if (xin.shape.size() != 3) throw std::runtime_error("conv1d_same_depthwise_k9 expects input [N, L, C]"); @@ -634,15 +647,15 @@ size_t CactusGraph::conv1d_same_depthwise_k9(size_t input, size_t weight) { } OpParams params{}; params.output_precision = xin.precision; - return add_node(OpType::CONV1D_SAME_DEPTHWISE_K9, {input, weight}, {xin.shape[0], xin.shape[1], C}, params); + return tag_backend(add_node(OpType::CONV1D_SAME_DEPTHWISE_K9, {input, weight}, {xin.shape[0], xin.shape[1], C}, params), backend); } -size_t CactusGraph::conv1d_same_depthwise_k9(size_t input, size_t weight, size_t bias) { - size_t node = conv1d_same_depthwise_k9(input, weight); +size_t CactusGraph::conv1d_same_depthwise_k9(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + size_t node = conv1d_same_depthwise_k9(input, weight, backend); return attach_conv_bias(node, bias, get_output_buffer(input).shape[2], "conv1d_same_depthwise_k9"); } -size_t CactusGraph::conv1d_pointwise(size_t input, size_t weight) { +size_t CactusGraph::conv1d_pointwise(size_t input, size_t weight, ComputeBackend backend) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); if (xin.shape.size() != 3) throw std::runtime_error("conv1d_pointwise expects input [N, L, C_in]"); @@ -657,15 +670,15 @@ size_t CactusGraph::conv1d_pointwise(size_t input, size_t weight) { } OpParams params{}; params.output_precision = xin.precision; - return add_node(OpType::CONV1D_POINTWISE, {input, weight}, {xin.shape[0], xin.shape[1], C_out}, params); + return tag_backend(add_node(OpType::CONV1D_POINTWISE, {input, weight}, {xin.shape[0], xin.shape[1], C_out}, params), backend); } -size_t CactusGraph::conv1d_pointwise(size_t input, size_t weight, size_t bias) { - size_t node = conv1d_pointwise(input, weight); +size_t CactusGraph::conv1d_pointwise(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + size_t node = conv1d_pointwise(input, weight, backend); return attach_conv_bias(node, bias, get_output_buffer(node).shape[2], "conv1d_pointwise"); } -size_t CactusGraph::conv2d_k3s2p1(size_t input, size_t weight) { +size_t CactusGraph::conv2d_k3s2p1(size_t input, size_t weight, ComputeBackend backend) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); @@ -694,15 +707,15 @@ size_t CactusGraph::conv2d_k3s2p1(size_t input, size_t weight) { OpParams params{}; params.output_precision = xin.precision; - return add_node(OpType::CONV2D_K3S2P1, {input, weight}, {N, C_out, H_out, W_out}, params); + return tag_backend(add_node(OpType::CONV2D_K3S2P1, {input, weight}, {N, C_out, H_out, W_out}, params), backend); } -size_t CactusGraph::conv2d_k3s2p1(size_t input, size_t weight, size_t bias) { - size_t node = conv2d_k3s2p1(input, weight); +size_t CactusGraph::conv2d_k3s2p1(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + size_t node = conv2d_k3s2p1(input, weight, backend); return attach_conv_bias(node, bias, get_output_buffer(node).shape[1], "conv2d_k3s2p1"); } -size_t CactusGraph::conv2d_depthwise_k3s2p1(size_t input, size_t weight) { +size_t CactusGraph::conv2d_depthwise_k3s2p1(size_t input, size_t weight, ComputeBackend backend) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); @@ -735,15 +748,15 @@ size_t CactusGraph::conv2d_depthwise_k3s2p1(size_t input, size_t weight) { const size_t W_out = (W - 1) / 2 + 1; OpParams params{}; params.output_precision = xin.precision; - return add_node(OpType::CONV2D_DEPTHWISE_K3S2P1, {input, weight}, {N, C, H_out, W_out}, params); + return tag_backend(add_node(OpType::CONV2D_DEPTHWISE_K3S2P1, {input, weight}, {N, C, H_out, W_out}, params), backend); } -size_t CactusGraph::conv2d_depthwise_k3s2p1(size_t input, size_t weight, size_t bias) { - size_t node = conv2d_depthwise_k3s2p1(input, weight); +size_t CactusGraph::conv2d_depthwise_k3s2p1(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + size_t node = conv2d_depthwise_k3s2p1(input, weight, backend); return attach_conv_bias(node, bias, get_output_buffer(node).shape[1], "conv2d_depthwise_k3s2p1"); } -size_t CactusGraph::conv2d_pointwise_1x1(size_t input, size_t weight) { +size_t CactusGraph::conv2d_pointwise_1x1(size_t input, size_t weight, ComputeBackend backend) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); @@ -777,22 +790,22 @@ size_t CactusGraph::conv2d_pointwise_1x1(size_t input, size_t weight) { OpParams params{}; params.output_precision = xin.precision; - return add_node(OpType::CONV2D_POINTWISE_1X1, {input, weight}, {N, C_out, H, W}, params); + return tag_backend(add_node(OpType::CONV2D_POINTWISE_1X1, {input, weight}, {N, C_out, H, W}, params), backend); } -size_t CactusGraph::conv2d_pointwise_1x1(size_t input, size_t weight, size_t bias) { - size_t node = conv2d_pointwise_1x1(input, weight); +size_t CactusGraph::conv2d_pointwise_1x1(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + size_t node = conv2d_pointwise_1x1(input, weight, backend); return attach_conv_bias(node, bias, get_output_buffer(node).shape[1], "conv2d_pointwise_1x1"); } -size_t CactusGraph::lstm_cell(size_t input, size_t h_prev, size_t c_prev, size_t weight_ih, size_t weight_hh, size_t bias_ih, size_t bias_hh) { +size_t CactusGraph::lstm_cell(size_t input, size_t h_prev, size_t c_prev, size_t weight_ih, size_t weight_hh, size_t bias_ih, size_t bias_hh, ComputeBackend backend) { const auto& h_buffer = get_output_buffer(h_prev); std::vector output_shape = {h_buffer.shape[0], h_buffer.shape[1], 2}; - return add_node(OpType::LSTM_CELL, {input, h_prev, c_prev, weight_ih, weight_hh, bias_ih, bias_hh}, output_shape, {}); + return tag_backend(add_node(OpType::LSTM_CELL, {input, h_prev, c_prev, weight_ih, weight_hh, bias_ih, bias_hh}, output_shape, {}), backend); } size_t CactusGraph::gated_deltanet_decode(size_t query, size_t key, size_t value, size_t gate_log, size_t beta, - size_t initial_state, float scale) { + size_t initial_state, float scale, ComputeBackend backend) { const auto& q = get_output_buffer(query); const auto& k = get_output_buffer(key); const auto& v = get_output_buffer(value); @@ -837,14 +850,14 @@ size_t CactusGraph::gated_deltanet_decode(size_t query, size_t key, size_t value params.scale = op_scale; params.num_kv_heads = Hq; params.output_precision = Precision::FP16; - return add_node(OpType::GATED_DELTANET_DECODE, - {query, key, value, gate_log, beta, initial_state}, - {B, static_cast(1 + K), Hv, V}, - params); + return tag_backend(add_node(OpType::GATED_DELTANET_DECODE, + {query, key, value, gate_log, beta, initial_state}, + {B, static_cast(1 + K), Hv, V}, + params), backend); } size_t CactusGraph::gated_deltanet_prefill(size_t query, size_t key, size_t value, size_t gate_log, size_t beta, - size_t initial_state, size_t chunk_size, float scale) { + size_t initial_state, size_t chunk_size, float scale, ComputeBackend backend) { const auto& q = get_output_buffer(query); const auto& k = get_output_buffer(key); const auto& v = get_output_buffer(value); @@ -891,13 +904,13 @@ size_t CactusGraph::gated_deltanet_prefill(size_t query, size_t key, size_t valu params.chunk_size = chunk_size; params.num_kv_heads = Hq; params.output_precision = Precision::FP16; - return add_node(OpType::GATED_DELTANET_PREFILL, - {query, key, value, gate_log, beta, initial_state}, - {B, T + K, Hv, V}, - params); + return tag_backend(add_node(OpType::GATED_DELTANET_PREFILL, + {query, key, value, gate_log, beta, initial_state}, + {B, T + K, Hv, V}, + params), backend); } -size_t CactusGraph::stft(size_t input, size_t weight, size_t stride, size_t num_fft_bins) { +size_t CactusGraph::stft(size_t input, size_t weight, size_t stride, size_t num_fft_bins, ComputeBackend backend) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); @@ -913,10 +926,10 @@ size_t CactusGraph::stft(size_t input, size_t weight, size_t stride, size_t num_ params.stride = stride; params.num_fft_bins = num_fft_bins; - return add_node(OpType::STFT, {input, weight}, {N, 2 * num_fft_bins, L_out}, params); + return tag_backend(add_node(OpType::STFT, {input, weight}, {N, 2 * num_fft_bins, L_out}, params), backend); } -size_t CactusGraph::concat(size_t input1, size_t input2, int axis) { +size_t CactusGraph::concat(size_t input1, size_t input2, int axis, ComputeBackend backend) { const auto& buffer1 = get_output_buffer(input1); const auto& buffer2 = get_output_buffer(input2); @@ -942,10 +955,10 @@ size_t CactusGraph::concat(size_t input1, size_t input2, int axis) { OpParams params; params.axis = axis; - return add_node(OpType::CONCAT, {input1, input2}, output_shape, params); + return tag_backend(add_node(OpType::CONCAT, {input1, input2}, output_shape, params), backend); } -size_t CactusGraph::cat(const std::vector& inputs, int axis) { +size_t CactusGraph::cat(const std::vector& inputs, int axis, ComputeBackend backend) { if (inputs.empty()) { throw std::runtime_error("Cat requires at least one input"); } @@ -974,10 +987,10 @@ size_t CactusGraph::cat(const std::vector& inputs, int axis) { OpParams params; params.axis = axis; - return add_node(OpType::CAT, inputs, output_shape, params); + return tag_backend(add_node(OpType::CAT, inputs, output_shape, params), backend); } -size_t CactusGraph::scatter_topk(size_t indices, size_t values, size_t num_classes) { +size_t CactusGraph::scatter_topk(size_t indices, size_t values, size_t num_classes, ComputeBackend backend) { const auto& indices_buffer = get_output_buffer(indices); const auto& values_buffer = get_output_buffer(values); @@ -993,17 +1006,17 @@ size_t CactusGraph::scatter_topk(size_t indices, size_t values, size_t num_class std::vector output_shape = {num_classes, indices_buffer.shape[0]}; OpParams params{.output_precision = Precision::FP32, .num_classes = num_classes}; - return add_node(OpType::SCATTER_TOPK, {indices, values}, output_shape, params); + return tag_backend(add_node(OpType::SCATTER_TOPK, {indices, values}, output_shape, params), backend); } size_t CactusGraph::sample(size_t logits, float temperature, float top_p, size_t top_k, - const std::unordered_map& logit_bias) { - return this->sample_with_options(logits, temperature, top_p, 0.15f, 1.1f, top_k, logit_bias); + const std::unordered_map& logit_bias, ComputeBackend backend) { + return this->sample_with_options(logits, temperature, top_p, 0.15f, 1.1f, top_k, logit_bias, backend); } size_t CactusGraph::sample_with_options(size_t logits, float temperature, float top_p, float min_p, float repetition_penalty, size_t top_k, - const std::unordered_map& logit_bias) { + const std::unordered_map& logit_bias, ComputeBackend backend) { const auto& logits_buffer = get_output_buffer(logits); if (logits_buffer.shape.empty()) { @@ -1029,7 +1042,7 @@ size_t CactusGraph::sample_with_options(size_t logits, float temperature, float } std::vector output_shape = {1}; - return add_node(OpType::SAMPLE, {logits}, output_shape, params); + return tag_backend(add_node(OpType::SAMPLE, {logits}, output_shape, params), backend); } static size_t scalar_val_op(CactusGraph& g, OpType op, size_t input, float value) { @@ -1039,57 +1052,57 @@ static size_t scalar_val_op(CactusGraph& g, OpType op, size_t input, float value return g.add_node(op, {input}, {}, params); } -size_t CactusGraph::scalar_add(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_ADD, input, value); } -size_t CactusGraph::scalar_subtract(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_SUBTRACT, input, value); } -size_t CactusGraph::scalar_multiply(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_MULTIPLY, input, value); } -size_t CactusGraph::scalar_divide(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_DIVIDE, input, value); } -size_t CactusGraph::scalar_not_equal(size_t input, float value) { return scalar_val_op(*this, OpType::SCALAR_NOT_EQUAL, input, value); } +size_t CactusGraph::scalar_add(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_ADD, input, value), backend); } +size_t CactusGraph::scalar_subtract(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_SUBTRACT, input, value), backend); } +size_t CactusGraph::scalar_multiply(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_MULTIPLY, input, value), backend); } +size_t CactusGraph::scalar_divide(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_DIVIDE, input, value), backend); } +size_t CactusGraph::scalar_not_equal(size_t input, float value, ComputeBackend backend) { return tag_backend(scalar_val_op(*this, OpType::SCALAR_NOT_EQUAL, input, value), backend); } -size_t CactusGraph::scalar_exp(size_t input) { - return add_node(OpType::SCALAR_EXP, {input}, {}); +size_t CactusGraph::scalar_exp(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SCALAR_EXP, {input}, {}), backend); } -size_t CactusGraph::scalar_sqrt(size_t input) { - return add_node(OpType::SCALAR_SQRT, {input}, {}); +size_t CactusGraph::scalar_sqrt(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SCALAR_SQRT, {input}, {}), backend); } -size_t CactusGraph::scalar_cos(size_t input) { - return add_node(OpType::SCALAR_COS, {input}, {}); +size_t CactusGraph::scalar_cos(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SCALAR_COS, {input}, {}), backend); } -size_t CactusGraph::scalar_sin(size_t input) { - return add_node(OpType::SCALAR_SIN, {input}, {}); +size_t CactusGraph::scalar_sin(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SCALAR_SIN, {input}, {}), backend); } -size_t CactusGraph::scalar_log(size_t input) { - return add_node(OpType::SCALAR_LOG, {input}, {}); +size_t CactusGraph::scalar_log(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SCALAR_LOG, {input}, {}), backend); } -size_t CactusGraph::relu(size_t input) { - return add_node(OpType::RELU, {input}, {}); +size_t CactusGraph::relu(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::RELU, {input}, {}), backend); } -size_t CactusGraph::silu(size_t input) { - return add_node(OpType::SILU, {input}, {}); +size_t CactusGraph::silu(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SILU, {input}, {}), backend); } -size_t CactusGraph::gelu(size_t input) { - return add_node(OpType::GELU, {input}, {}); +size_t CactusGraph::gelu(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::GELU, {input}, {}), backend); } -size_t CactusGraph::gelu_erf(size_t input) { - return add_node(OpType::GELU_ERF, {input}, {}); +size_t CactusGraph::gelu_erf(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::GELU_ERF, {input}, {}), backend); } -size_t CactusGraph::sigmoid(size_t input) { - return add_node(OpType::SIGMOID, {input}, {}); +size_t CactusGraph::sigmoid(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::SIGMOID, {input}, {}), backend); } -size_t CactusGraph::tanh(size_t input) { - return add_node(OpType::TANH, {input}, {}); +size_t CactusGraph::tanh(size_t input, ComputeBackend backend) { + return tag_backend(add_node(OpType::TANH, {input}, {}), backend); } -size_t CactusGraph::glu(size_t input, int axis) { +size_t CactusGraph::glu(size_t input, int axis, ComputeBackend backend) { const auto& xin = get_output_buffer(input); if (xin.shape.empty()) { throw std::runtime_error("glu expects non-scalar input"); @@ -1113,7 +1126,7 @@ size_t CactusGraph::glu(size_t input, int axis) { OpParams params{}; params.axis = actual_axis; params.output_precision = xin.precision; - return add_node(OpType::GLU, {input}, out_shape, params); + return tag_backend(add_node(OpType::GLU, {input}, out_shape, params), backend); } size_t CactusGraph::rope_gptj(size_t input, float theta, size_t position_offset, size_t rot_dim, ComputeBackend backend) { @@ -1125,7 +1138,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) { +size_t CactusGraph::gather(size_t tensor, size_t indices, ComputeBackend backend) { const auto& tensor_buffer = get_output_buffer(tensor); const auto& idx_shape = get_output_buffer(indices).shape; @@ -1141,10 +1154,10 @@ size_t CactusGraph::gather(size_t tensor, size_t indices) { OpParams params; params.output_precision = tensor_buffer.precision; - return add_node(OpType::GATHER, {tensor, indices}, output_shape, params); + return tag_backend(add_node(OpType::GATHER, {tensor, indices}, output_shape, params), backend); } -size_t CactusGraph::slice(size_t input, int axis, size_t start, size_t length) { +size_t CactusGraph::slice(size_t input, int axis, size_t start, size_t length, ComputeBackend backend) { const auto& input_buffer = get_output_buffer(input); if (input_buffer.shape.empty()) { throw std::runtime_error("Cannot slice a scalar tensor"); @@ -1166,10 +1179,10 @@ size_t CactusGraph::slice(size_t input, int axis, size_t start, size_t length) { params.slice_length = length; params.output_precision = input_buffer.precision; - return add_node(OpType::SLICE, {input}, output_shape, params); + return tag_backend(add_node(OpType::SLICE, {input}, output_shape, params), backend); } -size_t CactusGraph::embedding(size_t embedding_tensor, size_t indices) { +size_t CactusGraph::embedding(size_t embedding_tensor, size_t indices, ComputeBackend backend) { const auto& emb_buffer = get_output_buffer(embedding_tensor); const auto& idx_shape = get_output_buffer(indices).shape; @@ -1186,10 +1199,10 @@ size_t CactusGraph::embedding(size_t embedding_tensor, size_t indices) { OpParams params; params.output_precision = Precision::FP16; - return add_node(OpType::EMBEDDING, {embedding_tensor, indices}, output_shape, params); + return tag_backend(add_node(OpType::EMBEDDING, {embedding_tensor, indices}, output_shape, params), backend); } -size_t CactusGraph::bilinear_interpolation(size_t pos_embeds, size_t dst_height, size_t dst_width, bool align_corners) { +size_t CactusGraph::bilinear_interpolation(size_t pos_embeds, size_t dst_height, size_t dst_width, bool align_corners, ComputeBackend backend) { const auto& pos_embeds_buffer = get_output_buffer(pos_embeds); size_t embed_dim = pos_embeds_buffer.shape[1]; std::vector output_shape = {dst_height * dst_width, embed_dim}; @@ -1200,13 +1213,13 @@ size_t CactusGraph::bilinear_interpolation(size_t pos_embeds, size_t dst_height, params.align_corners = align_corners; params.output_precision = Precision::FP16; - return add_node(OpType::BILINEAR_INTERPOLATION, {pos_embeds}, output_shape, params); + return tag_backend(add_node(OpType::BILINEAR_INTERPOLATION, {pos_embeds}, output_shape, params), backend); } -size_t CactusGraph::precision_cast(size_t input, Precision target_precision) { +size_t CactusGraph::precision_cast(size_t input, Precision target_precision, ComputeBackend backend) { OpParams params{}; params.output_precision = target_precision; - return add_node(OpType::PRECISION_CAST, {input}, {}, params); + return tag_backend(add_node(OpType::PRECISION_CAST, {input}, {}, params), backend); } size_t CactusGraph::add_node(OpType op_type, const std::vector& inputs, const std::vector& output_shape, const OpParams& params) { @@ -1266,13 +1279,13 @@ void CactusGraph::resize_cache_slots(size_t node_id, size_t num_slots) { node.output_buffer.data.reset(); } -size_t CactusGraph::persistent(size_t source_node) { +size_t CactusGraph::persistent(size_t source_node, ComputeBackend backend) { const auto& source_buffer = get_output_buffer(source_node); OpParams params; params.output_precision = source_buffer.precision; size_t node_id = add_node(OpType::PERSISTENT, {source_node}, source_buffer.shape, params); persistent_node_ids_.insert(node_id); - return node_id; + return tag_backend(node_id, backend); } bool CactusGraph::is_populated(size_t persistent_node_id) const { @@ -1284,7 +1297,7 @@ void CactusGraph::invalidate_persistent(size_t persistent_node_id) { persistent_node_ids_.erase(persistent_node_id); } -size_t CactusGraph::altup_predict(size_t coefs, const size_t* streams, size_t num_streams) { +size_t CactusGraph::altup_predict(size_t coefs, const size_t* streams, size_t num_streams, ComputeBackend backend) { const auto& stream0_buf = get_output_buffer(streams[0]); size_t seq_len = stream0_buf.shape[0]; @@ -1296,10 +1309,10 @@ size_t CactusGraph::altup_predict(size_t coefs, const size_t* streams, size_t nu OpParams params; params.num_altup_inputs = num_streams; - return add_node(OpType::ALTUP_PREDICT, input_ids, {num_streams * seq_len, hidden_dim}, params); + return tag_backend(add_node(OpType::ALTUP_PREDICT, input_ids, {num_streams * seq_len, hidden_dim}, params), backend); } -size_t CactusGraph::altup_correct(size_t coefs, size_t innovation, const size_t* predictions, size_t num_predictions) { +size_t CactusGraph::altup_correct(size_t coefs, size_t innovation, const size_t* predictions, size_t num_predictions, ComputeBackend backend) { const auto& pred0_buf = get_output_buffer(predictions[0]); size_t seq_len = pred0_buf.shape[0]; @@ -1311,10 +1324,10 @@ size_t CactusGraph::altup_correct(size_t coefs, size_t innovation, const size_t* OpParams params; params.num_altup_inputs = num_predictions; - return add_node(OpType::ALTUP_CORRECT, input_ids, {num_predictions * seq_len, hidden_dim}, params); + return tag_backend(add_node(OpType::ALTUP_CORRECT, input_ids, {num_predictions * seq_len, hidden_dim}, params), backend); } -size_t CactusGraph::gaussian_topk(size_t input, float ppf) { +size_t CactusGraph::gaussian_topk(size_t input, float ppf, ComputeBackend backend) { const auto& in_buf = get_output_buffer(input); if (in_buf.precision != Precision::FP16) { throw std::runtime_error("gaussian_topk only supports FP16 input"); @@ -1322,37 +1335,38 @@ size_t CactusGraph::gaussian_topk(size_t input, float ppf) { OpParams params; params.scalar = ppf; params.output_precision = in_buf.precision; - return add_node(OpType::GAUSSIAN_TOPK, {input}, in_buf.shape, params); + return tag_backend(add_node(OpType::GAUSSIAN_TOPK, {input}, in_buf.shape, params), backend); } -size_t CactusGraph::leaky_relu(size_t input, float negative_slope) { +size_t CactusGraph::leaky_relu(size_t input, float negative_slope, ComputeBackend backend) { const auto& in_buf = get_output_buffer(input); OpParams params; params.scalar = negative_slope; - return add_node(OpType::LEAKY_RELU, {input}, in_buf.shape, params); + return tag_backend(add_node(OpType::LEAKY_RELU, {input}, in_buf.shape, params), backend); } -size_t CactusGraph::clamp(size_t input, float lo, float hi) { +size_t CactusGraph::clamp(size_t input, float lo, float hi, ComputeBackend backend) { const auto& in_buf = get_output_buffer(input); OpParams params; params.scalar = lo; params.scale = hi; - return add_node(OpType::CLAMP, {input}, in_buf.shape, params); + return tag_backend(add_node(OpType::CLAMP, {input}, in_buf.shape, params), backend); } size_t CactusGraph::bilstm_sequence(size_t input, size_t w_ih_fwd, size_t w_hh_fwd, size_t b_ih_fwd, size_t b_hh_fwd, - size_t w_ih_bwd, size_t w_hh_bwd, size_t b_ih_bwd, size_t b_hh_bwd) { + size_t w_ih_bwd, size_t w_hh_bwd, size_t b_ih_bwd, size_t b_hh_bwd, + ComputeBackend backend) { const auto& in_buf = get_output_buffer(input); const auto& w_ih = get_output_buffer(w_ih_fwd); size_t batch = in_buf.shape[0]; size_t seq_len = in_buf.shape[1]; size_t hidden_size = w_ih.shape[0] / 4; - return add_node(OpType::BILSTM_SEQUENCE, - {input, w_ih_fwd, w_hh_fwd, b_ih_fwd, b_hh_fwd, w_ih_bwd, w_hh_bwd, b_ih_bwd, b_hh_bwd}, - {batch, seq_len, 2 * hidden_size}, {}); + return tag_backend(add_node(OpType::BILSTM_SEQUENCE, + {input, w_ih_fwd, w_hh_fwd, b_ih_fwd, b_hh_fwd, w_ih_bwd, w_hh_bwd, b_ih_bwd, b_hh_bwd}, + {batch, seq_len, 2 * hidden_size}, {}), backend); } -size_t CactusGraph::maxpool1d(size_t input, size_t kernel_size, size_t stride) { +size_t CactusGraph::maxpool1d(size_t input, size_t kernel_size, size_t stride, ComputeBackend backend) { const auto& in_buf = get_output_buffer(input); size_t batch = in_buf.shape[0]; size_t channels = in_buf.shape[1]; @@ -1362,10 +1376,10 @@ size_t CactusGraph::maxpool1d(size_t input, size_t kernel_size, size_t stride) { OpParams params; params.kernel_size = kernel_size; params.stride = stride; - return add_node(OpType::MAXPOOL1D, {input}, {batch, channels, output_length}, params); + return tag_backend(add_node(OpType::MAXPOOL1D, {input}, {batch, channels, output_length}, params), backend); } -size_t CactusGraph::conv2d_k3s1p1(size_t input, size_t weight) { +size_t CactusGraph::conv2d_k3s1p1(size_t input, size_t weight, ComputeBackend backend) { const auto& xin = get_output_buffer(input); const auto& w = get_output_buffer(weight); @@ -1391,32 +1405,32 @@ size_t CactusGraph::conv2d_k3s1p1(size_t input, size_t weight) { OpParams params{}; params.output_precision = xin.precision; - return add_node(OpType::CONV2D_K3S1P1, {input, weight}, {N, C_out, H, W}, params); + return tag_backend(add_node(OpType::CONV2D_K3S1P1, {input, weight}, {N, C_out, H, W}, params), backend); } -size_t CactusGraph::conv2d_k3s1p1(size_t input, size_t weight, size_t bias) { - size_t node = conv2d_k3s1p1(input, weight); +size_t CactusGraph::conv2d_k3s1p1(size_t input, size_t weight, size_t bias, ComputeBackend backend) { + size_t node = conv2d_k3s1p1(input, weight, backend); return attach_conv_bias(node, bias, get_output_buffer(node).shape[1], "conv2d_k3s1p1"); } -size_t CactusGraph::stats_pool(size_t input) { +size_t CactusGraph::stats_pool(size_t input, ComputeBackend backend) { const auto& xin = get_output_buffer(input); size_t batch = xin.shape[0]; size_t features = 1; for (size_t i = 1; i < xin.shape.size() - 1; ++i) features *= xin.shape[i]; - return add_node(OpType::STATS_POOL, {input}, {batch, features * 2}); + return tag_backend(add_node(OpType::STATS_POOL, {input}, {batch, features * 2}), backend); } -size_t CactusGraph::weighted_stats_pool(size_t input, size_t weights) { +size_t CactusGraph::weighted_stats_pool(size_t input, size_t weights, ComputeBackend backend) { const auto& xin = get_output_buffer(input); size_t batch = xin.shape[0]; size_t features = 1; for (size_t i = 1; i < xin.shape.size() - 1; ++i) features *= xin.shape[i]; - return add_node(OpType::WEIGHTED_STATS_POOL, {input, weights}, {batch, features * 2}); + return tag_backend(add_node(OpType::WEIGHTED_STATS_POOL, {input, weights}, {batch, features * 2}), backend); } size_t CactusGraph::kv_cache_state(size_t max_seq_len, size_t num_kv_heads, size_t head_dim, - size_t window_size, size_t sink_size, size_t num_slots) { + size_t window_size, size_t sink_size, size_t num_slots, ComputeBackend backend) { if (num_slots == 0) num_slots = 1; bool fp16_cache = use_fp16_kv_cache_for_builder(); size_t total_elements = 0; @@ -1440,20 +1454,20 @@ size_t CactusGraph::kv_cache_state(size_t max_seq_len, size_t num_kv_heads, size params.output_precision = precision; size_t node_id = add_node(OpType::KV_CACHE_STATE, {}, {total_elements}, params); persistent_node_ids_.insert(node_id); - return node_id; + return tag_backend(node_id, backend); } size_t CactusGraph::kv_cache_append(size_t new_kv, size_t cache_state_node, - size_t window_size, size_t sink_size, size_t cache_slot) { + size_t window_size, size_t sink_size, size_t cache_slot, ComputeBackend backend) { OpParams params{}; params.window_size = window_size; params.cache_sink_size = sink_size; params.cache_slot = cache_slot; params.output_precision = Precision::FP32; - return add_node(OpType::KV_CACHE_APPEND, {new_kv, cache_state_node}, {1}, params); + return tag_backend(add_node(OpType::KV_CACHE_APPEND, {new_kv, cache_state_node}, {1}, params), backend); } -size_t CactusGraph::conv_cache_state(size_t ws, size_t hidden_dim) { +size_t CactusGraph::conv_cache_state(size_t ws, size_t hidden_dim, ComputeBackend backend) { size_t total_bytes = sizeof(uint64_t) * 8 + ws * hidden_dim * sizeof(__fp16); OpParams params{}; params.window_size = ws; @@ -1461,10 +1475,10 @@ size_t CactusGraph::conv_cache_state(size_t ws, size_t hidden_dim) { params.output_precision = Precision::INT8; size_t node_id = add_node(OpType::CONV_CACHE_STATE, {}, {total_bytes}, params); persistent_node_ids_.insert(node_id); - return node_id; + return tag_backend(node_id, backend); } -size_t CactusGraph::conv_cache_append(size_t new_data, size_t cache_state_node) { +size_t CactusGraph::conv_cache_append(size_t new_data, size_t cache_state_node, ComputeBackend backend) { const auto& cache_buf = get_output_buffer(cache_state_node); auto* raw = static_cast(cache_buf.get_data()); size_t ws, hd; @@ -1478,10 +1492,10 @@ size_t CactusGraph::conv_cache_append(size_t new_data, size_t cache_state_node) } OpParams params{}; params.output_precision = Precision::FP16; - return add_node(OpType::CONV_CACHE_APPEND, {new_data, cache_state_node}, {ws, hd}, params); + return tag_backend(add_node(OpType::CONV_CACHE_APPEND, {new_data, cache_state_node}, {ws, hd}, params), backend); } -size_t CactusGraph::conv_cache_initialize(size_t rows, size_t cache_state_node) { +size_t CactusGraph::conv_cache_initialize(size_t rows, size_t cache_state_node, ComputeBackend backend) { if (get_node_op_type(cache_state_node) != OpType::CONV_CACHE_STATE) { throw std::invalid_argument( "conv_cache_initialize target must be a CONV_CACHE_STATE node"); @@ -1506,18 +1520,18 @@ size_t CactusGraph::conv_cache_initialize(size_t rows, size_t cache_state_node) } OpParams params{}; params.output_precision = Precision::FP16; - return add_node(OpType::CONV_CACHE_INITIALIZE, {rows, cache_state_node}, {0}, params); + return tag_backend(add_node(OpType::CONV_CACHE_INITIALIZE, {rows, cache_state_node}, {0}, params), backend); } -size_t CactusGraph::recurrent_cache_state(const std::vector& shape, Precision precision) { +size_t CactusGraph::recurrent_cache_state(const std::vector& shape, Precision precision, ComputeBackend backend) { OpParams params{}; params.output_precision = precision; size_t node_id = add_node(OpType::RECURRENT_CACHE_STATE, {}, shape, params); persistent_node_ids_.insert(node_id); - return node_id; + return tag_backend(node_id, backend); } -size_t CactusGraph::recurrent_cache_write(size_t new_value, size_t cache_state) { +size_t CactusGraph::recurrent_cache_write(size_t new_value, size_t cache_state, ComputeBackend backend) { const auto& src_buf = get_output_buffer(new_value); const auto& dst_buf = get_output_buffer(cache_state); if (get_node_op_type(cache_state) != OpType::RECURRENT_CACHE_STATE) { @@ -1532,7 +1546,7 @@ size_t CactusGraph::recurrent_cache_write(size_t new_value, size_t cache_state) } OpParams params{}; params.output_precision = dst_buf.precision; - return add_node(OpType::RECURRENT_CACHE_WRITE, {new_value, cache_state}, {0}, params); + return tag_backend(add_node(OpType::RECURRENT_CACHE_WRITE, {new_value, cache_state}, {0}, params), backend); } size_t CactusGraph::image_preprocess( @@ -1541,7 +1555,8 @@ size_t CactusGraph::image_preprocess( int target_width, int target_height, int patch_size, int channels, float rescale_factor, - const float* mean, const float* std_dev) { + const float* mean, const float* std_dev, + ComputeBackend backend) { int tw = target_width > 0 ? target_width : src_width; int th = target_height > 0 ? target_height : src_height; @@ -1568,29 +1583,29 @@ size_t CactusGraph::image_preprocess( } params.output_precision = Precision::FP32; - return add_node(OpType::IMAGE_PREPROCESS, {pixel_input}, {num_patches, patch_dim}, params); + return tag_backend(add_node(OpType::IMAGE_PREPROCESS, {pixel_input}, {num_patches, patch_dim}, params), backend); } -size_t CactusGraph::rfft(size_t input) { +size_t CactusGraph::rfft(size_t input, ComputeBackend backend) { const auto& in_buf = get_output_buffer(input); size_t n = in_buf.total_size; size_t out_len = (n / 2 + 1) * 2; OpParams params{}; params.output_precision = in_buf.precision; - return add_node(OpType::RFFT, {input}, {out_len}, params); + return tag_backend(add_node(OpType::RFFT, {input}, {out_len}, params), backend); } -size_t CactusGraph::irfft(size_t input, size_t output_length) { +size_t CactusGraph::irfft(size_t input, size_t output_length, ComputeBackend backend) { const auto& in_buf = get_output_buffer(input); OpParams params{}; params.output_precision = in_buf.precision; - return add_node(OpType::IRFFT, {input}, {output_length}, params); + return tag_backend(add_node(OpType::IRFFT, {input}, {output_length}, params), backend); } size_t CactusGraph::mel_filter_bank( size_t num_frequency_bins, size_t num_mel_filters, float min_frequency, float max_frequency, size_t sampling_rate, - int norm_type, int scale_type) { + int norm_type, int scale_type, ComputeBackend backend) { OpParams params{}; params.num_mel_filters = num_mel_filters; params.min_frequency = min_frequency; @@ -1599,7 +1614,7 @@ size_t CactusGraph::mel_filter_bank( params.mel_norm_type = norm_type; params.mel_scale_type = scale_type; params.output_precision = Precision::FP32; - return add_node(OpType::MEL_FILTER_BANK, {}, {num_frequency_bins, num_mel_filters}, params); + return tag_backend(add_node(OpType::MEL_FILTER_BANK, {}, {num_frequency_bins, num_mel_filters}, params), backend); } size_t CactusGraph::spectrogram( @@ -1608,7 +1623,7 @@ size_t CactusGraph::spectrogram( float power, bool center, int pad_mode, float mel_floor, int log_mel_mode, float dither_val, float preemphasis, - bool remove_dc_offset) { + bool remove_dc_offset, ComputeBackend backend) { const auto& wav_buf = get_output_buffer(waveform); const auto& mel_buf = get_output_buffer(mel_filters_node); @@ -1634,13 +1649,14 @@ size_t CactusGraph::spectrogram( params.remove_dc_offset = remove_dc_offset; params.output_precision = Precision::FP32; - return add_node(OpType::SPECTROGRAM, {waveform, mel_filters_node}, {num_mel_bins, num_frames}, params); + return tag_backend(add_node(OpType::SPECTROGRAM, {waveform, mel_filters_node}, {num_mel_bins, num_frames}, params), backend); } size_t CactusGraph::attention_cached(size_t query, size_t key_new, size_t value_new, 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 window_size, size_t v_head_dim, size_t cache_slot, + ComputeBackend backend) { const auto& q_shape = get_output_buffer(query).shape; size_t batch = q_shape[0]; size_t seq_len = q_shape[1]; @@ -1655,7 +1671,7 @@ size_t CactusGraph::attention_cached(size_t query, size_t key_new, size_t value_ params.v_head_dim = v_head_dim; params.cache_slot = cache_slot; params.output_precision = Precision::FP16; - return add_node(OpType::ATTENTION_CACHED, - {query, key_new, value_new, k_cache_state, v_cache_state}, - {batch, seq_len, num_q_heads, out_v_dim}, params); + return tag_backend(add_node(OpType::ATTENTION_CACHED, + {query, key_new, value_new, k_cache_state, v_cache_state}, + {batch, seq_len, num_q_heads, out_v_dim}, params), backend); } diff --git a/cactus-graph/src/execute.cpp b/cactus-graph/src/execute.cpp index eb5a1ccba..fc19b2ccd 100644 --- a/cactus-graph/src/execute.cpp +++ b/cactus-graph/src/execute.cpp @@ -14,19 +14,24 @@ #include #include -static int g_backend_override = -1; +static int g_selected_backend = -1; + +ComputeBackend cactus_default_backend() { + if (g_selected_backend >= 0) return static_cast(g_selected_backend); + if (cactus_metal_available()) return ComputeBackend::METAL; + return ComputeBackend::CPU; +} + int cactus_backend_select(const char* backend) { if (!backend) return -1; - if (std::strcmp(backend, "auto") == 0) { g_backend_override = -1; return 0; } - if (std::strcmp(backend, "metal") == 0) { - if (!cactus_metal_available()) return -1; - g_backend_override = 1; - return 0; - } - if (std::strcmp(backend, "cpu") == 0) { g_backend_override = 0; return 0; } - return -1; + ComputeBackend selected; + if (std::strcmp(backend, "cpu") == 0) selected = ComputeBackend::CPU; + else if (std::strcmp(backend, "metal") == 0 && cactus_metal_available()) selected = ComputeBackend::METAL; + else return -1; + g_selected_backend = static_cast(selected); + return 0; } -bool cactus_backend_metal() { return g_backend_override != 0 && cactus_metal_available(); } + using ComputeFn = void(*)(GraphNode&, const nodes_vector&, const node_index_map_t&); @@ -482,6 +487,7 @@ static const __fp16* metal_conv_weight_f16(const BufferDesc& w, size_t rows, siz } static bool try_encode_metal(GraphNode& node, const nodes_vector& nodes, const node_index_map_t& map) { + if (node.params.backend != ComputeBackend::METAL) return false; BufferDesc& out = node.output_buffer; auto fp16 = [](const BufferDesc& b){ return b.precision == Precision::FP16; }; switch (node.op_type) { @@ -1276,6 +1282,7 @@ void CactusGraph::build_metal_retype_plan() { for (size_t i = 0; i < n; ++i) for (size_t id : nodes_[i]->input_ids) { long j = idxof(id); if (j >= 0) cons[(size_t)j].push_back(i); } auto interior = [&](const GraphNode& nd) { + if (nd.params.backend != ComputeBackend::METAL) return false; if (nd.output_buffer.precision != Precision::FP32) return false; if (nd.output_buffer.shape.size() > 8) return false; switch (nd.op_type) { @@ -1410,7 +1417,13 @@ void CactusGraph::execute(const std::string& profile_file) { return true; }; - bool metal_mode = cactus_backend_metal(); + bool metal_mode = false; + if (cactus_metal_available()) { + for (size_t i = 0; i < n; ++i) { + if (nodes_[i]->op_type == OpType::INPUT) continue; + if (nodes_[i]->params.backend == ComputeBackend::METAL) { metal_mode = true; break; } + } + } if (metal_mode && n < 100) { for (size_t i = 0; i < n; ++i) { OpType t = nodes_[i]->op_type; @@ -1428,6 +1441,8 @@ void CactusGraph::execute(const std::string& profile_file) { for (auto& np : nodes_) { sig ^= (uint64_t)np->op_type + 0x9e3779b97f4a7c15ull; sig *= 1099511628211ull; + sig ^= (uint64_t)np->params.backend + 0x9e3779b97f4a7c15ull; + sig *= 1099511628211ull; if (np->op_type != OpType::INPUT) continue; for (size_t d : np->output_buffer.shape) { sig ^= (uint64_t)d + 0x9e3779b97f4a7c15ull; @@ -2234,9 +2249,10 @@ void CactusGraph::soft_reset_keep_pool() { } void CactusGraph::prewarm_metal_quant_weights() { - if (!cactus_backend_metal()) return; + if (!cactus_metal_available()) return; for (auto& np : nodes_) { GraphNode& node = *np; + if (node.params.backend != ComputeBackend::METAL) continue; if (node.op_type != OpType::MATMUL || node.input_ids.size() < 2) continue; auto it = node_index_map_.find(node.input_ids[1]); if (it == node_index_map_.end()) continue; diff --git a/cactus-graph/src/graph_ffi.cpp b/cactus-graph/src/graph_ffi.cpp index 74ce9a8e6..998b905a1 100644 --- a/cactus-graph/src/graph_ffi.cpp +++ b/cactus-graph/src/graph_ffi.cpp @@ -437,6 +437,33 @@ int cactus_graph_pow(cactus_graph_t graph, cactus_node_t x, float exponent, cact } } +int32_t cactus_graph_default_backend(void) { + return static_cast(cactus_default_backend()); +} + +static bool valid_backend(int32_t backend) { + switch (static_cast(backend)) { + case ComputeBackend::CPU: + case ComputeBackend::METAL: + return true; + } + return false; +} + +int cactus_graph_set_node_backend(cactus_graph_t graph, cactus_node_t node, int32_t backend) { + if (!graph || !valid_backend(backend)) { + last_error_message = "Invalid args to cactus_graph_set_node_backend"; + return -1; + } + try { + as_graph(graph)->graph.set_node_backend(static_cast(node), static_cast(backend)); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_view(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) { last_error_message = "Invalid args to cactus_graph_view"; diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index 4f6e5a25f..147eeb98b 100644 --- a/cactus-graph/src/io.cpp +++ b/cactus-graph/src/io.cpp @@ -626,7 +626,7 @@ void CactusGraph::release_all_weight_pages() { } } -size_t CactusGraph::embedding(const std::string& filename, size_t indices) { +size_t CactusGraph::embedding(const std::string& filename, size_t indices, ComputeBackend backend) { auto mapped_file = std::make_unique(filename); const auto& shape = mapped_file->shape(); @@ -646,7 +646,7 @@ size_t CactusGraph::embedding(const std::string& filename, size_t indices) { OpParams params; params.output_precision = Precision::FP16; - return add_node(OpType::EMBEDDING, {embeddings_node, indices}, output_shape, params); + return tag_backend(add_node(OpType::EMBEDDING, {embeddings_node, indices}, output_shape, params), backend); } diff --git a/cactus-graph/src/metal_plan.cpp b/cactus-graph/src/metal_plan.cpp index 5a80117a1..e0fb923ae 100644 --- a/cactus-graph/src/metal_plan.cpp +++ b/cactus-graph/src/metal_plan.cpp @@ -70,6 +70,9 @@ MetalFusePlan* cactus_metal_plan_build( std::vector pinned(n, 0); for (size_t i = 0; i < n; ++i) if (pinned_ids.count(nodes[i]->id)) pinned[i] = 1; + std::vector cpu_only(n, 0); + for (size_t i = 0; i < n; ++i) + if (nodes[i]->params.backend != ComputeBackend::METAL) { cpu_only[i] = 1; pinned[i] = 1; } auto idxof = [&](size_t id) -> long { auto it = map.find(id); @@ -225,6 +228,7 @@ MetalFusePlan* cactus_metal_plan_build( }; auto add_cluster = [&](MetalCluster c, size_t anchor, const std::vector& cover) -> bool { if (banned && banned->count(anchor)) { release_scratch(c); return false; } + if (cpu_only[anchor]) { release_scratch(c); return false; } for (size_t v : cover) if (v != anchor && pinned[v]) { release_scratch(c); return false; } int32_t cid = (int32_t)plan->clusters.size(); plan->clusters.push_back(c); @@ -866,6 +870,8 @@ MetalFusePlan* cactus_metal_plan_build( } } + cands.erase(std::remove_if(cands.begin(), cands.end(), + [&](const AttnCand& cd) { return cpu_only[cd.anchor] != 0; }), cands.end()); if (!cands.empty()) { std::vector mark(n, 0); for (auto& cd : cands) { diff --git a/cactus-graph/src/ops_cache.cpp b/cactus-graph/src/ops_cache.cpp index f0ecd9d78..ed9db8edd 100644 --- a/cactus-graph/src/ops_cache.cpp +++ b/cactus-graph/src/ops_cache.cpp @@ -103,7 +103,7 @@ inline bool use_fp16_kv_cache() { constexpr size_t kInitialCacheEntries = 256; inline bool kv_cache_resident() { - return cactus_backend_metal(); + return cactus_default_backend() == ComputeBackend::METAL; } inline bool resize_cache_buffer(BufferDesc& buf, size_t new_max) { diff --git a/cactus-graph/src/ops_nn.cpp b/cactus-graph/src/ops_nn.cpp index b8f0eb033..3d192f52e 100644 --- a/cactus-graph/src/ops_nn.cpp +++ b/cactus-graph/src/ops_nn.cpp @@ -51,7 +51,7 @@ void compute_matmul_node(GraphNode& node, const std::vector(M), output); - else if (cactus_backend_metal()) + else if (node.params.backend == ComputeBackend::METAL && cactus_metal_available()) cactus_metal_quant_matmul(&mat, lhs, static_cast(M), output); else cactus_quant_matmul(&mat, lhs, static_cast(M), output); diff --git a/cactus-graph/src/param_io.cpp b/cactus-graph/src/param_io.cpp index b173c3395..e698638f7 100644 --- a/cactus-graph/src/param_io.cpp +++ b/cactus-graph/src/param_io.cpp @@ -254,7 +254,6 @@ void write_field(std::ostream& out, ParamField field, const OpParams& params) { case ParamField::Axis: write_i32(out, static_cast(params.axis)); break; case ParamField::NewShape: write_size_vector(out, params.new_shape); break; case ParamField::PretransposedRhs: write_u32(out, params.pretransposed_rhs ? 1u : 0u); break; - case ParamField::Backend: write_u32(out, static_cast(params.backend)); break; case ParamField::SliceStart: write_u64(out, static_cast(params.slice_start)); break; case ParamField::SliceLength: write_u64(out, static_cast(params.slice_length)); break; case ParamField::Epsilon: write_f32(out, params.epsilon); break; @@ -295,11 +294,12 @@ void write_field(std::ostream& out, ParamField field, const OpParams& params) { case ParamField::CacheSeqLen: write_u64(out, static_cast(params.cache_seq_len)); break; case ParamField::HeadDim: write_u64(out, static_cast(params.head_dim)); break; case ParamField::VHeadDim: write_u64(out, static_cast(params.v_head_dim)); break; + case ParamField::Backend: case ParamField::CachedKeysInt8Ptr: case ParamField::CachedValuesInt8Ptr: case ParamField::CachedKScalesPtr: case ParamField::CachedVScalesPtr: - throw std::runtime_error("Attempted to serialize runtime-only pointer field"); + throw std::runtime_error("Attempted to serialize runtime-only field"); case ParamField::MaxCacheSeqLen: write_u64(out, static_cast(params.max_cache_seq_len)); break; case ParamField::CacheSinkSize: write_u64(out, static_cast(params.cache_sink_size)); break; case ParamField::CacheNumSlots: write_u64(out, static_cast(params.cache_num_slots)); break; @@ -312,14 +312,9 @@ void read_field(std::istream& in, ParamField field, OpParams& params) { case ParamField::Axis: params.axis = static_cast(read_i32(in)); break; case ParamField::NewShape: params.new_shape = read_size_vector(in); break; case ParamField::PretransposedRhs: params.pretransposed_rhs = (read_u32(in) != 0); break; - case ParamField::Backend: { - uint32_t backend_val = read_u32(in); - if (backend_val > static_cast(ComputeBackend::METAL)) { - throw std::runtime_error("Graph file corrupted: invalid backend"); - } - params.backend = static_cast(backend_val); + case ParamField::Backend: + read_u32(in); break; - } case ParamField::SliceStart: params.slice_start = static_cast(read_u64(in)); break; case ParamField::SliceLength: params.slice_length = static_cast(read_u64(in)); break; case ParamField::Epsilon: params.epsilon = read_f32(in); break; @@ -380,7 +375,7 @@ void write_op_params(std::ostream& out, OpType op_type, const OpParams& params) std::vector fields; fields.reserve(schema.size()); for (const auto& spec : schema) { - if (spec.persistence == FieldPersistence::Persistent) { + if (spec.persistence == FieldPersistence::Persistent && spec.field != ParamField::Backend) { fields.push_back(spec.field); } } diff --git a/cactus-graph/tests/test_metal_parity.cpp b/cactus-graph/tests/test_metal_parity.cpp index 36a50a3eb..ae8f2d548 100644 --- a/cactus-graph/tests/test_metal_parity.cpp +++ b/cactus-graph/tests/test_metal_parity.cpp @@ -76,8 +76,7 @@ struct ParityCase { for (size_t i = 0; i < result.size(); ++i) result[i] = p[i]; } graph.hard_reset(); - cactus_backend_select("auto"); - return result; + return result; } bool check() { @@ -98,29 +97,53 @@ struct ParityCase { bool metal_present() { bool ok = cactus_backend_select("metal") == 0; - cactus_backend_select("auto"); return ok; } -bool parity_unary(size_t (CactusGraph::*fn)(size_t), float lo = -1.0f, float hi = 1.0f) { +bool parity_per_op_pins() { + auto run_pinned = [&](const char* global, ComputeBackend matmul_backend) { + ParityCase c; + size_t a = c.add_input({8, 32}); + size_t b = c.add_input({32, 16}); + c.build = [matmul_backend](CactusGraph& g, const std::vector& in) { + size_t m = g.matmul(in[0], in[1], false, matmul_backend); + size_t s = g.silu(m); + return g.scalar_add(s, 1.0f); + }; + return c.run(global); + }; + std::vector ref = run_pinned("cpu", ComputeBackend::CPU); + std::vector cpu_pin_under_metal = run_pinned("metal", ComputeBackend::CPU); + std::vector metal_pin_under_cpu = run_pinned("cpu", ComputeBackend::METAL); + if (ref.empty() || cpu_pin_under_metal.size() != ref.size() || metal_pin_under_cpu.size() != ref.size()) + return false; + for (size_t i = 0; i < ref.size(); ++i) { + float scale = std::max(1.0f, std::fabs(ref[i])); + if (std::fabs(ref[i] - cpu_pin_under_metal[i]) > 5e-2f * scale) return false; + if (std::fabs(ref[i] - metal_pin_under_cpu[i]) > 5e-2f * scale) return false; + } + return true; +} + +bool parity_unary(size_t (CactusGraph::*fn)(size_t, ComputeBackend), float lo = -1.0f, float hi = 1.0f) { ParityCase c; c.add_input({4, 33}, lo, hi); - c.build = [fn](CactusGraph& g, const std::vector& in) { return (g.*fn)(in[0]); }; + c.build = [fn](CactusGraph& g, const std::vector& in) { return (g.*fn)(in[0], cactus_default_backend()); }; return c.check(); } -bool parity_scalar(size_t (CactusGraph::*fn)(size_t, float), float p, float lo = -1.0f, float hi = 1.0f) { +bool parity_scalar(size_t (CactusGraph::*fn)(size_t, float, ComputeBackend), float p, float lo = -1.0f, float hi = 1.0f) { ParityCase c; c.add_input({4, 33}, lo, hi); - c.build = [fn, p](CactusGraph& g, const std::vector& in) { return (g.*fn)(in[0], p); }; + c.build = [fn, p](CactusGraph& g, const std::vector& in) { return (g.*fn)(in[0], p, cactus_default_backend()); }; return c.check(); } -bool parity_reduce(size_t (CactusGraph::*fn)(size_t, int)) { +bool parity_reduce(size_t (CactusGraph::*fn)(size_t, int, ComputeBackend)) { for (int axis = 0; axis < 3; ++axis) { ParityCase c; c.add_input({3, 5, 17}); - c.build = [fn, axis](CactusGraph& g, const std::vector& in) { return (g.*fn)(in[0], axis); }; + c.build = [fn, axis](CactusGraph& g, const std::vector& in) { return (g.*fn)(in[0], axis, cactus_default_backend()); }; if (!c.check()) return false; } return true; @@ -248,14 +271,14 @@ bool parity_conv1d_pointwise() { return c.check(); } -bool parity_conv2d(size_t (CactusGraph::*fn)(size_t, size_t, size_t), bool depthwise, bool pointwise) { +bool parity_conv2d(size_t (CactusGraph::*fn)(size_t, size_t, size_t, ComputeBackend), bool depthwise, bool pointwise) { ParityCase c; c.add_input({1, 4, 11, 13}); if (depthwise) c.add_input({4, 1, 3, 3}); else if (pointwise) c.add_input({6, 4, 1, 1}); else c.add_input({6, 4, 3, 3}); c.add_input({pointwise || depthwise ? (depthwise ? 4u : 6u) : 6u}); - c.build = [fn](CactusGraph& g, const std::vector& in) { return (g.*fn)(in[0], in[1], in[2]); }; + c.build = [fn](CactusGraph& g, const std::vector& in) { return (g.*fn)(in[0], in[1], in[2], cactus_default_backend()); }; return c.check(); } @@ -431,7 +454,6 @@ std::vector run_cached_attention(const char* backend, size_t ceiling, siz g.soft_reset(); } g.hard_reset(); - cactus_backend_select("auto"); return collected; } @@ -543,6 +565,7 @@ int main() { runner.run_test("equality_exact", parity_equality_exact()); runner.run_test("sliding_window_cache (ring wrap)", parity_sliding_window_cache()); runner.run_test("cache_growth", parity_cache_growth()); + runner.run_test("per_op_backend_pins", parity_per_op_pins()); runner.print_summary(); return runner.all_passed() ? 0 : 1; diff --git a/docs/cactus_graph.md b/docs/cactus_graph.md index 8c0f4d858..a07a48d51 100644 --- a/docs/cactus_graph.md +++ b/docs/cactus_graph.md @@ -489,10 +489,11 @@ size_t similarity = fixture.graph().divide(dot_product, fixture.graph().multiply 1. **Batch operations**: Process multiple samples together 2. **Pre-transpose weights**: Use `pretransposed_rhs=true` for matmul when possible 3. **Fused operations**: The framework automatically fuses compatible operations -4. **Backend selection**: Use the Metal backend for supported operations: +4. **Backend selection**: Every op accepts a backend (`ComputeBackend::CPU` or `ComputeBackend::METAL`). When omitted, it defaults to `cactus_default_backend()` — auto (best available), or the one forced with `cactus_set_backend("cpu"|"metal")`. Pin a single op: ```cpp size_t result = graph.matmul(a, b, false, ComputeBackend::METAL); ``` + Python mirrors this: `g.matmul(a, b, backend=Graph.METAL)`, default `backend=None` follows the global backend. ### Graph Construction 1. **Build once, execute many**: Construct the graph once, run with different inputs diff --git a/python/README.md b/python/README.md index 616f4344c..91d1a452b 100644 --- a/python/README.md +++ b/python/README.md @@ -361,6 +361,12 @@ print(y.numpy()) # [9. 144. 729. 2304.] Supported ops: `+`, `-`, `*`, `/`, `abs`, `pow`, `view`, `flatten`, `concat`, `cat`, `relu`, `sigmoid`, `tanh`, `gelu`, `softmax`. +Every op accepts a `backend=` kwarg (`Graph.CPU` or `Graph.METAL`) to pin that op; by default it follows the global backend — auto (best available), or the one forced with `cactus_set_backend("cpu"|"metal")`: + +```python +m = g.matmul(a, b, backend=Graph.METAL) +``` + ## Testing Run the full test suite: diff --git a/python/cactus/__init__.py b/python/cactus/__init__.py index 585fe7bf1..1af766aee 100644 --- a/python/cactus/__init__.py +++ b/python/cactus/__init__.py @@ -28,10 +28,8 @@ "cactus_stream_transcribe_start", "cactus_stream_transcribe_process", "cactus_stream_transcribe_stop", - "cactus_detect_language", "cactus_tokenize", "cactus_render_prompt", - "cactus_decode_tokens", "cactus_score_window", "cactus_rag_query", "cactus_index_init", diff --git a/python/cactus/bindings/cactus.py b/python/cactus/bindings/cactus.py index 638ff4070..7b6d60cb5 100644 --- a/python/cactus/bindings/cactus.py +++ b/python/cactus/bindings/cactus.py @@ -76,6 +76,10 @@ class cactus_tensor_info_t(ctypes.Structure): _lib.cactus_graph_destroy.restype = None _lib.cactus_graph_hard_reset.argtypes = [cactus_graph_t] _lib.cactus_graph_hard_reset.restype = ctypes.c_int +_lib.cactus_graph_default_backend.argtypes = [] +_lib.cactus_graph_default_backend.restype = ctypes.c_int32 +_lib.cactus_graph_set_node_backend.argtypes = [cactus_graph_t, cactus_node_t, ctypes.c_int32] +_lib.cactus_graph_set_node_backend.restype = ctypes.c_int _lib.cactus_graph_save.argtypes = [cactus_graph_t, ctypes.c_char_p] _lib.cactus_graph_save.restype = ctypes.c_int @@ -692,15 +696,6 @@ class cactus_tensor_info_t(ctypes.Structure): _lib.cactus_stream_transcribe_stop.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_size_t] _lib.cactus_stream_transcribe_stop.restype = ctypes.c_int -_bind_optional( - "cactus_detect_language", - [ - ctypes.c_void_p, ctypes.c_char_p, ctypes.c_char_p, ctypes.c_size_t, - ctypes.c_char_p, ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t, - ], - ctypes.c_int, -) - _lib.cactus_embed.argtypes = [ ctypes.c_void_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_float), ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t), ctypes.c_bool @@ -765,20 +760,6 @@ class cactus_tensor_info_t(ctypes.Structure): ctypes.c_int, ) -_bind_optional( - "cactus_decode_tokens", - [ - ctypes.c_void_p, - ctypes.POINTER(ctypes.c_uint32), - ctypes.c_size_t, - ctypes.c_float, - ctypes.c_float, - ctypes.c_size_t, - ctypes.POINTER(ctypes.c_uint32), - ], - ctypes.c_int, -) - _lib.cactus_score_window.argtypes = [ ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint32), @@ -956,7 +937,7 @@ def cactus_telemetry_shutdown(): def cactus_set_backend(backend): - """Select the inference backend ("auto", "cpu", or "metal"). Returns 0 on success, -1 on failure.""" + """Select the inference backend ("cpu" or "metal"); unset defaults to auto (best available). Returns 0 on success, -1 on failure.""" return _lib.cactus_set_backend(_enc(backend)) @@ -1137,22 +1118,6 @@ def cactus_stream_transcribe_stop(stream): return _from_json(buf) -def cactus_detect_language(model, audio_path, options=None, pcm_data=None): - """Detect the spoken language in audio. - - Returns: - A dict with detected language info. - """ - buf = ctypes.create_string_buffer(65536) - pcm_ptr, pcm_size = _prepare_pcm(pcm_data) - rc = _lib.cactus_detect_language( - model, _enc(audio_path), buf, len(buf), _to_json(options), pcm_ptr, pcm_size, - ) - if rc < 0: - raise RuntimeError(_err("Detect language failed")) - return _from_json(buf) - - def cactus_preprocess_audio_features(audio_path, model_type, mel_bins, capacity): """Compute mel spectrogram features from an audio file. @@ -1245,24 +1210,6 @@ def cactus_render_prompt(model, messages, options=None, tools=None): return buf.value.decode("utf-8", errors="ignore") -def cactus_decode_tokens(model, tokens, temperature=0.0, top_p=1.0, top_k=1): - """Decode the next token from a sequence. Returns the next token ID as an int.""" - if not tokens: - raise ValueError("tokens must be non-empty") - arr = (ctypes.c_uint32 * len(tokens))(*[int(t) for t in tokens]) - out = ctypes.c_uint32(0) - rc = _lib.cactus_decode_tokens( - model, arr, len(tokens), - ctypes.c_float(float(temperature)), - ctypes.c_float(float(top_p)), - ctypes.c_size_t(int(top_k)), - ctypes.byref(out), - ) - if rc < 0: - raise RuntimeError(_err("Token decode failed")) - return int(out.value) - - def cactus_score_window(model, tokens, start, end, context): """Score a window of tokens for log-probabilities. @@ -1617,57 +1564,57 @@ def execute(self): if rc != 0: raise RuntimeError(_err("graph_execute failed")) - def add(self, a, b): - return self._binary("cactus_graph_add", a, b) + def add(self, a, b, backend=None): + return self._binary("cactus_graph_add", a, b, backend=backend) - def add_clipped(self, a, b): - return self._binary("cactus_graph_add_clipped", a, b) + def add_clipped(self, a, b, backend=None): + return self._binary("cactus_graph_add_clipped", a, b, backend=backend) - def subtract(self, a, b): - return self._binary("cactus_graph_subtract", a, b) + def subtract(self, a, b, backend=None): + return self._binary("cactus_graph_subtract", a, b, backend=backend) - def multiply(self, a, b): - return self._binary("cactus_graph_multiply", a, b) + def multiply(self, a, b, backend=None): + return self._binary("cactus_graph_multiply", a, b, backend=backend) - def divide(self, a, b): - return self._binary("cactus_graph_divide", a, b) + def divide(self, a, b, backend=None): + return self._binary("cactus_graph_divide", a, b, backend=backend) - def not_equal(self, a, b): - return self._binary("cactus_graph_not_equal", a, b) + def not_equal(self, a, b, backend=None): + return self._binary("cactus_graph_not_equal", a, b, backend=backend) - def abs(self, x): + def abs(self, x, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_abs(self.h, cactus_node_t(x.id), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_abs failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def pow(self, x, exponent): + def pow(self, x, exponent, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_pow(self.h, cactus_node_t(x.id), ctypes.c_float(float(exponent)), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_pow failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def precision_cast(self, x, dtype): + def precision_cast(self, x, dtype, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_precision_cast(self.h, cactus_node_t(x.id), int(dtype), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_precision_cast failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def quantize_activations(self, x): + def quantize_activations(self, x, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_quantize_activations(self.h, cactus_node_t(x.id), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_quantize_activations failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def _scalar(self, fn_name, x, value=None): + def _scalar(self, fn_name, x, value=None, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() fn = getattr(_lib, fn_name) @@ -1677,42 +1624,42 @@ def _scalar(self, fn_name, x, value=None): rc = fn(self.h, cactus_node_t(x.id), ctypes.c_float(float(value)), ctypes.byref(out)) if rc != 0: raise RuntimeError(f"{fn_name} failed") - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def scalar_add(self, x, value): - return self._scalar("cactus_graph_scalar_add", x, value) + def scalar_add(self, x, value, backend=None): + return self._scalar("cactus_graph_scalar_add", x, value, backend=backend) - def scalar_subtract(self, x, value): - return self._scalar("cactus_graph_scalar_subtract", x, value) + def scalar_subtract(self, x, value, backend=None): + return self._scalar("cactus_graph_scalar_subtract", x, value, backend=backend) - def scalar_multiply(self, x, value): - return self._scalar("cactus_graph_scalar_multiply", x, value) + def scalar_multiply(self, x, value, backend=None): + return self._scalar("cactus_graph_scalar_multiply", x, value, backend=backend) - def scalar_divide(self, x, value): - return self._scalar("cactus_graph_scalar_divide", x, value) + def scalar_divide(self, x, value, backend=None): + return self._scalar("cactus_graph_scalar_divide", x, value, backend=backend) - def scalar_floor_divide(self, x, value): - return self._scalar("cactus_graph_scalar_floor_divide", x, value) + def scalar_floor_divide(self, x, value, backend=None): + return self._scalar("cactus_graph_scalar_floor_divide", x, value, backend=backend) - def scalar_not_equal(self, x, value): - return self._scalar("cactus_graph_scalar_not_equal", x, value) + def scalar_not_equal(self, x, value, backend=None): + return self._scalar("cactus_graph_scalar_not_equal", x, value, backend=backend) - def scalar_exp(self, x): - return self._scalar("cactus_graph_scalar_exp", x) + def scalar_exp(self, x, backend=None): + return self._scalar("cactus_graph_scalar_exp", x, backend=backend) - def scalar_sqrt(self, x): - return self._scalar("cactus_graph_scalar_sqrt", x) + def scalar_sqrt(self, x, backend=None): + return self._scalar("cactus_graph_scalar_sqrt", x, backend=backend) - def scalar_cos(self, x): - return self._scalar("cactus_graph_scalar_cos", x) + def scalar_cos(self, x, backend=None): + return self._scalar("cactus_graph_scalar_cos", x, backend=backend) - def scalar_sin(self, x): - return self._scalar("cactus_graph_scalar_sin", x) + def scalar_sin(self, x, backend=None): + return self._scalar("cactus_graph_scalar_sin", x, backend=backend) - def scalar_log(self, x): - return self._scalar("cactus_graph_scalar_log", x) + def scalar_log(self, x, backend=None): + return self._scalar("cactus_graph_scalar_log", x, backend=backend) - def clamp(self, x, lo, hi): + def clamp(self, x, lo, hi, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_clamp( @@ -1724,12 +1671,12 @@ def clamp(self, x, lo, hi): ) if rc != 0: raise RuntimeError(_err("graph_clamp failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def masked_select_prefix(self, x, mask): - return self._binary("cactus_graph_masked_select_prefix", x, mask) + def masked_select_prefix(self, x, mask, backend=None): + return self._binary("cactus_graph_masked_select_prefix", x, mask, backend=backend) - def masked_scatter(self, x, mask, source): + def masked_scatter(self, x, mask, source, backend=None): x = self._ensure_tensor(x) mask = self._ensure_tensor(mask) source = self._ensure_tensor(source) @@ -1743,9 +1690,9 @@ def masked_scatter(self, x, mask, source): ) if rc != 0: raise RuntimeError(_err("graph_masked_scatter failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def view(self, x, shape): + def view(self, x, shape, backend=None): x = self._ensure_tensor(x) shape = tuple(int(v) for v in shape) arr = (ctypes.c_size_t * len(shape))(*shape) @@ -1753,9 +1700,9 @@ def view(self, x, shape): rc = _lib.cactus_graph_view(self.h, cactus_node_t(x.id), arr, len(shape), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_view failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def reshape(self, x, shape): + def reshape(self, x, shape, backend=None): x = self._ensure_tensor(x) shape = tuple(int(v) for v in shape) arr = (ctypes.c_size_t * len(shape))(*shape) @@ -1763,9 +1710,9 @@ def reshape(self, x, shape): rc = _lib.cactus_graph_reshape(self.h, cactus_node_t(x.id), arr, len(shape), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_reshape failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def expand(self, x, shape): + def expand(self, x, shape, backend=None): x = self._ensure_tensor(x) shape = tuple(int(v) for v in shape) arr = (ctypes.c_size_t * len(shape))(*shape) @@ -1773,9 +1720,9 @@ def expand(self, x, shape): rc = _lib.cactus_graph_expand(self.h, cactus_node_t(x.id), arr, len(shape), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_expand failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def flatten(self, x, start_dim=0, end_dim=-1): + def flatten(self, x, start_dim=0, end_dim=-1, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_flatten( @@ -1787,9 +1734,9 @@ def flatten(self, x, start_dim=0, end_dim=-1): ) if rc != 0: raise RuntimeError(_err("graph_flatten failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def slice(self, x, axis, start, length=0): + def slice(self, x, axis, start, length=0, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_slice( @@ -1802,9 +1749,9 @@ def slice(self, x, axis, start, length=0): ) if rc != 0: raise RuntimeError(_err("graph_slice failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def index(self, x, index_value, axis=0): + def index(self, x, index_value, axis=0, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_index( @@ -1816,9 +1763,11 @@ def index(self, x, index_value, axis=0): ) if rc != 0: raise RuntimeError(_err("graph_index failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def transpose(self, x, backend=CPU): + def transpose(self, x, backend=None): + if backend is None: + backend = _lib.cactus_graph_default_backend() x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_transpose( @@ -1831,7 +1780,9 @@ def transpose(self, x, backend=CPU): raise RuntimeError(_err("graph_transpose failed")) return self._tensor_from_node(out.value) - def permute(self, x, permutation, backend=CPU): + def permute(self, x, permutation, backend=None): + if backend is None: + backend = _lib.cactus_graph_default_backend() x = self._ensure_tensor(x) permutation = tuple(int(v) for v in permutation) arr = (ctypes.c_size_t * len(permutation))(*permutation) @@ -1848,7 +1799,9 @@ def permute(self, x, permutation, backend=CPU): raise RuntimeError(_err("graph_transpose_n failed")) return self._tensor_from_node(out.value) - def matmul(self, a, b, pretransposed_rhs=False, backend=CPU, output_dtype=None): + def matmul(self, a, b, pretransposed_rhs=False, backend=None, output_dtype=None): + if backend is None: + backend = _lib.cactus_graph_default_backend() a = self._ensure_tensor(a) b = self._ensure_tensor(b) out = cactus_node_t() @@ -1862,9 +1815,12 @@ def matmul(self, a, b, pretransposed_rhs=False, backend=CPU, output_dtype=None): ) if rc != 0: raise RuntimeError(_err("graph_matmul failed")) - return self._tensor_from_node(out.value) + result = self._tensor_from_node(out.value) + if output_dtype is not None and int(output_dtype) != int(result.dtype): + result = self.precision_cast(result, int(output_dtype)) + return result - def gather(self, tensor, indices): + def gather(self, tensor, indices, backend=None): tensor = self._ensure_tensor(tensor) indices = self._ensure_tensor(indices) out = cactus_node_t() @@ -1876,9 +1832,9 @@ def gather(self, tensor, indices): ) if rc != 0: raise RuntimeError(_err("graph_gather failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def embedding_from_tensor(self, embedding_tensor, indices): + def embedding_from_tensor(self, embedding_tensor, indices, backend=None): embedding_tensor = self._ensure_tensor(embedding_tensor) indices = self._ensure_tensor(indices) out = cactus_node_t() @@ -1887,15 +1843,15 @@ def embedding_from_tensor(self, embedding_tensor, indices): ) if rc != 0: raise RuntimeError(_err("graph_embedding_from_tensor failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def embedding_from_file(self, filename, indices): + def embedding_from_file(self, filename, indices, backend=None): indices = self._ensure_tensor(indices) out = cactus_node_t() rc = _lib.cactus_graph_embedding_from_file(self.h, str(filename).encode(), cactus_node_t(indices.id), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_embedding_from_file failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def mmap_embeddings(self, filename): out = cactus_node_t() @@ -1917,7 +1873,7 @@ def bind_mmap_weights(self, tensor, filename): if rc != 0: raise RuntimeError(_err("graph_bind_mmap_weights failed")) - def bilinear_interpolation(self, pos_embeds, dst_height, dst_width): + def bilinear_interpolation(self, pos_embeds, dst_height, dst_width, backend=None): pos_embeds = self._ensure_tensor(pos_embeds) out = cactus_node_t() rc = _lib.cactus_graph_bilinear_interpolation( @@ -1925,7 +1881,7 @@ def bilinear_interpolation(self, pos_embeds, dst_height, dst_width): ) if rc != 0: raise RuntimeError(_err("graph_bilinear_interpolation failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def set_grouped_scales(self, tensor, group_size, num_groups, scales): tensor = self._ensure_tensor(tensor) @@ -1965,7 +1921,7 @@ def release_all_weight_pages(self): if rc != 0: raise RuntimeError(_err("graph_release_all_weight_pages failed")) - def concat(self, a, b, axis=0): + def concat(self, a, b, axis=0, backend=None): a = self._ensure_tensor(a) b = self._ensure_tensor(b) out = cactus_node_t() @@ -1978,9 +1934,9 @@ def concat(self, a, b, axis=0): ) if rc != 0: raise RuntimeError(_err("graph_concat failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def cat(self, tensors, axis=0): + def cat(self, tensors, axis=0, backend=None): tensors = [self._ensure_tensor(t) for t in tensors] if not tensors: raise ValueError("cat requires at least one tensor") @@ -1994,9 +1950,9 @@ def cat(self, tensors, axis=0): ) if rc != 0: raise RuntimeError(_err("graph_cat failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def groupnorm(self, x, weight, bias, num_groups, eps=1e-5): + def groupnorm(self, x, weight, bias, num_groups, eps=1e-5, backend=None): x = self._ensure_tensor(x) weight = self._ensure_tensor(weight) bias = self._ensure_tensor(bias) @@ -2012,12 +1968,12 @@ def groupnorm(self, x, weight, bias, num_groups, eps=1e-5): ) if rc != 0: raise RuntimeError(_err("graph_group_norm failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def group_norm(self, x, weight, bias, num_groups, eps=1e-5): - return self.groupnorm(x, weight, bias, num_groups, eps=eps) + def group_norm(self, x, weight, bias, num_groups, eps=1e-5, backend=None): + return self.groupnorm(x, weight, bias, num_groups, eps=eps, backend=backend) - def layernorm(self, x, weight, bias=None, eps=1e-5): + def layernorm(self, x, weight, bias=None, eps=1e-5, backend=None): x = self._ensure_tensor(x) weight = self._ensure_tensor(weight) has_bias = bias is not None @@ -2037,12 +1993,12 @@ def layernorm(self, x, weight, bias=None, eps=1e-5): ) if rc != 0: raise RuntimeError(_err("graph_layer_norm failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def layer_norm(self, x, weight, bias=None, eps=1e-5): - return self.layernorm(x, weight, bias=bias, eps=eps) + def layer_norm(self, x, weight, bias=None, eps=1e-5, backend=None): + return self.layernorm(x, weight, bias=bias, eps=eps, backend=backend) - def batchnorm(self, x, weight, bias, running_mean, running_var, axis=1, eps=1e-5): + def batchnorm(self, x, weight, bias, running_mean, running_var, axis=1, eps=1e-5, backend=None): x = self._ensure_tensor(x) weight = self._ensure_tensor(weight) bias = self._ensure_tensor(bias) @@ -2062,12 +2018,12 @@ def batchnorm(self, x, weight, bias, running_mean, running_var, axis=1, eps=1e-5 ) if rc != 0: raise RuntimeError(_err("graph_batchnorm failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def batch_norm(self, x, weight, bias, running_mean, running_var, axis=1, eps=1e-5): - return self.batchnorm(x, weight, bias, running_mean, running_var, axis=axis, eps=eps) + def batch_norm(self, x, weight, bias, running_mean, running_var, axis=1, eps=1e-5, backend=None): + return self.batchnorm(x, weight, bias, running_mean, running_var, axis=axis, eps=eps, backend=backend) - def rms_norm(self, x, weight, eps=1e-5): + def rms_norm(self, x, weight, eps=1e-5, backend=None): x = self._ensure_tensor(x) weight = self._ensure_tensor(weight) out = cactus_node_t() @@ -2080,17 +2036,19 @@ def rms_norm(self, x, weight, eps=1e-5): ) if rc != 0: raise RuntimeError(_err("graph_rms_norm failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def topk(self, x, k): + def topk(self, x, k, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_topk(self.h, cactus_node_t(x.id), ctypes.c_size_t(int(k)), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_topk failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def rope(self, x, theta, position_offset=0, backend=CPU): + def rope(self, x, theta, position_offset=0, backend=None): + if backend is None: + backend = _lib.cactus_graph_default_backend() x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_rope( @@ -2101,7 +2059,9 @@ def rope(self, x, theta, position_offset=0, backend=CPU): raise RuntimeError(_err("graph_rope failed")) return self._tensor_from_node(out.value) - def rope_gptj(self, x, theta, position_offset=0, rot_dim=0, backend=CPU): + def rope_gptj(self, x, theta, position_offset=0, rot_dim=0, backend=None): + if backend is None: + backend = _lib.cactus_graph_default_backend() x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_rope_gptj( @@ -2112,33 +2072,33 @@ def rope_gptj(self, x, theta, position_offset=0, rot_dim=0, backend=CPU): raise RuntimeError(_err("graph_rope_gptj failed")) return self._tensor_from_node(out.value) - def _reduce(self, fn_name, x, axis): + def _reduce(self, fn_name, x, axis, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = getattr(_lib, fn_name)(self.h, cactus_node_t(x.id), ctypes.c_int32(int(axis)), ctypes.byref(out)) if rc != 0: raise RuntimeError(f"{fn_name} failed") - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def sum(self, x, axis): - return self._reduce("cactus_graph_sum", x, axis) + def sum(self, x, axis, backend=None): + return self._reduce("cactus_graph_sum", x, axis, backend=backend) - def mean(self, x, axis): - return self._reduce("cactus_graph_mean", x, axis) + def mean(self, x, axis, backend=None): + return self._reduce("cactus_graph_mean", x, axis, backend=backend) - def variance(self, x, axis): - return self._reduce("cactus_graph_variance", x, axis) + def variance(self, x, axis, backend=None): + return self._reduce("cactus_graph_variance", x, axis, backend=backend) - def min(self, x, axis): - return self._reduce("cactus_graph_min", x, axis) + def min(self, x, axis, backend=None): + return self._reduce("cactus_graph_min", x, axis, backend=backend) - def max(self, x, axis): - return self._reduce("cactus_graph_max", x, axis) + def max(self, x, axis, backend=None): + return self._reduce("cactus_graph_max", x, axis, backend=backend) - def cumsum(self, x, axis): - return self._reduce("cactus_graph_cumsum", x, axis) + def cumsum(self, x, axis, backend=None): + return self._reduce("cactus_graph_cumsum", x, axis, backend=backend) - def softmax(self, x, axis=-1): + def softmax(self, x, axis=-1, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_softmax( @@ -2149,10 +2109,12 @@ def softmax(self, x, axis=-1): ) if rc != 0: raise RuntimeError(_err("graph_softmax failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def attention(self, query, key, value, scale, is_causal=True, position_offset=0, window_size=0, - backend=CPU, mask=None, additive_mask=False): + backend=None, mask=None, additive_mask=False): + if backend is None: + backend = _lib.cactus_graph_default_backend() query = self._ensure_tensor(query) key = self._ensure_tensor(key) value = self._ensure_tensor(value) @@ -2180,7 +2142,7 @@ def attention(self, query, key, value, scale, is_causal=True, position_offset=0, raise RuntimeError(_err("graph_attention failed")) return self._tensor_from_node(out.value) - def kv_cache_state(self, max_seq_len, num_kv_heads, head_dim, window_size=0, sink_size=0, num_slots=1): + def kv_cache_state(self, max_seq_len, num_kv_heads, head_dim, window_size=0, sink_size=0, num_slots=1, backend=None): out = cactus_node_t() rc = _lib.cactus_graph_kv_cache_state( self.h, @@ -2194,9 +2156,9 @@ def kv_cache_state(self, max_seq_len, num_kv_heads, head_dim, window_size=0, sin ) if rc != 0: raise RuntimeError(_err("graph_kv_cache_state failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def kv_cache_append(self, new_kv, cache_state, window_size=0, sink_size=0): + def kv_cache_append(self, new_kv, cache_state, window_size=0, sink_size=0, backend=None): new_kv = self._ensure_tensor(new_kv) cache_state = self._ensure_tensor(cache_state) out = cactus_node_t() @@ -2210,7 +2172,7 @@ def kv_cache_append(self, new_kv, cache_state, window_size=0, sink_size=0): ) if rc != 0: raise RuntimeError(_err("graph_kv_cache_append failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def attention_cached( self, @@ -2223,6 +2185,7 @@ def attention_cached( position_offset=0, window_size=0, v_head_dim=0, + backend=None, ): query = self._ensure_tensor(query) key_new = self._ensure_tensor(key_new) @@ -2245,9 +2208,9 @@ def attention_cached( ) if rc != 0: raise RuntimeError(_err("graph_attention_cached failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def conv_cache_state(self, window_size, hidden_dim): + def conv_cache_state(self, window_size, hidden_dim, backend=None): out = cactus_node_t() rc = _lib.cactus_graph_conv_cache_state( self.h, @@ -2257,9 +2220,9 @@ def conv_cache_state(self, window_size, hidden_dim): ) if rc != 0: raise RuntimeError(_err("graph_conv_cache_state failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def conv_cache_append(self, new_data, cache_state): + def conv_cache_append(self, new_data, cache_state, backend=None): new_data = self._ensure_tensor(new_data) cache_state = self._ensure_tensor(cache_state) out = cactus_node_t() @@ -2271,9 +2234,9 @@ def conv_cache_append(self, new_data, cache_state): ) if rc != 0: raise RuntimeError(_err("graph_conv_cache_append failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def conv_cache_initialize(self, rows, cache_state): + def conv_cache_initialize(self, rows, cache_state, backend=None): rows = self._ensure_tensor(rows) cache_state = self._ensure_tensor(cache_state) out = cactus_node_t() @@ -2285,9 +2248,9 @@ def conv_cache_initialize(self, rows, cache_state): ) if rc != 0: raise RuntimeError(_err("graph_conv_cache_initialize failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def recurrent_cache_state(self, shape, dtype=FP16): + def recurrent_cache_state(self, shape, dtype=FP16, backend=None): shape = tuple(int(x) for x in shape) arr = (ctypes.c_size_t * len(shape))(*shape) out = cactus_node_t() @@ -2300,9 +2263,9 @@ def recurrent_cache_state(self, shape, dtype=FP16): ) if rc != 0: raise RuntimeError(_err("graph_recurrent_cache_state failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def recurrent_cache_write(self, new_value, cache_input): + def recurrent_cache_write(self, new_value, cache_input, backend=None): new_value = self._ensure_tensor(new_value) cache_input = self._ensure_tensor(cache_input) out = cactus_node_t() @@ -2314,9 +2277,9 @@ def recurrent_cache_write(self, new_value, cache_input): ) if rc != 0: raise RuntimeError(_err("graph_recurrent_cache_write failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def rel_pos_bias(self, query, relative_key, scale): + def rel_pos_bias(self, query, relative_key, scale, backend=None): query = self._ensure_tensor(query) relative_key = self._ensure_tensor(relative_key) out = cactus_node_t() @@ -2325,11 +2288,11 @@ def rel_pos_bias(self, query, relative_key, scale): ) if rc != 0: raise RuntimeError(_err("graph_rel_pos_bias failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def attention_int8_hybrid(self, query, key_new, value_new, scale, position_offset, cached_keys, cached_values, k_scales, v_scales, - cache_len, num_kv_heads, head_dim, window_size=0): + cache_len, num_kv_heads, head_dim, window_size=0, backend=None): query = self._ensure_tensor(query) key_new = self._ensure_tensor(key_new) value_new = self._ensure_tensor(value_new) @@ -2357,65 +2320,65 @@ def attention_int8_hybrid(self, query, key_new, value_new, scale, position_offse ) if rc != 0: raise RuntimeError("graph_attention_int8_hybrid failed") - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def relu(self, x): + def relu(self, x, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_relu(self.h, cactus_node_t(x.id), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_relu failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def silu(self, x): + def silu(self, x, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_silu(self.h, cactus_node_t(x.id), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_silu failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def gelu(self, x): + def gelu(self, x, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_gelu(self.h, cactus_node_t(x.id), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_gelu failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def gelu_erf(self, x): + def gelu_erf(self, x, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_gelu_erf(self.h, cactus_node_t(x.id), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_gelu_erf failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def sigmoid(self, x): + def sigmoid(self, x, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_sigmoid(self.h, cactus_node_t(x.id), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_sigmoid failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def tanh(self, x): + def tanh(self, x, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_tanh(self.h, cactus_node_t(x.id), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_tanh failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def glu(self, x, axis=-1): + def glu(self, x, axis=-1, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_glu(self.h, cactus_node_t(x.id), ctypes.c_int32(int(axis)), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_glu failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def conv1d_causal(self, x, weight, kernel_size, dilation): + def conv1d_causal(self, x, weight, kernel_size, dilation, backend=None): x = self._ensure_tensor(x) weight = self._ensure_tensor(weight) out = cactus_node_t() @@ -2425,9 +2388,9 @@ def conv1d_causal(self, x, weight, kernel_size, dilation): ) if rc != 0: raise RuntimeError("graph_conv1d_causal failed") - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def conv1d_k3(self, x, weight, stride=1): + def conv1d_k3(self, x, weight, stride=1, backend=None): x = self._ensure_tensor(x) weight = self._ensure_tensor(weight) out = cactus_node_t() @@ -2436,9 +2399,9 @@ def conv1d_k3(self, x, weight, stride=1): ) if rc != 0: raise RuntimeError("graph_conv1d_k3 failed") - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def conv1d_k7s3(self, x, weight, bias): + def conv1d_k7s3(self, x, weight, bias, backend=None): x = self._ensure_tensor(x) weight = self._ensure_tensor(weight) bias = self._ensure_tensor(bias) @@ -2448,30 +2411,30 @@ def conv1d_k7s3(self, x, weight, bias): ) if rc != 0: raise RuntimeError("graph_conv1d_k7s3 failed") - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def conv1d(self, x, weight, bias=None, stride=1): - return self._conv_with_optional_bias("cactus_graph_conv1d", x, weight, bias, ctypes.c_size_t(int(stride))) + def conv1d(self, x, weight, bias=None, stride=1, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv1d", x, weight, bias, ctypes.c_size_t(int(stride)), backend=backend) - def conv1d_same_depthwise_k9(self, x, weight, bias=None): - return self._conv_with_optional_bias("cactus_graph_conv1d_same_depthwise_k9", x, weight, bias) + def conv1d_same_depthwise_k9(self, x, weight, bias=None, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv1d_same_depthwise_k9", x, weight, bias, backend=backend) - def conv1d_pointwise(self, x, weight, bias=None): - return self._conv_with_optional_bias("cactus_graph_conv1d_pointwise", x, weight, bias) + def conv1d_pointwise(self, x, weight, bias=None, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv1d_pointwise", x, weight, bias, backend=backend) - def conv2d_k3s2p1(self, x, weight, bias=None): - return self._conv_with_optional_bias("cactus_graph_conv2d_k3s2p1", x, weight, bias) + def conv2d_k3s2p1(self, x, weight, bias=None, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv2d_k3s2p1", x, weight, bias, backend=backend) - def conv2d_depthwise_k3s2p1(self, x, weight, bias=None): - return self._conv_with_optional_bias("cactus_graph_conv2d_depthwise_k3s2p1", x, weight, bias) + def conv2d_depthwise_k3s2p1(self, x, weight, bias=None, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv2d_depthwise_k3s2p1", x, weight, bias, backend=backend) - def conv2d_pointwise_1x1(self, x, weight, bias=None): - return self._conv_with_optional_bias("cactus_graph_conv2d_pointwise_1x1", x, weight, bias) + def conv2d_pointwise_1x1(self, x, weight, bias=None, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv2d_pointwise_1x1", x, weight, bias, backend=backend) - def conv2d_k3s1p1(self, x, weight, bias=None): - return self._conv_with_optional_bias("cactus_graph_conv2d_k3s1p1", x, weight, bias) + def conv2d_k3s1p1(self, x, weight, bias=None, backend=None): + return self._conv_with_optional_bias("cactus_graph_conv2d_k3s1p1", x, weight, bias, backend=backend) - def conv2d(self, x, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): + def conv2d(self, x, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, backend=None): return self._conv_with_optional_bias( "cactus_graph_conv2d", x, @@ -2481,9 +2444,10 @@ def conv2d(self, x, weight, bias=None, stride=1, padding=0, dilation=1, groups=1 ctypes.c_size_t(int(padding)), ctypes.c_size_t(int(dilation)), ctypes.c_size_t(int(groups)), + backend=backend, ) - def _conv_with_optional_bias(self, fn_name, x, weight, bias=None, *extra): + def _conv_with_optional_bias(self, fn_name, x, weight, bias=None, *extra, backend=None): x = self._ensure_tensor(x) weight = self._ensure_tensor(weight) has_bias = bias is not None @@ -2494,17 +2458,17 @@ def _conv_with_optional_bias(self, fn_name, x, weight, bias=None, *extra): ) if rc != 0: raise RuntimeError(f"{fn_name} failed") - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def lstm_cell(self, input, h_prev, c_prev, weight_ih, weight_hh, bias_ih, bias_hh): + def lstm_cell(self, input, h_prev, c_prev, weight_ih, weight_hh, bias_ih, bias_hh, backend=None): tensors = [self._ensure_tensor(t) for t in (input, h_prev, c_prev, weight_ih, weight_hh, bias_ih, bias_hh)] out = cactus_node_t() rc = _lib.cactus_graph_lstm_cell(self.h, *(cactus_node_t(t.id) for t in tensors), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_lstm_cell failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def gated_deltanet_decode(self, query, key, value, gate_log, beta, initial_state, scale): + def gated_deltanet_decode(self, query, key, value, gate_log, beta, initial_state, scale, backend=None): tensors = [self._ensure_tensor(t) for t in (query, key, value, gate_log, beta, initial_state)] out = cactus_node_t() rc = _lib.cactus_graph_gated_deltanet_decode( @@ -2512,9 +2476,9 @@ def gated_deltanet_decode(self, query, key, value, gate_log, beta, initial_state ) if rc != 0: raise RuntimeError(_err("graph_gated_deltanet_decode failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def gated_deltanet_prefill(self, query, key, value, gate_log, beta, initial_state, chunk_size, scale): + def gated_deltanet_prefill(self, query, key, value, gate_log, beta, initial_state, chunk_size, scale, backend=None): tensors = [self._ensure_tensor(t) for t in (query, key, value, gate_log, beta, initial_state)] out = cactus_node_t() rc = _lib.cactus_graph_gated_deltanet_prefill( @@ -2522,9 +2486,9 @@ def gated_deltanet_prefill(self, query, key, value, gate_log, beta, initial_stat ) if rc != 0: raise RuntimeError(_err("graph_gated_deltanet_prefill failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def stft(self, x, weight, stride, num_fft_bins): + def stft(self, x, weight, stride, num_fft_bins, backend=None): x = self._ensure_tensor(x) weight = self._ensure_tensor(weight) out = cactus_node_t() @@ -2534,9 +2498,9 @@ def stft(self, x, weight, stride, num_fft_bins): ) if rc != 0: raise RuntimeError(_err("graph_stft failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def altup_predict(self, coefs, streams): + def altup_predict(self, coefs, streams, backend=None): coefs = self._ensure_tensor(coefs) streams = [self._ensure_tensor(t) for t in streams] ids = (cactus_node_t * len(streams))(*(cactus_node_t(t.id) for t in streams)) @@ -2544,9 +2508,9 @@ def altup_predict(self, coefs, streams): rc = _lib.cactus_graph_altup_predict(self.h, cactus_node_t(coefs.id), ids, ctypes.c_size_t(len(streams)), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_altup_predict failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def altup_correct(self, coefs, innovation, predictions): + def altup_correct(self, coefs, innovation, predictions, backend=None): coefs = self._ensure_tensor(coefs) innovation = self._ensure_tensor(innovation) predictions = [self._ensure_tensor(t) for t in predictions] @@ -2557,18 +2521,18 @@ def altup_correct(self, coefs, innovation, predictions): ) if rc != 0: raise RuntimeError(_err("graph_altup_correct failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def gaussian_topk(self, x, ppf): + def gaussian_topk(self, x, ppf, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_gaussian_topk(self.h, cactus_node_t(x.id), ctypes.c_float(float(ppf)), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_gaussian_topk failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def moe_layer_gated(self, hidden, routing_probs, topk_indices, w1_weights, w3_weights, w2_weights, - num_experts, num_experts_per_tok, normalize_routing=True, epsilon=1e-6, routed_scaling_factor=1.0, activation=0): + num_experts, num_experts_per_tok, normalize_routing=True, epsilon=1e-6, routed_scaling_factor=1.0, activation=0, backend=None): hidden = self._ensure_tensor(hidden) routing_probs = self._ensure_tensor(routing_probs) topk_indices = self._ensure_tensor(topk_indices) @@ -2584,9 +2548,9 @@ def moe_layer_gated(self, hidden, routing_probs, topk_indices, w1_weights, w3_we ) if rc != 0: raise RuntimeError(_err("graph_moe_layer_gated failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def dense_mlp_tq_fused(self, hidden, gate_weight, up_weight, down_weight, product_scale=1.0): + def dense_mlp_tq_fused(self, hidden, gate_weight, up_weight, down_weight, product_scale=1.0, backend=None): hidden = self._ensure_tensor(hidden) gate_weight = self._ensure_tensor(gate_weight) up_weight = self._ensure_tensor(up_weight) @@ -2603,11 +2567,11 @@ def dense_mlp_tq_fused(self, hidden, gate_weight, up_weight, down_weight, produc ) if rc != 0: raise RuntimeError(_err("graph_dense_mlp_tq_fused failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def moe_layer_ungated(self, hidden, routing_probs, topk_indices, w1_weights, w2_weights, num_experts, num_experts_per_tok, normalize_routing=True, epsilon=1e-6, - routed_scaling_factor=1.0, activation=ACT_GELU): + routed_scaling_factor=1.0, activation=ACT_GELU, backend=None): hidden = self._ensure_tensor(hidden) routing_probs = self._ensure_tensor(routing_probs) topk_indices = self._ensure_tensor(topk_indices) @@ -2622,9 +2586,9 @@ def moe_layer_ungated(self, hidden, routing_probs, topk_indices, w1_weights, w2_ ) if rc != 0: raise RuntimeError(_err("graph_moe_layer_ungated failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def sample(self, logits, temperature=0.6, top_p=0.95, top_k=20): + def sample(self, logits, temperature=0.6, top_p=0.95, top_k=20, backend=None): logits = self._ensure_tensor(logits) out = cactus_node_t() rc = _lib.cactus_graph_sample( @@ -2633,9 +2597,9 @@ def sample(self, logits, temperature=0.6, top_p=0.95, top_k=20): ) if rc != 0: raise RuntimeError(_err("graph_sample failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def scatter_topk(self, indices, values, num_classes): + def scatter_topk(self, indices, values, num_classes, backend=None): indices = self._ensure_tensor(indices) values = self._ensure_tensor(values) out = cactus_node_t() @@ -2644,15 +2608,15 @@ def scatter_topk(self, indices, values, num_classes): ) if rc != 0: raise RuntimeError(_err("graph_scatter_topk failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def persistent(self, source_node): + def persistent(self, source_node, backend=None): source_node = self._ensure_tensor(source_node) out = cactus_node_t() rc = _lib.cactus_graph_persistent(self.h, cactus_node_t(source_node.id), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_persistent failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def is_populated(self, persistent_node): persistent_node = self._ensure_tensor(persistent_node) @@ -2670,16 +2634,16 @@ def invalidate_persistent(self, persistent_node): # ── Audio / signal processing ─────────────────────────────────── - def rfft(self, x): + def rfft(self, x, backend=None): """Real-to-complex FFT.""" x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_rfft(self.h, cactus_node_t(x.id), ctypes.byref(out)) if rc != 0: raise RuntimeError(_err("graph_rfft failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def irfft(self, x, output_length): + def irfft(self, x, output_length, backend=None): """Complex-to-real inverse FFT.""" x = self._ensure_tensor(x) out = cactus_node_t() @@ -2689,11 +2653,11 @@ def irfft(self, x, output_length): ) if rc != 0: raise RuntimeError(_err("graph_irfft failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def mel_filter_bank(self, num_frequency_bins, num_mel_filters, min_frequency, max_frequency, sampling_rate, - norm_type=0, scale_type=0): + norm_type=0, scale_type=0, backend=None): """Generate a mel-scale filter bank tensor.""" out = cactus_node_t() rc = _lib.cactus_graph_mel_filter_bank( @@ -2709,12 +2673,12 @@ def mel_filter_bank(self, num_frequency_bins, num_mel_filters, ) if rc != 0: raise RuntimeError(_err("graph_mel_filter_bank failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def spectrogram(self, waveform, mel_filters, frame_length, hop_length, fft_length, power=2.0, center=True, pad_mode=0, mel_floor=1e-10, log_mel_mode=0, dither=0.0, - preemphasis=0.0, remove_dc_offset=False): + preemphasis=0.0, remove_dc_offset=False, backend=None): """Compute a spectrogram from a waveform tensor.""" waveform = self._ensure_tensor(waveform) mel_filters = self._ensure_tensor(mel_filters) @@ -2738,13 +2702,13 @@ def spectrogram(self, waveform, mel_filters, frame_length, hop_length, ) if rc != 0: raise RuntimeError(_err("graph_spectrogram failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) # ── Image processing ───────────────────────────────────────────── def image_preprocess(self, pixel_input, src_width, src_height, target_width, target_height, patch_size, channels, - rescale_factor, mean, std_dev): + rescale_factor, mean, std_dev, backend=None): """Resize, normalize, and patch an image tensor. *mean* and *std_dev* are per-channel float sequences (length = channels). @@ -2766,7 +2730,7 @@ def image_preprocess(self, pixel_input, src_width, src_height, ) if rc != 0: raise RuntimeError(_err("graph_image_preprocess failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) # ── Introspection ──────────────────────────────────────────────── @@ -2774,14 +2738,22 @@ def output_info(self, x): x = self._ensure_tensor(x) return self._get_output_info(x.id) - def _binary(self, fn_name, a, b): + def _apply_backend(self, tensor, backend): + if backend is None: + return tensor + rc = _lib.cactus_graph_set_node_backend(self.h, cactus_node_t(tensor.id), ctypes.c_int32(int(backend))) + if rc != 0: + raise RuntimeError(_err("set_node_backend failed")) + return tensor + + def _binary(self, fn_name, a, b, backend=None): a = self._ensure_tensor(a) b = self._ensure_tensor(b) out = cactus_node_t() rc = getattr(_lib, fn_name)(self.h, cactus_node_t(a.id), cactus_node_t(b.id), ctypes.byref(out)) if rc != 0: raise RuntimeError(f"{fn_name} failed") - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def _ensure_tensor(self, x): if not isinstance(x, Tensor): @@ -2934,10 +2906,10 @@ def slice(self, axis, start, length=0): def index(self, index_value, axis=0): return self.g.index(self, index_value, axis=axis) - def transpose(self, backend=Graph.CPU): + def transpose(self, backend=None): return self.g.transpose(self, backend=backend) - def permute(self, permutation, backend=Graph.CPU): + def permute(self, permutation, backend=None): return self.g.permute(self, permutation, backend=backend) def concat(self, other, axis=0): @@ -2973,7 +2945,7 @@ def softmax(self, axis=-1): def glu(self, axis=-1): return self.g.glu(self, axis=axis) - def matmul(self, other, pretransposed_rhs=False, backend=Graph.CPU, output_dtype=None): + def matmul(self, other, pretransposed_rhs=False, backend=None, output_dtype=None): return self.g.matmul( self, other, diff --git a/python/cactus/cli/__init__.py b/python/cactus/cli/__init__.py index 6cf394f02..fd7a005bc 100644 --- a/python/cactus/cli/__init__.py +++ b/python/cactus/cli/__init__.py @@ -45,8 +45,8 @@ def _build_parent(): def _engine_test_parent(): """Args shared by `test` and `benchmark`.""" p = argparse.ArgumentParser(add_help=False) - p.add_argument("--backend", choices=["auto", "cpu", "metal"], default="auto", - help="Inference backend, default auto (Metal on Apple Silicon GPUs)") + p.add_argument("--backend", choices=["cpu", "metal"], default=None, + help="Inference backend (default: auto)") p.add_argument("--model", dest="model_id", default=None, type=_hf_id_or_path, help=f"HF model ID under test (default: {DEFAULT_TEST_MODEL_ID})") @@ -180,7 +180,7 @@ def create_parser(): --model default: {DEFAULT_TEST_MODEL_ID} --transcription-model default: {DEFAULT_TEST_TRANSCRIPTION_MODEL_ID} --bits 1|2|3|4 CQ quantization (default: 4) - --backend auto|cpu|metal inference backend (default: auto) + --backend cpu|metal inference backend (default: auto) --ios run on connected iPhone --android run on connected Android @@ -289,8 +289,8 @@ def create_parser(): help="Confidence threshold below which local completions may hand off to cloud") run_parser.add_argument("--cloud-timeout-ms", type=_non_negative_int, default=None, help="Maximum time to wait for cloud handoff before falling back locally") - run_parser.add_argument("--backend", choices=["auto", "cpu", "metal"], default="auto", - help="Inference backend, default auto (Metal on Apple Silicon GPUs)") + run_parser.add_argument("--backend", choices=["cpu", "metal"], default=None, + help="Inference backend (default: auto)") transcribe_parser = subparsers.add_parser("transcribe", help="Transcribe audio with a model", parents=[_telemetry_parent(), _build_parent()]) @@ -319,8 +319,8 @@ def create_parser(): help="Maximum time to wait for cloud handoff before falling back locally") serve_parser.add_argument("--no-access-log", action="store_true", help="Disable per-request HTTP access logging") - serve_parser.add_argument("--backend", choices=["auto", "cpu", "metal"], default="auto", - help="Inference backend, default auto (Metal on Apple Silicon GPUs)") + serve_parser.add_argument("--backend", choices=["cpu", "metal"], default=None, + help="Inference backend (default: auto)") code_parser = subparsers.add_parser("code", help="Run the Cactus coding agent (TUI / print mode)", parents=[_build_parent()]) @@ -338,8 +338,8 @@ def create_parser(): help="Confidence threshold below which completions hand off to cloud") code_parser.add_argument("--cloud-timeout-ms", type=_non_negative_int, default=None, help="Maximum time to wait for cloud handoff before falling back locally") - code_parser.add_argument("--backend", choices=["auto", "cpu", "metal"], default="auto", - help="Inference backend, default auto (Metal on Apple Silicon GPUs)") + code_parser.add_argument("--backend", choices=["cpu", "metal"], default=None, + help="Inference backend (default: auto)") code_parser.add_argument("agent_args", nargs=argparse.REMAINDER, help="Arguments passed through to the coding agent (prefix with -- )") diff --git a/python/cactus/cli/code.py b/python/cactus/cli/code.py index 7d38b7077..090528544 100644 --- a/python/cactus/cli/code.py +++ b/python/cactus/cli/code.py @@ -145,8 +145,9 @@ def cmd_code(args) -> int: sys.executable, "-m", "cactus", "serve", serve_model, "--host", args.host, "--port", str(args.port), "--no-access-log", "--bits", str(getattr(args, "bits", 4)), - "--backend", getattr(args, "backend", "auto"), ] + if getattr(args, "backend", None): + serve_cmd += ["--backend", args.backend] if getattr(args, "token", None): serve_cmd += ["--token", args.token] if getattr(args, "reconvert", False): diff --git a/python/cactus/cli/serve.py b/python/cactus/cli/serve.py index 67ee60e11..f244743de 100644 --- a/python/cactus/cli/serve.py +++ b/python/cactus/cli/serve.py @@ -54,8 +54,8 @@ def cmd_serve(args): print_color(RED, "Error: uvicorn not installed. Install the serve extra or run `pip install fastapi uvicorn python-multipart`.") return 1 - backend = getattr(args, "backend", "auto") - if backend and backend != "auto": + backend = getattr(args, "backend", None) + if backend: from ..bindings.cactus import cactus_set_backend if cactus_set_backend(backend) != 0: print_color(YELLOW, f"Backend '{backend}' not available; using auto") diff --git a/python/cactus/cli/test.py b/python/cactus/cli/test.py index 50c50ed6c..720f9268b 100644 --- a/python/cactus/cli/test.py +++ b/python/cactus/cli/test.py @@ -31,7 +31,8 @@ def _component_args(component, args): if component == "engine": cmd.extend(["--model", args.model_id]) cmd.extend(["--transcription-model", args.transcription_model_id]) - cmd.extend(["--backend", getattr(args, "backend", "auto")]) + if getattr(args, "backend", None): + cmd.extend(["--backend", args.backend]) if args.android: cmd.append("--android") if args.ios: diff --git a/python/cactus/transpile/runtime_compat.py b/python/cactus/transpile/runtime_compat.py index 028682a96..85acfdd0b 100644 --- a/python/cactus/transpile/runtime_compat.py +++ b/python/cactus/transpile/runtime_compat.py @@ -110,7 +110,9 @@ def _approx_nonzero_mask(self, tensor): orig_pow = Graph.pow orig_expand = Graph.expand - def matmul(self, a, b, pretransposed_rhs=False, backend=Graph.CPU, output_dtype=None): + def matmul(self, a, b, pretransposed_rhs=False, backend=None, output_dtype=None): + if backend is None: + backend = _lib.cactus_graph_default_backend() a = _ensure_fp16_activation(self, a) b = _ensure_fp16_activation(self, b) out = cactus_node_t() @@ -129,7 +131,7 @@ def matmul(self, a, b, pretransposed_rhs=False, backend=Graph.CPU, output_dtype= result = self.precision_cast(result, int(output_dtype)) return result - def gather(self, tensor, indices, axis=0): + def gather(self, tensor, indices, axis=0, backend=None): tensor = self._ensure_tensor(tensor) indices = self._ensure_tensor(indices) if int(axis) != 0: @@ -145,9 +147,9 @@ def gather(self, tensor, indices, axis=0): ) if rc != 0: raise RuntimeError("graph_gather failed") - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) - def conv2d(self, x, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): + def conv2d(self, x, weight, bias=None, stride=1, padding=0, dilation=1, groups=1, backend=None): def _pair(value: Any) -> tuple[int, int]: if isinstance(value, (tuple, list)): if len(value) != 2: @@ -171,6 +173,7 @@ def _pair(value: Any) -> tuple[int, int]: padding=padding, dilation=dilation, groups=groups, + backend=backend, ) if ( @@ -180,9 +183,9 @@ def _pair(value: Any) -> tuple[int, int]: and dilation_hw == (1, 1) ): if groups_int == 1: - return self.conv2d_k3s2p1(x, weight, bias=bias) + return self.conv2d_k3s2p1(x, weight, bias=bias, backend=backend) if len(weight.shape) >= 1 and groups_int == int(weight.shape[0]): - return self.conv2d_depthwise_k3s2p1(x, weight, bias=bias) + return self.conv2d_depthwise_k3s2p1(x, weight, bias=bias, backend=backend) if ( kernel_hw == (3, 3) @@ -192,7 +195,7 @@ def _pair(value: Any) -> tuple[int, int]: and groups_int == 1 and _has_symbol("cactus_graph_conv2d_k3s1p1") ): - return self.conv2d_k3s1p1(x, weight, bias=bias) + return self.conv2d_k3s1p1(x, weight, bias=bias, backend=backend) if ( kernel_hw == (1, 1) @@ -201,7 +204,7 @@ def _pair(value: Any) -> tuple[int, int]: and dilation_hw == (1, 1) and groups_int == 1 ): - return self.conv2d_pointwise_1x1(x, weight, bias=bias) + return self.conv2d_pointwise_1x1(x, weight, bias=bias, backend=backend) raise NotImplementedError( "v2 runtime does not expose generic conv2d for this configuration: " @@ -209,91 +212,91 @@ def _pair(value: Any) -> tuple[int, int]: f"dilation={dilation_hw} groups={groups_int}" ) - def scalar_not_equal(self, x, value): + def scalar_not_equal(self, x, value, backend=None): if _has_symbol("cactus_graph_scalar_not_equal"): - return orig_scalar_not_equal(self, x, value) + return orig_scalar_not_equal(self, x, value, backend=backend) x = _ensure_compare_tensor(self, x) delta = self.scalar_add(x, -float(value)) return _approx_nonzero_mask(self, delta) - def not_equal(self, a, b): + def not_equal(self, a, b, backend=None): if _has_symbol("cactus_graph_not_equal"): - return orig_not_equal(self, a, b) + return orig_not_equal(self, a, b, backend=backend) a = _ensure_compare_tensor(self, a) b = _ensure_compare_tensor(self, b) delta = self.subtract(a, b) return _approx_nonzero_mask(self, delta) - def transpose(self, x, backend=Graph.CPU): + def transpose(self, x, backend=None): return orig_transpose(self, _ensure_scalar_tensor(self, x), backend=backend) - def permute(self, x, permutation, backend=Graph.CPU): + def permute(self, x, permutation, backend=None): return orig_permute(self, _ensure_scalar_tensor(self, x), permutation, backend=backend) - def scalar_add(self, x, value): - return orig_scalar_add(self, self._ensure_tensor(x), value) + def scalar_add(self, x, value, backend=None): + return orig_scalar_add(self, self._ensure_tensor(x), value, backend=backend) - def scalar_subtract(self, x, value): - return orig_scalar_subtract(self, self._ensure_tensor(x), value) + def scalar_subtract(self, x, value, backend=None): + return orig_scalar_subtract(self, self._ensure_tensor(x), value, backend=backend) - def scalar_multiply(self, x, value): - return orig_scalar_multiply(self, self._ensure_tensor(x), value) + def scalar_multiply(self, x, value, backend=None): + return orig_scalar_multiply(self, self._ensure_tensor(x), value, backend=backend) - def scalar_divide(self, x, value): - return orig_scalar_divide(self, self._ensure_tensor(x), value) + def scalar_divide(self, x, value, backend=None): + return orig_scalar_divide(self, self._ensure_tensor(x), value, backend=backend) - def scalar_exp(self, x): - return orig_scalar_exp(self, self._ensure_tensor(x)) + def scalar_exp(self, x, backend=None): + return orig_scalar_exp(self, self._ensure_tensor(x), backend=backend) - def scalar_sqrt(self, x): - return orig_scalar_sqrt(self, self._ensure_tensor(x)) + def scalar_sqrt(self, x, backend=None): + return orig_scalar_sqrt(self, self._ensure_tensor(x), backend=backend) - def scalar_cos(self, x): - return orig_scalar_cos(self, self._ensure_tensor(x)) + def scalar_cos(self, x, backend=None): + return orig_scalar_cos(self, self._ensure_tensor(x), backend=backend) - def scalar_sin(self, x): - return orig_scalar_sin(self, self._ensure_tensor(x)) + def scalar_sin(self, x, backend=None): + return orig_scalar_sin(self, self._ensure_tensor(x), backend=backend) - def scalar_log(self, x): - return orig_scalar_log(self, self._ensure_tensor(x)) + def scalar_log(self, x, backend=None): + return orig_scalar_log(self, self._ensure_tensor(x), backend=backend) - def clamp(self, x, lo, hi): + def clamp(self, x, lo, hi, backend=None): if orig_clamp is None: raise RuntimeError("Cactus runtime is missing required symbol: cactus_graph_clamp") - return orig_clamp(self, _ensure_scalar_tensor(self, x), lo, hi) + return orig_clamp(self, _ensure_scalar_tensor(self, x), lo, hi, backend=backend) - def add(self, a, b): - return orig_add(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b)) + def add(self, a, b, backend=None): + return orig_add(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b), backend=backend) - def subtract(self, a, b): - return orig_subtract(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b)) + def subtract(self, a, b, backend=None): + return orig_subtract(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b), backend=backend) - def multiply(self, a, b): - return orig_multiply(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b)) + def multiply(self, a, b, backend=None): + return orig_multiply(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b), backend=backend) - def divide(self, a, b): - return orig_divide(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b)) + def divide(self, a, b, backend=None): + return orig_divide(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b), backend=backend) - def concat(self, a, b, axis=0): - return orig_concat(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b), axis=axis) + def concat(self, a, b, axis=0, backend=None): + return orig_concat(self, _ensure_scalar_tensor(self, a), _ensure_scalar_tensor(self, b), axis=axis, backend=backend) - def cat(self, tensors, axis=0): + def cat(self, tensors, axis=0, backend=None): legalized = [_ensure_scalar_tensor(self, tensor) for tensor in tensors] - return orig_cat(self, legalized, axis=axis) + return orig_cat(self, legalized, axis=axis, backend=backend) - def abs(self, x): - return orig_abs(self, self._ensure_tensor(x)) + def abs(self, x, backend=None): + return orig_abs(self, self._ensure_tensor(x), backend=backend) - def pow(self, x, exponent): - return orig_pow(self, self._ensure_tensor(x), exponent) + def pow(self, x, exponent, backend=None): + return orig_pow(self, self._ensure_tensor(x), exponent, backend=backend) - def expand(self, x, shape): + def expand(self, x, shape, backend=None): x = self._ensure_tensor(x) shape = tuple(int(v) for v in shape) if tuple(int(dim) for dim in x.shape) == shape: return x if _has_symbol("cactus_graph_expand"): - return orig_expand(self, x, shape) + return orig_expand(self, x, shape, backend=backend) # v2 builds do not always expose the generic expand symbol. For # broadcast-only expands, adding an embedded zero tensor with the target @@ -312,7 +315,7 @@ def expand(self, x, shape): zero = self.input(shape, dtype=x_dtype) self.set_input(zero, np.zeros(shape, dtype=dtype_map[x_dtype]), dtype=x_dtype) self.mark_embedded_input(zero) - return self.add(x, zero) + return self.add(x, zero, backend=backend) Graph.matmul = matmul Graph.gather = gather From 9c7d4a6d9001ef06c9dfd83571178a7c183ffced Mon Sep 17 00:00:00 2001 From: jakmro Date: Thu, 9 Jul 2026 03:56:40 +0200 Subject: [PATCH 4/4] fixes --- README.md | 10 ++--- cactus-graph/cactus_graph.h | 13 +++---- cactus-graph/src/graph_ffi.cpp | 32 +++++++--------- python/cactus/bindings/cactus.py | 46 +++++++---------------- python/cactus/transpile/runtime_compat.py | 6 +-- 5 files changed, 40 insertions(+), 67 deletions(-) diff --git a/README.md b/README.md index c064037a8..38deee9cc 100644 --- a/README.md +++ b/README.md @@ -208,7 +208,7 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ │ │ cactus run [model|path] run a model (downloads if needed) │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --backend cpu|metal inference backend (default: auto) │ +│ --backend cpu|metal inference backend (default: auto) │ │ --image image file for VLM inference │ │ --audio audio file for audio chat │ │ --system system prompt │ @@ -242,7 +242,7 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ --host bind address (default: 127.0.0.1) │ │ --port port (default: 8080) │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --backend cpu|metal inference backend (default: auto) │ +│ --backend cpu|metal inference backend (default: auto) │ │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild from source │ │ --no-cloud-handoff disable automatic cloud handoff │ @@ -252,7 +252,7 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ cactus code run the AI coding agent (TUI / print) │ │ --serve-model auto-start a server with this model │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --backend cpu|metal inference backend (default: auto) │ +│ --backend cpu|metal inference backend (default: auto) │ │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild from source │ │ --host server address (default: 127.0.0.1) │ @@ -276,7 +276,7 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ --model default: google/gemma-4-E2B-it │ │ --transcription-model default: nvidia/parakeet-tdt-0.6b-v3 │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --backend cpu|metal inference backend (default: auto) │ +│ --backend cpu|metal inference backend (default: auto) │ │ --token HuggingFace token (gated models) │ │ --reconvert force local rebuild of test models │ │ --suite run a single test suite by name │ @@ -291,7 +291,7 @@ Command: `cactus benchmark` [optional `--ios` or `--android`] │ --model default: google/gemma-4-E2B-it │ │ --transcription-model default: nvidia/parakeet-tdt-0.6b-v3 │ │ --bits 1|2|3|4 CQ quantization (default: 4) │ -│ --backend cpu|metal inference backend (default: auto) │ +│ --backend cpu|metal inference backend (default: auto) │ │ --ios run on connected iPhone │ │ --android run on connected Android │ │ │ diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index d43d11532..279d1be02 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -919,7 +919,6 @@ CACTUS_FFI_EXPORT cactus_graph_t cactus_graph_create(void); CACTUS_FFI_EXPORT void cactus_graph_destroy(cactus_graph_t graph); CACTUS_FFI_EXPORT int cactus_graph_hard_reset(cactus_graph_t graph); -CACTUS_FFI_EXPORT int32_t cactus_graph_default_backend(void); CACTUS_FFI_EXPORT int cactus_graph_set_node_backend( cactus_graph_t graph, cactus_node_t node, int32_t backend); @@ -983,9 +982,9 @@ 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_transpose( - cactus_graph_t graph, cactus_node_t x, int32_t backend, cactus_node_t* out); + 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, int32_t backend, cactus_node_t* out); + 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_index( @@ -1005,7 +1004,7 @@ CACTUS_FFI_EXPORT int cactus_graph_cat( cactus_graph_t graph, const cactus_node_t* nodes, size_t count, int32_t axis, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_matmul( - cactus_graph_t graph, cactus_node_t a, cactus_node_t b, bool pretransposed_rhs, int32_t backend, cactus_node_t* out); + 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_embedding_from_tensor( @@ -1046,12 +1045,12 @@ CACTUS_FFI_EXPORT int cactus_graph_topk(cactus_graph_t graph, cactus_node_t inpu CACTUS_FFI_EXPORT int cactus_graph_rms_norm( cactus_graph_t graph, cactus_node_t input, cactus_node_t weight, float epsilon, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_rope( - cactus_graph_t graph, cactus_node_t input, float theta, size_t position_offset, int32_t backend, cactus_node_t* out); + cactus_graph_t graph, cactus_node_t input, float theta, size_t position_offset, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_rope_gptj( - cactus_graph_t graph, cactus_node_t input, float theta, size_t position_offset, size_t rot_dim, int32_t backend, cactus_node_t* out); + cactus_graph_t graph, cactus_node_t input, float theta, size_t position_offset, size_t rot_dim, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_softmax(cactus_graph_t graph, cactus_node_t input, int32_t axis, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_attention( - cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value, float scale, bool is_causal, size_t position_offset, size_t window_size, int32_t backend, bool use_mask, cactus_node_t mask, bool additive_mask, cactus_node_t* out); + cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value, float scale, bool is_causal, size_t position_offset, size_t window_size, bool use_mask, cactus_node_t mask, bool additive_mask, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_rel_pos_bias( cactus_graph_t graph, cactus_node_t query, cactus_node_t relative_key, float scale, cactus_node_t* out); CACTUS_FFI_EXPORT int cactus_graph_attention_int8_hybrid( diff --git a/cactus-graph/src/graph_ffi.cpp b/cactus-graph/src/graph_ffi.cpp index 998b905a1..99cc069bb 100644 --- a/cactus-graph/src/graph_ffi.cpp +++ b/cactus-graph/src/graph_ffi.cpp @@ -437,10 +437,6 @@ int cactus_graph_pow(cactus_graph_t graph, cactus_node_t x, float exponent, cact } } -int32_t cactus_graph_default_backend(void) { - return static_cast(cactus_default_backend()); -} - static bool valid_backend(int32_t backend) { switch (static_cast(backend)) { case ComputeBackend::CPU: @@ -505,10 +501,10 @@ int cactus_graph_reshape(cactus_graph_t graph, cactus_node_t x, const size_t* sh } } -int cactus_graph_transpose(cactus_graph_t graph, cactus_node_t x, int32_t backend, cactus_node_t* out) { +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 { - *out = static_cast(as_graph(graph)->graph.transpose(static_cast(x), static_cast(backend))); + *out = static_cast(as_graph(graph)->graph.transpose(static_cast(x))); return 0; } catch (const std::exception& e) { last_error_message = e.what(); @@ -516,11 +512,11 @@ int cactus_graph_transpose(cactus_graph_t graph, cactus_node_t x, int32_t backen } } -int cactus_graph_transpose_n(cactus_graph_t graph, cactus_node_t x, const size_t* permutation, size_t rank, int32_t backend, cactus_node_t* out) { +int cactus_graph_transpose_n(cactus_graph_t graph, cactus_node_t x, const size_t* permutation, size_t rank, cactus_node_t* out) { if (!graph || !permutation || rank == 0 || !out) return fail_invalid("Invalid args to cactus_graph_transpose_n"); try { std::vector p(permutation, permutation + rank); - *out = static_cast(as_graph(graph)->graph.transposeN(static_cast(x), p, static_cast(backend))); + *out = static_cast(as_graph(graph)->graph.transposeN(static_cast(x), p)); return 0; } catch (const std::exception& e) { last_error_message = e.what(); @@ -649,10 +645,10 @@ int cactus_graph_cat(cactus_graph_t graph, const cactus_node_t* nodes, size_t co } } -int cactus_graph_matmul(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, bool pretransposed_rhs, int32_t backend, cactus_node_t* out) { +int cactus_graph_matmul(cactus_graph_t graph, cactus_node_t a, cactus_node_t b, bool pretransposed_rhs, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_matmul"); try { - *out = static_cast(as_graph(graph)->graph.matmul(static_cast(a), static_cast(b), pretransposed_rhs, static_cast(backend))); + *out = static_cast(as_graph(graph)->graph.matmul(static_cast(a), static_cast(b), pretransposed_rhs)); return 0; } catch (const std::exception& e) { last_error_message = e.what(); @@ -908,10 +904,10 @@ int cactus_graph_rms_norm(cactus_graph_t graph, cactus_node_t input, cactus_node } } -int cactus_graph_rope(cactus_graph_t graph, cactus_node_t input, float theta, size_t position_offset, int32_t backend, cactus_node_t* out) { +int cactus_graph_rope(cactus_graph_t graph, cactus_node_t input, float theta, size_t position_offset, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_rope"); try { - *out = static_cast(as_graph(graph)->graph.rope(static_cast(input), theta, position_offset, static_cast(backend))); + *out = static_cast(as_graph(graph)->graph.rope(static_cast(input), theta, position_offset)); return 0; } catch (const std::exception& e) { last_error_message = e.what(); @@ -919,10 +915,10 @@ int cactus_graph_rope(cactus_graph_t graph, cactus_node_t input, float theta, si } } -int cactus_graph_rope_gptj(cactus_graph_t graph, cactus_node_t input, float theta, size_t position_offset, size_t rot_dim, int32_t backend, cactus_node_t* out) { +int cactus_graph_rope_gptj(cactus_graph_t graph, cactus_node_t input, float theta, size_t position_offset, size_t rot_dim, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_rope_gptj"); try { - *out = static_cast(as_graph(graph)->graph.rope_gptj(static_cast(input), theta, position_offset, rot_dim, static_cast(backend))); + *out = static_cast(as_graph(graph)->graph.rope_gptj(static_cast(input), theta, position_offset, rot_dim)); return 0; } catch (const std::exception& e) { last_error_message = e.what(); @@ -941,15 +937,15 @@ int cactus_graph_softmax(cactus_graph_t graph, cactus_node_t input, int32_t axis } } -int cactus_graph_attention(cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value, float scale, bool is_causal, size_t position_offset, size_t window_size, int32_t backend, bool use_mask, cactus_node_t mask, bool additive_mask, cactus_node_t* out) { +int cactus_graph_attention(cactus_graph_t graph, cactus_node_t query, cactus_node_t key, cactus_node_t value, float scale, bool is_causal, size_t position_offset, size_t window_size, bool use_mask, cactus_node_t mask, bool additive_mask, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_attention"); try { if (use_mask) { - *out = static_cast(as_graph(graph)->graph.attention_masked(static_cast(query), static_cast(key), static_cast(value), static_cast(mask), scale, is_causal, static_cast(backend), additive_mask, position_offset, window_size)); + *out = static_cast(as_graph(graph)->graph.attention_masked(static_cast(query), static_cast(key), static_cast(value), static_cast(mask), scale, is_causal, cactus_default_backend(), additive_mask, position_offset, window_size)); } else if (window_size > 0 || position_offset > 0) { - *out = static_cast(as_graph(graph)->graph.attention(static_cast(query), static_cast(key), static_cast(value), scale, position_offset, window_size, static_cast(backend))); + *out = static_cast(as_graph(graph)->graph.attention(static_cast(query), static_cast(key), static_cast(value), scale, position_offset, window_size)); } else { - *out = static_cast(as_graph(graph)->graph.attention(static_cast(query), static_cast(key), static_cast(value), scale, is_causal, static_cast(backend))); + *out = static_cast(as_graph(graph)->graph.attention(static_cast(query), static_cast(key), static_cast(value), scale, is_causal)); } return 0; } catch (const std::exception& e) { diff --git a/python/cactus/bindings/cactus.py b/python/cactus/bindings/cactus.py index 7b6d60cb5..5901d6a4c 100644 --- a/python/cactus/bindings/cactus.py +++ b/python/cactus/bindings/cactus.py @@ -76,8 +76,6 @@ class cactus_tensor_info_t(ctypes.Structure): _lib.cactus_graph_destroy.restype = None _lib.cactus_graph_hard_reset.argtypes = [cactus_graph_t] _lib.cactus_graph_hard_reset.restype = ctypes.c_int -_lib.cactus_graph_default_backend.argtypes = [] -_lib.cactus_graph_default_backend.restype = ctypes.c_int32 _lib.cactus_graph_set_node_backend.argtypes = [cactus_graph_t, cactus_node_t, ctypes.c_int32] _lib.cactus_graph_set_node_backend.restype = ctypes.c_int @@ -243,11 +241,11 @@ class cactus_tensor_info_t(ctypes.Structure): ctypes.c_int, ) _lib.cactus_graph_transpose.argtypes = [ - cactus_graph_t, cactus_node_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) + cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) ] _lib.cactus_graph_transpose.restype = ctypes.c_int _lib.cactus_graph_transpose_n.argtypes = [ - cactus_graph_t, cactus_node_t, ctypes.POINTER(ctypes.c_size_t), ctypes.c_size_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) + cactus_graph_t, cactus_node_t, ctypes.POINTER(ctypes.c_size_t), ctypes.c_size_t, ctypes.POINTER(cactus_node_t) ] _lib.cactus_graph_transpose_n.restype = ctypes.c_int _lib.cactus_graph_slice.argtypes = [ @@ -270,7 +268,7 @@ class cactus_tensor_info_t(ctypes.Structure): ] _lib.cactus_graph_cat.restype = ctypes.c_int _lib.cactus_graph_matmul.argtypes = [ - cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, ctypes.c_int32, ctypes.POINTER(cactus_node_t) + cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_bool, ctypes.POINTER(cactus_node_t) ] _lib.cactus_graph_matmul.restype = ctypes.c_int _lib.cactus_graph_gather.argtypes = [ @@ -394,11 +392,11 @@ class cactus_tensor_info_t(ctypes.Structure): ] _lib.cactus_graph_topk.restype = ctypes.c_int _lib.cactus_graph_rope.argtypes = [ - cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.c_size_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) + cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) ] _lib.cactus_graph_rope.restype = ctypes.c_int _lib.cactus_graph_rope_gptj.argtypes = [ - cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.c_size_t, ctypes.c_size_t, ctypes.c_int32, ctypes.POINTER(cactus_node_t) + cactus_graph_t, cactus_node_t, ctypes.c_float, ctypes.c_size_t, ctypes.c_size_t, ctypes.POINTER(cactus_node_t) ] _lib.cactus_graph_rope_gptj.restype = ctypes.c_int _lib.cactus_graph_softmax.argtypes = [ @@ -407,7 +405,7 @@ class cactus_tensor_info_t(ctypes.Structure): _lib.cactus_graph_softmax.restype = ctypes.c_int _lib.cactus_graph_attention.argtypes = [ cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.c_float, ctypes.c_bool, - ctypes.c_size_t, ctypes.c_size_t, ctypes.c_int32, ctypes.c_bool, cactus_node_t, ctypes.c_bool, + ctypes.c_size_t, ctypes.c_size_t, ctypes.c_bool, cactus_node_t, ctypes.c_bool, ctypes.POINTER(cactus_node_t) ] _lib.cactus_graph_attention.restype = ctypes.c_int @@ -1766,23 +1764,18 @@ def index(self, x, index_value, axis=0, backend=None): return self._apply_backend(self._tensor_from_node(out.value), backend) def transpose(self, x, backend=None): - if backend is None: - backend = _lib.cactus_graph_default_backend() x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_transpose( self.h, cactus_node_t(x.id), - ctypes.c_int32(int(backend)), ctypes.byref(out), ) if rc != 0: raise RuntimeError(_err("graph_transpose failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def permute(self, x, permutation, backend=None): - if backend is None: - backend = _lib.cactus_graph_default_backend() x = self._ensure_tensor(x) permutation = tuple(int(v) for v in permutation) arr = (ctypes.c_size_t * len(permutation))(*permutation) @@ -1792,16 +1785,13 @@ def permute(self, x, permutation, backend=None): cactus_node_t(x.id), arr, len(permutation), - ctypes.c_int32(int(backend)), ctypes.byref(out), ) if rc != 0: raise RuntimeError(_err("graph_transpose_n failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def matmul(self, a, b, pretransposed_rhs=False, backend=None, output_dtype=None): - if backend is None: - backend = _lib.cactus_graph_default_backend() a = self._ensure_tensor(a) b = self._ensure_tensor(b) out = cactus_node_t() @@ -1810,12 +1800,11 @@ def matmul(self, a, b, pretransposed_rhs=False, backend=None, output_dtype=None) cactus_node_t(a.id), cactus_node_t(b.id), ctypes.c_bool(bool(pretransposed_rhs)), - ctypes.c_int32(int(backend)), ctypes.byref(out), ) if rc != 0: raise RuntimeError(_err("graph_matmul failed")) - result = self._tensor_from_node(out.value) + result = self._apply_backend(self._tensor_from_node(out.value), backend) if output_dtype is not None and int(output_dtype) != int(result.dtype): result = self.precision_cast(result, int(output_dtype)) return result @@ -2047,30 +2036,26 @@ def topk(self, x, k, backend=None): return self._apply_backend(self._tensor_from_node(out.value), backend) def rope(self, x, theta, position_offset=0, backend=None): - if backend is None: - backend = _lib.cactus_graph_default_backend() x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_rope( self.h, cactus_node_t(x.id), ctypes.c_float(float(theta)), ctypes.c_size_t(int(position_offset)), - ctypes.c_int32(int(backend)), ctypes.byref(out) + ctypes.byref(out) ) if rc != 0: raise RuntimeError(_err("graph_rope failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def rope_gptj(self, x, theta, position_offset=0, rot_dim=0, backend=None): - if backend is None: - backend = _lib.cactus_graph_default_backend() x = self._ensure_tensor(x) out = cactus_node_t() rc = _lib.cactus_graph_rope_gptj( self.h, cactus_node_t(x.id), ctypes.c_float(float(theta)), ctypes.c_size_t(int(position_offset)), - ctypes.c_size_t(int(rot_dim)), ctypes.c_int32(int(backend)), ctypes.byref(out) + ctypes.c_size_t(int(rot_dim)), ctypes.byref(out) ) if rc != 0: raise RuntimeError(_err("graph_rope_gptj failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def _reduce(self, fn_name, x, axis, backend=None): x = self._ensure_tensor(x) @@ -2113,8 +2098,6 @@ def softmax(self, x, axis=-1, backend=None): def attention(self, query, key, value, scale, is_causal=True, position_offset=0, window_size=0, backend=None, mask=None, additive_mask=False): - if backend is None: - backend = _lib.cactus_graph_default_backend() query = self._ensure_tensor(query) key = self._ensure_tensor(key) value = self._ensure_tensor(value) @@ -2132,7 +2115,6 @@ def attention(self, query, key, value, scale, is_causal=True, position_offset=0, ctypes.c_bool(bool(is_causal)), ctypes.c_size_t(int(position_offset)), ctypes.c_size_t(int(window_size)), - ctypes.c_int32(int(backend)), ctypes.c_bool(use_mask), mask_node, ctypes.c_bool(bool(additive_mask)), @@ -2140,7 +2122,7 @@ def attention(self, query, key, value, scale, is_causal=True, position_offset=0, ) if rc != 0: raise RuntimeError(_err("graph_attention failed")) - return self._tensor_from_node(out.value) + return self._apply_backend(self._tensor_from_node(out.value), backend) def kv_cache_state(self, max_seq_len, num_kv_heads, head_dim, window_size=0, sink_size=0, num_slots=1, backend=None): out = cactus_node_t() diff --git a/python/cactus/transpile/runtime_compat.py b/python/cactus/transpile/runtime_compat.py index 85acfdd0b..1fd53f0b4 100644 --- a/python/cactus/transpile/runtime_compat.py +++ b/python/cactus/transpile/runtime_compat.py @@ -111,8 +111,6 @@ def _approx_nonzero_mask(self, tensor): orig_expand = Graph.expand def matmul(self, a, b, pretransposed_rhs=False, backend=None, output_dtype=None): - if backend is None: - backend = _lib.cactus_graph_default_backend() a = _ensure_fp16_activation(self, a) b = _ensure_fp16_activation(self, b) out = cactus_node_t() @@ -121,12 +119,11 @@ def matmul(self, a, b, pretransposed_rhs=False, backend=None, output_dtype=None) cactus_node_t(a.id), cactus_node_t(b.id), ctypes.c_bool(bool(pretransposed_rhs)), - ctypes.c_int32(int(backend)), ctypes.byref(out), ) if rc != 0: raise RuntimeError(_err("graph_matmul failed")) - result = self._tensor_from_node(out.value) + result = self._apply_backend(self._tensor_from_node(out.value), backend) if output_dtype is not None and int(output_dtype) != int(result.dtype): result = self.precision_cast(result, int(output_dtype)) return result @@ -356,7 +353,6 @@ def expand(self, x, shape, backend=None): _cactus_module.cactus_node_t, _cactus_module.cactus_node_t, ctypes.c_bool, - ctypes.c_int32, ctypes.POINTER(_cactus_module.cactus_node_t), ] _lib_obj.cactus_graph_matmul.restype = ctypes.c_int