diff --git a/cactus-engine/src/complete.cpp b/cactus-engine/src/complete.cpp index e8fe2e33d..14be4b671 100644 --- a/cactus-engine/src/complete.cpp +++ b/cactus-engine/src/complete.cpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include @@ -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; } diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 73fc6d47a..c7f7cc216 100644 --- a/cactus-engine/src/engine.h +++ b/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& token_ids, const std::vector& 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& 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, diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 15788bb43..d69357583 100644 --- a/cactus-engine/src/model.cpp +++ b/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(desc.shape.back()) != static_cast(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(comp.graph->get_output(node)); if (!data) return; for (size_t row = 0; row < rows; ++row) { @@ -1049,7 +1053,11 @@ std::vector> Model::decode_batch(const std::vectorgraph && !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& 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::vectorgraph->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 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(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>> backups; @@ -2204,12 +2227,14 @@ bool Model::prefill_and_sample_first_token(const std::vector& 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& tokens, } record_sampled_token(out_token); maybe_roll_compact(); + trim_prefill_components(); return true; } @@ -2239,8 +2265,9 @@ void Model::prefill(const std::vector& tokens, size_t /*chunk_size*/, } ChunkedPrefillResult chunked = run_chunked_prefill(tokens, cache_total_seq_len_, get_prefill_chunk_size(), prepare_decode); cache_total_seq_len_ += chunked.logical_tokens; + 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& 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& tokens, diff --git a/cactus-engine/src/tokenizer.cpp b/cactus-engine/src/tokenizer.cpp index d792578db..879e3bebf 100644 --- a/cactus-engine/src/tokenizer.cpp +++ b/cactus-engine/src/tokenizer.cpp @@ -13,9 +13,11 @@ namespace engine { namespace { -std::string format_needle_query_text(const std::vector& messages) { +std::string format_needle_query_text(const std::vector& 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") { @@ -23,13 +25,38 @@ std::string format_needle_query_text(const std::vector& messages) { 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& messages, bool add_generation_prompt, @@ -538,8 +571,18 @@ std::string Tokenizer::format_qwen_style(const std::vector& message } long last_query_index = static_cast(n) - 1; - for (long i = static_cast(n) - 1; i >= 0; --i) { - if (messages[i].role == "user") { last_query_index = i; break; } + for (long q = static_cast(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 = ""; + const std::string close_tag = ""; + 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& message content = lstrip_newlines(content.substr(tpos + 8)); } result += "<|im_start|>assistant\n"; - if (template_has_thinking && static_cast(i) > last_query_index && (i == n - 1 || !reasoning.empty())) { + if (template_has_thinking && static_cast(i) > last_query_index) { result += "\n" + reasoning + "\n\n\n" + lstrip_newlines(content); } else { result += content; @@ -670,10 +713,10 @@ std::string Tokenizer::format_lfm2_style(const std::vector& message return result; } -std::string Tokenizer::format_needle_style(const std::vector& messages, bool /*add_generation_prompt*/, +std::string Tokenizer::format_needle_style(const std::vector& 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) + "" + serialized_tools + ""; + return format_needle_query_text(messages, add_generation_prompt) + "" + serialized_tools + ""; } std::string Tokenizer::format_gemma4_style(const std::vector& messages, bool add_generation_prompt, diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 279d1be02..394ebb0c5 100644 --- a/cactus-graph/cactus_graph.h +++ b/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,6 +744,7 @@ 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& node_ids); @@ -748,7 +752,7 @@ class CactusGraph { 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 persistent_node_ids_; std::unordered_set populated_node_ids_; + std::vector node_is_constant_; + bool constants_cached_ = false; std::unordered_set embedded_input_node_ids_; std::unordered_set retained_output_node_ids_; void build_metal_retype_plan(); - void invalidate_metal_state(); std::unordered_map metal_plans_; std::unordered_map> 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); diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp index cc52bda29..fe7603980 100644 --- a/cactus-graph/src/builder.cpp +++ b/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); } diff --git a/cactus-graph/src/core.cpp b/cactus-graph/src/core.cpp index f850c99bc..b38baaf92 100644 --- a/cactus-graph/src/core.cpp +++ b/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(); } diff --git a/cactus-graph/src/execute.cpp b/cactus-graph/src/execute.cpp index fc19b2ccd..a4de7b853 100644 --- a/cactus-graph/src/execute.cpp +++ b/cactus-graph/src/execute.cpp @@ -67,6 +67,7 @@ DECLARE_COMPUTE(compute_groupnorm_node); DECLARE_COMPUTE(compute_rope_gptj_node); DECLARE_COMPUTE(compute_lstm_cell_node); DECLARE_COMPUTE(compute_gated_deltanet_decode_node); +DECLARE_COMPUTE(compute_gated_deltanet_gate_log_node); DECLARE_COMPUTE(compute_gated_deltanet_prefill_node); DECLARE_COMPUTE(compute_stft_node); DECLARE_COMPUTE(compute_altup_predict_node); @@ -107,7 +108,7 @@ DECLARE_COMPUTE(compute_spectrogram_node); extern void shrink_thread_local_buffers(); #undef DECLARE_COMPUTE -static constexpr int OP_TYPE_COUNT = static_cast(OpType::CONV_CACHE_INITIALIZE) + 1; +static constexpr int OP_TYPE_COUNT = static_cast(OpType::GATED_DELTANET_GATE_LOG) + 1; static_assert(OP_TYPE_COUNT <= 256, "OpType dispatch table overflow"); static ComputeFn dispatch_flat[OP_TYPE_COUNT] = {}; @@ -186,6 +187,7 @@ static bool init_dispatch() { dispatch_flat[static_cast(OpType::PERSISTENT)] = compute_persistent_node; dispatch_flat[static_cast(OpType::LSTM_CELL)] = compute_lstm_cell_node; dispatch_flat[static_cast(OpType::GATED_DELTANET_DECODE)] = compute_gated_deltanet_decode_node; + dispatch_flat[static_cast(OpType::GATED_DELTANET_GATE_LOG)] = compute_gated_deltanet_gate_log_node; dispatch_flat[static_cast(OpType::GATED_DELTANET_PREFILL)] = compute_gated_deltanet_prefill_node; dispatch_flat[static_cast(OpType::STFT)] = compute_stft_node; dispatch_flat[static_cast(OpType::ALTUP_PREDICT)] = compute_altup_predict_node; @@ -253,8 +255,11 @@ static const char* op_type_names[] = { "NOT_EQUAL", "SCALAR_NOT_EQUAL", "RECURRENT_CACHE_STATE", "RECURRENT_CACHE_WRITE", - "CONV_CACHE_INITIALIZE" + "CONV_CACHE_INITIALIZE", + "GATED_DELTANET_GATE_LOG" }; +static_assert(sizeof(op_type_names) / sizeof(op_type_names[0]) == OP_TYPE_COUNT, + "op_type_names must have exactly one entry per OpType"); static const char* get_op_name(OpType op) { return op_type_names[static_cast(op)]; @@ -381,6 +386,8 @@ std::vector infer_output_shape(const GraphNode& node, const nodes_vector out[axis] = node.params.slice_length; return out; } + case OpType::GATED_DELTANET_GATE_LOG: + return in(0); default: { std::vector out = node.output_buffer.shape; if (out.empty()) return in(0); @@ -486,7 +493,8 @@ static const __fp16* metal_conv_weight_f16(const BufferDesc& w, size_t rows, siz return e.data.data(); } -static bool try_encode_metal(GraphNode& node, const nodes_vector& nodes, const node_index_map_t& map) { +static bool try_encode_metal(GraphNode& node, const nodes_vector& nodes, const node_index_map_t& map, + uint32_t prefill_valid_len) { if (node.params.backend != ComputeBackend::METAL) return false; BufferDesc& out = node.output_buffer; auto fp16 = [](const BufferDesc& b){ return b.precision == Precision::FP16; }; @@ -500,7 +508,12 @@ static bool try_encode_metal(GraphNode& node, const nodes_vector& nodes, const n CactusQuantMatrix mat = rhs.to_cq_matrix(); if (M == 1 && !cactus_graph_prefill_consistent()) return cactus_metal_encode_quant_matmul(out.get_data(), lhs.get_data(), &mat); - return cactus_metal_encode_quant_matmul_m(out.get_data(), lhs.get_data(), &mat, (uint32_t)M); + uint32_t Mq = (uint32_t)M; + if (prefill_valid_len && M > 1 && prefill_valid_len < M) { + uint32_t mv = ((prefill_valid_len + 31u) / 32u) * 32u; + if (mv < Mq) Mq = mv; + } + return cactus_metal_encode_quant_matmul_m(out.get_data(), lhs.get_data(), &mat, Mq); } if (!fp16(rhs) || !fp16(out) || rhs.shape.size() != 2) return false; size_t K = lhs.shape.back(); @@ -868,6 +881,17 @@ static bool try_encode_metal(GraphNode& node, const nodes_vector& nodes, const n for (size_t i = (size_t)axis + 1; i < in.shape.size(); ++i) inner *= in.shape[i]; return cactus_metal_encode_glu(out.get_data(), in.get_data(), split, inner, out.total_size); } + case OpType::GATED_DELTANET_GATE_LOG: { + if (node.input_ids.size() != 3) return false; + const auto& a = get_input(node, 0, nodes, map); + const auto& dt = get_input(node, 1, nodes, map); + const auto& al = get_input(node, 2, nodes, map); + if (!fp16(a) || !fp16(dt) || !fp16(al) || !fp16(out) || out.shape.empty()) return false; + size_t H = out.shape.back(); + if (H == 0) return false; + return cactus_metal_encode_deltanet_gate_log(out.get_data(), a.get_data(), + dt.get_data(), al.get_data(), out.total_size / H, H); + } case OpType::SOFTMAX: { const auto& in = get_input(node, 0, nodes, map); if (!fp16(in) || !fp16(out) || out.shape.empty()) return false; @@ -912,7 +936,8 @@ static bool try_encode_metal(GraphNode& node, const nodes_vector& nodes, const n if (out.total_size != B * (T + K) * Hv * V) return false; return cactus_metal_encode_deltanet_prefill(out.get_data(), q.get_data(), k.get_data(), v.get_data(), g.get_data(), b.get_data(), st.get_data(), - (uint32_t)B, (uint32_t)T, (uint32_t)Hq, (uint32_t)Hv, (uint32_t)K, (uint32_t)V, node.params.scale); + (uint32_t)B, (uint32_t)T, (uint32_t)Hq, (uint32_t)Hv, (uint32_t)K, (uint32_t)V, + node.params.scale, prefill_valid_len); } case OpType::RECURRENT_CACHE_WRITE: { if (node.input_ids.size() < 2) return false; @@ -1115,6 +1140,7 @@ static bool try_encode_metal(GraphNode& node, const nodes_vector& nodes, const n size_t current_len=km[0], max_len=km[1], kv_heads=km[2], hdim=km[3], sink=km[4], num_slots=km[5]; if (num_slots != 1 || kv_heads == 0 || hdim == 0) return false; size_t new_seq_len = new_kv.total_size / (kv_heads * hdim); + if (prefill_valid_len && prefill_valid_len < new_seq_len) new_seq_len = prefill_valid_len; size_t ceiling = cache_node.params.max_cache_seq_len, ws = node.params.window_size; bool sliding = ws > 0 && ws < ceiling; size_t new_total = current_len + new_seq_len; @@ -1213,6 +1239,7 @@ static bool try_encode_metal(GraphNode& node, const nodes_vector& nodes, const n uint64_t head = cm[0], count = cm[1], ws = cm[2], hd = cm[3]; if (ws == 0 || hd == 0 || head >= ws) return false; size_t num_rows = new_data.total_size / hd; + if (prefill_valid_len && prefill_valid_len < num_rows) num_rows = prefill_valid_len; if (num_rows == 0 || out.total_size != ws * hd) return false; uint32_t nnew = (uint32_t)std::min(num_rows, ws); uint64_t count_new = std::min(ws, count + num_rows); @@ -1342,7 +1369,7 @@ void CactusGraph::build_metal_retype_plan() { } } -void CactusGraph::execute(const std::string& profile_file) { +void CactusGraph::execute(const std::string& profile_file, uint32_t prefill_valid_len) { cactus_graph_mark_unadjusted(); BufferPool& pool = buffer_pool_; const size_t n = nodes_.size(); @@ -1434,6 +1461,38 @@ void CactusGraph::execute(const std::string& profile_file) { } } if (metal_mode && !need_debug && !metal_retype_built_) build_metal_retype_plan(); + if (node_is_constant_.size() != n) { + node_is_constant_.assign(n, 0); + constants_cached_ = false; + auto runtime_source = [&](const GraphNode& nd) { + switch (nd.op_type) { + case OpType::INPUT: + case OpType::KV_CACHE_STATE: + case OpType::CONV_CACHE_STATE: + case OpType::RECURRENT_CACHE_STATE: + case OpType::KV_CACHE_APPEND: + case OpType::CONV_CACHE_APPEND: + case OpType::RECURRENT_CACHE_WRITE: + case OpType::SAMPLE: + return true; + default: + return false; + } + }; + for (size_t i = 0; i < n; ++i) { + GraphNode& nd = *nodes_[i]; + if (runtime_source(nd)) { node_is_constant_[i] = 0; continue; } + bool cst = true; + for (size_t id : nd.input_ids) { + auto it = node_index_map_.find(id); + if (it == node_index_map_.end() || !node_is_constant_[it->second]) { cst = false; break; } + } + if (nd.input_ids.empty()) + cst = nd.op_type != OpType::INPUT && nd.output_buffer.get_data() != nullptr; + node_is_constant_[i] = cst ? 1 : 0; + } + } + MetalFusePlan* fplan = nullptr; if (metal_mode && !need_debug) { uint64_t sig = 1469598103934665603ull; @@ -1453,6 +1512,8 @@ void CactusGraph::execute(const std::string& profile_file) { if (it == metal_plans_.end()) { std::unordered_set pinned(retained_output_node_ids_.begin(), retained_output_node_ids_.end()); pinned.insert(persistent_node_ids_.begin(), persistent_node_ids_.end()); + for (size_t ci = 0; ci < n; ++ci) + if (node_is_constant_[ci]) pinned.insert(nodes_[ci]->id); static const std::vector no_retype; const std::vector& rt = (!metal_retype_disabled_ && metal_retype_plan_.size() == n) ? metal_retype_plan_ : no_retype; @@ -1571,6 +1632,7 @@ void CactusGraph::execute(const std::string& profile_file) { transient_ok[i] = !retained_output_node_ids_.count(id) && !persistent_node_ids_.count(id) && !keep_until_graph_cleanup[i] + && !node_is_constant_[i] && nodes_[i]->output_buffer.byte_size >= (256u << 10); } } @@ -1599,8 +1661,9 @@ void CactusGraph::execute(const std::string& profile_file) { size_t since_flush = 0, since_recycle = 0; if (fplan) cactus_metal_plan_fold(fplan, nodes_); std::vector metal_live(n, 0); + static const size_t recycle_cadence = (size_t)get_env_int("CACTUS_METAL_RECYCLE", 256); auto maybe_recycle = [&]() { - if (!transient_acts || ++since_recycle < 256 || transient_dead.empty()) return; + if (!transient_acts || ++since_recycle < recycle_cadence || transient_dead.empty()) return; for (void* p : transient_dead) cactus_metal_free_shared(p); transient_dead.clear(); cactus_metal_session_sync(); @@ -1647,6 +1710,10 @@ void CactusGraph::execute(const std::string& profile_file) { for (size_t r : release_after[i]) metal_release(r); continue; } + if (constants_cached_ && node_is_constant_[i] && !aliases_input(*node)) { + for (size_t r : release_after[i]) metal_release(r); + continue; + } int32_t fact = fplan ? cactus_metal_plan_action(fplan, i) : -1; if (fact == -3) { assign_persistent_act(*node); @@ -1696,7 +1763,7 @@ void CactusGraph::execute(const std::string& profile_file) { metal_plan_banned_[metal_plan_sig_].insert(i); metal_plans_.erase(metal_plan_sig_); cactus_metal_plan_free(fplan); - execute(profile_file); + execute(profile_file, prefill_valid_len); return; } std::vector> flipped; @@ -1713,13 +1780,13 @@ void CactusGraph::execute(const std::string& profile_file) { bool force_cpu = (ot == OpType::PRECISION_CAST && node->output_buffer.total_size <= 8 && !input_metal_live(*node)) || (ot == OpType::EMBEDDING && input_metal_live(*node)); - bool encoded = !force_cpu && try_encode_metal(*node, nodes_, node_index_map_); + bool encoded = !force_cpu && try_encode_metal(*node, nodes_, node_index_map_, prefill_valid_len); for (auto& fp : flipped) fp.first->precision = fp.second; if (!encoded && rplan && rplan[i] == 2) { metal_retype_disabled_ = true; metal_guard.armed = false; metal_abort_cleanup(); - execute(profile_file); + execute(profile_file, prefill_valid_len); return; } if (!encoded) { @@ -1741,6 +1808,7 @@ void CactusGraph::execute(const std::string& profile_file) { cactus_metal_set_active(false); cactus_metal_session_end(); cactus_graph_metal_tail_commit(); + constants_cached_ = true; return; } @@ -2156,10 +2224,13 @@ void CactusGraph::invalidate_metal_state() { metal_persistent_acts_.clear(); metal_retype_plan_.clear(); metal_retype_built_ = false; + constants_cached_ = false; } void CactusGraph::hard_reset() { invalidate_metal_state(); + node_is_constant_.clear(); + constants_cached_ = false; nodes_.clear(); node_index_map_.clear(); mapped_files_.clear(); diff --git a/cactus-graph/src/graph_ffi.cpp b/cactus-graph/src/graph_ffi.cpp index 99cc069bb..c5316d7ba 100644 --- a/cactus-graph/src/graph_ffi.cpp +++ b/cactus-graph/src/graph_ffi.cpp @@ -779,6 +779,18 @@ int cactus_graph_relu(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) } } +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) { + if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_gated_deltanet_gate_log"); + try { + *out = static_cast(as_graph(graph)->graph.gated_deltanet_gate_log( + static_cast(a_logits), static_cast(dt_bias), static_cast(a_log))); + return 0; + } catch (const std::exception& e) { + last_error_message = e.what(); + return -1; + } +} + int cactus_graph_silu(cactus_graph_t graph, cactus_node_t x, cactus_node_t* out) { if (!graph || !out) return fail_invalid("Invalid args to cactus_graph_silu"); try { diff --git a/cactus-graph/src/io.cpp b/cactus-graph/src/io.cpp index 147eeb98b..d54d28d1e 100644 --- a/cactus-graph/src/io.cpp +++ b/cactus-graph/src/io.cpp @@ -26,6 +26,7 @@ namespace { constexpr uint32_t FLAG_ORTHOGONAL_ROTATION = 1 << 1; constexpr uint32_t FLAG_INTERLEAVED_4ROW = 1 << 2; constexpr uint32_t FLAG_EXTENDED_SHAPE = 1 << 4; + constexpr uint32_t FLAG_NO_ROTATION = 1 << 5; constexpr size_t HEADER_SIZE = 84; inline size_t align_offset(size_t offset, size_t alignment) { @@ -241,7 +242,7 @@ namespace { GraphFile::NodeEntry node; node.index = read_u32(in); uint32_t op_type_val = read_u32(in); - if (op_type_val > static_cast(OpType::CONV_CACHE_INITIALIZE)) { + if (op_type_val > static_cast(OpType::GATED_DELTANET_GATE_LOG)) { throw std::runtime_error("Graph file corrupted: invalid op type"); } node.op_type = static_cast(op_type_val); @@ -435,6 +436,8 @@ size_t CactusGraph::mmap_embeddings(const std::string& filename) { if (mapped_file->is_orthogonal_rotation()) { buffer.cq_rotation = reinterpret_cast(scales_base + off); buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; + } else if (mapped_file->is_no_rotation()) { + buffer.cq_flags = CACTUS_QUANT_FLAG_NO_ROTATION; } else { buffer.cq_left_signs = reinterpret_cast(scales_base + off); off += gs; @@ -500,6 +503,8 @@ size_t CactusGraph::mmap_weights(const std::string& filename) { if (mapped_file->is_orthogonal_rotation()) { buffer.cq_rotation = reinterpret_cast(scales_base + off); buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; + } else if (mapped_file->is_no_rotation()) { + buffer.cq_flags = CACTUS_QUANT_FLAG_NO_ROTATION; } else { buffer.cq_left_signs = reinterpret_cast(scales_base + off); off += gs; @@ -583,6 +588,8 @@ void CactusGraph::bind_mmap_weights(size_t node_id, const std::string& filename) if (mapped_file->is_orthogonal_rotation()) { buffer.cq_rotation = reinterpret_cast(scales_base + off); buffer.cq_flags = CACTUS_QUANT_FLAG_ORTHOGONAL; + } else if (mapped_file->is_no_rotation()) { + buffer.cq_flags = CACTUS_QUANT_FLAG_NO_ROTATION; } else { buffer.cq_left_signs = reinterpret_cast(scales_base + off); off += gs; @@ -905,12 +912,14 @@ MappedFile::MappedFile(MappedFile&& other) noexcept alignment_(other.alignment_), is_orthogonal_rotation_(other.is_orthogonal_rotation_), is_interleaved_4row_(other.is_interleaved_4row_), + is_no_rotation_(other.is_no_rotation_), original_N_(other.original_N_) { other.fd_ = -1; other.mapped_data_ = nullptr; other.file_size_ = 0; other.is_orthogonal_rotation_ = false; other.is_interleaved_4row_ = false; + other.is_no_rotation_ = false; other.original_N_ = 0; } @@ -937,6 +946,7 @@ MappedFile& MappedFile::operator=(MappedFile&& other) noexcept { alignment_ = other.alignment_; is_orthogonal_rotation_ = other.is_orthogonal_rotation_; is_interleaved_4row_ = other.is_interleaved_4row_; + is_no_rotation_ = other.is_no_rotation_; original_N_ = other.original_N_; other.fd_ = -1; other.mapped_data_ = nullptr; @@ -995,6 +1005,7 @@ void MappedFile::parse_header() { offset += sizeof(uint32_t); is_orthogonal_rotation_ = (flags & FLAG_ORTHOGONAL_ROTATION) != 0; is_interleaved_4row_ = (flags & FLAG_INTERLEAVED_4ROW) != 0; + is_no_rotation_ = (flags & FLAG_NO_ROTATION) != 0; alignment_ = *reinterpret_cast(ptr + offset); offset += sizeof(uint32_t); diff --git a/cactus-graph/src/metal_plan.cpp b/cactus-graph/src/metal_plan.cpp index e0fb923ae..bd4bd1ab0 100644 --- a/cactus-graph/src/metal_plan.cpp +++ b/cactus-graph/src/metal_plan.cpp @@ -1237,6 +1237,65 @@ MetalFusePlan* cactus_metal_plan_build( collect_chain((size_t)idxof(nd0.input_ids[0]), sm, cover); if (add_cluster(c, i, cover)) plan->action[(size_t)sfm] = -3; } + for (size_t i = 0; i < n; ++i) { + if (plan->action[i] != -1) continue; + GraphNode& dn = *nodes[i]; + if (dn.op_type != OpType::DIVIDE || dn.input_ids.size() != 2) continue; + if (dn.output_buffer.precision != Precision::FP16 || retyped(i)) continue; + long xs = idxof(dn.input_ids[0]); + if (xs < 0 || nodes[(size_t)xs]->output_buffer.precision != Precision::FP16) continue; + std::vector cov; + auto strip1 = [&](long j) -> long { + while (j >= 0 && passthrough(*nodes[(size_t)j]) + && !nodes[(size_t)j]->input_ids.empty() + && cons[(size_t)j].size() == 1 && plan->action[(size_t)j] == -1) { + cov.push_back((size_t)j); + j = idxof(nodes[(size_t)j]->input_ids[0]); + } + return j; + }; + long sq = strip1(idxof(dn.input_ids[1])); + if (sq < 0 || nodes[(size_t)sq]->op_type != OpType::SCALAR_SQRT + || cons[(size_t)sq].size() != 1 || plan->action[(size_t)sq] != -1 || retyped((size_t)sq)) continue; + float eps = 0.f; + long j = strip1(nodes[(size_t)sq]->input_ids.empty() ? -1 : idxof(nodes[(size_t)sq]->input_ids[0])); + if (j < 0 || nodes[(size_t)j]->op_type != OpType::SCALAR_ADD + || cons[(size_t)j].size() != 1 || plan->action[(size_t)j] != -1 || retyped((size_t)j)) continue; + eps = nodes[(size_t)j]->params.scalar; + cov.push_back((size_t)sq); + cov.push_back((size_t)j); + long sm = strip1(nodes[(size_t)j]->input_ids.empty() ? -1 : idxof(nodes[(size_t)j]->input_ids[0])); + if (sm < 0 || nodes[(size_t)sm]->op_type != OpType::SUM + || cons[(size_t)sm].size() != 1 || plan->action[(size_t)sm] != -1 || retyped((size_t)sm)) continue; + cov.push_back((size_t)sm); + long mu = strip1(nodes[(size_t)sm]->input_ids.empty() ? -1 : idxof(nodes[(size_t)sm]->input_ids[0])); + if (mu < 0 || nodes[(size_t)mu]->op_type != OpType::MULTIPLY + || cons[(size_t)mu].size() != 1 || plan->action[(size_t)mu] != -1 || retyped((size_t)mu)) continue; + GraphNode& mun = *nodes[(size_t)mu]; + if (mun.input_ids.size() != 2) continue; + long m0 = idxof(mun.input_ids[0]), m1 = idxof(mun.input_ids[1]); + auto base_src = [&](long b) -> long { + while (b >= 0 && passthrough(*nodes[(size_t)b]) && !nodes[(size_t)b]->input_ids.empty()) + b = idxof(nodes[(size_t)b]->input_ids[0]); + return b; + }; + long bx = base_src(xs); + if (base_src(m0) != bx || base_src(m1) != bx) continue; + cov.push_back((size_t)mu); + size_t dim = dn.output_buffer.shape.empty() ? 0 : dn.output_buffer.shape.back(); + if (!dim || dn.output_buffer.total_size % dim) continue; + size_t drows = dn.output_buffer.total_size / dim; + if (nodes[(size_t)sq]->output_buffer.total_size != drows) continue; + MetalCluster c; + c.rule = 20; + c.a0 = (size_t)xs; + c.u0 = (uint32_t)drows; + c.u1 = (uint32_t)dim; + c.f0 = eps; + c.b4 = i; + add_cluster(c, i, cov); + } + for (size_t i = 0; i < n; ++i) { if (plan->action[i] != -1) continue; int kind, code; float p0, p1; @@ -1244,10 +1303,26 @@ MetalFusePlan* cactus_metal_plan_build( if (!prec_ok(nodes[i]->output_buffer.precision) || retyped(i)) continue; size_t total = nodes[i]->output_buffer.total_size; std::vector chain{i}; + std::vector hop_cover; + std::vector links{-1}; size_t cur = i; + const size_t head_inner = nodes[i]->output_buffer.shape.empty() + ? 1 : nodes[i]->output_buffer.shape.back(); while (chain.size() < 12) { if (cons[cur].size() != 1) break; size_t nxt = cons[cur][0]; + std::vector hops; + size_t link = cur; + while (passthrough(*nodes[nxt]) && plan->action[nxt] == -1 + && !retyped(nxt) + && nodes[nxt]->output_buffer.total_size == total + && (nodes[nxt]->output_buffer.shape.empty() + ? 1 : nodes[nxt]->output_buffer.shape.back()) == head_inner + && cons[nxt].size() == 1) { + hops.push_back(nxt); + link = nxt; + nxt = cons[nxt][0]; + } if (plan->action[nxt] != -1) break; int k2, c2; float q0, q1; if (!ew_kind(*nodes[nxt], &k2, &c2, &q0, &q1)) break; @@ -1255,8 +1330,10 @@ MetalFusePlan* cactus_metal_plan_build( if (nodes[nxt]->output_buffer.total_size != total) break; bool feeds = false; for (size_t a = 0; a < nodes[nxt]->input_ids.size() && a < 2; ++a) - if (idxof(nodes[nxt]->input_ids[a]) == (long)cur) feeds = true; + if (idxof(nodes[nxt]->input_ids[a]) == (long)link) feeds = true; if (!feeds) break; + hop_cover.insert(hop_cover.end(), hops.begin(), hops.end()); + links.push_back((long)link); chain.push_back(nxt); cur = nxt; } @@ -1277,7 +1354,7 @@ MetalFusePlan* cactus_metal_plan_build( const GraphNode& cn = *nodes[chain[ci]]; int k2, c2; float q0, q1; ew_kind(cn, &k2, &c2, &q0, &q1); - long prev = ci == 0 ? -1 : (long)chain[ci - 1]; + long prev = ci == 0 ? -1 : links[ci]; if (ci == 0) { long hi2 = idxof(cn.input_ids[0]); if (hi2 < 0 || !prec_ok(nodes[(size_t)hi2]->output_buffer.precision) @@ -1347,6 +1424,7 @@ MetalFusePlan* cactus_metal_plan_build( c.b4 = chain.back(); plan->blobs.push_back(std::move(blob)); std::vector cover(chain.begin(), chain.end() - 1); + cover.insert(cover.end(), hop_cover.begin(), hop_cover.end()); add_cluster(c, chain.back(), cover); } } @@ -1800,6 +1878,13 @@ bool cactus_metal_plan_encode(MetalFusePlan* p, int32_t cid, if (!lg || !probs || !tk) return false; return cactus_metal_encode_softmax_topk(probs, tk, lg, c.u0, c.u1, c.b0, c.f0); } + case 20: { + GraphNode& anchor = *nodes[c.b4]; + const void* x = nodes[c.a0]->output_buffer.get_data(); + void* out = anchor.output_buffer.get_data(); + if (!x || !out) return false; + return cactus_metal_encode_l2_norm(out, x, c.u0, c.u1, c.f0); + } case 18: { GraphNode& anchor = *nodes[c.b4]; const void* a = nodes[c.a0]->output_buffer.get_data(); diff --git a/cactus-graph/src/ops_nn.cpp b/cactus-graph/src/ops_nn.cpp index 3d192f52e..507fc5e16 100644 --- a/cactus-graph/src/ops_nn.cpp +++ b/cactus-graph/src/ops_nn.cpp @@ -831,6 +831,27 @@ void compute_layernorm_node(GraphNode& node, const std::vector>& nodes, + const std::unordered_map& node_index_map) { + const auto& a = get_input(node, 0, nodes, node_index_map); + const auto& dt = get_input(node, 1, nodes, node_index_map); + const auto& al = get_input(node, 2, nodes, node_index_map); + const __fp16* ap = a.data_as<__fp16>(); + const __fp16* dp = dt.data_as<__fp16>(); + const __fp16* lp = al.data_as<__fp16>(); + __fp16* op = node.output_buffer.data_as<__fp16>(); + size_t H = a.shape.empty() ? a.total_size : a.shape.back(); + size_t rows = H ? a.total_size / H : 0; + for (size_t r = 0; r < rows; ++r) { + for (size_t h = 0; h < H; ++h) { + float x = (float)ap[r*H+h] + (float)dp[h]; + float sp = x > 20.0f ? x : (x < -20.0f ? std::exp(x) : std::log1p(std::exp(x))); + op[r*H+h] = (__fp16)(-std::exp((float)lp[h]) * sp); + } + } +} + void compute_glu_node(GraphNode& node, const std::vector>& nodes, const std::unordered_map& node_index_map) { const auto& X = get_input(node, 0, nodes, node_index_map); diff --git a/cactus-graph/src/ops_tensor.cpp b/cactus-graph/src/ops_tensor.cpp index 1afb75563..89eb5935f 100644 --- a/cactus-graph/src/ops_tensor.cpp +++ b/cactus-graph/src/ops_tensor.cpp @@ -196,6 +196,7 @@ void compute_embedding_node(GraphNode& node, const std::vector 16) { diff --git a/cactus-graph/tests/test_nn.cpp b/cactus-graph/tests/test_nn.cpp index ffb26372e..073cb074b 100644 --- a/cactus-graph/tests/test_nn.cpp +++ b/cactus-graph/tests/test_nn.cpp @@ -53,7 +53,8 @@ static void write_test_cq_weights( const std::vector<__fp16>& norms, const std::vector& left_signs, const std::vector& right_signs, - const std::vector& permutation) { + const std::vector& permutation, + uint32_t extra_flags = 0) { constexpr uint32_t CACTUS_MAGIC = 0x54434143; constexpr uint32_t FLAG_HAS_SCALES = 1 << 0; constexpr size_t HEADER_SIZE = 84; @@ -80,7 +81,7 @@ static void write_test_cq_weights( }; write_u32(CACTUS_MAGIC); - write_u32(FLAG_HAS_SCALES); + write_u32(FLAG_HAS_SCALES | extra_flags); write_u32(alignment); write_u32(2); write_u64(N); @@ -111,6 +112,74 @@ static void write_test_cq_weights( if (!file) throw std::runtime_error("failed writing test CQ weights"); } +bool test_matmul_cq_no_rotation() { + const size_t M = 2, K = 256, N = 8; + const size_t gs = 128; + const size_t ng = K / gs; + constexpr uint32_t FLAG_NO_ROTATION_FILE = 1u << 5; + + std::mt19937 gen(99); + std::uniform_real_distribution dist(-1.f, 1.f); + + std::vector<__fp16> A(M * K); + for (auto& v : A) v = static_cast<__fp16>(dist(gen)); + + uint32_t pgb = cactus_quant_packed_group_bytes(1, gs); + std::vector packed(N * ng * pgb); + for (auto& v : packed) v = static_cast(gen() & 0xFF); + + std::vector<__fp16> codebook = {static_cast<__fp16>(-1.f), static_cast<__fp16>(1.f)}; + std::vector<__fp16> input_scale(K, static_cast<__fp16>(1.f)); + std::vector<__fp16> input_scale_recip(K, static_cast<__fp16>(1.f)); + std::vector<__fp16> norms(N * ng); + for (auto& v : norms) v = static_cast<__fp16>(0.01f + std::abs(dist(gen)) * 0.05f); + + CactusQuantMatrix mat{ + .bits = 1, .K = static_cast(K), .N = static_cast(N), + .group_size = static_cast(gs), .num_groups = static_cast(ng), + .flags = CACTUS_QUANT_FLAG_NO_ROTATION, + .codebook = codebook.data(), + .input_scale = input_scale.data(), + .input_scale_recip = input_scale_recip.data(), + .norms = norms.data(), + .packed_indices = packed.data(), + .left_signs = nullptr, + .right_signs = nullptr, + .permutation = nullptr, + .rotation = nullptr, + .expanded = nullptr, + .norm_f32 = nullptr, + }; + + std::vector<__fp16> direct(M * N, static_cast<__fp16>(0)); + cactus_quant_matmul(&mat, A.data(), static_cast(M), direct.data()); + + auto path = std::filesystem::temp_directory_path() / "cactus_graph_cq1_norot_matmul.weights"; + write_test_cq_weights(path, 1, K, N, gs, packed, codebook, input_scale, + input_scale_recip, norms, {}, {}, {}, FLAG_NO_ROTATION_FILE); + + CactusGraph g; + size_t ia = g.input({M, K}, Precision::FP16); + size_t iw = g.mmap_weights(path.string()); + size_t out = g.matmul(ia, iw, true); + g.set_input(ia, A.data(), Precision::FP16); + g.execute(); + __fp16* graph_out = static_cast<__fp16*>(g.get_output(out)); + + bool ok = true; + for (size_t i = 0; i < M * N; i++) { + float actual = static_cast(graph_out[i]); + float expected = static_cast(direct[i]); + if (!std::isfinite(actual) || std::abs(actual - expected) > 1e-3f) { + ok = false; + break; + } + } + g.hard_reset(); + std::filesystem::remove(path); + return ok; +} + bool test_matmul_cq() { // Test graph-level CQ matmul dispatch for every supported bit width. const size_t M = 2, K = 128, N = 8; @@ -710,6 +779,7 @@ int main() { runner.run_test("Matrix Multiplication", test_matrix_multiplication()); runner.run_test("MatMul CQ", test_matmul_cq()); + runner.run_test("MatMul CQ (no rotation)", test_matmul_cq_no_rotation()); runner.run_test("Transpose", test_transpose()); runner.run_test("RMS Norm", test_rms_norm()); runner.run_test("Softmax", test_softmax()); diff --git a/cactus-graph/tests/test_ops.cpp b/cactus-graph/tests/test_ops.cpp index 009f7d1c4..06e4ce028 100644 --- a/cactus-graph/tests/test_ops.cpp +++ b/cactus-graph/tests/test_ops.cpp @@ -477,7 +477,7 @@ bool test_cq_embedding_operation() { bits, hidden_dim, group_size, num_groups, row, packed.data(), codebook.data(), norms.data(), input_scale_recip.data(), left_signs.data(), right_signs.data(), - permutation.data(), expected.data() + row * hidden_dim); + permutation.data(), 0u, expected.data() + row * hidden_dim); } CactusGraph graph; diff --git a/cactus-kernels/cactus_kernels.h b/cactus-kernels/cactus_kernels.h index 4260e9e2a..b9e96cae7 100644 --- a/cactus-kernels/cactus_kernels.h +++ b/cactus-kernels/cactus_kernels.h @@ -201,6 +201,7 @@ void cactus_matmul_f16( enum CactusQuantFlags : uint32_t { CACTUS_QUANT_FLAG_ORTHOGONAL = 1u << 2, CACTUS_QUANT_FLAG_INTERLEAVED_4ROW = 1u << 3, + CACTUS_QUANT_FLAG_NO_ROTATION = 1u << 4, }; struct CactusQuantMatrix { @@ -322,6 +323,7 @@ void cactus_quant_dequantize_hadamard_embedding_row( const int8_t* left_signs, const int8_t* right_signs, const uint32_t* permutation, + uint32_t flags, __fp16* out_row); void cactus_quant_dequantize_orthogonal_embedding_row( diff --git a/cactus-kernels/metal_backend.h b/cactus-kernels/metal_backend.h index 9505aad48..9213eb828 100644 --- a/cactus-kernels/metal_backend.h +++ b/cactus-kernels/metal_backend.h @@ -42,6 +42,9 @@ bool cactus_metal_encode_unary(int op_type, void* out, const void* in, size_t n) bool cactus_metal_encode_swiglu(void* out, const void* gate, const void* up, size_t n, float scale); bool cactus_metal_encode_rms_norm(void* out, const void* in, const void* weight, size_t rows, size_t dim, float eps); +bool cactus_metal_encode_deltanet_gate_log(void* out, const void* a, const void* dt, + const void* al, size_t rows, size_t H); +bool cactus_metal_encode_l2_norm(void* out, const void* in, size_t rows, size_t dim, float eps); bool cactus_metal_encode_rms_norm_add(void* out, const void* in, const void* weight, const void* res, size_t rows, size_t dim, float eps); bool cactus_metal_encode_rms_norm_add_rms(void* h_out, void* xn_out, const void* in, const void* w1, @@ -78,7 +81,7 @@ bool cactus_metal_encode_deltanet_decode(void* out, const void* q, const void* k bool cactus_metal_encode_deltanet_prefill(void* out, const void* q, const void* k, const void* v, const void* g, const void* b, const void* s, uint32_t B, uint32_t T, uint32_t Hq, uint32_t Hv, - uint32_t K, uint32_t V, float scale); + uint32_t K, uint32_t V, float scale, uint32_t valid_T = 0); bool cactus_metal_encode_rms2_add_clip(void* out, const void* a, const void* wa, const void* b, const void* wb, size_t dim, float eps_a, float eps_b); diff --git a/cactus-kernels/src/cactus_kernels.metal b/cactus-kernels/src/cactus_kernels.metal index 01c3b8b4d..31c6388f2 100644 --- a/cactus-kernels/src/cactus_kernels.metal +++ b/cactus-kernels/src/cactus_kernels.metal @@ -471,7 +471,20 @@ kernel void cq_gemv_mr_lowbit( for (uint r=0;r>2)*num_groups+g)*4u*(size_t)pgb + byte_off*4u; + uint rr = n&3u; + w = (uint)pan[rr] | ((uint)pan[4u+rr]<<8); + } else { + w = *(device const ushort*)(packed + ((size_t)n*num_groups+g)*pgb + byte_off); + } + w0=cb[w&1u]; w1=cb[(w>>1)&1u]; w2=cb[(w>>2)&1u]; w3=cb[(w>>3)&1u]; + w4=cb[(w>>4)&1u]; w5=cb[(w>>5)&1u]; w6=cb[(w>>6)&1u]; w7=cb[(w>>7)&1u]; + w8=cb[(w>>8)&1u]; w9=cb[(w>>9)&1u]; wa=cb[(w>>10)&1u]; wb=cb[(w>>11)&1u]; + wc=cb[(w>>12)&1u]; wd=cb[(w>>13)&1u]; we=cb[(w>>14)&1u]; wf=cb[(w>>15)&1u]; + } else if (bits == 2u) { uint w; if (il != 0u) { device const uchar* pan = packed + ((size_t)(n>>2)*num_groups+g)*4u*(size_t)pgb + byte_off*4u; @@ -520,6 +533,156 @@ kernel void cq_gemv_mr_lowbit( } } +kernel void cq1_norot_gemv( + device const half* x [[buffer(0)]], + device const uchar* packed [[buffer(1)]], + device const half* codebook [[buffer(2)]], + device const half* norms [[buffer(3)]], + device half* y [[buffer(4)]], + constant uint& gs [[buffer(5)]], + constant uint& num_groups [[buffer(6)]], + constant uint& pgb [[buffer(7)]], + constant uint& N [[buffer(8)]], + constant uint& il [[buffer(9)]], + uint tg [[threadgroup_position_in_grid]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]]) +{ + (void)il; + const float cbscale = (float)codebook[1]; + const uint K = num_groups*gs; + const uint out_row = tg*8u + sgid*4u; + if (out_row >= N) return; + const uint rows = min(4u, N - out_row); + + float xv[32]; + float result[4] = {0.f, 0.f, 0.f, 0.f}; + + device const uchar* ws = packed + (size_t)out_row*num_groups*pgb + lane*4u; + device const half* xp = x + lane*32u; + + for (uint k = lane*32u; k < K; k += 1024u) { + float sum = 0.f; + #pragma clang loop unroll(full) + for (uint i = 0; i < 32u; ++i) { + float v = (float)xp[i]; + xv[i] = v; + sum += v; + } + const uint g = k / gs; + #pragma clang loop unroll(full) + for (uint row = 0; row < 4u; ++row) { + if (row >= rows) break; + device const uchar* wl = ws + (size_t)row*num_groups*pgb; + float accum = 0.f; + #pragma clang loop unroll(full) + for (uint i = 0; i < 4u; ++i) { + uint wb = wl[i]; + accum += select(0.f, xv[8u*i ], bool(wb & 0x01u)); + accum += select(0.f, xv[8u*i + 1], bool(wb & 0x02u)); + accum += select(0.f, xv[8u*i + 2], bool(wb & 0x04u)); + accum += select(0.f, xv[8u*i + 3], bool(wb & 0x08u)); + accum += select(0.f, xv[8u*i + 4], bool(wb & 0x10u)); + accum += select(0.f, xv[8u*i + 5], bool(wb & 0x20u)); + accum += select(0.f, xv[8u*i + 6], bool(wb & 0x40u)); + accum += select(0.f, xv[8u*i + 7], bool(wb & 0x80u)); + } + float nrm = (float)norms[(size_t)(out_row+row)*num_groups + g]; + result[row] += nrm*(2.0f*accum - sum); + } + ws += 128u; + xp += 1024u; + } + #pragma clang loop unroll(full) + for (uint row = 0; row < 4u; ++row) { + float tot = simd_sum(result[row]); + if (lane == 0 && row < rows) y[out_row+row] = (half)(cbscale*tot); + } +} + +kernel void cq1_norot_gemm_mma( + device const half* x [[buffer(0)]], + device const uchar* packed [[buffer(1)]], + device const half* norms [[buffer(2)]], + device half* y [[buffer(3)]], + constant uint& gs [[buffer(4)]], + constant uint& num_groups [[buffer(5)]], + constant uint& pgb [[buffer(6)]], + constant uint& N [[buffer(7)]], + constant uint& M [[buffer(8)]], + constant float& cbscale [[buffer(9)]], + uint2 tg [[threadgroup_position_in_grid]], + uint tl [[thread_index_in_threadgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) +{ + threadgroup float pool[2048]; + threadgroup half* sa = (threadgroup half*)pool; + threadgroup float* Cs = pool; + const uint NK=64u; + uint K = num_groups*gs; + uint r0 = tg.y*64u; + uint r1 = tg.x*32u; + uint lr0 = tl/2u, il0 = tl%2u; + uint nrow = r0+lr0; + uint ny = (lr0%8u), nyb = (lr0/8u); + simdgroup_matrix ma[4], mb[2]; + simdgroup_matrix mc[8]; + for (uint i=0;i<8u;i++) mc[i]=make_filled_simdgroup_matrix(0.f); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint loop_k=0; loop_k>3)); + #pragma clang loop unroll(full) + for (uint w=0; w<2u; w++) { + uint koff = 32u*il0 + 16u*w; + uint w16 = wp[w]; + #pragma clang loop unroll(full) + for (uint i=0;i<16u;i++) { + uint kk = koff + i; + float v = (loop_k+kk>3) + nyb; + sa[64u*ib + 8u*(kk&7u) + ny] = (half)v; + } + } + } else { + #pragma clang loop unroll(full) + for (uint kk=32u*il0; kk<32u*il0+32u; kk++) { + uint ib = 8u*(kk>>3) + nyb; + sa[64u*ib + 8u*(kk&7u) + ny] = (half)0; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup const half* lsma = sa + 4u*64u*(sg%2u); + device const half* ax = x + (size_t)(r1 + 16u*(sg/2u))*K + loop_k; + #pragma clang loop unroll(full) + for (uint ik=0; ik>1u), n_sg = 32u*(sg&1u); + for (uint i=0;i<8u;i++) + simdgroup_store(mc[i], &Cs[(m_sg+8u*(i/4u))*64u + n_sg+8u*(i%4u)], 64u); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint idx=tl; idx<32u*64u; idx+=128u){ + uint m=idx/64u, n=idx%64u; + if (r1+m=n) return; + uint h = i % H; + float x = (float)a[i] + (float)dt[h]; + float sp = x > 20.0f ? x : (x < -20.0f ? exp(x) : log(1.0f + exp(x))); + y[i] = (half)(-exp((float)al[h]) * sp); +} +kernel void l2_norm_rows_f16(device const half* in [[buffer(0)]], device half* y [[buffer(1)]], + constant uint& dim [[buffer(2)]], constant float& eps [[buffer(3)]], + uint row [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], + uint nt [[threads_per_threadgroup]], threadgroup float* red [[threadgroup(0)]]) { + device const half* x = in + (size_t)row*dim; + device half* o = y + (size_t)row*dim; + float partial=0; + for (uint i=t;i0; s>>=1){ if (t(cactus_quant_input_scale_recip1(W, offset)); + code_basis[j] = static_cast<__fp16>(static_cast(x_group[j]) * scale); + } + return; + } __fp16 tmp[256]; __fp16* work = (W.permutation == nullptr) ? code_basis : tmp; @@ -1047,6 +1061,89 @@ void cactus_quant_2bit_gemm(const CactusQuantMatrix* W, const __fp16* A, uint32_ +alignas(16) static const uint8_t tq1_bitsel[16] = { + 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128}; +alignas(16) static const uint8_t tq1_pair_shuf[8][16] = { + {0,0,0,0,0,0,0,0, 1,1,1,1,1,1,1,1}, + {2,2,2,2,2,2,2,2, 3,3,3,3,3,3,3,3}, + {4,4,4,4,4,4,4,4, 5,5,5,5,5,5,5,5}, + {6,6,6,6,6,6,6,6, 7,7,7,7,7,7,7,7}, + {8,8,8,8,8,8,8,8, 9,9,9,9,9,9,9,9}, + {10,10,10,10,10,10,10,10, 11,11,11,11,11,11,11,11}, + {12,12,12,12,12,12,12,12, 13,13,13,13,13,13,13,13}, + {14,14,14,14,14,14,14,14, 15,15,15,15,15,15,15,15}, +}; + +__attribute__((always_inline)) static inline int32_t tq1_masked_group_dot( + const uint8_t* rowbits, const int8_t* act, int32_t s_all) { + const uint8x16_t bitsel = vld1q_u8(tq1_bitsel); + const uint8x16_t rb = vld1q_u8(rowbits); + int16x8_t acc16 = vdupq_n_s16(0); + for (uint32_t v = 0; v < 8; ++v) { + uint8x16_t mask = vtstq_u8(vqtbl1q_u8(rb, vld1q_u8(tq1_pair_shuf[v])), bitsel); + int8x16_t sel = vandq_s8(vreinterpretq_s8_u8(mask), vld1q_s8(act + v * 16)); + acc16 = vpadalq_s8(acc16, sel); + } + return 2 * static_cast(vaddlvq_s16(acc16)) - s_all; +} + +static bool cactus_quant_masked_gemv_cq1( + const CactusQuantMatrix* W, + const __fp16* code_basis, + __fp16* y) { + if (W->group_size != 128) return false; + if (static_cast(W->codebook[0]) != -1.0f || + static_cast(W->codebook[1]) != 1.0f) return false; + + const uint32_t gs = W->group_size; + const uint32_t pgb = cactus_quant_packed_group_bytes(1, gs); + const uint32_t groups = W->num_groups; + + thread_local std::vector act_i8_buf; + thread_local std::vector act_scales_buf; + thread_local std::vector s_all_buf; + if (act_i8_buf.size() < W->K) act_i8_buf.resize(W->K); + if (act_scales_buf.size() < groups) act_scales_buf.resize(groups); + if (s_all_buf.size() < groups) s_all_buf.resize(groups); + for (uint32_t g = 0; g < groups; ++g) { + act_scales_buf[g] = tq_quantize_group_i8( + code_basis + static_cast(g) * gs, + act_i8_buf.data() + static_cast(g) * gs, gs); + const int8_t* ag = act_i8_buf.data() + static_cast(g) * gs; + int16x8_t s16 = vdupq_n_s16(0); + for (uint32_t v = 0; v < 8; ++v) s16 = vpadalq_s8(s16, vld1q_s8(ag + v * 16)); + s_all_buf[g] = vaddlvq_s16(s16); + } + const int8_t* act_i8 = act_i8_buf.data(); + const float* act_scales = act_scales_buf.data(); + const int32_t* s_all = s_all_buf.data(); + + const size_t n_blocks = (W->N + 3) / 4; + cactus_quant_parallel_ranges(n_blocks, 64, [&](size_t block_start, size_t block_end) { + for (size_t nb = block_start; nb < block_end; ++nb) { + const size_t n0 = nb * 4; + const size_t rows = std::min(4, W->N - n0); + float acc[4] = {}; + for (uint32_t g = 0; g < groups; ++g) { + const int8_t* ag = act_i8 + static_cast(g) * gs; + const float scale = act_scales[g]; + const int32_t sa = s_all[g]; + for (size_t r = 0; r < rows; ++r) { + const uint8_t* rowbits = + W->packed_indices + ((n0 + r) * groups + g) * pgb; + const float norm = + static_cast(W->norms[(n0 + r) * groups + g]); + acc[r] += norm * scale * + static_cast(tq1_masked_group_dot(rowbits, ag, sa)); + } + } + for (size_t r = 0; r < rows; ++r) + y[n0 + r] = static_cast<__fp16>(acc[r]); + } + }); + return true; +} + struct CactusTQ1ScaledDecoder { uint8x8_t cb_bytes; @@ -1097,6 +1194,10 @@ void cactus_quant_1bit_gemv( thread_local std::vector<__fp16> code_basis_buf; if (code_basis_buf.size() < W->K) code_basis_buf.resize(W->K); cactus_quant_transform_hadamard_activations(*W, x, 1, code_basis_buf.data()); + if ((W->flags & CACTUS_QUANT_FLAG_NO_ROTATION) && + cactus_quant_masked_gemv_cq1(W, code_basis_buf.data(), y)) { + return; + } cactus_quant_sdot_gemv_int8<1>(W, code_basis_buf.data(), y); return; } @@ -1902,6 +2003,7 @@ void cactus_quant_dequantize_hadamard_embedding_row( const int8_t* left_signs, const int8_t* right_signs, const uint32_t* permutation, + uint32_t flags, __fp16* out_row) { if (!packed_base || !codebook || !norms || !out_row || bits == 0 || bits > 4) return; if (hidden_dim == 0 || group_size == 0 || num_groups == 0) return; @@ -1909,6 +2011,22 @@ void cactus_quant_dequantize_hadamard_embedding_row( if (group_size > 256) return; const uint32_t packed_group_bytes = cactus_quant_packed_group_bytes(bits, group_size); + if (flags & CACTUS_QUANT_FLAG_NO_ROTATION) { + for (uint32_t g = 0; g < num_groups; ++g) { + const uint8_t* packed = packed_base + (static_cast(row) * num_groups + g) * packed_group_bytes; + const float norm = static_cast(norms[static_cast(row) * num_groups + g]); + for (uint32_t k = 0; k < group_size; ++k) { + uint8_t idx = static_cast(tq_extract_idx(packed, k, bits)); + uint32_t col = g * group_size + k; + float scale = input_scale_recip ? static_cast(input_scale_recip[col]) : 1.0f; + out_row[col] = static_cast<__fp16>(static_cast(codebook[idx]) * norm * scale); + } + } + for (uint32_t col = num_groups * group_size; col < hidden_dim; ++col) { + out_row[col] = static_cast<__fp16>(0); + } + return; + } std::vector rotated(group_size); for (uint32_t g = 0; g < num_groups; ++g) { std::fill(rotated.begin(), rotated.end(), 0.0f); diff --git a/cactus-kernels/src/metal_backend.mm b/cactus-kernels/src/metal_backend.mm index 0c4a5f236..0cbbfea98 100644 --- a/cactus-kernels/src/metal_backend.mm +++ b/cactus-kernels/src/metal_backend.mm @@ -37,7 +37,7 @@ id psoTsimd = nil; id psoTbatch = nil; id psoMega = nil, psoSwiT = nil, psoCat = nil; - id psoG = nil, psoGmr = nil, psoGmrLow = nil; + id psoG = nil, psoGmr = nil, psoGmrLow = nil, psoG1N = nil, psoG1NM = nil, psoL2N = nil, psoGateLog = nil; id psoActQ = nil, psoG2i8 = nil; id psoTm = nil; id psoGmma = nil, psoGdense = nil; @@ -91,6 +91,10 @@ psoTm=pso("cq4_transform_m"); psoGmma=pso("cq4_gemm_mma"); psoGdense=pso("cq4_gemm_dense_f16"); psoGmrLow=pso("cq_gemv_mr_lowbit"); + psoG1N=pso("cq1_norot_gemv"); + psoG1NM=pso("cq1_norot_gemm_mma"); + psoL2N=pso("l2_norm_rows_f16"); + psoGateLog=pso("deltanet_gate_log_f16"); psoActQ=pso("cq_act_quant_i8"); psoG2i8=pso("cq2_gemv_i8"); psoRotW=pso("lmhead_rotate_wide"); psoEmbOW=pso("emb_ortho_wide"); psoRmsAddRms=pso("rms_norm_add_rms_f16"); psoEmbH=pso("emb_hadamard"); @@ -190,6 +194,7 @@ static uint64_t resident_key(const CactusQuantMatrix* W) { } id g_attn_scores = nil; std::unordered_map g_resident_emb; +std::unordered_map, id>> g_deint_packed; std::mutex g_resident_mu; struct MoESet4 { @@ -427,6 +432,7 @@ void cactus_metal_invalidate_host_wraps() { g_readonly.clear(); g_resident.clear(); g_resident_emb.clear(); + g_deint_packed.clear(); g_mps_mats.clear(); g_moe_cat.clear(); } @@ -709,13 +715,60 @@ static inline bool quant_fast_eligible(const CactusQuantMatrix* W) { } static inline bool quant_lowbit_eligible(const CactusQuantMatrix* W) { - return (W->bits == 2 || W->bits == 3) && W->group_size == 128 && (W->N % 16) == 0 && + return (W->bits == 1 || W->bits == 2 || W->bits == 3) && W->group_size == 128 && (W->N % 16) == 0 && (W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) && !(W->flags & CACTUS_QUANT_FLAG_ORTHOGONAL) && W->input_scale_recip && W->left_signs && W->right_signs && W->permutation && W->codebook && W->norms && W->packed_indices; } +static inline bool quant_cq1_norot_eligible(const CactusQuantMatrix* W) { + return W->bits == 1 && W->group_size == 128 && (W->N % 16) == 0 && + (W->flags & CACTUS_QUANT_FLAG_NO_ROTATION) && + !(W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) && + W->codebook && W->norms && W->packed_indices && + (float)W->codebook[0] == -1.0f && (float)W->codebook[1] == 1.0f; +} + +static inline bool quant_cq1_interleaved_sym_eligible(const CactusQuantMatrix* W) { + return W->bits == 1 && W->group_size == 128 && (W->N % 16) == 0 && (W->N % 4) == 0 && + (W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW) && + !(W->flags & CACTUS_QUANT_FLAG_ORTHOGONAL) && + W->codebook && W->norms && W->packed_indices && + W->input_scale_recip && W->left_signs && W->right_signs && W->permutation && + (float)W->codebook[0] == -(float)W->codebook[1] && (float)W->codebook[0] < 0.0f; +} + +static bool deint_bufs(const CactusQuantMatrix* W, id* out_packed, id* out_norms) { + std::lock_guard lk(g_resident_mu); + auto it = g_deint_packed.find(W->packed_indices); + if (it != g_deint_packed.end()) { + *out_packed = it->second.first; *out_norms = it->second.second; + return *out_packed != nil && *out_norms != nil; + } + const uint32_t N = W->N, ng = W->num_groups, nb = N / 4; + id pb = owned_shared((size_t)N * ng * 16); + id nb_buf = owned_shared((size_t)N * ng * sizeof(__fp16)); + if (!pb || !nb_buf) return false; + uint8_t* pdst = (uint8_t*)[pb contents]; + const uint8_t* psrc = (const uint8_t*)W->packed_indices; + __fp16* ndst = (__fp16*)[nb_buf contents]; + const __fp16* nsrc = W->norms; + for (uint32_t b = 0; b < nb; ++b) + for (uint32_t g = 0; g < ng; ++g) { + for (uint32_t r = 0; r < 4; ++r) + ndst[((size_t)(b * 4 + r) * ng + g)] = nsrc[(size_t)(b * ng + g) * 4 + r]; + for (uint32_t c = 0; c < 4; ++c) + for (uint32_t bc = 0; bc < 4; ++bc) + for (uint32_t r = 0; r < 4; ++r) + pdst[((size_t)(b * 4 + r) * ng + g) * 16 + c * 4 + bc] = + psrc[(((size_t)(b * ng + g) * 4 + c) * 4 + bc) * 4 + r]; + } + g_deint_packed.emplace(W->packed_indices, std::make_pair(pb, nb_buf)); + *out_packed = pb; *out_norms = nb_buf; + return true; +} + void cactus_metal_trim_prefill_cache() { std::lock_guard lk(g_resident_mu); g_mps_mats.clear(); @@ -737,6 +790,53 @@ bool cactus_metal_encode_quant_matmul(void* out, const void* lhs, const CactusQu } if (ortho_code) return cactus_metal_encode_quant_matmul_ortho(out, lhs, ortho_code, W); } + if (ctx().psoG1N && quant_cq1_norot_eligible(W)) { + ResW& rw = resident(W); + if (!rw.ok) return false; + const uint32_t pgbw = (gs + 7u) / 8u; + ensureEncoder(); + [g_enc setComputePipelineState:ctx().psoG1N]; + setBufAt(lhs, (size_t)K*2, 0); + [g_enc setBuffer:rw.packed offset:rw.packed_off atIndex:1]; + [g_enc setBuffer:rw.codebook offset:rw.cb_off atIndex:2]; + [g_enc setBuffer:rw.norms offset:rw.norms_off atIndex:3]; + setBufAt(out, (size_t)N*2, 4); + [g_enc setBytes:&gs length:4 atIndex:5]; [g_enc setBytes:&ng length:4 atIndex:6]; + [g_enc setBytes:&pgbw length:4 atIndex:7]; [g_enc setBytes:&N length:4 atIndex:8]; + [g_enc setBytes:&rw.il length:4 atIndex:9]; + [g_enc dispatchThreadgroups:MTLSizeMake((N+7u)/8u,1,1) + threadsPerThreadgroup:MTLSizeMake(64,1,1)]; + return true; + } + if (ctx().psoG1N && quant_cq1_interleaved_sym_eligible(W)) { + ResW& rw = resident(W); + id pb = nil, nb = nil; + if (!rw.ok || !rw.recip || !rw.lsign || !rw.rsign || !rw.perm || !rw.codebook || !deint_bufs(W, &pb, &nb)) return false; + size_t code_bytes = (size_t)ng*gs*sizeof(__fp16); + if (!g_code_buf || (size_t)g_code_buf.length < code_bytes) g_code_buf = owned_shared(code_bytes); + if (!g_code_buf) return false; + ensureEncoder(); + bool simdT = (gs==128u && ctx().psoTsimd); + [g_enc setComputePipelineState:(simdT?ctx().psoTsimd:ctx().psoT)]; + setBufAt(lhs, (size_t)K*2, 0); [g_enc setBuffer:rw.recip offset:rw.rc_off atIndex:1]; + [g_enc setBuffer:rw.lsign offset:rw.ls_off atIndex:2]; [g_enc setBuffer:rw.rsign offset:rw.rs_off atIndex:3]; + [g_enc setBuffer:rw.perm offset:rw.pm_off atIndex:4]; [g_enc setBuffer:g_code_buf offset:0 atIndex:5]; + [g_enc setBytes:&gs length:4 atIndex:6]; [g_enc setThreadgroupMemoryLength:gs*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(ng,1,1) threadsPerThreadgroup:MTLSizeMake(simdT?32u:(gs>1024u?1024u:gs),1,1)]; + const uint32_t pgbw = (gs + 7u) / 8u, il0 = 0u; + [g_enc setComputePipelineState:ctx().psoG1N]; + [g_enc setBuffer:g_code_buf offset:0 atIndex:0]; + [g_enc setBuffer:pb offset:0 atIndex:1]; + [g_enc setBuffer:rw.codebook offset:rw.cb_off atIndex:2]; + [g_enc setBuffer:nb offset:0 atIndex:3]; + setBufAt(out, (size_t)N*2, 4); + [g_enc setBytes:&gs length:4 atIndex:5]; [g_enc setBytes:&ng length:4 atIndex:6]; + [g_enc setBytes:&pgbw length:4 atIndex:7]; [g_enc setBytes:&N length:4 atIndex:8]; + [g_enc setBytes:&il0 length:4 atIndex:9]; + [g_enc dispatchThreadgroups:MTLSizeMake((N+7u)/8u,1,1) + threadsPerThreadgroup:MTLSizeMake(64,1,1)]; + return true; + } const bool fast = quant_fast_eligible(W); const bool lowbit = !fast && ctx().psoGmrLow && quant_lowbit_eligible(W); if (!fast && !lowbit) return false; @@ -913,9 +1013,38 @@ bool cactus_metal_encode_gemv_cat(void* const* outs, const void* const* codes, return true; } +bool cactus_metal_encode_deltanet_gate_log(void* out, const void* a, const void* dt, + const void* al, size_t rows, size_t H) { + if (!ctx().ok || !ctx().psoGateLog) return false; + ensureEncoder(); + uint32_t h=(uint32_t)H, n=(uint32_t)(rows*H); + [g_enc setComputePipelineState:ctx().psoGateLog]; + setBufAt(a, rows*H*2, 0); setBufAt(dt, H*2, 1); setBufAt(al, H*2, 2); setBufAt(out, rows*H*2, 3); + [g_enc setBytes:&h length:4 atIndex:4]; [g_enc setBytes:&n length:4 atIndex:5]; + [g_enc dispatchThreads:MTLSizeMake(n,1,1) threadsPerThreadgroup:MTLSizeMake(256,1,1)]; + return true; +} + +bool cactus_metal_encode_l2_norm(void* out, const void* in, size_t rows, size_t dim, float eps) { + if (!ctx().ok || !ctx().psoL2N) return false; + ensureEncoder(); + uint32_t d=(uint32_t)dim; float e=eps; + [g_enc setComputePipelineState:ctx().psoL2N]; + setBufAt(in, rows*dim*2, 0); setBufAt(out, rows*dim*2, 1); + [g_enc setBytes:&d length:4 atIndex:2]; [g_enc setBytes:&e length:4 atIndex:3]; + [g_enc setThreadgroupMemoryLength:128*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake(rows,1,1) threadsPerThreadgroup:MTLSizeMake(128,1,1)]; + return true; +} + bool cactus_metal_prewarm_quant(const CactusQuantMatrix* W) { if (!ctx().ok) return false; - if (!quant_fast_eligible(W) && !quant_lowbit_eligible(W)) return false; + if (quant_cq1_interleaved_sym_eligible(W)) { + id pb = nil, nb = nil; + return resident(W).ok && deint_bufs(W, &pb, &nb); + } + if (!quant_fast_eligible(W) && !quant_lowbit_eligible(W) && + !quant_cq1_norot_eligible(W)) return false; return resident(W).ok; } @@ -1156,9 +1285,10 @@ bool cactus_metal_encode_deltanet_decode(void* out, const void* q, const void* k bool cactus_metal_encode_deltanet_prefill(void* out, const void* q, const void* k, const void* v, const void* g, const void* b, const void* s, uint32_t B, uint32_t T, uint32_t Hq, uint32_t Hv, - uint32_t K, uint32_t V, float scale) { + uint32_t K, uint32_t V, float scale, uint32_t valid_T) { if (!ctx().ok || !ctx().psoDeltanetPre) return false; if (B == 0 || T == 0 || Hq == 0 || Hv == 0 || K == 0 || V == 0 || V > 1024 || (Hv % Hq) != 0) return false; + if (valid_T == 0 || valid_T > T) valid_T = T; ensureEncoder(); id scratch = recycled((size_t)B*Hv*K*V*sizeof(float)); [g_enc setComputePipelineState:ctx().psoDeltanetPre]; @@ -1173,6 +1303,7 @@ bool cactus_metal_encode_deltanet_prefill(void* out, const void* q, const void* [g_enc setBytes:&T length:4 atIndex:8]; [g_enc setBytes:&Hq length:4 atIndex:9]; [g_enc setBytes:&Hv length:4 atIndex:10]; [g_enc setBytes:&K length:4 atIndex:11]; [g_enc setBytes:&V length:4 atIndex:12]; [g_enc setBytes:&scale length:4 atIndex:13]; + [g_enc setBytes:&valid_T length:4 atIndex:14]; [g_enc setThreadgroupMemoryLength:(size_t)K*4 atIndex:0]; [g_enc dispatchThreadgroups:MTLSizeMake(Hv, B, 1) threadsPerThreadgroup:MTLSizeMake(V, 1, 1)]; return true; @@ -1280,6 +1411,49 @@ bool cactus_metal_encode_quant_matmul_m(void* out, const void* lhs, const Cactus if (!ctx().ok || !W) return false; if (M == 1 && (W->flags & CACTUS_QUANT_FLAG_ORTHOGONAL)) return cactus_metal_encode_quant_matmul(out, lhs, W); + const bool cq1_norot = quant_cq1_norot_eligible(W) && + !(W->flags & CACTUS_QUANT_FLAG_INTERLEAVED_4ROW); + const bool cq1_sym = !cq1_norot && quant_cq1_interleaved_sym_eligible(W); + if (ctx().psoG1NM && (cq1_norot || cq1_sym)) { + if (M == 1) return cactus_metal_encode_quant_matmul(out, lhs, W); + if ((M % 32u) == 0u) { + const uint32_t gs1=W->group_size, K1=W->K, N1=W->N, ng1=W->num_groups; + const uint32_t pgb1=(gs1+7u)/8u; + ResW& rw = resident(W); + if (!rw.ok) return false; + id packed = rw.packed, norms = rw.norms; size_t packed_off = rw.packed_off, norms_off = rw.norms_off; + float cbscale = 1.0f; + ensureEncoder(); + if (cq1_sym) { + id pb = nil, nb = nil; + if (!rw.recip || !rw.lsign || !rw.rsign || !rw.perm || !rw.codebook || !deint_bufs(W, &pb, &nb)) return false; + packed = pb; packed_off = 0; norms = nb; norms_off = 0; + cbscale = (float)W->codebook[1]; + size_t cbytes = (size_t)M*ng1*gs1*sizeof(__fp16); + if (!g_code_buf_m || (size_t)g_code_buf_m.length < cbytes) g_code_buf_m = owned_shared(cbytes); + if (!g_code_buf_m) return false; + [g_enc setComputePipelineState:ctx().psoTm]; + setBufAt(lhs, (size_t)M*K1*2, 0); [g_enc setBuffer:rw.recip offset:rw.rc_off atIndex:1]; + [g_enc setBuffer:rw.lsign offset:rw.ls_off atIndex:2]; [g_enc setBuffer:rw.rsign offset:rw.rs_off atIndex:3]; + [g_enc setBuffer:rw.perm offset:rw.pm_off atIndex:4]; [g_enc setBuffer:g_code_buf_m offset:0 atIndex:5]; + [g_enc setBytes:&gs1 length:4 atIndex:6]; [g_enc setBytes:&K1 length:4 atIndex:7]; + [g_enc setThreadgroupMemoryLength:gs1*sizeof(float) atIndex:0]; + [g_enc dispatchThreadgroups:MTLSizeMake((size_t)ng1*M,1,1) threadsPerThreadgroup:MTLSizeMake(gs1>1024u?1024u:gs1,1,1)]; + } + [g_enc setComputePipelineState:ctx().psoG1NM]; + if (cq1_sym) [g_enc setBuffer:g_code_buf_m offset:0 atIndex:0]; + else setBufAt(lhs, (size_t)M*K1*2, 0); + [g_enc setBuffer:packed offset:packed_off atIndex:1]; + [g_enc setBuffer:norms offset:norms_off atIndex:2]; + setBufAt(out, (size_t)M*N1*2, 3); + [g_enc setBytes:&gs1 length:4 atIndex:4]; [g_enc setBytes:&ng1 length:4 atIndex:5]; + [g_enc setBytes:&pgb1 length:4 atIndex:6]; [g_enc setBytes:&N1 length:4 atIndex:7]; + [g_enc setBytes:&M length:4 atIndex:8]; [g_enc setBytes:&cbscale length:4 atIndex:9]; + [g_enc dispatchThreadgroups:MTLSizeMake(M/32u, (N1+63u)/64u, 1) + threadsPerThreadgroup:MTLSizeMake(128,1,1)]; + return true; + } + } const uint32_t gs=W->group_size, K=W->K, N=W->N, ng=W->num_groups; const uint32_t bits=W->bits, pgb=(gs*bits+7u)/8u; diff --git a/cactus-kernels/src/metal_backend_stub.cpp b/cactus-kernels/src/metal_backend_stub.cpp index 9ecafc05b..0c616a0ef 100644 --- a/cactus-kernels/src/metal_backend_stub.cpp +++ b/cactus-kernels/src/metal_backend_stub.cpp @@ -120,7 +120,7 @@ bool cactus_metal_encode_rope_pair_rms(void*, const void*, const void*, const vo bool cactus_metal_encode_deltanet_decode(void*, const void*, const void*, const void*, const void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, float) { return false; } bool cactus_metal_encode_deltanet_prefill(void*, const void*, const void*, const void*, - const void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, float) { return false; } + const void*, const void*, const void*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, float, uint32_t) { return false; } bool cactus_metal_encode_rms2_add_clip(void*, const void*, const void*, const void*, const void*, size_t, float, float) { return false; } bool cactus_metal_encode_rms_norm_scale(void*, const void*, const void*, @@ -133,4 +133,6 @@ bool cactus_metal_moe_cq4_build(const CactusQuantMatrix*, const CactusQuantMatri const CactusQuantMatrix*, uint32_t) { return false; } bool cactus_metal_encode_moe_gated_cq4(void*, const void*, const void*, const void*, const CactusQuantMatrix*, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, float, float) { return false; } +bool cactus_metal_encode_deltanet_gate_log(void*, const void*, const void*, const void*, size_t, size_t) { return false; } +bool cactus_metal_encode_l2_norm(void*, const void*, size_t, size_t, float) { return false; } #endif diff --git a/cactus-kernels/tests/test_matmul.cpp b/cactus-kernels/tests/test_matmul.cpp index 74f917974..79c8738e5 100644 --- a/cactus-kernels/tests/test_matmul.cpp +++ b/cactus-kernels/tests/test_matmul.cpp @@ -316,6 +316,118 @@ bool test_cq_correctness(uint32_t bits) { return true; } +bool test_cq1_no_rotation_exactness() { + const uint32_t K = 1024, N = 64, gs = 128, ng = K / gs; + std::mt19937 gen(321); + std::uniform_real_distribution dist(-1.f, 1.f); + + std::vector<__fp16> codebook = {static_cast<__fp16>(-1.f), static_cast<__fp16>(1.f)}; + std::vector<__fp16> input_scale(K, static_cast<__fp16>(1.f)); + std::vector<__fp16> input_scale_recip(K, static_cast<__fp16>(1.f)); + std::vector<__fp16> norms(size_t(N) * ng); + for (auto& v : norms) v = static_cast<__fp16>(0.01f + std::abs(dist(gen)) * 0.05f); + uint32_t pgb = cactus_quant_packed_group_bytes(1, gs); + std::vector packed(size_t(N) * ng * pgb); + for (auto& v : packed) v = static_cast(gen() & 0xFF); + + CactusQuantMatrix mat = {}; + mat.bits = 1; + mat.K = K; + mat.N = N; + mat.group_size = gs; + mat.num_groups = ng; + mat.flags = CACTUS_QUANT_FLAG_NO_ROTATION; + mat.codebook = codebook.data(); + mat.input_scale = input_scale.data(); + mat.input_scale_recip = input_scale_recip.data(); + mat.norms = norms.data(); + mat.packed_indices = packed.data(); + + std::vector x_f32(K); + for (auto& v : x_f32) v = dist(gen); + + std::vector ref(N, 0.f); + for (uint32_t n = 0; n < N; ++n) { + for (uint32_t g = 0; g < ng; ++g) { + const uint8_t* pr = packed.data() + (size_t(n) * ng + g) * pgb; + float gsum = 0.f; + for (uint32_t k = 0; k < gs; ++k) { + uint8_t idx = unpack_index(pr, 1, k); + gsum += x_f32[g * gs + k] * (idx ? 1.f : -1.f); + } + ref[n] += static_cast(norms[size_t(n) * ng + g]) * gsum; + } + } + + std::vector<__fp16> x_f16(K), y(N, static_cast<__fp16>(0)); + for (size_t i = 0; i < K; i++) x_f16[i] = static_cast<__fp16>(x_f32[i]); + cactus_quant_matmul(&mat, x_f16.data(), 1, y.data()); + + double mse = compute_mse(ref.data(), y.data(), N); + double threshold = 1e-3; + if (mse > threshold) { + std::cerr << " cq1-no-rotation MSE=" << mse << " > " << threshold << "\n"; + return false; + } + return true; +} + +bool test_cq1_no_rotation_gemm() { + const uint32_t M = 33, K = 1024, N = 80, gs = 128, ng = K / gs; + std::mt19937 gen(3210); + std::uniform_real_distribution dist(-1.f, 1.f); + + std::vector<__fp16> codebook = {static_cast<__fp16>(-1.f), static_cast<__fp16>(1.f)}; + std::vector<__fp16> input_scale(K, static_cast<__fp16>(1.f)); + std::vector<__fp16> input_scale_recip(K, static_cast<__fp16>(1.f)); + std::vector<__fp16> norms(size_t(N) * ng); + for (auto& v : norms) v = static_cast<__fp16>(0.01f + std::abs(dist(gen)) * 0.05f); + uint32_t pgb = cactus_quant_packed_group_bytes(1, gs); + std::vector packed(size_t(N) * ng * pgb); + for (auto& v : packed) v = static_cast(gen() & 0xFF); + + CactusQuantMatrix mat = {}; + mat.bits = 1; + mat.K = K; + mat.N = N; + mat.group_size = gs; + mat.num_groups = ng; + mat.flags = CACTUS_QUANT_FLAG_NO_ROTATION; + mat.codebook = codebook.data(); + mat.input_scale = input_scale.data(); + mat.input_scale_recip = input_scale_recip.data(); + mat.norms = norms.data(); + mat.packed_indices = packed.data(); + + std::vector x_f32(size_t(M) * K); + for (auto& v : x_f32) v = dist(gen); + + std::vector ref(size_t(M) * N, 0.f); + for (uint32_t m = 0; m < M; ++m) + for (uint32_t n = 0; n < N; ++n) + for (uint32_t g = 0; g < ng; ++g) { + const uint8_t* pr = packed.data() + (size_t(n) * ng + g) * pgb; + float gsum = 0.f; + for (uint32_t k = 0; k < gs; ++k) { + uint8_t idx = unpack_index(pr, 1, k); + gsum += x_f32[size_t(m) * K + g * gs + k] * (idx ? 1.f : -1.f); + } + ref[size_t(m) * N + n] += static_cast(norms[size_t(n) * ng + g]) * gsum; + } + + std::vector<__fp16> x_f16(size_t(M) * K), y(size_t(M) * N, static_cast<__fp16>(0)); + for (size_t i = 0; i < x_f16.size(); i++) x_f16[i] = static_cast<__fp16>(x_f32[i]); + cactus_quant_matmul(&mat, x_f16.data(), M, y.data()); + + double mse = compute_mse(ref.data(), y.data(), size_t(M) * N); + double threshold = 1e-3; + if (mse > threshold) { + std::cerr << " cq1-no-rotation-gemm MSE=" << mse << " > " << threshold << "\n"; + return false; + } + return true; +} + bool run_benchmarks() { auto bench = [](const char* label, size_t M, size_t K, size_t N, auto fn) { fn(); @@ -581,6 +693,8 @@ int main() { runner.run_test("matmul_cq2", test_cq_correctness(2)); runner.run_test("matmul_cq3", test_cq_correctness(3)); runner.run_test("matmul_cq4", test_cq_correctness(4)); + runner.run_test("matmul_cq1_norot", test_cq1_no_rotation_exactness()); + runner.run_test("matmul_cq1_norot_gemm", test_cq1_no_rotation_gemm()); { double m1 = 0; runner.run_test("matmul_cq4_il", test_cq4_interleaved(m1)); diff --git a/python/cactus/bindings/cactus.py b/python/cactus/bindings/cactus.py index 5901d6a4c..f78aedd88 100644 --- a/python/cactus/bindings/cactus.py +++ b/python/cactus/bindings/cactus.py @@ -350,6 +350,10 @@ class cactus_tensor_info_t(ctypes.Structure): cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) ] _lib.cactus_graph_silu.restype = ctypes.c_int +_lib.cactus_graph_gated_deltanet_gate_log.argtypes = [ + cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.POINTER(cactus_node_t) +] +_lib.cactus_graph_gated_deltanet_gate_log.restype = ctypes.c_int _lib.cactus_graph_gelu.argtypes = [ cactus_graph_t, cactus_node_t, ctypes.POINTER(cactus_node_t) ] @@ -2312,6 +2316,18 @@ def relu(self, x, backend=None): raise RuntimeError(_err("graph_relu failed")) return self._apply_backend(self._tensor_from_node(out.value), backend) + def gated_deltanet_gate_log(self, a_logits, dt_bias, a_log, backend=None): + a_logits = self._ensure_tensor(a_logits) + dt_bias = self._ensure_tensor(dt_bias) + a_log = self._ensure_tensor(a_log) + out = cactus_node_t() + rc = _lib.cactus_graph_gated_deltanet_gate_log( + self.h, cactus_node_t(a_logits.id), cactus_node_t(dt_bias.id), + cactus_node_t(a_log.id), ctypes.byref(out)) + if rc != 0: + raise RuntimeError(_err("graph_gated_deltanet_gate_log failed")) + return self._apply_backend(self._tensor_from_node(out.value), backend) + def silu(self, x, backend=None): x = self._ensure_tensor(x) out = cactus_node_t() diff --git a/python/cactus/cli/convert.py b/python/cactus/cli/convert.py index cf2279326..efd8c7d82 100644 --- a/python/cactus/cli/convert.py +++ b/python/cactus/cli/convert.py @@ -46,6 +46,14 @@ def cmd_convert(args): then build the runtime graph. Pass --weights-only to stop after the weight conversion. """ from .model import ensure_weights + from cactus.transpile.model_profiles import profile_for_model_id + + profile = profile_for_model_id(args.model_id) + if profile is not None: + if profile.default_skip_model_load: + args.skip_model_load = True + if profile.default_low_memory_load: + args.low_memory_load = True source_model_id = args.model_id merged_dir = None diff --git a/python/cactus/convert/cli.py b/python/cactus/convert/cli.py index 3a494f961..e6fd0b577 100644 --- a/python/cactus/convert/cli.py +++ b/python/cactus/convert/cli.py @@ -32,7 +32,7 @@ from .model_adapters.detection import SUPPORTED_FAMILIES, detect_family from .model_adapters.adapters import adapter_for_family from .model_adapters.nemo import ensure_parakeet_tdt_nemo_source -from .quantization.cq import quantize_hadamard, quantize_orthogonal, write_cq_tensor +from .quantization.cq import quantize_binary_repack, quantize_hadamard, quantize_orthogonal, write_cq_tensor from .compat import patch_transformers_import_compat ALPHA = 0.25 @@ -523,23 +523,31 @@ def convert(args: argparse.Namespace) -> None: if args.strict: raise RuntimeError(f"{name}: {hessian_missing_reason} ({module_name})") hessian = hessians_np.get(module_name) - input_scale = None - if emit_policy.rotation == "orthogonal" or name.endswith("embed_tokens_per_layer.weight"): - input_scale = _row_table_scale(emit_tensor, scale_token_ids) - else: - input_scale = _input_scale_from_diag(diag_np.get(module_name), emit_tensor) - if emit_policy.rotation == "orthogonal": - cq = quantize_orthogonal(emit_tensor, bits=int(emit_policy.bits or 4), input_scale=input_scale) - else: - cq = quantize_hadamard( - emit_tensor, - bits=int(emit_policy.bits or requested_bits), - hessian=hessian, - use_gptq=emit_policy.use_gptq, - input_scale=input_scale, - ) + cq = None + if requested_bits == 1: + cq = quantize_binary_repack(emit_tensor) + if cq is not None: + print(f" repack: {out_path.name} is QAT-binary -> lossless 1-bit repack (no rotation)") + precision = "CQ1" + emit_policy = replace(emit_policy, precision="CQ1", bits=1) + if cq is None: + input_scale = None + if emit_policy.rotation == "orthogonal" or name.endswith("embed_tokens_per_layer.weight"): + input_scale = _row_table_scale(emit_tensor, scale_token_ids) + else: + input_scale = _input_scale_from_diag(diag_np.get(module_name), emit_tensor) + if emit_policy.rotation == "orthogonal": + cq = quantize_orthogonal(emit_tensor, bits=int(emit_policy.bits or 4), input_scale=input_scale) + else: + cq = quantize_hadamard( + emit_tensor, + bits=int(emit_policy.bits or requested_bits), + hessian=hessian, + use_gptq=emit_policy.use_gptq, + input_scale=input_scale, + ) cq = _scale_cq_norms(cq, adapter.scale_factor(out_path.name)) - if getattr(emit_policy, "layout", "row_major") == "interleaved_4row": + if getattr(emit_policy, "layout", "row_major") == "interleaved_4row" and cq.rotation_family != "none": cq = replace(cq, interleaved_4row=True) write_cq_tensor(out_path, cq) gptq_used = cq.gptq_used diff --git a/python/cactus/convert/export/qdq.py b/python/cactus/convert/export/qdq.py index 5c73081d6..cffc1372b 100644 --- a/python/cactus/convert/export/qdq.py +++ b/python/cactus/convert/export/qdq.py @@ -25,6 +25,7 @@ FLAG_ORTHOGONAL_ROTATION = 1 << 1 FLAG_INTERLEAVED_4ROW = 1 << 2 FLAG_INTERLEAVED = 1 << 3 +FLAG_NO_ROTATION = 1 << 5 PRECISION_INT8 = 0 PRECISION_FP16 = 1 @@ -306,6 +307,17 @@ def parse_normal_metadata(blob: bytes, header: CactusHeader): return torch.from_numpy(codebook), torch.from_numpy(input_scale), torch.from_numpy(norms), torch.from_numpy(left), torch.from_numpy(right), torch.from_numpy(perm) +def parse_no_rotation_metadata(blob: bytes, header: CactusHeader): + n, k = header.shape + pos = 0 + codebook = np.frombuffer(blob, dtype=np.float16, count=1 << header.bits, offset=pos).astype(np.float32) + pos += (1 << header.bits) * 2 + input_scale = np.frombuffer(blob, dtype=np.float16, count=k, offset=pos).astype(np.float32) + pos += k * 4 + norms = np.frombuffer(blob, dtype=np.float16, count=n * header.num_groups, offset=pos).astype(np.float32).reshape(n, header.num_groups) + return torch.from_numpy(codebook), torch.from_numpy(input_scale), torch.from_numpy(norms) + + def parse_orthogonal_metadata(blob: bytes, header: CactusHeader): n, k = header.shape pos = 0 @@ -351,15 +363,19 @@ def dequantize_cq_file(path: Path, header: CactusHeader, out_dtype: torch.dtype, out[start:end] = (recon / scale).to(out_dtype) return out packed_group_bytes = math.ceil(header.group_size * bits / 8) - codebook, input_scale, norms, left, right, perm = parse_normal_metadata(scales_blob, header) - signs = set(int(x) for x in left.tolist()) | set(int(x) for x in right.tolist()) - if not signs.issubset({-1, 1}): - raise ValueError(f"{path}: signs must be +/-1, got {sorted(signs)}") - if sorted(int(x) for x in perm.tolist()) != list(range(header.group_size)): - raise ValueError(f"{path}: permutation is not bijective") - base_h = torch.from_numpy((hadamard(header.group_size, dtype=float) / math.sqrt(header.group_size)).astype(np.float32)) - rotation = (left.float().unsqueeze(1) * base_h * right.float().unsqueeze(0))[:, perm.long()].contiguous() - rt = rotation.T.contiguous() + if header.flags & FLAG_NO_ROTATION: + codebook, input_scale, norms = parse_no_rotation_metadata(scales_blob, header) + rt = torch.eye(header.group_size, dtype=torch.float32) + else: + codebook, input_scale, norms, left, right, perm = parse_normal_metadata(scales_blob, header) + signs = set(int(x) for x in left.tolist()) | set(int(x) for x in right.tolist()) + if not signs.issubset({-1, 1}): + raise ValueError(f"{path}: signs must be +/-1, got {sorted(signs)}") + if sorted(int(x) for x in perm.tolist()) != list(range(header.group_size)): + raise ValueError(f"{path}: permutation is not bijective") + base_h = torch.from_numpy((hadamard(header.group_size, dtype=float) / math.sqrt(header.group_size)).astype(np.float32)) + rotation = (left.float().unsqueeze(1) * base_h * right.float().unsqueeze(0))[:, perm.long()].contiguous() + rt = rotation.T.contiguous() scale = input_scale.float().unsqueeze(0) codebook = codebook.float() diff --git a/python/cactus/convert/quantization/cq.py b/python/cactus/convert/quantization/cq.py index 2d3942389..1e38f1d8b 100644 --- a/python/cactus/convert/quantization/cq.py +++ b/python/cactus/convert/quantization/cq.py @@ -24,6 +24,7 @@ PRECISION_CQ = {1: 3, 2: 4, 3: 5, 4: 6} FLAG_ORTHOGONAL_ROTATION = 1 << 1 FLAG_INTERLEAVED_4ROW = 1 << 2 +FLAG_NO_ROTATION = 1 << 5 @dataclass(frozen=True) @@ -37,6 +38,7 @@ class CQTensor: seed: int = 1234 gptq_used: bool = False interleaved_4row: bool = False + codebook: np.ndarray | None = None _CODEBOOK_CACHE: dict[tuple[int, int], np.ndarray] = {} @@ -311,6 +313,39 @@ def quantize_orthogonal(weight, bits: int = 4, seed: int = 1234, input_scale: np return CQTensor(indices=indices, norms=norms.astype(np.float16).reshape(n, 1), input_scale=input_scale.astype(np.float16), bits=bits, group_size=k, rotation_family="orthogonal") +def quantize_binary_repack(weight, group_size: int = GROUP_SIZE) -> CQTensor | None: + w = _as_numpy_fp32(weight) + if w.ndim != 2: + return None + n, k = w.shape + if k == 0 or k % group_size != 0: + return None + groups = k // group_size + indices = np.zeros((n, k), dtype=np.uint8) + norms = np.zeros((n, groups), dtype=np.float16) + for start in range(0, n, 4096): + stop = min(start + 4096, n) + block = w[start:stop].reshape(stop - start, groups, group_size) + mags = np.abs(block) + gmax = mags.max(axis=2) + if not np.all((gmax == mags.min(axis=2)) | (gmax == 0.0)): + return None + s16 = gmax.astype(np.float16) + if not np.array_equal(s16.astype(np.float32), gmax): + return None + norms[start:stop] = s16 + indices[start:stop] = (block > 0).reshape(stop - start, k) + return CQTensor( + indices=indices, + norms=norms, + input_scale=np.ones(k, dtype=np.float16), + bits=1, + group_size=group_size, + rotation_family="none", + codebook=np.array([-1.0, 1.0], dtype=np.float16), + ) + + _EMBEDDING_TENSOR_STEMS = frozenset({ "token_embeddings", "embed_tokens", @@ -343,7 +378,7 @@ def write_cq_tensor(out_path: Path, cq: CQTensor) -> None: cq.interleaved_4row or ( cq.bits in (1, 2, 3, 4) - and cq.rotation_family != "orthogonal" + and cq.rotation_family not in ("orthogonal", "none") and n % 4 == 0 and group_size % 32 == 0 and group_size <= 256 @@ -351,9 +386,14 @@ def write_cq_tensor(out_path: Path, cq: CQTensor) -> None: ) ) - codebook_f32 = make_codebook(group_size, cq.bits).astype(np.float32) + if cq.codebook is not None: + codebook_f32 = cq.codebook.astype(np.float32) + if codebook_f32.shape != (1 << cq.bits,): + raise ValueError(f"explicit codebook must have {1 << cq.bits} entries, got {codebook_f32.shape}") + else: + codebook_f32 = make_codebook(group_size, cq.bits).astype(np.float32) norms_f32 = cq.norms.astype(np.float32, copy=True) - if interleaved and cq.rotation_family != "orthogonal": + if interleaved and cq.rotation_family not in ("orthogonal", "none"): cb_max = float(np.max(np.abs(codebook_f32))) cb_factor = cb_max / 127.0 if cb_max > 1e-12 else 1.0 / 127.0 codebook_f32 = codebook_f32 / cb_factor @@ -377,6 +417,8 @@ def write_cq_tensor(out_path: Path, cq: CQTensor) -> None: if cq.rotation_family == "orthogonal": flags |= FLAG_ORTHOGONAL_ROTATION parts.append(make_orthogonal_rotation(group_size, cq.seed).astype(np.float16).tobytes()) + elif cq.rotation_family == "none": + flags |= FLAG_NO_ROTATION else: left, right, perm = make_hadamard_components(group_size, cq.seed) parts.extend([left.tobytes(), right.tobytes(), perm.astype(" ) b_logits = g.reshape(b_logits, (batch_size, seq_len, int(b_weight.shape[0]))) - if dt_bias is not None: + if (dt_bias is not None and a_log is not None + and a_logits.dtype == Graph.FP16 + and dt_bias.dtype == Graph.FP16 + and a_log.dtype == Graph.FP16): dt_bias_2d = g.reshape(dt_bias, (1, int(dt_bias.shape[0]))) - a_logits, dt_bias_2d = _legalize_elementwise_binary_inputs(g, a_logits, dt_bias_2d) - a_logits = g.add(a_logits, dt_bias_2d) - a_softplus = _lower_softplus(g, a_logits) - - if a_log is not None: a_log_2d = g.reshape(a_log, (1, int(a_log.shape[0]))) - neg_exp_a = g.scalar_multiply(g.scalar_exp(a_log_2d), -1.0) - neg_exp_a, a_softplus = _legalize_elementwise_binary_inputs(g, neg_exp_a, a_softplus) - gate_log = g.multiply(neg_exp_a, a_softplus) + gate_log = g.gated_deltanet_gate_log(a_logits, dt_bias_2d, a_log_2d) else: - gate_log = g.scalar_multiply(a_softplus, -1.0) + if dt_bias is not None: + dt_bias_2d = g.reshape(dt_bias, (1, int(dt_bias.shape[0]))) + a_logits, dt_bias_2d = _legalize_elementwise_binary_inputs(g, a_logits, dt_bias_2d) + a_logits = g.add(a_logits, dt_bias_2d) + a_softplus = _lower_softplus(g, a_logits) + if a_log is not None: + a_log_2d = g.reshape(a_log, (1, int(a_log.shape[0]))) + neg_exp_a = g.scalar_multiply(g.scalar_exp(a_log_2d), -1.0) + neg_exp_a, a_softplus = _legalize_elementwise_binary_inputs(g, neg_exp_a, a_softplus) + gate_log = g.multiply(neg_exp_a, a_softplus) + else: + gate_log = g.scalar_multiply(a_softplus, -1.0) beta = g.sigmoid(b_logits) cache_layer_key, initial_state = _lower_gated_deltanet_initial_state( diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index ba9e3bf04..33b4c0227 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -4863,6 +4863,13 @@ def _require(component: str) -> None: return specs +def _qwen35_rotary_embeddings(backbone, hidden_states, position_ids): + rope_device = backbone.rotary_emb.inv_freq.device + rope_x = torch.empty(0, dtype=hidden_states.dtype, device=rope_device) + cos, sin = backbone.rotary_emb(rope_x, position_ids.to(rope_device)) + return cos.to(hidden_states.device), sin.to(hidden_states.device) + + class Qwen35CausalLMLogitsAdapter(torch.nn.Module): def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): super().__init__() @@ -4883,7 +4890,7 @@ def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[tor if self.pad_token_id is not None else None ) - base_position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).view(1, 1, -1) + base_position_ids = torch.arange(inputs_embeds.shape[1], device=input_ids.device).view(1, 1, -1) position_ids = torch.cat( (base_position_ids, base_position_ids, base_position_ids, base_position_ids), dim=0, @@ -4902,7 +4909,7 @@ def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[tor hidden_states = inputs_embeds checkpoints: list[torch.Tensor] = [] - position_embeddings = self.backbone.rotary_emb(hidden_states, multimodal_position_ids) + position_embeddings = _qwen35_rotary_embeddings(self.backbone, hidden_states, multimodal_position_ids) for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): hidden_states = decoder_layer( @@ -4973,7 +4980,7 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> torch. linear_attn_mask = None hidden_states = inputs_embeds - position_embeddings = self.backbone.rotary_emb(hidden_states, multimodal_position_ids) + position_embeddings = _qwen35_rotary_embeddings(self.backbone, hidden_states, multimodal_position_ids) for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): hidden_states = decoder_layer( hidden_states, @@ -5186,7 +5193,7 @@ def forward( position_ids=text_position_ids, ) hidden_states = inputs_embeds - position_embeddings = self.backbone.rotary_emb(hidden_states, multimodal_position_ids) + position_embeddings = _qwen35_rotary_embeddings(self.backbone, hidden_states, multimodal_position_ids) for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): hidden_states = decoder_layer( hidden_states, @@ -5279,7 +5286,7 @@ def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> to position_ids=text_position_ids, ) hidden_states = inputs_embeds - position_embeddings = self.backbone.rotary_emb(hidden_states, multimodal_position_ids) + position_embeddings = _qwen35_rotary_embeddings(self.backbone, hidden_states, multimodal_position_ids) for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): hidden_states = decoder_layer( hidden_states, @@ -5356,7 +5363,7 @@ def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[tor hidden_states = inputs_embeds checkpoints: list[torch.Tensor] = [] - position_embeddings = self.backbone.rotary_emb(hidden_states, position_ids) + position_embeddings = _qwen35_rotary_embeddings(self.backbone, hidden_states, position_ids) for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): hidden_states = decoder_layer( @@ -5426,7 +5433,7 @@ def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> torch. hidden_states = inputs_embeds text_position_ids = position_ids.to(dtype=torch.int64) - position_embeddings = self.backbone.rotary_emb(hidden_states, text_position_ids) + position_embeddings = _qwen35_rotary_embeddings(self.backbone, hidden_states, text_position_ids) for decoder_layer in self.backbone.layers[: self.backbone.config.num_hidden_layers]: hidden_states = decoder_layer( hidden_states, @@ -5521,7 +5528,7 @@ def _encode_hidden_states(self, inputs_embeds: torch.Tensor, position_ids: torch hidden_states = inputs_embeds text_position_ids = position_ids.to(dtype=torch.int64) - position_embeddings = self.backbone.rotary_emb(hidden_states, text_position_ids) + position_embeddings = _qwen35_rotary_embeddings(self.backbone, hidden_states, text_position_ids) for decoder_layer in self.backbone.layers[: self.backbone.config.num_hidden_layers]: hidden_states = decoder_layer( hidden_states, @@ -5604,10 +5611,12 @@ def _build_qwen_causal_lm_component_specs( chunk_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_media_step", "decoder_prefill_chunk", "decoder_embed_chunk"} wants_chunked = bool(chunk_components & requested_set) - if wants_chunked and family != "qwen3": + if wants_chunked and family not in {"qwen3", "qwen3_5"}: raise UnsupportedComponentsError( - f"text chunked-prefill components are only supported for the qwen3 family, got {family}" + f"text chunked-prefill components are only supported for the qwen3/qwen3_5 families, got {family}" ) + if wants_chunked and family == "qwen3_5" and "decoder_embed_chunk" in requested_set: + raise UnsupportedComponentsError("decoder_embed_chunk is not supported for qwen3_5") specs: list[ComponentModuleSpec] = [] if "decoder" in requested_set: @@ -5644,6 +5653,84 @@ def _build_qwen_causal_lm_component_specs( dynamic_batch_axis=0 if dynamic_batch else None, )) + if wants_chunked and family == "qwen3_5": + lm_encoder_step = Qwen35LMEncoderStepAdapter(model).eval() + lm_encoder_text_chunk = Qwen35LMEncoderTextChunkAdapter(model).eval() + decoder_media_step = Qwen35EmbedsCausalLMStepAdapter(model).eval() + decoder_prefill_chunk = Qwen35EmbedsCausalLMPrefillChunkAdapter(model).eval() + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) + prefill_chunk_size = max(1, int(os.environ.get("CACTUS_QWEN_PREFILL_CHUNK", "512") or "512")) + chunk_input_ids = _tile_to_length(input_ids, prefill_chunk_size) + chunk_position_ids = torch.arange( + prefill_chunk_size, + dtype=torch.long, + device=input_ids.device, + ).unsqueeze(0).expand(int(input_ids.shape[0]), -1).contiguous() + step_input_ids = input_ids[:, :1] + step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) + with torch.no_grad(): + chunk_embeds, chunk_pos_out = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) + step_embeds, step_pos_out = lm_encoder_step(step_input_ids, step_position_ids) + cache_meta = { + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "use_internal_gated_deltanet_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + } + if "lm_encoder_step" in requested_set: + specs.append(ComponentModuleSpec( + component="lm_encoder_step", + module=lm_encoder_step, + example_inputs=(step_input_ids, step_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("inputs_embeds", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, + metadata={"family": family, "task": "causal_lm_logits"}, + )) + if "lm_encoder_text_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="lm_encoder_text_chunk", + module=lm_encoder_text_chunk, + example_inputs=(chunk_input_ids, chunk_position_ids), + input_keys=("input_ids", "position_ids"), + output_keys=("inputs_embeds", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder_text_chunk"}, + metadata={"family": family, "task": "causal_lm_logits"}, + )) + if "decoder_prefill_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_prefill_chunk", + module=decoder_prefill_chunk, + example_inputs=(chunk_embeds, chunk_pos_out), + input_keys=("inputs_embeds", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_prefill_chunk", + **cache_meta, + "prefill_chunk_size": prefill_chunk_size, + }, + metadata={"family": family, "task": "causal_lm_logits"}, + )) + if "decoder_media_step" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_media_step", + module=decoder_media_step, + example_inputs=(step_embeds, step_pos_out), + input_keys=("inputs_embeds", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + "component": "decoder_media_step", + **cache_meta, + "cache_num_slots": (max_slots if dynamic_batch else 1), + }, + metadata={"family": family, "task": "causal_lm_logits"}, + dynamic_batch_axis=0 if dynamic_batch else None, + )) + return specs + if wants_chunked: lm_encoder_step = Qwen3LMEncoderStepAdapter(model).eval() lm_encoder_text_chunk = Qwen3LMEncoderTextChunkAdapter(model).eval() diff --git a/python/cactus/transpile/model_profiles.py b/python/cactus/transpile/model_profiles.py index 7d132182b..174dcd0f8 100644 --- a/python/cactus/transpile/model_profiles.py +++ b/python/cactus/transpile/model_profiles.py @@ -36,6 +36,8 @@ class ModelProfile: force_component_pipeline: bool = False aliases: tuple[tuple[str, str], ...] = () regex_aliases: tuple[tuple[str, str], ...] = () + default_skip_model_load: bool = False + default_low_memory_load: bool = False GEMMA4_PROFILE = ModelProfile( @@ -225,11 +227,33 @@ class ModelProfile: ) +BONSAI_PROFILE = ModelProfile( + family="bonsai", + model_id_markers=("bonsai",), + avoid_native_loader=True, + stop_tokens=("<|im_end|>",), + prompt_style="qwen_chat", + default_task="causal_lm_logits", + default_components=( + "decoder", + "decoder_step", + "decoder_media_step", + "decoder_prefill_chunk", + "lm_encoder_step", + "lm_encoder_text_chunk", + ), + force_component_pipeline=True, + default_skip_model_load=True, + default_low_memory_load=True, +) + + PROFILES: tuple[ModelProfile, ...] = ( GEMMA4_PROFILE, LFM2_PROFILE, PARAKEET_TDT_PROFILE, WHISPER_PROFILE, + BONSAI_PROFILE, QWEN_PROFILE, NEEDLE_PROFILE, )