From c3871ae690ac5942ad24cf1b492bb4711385a745 Mon Sep 17 00:00:00 2001 From: Karen Mosoyan Date: Wed, 10 Jun 2026 20:31:08 -0700 Subject: [PATCH] fixed needle rope bug, moved repeated rope table instantiations from components into weights --- cactus-engine/tests/test_needle.cpp | 97 +++++++++++++++++++++++ python/cactus/transpile/hf_model.py | 49 ++++++++++++ python/cactus/transpile/model_adapters.py | 11 +++ python/cactus/transpile/optimize_graph.py | 12 ++- 4 files changed, 167 insertions(+), 2 deletions(-) create mode 100644 cactus-engine/tests/test_needle.cpp diff --git a/cactus-engine/tests/test_needle.cpp b/cactus-engine/tests/test_needle.cpp new file mode 100644 index 000000000..905672a09 --- /dev/null +++ b/cactus-engine/tests/test_needle.cpp @@ -0,0 +1,97 @@ +#include "test_utils.h" + +#include +#include +#include +#include +#include +#include + +static const char* g_model_path = std::getenv("CACTUS_TEST_MODEL"); + +static bool model_is_needle(const std::string& bundle_path) { + std::ifstream f(bundle_path + "/config.json"); + if (!f) return false; + std::stringstream ss; + ss << f.rdbuf(); + const std::string content = ss.str(); + const size_t key = content.find("\"model_type\""); + if (key == std::string::npos) return false; + const size_t colon = content.find(':', key); + if (colon == std::string::npos) return false; + const size_t q1 = content.find('"', colon); + if (q1 == std::string::npos) return false; + const size_t q2 = content.find('"', q1 + 1); + if (q2 == std::string::npos) return false; + return content.substr(q1 + 1, q2 - q1 - 1) == "needle"; +} + +bool test_needle_tool_call() { + std::cout << "\n╔══════════════════════════════════════════╗\n" + << "║ NEEDLE TOOL CALL TEST ║\n" + << "╚══════════════════════════════════════════╝\n"; + if (!g_model_path) { + std::cout << " [WARN] CACTUS_TEST_MODEL not set; skipping\n"; + return true; + } + if (!model_is_needle(g_model_path)) { + std::cout << " [SKIP] model under test is not needle; skipping needle tool-call test\n"; + return true; + } + + cactus_model_t model = cactus_init(g_model_path, nullptr, false); + if (!model) { + std::cerr << " [✗] Failed to initialize needle model at " << g_model_path << "\n"; + return false; + } + + const char* messages = R"([ + {"role": "user", "content": "What's the weather in San Francisco?"} + ])"; + + const char* tools = R"([{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a city.", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string", "description": "City name."} + }, + "required": ["location"] + } + } + }])"; + + const char* options = R"({"max_tokens": 64, "force_tools": true, "telemetry_enabled": false, "auto_handoff": false})"; + + char response[1 << 15] = {0}; + int rc = cactus_complete(model, messages, response, sizeof(response), options, tools, + nullptr, nullptr, nullptr, 0); + std::string r(response); + + std::cout << "├─ rc: " << rc << "\n"; + std::cout << "├─ Response: " << r << "\n"; + + const bool has_function_calls = + r.find("\"function_calls\":[{") != std::string::npos; + const bool calls_tool = r.find("get_weather") != std::string::npos; + const bool has_arg = r.find("San Francisco") != std::string::npos; + std::cout << "├─ function_calls present: " << (has_function_calls ? "YES" : "NO") << "\n"; + std::cout << "├─ Calls get_weather: " << (calls_tool ? "YES" : "NO") << "\n"; + std::cout << "├─ Location arg present: " << (has_arg ? "YES" : "NO") << "\n"; + + const bool ok = rc > 0 && has_function_calls && calls_tool && has_arg; + std::cout << "└─ Status: " << (ok ? "PASSED ✓" : "FAILED ✗") << "\n"; + + cactus_destroy(model); + return ok; +} + +int main() { + bool ok = true; + ok &= test_needle_tool_call(); + std::cout << "\n" << (ok ? "✓ needle tool-call test passed" : "✗ needle tool-call test failed") << "\n"; + return ok ? 0 : 1; +} diff --git a/python/cactus/transpile/hf_model.py b/python/cactus/transpile/hf_model.py index f0edbd570..c0ee2183d 100644 --- a/python/cactus/transpile/hf_model.py +++ b/python/cactus/transpile/hf_model.py @@ -3,6 +3,7 @@ import argparse import copy import gc +import hashlib import itertools import json import os @@ -24,6 +25,7 @@ sys.path.insert(0, str(PYTHON_ROOT)) from cactus.transpile.runtime_compat import Graph +from cactus.convert.cactus_adapters.tensor_io import save_tensor_with_header from cactus.convert.model_adapters.nemo import ensure_parakeet_tdt_nemo_source from cactus.transpile.audio_preprocess import generic_log_mel_features as _generic_log_mel_features from cactus.transpile.audio_preprocess import load_audio_waveform as _load_audio_waveform @@ -238,6 +240,46 @@ def _binding_entries_by_node_id( return result +def _externalize_shared_rope_tables( + transpiled_graph: TranspiledGraph, + *, + weights_root: Path, + registry: dict[str, str], +) -> int: + value_ids = getattr(transpiled_graph, "bound_constant_value_ids", {}) or {} + already_bound = { + int(binding["node_id"]) + for binding in transpiled_graph.bound_constant_bindings + if isinstance(binding, dict) and "node_id" in binding + } + externalized = 0 + for constant in transpiled_graph.bound_constants: + node_id = int(constant.id) + if node_id in already_bound: + continue + value_id = str(value_ids.get(node_id, "")) + if not value_id.startswith("c_rope_table_"): + continue + data = np.ascontiguousarray(constant.numpy()) + digest = hashlib.sha1(data.tobytes()).hexdigest()[:16] + filename = registry.get(digest) + if filename is None: + filename = f"rope_table_{digest}.weights" + save_tensor_with_header(data, weights_root / filename, precision="FP16") + registry[digest] = filename + transpiled_graph.bound_constant_bindings.append( + { + "node_id": node_id, + "value_id": value_id, + "path": filename, + "kind": "weight", + "source_name": value_id, + } + ) + externalized += 1 + return externalized + + def _embed_materialized_bound_constants(transpiled_graph: TranspiledGraph) -> int: """Serialize small graph-local constants into the .cactus graph itself. @@ -335,6 +377,8 @@ def _write_component_bundle( component_order.extend(extra_components) manifest_components: list[dict[str, object]] = [] + + rope_table_registry: dict[str, str] = {} for component in component_order: raw_graph = raw_component_graphs.get(component) optimized_graph = optimized_component_graphs.get(component) @@ -375,6 +419,11 @@ def _write_component_bundle( graph_relpath = Path(component) / graph_filename component_dir.mkdir(parents=True, exist_ok=True) shutil.rmtree(component_dir / "bound_constants", ignore_errors=True) + _externalize_shared_rope_tables( + transpiled_graph, + weights_root=artifact_dir, + registry=rope_table_registry, + ) _embed_materialized_bound_constants(transpiled_graph) transpiled_graph.graph.save(bundle_dir / graph_relpath) diff --git a/python/cactus/transpile/model_adapters.py b/python/cactus/transpile/model_adapters.py index 701fddd98..071687b39 100644 --- a/python/cactus/transpile/model_adapters.py +++ b/python/cactus/transpile/model_adapters.py @@ -5310,6 +5310,7 @@ def _resolve_decoder_start_token_id(model: torch.nn.Module, *, fallback: int = 0 return int(value) return int(fallback) +NEEDLE_CONTEXT_LENGTH = 1024 def _build_needle_causal_lm_component_specs( model: torch.nn.Module, @@ -5325,6 +5326,16 @@ def _build_needle_causal_lm_component_specs( raise RuntimeError(f"{type(model).__name__} does not expose {method_name}") pad_token_id = _resolve_model_pad_token_id(model) + current_len = int(input_ids.shape[1]) + if current_len < NEEDLE_CONTEXT_LENGTH: + pad_value = int(pad_token_id) if pad_token_id is not None else 0 + padding = torch.full( + (int(input_ids.shape[0]), NEEDLE_CONTEXT_LENGTH - current_len), + pad_value, + dtype=input_ids.dtype, + device=input_ids.device, + ) + input_ids = torch.cat([input_ids, padding], dim=1) if pad_token_id is None: attention_mask = torch.ones_like(input_ids, dtype=torch.int64) else: diff --git a/python/cactus/transpile/optimize_graph.py b/python/cactus/transpile/optimize_graph.py index c700c7128..25af1b8e5 100644 --- a/python/cactus/transpile/optimize_graph.py +++ b/python/cactus/transpile/optimize_graph.py @@ -23,7 +23,7 @@ from cactus.transpile.fusion.common import strip_layout_passthrough from cactus.transpile.fusion.common import strip_passthrough from cactus.transpile.fusion.linear import match_linear -from cactus.transpile.fusion.rope import _extract_rope_angle_source, _find_constant_ancestor +from cactus.transpile.fusion.rope import _extract_rope_angle_source from cactus.transpile.graph_ir import IRGraph from cactus.transpile.graph_ir import IRNode from cactus.transpile.graph_ir import IRValue @@ -414,6 +414,14 @@ def precompute_rope_tables(graph: IRGraph) -> bool: return changed +def _stored_inv_freq_constant(graph: IRGraph, value_id: str) -> str | None: + stripped = strip_layout_passthrough(graph, value_id) + constant = graph.constants.get(stripped) + if isinstance(constant, torch.Tensor): + return stripped + return None + + def _match_rope_angle_table_source(graph: IRGraph, value_id: str, index_shape: tuple[int, ...]) -> tuple[str, str] | None: cat_node = producer(graph, strip_passthrough(graph, value_id)) if cat_node is None or cat_node.op != "cat" or len(cat_node.inputs) != 2: @@ -430,7 +438,7 @@ def _match_rope_angle_table_source(graph: IRGraph, value_id: str, index_shape: t inv_freq_value_id: str | None = None position_value_id: str | None = None for input_id in angle_node.inputs: - const_id = _find_constant_ancestor(graph, input_id) + const_id = _stored_inv_freq_constant(graph, input_id) if const_id is not None: inv_freq_value_id = const_id continue