From 363dfb3cfab0b723524e58b023d395c6d30df5f4 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Wed, 10 Jun 2026 22:56:14 -0700 Subject: [PATCH 01/13] Report completion_token_ids from cactus_benchmark_tokens Lets benchmark runs compare generated tokens across engine/bundle variants without a separate text-capture path. Signed-off-by: Noah Cylich --- cactus-engine/src/complete.cpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/cactus-engine/src/complete.cpp b/cactus-engine/src/complete.cpp index 84f3babbc..1c04e7d49 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; } @@ -1400,6 +1403,12 @@ int cactus_benchmark_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() << "}"; From 83fa8ad41e4879a570dbd7196b82c30e096c5f9c Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Wed, 10 Jun 2026 22:56:14 -0700 Subject: [PATCH 02/13] engine: pad the prefill tail chunk and roll the KV caches back MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prompt remainder (L mod 128) ran one token at a time at decode speed (~17.7ms/token on M4): a 1023-token prompt paid 127 decode steps, 3.3x slower prefill than 1024 (gemma-prefill-bug.md). Global-attention models already pad the tail chunk, but a pad through a sliding-window cache evicts real window rows, which is why sliding bundles were excluded — the original chunked-prefill garbled-output bug. Run the remainder as one padded chunk instead. The sliding caches compact chronologically on eviction, so a padded append just over-evicts up to pad_count of the oldest window rows: snapshot them before the padded execute and reinsert them after (rollback also drops the pad rows), leaving the persistent cache state identical to an unpadded prefill. The chunk graph only emits the last row's logits — a pad's prediction — so the last real token is rolled back too and re-run through the step graph, which is the same mechanism the global-attention padding path uses. Worst case tail cost drops from 127 decode steps (~2.2s) to one padded chunk and one decode step (~150ms), with no bundle changes: existing sliding-window bundles speed up as-is. The path applies to any sliding-window bundle with a chunked prefill graph; it skips graphs whose window is too small for the pad bound and bundles with recurrent or conv caches (no rollback support), and those plus the global-attention and no-chunk paths behave exactly as before. Telemetry: prefill_tail_chunk_tokens / prefill_tail_padding_tokens added; prefill_scalar_tail_tokens is now at most 1 on the padded path. Kill switch for A/B: CACTUS_DISABLE_PREFILL_TAIL_PAD=1. Note: cloud-handoff probe hiddens are only captured in run_step, so probe rollouts see at most one sample on such bundles. Signed-off-by: Noah Cylich --- cactus-engine/src/complete.cpp | 2 + cactus-engine/src/engine.h | 10 +++ cactus-engine/src/model.cpp | 135 +++++++++++++++++++++++---------- cactus-graph/cactus_graph.h | 4 + cactus-graph/src/builder.cpp | 4 + cactus-graph/src/ops_cache.cpp | 93 +++++++++++++++++++++++ 6 files changed, 210 insertions(+), 38 deletions(-) diff --git a/cactus-engine/src/complete.cpp b/cactus-engine/src/complete.cpp index 1c04e7d49..e3ff46097 100644 --- a/cactus-engine/src/complete.cpp +++ b/cactus-engine/src/complete.cpp @@ -1400,6 +1400,8 @@ 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 << "," diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index e4d361e8b..4ddc5a81a 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -684,6 +684,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 +783,15 @@ 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); + bool padded_tail_window_ok(const Component& comp, size_t chunk_tokens) const; + 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..8522348af 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -563,6 +563,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 +1062,61 @@ 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; +bool Model::padded_tail_window_ok(const Component& comp, size_t chunk_tokens) const { + for (const auto& state : comp.cache_states) { + for (int node_id : {state.key_node_id, state.value_node_id}) { + if (node_id < 0) continue; + if (comp.graph->get_node_op_type(static_cast(node_id)) != OpType::KV_CACHE_STATE) continue; + size_t window = comp.graph->get_node_window_size(static_cast(node_id)); + size_t sink = comp.graph->get_node_sink_size(static_cast(node_id)); + if (window > 0 && (window < chunk_tokens * 4 || window <= chunk_tokens + sink)) return false; + } + } + return true; +} + +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; @@ -1100,16 +1153,18 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector= padding_cutoff; + 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 use_padded_tail = !pad_tail && !prefill_tail_pad_disabled_ + && has_sliding_window_cache && !has_recurrent_state && !has_conv_state + && tail_tokens >= 2 && padded_tail_window_ok(*decoder_prefill_, effective_chunk); 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 +1176,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 +1594,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 +1671,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); 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; From f463d8409e11cc115afe5595b304fa2205516ea1 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Wed, 10 Jun 2026 22:56:14 -0700 Subject: [PATCH 03/13] Test padded-append rollback and chunked-prefill padding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Graph level: append(real + pads) followed by rollback must leave the cache byte-identical to append(real) — covered for no-eviction, eviction with overshoot, pads-only eviction, and empty-cache cases. Engine level: the padded tail engages with correct telemetry, is deterministic, and the CACTUS_DISABLE_PREFILL_TAIL_PAD kill switch restores the scalar tail. Signed-off-by: Noah Cylich --- cactus-engine/tests/test_llm.cpp | 88 +++++++++++++++++++++++++++++++ cactus-graph/tests/test_cache.cpp | 75 ++++++++++++++++++++++++++ 2 files changed, 163 insertions(+) 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/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()); From 610bdb8c73c430d4a7b03f23429633c193cab529 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Wed, 10 Jun 2026 22:58:05 -0700 Subject: [PATCH 04/13] transpile: add Gemma3 chunked-prefill component pipeline Gemma3 previously only had a single full-graph causal-LM capture, so the engine had no CACHED_STEP route and prefilled token-by-token. Mirror the qwen text component set (lm_encoder_step, lm_encoder_text_chunk, decoder_prefill_chunk, decoder_step) with internal KV caches and per-layer-type rotary/masks. A family-neutral optimize pass generalizes the gemma4-only hint assignment so sliding layers get window_size = sliding_window on their KV cache states. Signed-off-by: Noah Cylich --- python/cactus/transpile/model_adapters.py | 269 ++++++++++++++++++++-- python/cactus/transpile/optimize_graph.py | 26 +++ 2 files changed, 271 insertions(+), 24 deletions(-) diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index a2ae0a36a..a7b831fef 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -2054,6 +2054,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 +2129,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 +3472,129 @@ 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")) + prefill_chunk_size = min(prefill_chunk_size, int(input_ids.shape[1])) + + step_input_ids = input_ids[:, :1] + step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) + chunk_input_ids = input_ids[:, :prefill_chunk_size].contiguous() + 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, *, @@ -5897,6 +6110,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/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"}: From ff4390f2453705c9088c5f018286c5d4fe484d7e Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Wed, 10 Jun 2026 23:17:24 -0700 Subject: [PATCH 05/13] convert+engine: gemma3 family with fp16-range weight rescale Gemma3 text checkpoints fell through to the generic family, so bundles carried model_type=generic, which the engine treats as gemma4 and rejects for missing gemma4-only config fields. Worse, gemma3's residual stream peaks far above fp16 max (~103K on gemma-3-270m), so unscaled weights produce garbage on the fp16-activation engine. Add a gemma3 family that keeps generic naming/policy but mirrors the gemma4 16x residual rescale: embeddings (and the tied lm head) shrink 16x, gate/up restore MLP scale, and the norm files fold the rescale through gemma3's RMSNorm +1 offset, which the transpiled graph re-applies. The engine maps model_type=gemma3 onto its non-strict GEMMA type. Signed-off-by: Noah Cylich --- cactus-engine/src/engine.h | 3 +- cactus-engine/src/model.cpp | 2 +- cactus-engine/src/tokenizer.cpp | 36 ++++++++++++++++++- .../convert/cactus_adapters/tensor_io.py | 13 ++++++- .../cactus/convert/model_adapters/adapters.py | 10 +++++- .../convert/model_adapters/detection.py | 7 +++- .../cactus/convert/model_adapters/naming.py | 16 ++++++++- 7 files changed, 80 insertions(+), 7 deletions(-) diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 4ddc5a81a..6f0d12cbf 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 { diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 8522348af..7d07b4f34 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -3524,7 +3524,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/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: From a584bfdf002bc6a1a8afaf8f84c63e30771d2927 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 9 Jun 2026 17:29:57 -0700 Subject: [PATCH 06/13] engine: warn when a text bundle prefills token-by-token If an LM bundle provides no decoder_prefill_chunk component, prompts prefill one token at a time and prefill throughput collapses to roughly decode speed - silently. Emit a one-time warning at model load naming the fallback and the fix (re-transpile with the current converter); audio/whisper routes are exempt. Raise the benchmark fixture's log level to WARN so the warning is visible during benchmarking, where an unnoticed fallback invalidates prefill numbers (observed with a qwen3 bundle whose manifest omitted the chunk components: 21 tps prefill vs 62 with them). Signed-off-by: Noah Cylich --- cactus-engine/src/model.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 7d07b4f34..f3e3262c4 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). Re-transpile with the current converter for chunked prefill."); } if (components_.count("decoder_embed_chunk") && components_.at("decoder_embed_chunk").graph) { decoder_embed_ = &components_.at("decoder_embed_chunk"); From d3f6c4de6f8196e0d3a47c6e69b5cfbfc45e8bc8 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 9 Jun 2026 17:57:22 -0700 Subject: [PATCH 07/13] transpile: honor the inferred component plan for text causal-lm models build_component_module_specs only ever received explicit --components; the component plan that infer_component_plan_from_config computes was consulted for task inference and then dropped, so default text-qwen3 transpiles emitted the monolithic decoder/decoder_step pair instead of the chunked-prefill pipeline - prefill collapsed to decode speed (~21 vs ~62 tps at 512 tokens on Pixel 10a). Pass the plan's components through when --components is absent. Spec builders reject inapplicable plans with the new UnsupportedComponentsError and the call site falls back to builder defaults with a warning; explicit --components and unrelated failures still raise. Verified: fresh default transpile of Qwen/Qwen3-0.6B emits lm_encoder_text_chunk/decoder_prefill_chunk/lm_encoder_step/ decoder_media_step/decoder_step and benches at 229 prefill / 61 decode tps on an M4 Pro (previously 67/51 token-by-token); transpile test suite passes. Signed-off-by: Noah Cylich --- python/cactus/transpile/hf_model.py | 27 ++++++++++++++++++++--- python/cactus/transpile/model_adapters.py | 6 ++++- 2 files changed, 29 insertions(+), 4 deletions(-) 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 a7b831fef..5dba717cb 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 @@ -4698,7 +4702,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}" ) From 2a1ed3eb7bc62c6e35cf59e3b1210cf5b8b9395b Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 9 Jun 2026 18:02:09 -0700 Subject: [PATCH 08/13] engine: drop re-transpile advice from the prefill fallback warning The remedy only applies to families whose converter emits chunked components (currently qwen3 text); for others (e.g. base LFM2) a re-transpile changes nothing and the advice misleads. Signed-off-by: Noah Cylich --- cactus-engine/src/model.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index f3e3262c4..d3f543b1e 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -557,7 +557,7 @@ bool Model::init(const std::string& bundle_dir, size_t context_size, 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). Re-transpile with the current converter for chunked prefill."); + 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"); From 1ca3a3fa4e3d6d04aebf5d41bfb649bf5a4fef89 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 9 Jun 2026 18:12:26 -0700 Subject: [PATCH 09/13] transpile: chunked prefill for text-only LFM2 conversions The lfm2 chunk adapters (lm_encoder_step, lm_encoder_text_chunk, decoder_prefill_chunk and the embeds-based decoder_step) already power the LFM2-VL pipeline and resolve the language backbone generically, so text-only Lfm2ForCausalLM can reuse them directly. Route chunk-component requests in the lfm2 text spec builder to a new chunked builder and make the component plan select the chunked set for the Lfm2ForCausalLM architecture, matching the qwen3 text default. Verified on LFM2-350M: transpile with reference compare passes, the bundle prefills 565 tps vs 134 token-by-token on an M4 Pro (4.2x) with decode unchanged (141 vs 143), no prefill-fallback warning, and greedy generation stays coherent and factually correct across spot checks. Signed-off-by: Noah Cylich --- python/cactus/transpile/component_plan.py | 7 ++ python/cactus/transpile/model_adapters.py | 93 +++++++++++++++++++++++ 2 files changed, 100 insertions(+) diff --git a/python/cactus/transpile/component_plan.py b/python/cactus/transpile/component_plan.py index 5c5df4114..bf0ff4135 100644 --- a/python/cactus/transpile/component_plan.py +++ b/python/cactus/transpile/component_plan.py @@ -37,6 +37,13 @@ def _plan_from_profile( "decoder_prefill_chunk", "decoder_embed_chunk", "decoder_media_step"), force_component_pipeline=True, ) + if any("lfm2forcausallm" 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"), + force_component_pipeline=True, + ) return ComponentPlan( task="causal_lm_logits", components=("decoder", "decoder_step"), diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index 5dba717cb..7120922b7 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -5783,6 +5783,15 @@ def _build_lfm2_causal_lm_component_specs( requested = tuple(components or ("decoder", "decoder_step")) requested_set = set(requested) + chunk_components = {"lm_encoder_step", "lm_encoder_text_chunk", "decoder_prefill_chunk"} + if chunk_components & requested_set: + return _build_lfm2_text_chunked_component_specs( + model, + input_ids=input_ids, + weights_dir=weights_dir, + requested_set=requested_set, + cache_context_length=cache_context_length, + ) common_graph_meta = { "weights_dir": weights_dir, "task": "causal_lm_logits", @@ -5823,6 +5832,90 @@ def _build_lfm2_causal_lm_component_specs( return specs +def _build_lfm2_text_chunked_component_specs( + model: torch.nn.Module, + *, + input_ids: torch.Tensor, + weights_dir: str | None, + requested_set: set[str], + cache_context_length: str | int | 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() + + 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])) + 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() + step_input_ids = input_ids[:, :1] + step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) + with torch.no_grad(): + chunk_decoder_inputs = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) + + common_graph_meta = { + "weights_dir": weights_dir, + "task": "causal_lm_logits", + "adapter_family": "lfm2", + } + 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, + } + metadata = {"family": "lfm2", "task": "causal_lm_logits"} + + 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=chunk_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_step" in requested_set: + specs.append(ComponentModuleSpec( + component="decoder_step", + module=decoder_last_token, + example_inputs=tuple(tensor[:, :1, ...] for tensor in chunk_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 + + def _build_lfm2_vl_multimodal_component_specs( model: torch.nn.Module, *, From 5b20bc9780d93ceefa31238a6b9e2b2f93521ebb Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 9 Jun 2026 18:37:03 -0700 Subject: [PATCH 10/13] transpile: unify LFM2 text and VL chunked component spec emission The four chunked component specs (lm_encoder_step, lm_encoder_text_chunk, decoder_prefill_chunk, decoder_step) were emitted twice with identical structure: once by the text-only builder and once inline in the LFM2-VL multimodal builder. Generalize the text builder into _lfm2_chunked_pipeline_specs, parameterized by graph/spec metadata and an optional decoder-inputs provider, and emit both paths through it. The VL path keeps its lazy vision-encoder + lm_encoder decoder-input computation via a closure; the text path keeps deriving decoder inputs from lm_encoder_text_chunk. Net -54 lines, no behavior change: transpiled artifacts for LFM2-350M and LFM2-VL-450M are byte-identical before and after. Signed-off-by: Noah Cylich --- python/cactus/transpile/model_adapters.py | 197 +++++++--------------- 1 file changed, 65 insertions(+), 132 deletions(-) diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index 7120922b7..56e99f940 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -5777,26 +5777,29 @@ 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 = { + "weights_dir": weights_dir, + "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 _build_lfm2_text_chunked_component_specs( + return _lfm2_chunked_pipeline_specs( model, + requested_set=requested_set, input_ids=input_ids, weights_dir=weights_dir, - requested_set=requested_set, cache_context_length=cache_context_length, + common_graph_meta=common_graph_meta, + metadata=metadata, ) - common_graph_meta = { - "weights_dir": weights_dir, - "task": "causal_lm_logits", - "adapter_family": "lfm2", - } + + 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( @@ -5806,7 +5809,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] @@ -5827,22 +5830,26 @@ 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 _build_lfm2_text_chunked_component_specs( +def _lfm2_chunked_pipeline_specs( model: torch.nn.Module, *, + requested_set: set[str], input_ids: torch.Tensor, weights_dir: str | None, - requested_set: set[str], 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")) prefill_chunk = min(prefill_chunk, int(input_ids.shape[1])) @@ -5852,21 +5859,21 @@ def _build_lfm2_text_chunked_component_specs( ).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_decoder_inputs = lm_encoder_text_chunk(chunk_input_ids, chunk_position_ids) - common_graph_meta = { - "weights_dir": weights_dir, - "task": "causal_lm_logits", - "adapter_family": "lfm2", - } + 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, } - metadata = {"family": "lfm2", "task": "causal_lm_logits"} specs: list[ComponentModuleSpec] = [] if "lm_encoder_step" in requested_set: @@ -5897,17 +5904,27 @@ def _build_lfm2_text_chunked_component_specs( specs.append(ComponentModuleSpec( component="decoder_prefill_chunk", module=decoder_last_token, - example_inputs=chunk_decoder_inputs, + 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 chunk_decoder_inputs), + 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}, @@ -5941,8 +5958,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) @@ -5995,11 +6010,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 @@ -6014,19 +6025,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( @@ -6036,7 +6048,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: @@ -6048,7 +6060,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: @@ -6060,7 +6072,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: @@ -6074,7 +6086,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: @@ -6086,97 +6098,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 From dc3614e9df6cbeef734ee3d4ff707829d0b107ab Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 9 Jun 2026 19:13:24 -0700 Subject: [PATCH 11/13] transpile: declare text-only component plans on model profiles Move the qwen3/lfm2 text-only chunked component plans out of _plan_from_profile's hardcoded architecture conditionals into a text_component_plans field on ModelProfile, matching how multimodal pipelines are declared via default_components. The plan layer now consults only the matched profile's own declarations, so one family's architecture markers can no longer select another family's components; real checkpoints are unaffected since model_type and architectures always agree on family. Signed-off-by: Noah Cylich --- python/cactus/transpile/component_plan.py | 21 +++++++-------------- python/cactus/transpile/model_profiles.py | 9 +++++++++ 2 files changed, 16 insertions(+), 14 deletions(-) diff --git a/python/cactus/transpile/component_plan.py b/python/cactus/transpile/component_plan.py index bf0ff4135..c9e63cabd 100644 --- a/python/cactus/transpile/component_plan.py +++ b/python/cactus/transpile/component_plan.py @@ -30,20 +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, - ) - if any("lfm2forcausallm" 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"), - 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/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, ) From 234211098efceb6f4f533c56b94006b13149905f Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Thu, 11 Jun 2026 00:13:45 -0700 Subject: [PATCH 12/13] engine: gate tail padding on conv caches instead of the lfm2_vl family MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Text-only LFM2 bundles now ship chunked prefill graphs, so they reach the tail-padding decision for the first time — and the family gate only covered lfm2_vl. Pads appended to a conv ring displace exactly the rows the conv kernel's lookback reads and no clamp removes them: padded prefill produced corrupted output on LFM2-350M (tokens diverge from position 1 vs the scalar tail). Gating on conv-cache presence fixes that and subsumes the lfm2_vl special case. Signed-off-by: Noah Cylich --- cactus-engine/src/engine.h | 1 - cactus-engine/src/model.cpp | 27 ++++++++++----------------- 2 files changed, 10 insertions(+), 18 deletions(-) diff --git a/cactus-engine/src/engine.h b/cactus-engine/src/engine.h index 6f0d12cbf..aeed9d0e5 100644 --- a/cactus-engine/src/engine.h +++ b/cactus-engine/src/engine.h @@ -789,7 +789,6 @@ class Model { 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); - bool padded_tail_window_ok(const Component& comp, size_t chunk_tokens) const; 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); diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index d3f543b1e..50e766c8b 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -1064,19 +1064,6 @@ size_t Model::component_output_tokens(const Component& comp, const std::string& return 0; } -bool Model::padded_tail_window_ok(const Component& comp, size_t chunk_tokens) const { - for (const auto& state : comp.cache_states) { - for (int node_id : {state.key_node_id, state.value_node_id}) { - if (node_id < 0) continue; - if (comp.graph->get_node_op_type(static_cast(node_id)) != OpType::KV_CACHE_STATE) continue; - size_t window = comp.graph->get_node_window_size(static_cast(node_id)); - size_t sink = comp.graph->get_node_sink_size(static_cast(node_id)); - if (window > 0 && (window < chunk_tokens * 4 || window <= chunk_tokens + sink)) return false; - } - } - return true; -} - 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) { @@ -1151,16 +1138,22 @@ 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 has_conv_state = any_cache_node([&](size_t id) { - return decoder_prefill_->graph->get_node_op_type(id) == OpType::CONV_CACHE_STATE; + 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 >= 2 && padded_tail_window_ok(*decoder_prefill_, effective_chunk); + && tail_tokens > 8 && !padded_window_too_small; const size_t executable_tokens = whole_chunks_end + (pad_tail ? effective_chunk : 0); if (executable_tokens == 0 && !use_padded_tail) { result.scalar_tail_tokens = tail_tokens; From e8ef8cd66764b4d894a4a4330b03046464bc5f0a Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Thu, 11 Jun 2026 19:21:51 -0700 Subject: [PATCH 13/13] transpile: decouple prefill chunk size from the calibration prompt The chunk components sliced their example input from the calibration prompt, so a short prompt silently shrank the bundle's prefill chunk (default convert: 46 tokens instead of 128). Tile the trace input up to the configured chunk size instead. Signed-off-by: Noah Cylich --- python/cactus/transpile/model_adapters.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index 56e99f940..38eef3b3d 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -216,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"): @@ -3531,11 +3536,10 @@ def _build_gemma3_causal_lm_component_specs( "cache_sink_size": 4, } prefill_chunk_size = max(2, int(os.environ.get("CACTUS_GEMMA3_PREFILL_CHUNK", "128") or "128")) - prefill_chunk_size = min(prefill_chunk_size, int(input_ids.shape[1])) step_input_ids = input_ids[:, :1] step_position_ids = torch.zeros_like(step_input_ids, dtype=torch.int64) - 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.int64, device=input_ids.device, ).unsqueeze(0) @@ -4747,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, @@ -4916,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, @@ -5852,8 +5854,7 @@ def _lfm2_chunked_pipeline_specs( 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")) - prefill_chunk = min(prefill_chunk, int(input_ids.shape[1])) - chunk_input_ids = input_ids[:, :prefill_chunk].contiguous() + 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()