diff --git a/cactus-engine/src/complete.cpp b/cactus-engine/src/complete.cpp index 84f3babbc..e3ff46097 100644 --- a/cactus-engine/src/complete.cpp +++ b/cactus-engine/src/complete.cpp @@ -1356,9 +1356,12 @@ int cactus_benchmark_tokens( sample_peak_ram(); auto first_token_time = std::chrono::high_resolution_clock::now(); + std::vector completion_tokens; + if (decode_token_len > 0) completion_tokens.push_back(next_token); size_t generated = decode_token_len > 0 ? 1 : 0; while (generated < decode_token_len) { next_token = handle->model->decode({next_token}, 0.0f, 1.0f, 1); + completion_tokens.push_back(next_token); sample_peak_ram(); ++generated; } @@ -1397,9 +1400,17 @@ int cactus_benchmark_tokens( << "\"first_token_from_prefill\":" << (first_token_from_prefill ? "true" : "false") << "," << "\"prefill_padding_tokens\":" << handle->model->last_prefill_padding_tokens() << "," << "\"prefill_scalar_tail_tokens\":" << handle->model->last_prefill_scalar_tail_tokens() << "," + << "\"prefill_tail_chunk_tokens\":" << handle->model->last_prefill_tail_chunk_tokens() << "," + << "\"prefill_tail_padding_tokens\":" << handle->model->last_prefill_tail_padding_tokens() << "," << "\"decode_tps\":" << std::fixed << std::setprecision(2) << decode_tps << "," << "\"prompt_tokens\":" << prompt_token_len << "," << "\"completion_tokens\":" << generated << "," + << "\"completion_token_ids\":["; + for (size_t i = 0; i < completion_tokens.size(); ++i) { + if (i) oss << ","; + oss << completion_tokens[i]; + } + oss << "]," << "\"peak_ram_usage_mb\":" << std::fixed << std::setprecision(2) << peak_ram_usage_mb << "," << "\"ram_usage_mb\":" << std::fixed << std::setprecision(2) << get_ram_usage_mb() << "}"; diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index e4d361e8b..aeed9d0e5 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -349,7 +349,7 @@ class Tokenizer { size_t get_image_soft_token_count() const { return image_soft_token_count_; } protected: - enum class ModelType { UNKNOWN, GEMMA4, QWEN, LFM2, NEEDLE }; + enum class ModelType { UNKNOWN, GEMMA4, GEMMA, QWEN, LFM2, NEEDLE }; ModelType model_type_ = ModelType::UNKNOWN; enum class ModelVariant { DEFAULT, VLM, EXTRACT, RAG}; ModelVariant model_variant_ = ModelVariant::DEFAULT; @@ -374,6 +374,7 @@ class Tokenizer { std::string format_qwen_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported = false) const; std::string format_lfm2_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json, bool enable_thinking_if_supported = false) const; std::string format_needle_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const; + std::string format_gemma_style(const std::vector& messages, bool add_generation_prompt, const std::string& tools_json) const; }; class BPETokenizer : public Tokenizer { @@ -684,6 +685,8 @@ class Model { double last_prefill_cache_copy_ms() const { return last_prefill_cache_copy_ms_; } size_t last_prefill_padding_tokens() const { return last_prefill_padding_tokens_; } size_t last_prefill_scalar_tail_tokens() const { return last_prefill_scalar_tail_tokens_; } + size_t last_prefill_tail_chunk_tokens() const { return last_prefill_tail_chunk_tokens_; } + size_t last_prefill_tail_padding_tokens() const { return last_prefill_tail_padding_tokens_; } bool load_npu_audio_encoder(const std::string& model_path); bool has_npu_audio_encoder() const { return npu_audio_encoder_ != nullptr; } @@ -781,10 +784,14 @@ class Model { void move_cache_states(Component& source, Component& target, size_t logical_current = std::numeric_limits::max()); void set_cache_current_len(Component& comp, size_t len); void reset_component_cache_states(Component& comp); + void reset_prefill_stats(); size_t component_chunk_tokens(const Component& comp, const std::string& input_name) const; size_t component_output_tokens(const Component& comp, const std::string& output_name) const; ChunkedPrefillResult run_chunked_prefill(const std::vector& tokens, size_t start_position, 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); void run_full_context_text(); uint32_t argmax_component_logits(Component& comp, size_t logit_row = std::numeric_limits::max(), float* out_uncertainty = nullptr); @@ -829,6 +836,7 @@ class Model { std::string encoder_cross_kv_source_kind_; bool encoder_cross_kv_ready_ = false; size_t encoder_cross_kv_source_len_ = 0; + bool prefill_tail_pad_disabled_ = false; std::string family_; std::string npu_audio_encoder_mlpackage_; @@ -860,6 +868,8 @@ class Model { double last_prefill_cache_copy_ms_ = 0.0; size_t last_prefill_padding_tokens_ = 0; size_t last_prefill_scalar_tail_tokens_ = 0; + size_t last_prefill_tail_chunk_tokens_ = 0; + size_t last_prefill_tail_padding_tokens_ = 0; std::vector context_tokens_; static constexpr size_t MAX_TOKEN_HISTORY = 128; diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 53a3fbb28..50e766c8b 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -556,6 +556,8 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, if (components_.count("decoder_prefill_chunk") && components_.at("decoder_prefill_chunk").graph) { decoder_prefill_ = &components_.at("decoder_prefill_chunk"); decoder_prefill_chunk_ = decoder_prefill_; + } else if (decode_route_ != DecodeRoute::ENCODER_CROSS_KV_STEP && !components_.count("audio_encoder")) { + CACTUS_LOG_WARN("model", "Bundle has no decoder_prefill_chunk component; prompts will prefill token-by-token (prefill speed ~= decode speed)."); } if (components_.count("decoder_embed_chunk") && components_.at("decoder_embed_chunk").graph) { decoder_embed_ = &components_.at("decoder_embed_chunk"); @@ -563,6 +565,9 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, if (components_.count("lm_encoder_text_chunk") && components_.at("lm_encoder_text_chunk").graph) { prefill_encoder_ = &components_.at("lm_encoder_text_chunk"); } + if (const char* env = std::getenv("CACTUS_DISABLE_PREFILL_TAIL_PAD")) { + prefill_tail_pad_disabled_ = std::atoi(env) != 0; + } vision_encoder_ = components_.count("vision_encoder") ? &components_.at("vision_encoder") : nullptr; audio_encoder_ = components_.count("audio_encoder") ? &components_.at("audio_encoder") : nullptr; lm_encoder_media_step_ = components_.count("lm_encoder_media_step") ? &components_.at("lm_encoder_media_step") : nullptr; @@ -1059,11 +1064,48 @@ size_t Model::component_output_tokens(const Component& comp, const std::string& return 0; } -Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector& tokens, size_t start_position, size_t chunk_size, bool prepare_decode) { - ChunkedPrefillResult result; +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) { + 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); + } + if (enc_comp && encoder_chunk > 0) { + for (size_t chunk_offset = 0; chunk_offset < chunk_tokens; chunk_offset += encoder_chunk) { + for (size_t i = 0; i < enc_comp->input_buffers.size(); ++i) { + std::fill(enc_comp->input_buffers[i].begin(), enc_comp->input_buffers[i].end(), 0); + } + for (size_t i = 0; i < encoder_chunk; ++i) { + size_t index = processed + chunk_offset + i; + uint32_t token = index < tokens.size() ? tokens[index] : static_cast(config_.pad_token_id); + write_int_input_at(*enc_comp, "input_ids", i, static_cast(token)); + write_int_input_at(*enc_comp, "position_ids", i, static_cast(start_position + processed + chunk_offset + i)); + } + enc_comp->graph->execute(); + copy_component_outputs_to_chunk_inputs_range(*enc_comp, chunk_comp, chunk_offset); + } + } else { + for (size_t i = 0; i < chunk_tokens; ++i) { + size_t index = processed + i; + uint32_t token = index < tokens.size() ? tokens[index] : static_cast(config_.pad_token_id); + run_encoder_step(token, start_position + processed + i); + copy_component_outputs_to_chunk_inputs(*encoder_, chunk_comp, i); + } + } + chunk_comp.graph->execute(); +} + +void Model::reset_prefill_stats() { last_prefill_cache_copy_ms_ = 0.0; last_prefill_padding_tokens_ = 0; last_prefill_scalar_tail_tokens_ = 0; + last_prefill_tail_chunk_tokens_ = 0; + last_prefill_tail_padding_tokens_ = 0; +} + +Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector& tokens, size_t start_position, size_t chunk_size, bool prepare_decode) { + ChunkedPrefillResult result; + reset_prefill_stats(); if (decode_route_ != DecodeRoute::CACHED_STEP || !encoder_ || !decoder_ || !decoder_prefill_) return result; if (start_position != 0) return result; if (!load_component_graph(*decoder_prefill_)) return result; @@ -1096,20 +1138,28 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector(1, effective_chunk / 16); - const bool pad_tail = family_ != "lfm2_vl" + const bool has_conv_state = any_cache_node([&](size_t id) { + return decoder_prefill_->graph->get_node_op_type(id) == OpType::CONV_CACHE_STATE; + }); + const bool pad_tail = !has_conv_state && !has_recurrent_state && !has_sliding_window_cache && tail_tokens >= padding_cutoff; + const bool padded_window_too_small = any_cache_node([&](size_t id) { + if (decoder_prefill_->graph->get_node_op_type(id) != OpType::KV_CACHE_STATE) return false; + size_t window = decoder_prefill_->graph->get_node_window_size(id); + size_t sink = decoder_prefill_->graph->get_node_sink_size(id); + return window > 0 && (window < effective_chunk * 4 || window <= effective_chunk + sink); + }); + const bool use_padded_tail = !pad_tail && !prefill_tail_pad_disabled_ + && has_sliding_window_cache && !has_recurrent_state && !has_conv_state + && tail_tokens > 8 && !padded_window_too_small; const size_t executable_tokens = whole_chunks_end + (pad_tail ? effective_chunk : 0); - if (executable_tokens == 0) { + if (executable_tokens == 0 && !use_padded_tail) { result.scalar_tail_tokens = tail_tokens; last_prefill_scalar_tail_tokens_ = tail_tokens; return result; } - result.padding_tokens = executable_tokens > tokens.size() ? executable_tokens - tokens.size() : 0; - result.scalar_tail_tokens = tokens.size() - std::min(tokens.size(), executable_tokens); - last_prefill_padding_tokens_ = result.padding_tokens; - last_prefill_scalar_tail_tokens_ = result.scalar_tail_tokens; size_t encoder_chunk = 0; if (prefill_encoder_ && input_index(*prefill_encoder_, "input_ids") >= 0 && input_index(*prefill_encoder_, "position_ids") >= 0) { @@ -1121,39 +1171,47 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vectorinput_buffers.size(); ++i) { - std::fill(decoder_prefill_->input_buffers[i].begin(), decoder_prefill_->input_buffers[i].end(), 0); - } - if (encoder_chunk > 0) { - for (size_t chunk_offset = 0; chunk_offset < effective_chunk; chunk_offset += encoder_chunk) { - for (size_t i = 0; i < prefill_encoder_->input_buffers.size(); ++i) { - std::fill(prefill_encoder_->input_buffers[i].begin(), prefill_encoder_->input_buffers[i].end(), 0); - } - for (size_t i = 0; i < encoder_chunk; ++i) { - size_t index = processed + chunk_offset + i; - uint32_t token = index < tokens.size() ? tokens[index] : static_cast(config_.pad_token_id); - write_int_input_at(*prefill_encoder_, "input_ids", i, static_cast(token)); - write_int_input_at(*prefill_encoder_, "position_ids", i, static_cast(start_position + processed + chunk_offset + i)); - } - prefill_encoder_->graph->execute(); - copy_component_outputs_to_chunk_inputs_range(*prefill_encoder_, *decoder_prefill_, chunk_offset); - } - } else { - for (size_t i = 0; i < effective_chunk; ++i) { - size_t index = processed + i; - uint32_t token = index < tokens.size() ? tokens[index] : static_cast(config_.pad_token_id); - run_encoder_step(token, start_position + processed + i); - copy_component_outputs_to_chunk_inputs(*encoder_, *decoder_prefill_, i); + execute_prefill_chunk(*decoder_prefill_, prefill_encoder_, encoder_chunk, + effective_chunk, tokens, processed, start_position); + processed += effective_chunk; + } + + size_t tail_executed = 0; + size_t tail_padding = 0; + if (use_padded_tail) { + const size_t pads = effective_chunk - tail_tokens; + const size_t kept_real = tail_tokens - 1; + std::vector>> backups; + for (const auto& state : decoder_prefill_->cache_states) { + for (int node_id : {state.key_node_id, state.value_node_id}) { + if (node_id < 0) continue; + size_t id = static_cast(node_id); + if (decoder_prefill_->graph->get_node_op_type(id) != OpType::KV_CACHE_STATE) continue; + backups.emplace_back(id, decoder_prefill_->graph->snapshot_cache_padded_append(id, kept_real, pads + 1)); } } - decoder_prefill_->graph->execute(); - processed += effective_chunk; + execute_prefill_chunk(*decoder_prefill_, prefill_encoder_, encoder_chunk, + effective_chunk, tokens, processed, start_position); + for (auto& [id, backup] : backups) { + decoder_prefill_->graph->rollback_cache_padded_append(id, kept_real, pads + 1, backup); + } + processed += kept_real; + tail_executed = kept_real; + tail_padding = pads; } + result.executed_tokens = processed; result.logical_tokens = std::min(tokens.size(), processed); if (result.logical_tokens > 0) { result.last_logit_row = (result.logical_tokens - 1) % effective_chunk; } + result.padding_tokens = processed > tokens.size() ? processed - tokens.size() : 0; + result.scalar_tail_tokens = tokens.size() - result.logical_tokens; + last_prefill_padding_tokens_ = result.padding_tokens; + last_prefill_scalar_tail_tokens_ = result.scalar_tail_tokens; + last_prefill_tail_chunk_tokens_ = tail_executed; + last_prefill_tail_padding_tokens_ = tail_padding; + if (processed > 0 && prepare_decode) { for (size_t i = 0; i < decoder_->input_buffers.size(); ++i) { std::fill(decoder_->input_buffers[i].begin(), decoder_->input_buffers[i].end(), 0); @@ -1531,9 +1589,7 @@ uint32_t Model::argmax_last_logits(float* out_uncertainty) { } bool Model::prefill_and_sample_first_token(const std::vector& tokens, uint32_t& out_token) { - last_prefill_cache_copy_ms_ = 0.0; - last_prefill_padding_tokens_ = 0; - last_prefill_scalar_tail_tokens_ = 0; + reset_prefill_stats(); if (tokens.empty() || !decoder_ || cache_total_seq_len_ != 0) { return false; } @@ -1610,9 +1666,7 @@ bool Model::prefill_and_sample_first_token(const std::vector& tokens, } void Model::prefill(const std::vector& tokens, size_t /*chunk_size*/, const std::string& /*profile_file*/, bool prepare_decode) { - last_prefill_cache_copy_ms_ = 0.0; - last_prefill_padding_tokens_ = 0; - last_prefill_scalar_tail_tokens_ = 0; + reset_prefill_stats(); if (decode_route_ == DecodeRoute::ENCODER_CROSS_KV_STEP && encoder_cross_kv_source_kind_ == "text_tokens") { (void)prepare_decode; prepare_encoder_cross_kv_from_text(tokens); @@ -3465,7 +3519,7 @@ bool Config::from_json(const std::string& config_path) { std::transform(mt.begin(), mt.end(), mt.begin(), ::tolower); if (mt == "qwen") model_type = ModelType::QWEN; else if (mt == "qwen3p5" || mt == "qwen3_5") model_type = ModelType::QWEN3P5; - else if (mt == "gemma") model_type = ModelType::GEMMA; + else if (mt == "gemma" || mt == "gemma3") model_type = ModelType::GEMMA; else if (mt == "gemma3n") model_type = ModelType::GEMMA3N; else if (mt == "lfm2") model_type = ModelType::LFM2; else if (mt == "whisper") model_type = ModelType::WHISPER; diff --git a/cactus-engine/src/tokenizer.cpp b/cactus-engine/src/tokenizer.cpp index a4c36fe00..746a15aa7 100644 --- a/cactus-engine/src/tokenizer.cpp +++ b/cactus-engine/src/tokenizer.cpp @@ -371,8 +371,10 @@ void Tokenizer::detect_model_type(const std::string& config_path) { model_type_ = ModelType::QWEN; } else if (lower_line.find("lfm2") != std::string::npos) { model_type_ = ModelType::LFM2; - } else if (lower_line.find("gemma") != std::string::npos) { + } else if (lower_line.find("gemma4") != std::string::npos) { model_type_ = ModelType::GEMMA4; + } else if (lower_line.find("gemma") != std::string::npos) { + model_type_ = ModelType::GEMMA; } } else if (lower_line.find("model_variant") != std::string::npos) { if (lower_line.find("vlm") != std::string::npos) { model_variant_ = ModelVariant::VLM; } @@ -404,6 +406,9 @@ std::string Tokenizer::get_default_stop_sequence() const { if (model_type_ == ModelType::QWEN || model_type_ == ModelType::LFM2) { return "<|im_end|>"; } + if (model_type_ == ModelType::GEMMA) { + return ""; + } return ""; } @@ -422,9 +427,38 @@ std::string Tokenizer::format_chat_prompt(const std::vector& messag if (model_type_ == ModelType::NEEDLE) { return format_needle_style(messages, add_generation_prompt, tools_json); } + if (model_type_ == ModelType::GEMMA) { + return format_gemma_style(messages, add_generation_prompt, tools_json); + } return format_gemma4_style(messages, add_generation_prompt, tools_json, enable_thinking_if_supported); } +std::string Tokenizer::format_gemma_style(const std::vector& messages, bool add_generation_prompt, + const std::string& tools_json) const { + std::string result = ""; + std::string pending_context = tools_json; + for (const auto& msg : messages) { + if (msg.role == "system" || msg.role == "developer") { + pending_context += msg.content + "\n\n"; + continue; + } + std::string role = msg.role == "assistant" ? "model" : "user"; + std::string content = msg.content; + for (const auto& tc : msg.tool_calls) { + content += "\ncall:" + tc.name + "(" + tc.arguments + ")"; + } + if (role == "user" && !pending_context.empty()) { + content = pending_context + content; + pending_context.clear(); + } + result += "" + role + "\n" + content + "\n"; + } + if (add_generation_prompt) { + result += "model\n"; + } + return result; +} + namespace { std::string strip_newlines(const std::string& s) { size_t a = s.find_first_not_of('\n'); diff --git a/cactus-engine/tests/test_llm.cpp b/cactus-engine/tests/test_llm.cpp index a7fe93cb0..fd71b308c 100644 --- a/cactus-engine/tests/test_llm.cpp +++ b/cactus-engine/tests/test_llm.cpp @@ -702,6 +702,93 @@ bool test_multiturn_thinking_persist() { return prefix_ok && mentions_alice; } +static std::string benchmark_tokens_json(cactus_model_t model, const std::vector& ids, size_t max_new) { + std::vector response(1 << 16, 0); + int rc = cactus_benchmark_tokens(model, ids.data(), ids.size(), max_new, response.data(), response.size()); + return rc < 0 ? std::string() : std::string(response.data()); +} + +static std::string completion_ids_field(const std::string& json) { + size_t start = json.find("\"completion_token_ids\":["); + if (start == std::string::npos) return std::string(); + size_t end = json.find(']', start); + return end == std::string::npos ? std::string() : json.substr(start, end - start + 1); +} + +static std::string first_completion_id(const std::string& json) { + std::string ids = completion_ids_field(json); + size_t open = ids.find('['); + if (open == std::string::npos) return std::string(); + size_t end = ids.find_first_of(",]", open + 1); + return end == std::string::npos ? std::string() : ids.substr(open + 1, end - open - 1); +} + +bool test_chunked_prefill_padding() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║" << std::setw(42) << std::left << " CHUNKED PREFILL PADDING TEST" << "║\n" + << "╚══════════════════════════════════════════╝\n"; + if (!g_model_path) { std::cout << " [WARN] CACTUS_TEST_MODEL not set; skipping\n"; return true; } + + for (size_t prompt_len : {size_t(95), size_t(600)}) { + std::vector ids(prompt_len); + for (size_t i = 0; i < ids.size(); i++) ids[i] = static_cast(100 + (i % 200)); + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { std::cerr << " [✗] model init failed\n"; return false; } + std::string padded = benchmark_tokens_json(model, ids, 4); + std::string padded_again = benchmark_tokens_json(model, ids, 4); + cactus_destroy(model); + if (padded.empty() || padded.find("\"success\":true") == std::string::npos) { + std::cerr << " [✗] padded benchmark failed at len " << prompt_len << "\n"; + return false; + } + long tail_chunk = static_cast(EngineTestUtils::json_number(padded, "prefill_tail_chunk_tokens", -1)); + long tail_pads = static_cast(EngineTestUtils::json_number(padded, "prefill_tail_padding_tokens", -1)); + long scalar = static_cast(EngineTestUtils::json_number(padded, "prefill_scalar_tail_tokens", -1)); + std::cout << " len " << prompt_len << ": tail_chunk=" << tail_chunk + << " pads=" << tail_pads << " scalar=" << scalar << "\n"; + if (tail_chunk <= 0) { + std::cout << " [WARN] padded tail did not engage (no sliding caches?); skipping\n"; + return true; + } + if (tail_pads <= 0 || scalar > 1) { + std::cerr << " [✗] unexpected padding telemetry\n"; + return false; + } + if (completion_ids_field(padded).empty() + || completion_ids_field(padded) != completion_ids_field(padded_again)) { + std::cerr << " [✗] padded prefill is not deterministic\n"; + return false; + } + + setenv("CACTUS_DISABLE_PREFILL_TAIL_PAD", "1", 1); + model = cactus_init(g_model_path, nullptr, false); + std::string scalar_run = model ? benchmark_tokens_json(model, ids, 4) : std::string(); + if (model) cactus_destroy(model); + unsetenv("CACTUS_DISABLE_PREFILL_TAIL_PAD"); + if (scalar_run.empty() || scalar_run.find("\"success\":true") == std::string::npos) { + std::cerr << " [✗] kill-switch benchmark failed\n"; + return false; + } + long off_tail = static_cast(EngineTestUtils::json_number(scalar_run, "prefill_tail_chunk_tokens", -1)); + long off_scalar = static_cast(EngineTestUtils::json_number(scalar_run, "prefill_scalar_tail_tokens", -1)); + const size_t chunk_size = static_cast(tail_chunk + tail_pads + 1); + if (off_tail != 0 || off_scalar != static_cast(prompt_len % chunk_size)) { + std::cerr << " [✗] kill switch did not restore the scalar tail (tail=" << off_tail + << " scalar=" << off_scalar << ")\n"; + return false; + } + if (first_completion_id(padded).empty() + || first_completion_id(padded) != first_completion_id(scalar_run)) { + std::cerr << " [✗] padded prefill diverged from the scalar tail\n" + << " padded: " << completion_ids_field(padded) << "\n" + << " scalar: " << completion_ids_field(scalar_run) << "\n"; + return false; + } + } + return true; +} + int main() { TestUtils::TestRunner runner("LLM Tests"); runner.run_test("1k_context", test_1k_context()); @@ -710,6 +797,7 @@ int main() { runner.run_test("prefill_idempotent_reuse", test_prefill_idempotent_reuse()); runner.run_test("prefill_prefix_extension_reuse", test_prefill_prefix_extension_reuse()); runner.run_test("prefill_invalidated_on_message_change", test_prefill_invalidated_on_message_change()); + runner.run_test("chunked_prefill_padding", test_chunked_prefill_padding()); runner.run_test("tool_calls", test_tool_call()); runner.run_test("tool_multiple_tool_call_invocations", test_multiple_tool_call_invocations()); runner.run_test("tool_calls_with_three_tools", test_tool_call_with_three_tools()); diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 9780cedff..68a3ae46d 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -708,8 +708,12 @@ class CactusGraph { const BufferDesc& get_output_buffer(size_t node_id) const; OpType get_node_op_type(size_t node_id) const; size_t get_node_window_size(size_t node_id) const; + size_t get_node_sink_size(size_t node_id) const; void steal_cache_buffer(size_t dst_node, CactusGraph& src, size_t src_node); void shrink_cache_buffer(size_t node_id, size_t new_capacity); + std::vector snapshot_cache_padded_append(size_t node_id, size_t real_tokens, size_t pad_tokens) const; + void rollback_cache_padded_append(size_t node_id, size_t real_tokens, size_t pad_tokens, + const std::vector& backup); void allocate_buffers(); size_t get_node_count() const; diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp index 0ae9abdfc..35b9b1d3f 100644 --- a/cactus-graph/src/builder.cpp +++ b/cactus-graph/src/builder.cpp @@ -1247,6 +1247,10 @@ size_t CactusGraph::get_node_window_size(size_t node_id) const { return nodes_[node_index_map_.at(node_id)]->params.window_size; } +size_t CactusGraph::get_node_sink_size(size_t node_id) const { + return nodes_[node_index_map_.at(node_id)]->params.cache_sink_size; +} + size_t CactusGraph::persistent(size_t source_node) { const auto& source_buffer = get_output_buffer(source_node); OpParams params; diff --git a/cactus-graph/src/ops_cache.cpp b/cactus-graph/src/ops_cache.cpp index 2ed462a23..24c0998dc 100644 --- a/cactus-graph/src/ops_cache.cpp +++ b/cactus-graph/src/ops_cache.cpp @@ -5,6 +5,7 @@ #include #include #include +#include namespace { @@ -532,6 +533,98 @@ void CactusGraph::steal_cache_buffer(size_t dst_node, CactusGraph& src, size_t s dst->output_buffer = std::move(s->output_buffer); } +namespace { + +struct PaddedAppendBackup { + uint64_t overshoot; + uint64_t keep_sink; + uint64_t kept_padded; +}; + +struct CacheRowRegion { + size_t offset; + size_t row_bytes; +}; + +std::vector cache_row_regions(const BufferDesc& buf, const CacheMetadata* meta) { + const auto* base = static_cast(buf.get_data()); + const size_t stride = meta->num_kv_heads * meta->head_dim; + if (buf.precision == Precision::FP16) { + return {{static_cast(reinterpret_cast(get_fp16_data(buf)) - base), stride * sizeof(__fp16)}}; + } + const size_t num_groups = (meta->head_dim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; + return { + {static_cast(reinterpret_cast(get_int8_data(buf)) - base), stride}, + {static_cast(reinterpret_cast(get_scales(buf, meta->max_seq_len, meta->num_kv_heads, meta->head_dim)) - base), + meta->num_kv_heads * num_groups * sizeof(float)}, + }; +} + +} // namespace + +std::vector CactusGraph::snapshot_cache_padded_append(size_t node_id, size_t real_tokens, size_t pad_tokens) const { + const auto& node = *nodes_[node_index_map_.at(node_id)]; + const auto& buf = node.output_buffer; + if (node.op_type != OpType::KV_CACHE_STATE || !buf.get_data() || pad_tokens == 0) return {}; + const size_t window = node.params.window_size; + const size_t ceiling = node.params.max_cache_seq_len; + if (window == 0 || window >= ceiling) return {}; + const auto* meta = get_meta(buf); + const size_t len0 = meta->current_seq_len; + const size_t appended = real_tokens + pad_tokens; + if (len0 == 0 || len0 + appended <= window) return {}; + const size_t keep_sink = std::min({static_cast(meta->sink_size), len0, window}); + const size_t tail_capacity = window - keep_sink; + if (appended >= tail_capacity) { + throw std::runtime_error("padded cache append larger than the attention window is not supported"); + } + auto kept_after = [&](size_t n) { return std::min(tail_capacity - n, len0 - keep_sink); }; + const size_t kept_real = kept_after(real_tokens); + const size_t kept_padded = kept_after(appended); + const size_t overshoot = kept_real - kept_padded; + if (overshoot == 0) return {}; + + const size_t first_saved_row = len0 - kept_real; + std::vector backup(sizeof(PaddedAppendBackup)); + auto* header = reinterpret_cast(backup.data()); + header->overshoot = overshoot; + header->keep_sink = keep_sink; + header->kept_padded = kept_padded; + const auto* base = static_cast(buf.get_data()); + for (const auto& region : cache_row_regions(buf, meta)) { + const uint8_t* rows = base + region.offset + first_saved_row * region.row_bytes; + backup.insert(backup.end(), rows, rows + overshoot * region.row_bytes); + } + return backup; +} + +void CactusGraph::rollback_cache_padded_append(size_t node_id, size_t real_tokens, size_t pad_tokens, + const std::vector& backup) { + auto& node = *nodes_[node_index_map_.at(node_id)]; + auto& buf = node.output_buffer; + if (node.op_type != OpType::KV_CACHE_STATE || !buf.get_data() || pad_tokens == 0) return; + auto* meta = get_meta(buf); + if (backup.empty()) { + meta->current_seq_len = meta->current_seq_len >= pad_tokens ? meta->current_seq_len - pad_tokens : 0; + return; + } + const auto* header = reinterpret_cast(backup.data()); + const size_t overshoot = header->overshoot; + const size_t keep_sink = header->keep_sink; + const size_t kept_padded = header->kept_padded; + const size_t move_rows = kept_padded + real_tokens; + auto* base = static_cast(buf.get_data()); + const uint8_t* saved = backup.data() + sizeof(PaddedAppendBackup); + for (const auto& region : cache_row_regions(buf, meta)) { + uint8_t* rows = base + region.offset; + std::memmove(rows + (keep_sink + overshoot) * region.row_bytes, + rows + keep_sink * region.row_bytes, move_rows * region.row_bytes); + std::memcpy(rows + keep_sink * region.row_bytes, saved, overshoot * region.row_bytes); + saved += overshoot * region.row_bytes; + } + meta->current_seq_len = keep_sink + overshoot + kept_padded + real_tokens; +} + void CactusGraph::shrink_cache_buffer(size_t node_id, size_t new_capacity) { auto& buf = nodes_[node_index_map_.at(node_id)]->output_buffer; if (!buf.get_data()) return; diff --git a/cactus-graph/tests/test_cache.cpp b/cactus-graph/tests/test_cache.cpp index 07429b8c2..6ed754dd9 100644 --- a/cactus-graph/tests/test_cache.cpp +++ b/cactus-graph/tests/test_cache.cpp @@ -189,6 +189,77 @@ bool test_kv_cache_append_full_window_eviction() { return true; } +static bool padded_rollback_matches_exact_append(size_t prefix, size_t real, size_t pads) { + const size_t kv_heads = 1, head_dim = 16, max_seq = 64, window = 16, sink = 2; + const size_t stride = kv_heads * head_dim; + + auto append = [&](CactusGraph& g, size_t cache_node, size_t real_tokens, size_t pad_tokens, float base) { + const size_t tokens = real_tokens + pad_tokens; + const size_t elements = tokens * stride; + size_t kv_input = g.input({elements}, Precision::FP16); + std::vector<__fp16> data(elements); + for (size_t t = 0; t < tokens; t++) { + float value = t < real_tokens ? base + static_cast(t) : 9999.0f; + for (size_t j = 0; j < stride; j++) data[t * stride + j] = static_cast<__fp16>(value); + } + g.set_input(kv_input, data.data(), Precision::FP16); + g.kv_cache_append(kv_input, cache_node, window, sink); + g.execute(); + g.soft_reset(); + }; + + CactusGraph g_exact, g_padded; + size_t exact_node = g_exact.kv_cache_state(max_seq, kv_heads, head_dim, window, sink); + size_t padded_node = g_padded.kv_cache_state(max_seq, kv_heads, head_dim, window, sink); + if (prefix > 0) { + append(g_exact, exact_node, prefix, 0, 1.0f); + append(g_padded, padded_node, prefix, 0, 1.0f); + } + append(g_exact, exact_node, real, 0, 100.0f); + auto backup = g_padded.snapshot_cache_padded_append(padded_node, real, pads); + append(g_padded, padded_node, real, pads, 100.0f); + g_padded.rollback_cache_padded_append(padded_node, real, pads, backup); + + auto* exact_raw = static_cast(g_exact.get_output(exact_node)); + auto* padded_raw = static_cast(g_padded.get_output(padded_node)); + uint64_t exact_len = *reinterpret_cast(exact_raw); + uint64_t padded_len = *reinterpret_cast(padded_raw); + if (exact_len != padded_len) { + std::cerr << " current_seq_len mismatch: " << exact_len << " != " << padded_len << "\n"; + return false; + } + uint64_t buffer_max_seq = *reinterpret_cast(exact_raw + 8); + const size_t meta_bytes = 64; + if (std::memcmp(exact_raw + meta_bytes, padded_raw + meta_bytes, exact_len * stride) != 0) { + std::cerr << " cache data rows differ after rollback\n"; + return false; + } + const size_t num_groups = (head_dim + KV_QUANT_GROUP_SIZE - 1) / KV_QUANT_GROUP_SIZE; + const size_t scales_offset = meta_bytes + buffer_max_seq * stride; + if (std::memcmp(exact_raw + scales_offset, padded_raw + scales_offset, + exact_len * kv_heads * num_groups * sizeof(float)) != 0) { + std::cerr << " cache scale rows differ after rollback\n"; + return false; + } + return true; +} + +bool test_padded_rollback_no_eviction() { + return padded_rollback_matches_exact_append(4, 3, 5); +} + +bool test_padded_rollback_eviction_overshoot() { + return padded_rollback_matches_exact_append(16, 3, 6); +} + +bool test_padded_rollback_only_pads_evict() { + return padded_rollback_matches_exact_append(10, 2, 8); +} + +bool test_padded_rollback_empty_cache() { + return padded_rollback_matches_exact_append(0, 5, 11); +} + bool test_attention_cached_basic() { const size_t b = 1, s = 1, h = 2, kv = 2, d = 16; const size_t max_seq = 64; @@ -848,6 +919,10 @@ int main() { runner.run_test("KV Cache Append Multiple", test_kv_cache_append_multiple()); runner.run_test("KV Cache Append Eviction", test_kv_cache_append_eviction()); runner.run_test("KV Cache Append Full Window Eviction", test_kv_cache_append_full_window_eviction()); + runner.run_test("Padded Rollback No Eviction", test_padded_rollback_no_eviction()); + runner.run_test("Padded Rollback Eviction Overshoot", test_padded_rollback_eviction_overshoot()); + runner.run_test("Padded Rollback Only Pads Evict", test_padded_rollback_only_pads_evict()); + runner.run_test("Padded Rollback Empty Cache", test_padded_rollback_empty_cache()); runner.run_test("Attention Cached Basic", test_attention_cached_basic()); runner.run_test("Attention Cached Multistep", test_attention_cached_multistep()); runner.run_test("KV Cache Invalidate", test_kv_cache_invalidate()); diff --git a/python/cactus/convert/cactus_adapters/tensor_io.py b/python/cactus/convert/cactus_adapters/tensor_io.py index 3ce614cc7..5afb0df6f 100644 --- a/python/cactus/convert/cactus_adapters/tensor_io.py +++ b/python/cactus/convert/cactus_adapters/tensor_io.py @@ -3,6 +3,7 @@ from pathlib import Path from typing import Optional, Dict, Any, List from .weight_patterns import EMBED_NAMES +from ..model_adapters.naming import GEMMA4_WEIGHT_SCALE try: import torch @@ -270,8 +271,18 @@ def save_tensor_with_header(tensor, output_path, precision='INT8', transpose=Fal if model_type == 'gemma' and 'norm' in str(output_path): data = data + 1.0 + if model_type == 'gemma3': + filename = output_path.name + if any(x in filename for x in ['input_norm', 'post_attn_norm', 'pre_ffn_norm', 'post_ffn_norm']): + data = (data + 1.0) / GEMMA4_WEIGHT_SCALE - 1.0 + elif filename == 'output_norm.weights': + data = (data + 1.0) * GEMMA4_WEIGHT_SCALE - 1.0 + elif any(x in filename for x in ['ffn_gate', 'ffn_up']): + data = data * GEMMA4_WEIGHT_SCALE + elif filename == 'token_embeddings.weights': + data = data / GEMMA4_WEIGHT_SCALE + if model_type == 'gemma4': - GEMMA4_WEIGHT_SCALE = 16.0 filename = output_path.name is_audio_weight = filename.startswith('audio_') if any(x in filename for x in ['input_norm', 'post_attn_norm', 'pre_ffn_norm', 'post_ffn_norm', diff --git a/python/cactus/convert/model_adapters/adapters.py b/python/cactus/convert/model_adapters/adapters.py index 61ee9a9b2..a29592d6a 100644 --- a/python/cactus/convert/model_adapters/adapters.py +++ b/python/cactus/convert/model_adapters/adapters.py @@ -21,7 +21,7 @@ extract_vision_config, extract_whisper_config, ) -from .naming import NameMatch, cactus_name_for_tensor, gemma4_scale_factor, restore_hf_key_for_family +from .naming import NameMatch, cactus_name_for_tensor, gemma3_scale_factor, gemma4_scale_factor, restore_hf_key_for_family from .policy import TensorPolicy, policy_for_tensor from ..compat import patch_transformers_import_compat @@ -495,8 +495,16 @@ def module_target_name(self, source_name: str, _model: Any) -> str | None: return super().module_target_name(source_name, _model) +class Gemma3Adapter(FamilyAdapter): + family = "gemma3" + + def scale_factor(self, output_name: str) -> float: + return gemma3_scale_factor(output_name) + + ADAPTERS: dict[str, FamilyAdapter] = { "generic": FamilyAdapter(), + "gemma3": Gemma3Adapter(), "gemma4": Gemma4Adapter(), "qwen": QwenAdapter(), "lfm2": Lfm2Adapter(), diff --git a/python/cactus/convert/model_adapters/detection.py b/python/cactus/convert/model_adapters/detection.py index 7a72a4edf..69cb2aa09 100644 --- a/python/cactus/convert/model_adapters/detection.py +++ b/python/cactus/convert/model_adapters/detection.py @@ -3,7 +3,7 @@ from ..cactus_adapters.config_utils import cfg_get, detect_model_type -SUPPORTED_FAMILIES = {"auto", "gemma4", "qwen", "lfm2", "whisper", "parakeet", "parakeet_tdt", "moonshine", "nomic", "needle", "generic"} +SUPPORTED_FAMILIES = {"auto", "gemma3", "gemma4", "qwen", "lfm2", "whisper", "parakeet", "parakeet_tdt", "moonshine", "nomic", "needle", "generic"} def detect_family(config, requested: str = "auto") -> str: @@ -12,6 +12,11 @@ def detect_family(config, requested: str = "auto") -> str: text_config = cfg_get(config, "text_config", None) base = text_config if text_config is not None else config detected = detect_model_type(base, config) + if detected == "gemma": + raw = str(cfg_get(base, "model_type", cfg_get(config, "model_type", "")) or "").lower() + if "gemma3" in raw: + return "gemma3" + return "generic" if detected in {"gemma4", "qwen", "lfm2", "whisper", "moonshine", "parakeet", "parakeet_tdt", "nomic", "needle"}: return detected return "generic" diff --git a/python/cactus/convert/model_adapters/naming.py b/python/cactus/convert/model_adapters/naming.py index b85844152..d854b1a02 100644 --- a/python/cactus/convert/model_adapters/naming.py +++ b/python/cactus/convert/model_adapters/naming.py @@ -81,11 +81,25 @@ def restore_hf_key_for_family(key: str, family: str) -> str: return key -def gemma4_scale_factor(out_name: str) -> float: +def _scale_basename(out_name: str) -> str: base = out_name.removesuffix(".weights").removesuffix(".bias") parts = base.split("_", 2) if len(parts) == 3 and parts[0] == "layer" and parts[1].isdigit(): base = parts[2] + return base + + +def gemma3_scale_factor(out_name: str) -> float: + base = _scale_basename(out_name) + if base in {"ffn_gate", "ffn_up"}: + return GEMMA4_WEIGHT_SCALE + if base == "token_embeddings": + return 1.0 / GEMMA4_WEIGHT_SCALE + return 1.0 + + +def gemma4_scale_factor(out_name: str) -> float: + base = _scale_basename(out_name) if base in GEMMA4_MULT_BASENAMES: return GEMMA4_WEIGHT_SCALE if base in GEMMA4_DIV_BASENAMES: diff --git a/python/cactus/transpile/component_plan.py b/python/cactus/transpile/component_plan.py index 5c5df4114..c9e63cabd 100644 --- a/python/cactus/transpile/component_plan.py +++ b/python/cactus/transpile/component_plan.py @@ -30,13 +30,13 @@ def _plan_from_profile( has_vision = _has_dict_config(config, "vision_config", "visual_config", "image_config") if not has_vision: architectures = [str(a).lower() for a in (config.get("architectures") or ())] - if any("qwen3forcausallm" in a for a in architectures): - return ComponentPlan( - task="causal_lm_logits", - components=("decoder_step", "lm_encoder_step", "lm_encoder_text_chunk", - "decoder_prefill_chunk", "decoder_embed_chunk", "decoder_media_step"), - force_component_pipeline=True, - ) + for arch_marker, components in profile.text_component_plans: + if any(arch_marker in a for a in architectures): + return ComponentPlan( + task="causal_lm_logits", + components=components, + force_component_pipeline=True, + ) return ComponentPlan( task="causal_lm_logits", components=("decoder", "decoder_step"), diff --git a/python/cactus/transpile/hf_model.py b/python/cactus/transpile/hf_model.py index 62693019c..6fea22f23 100644 --- a/python/cactus/transpile/hf_model.py +++ b/python/cactus/transpile/hf_model.py @@ -47,6 +47,7 @@ from cactus.transpile.lower import _lower_ir_node from cactus.transpile.lower import _lookup_weight_binding from cactus.transpile.media_limits import resize_static_image +from cactus.transpile.model_adapters import UnsupportedComponentsError from cactus.transpile.model_adapters import build_component_module_specs from cactus.transpile.model_adapters import canonicalize_model_interface from cactus.transpile.model_profiles import multimodal_context_tokens_for_model_type @@ -3324,15 +3325,35 @@ def main() -> int: for component in str(args.components).split(",") if component.strip() ) - component_specs = build_component_module_specs( - model, + plan_derived_components = False + if requested_components is None and task == "causal_lm_logits": + inferred_plan = infer_component_plan_from_config( + _load_config_json(args.model_id), model_id=args.model_id + ) + if ( + inferred_plan is not None + and inferred_plan.task == task + and inferred_plan.force_component_pipeline + and inferred_plan.components + ): + requested_components = tuple(inferred_plan.components) + plan_derived_components = True + spec_kwargs = dict( task=task, named_tensors=_named_tensor_store(prepared), weights_dir=weights_dir, inputs_metadata=prepared.metadata, - components=requested_components, cache_context_length=args.cache_context_length, ) + try: + component_specs = build_component_module_specs( + model, components=requested_components, **spec_kwargs + ) + except UnsupportedComponentsError as exc: + if not plan_derived_components: + raise + print(f"warning: dropping inferred component plan ({exc}); using builder defaults", file=sys.stderr) + component_specs = build_component_module_specs(model, components=None, **spec_kwargs) use_component_pipeline = False if args.component_pipeline == "on": if component_specs is None: diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index a2ae0a36a..38eef3b3d 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -29,6 +29,10 @@ _GEMMA4_CREATE_BIDIRECTIONAL_MASK = None +class UnsupportedComponentsError(RuntimeError): + pass + + @dataclass class CanonicalizedModel: module: torch.nn.Module @@ -212,6 +216,11 @@ def _select_last_non_pad_token( return _select_last_active_token(hidden_or_logits, attention_mask) +def _tile_to_length(tensor: torch.Tensor, length: int) -> torch.Tensor: + reps = -(-length // int(tensor.shape[1])) + return torch.cat([tensor] * reps, dim=1)[:, :length].contiguous() + + def _resolve_model_pad_token_id(model: torch.nn.Module) -> int | None: config = getattr(model, "config", None) for attr_name in ("pad_token_id", "eos_token_id", "bos_token_id"): @@ -2054,6 +2063,44 @@ def get_transpile_metadata(self): } +def _gemma3_text_backbone_forward(backbone, inputs_embeds, causal_mask_mapping, position_ids, + checkpoints=None): + hidden_states = inputs_embeds + layer_types = backbone.config.layer_types + position_embeddings = { + layer_type: backbone.rotary_emb(hidden_states, position_ids, layer_type) + for layer_type in set(layer_types) + } + for i, decoder_layer in enumerate(backbone.layers[: backbone.config.num_hidden_layers]): + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask_mapping[layer_types[i]], + position_embeddings=position_embeddings[layer_types[i]], + position_ids=position_ids, + past_key_values=None, + ) + if checkpoints is not None: + checkpoints.append(hidden_states) + return backbone.norm(hidden_states) + + +def _gemma3_decoder_graph_meta(model, backbone, adapter_type, input_names): + sliding_window = getattr(backbone.config, "sliding_window", None) + return { + "graph": { + **_transpile_graph_meta( + model, + adapter_family="gemma3", + adapter_type=adapter_type, + input_names=input_names, + ), + "num_hidden_layers": int(backbone.config.num_hidden_layers), + "layer_types": tuple(getattr(backbone.config, "layer_types", [])), + "sliding_window": None if sliding_window is None else int(sliding_window), + } + } + + class Gemma3CausalLMLogitsAdapter(torch.nn.Module): def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): super().__init__() @@ -2091,45 +2138,97 @@ def debug_forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, list[tor ), } - hidden_states = inputs_embeds checkpoints: list[torch.Tensor] = [] - position_embeddings = { - layer_type: self.backbone.rotary_emb(hidden_states, position_ids, layer_type) - for layer_type in self.backbone.config.layer_types - } - - for i, decoder_layer in enumerate(self.backbone.layers[: self.backbone.config.num_hidden_layers]): - hidden_states = decoder_layer( - hidden_states, - attention_mask=causal_mask_mapping[self.backbone.config.layer_types[i]], - position_embeddings=position_embeddings[self.backbone.config.layer_types[i]], - position_ids=position_ids, - past_key_values=None, - ) - checkpoints.append(hidden_states) - - hidden_states = self.backbone.norm(hidden_states) + hidden_states = _gemma3_text_backbone_forward( + self.backbone, inputs_embeds, causal_mask_mapping, position_ids, checkpoints=checkpoints, + ) checkpoints.append(hidden_states) return self.model.lm_head(hidden_states), checkpoints def get_transpile_metadata(self): - sliding_window = getattr(self.backbone.config, "sliding_window", None) - layer_types = list(getattr(self.backbone.config, "layer_types", [])) + return _gemma3_decoder_graph_meta(self.model, self.backbone, type(self).__name__, ("input_ids",)) + + +class Gemma3LMEncoderStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.backbone = model.model + + def forward(self, input_ids: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + inputs_embeds = self.backbone.embed_tokens(input_ids) + return inputs_embeds.to(torch.float16), position_ids.to(dtype=torch.int64) + + def get_transpile_metadata(self): return { "graph": { **_transpile_graph_meta( - self.model, + self.backbone, adapter_family="gemma3", adapter_type=type(self).__name__, - input_names=("input_ids",), + input_names=("input_ids", "position_ids"), ), - "num_hidden_layers": int(self.backbone.config.num_hidden_layers), - "layer_types": tuple(layer_types), - "sliding_window": None if sliding_window is None else int(sliding_window), } } +class Gemma3EmbedsCausalLMStepAdapter(torch.nn.Module): + def __init__(self, model: torch.nn.Module): + super().__init__() + self.model = model + self.backbone = model.model + + def _additive_mask_mapping(self, inputs_embeds: torch.Tensor) -> dict[str, torch.Tensor]: + seq_len = int(inputs_embeds.shape[1]) + zero_mask = torch.zeros( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + return {layer_type: zero_mask for layer_type in set(self.backbone.config.layer_types)} + + def forward(self, inputs_embeds: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + hidden_states = inputs_embeds.to(self.backbone.norm.weight.dtype) + causal_mask_mapping = self._additive_mask_mapping(hidden_states) + text_position_ids = position_ids.to(dtype=torch.int64) + hidden_states = _gemma3_text_backbone_forward( + self.backbone, hidden_states, causal_mask_mapping, text_position_ids, + ) + logits = self.model.lm_head(hidden_states[:, -1:, :]) + return _gemma4_apply_final_logit_softcapping(self.model, logits) + + def get_transpile_metadata(self): + return _gemma3_decoder_graph_meta(self.model, self.backbone, type(self).__name__, ("inputs_embeds", "position_ids")) + + +class Gemma3EmbedsCausalLMPrefillChunkAdapter(Gemma3EmbedsCausalLMStepAdapter): + def _additive_mask_mapping(self, inputs_embeds: torch.Tensor) -> dict[str, torch.Tensor]: + seq_len = int(inputs_embeds.shape[1]) + blocked_values = torch.ones( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) * torch.finfo(inputs_embeds.dtype).min + allowed_values = torch.zeros( + (1, 1, seq_len, seq_len), + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + causal = torch.tril( + torch.ones((seq_len, seq_len), dtype=torch.bool, device=inputs_embeds.device), + ).view(1, 1, seq_len, seq_len) + mapping = {"full_attention": torch.where(causal, allowed_values, blocked_values)} + if "sliding_attention" in set(self.backbone.config.layer_types): + window = int(self.backbone.config.sliding_window) + below_window = torch.tril( + torch.ones((seq_len, seq_len), dtype=torch.bool, device=inputs_embeds.device), + diagonal=-window, + ).view(1, 1, seq_len, seq_len) + mapping["sliding_attention"] = torch.where( + below_window, blocked_values, mapping["full_attention"], + ) + return mapping + + class Gemma4CausalLMLogitsAdapter(torch.nn.Module): def __init__(self, model: torch.nn.Module, *, pad_token_id: int | None = None): super().__init__() @@ -3382,6 +3481,128 @@ def forward( ) +def _build_gemma3_causal_lm_component_specs( + model: torch.nn.Module, + *, + named_tensors: dict[str, torch.Tensor], + weights_dir: str | None = None, + components: tuple[str, ...] | None = None, + cache_context_length: str | int | None = None, +) -> list[ComponentModuleSpec] | None: + input_ids = named_tensors.get("input_ids") + if input_ids is None: + return None + pad_token_id = _resolve_model_pad_token_id(model) + requested = tuple(components or ( + "decoder", + "lm_encoder_step", + "lm_encoder_text_chunk", + "decoder_prefill_chunk", + "decoder_step", + )) + requested_set = set(requested) + common_graph_meta = { + "weights_dir": weights_dir, + "task": "causal_lm_logits", + "adapter_family": "gemma3", + } + metadata = {"family": "gemma3", "task": "causal_lm_logits"} + + specs: list[ComponentModuleSpec] = [] + if "decoder" in requested_set: + decoder = Gemma3CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() + specs.append(ComponentModuleSpec( + component="decoder", + module=decoder, + example_inputs=(input_ids,), + input_keys=("input_ids",), + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "decoder"}, + metadata=metadata, + )) + + cached_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_prefill_chunk", "decoder_step"} + if not (cached_components & requested_set): + return specs + + lm_encoder_step = Gemma3LMEncoderStepAdapter(model).eval() + lm_encoder_text_chunk = Gemma3LMEncoderStepAdapter(model).eval() + decoder_step = Gemma3EmbedsCausalLMStepAdapter(model).eval() + decoder_prefill_chunk = Gemma3EmbedsCausalLMPrefillChunkAdapter(model).eval() + max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) + cached_graph_meta = { + "use_internal_kv_cache": True, + "max_cache_seq_len": max_cache_seq_len, + "cache_sink_size": 4, + } + prefill_chunk_size = max(2, int(os.environ.get("CACTUS_GEMMA3_PREFILL_CHUNK", "128") or "128")) + + step_input_ids = input_ids[:, :1] + step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) + chunk_input_ids = _tile_to_length(input_ids, prefill_chunk_size) + chunk_position_ids = torch.arange( + prefill_chunk_size, dtype=torch.int64, device=input_ids.device, + ).unsqueeze(0) + with torch.no_grad(): + step_embeds, step_pos_out = lm_encoder_step(step_input_ids, step_position_ids) + chunk_embeds, chunk_pos_out = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) + + 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=metadata, + )) + 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", + "encoder_chunk_size": prefill_chunk_size, + }, + metadata=metadata, + )) + 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, + **cached_graph_meta, + "component": "decoder_prefill_chunk", + "prefill_chunk_size": prefill_chunk_size, + }, + metadata=metadata, + )) + if "decoder_step" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_step", + module=decoder_step, + example_inputs=(step_embeds, step_pos_out), + input_keys=("inputs_embeds", "position_ids"), + output_keys=("logits",), + graph_meta={ + **common_graph_meta, + **cached_graph_meta, + "component": "decoder_step", + }, + metadata=metadata, + )) + return specs + + def _build_gemma4_multimodal_component_specs( model: torch.nn.Module, *, @@ -4485,7 +4706,7 @@ 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": - raise RuntimeError( + raise UnsupportedComponentsError( f"text chunked-prefill components are only supported for the qwen3 family, got {family}" ) @@ -4530,8 +4751,7 @@ def _build_qwen_causal_lm_component_specs( decoder_embed_chunk = Qwen3EmbedsCausalLMEmbedChunkAdapter(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", "128") or "128")) - prefill_chunk_size = min(prefill_chunk_size, int(input_ids.shape[1])) - chunk_input_ids = input_ids[:, :prefill_chunk_size].contiguous() + chunk_input_ids = _tile_to_length(input_ids, prefill_chunk_size) chunk_position_ids = torch.arange( prefill_chunk_size, dtype=torch.long, @@ -4699,8 +4919,7 @@ def _require(component: str) -> None: image_features = vision_encoder(pixel_values) decoder_inputs = lm_encoder(input_ids, attention_mask, position_ids, image_features) prefill_chunk_size = max(1, int(os.environ.get("CACTUS_QWEN_PREFILL_CHUNK", "128") or "128")) - prefill_chunk_size = min(prefill_chunk_size, int(input_ids.shape[1])) - chunk_input_ids = input_ids[:, :prefill_chunk_size].contiguous() + chunk_input_ids = _tile_to_length(input_ids, prefill_chunk_size) chunk_position_ids = torch.arange( prefill_chunk_size, dtype=torch.long, @@ -5560,10 +5779,6 @@ def _build_lfm2_causal_lm_component_specs( input_ids = named_tensors.get("input_ids") if input_ids is None: return None - pad_token_id = _resolve_model_pad_token_id(model) - decoder = Lfm2CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() - decoder_step = Lfm2CausalLMStepAdapter(model, pad_token_id=pad_token_id).eval() - requested = tuple(components or ("decoder", "decoder_step")) requested_set = set(requested) common_graph_meta = { @@ -5571,6 +5786,22 @@ def _build_lfm2_causal_lm_component_specs( "task": "causal_lm_logits", "adapter_family": "lfm2", } + metadata = {"family": "lfm2", "task": "causal_lm_logits"} + chunk_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_prefill_chunk"} + if chunk_components & requested_set: + return _lfm2_chunked_pipeline_specs( + model, + requested_set=requested_set, + input_ids=input_ids, + weights_dir=weights_dir, + cache_context_length=cache_context_length, + common_graph_meta=common_graph_meta, + metadata=metadata, + ) + + pad_token_id = _resolve_model_pad_token_id(model) + decoder = Lfm2CausalLMLogitsAdapter(model, pad_token_id=pad_token_id).eval() + decoder_step = Lfm2CausalLMStepAdapter(model, pad_token_id=pad_token_id).eval() specs: list[ComponentModuleSpec] = [] if "decoder" in requested_set: specs.append(ComponentModuleSpec( @@ -5580,7 +5811,7 @@ def _build_lfm2_causal_lm_component_specs( input_keys=("input_ids",), output_keys=("logits",), graph_meta={**common_graph_meta, "component": "decoder"}, - metadata={"family": "lfm2", "task": "causal_lm_logits"}, + metadata=metadata, )) if "decoder_step" in requested_set: step_input_ids = input_ids[:, :1] @@ -5601,7 +5832,104 @@ def _build_lfm2_causal_lm_component_specs( "max_cache_seq_len": max_cache_seq_len, "cache_sink_size": 4, }, - metadata={"family": "lfm2", "task": "causal_lm_logits"}, + metadata=metadata, + )) + return specs + + +def _lfm2_chunked_pipeline_specs( + model: torch.nn.Module, + *, + requested_set: set[str], + input_ids: torch.Tensor, + weights_dir: str | None, + cache_context_length: str | int | None, + common_graph_meta: dict[str, object], + metadata: dict[str, str], + decoder_inputs_fn: Callable[[], tuple[torch.Tensor, ...]] | None = None, +) -> list[ComponentModuleSpec]: + lm_encoder_step = Lfm2VlLMEncoderStepAdapter(model, weights_dir=weights_dir).eval() + lm_encoder_text_chunk = Lfm2VlLMEncoderTextChunkAdapter(model, weights_dir=weights_dir).eval() + decoder_last_token = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=True).eval() + decoder_embed = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=False, return_hidden=True).eval() + + prefill_chunk = max(2, int(os.environ.get("CACTUS_LFM2_PREFILL_CHUNK", "128") or "128")) + chunk_input_ids = _tile_to_length(input_ids, prefill_chunk) + chunk_position_ids = torch.arange( + prefill_chunk, 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) + + decoder_inputs: tuple[torch.Tensor, ...] | None = None + if {"decoder_prefill_chunk", "decoder_embed_chunk", "decoder_step"} & requested_set: + if decoder_inputs_fn is not None: + decoder_inputs = decoder_inputs_fn() + else: + with torch.no_grad(): + decoder_inputs = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) + + cache_graph_meta = { + "use_internal_kv_cache": True, + "use_internal_conv_cache": True, + "max_cache_seq_len": _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512), + "cache_sink_size": 4, + } + + specs: list[ComponentModuleSpec] = [] + 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", "attention_mask", "position_ids"), + graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, + metadata=metadata, + )) + 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", "attention_mask", "position_ids"), + graph_meta={ + **common_graph_meta, + "component": "lm_encoder_text_chunk", + "encoder_chunk_size": prefill_chunk, + }, + metadata=metadata, + )) + if "decoder_prefill_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_prefill_chunk", + module=decoder_last_token, + example_inputs=tuple(tensor[:, :prefill_chunk, ...] for tensor in decoder_inputs), + input_keys=("inputs_embeds", "attention_mask", "position_ids"), + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "decoder_prefill_chunk", **cache_graph_meta}, + metadata=metadata, + )) + if "decoder_embed_chunk" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_embed_chunk", + module=decoder_embed, + example_inputs=tuple(tensor[:, :prefill_chunk, ...] for tensor in decoder_inputs), + input_keys=("inputs_embeds", "attention_mask", "position_ids"), + output_keys=("last_hidden_state",), + graph_meta={**common_graph_meta, "component": "decoder_embed_chunk", **cache_graph_meta}, + metadata=metadata, + )) + if "decoder_step" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_step", + module=decoder_last_token, + example_inputs=tuple(tensor[:, :1, ...] for tensor in decoder_inputs), + input_keys=("inputs_embeds", "attention_mask", "position_ids"), + output_keys=("logits",), + graph_meta={**common_graph_meta, "component": "decoder_step", **cache_graph_meta}, + metadata=metadata, )) return specs @@ -5631,8 +5959,6 @@ def _build_lfm2_vl_multimodal_component_specs( dtype=attention_mask.dtype, device=attention_mask.device, ) - prefill_chunk = max(2, int(os.environ.get("CACTUS_LFM2_PREFILL_CHUNK", "128") or "128")) - prefill_chunk = min(prefill_chunk, int(input_ids.shape[1])) requested_components = tuple(components or ("vision_encoder", "lm_encoder", "text_lm_encoder", "decoder")) requested_set = set(requested_components) @@ -5685,11 +6011,7 @@ def _require(component: str) -> None: ).eval() lm_encoder = Lfm2VlLMEncoderAdapter(model, input_ids=input_ids, weights_dir=weights_dir).eval() text_lm_encoder = Lfm2VlTextLMEncoderAdapter(model, weights_dir=weights_dir).eval() - lm_encoder_step = Lfm2VlLMEncoderStepAdapter(model, weights_dir=weights_dir).eval() - lm_encoder_text_chunk = Lfm2VlLMEncoderTextChunkAdapter(model, weights_dir=weights_dir).eval() decoder = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=False).eval() - decoder_last_token = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=True).eval() - decoder_embed = Lfm2VlDecoderAdapter(model, weights_dir=weights_dir, last_token_only=False, return_hidden=True).eval() image_features: torch.Tensor | None = None decoder_inputs: tuple[torch.Tensor, ...] | None = None @@ -5704,19 +6026,20 @@ def _require(component: str) -> None: if "text_lm_encoder" in expanded_components: text_decoder_inputs = text_lm_encoder(text_input_ids, text_attention_mask) - chunk_input_ids = input_ids[:, :prefill_chunk].contiguous() - chunk_position_ids = torch.arange( - prefill_chunk, - dtype=torch.long, - device=input_ids.device, - ).unsqueeze(0).expand(int(input_ids.shape[0]), -1).contiguous() + def _ensure_decoder_inputs() -> tuple[torch.Tensor, ...]: + nonlocal decoder_inputs, image_features + if decoder_inputs is None: + if image_features is None: + image_features = vision_encoder(pixel_values, spatial_shapes, pixel_attention_mask) + decoder_inputs = lm_encoder(input_ids, attention_mask, image_features) + return decoder_inputs common_graph_meta = { "weights_dir": weights_dir, "task": "multimodal_causal_lm_logits", "adapter_family": "lfm2_vl", } - max_cache_seq_len = _max_cache_seq_len(model, input_ids, cache_context_length, fallback_extra_tokens=512) + metadata = {"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"} specs: list[ComponentModuleSpec] = [] if "vision_encoder" in expanded_components: specs.append(ComponentModuleSpec( @@ -5726,7 +6049,7 @@ def _require(component: str) -> None: input_keys=("pixel_values", "spatial_shapes", "pixel_attention_mask"), output_keys=("image_features",), graph_meta={**common_graph_meta, "component": "vision_encoder"}, - metadata={"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"}, + metadata=metadata, )) if "lm_encoder" in expanded_components: if image_features is None: @@ -5738,7 +6061,7 @@ def _require(component: str) -> None: input_keys=("input_ids", "attention_mask", "image_features"), output_keys=("inputs_embeds", "attention_mask", "position_ids"), graph_meta={**common_graph_meta, "component": "lm_encoder"}, - metadata={"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"}, + metadata=metadata, )) if "text_lm_encoder" in expanded_components: if text_decoder_inputs is None: @@ -5750,7 +6073,7 @@ def _require(component: str) -> None: input_keys=("input_ids", "attention_mask"), output_keys=("inputs_embeds", "attention_mask", "position_ids"), graph_meta={**common_graph_meta, "component": "text_lm_encoder"}, - metadata={"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"}, + metadata=metadata, )) if "decoder" in expanded_components: if decoder_inputs is None: @@ -5764,7 +6087,7 @@ def _require(component: str) -> None: input_keys=("inputs_embeds", "attention_mask", "position_ids"), output_keys=("logits",), graph_meta={**common_graph_meta, "component": "decoder"}, - metadata={"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"}, + metadata=metadata, )) if "text_decoder" in expanded_components: if text_decoder_inputs is None: @@ -5776,97 +6099,18 @@ def _require(component: str) -> None: input_keys=("inputs_embeds", "attention_mask", "position_ids"), output_keys=("logits",), graph_meta={**common_graph_meta, "component": "text_decoder"}, - metadata={"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"}, - )) - if "lm_encoder_step" in expanded_components: - step_input_ids = input_ids[:, :1] - step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) - 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", "attention_mask", "position_ids"), - graph_meta={**common_graph_meta, "component": "lm_encoder_step"}, - metadata={"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"}, - )) - if "lm_encoder_text_chunk" in expanded_components: - 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", "attention_mask", "position_ids"), - graph_meta={ - **common_graph_meta, - "component": "lm_encoder_text_chunk", - "encoder_chunk_size": prefill_chunk, - }, - metadata={"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"}, - )) - if "decoder_prefill_chunk" in expanded_components: - if decoder_inputs is None: - if image_features is None: - image_features = vision_encoder(pixel_values, spatial_shapes, pixel_attention_mask) - decoder_inputs = lm_encoder(input_ids, attention_mask, image_features) - specs.append(ComponentModuleSpec( - component="decoder_prefill_chunk", - module=decoder_last_token, - example_inputs=tuple(tensor[:, :prefill_chunk, ...] for tensor in decoder_inputs), - input_keys=("inputs_embeds", "attention_mask", "position_ids"), - output_keys=("logits",), - graph_meta={ - **common_graph_meta, - "component": "decoder_prefill_chunk", - "use_internal_kv_cache": True, - "use_internal_conv_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 4, - }, - metadata={"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"}, - )) - if "decoder_embed_chunk" in expanded_components: - if decoder_inputs is None: - if image_features is None: - image_features = vision_encoder(pixel_values, spatial_shapes, pixel_attention_mask) - decoder_inputs = lm_encoder(input_ids, attention_mask, image_features) - specs.append(ComponentModuleSpec( - component="decoder_embed_chunk", - module=decoder_embed, - example_inputs=tuple(tensor[:, :prefill_chunk, ...] for tensor in decoder_inputs), - input_keys=("inputs_embeds", "attention_mask", "position_ids"), - output_keys=("last_hidden_state",), - graph_meta={ - **common_graph_meta, - "component": "decoder_embed_chunk", - "use_internal_kv_cache": True, - "use_internal_conv_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 4, - }, - metadata={"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"}, - )) - if "decoder_step" in expanded_components: - if decoder_inputs is None: - if image_features is None: - image_features = vision_encoder(pixel_values, spatial_shapes, pixel_attention_mask) - decoder_inputs = lm_encoder(input_ids, attention_mask, image_features) - specs.append(ComponentModuleSpec( - component="decoder_step", - module=decoder_last_token, - example_inputs=tuple(tensor[:, :1, ...] for tensor in decoder_inputs), - input_keys=("inputs_embeds", "attention_mask", "position_ids"), - output_keys=("logits",), - graph_meta={ - **common_graph_meta, - "component": "decoder_step", - "use_internal_kv_cache": True, - "use_internal_conv_cache": True, - "max_cache_seq_len": max_cache_seq_len, - "cache_sink_size": 4, - }, - metadata={"family": "lfm2_vl", "task": "multimodal_causal_lm_logits"}, + metadata=metadata, )) + specs.extend(_lfm2_chunked_pipeline_specs( + model, + requested_set=set(expanded_components), + input_ids=input_ids, + weights_dir=weights_dir, + cache_context_length=cache_context_length, + common_graph_meta=common_graph_meta, + metadata=metadata, + decoder_inputs_fn=_ensure_decoder_inputs, + )) return specs @@ -5897,6 +6141,14 @@ def build_component_module_specs( components=components, cache_context_length=cache_context_length, ) + if family == "gemma3" and task == "causal_lm_logits": + return _build_gemma3_causal_lm_component_specs( + model, + named_tensors=named_tensors, + weights_dir=weights_dir, + components=components, + cache_context_length=cache_context_length, + ) if family == "gemma4" and task == "multimodal_causal_lm_logits": return _build_gemma4_multimodal_component_specs( model, diff --git a/python/cactus/transpile/model_profiles.py b/python/cactus/transpile/model_profiles.py index 9d94c9cc6..75580eb47 100644 --- a/python/cactus/transpile/model_profiles.py +++ b/python/cactus/transpile/model_profiles.py @@ -29,6 +29,7 @@ class ModelProfile: text_only_component: str | None = None default_task: str | None = None default_components: tuple[str, ...] = () + text_component_plans: tuple[tuple[str, tuple[str, ...]], ...] = () default_max_new_tokens: int | None = None needs_image: bool = False needs_audio: bool = False @@ -90,6 +91,10 @@ class ModelProfile: text_only_component="text_lm_encoder", default_task="multimodal_causal_lm_logits", default_components=("vision_encoder", "lm_encoder", "decoder"), + text_component_plans=( + ("lfm2forcausallm", ("decoder_step", "lm_encoder_step", "lm_encoder_text_chunk", + "decoder_prefill_chunk")), + ), needs_image=True, force_component_pipeline=True, ) @@ -197,6 +202,10 @@ class ModelProfile: "decoder_media_step", "decoder_step", ), + text_component_plans=( + ("qwen3forcausallm", ("decoder_step", "lm_encoder_step", "lm_encoder_text_chunk", + "decoder_prefill_chunk", "decoder_embed_chunk", "decoder_media_step")), + ), needs_image=True, force_component_pipeline=True, ) diff --git a/python/cactus/transpile/optimize_graph.py b/python/cactus/transpile/optimize_graph.py index 82d49fccf..de3116e22 100644 --- a/python/cactus/transpile/optimize_graph.py +++ b/python/cactus/transpile/optimize_graph.py @@ -99,6 +99,8 @@ def optimize_graph(graph: IRGraph, *, max_passes: int = 8, config: FusionConfig changed = True if normalize_gemma4_decoder_attention_semantics(graph): changed = True + if normalize_cached_decoder_attention_hints(graph): + changed = True if config.enable_dense_mlp_tq_fused and fuse_dense_mlp_tq(graph): changed = True if config.enable_conv_module and fuse_conv_modules(graph): @@ -812,6 +814,30 @@ def normalize_gemma4_decoder_attention_semantics(graph: IRGraph) -> bool: return changed +def normalize_cached_decoder_attention_hints(graph: IRGraph) -> bool: + if _is_gemma4_graph(graph): + return False + if not bool(graph.meta.get("use_internal_kv_cache", False)): + return False + changed = _assign_gemma4_decoder_attention_hints_from_graph_meta(graph) + sliding_window = _graph_sliding_window(graph) + if sliding_window is not None and sliding_window > 0: + for node_id in graph.order: + node = graph.nodes.get(node_id) + if node is None or node.op not in {"attention", "scaled_dot_product_attention", "attention_block"}: + continue + layer_type = str(node.meta.get("attention_layer_type") or "").strip().lower() + if layer_type not in {"sliding", "sliding_attention"}: + continue + if int(node.attrs.get("window_size", 0) or 0) == 0: + node.attrs["window_size"] = int(sliding_window) + node.meta.setdefault("window_size_source", "layer_type_config") + changed = True + if changed: + rebuild_graph(graph) + return changed + + def _assign_gemma4_decoder_attention_hints_from_graph_meta(graph: IRGraph) -> bool: component = str(graph.meta.get("component", "") or "").strip().lower() if component not in {"decoder", "decoder_step", "decoder_prefill_chunk"}: