Bonsai masked gemv

cactus-engine/src/complete.cpp

@@ -16,6 +16,7 @@
 #include <iostream>
 #include <memory>
 #include <sstream>
+#include <cctype>
 #include <unordered_set>
 #include <vector>
 
@@ -506,11 +507,6 @@ bool prompt_context_matches(
     if (handle->processed_tokens.empty()) {
         return false;
     }
-    if (handle->model &&
-        handle->model->get_config().model_type == Config::ModelType::NEEDLE &&
-        prompt.tokens.size() != handle->processed_tokens.size() + 1) {
-        return false;
-    }
     if (prompt.context_token_count < handle->processed_tokens.size()) {
         return false;
     }

cactus-engine/src/engine.h

@@ -762,6 +762,7 @@ class Model {
     void unload_component_graph(Component& comp);
     bool bind_runtime_buffers(Component& comp);
     void run_step(uint32_t token_id, size_t position, bool read_logits, bool use_fused = true);
+    void trim_prefill_components();
     void run_step_batch(const std::vector<uint32_t>& token_ids, const std::vector<size_t>& positions);
     void set_component_batch(Component& comp, size_t batch);
     size_t decoder_cache_num_slots();
@@ -796,7 +797,7 @@ class Model {
                                              size_t chunk_size, bool prepare_decode);
     void execute_prefill_chunk(Component& chunk_comp, Component* enc_comp, size_t encoder_chunk,
                                size_t chunk_tokens, const std::vector<uint32_t>& tokens,
-                               size_t processed, size_t start_position);
+                               size_t processed, size_t start_position, uint32_t prefill_valid_len = 0);
     void run_full_context_text();
     void prepare_sampling_context(float repetition_penalty);
     uint32_t sample_component_logits(Component& comp, float temperature, float top_p, size_t top_k,

cactus-engine/src/model.cpp

@@ -1,3 +1,4 @@
+#include "metal_backend.h"
 #include "engine.h"
 #include "cactus_graph.h"
 #include "cactus_kernels.h"
@@ -354,6 +355,9 @@ void Model::maybe_capture_handoff_probe_hidden(const Component& comp, const std:
         static_cast<size_t>(desc.shape.back()) != static_cast<size_t>(handoff_probe_feat_dim_)) return;
     size_t rows = desc.total_size / handoff_probe_feat_dim_;
     if (rows == 0) return;
+    if (cactus_default_backend() == ComputeBackend::METAL && cactus_metal_available()) {
+        cactus_metal_session_sync();
+    }
     const auto* data = static_cast<const uint8_t*>(comp.graph->get_output(node));
     if (!data) return;
     for (size_t row = 0; row < rows; ++row) {
@@ -1049,7 +1053,11 @@ std::vector<std::vector<uint32_t>> Model::decode_batch(const std::vector<uint32_
 bool Model::supports_dynamic_batch() {
     if (!decoder_) return false;
     if (!decoder_->graph && !load_component_graph(*decoder_)) return false;
-    return decoder_->graph->has_dynamic_shapes();
+    if (!decoder_->graph->has_dynamic_shapes()) return false;
+    for (const auto& np : decoder_->graph->nodes_) {
+        if (np->op_type == OpType::CONV_CACHE_STATE || np->op_type == OpType::RECURRENT_CACHE_STATE) return false;
+    }
+    return true;
 }
 
 void Model::set_decode_slots(size_t num_slots) {
@@ -1254,7 +1262,7 @@ size_t Model::component_output_tokens(const Component& comp, const std::string&
 
 void Model::execute_prefill_chunk(Component& chunk_comp, Component* enc_comp, size_t encoder_chunk,
                                   size_t chunk_tokens, const std::vector<uint32_t>& tokens,
-                                  size_t processed, size_t start_position) {
+                                  size_t processed, size_t start_position, uint32_t prefill_valid_len) {
     for (size_t i = 0; i < chunk_comp.input_buffers.size(); ++i) {
         std::fill(chunk_comp.input_buffers[i].begin(), chunk_comp.input_buffers[i].end(), 0);
     }
@@ -1280,7 +1288,7 @@ void Model::execute_prefill_chunk(Component& chunk_comp, Component* enc_comp, si
             copy_component_outputs_to_chunk_inputs(*encoder_, chunk_comp, i);
         }
     }
-    chunk_comp.graph->execute();
+    chunk_comp.graph->execute("", prefill_valid_len);
 }
 
 void Model::reset_prefill_stats() {
@@ -1317,7 +1325,11 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector<uint32_
     const bool has_recurrent_state = any_cache_node([&](size_t id) {
         return decoder_prefill_->graph->get_node_op_type(id) == OpType::RECURRENT_CACHE_STATE;
     });
-    if (has_recurrent_state && whole_chunks_end > effective_chunk) {
+    static const bool multi_chunk_recurrent = []{
+        const char* v = std::getenv("CACTUS_PREFILL_MULTI_CHUNK");
+        return !v || v[0] != '0';
+    }();
+    if (has_recurrent_state && !multi_chunk_recurrent && whole_chunks_end > effective_chunk) {
         whole_chunks_end = effective_chunk;
     }
     const bool has_sliding_window_cache = any_cache_node([&](size_t id) {
@@ -1342,8 +1354,11 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector<uint32_
     const bool use_padded_tail = !pad_tail && !prefill_tail_pad_disabled_
         && has_sliding_window_cache && !has_recurrent_state && !has_conv_state
         && tail_tokens > 8 && !padded_window_too_small;
+    const bool metal_prefill = cactus_default_backend() == ComputeBackend::METAL && cactus_metal_available();
+    const bool use_recurrent_tail = !pad_tail && !use_padded_tail && !prefill_tail_pad_disabled_
+        && (has_recurrent_state || has_conv_state) && tail_tokens > 8 && metal_prefill;
     const size_t executable_tokens = whole_chunks_end + (pad_tail ? effective_chunk : 0);
-    if (executable_tokens == 0 && !use_padded_tail) {
+    if (executable_tokens == 0 && !use_padded_tail && !use_recurrent_tail) {
         result.scalar_tail_tokens = tail_tokens;
         last_prefill_scalar_tail_tokens_ = tail_tokens;
         return result;
@@ -1366,7 +1381,15 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector<uint32_
 
     size_t tail_executed = 0;
     size_t tail_padding = 0;
-    if (use_padded_tail) {
+    if (use_recurrent_tail) {
+        const size_t chunk_real = tail_tokens - 1;
+        execute_prefill_chunk(*decoder_prefill_, prefill_encoder_, encoder_chunk,
+                              effective_chunk, tokens, processed, start_position,
+                              static_cast<uint32_t>(chunk_real));
+        processed += chunk_real;
+        tail_executed = chunk_real;
+        tail_padding = effective_chunk - tail_tokens;
+    } else if (use_padded_tail) {
         const size_t pads = effective_chunk - tail_tokens;
         const size_t kept_real = tail_tokens - 1;
         std::vector<std::pair<size_t, std::vector<uint8_t>>> backups;
@@ -2204,12 +2227,14 @@ bool Model::prefill_and_sample_first_token(const std::vector<uint32_t>& tokens,
             record_sampled_token(out_token);
             last_prefill_scalar_tail_tokens_ = 1;
             maybe_roll_compact();
+            trim_prefill_components();
             return true;
         }
         cache_total_seq_len_ += chunked.logical_tokens;
     }
+    const bool tail_fused = decode_route_ == DecodeRoute::DIRECT_DECODER_STEP;
     for (size_t i = chunked.logical_tokens; i < tokens.size(); ++i) {
-        run_step(tokens[i], cache_total_seq_len_, i + 1 == tokens.size(), /*use_fused=*/false);
+        run_step(tokens[i], cache_total_seq_len_, i + 1 == tokens.size(), /*use_fused=*/tail_fused);
         ++cache_total_seq_len_;
     }
     last_prefill_scalar_tail_tokens_ = tokens.size() - chunked.logical_tokens;
@@ -2220,6 +2245,7 @@ bool Model::prefill_and_sample_first_token(const std::vector<uint32_t>& tokens,
     }
     record_sampled_token(out_token);
     maybe_roll_compact();
+    trim_prefill_components();
     return true;
 }
 
@@ -2239,8 +2265,9 @@ void Model::prefill(const std::vector<uint32_t>& tokens, size_t /*chunk_size*/,
     }
     ChunkedPrefillResult chunked = run_chunked_prefill(tokens, cache_total_seq_len_, get_prefill_chunk_size(), prepare_decode);
     cache_total_seq_len_ += chunked.logical_tokens;
+    const bool prefill_fused = decode_route_ == DecodeRoute::DIRECT_DECODER_STEP;
     for (size_t i = chunked.logical_tokens; i < tokens.size(); ++i) {
-        run_step(tokens[i], cache_total_seq_len_, /*read_logits=*/false, /*use_fused=*/false);
+        run_step(tokens[i], cache_total_seq_len_, /*read_logits=*/false, /*use_fused=*/prefill_fused);
         ++cache_total_seq_len_;
     }
     cache_token_ids_.insert(cache_token_ids_.end(), tokens.begin(), tokens.end());
@@ -2250,6 +2277,16 @@ void Model::prefill(const std::vector<uint32_t>& tokens, size_t /*chunk_size*/,
         // After the prompt reaches full length here -- never mid-chunk -- bound it to target_len.
         maybe_roll_compact();
     }
+    trim_prefill_components();
+}
+
+void Model::trim_prefill_components() {
+    for (Component* c : {decoder_prefill_, prefill_encoder_}) {
+        if (!c || !c->graph || c == decoder_ || c == encoder_) continue;
+        c->graph->invalidate_metal_state();
+        c->graph->release_runtime_buffers();
+    }
+    ::cactus_metal_trim_prefill_cache();
 }
 
 void Model::prefill_with_images(const std::vector<uint32_t>& tokens,

cactus-engine/src/tokenizer.cpp

@@ -13,23 +13,50 @@ namespace engine {
 
 namespace {
 
-std::string format_needle_query_text(const std::vector<ChatMessage>& messages) {
+std::string format_needle_query_text(const std::vector<ChatMessage>& messages, bool add_generation_prompt) {
     std::string system_text;
     std::string user_query;
+    size_t user_count = 0;
+    bool has_non_legacy_turn = false;
 
     for (const auto& msg : messages) {
         if (msg.role == "system" || msg.role == "developer") {
             if (!system_text.empty()) system_text += "\n";
             system_text += msg.content;
         } else if (msg.role == "user") {
             user_query = msg.content;
+            user_count++;
+        } else {
+            has_non_legacy_turn = true;
         }
     }
 
-    if (user_query.empty() && !messages.empty()) user_query = messages.back().content;
-    if (system_text.empty()) return user_query;
-    if (user_query.empty()) return system_text;
-    return system_text + "\n\n" + user_query;
+    if (!has_non_legacy_turn && user_count <= 1) {
+        if (user_query.empty() && !messages.empty()) user_query = messages.back().content;
+        if (system_text.empty()) return user_query;
+        if (user_query.empty()) return system_text;
+        return system_text + "\n\n" + user_query;
+    }
+
+    std::string transcript = system_text;
+    for (const auto& msg : messages) {
+        if (msg.role == "system" || msg.role == "developer") continue;
+        if (!transcript.empty()) transcript += "\n\n";
+        if (msg.role == "assistant") {
+            transcript += "Assistant: " + msg.content;
+        } else if (msg.role == "tool") {
+            transcript += "Tool";
+            if (!msg.name.empty()) transcript += " " + msg.name;
+            transcript += ": " + msg.content;
+        } else {
+            transcript += "User: " + msg.content;
+        }
+    }
+    if (add_generation_prompt) {
+        if (!transcript.empty()) transcript += "\n\n";
+        transcript += "Assistant:";
+    }
+    return transcript;
 }
 
 std::string format_tool_call_for_prompt(const std::string& name, const std::string& arguments, bool gemma4) {
@@ -516,6 +543,12 @@ std::string lstrip_newlines(const std::string& s) {
     size_t a = s.find_first_not_of('\n');
     return a == std::string::npos ? "" : s.substr(a);
 }
+std::string trim_whitespace(const std::string& s) {
+    size_t a = s.find_first_not_of(" \t\n\r\f\v");
+    if (a == std::string::npos) return "";
+    size_t b = s.find_last_not_of(" \t\n\r\f\v");
+    return s.substr(a, b - a + 1);
+}
 }  // namespace
 
 std::string Tokenizer::format_qwen_style(const std::vector<ChatMessage>& messages, bool add_generation_prompt,
@@ -538,8 +571,18 @@ std::string Tokenizer::format_qwen_style(const std::vector<ChatMessage>& message
     }
 
     long last_query_index = static_cast<long>(n) - 1;
-    for (long i = static_cast<long>(n) - 1; i >= 0; --i) {
-        if (messages[i].role == "user") { last_query_index = i; break; }
+    for (long q = static_cast<long>(n) - 1; q >= 0; --q) {
+        const auto& msg = messages[q];
+        if (msg.role == "system" || msg.role == "developer" || msg.role == "assistant" || msg.role == "tool") continue;
+        const std::string trimmed = trim_whitespace(msg.content);
+        const std::string open_tag = "<tool_response>";
+        const std::string close_tag = "</tool_response>";
+        const bool is_tool_response = trimmed.size() >= open_tag.size() + close_tag.size()
+            && trimmed.compare(0, open_tag.size(), open_tag) == 0
+            && trimmed.compare(trimmed.size() - close_tag.size(), close_tag.size(), close_tag) == 0;
+        if (is_tool_response) continue;
+        last_query_index = q;
+        break;
     }
 
     for (size_t i = first; i < n; i++) {
@@ -571,7 +614,7 @@ std::string Tokenizer::format_qwen_style(const std::vector<ChatMessage>& message
                 content = lstrip_newlines(content.substr(tpos + 8));
             }
             result += "<|im_start|>assistant\n";
-            if (template_has_thinking && static_cast<long>(i) > last_query_index && (i == n - 1 || !reasoning.empty())) {
+            if (template_has_thinking && static_cast<long>(i) > last_query_index) {
                 result += "<think>\n" + reasoning + "\n</think>\n\n" + lstrip_newlines(content);
             } else {
                 result += content;
@@ -670,10 +713,10 @@ std::string Tokenizer::format_lfm2_style(const std::vector<ChatMessage>& message
     return result;
 }
 
-std::string Tokenizer::format_needle_style(const std::vector<ChatMessage>& messages, bool /*add_generation_prompt*/,
+std::string Tokenizer::format_needle_style(const std::vector<ChatMessage>& messages, bool add_generation_prompt,
                                            const std::string& tools_json) const {
     std::string serialized_tools = tools_json.empty() ? "[]" : tools_json;
-    return format_needle_query_text(messages) + "<tools>" + serialized_tools + "</s>";
+    return format_needle_query_text(messages, add_generation_prompt) + "<tools>" + serialized_tools + "</s>";
 }
 
 std::string Tokenizer::format_gemma4_style(const std::vector<ChatMessage>& messages, bool add_generation_prompt,

cactus-graph/cactus_graph.h

@@ -122,7 +122,8 @@ enum class OpType {
     SCALAR_NOT_EQUAL,
     RECURRENT_CACHE_STATE,
     RECURRENT_CACHE_WRITE,
-    CONV_CACHE_INITIALIZE
+    CONV_CACHE_INITIALIZE,
+    GATED_DELTANET_GATE_LOG
 };
 
 struct PrecisionTraits {
@@ -533,6 +534,8 @@ class CactusGraph {
     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 gated_deltanet_gate_log(size_t a_logits, size_t dt_bias, size_t a_log,
+                                   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());
@@ -741,14 +744,15 @@ class CactusGraph {
     void prefetch_weight_pages(size_t node_id);
     void release_all_weight_pages();
     void release_runtime_buffers();
+    void invalidate_metal_state();
     void clear_buffer_pool();
     void retain_outputs(const std::vector<int>& node_ids);
 
     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);
 
-    void execute(const std::string& profile_file = "");
+    void execute(const std::string& profile_file = "", uint32_t prefill_valid_len = 0);
     bool extract_ple_pathway(FusedEmbedCtx& ctx) const;
     void hard_reset();
     void soft_reset();
@@ -801,10 +805,11 @@ class CactusGraph {
     bool runtime_shapes_dirty_ = false;
     std::unordered_set<size_t> persistent_node_ids_;
     std::unordered_set<size_t> populated_node_ids_;
+    std::vector<uint8_t> node_is_constant_;
+    bool constants_cached_ = false;
     std::unordered_set<size_t> embedded_input_node_ids_;
     std::unordered_set<size_t> retained_output_node_ids_;
     void build_metal_retype_plan();
-    void invalidate_metal_state();
     std::unordered_map<uint64_t, struct MetalFusePlan*> metal_plans_;
     std::unordered_map<uint64_t, std::unordered_set<size_t>> metal_plan_banned_;
     uint64_t metal_plan_sig_ = 0;
@@ -862,6 +867,7 @@ namespace GraphFile {
         const void* scales_data() const;
         bool is_orthogonal_rotation() const { return is_orthogonal_rotation_; }
         bool is_interleaved_4row() const { return is_interleaved_4row_; }
+        bool is_no_rotation() const { return is_no_rotation_; }
         size_t original_N() const { return original_N_; }
         void* data();
         const void* data() const;
@@ -883,6 +889,7 @@ namespace GraphFile {
         uint32_t alignment_ = 32;
         bool is_orthogonal_rotation_ = false;
         bool is_interleaved_4row_ = false;
+        bool is_no_rotation_ = false;
         size_t original_N_ = 0;
 
         void parse_header();
@@ -1025,6 +1032,7 @@ CACTUS_FFI_EXPORT int cactus_graph_release_all_weight_pages(cactus_graph_t graph
 
 CACTUS_FFI_EXPORT int cactus_graph_relu(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out);
 CACTUS_FFI_EXPORT int cactus_graph_silu(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out);
+CACTUS_FFI_EXPORT int cactus_graph_gated_deltanet_gate_log(cactus_graph_t graph, cactus_node_t a_logits, cactus_node_t dt_bias, cactus_node_t a_log, cactus_node_t* out);
 CACTUS_FFI_EXPORT int cactus_graph_gelu(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out);
 CACTUS_FFI_EXPORT int cactus_graph_gelu_erf(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out);
 CACTUS_FFI_EXPORT int cactus_graph_sigmoid(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out);

cactus-graph/src/builder.cpp

@@ -1086,6 +1086,11 @@ size_t CactusGraph::silu(size_t input, ComputeBackend backend) {
     return tag_backend(add_node(OpType::SILU, {input}, {}), backend);
 }
 
+size_t CactusGraph::gated_deltanet_gate_log(size_t a_logits, size_t dt_bias, size_t a_log,
+                                            ComputeBackend backend) {
+    return tag_backend(add_node(OpType::GATED_DELTANET_GATE_LOG, {a_logits, dt_bias, a_log}, {}), backend);
+}
+
 size_t CactusGraph::gelu(size_t input, ComputeBackend backend) {
     return tag_backend(add_node(OpType::GELU, {input}, {}), backend);
 }

cactus-graph/src/core.cpp

@@ -325,6 +325,7 @@ void CactusGraph::release_runtime_buffers() {
         if (persistent_node_ids_.count(node->id)) continue;
         node->output_buffer.release_memory(buffer_pool_);
     }
+    constants_cached_ = false;
     buffer_pool_.clear();
     shrink_thread_local_buffers();
 }