cactus-react-native 1.2.1 → 1.5.0

package/ios/cactus.xcframework/ios-arm64/cactus.framework/Headers/graph.h

@@ -8,15 +8,99 @@
 #include <cstring>
 #include <stdexcept>
 #include <string>
+#include <mutex>
+#include <sstream>
+#include <iostream>
+#include <arm_neon.h>
+
+namespace cactus {
+
+enum class LogLevel {
+    DEBUG = 0,
+    INFO = 1,
+    WARN = 2,
+    ERROR = 3,
+    NONE = 4
+};
+
+class Logger {
+public:
+    static Logger& instance() {
+        static Logger logger;
+        return logger;
+    }
+
+    void set_level(LogLevel level) { min_level_ = level; }
+    LogLevel get_level() const { return min_level_; }
+
+    void set_callback(std::function<void(LogLevel, const std::string&, const std::string&)> cb) {
+        std::lock_guard<std::mutex> lock(mutex_);
+        callback_ = cb;
+    }
+
+    void log(LogLevel level, const std::string& component, const std::string& message) {
+        if (level < min_level_) return;
+
+        std::lock_guard<std::mutex> lock(mutex_);
+
+        if (callback_) {
+            callback_(level, component, message);
+        } else {
+            std::cerr << "[" << level_string(level) << "] [" << component << "] " << message << std::endl;
+        }
+
+        if (level == LogLevel::ERROR) {
+            last_error_ = "[" + component + "] " + message;
+        }
+    }
+
+    const std::string& last_error() const { return last_error_; }
+    void clear_error() { last_error_.clear(); }
+
+private:
+    Logger() : min_level_(LogLevel::WARN) {}
+
+    static const char* level_string(LogLevel level) {
+        switch (level) {
+            case LogLevel::DEBUG: return "DEBUG";
+            case LogLevel::INFO:  return "INFO";
+            case LogLevel::WARN:  return "WARN";
+            case LogLevel::ERROR: return "ERROR";
+            default: return "?";
+        }
+    }
+
+    LogLevel min_level_;
+    std::mutex mutex_;
+    std::string last_error_;
+    std::function<void(LogLevel, const std::string&, const std::string&)> callback_;
+};
+
+} // namespace cactus
+
+#define CACTUS_LOG(level, component, msg) \
+    do { \
+        if (static_cast<int>(level) >= static_cast<int>(cactus::Logger::instance().get_level())) { \
+            std::ostringstream _cactus_log_ss; \
+            _cactus_log_ss << msg; \
+            cactus::Logger::instance().log(level, component, _cactus_log_ss.str()); \
+        } \
+    } while(0)
+
+#define CACTUS_LOG_DEBUG(component, msg) CACTUS_LOG(cactus::LogLevel::DEBUG, component, msg)
+#define CACTUS_LOG_INFO(component, msg)  CACTUS_LOG(cactus::LogLevel::INFO, component, msg)
+#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)
 
 namespace GraphFile {
     class MappedFile;
 }
 
 enum class Precision {
-    INT8, 
+    INT8,
     FP16,
-    FP32
+    FP32,
+    INT4 
 };
 
 enum class ComputeBackend {
@@ -30,7 +114,7 @@ enum class OpType {
     MATMUL, TRANSPOSE, RESHAPE, SLICE, GATHER, EMBEDDING,
     BILINEAR_INTERPOLATION,
     SUM, MEAN, VARIANCE, MIN, MAX,
-    RMS_NORM, ROPE, SOFTMAX, ATTENTION, CONV1D_CAUSAL, CONV1D_K3,
+    RMS_NORM, ROPE, SOFTMAX, ATTENTION, ATTENTION_INT8_HYBRID, CONV1D_CAUSAL, CONV1D_K3,
     SCALAR_ADD, SCALAR_SUBTRACT, SCALAR_MULTIPLY, SCALAR_DIVIDE, SCALAR_EXP, SCALAR_SQRT, SCALAR_COS, SCALAR_SIN,
     SILU, GELU, GELU_ERF,
     SAMPLE, CONCAT,
@@ -40,27 +124,38 @@ enum class OpType {
 };
 
 struct PrecisionTraits {
+    // Returns in-memory element size (INT4 unpacks to INT8, so returns 1)
     static constexpr size_t size_of(Precision prec) {
         switch (prec) {
             case Precision::INT8: return 1;
             case Precision::FP16: return 2;
             case Precision::FP32: return 4;
+            case Precision::INT4: return 1; 
         }
         return 1;
     }
-    
+
+    static constexpr size_t packed_size_of(Precision prec, size_t count) {
+        switch (prec) {
+            case Precision::INT4: return (count + 1) / 2;  
+            default: return count * size_of(prec);
+        }
+    }
+
     static constexpr bool is_integer(Precision prec) {
         switch (prec) {
             case Precision::INT8: return true;
+            case Precision::INT4: return true;
             case Precision::FP16: return false;
             case Precision::FP32: return false;
         }
         return true;
     }
-    
+
     static constexpr bool is_floating_point(Precision prec) {
         switch (prec) {
             case Precision::INT8: return false;
+            case Precision::INT4: return false;
             case Precision::FP16: return true;
             case Precision::FP32: return true;
         }
@@ -71,8 +166,6 @@ struct PrecisionTraits {
 namespace Quantization {
     void int8_to_fp32(const int8_t* src, float* dst, size_t count, float scale = 1.0f);
     void fp32_to_int8(const float* src, int8_t* dst, size_t count, float scale = 1.0f);
-    void dynamic_quantize_fp32_to_int8(const float* src, int8_t* dst, size_t count, 
-                                       float* computed_scale);
     void fp16_to_fp32(const __fp16* src, float* dst, size_t count);
     void fp32_to_fp16(const float* src, __fp16* dst, size_t count);
     void int8_to_fp16(const int8_t* src, __fp16* dst, size_t count, float scale = 1.0f);
@@ -106,10 +199,17 @@ struct BufferDesc {
     void* external_data;
     char* pooled_data;
     Precision precision;
-    float quantization_scale;
+
+    size_t group_size = 0;
+    size_t num_groups = 0;
+    void* scales_data = nullptr;
+    std::unique_ptr<char[]> owned_scales;
+
+    const void* packed_int4_data = nullptr;  
+    size_t packed_int4_size = 0; 
 
     BufferDesc();
-    BufferDesc(const std::vector<size_t>& s, Precision prec = Precision::INT8, float scale = 1.0f);
+    BufferDesc(const std::vector<size_t>& s, Precision prec = Precision::INT8);
     ~BufferDesc();
 
     BufferDesc(BufferDesc&& other) noexcept;
@@ -127,6 +227,28 @@ struct BufferDesc {
     template<typename T>
     const T* data_as() const { return static_cast<const T*>(get_data()); }
 
+    const __fp16* scales_as_fp16() const {
+        return reinterpret_cast<const __fp16*>(scales_data);
+    }
+    bool is_grouped_int8() const {
+        return precision == Precision::INT8 && group_size > 0;
+    }
+    bool is_packed_int4() const {
+        return packed_int4_data != nullptr && packed_int4_size > 0;
+    }
+    const uint8_t* packed_int4_as_uint8() const {
+        return reinterpret_cast<const uint8_t*>(packed_int4_data);
+    }
+    void set_grouped_scales(size_t gs, size_t ng, void* scales_ptr) {
+        group_size = gs;
+        num_groups = ng;
+        scales_data = scales_ptr;
+    }
+    void set_packed_int4(const void* packed_data, size_t packed_size) {
+        packed_int4_data = packed_data;
+        packed_int4_size = packed_size;
+    }
+
     void allocate();
     void allocate_from_pool(BufferPool& pool);
     void release_to_pool(BufferPool& pool);
@@ -144,7 +266,7 @@ struct OpParams {
     size_t slice_start = 0;
     size_t slice_length = 0;
     size_t window_size = 0;
-    bool is_causal = true;  // Default to causal for backward compatibility
+    bool is_causal = true;  
     std::vector<size_t> new_shape;
     std::vector<size_t> permutation;
     Precision output_precision = Precision::INT8;
@@ -158,10 +280,21 @@ struct OpParams {
     size_t top_k = 0;
     size_t random_seed = 0;
     
-    size_t index_value = 0;  // For INDEX operation
-    size_t num_classes = 0;  // For scatter operations
+    size_t index_value = 0;  
+    size_t num_classes = 0; 
     size_t dst_height = 0;
-    size_t dst_width = 0;   
+    size_t dst_width = 0;
+
+    std::vector<float> bias_values;
+    std::vector<uint32_t> bias_indices;
+
+    const int8_t* cached_keys_int8 = nullptr;
+    const int8_t* cached_values_int8 = nullptr;
+    const float* cached_k_scales = nullptr;
+    const float* cached_v_scales = nullptr;
+    size_t cache_seq_len = 0;
+    size_t num_kv_heads = 0;
+    size_t head_dim = 0;
 };
 
 struct GraphNode {
@@ -241,7 +374,7 @@ public:
     size_t precision_cast(size_t input, Precision target_precision);
     
     size_t add(size_t input1, size_t input2);
-    size_t add_clipped(size_t input1, size_t input2);  // For FP16 residual connections (Gemma)
+    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);
@@ -276,8 +409,12 @@ public:
     size_t gather(size_t embeddings, size_t indices);
     size_t mmap_embeddings(const std::string& filename);
     size_t mmap_weights(const std::string& filename);
-    size_t load_weights(const std::string& filename); 
-    void set_quantization_scale(size_t node_id, float scale);
+    size_t load_weights(const std::string& filename);
+    void set_grouped_scales(size_t node_id, size_t group_size, size_t num_groups, void* scales_ptr);
+
+    void release_weight_pages(size_t node_id);
+    void prefetch_weight_pages(size_t node_id);
+    void release_all_weight_pages();
     size_t embedding(const std::string& filename, size_t indices);
     size_t embedding(size_t embedding_tensor, size_t indices);
     size_t bilinear_interpolation(size_t pos_embeds, size_t dst_height, size_t dst_width);
@@ -291,10 +428,16 @@ public:
     size_t attention(size_t query, size_t key, size_t value, float scale, size_t position_offset, ComputeBackend backend = ComputeBackend::CPU);
     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 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 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 sample(size_t logits, float temperature = 0.6f, float top_p = 0.95f, size_t top_k = 20);
+    size_t sample(size_t logits, float temperature = 0.6f, float top_p = 0.95f, size_t top_k = 20,
+                  const std::unordered_map<uint32_t, float>& logit_bias = {});
     
     size_t concat(size_t input1, size_t input2, int axis = 0);
     size_t scatter_topk(size_t indices, size_t values, size_t num_classes);
@@ -306,6 +449,8 @@ public:
     void execute(const std::string& profile_file = "");
     void hard_reset();
     void soft_reset();
+    void soft_reset_keep_pool();
+    void set_prefill_mode(bool enabled) { prefill_mode_ = enabled; }
 
     void register_debug_node(uint32_t layer_idx, const std::string& name, size_t node_id);
     void capture_debug_node(uint32_t layer_idx, const std::string& name, size_t node_id);
@@ -324,8 +469,10 @@ private:
     size_t next_node_id_;
     std::vector<std::unique_ptr<GraphFile::MappedFile>> mapped_files_;
     std::unordered_map<std::string, size_t> weight_cache_;
+    std::unordered_map<size_t, size_t> node_to_mapped_file_;
     std::vector<DebugNodeEntry> debug_nodes_;
     BufferPool buffer_pool_;
+    bool prefill_mode_ = false;
 };
 
 
@@ -344,25 +491,36 @@ namespace GraphFile {
     public:
         MappedFile(const std::string& filename);
         ~MappedFile();
-        
+
         MappedFile(const MappedFile&) = delete;
         MappedFile& operator=(const MappedFile&) = delete;
         MappedFile(MappedFile&& other) noexcept;
         MappedFile& operator=(MappedFile&& other) noexcept;
-        
+
         const std::vector<size_t>& shape() const;
         Precision precision() const;
+        Precision effective_precision() const {
+            return is_int4_ ? Precision::INT8 : precision_;
+        }
         size_t byte_size() const;
-        float quantization_scale() const;
-        
+
+        size_t group_size() const { return group_size_; }
+        size_t num_groups() const { return num_groups_; }
+        const void* scales_data() const;
+        const void* raw_packed_data() const;  // Get raw mmap'd data without unpacking (for INT4)
+        bool is_int4() const { return is_int4_; }
+
         void* data();
         const void* data() const;
-        
+
         template<typename T>
         const T* typed_data() const;
-        
+
         LoadedNode load_into_graph(CactusGraph& graph) const;
-        
+
+        void release_pages();
+        void prefetch_pages();
+
     private:
         int fd_;
         void* mapped_data_;
@@ -370,10 +528,19 @@ namespace GraphFile {
         std::vector<size_t> shape_;
         Precision precision_;
         size_t byte_size_;
-        float quantization_scale_;
+        size_t group_size_ = 0;
+        size_t num_groups_ = 0;
+        size_t scales_offset_ = 0;
+        size_t scales_bytes_ = 0;
+        uint32_t version_ = 1;
+        uint32_t alignment_ = 32;
+        bool is_int4_ = false;
+        mutable std::unique_ptr<int8_t[]> unpacked_int4_data_;
         void parse_header();
+        void apply_madvise_hints();
+        void unpack_int4_if_needed() const;
     };
-    
+
     MappedFile mmap_load(const std::string& filename);
 }