fixed needle rope bug, moved repeated rope table instantiations from …

cactus-engine/tests/test_needle.cpp

@@ -0,0 +1,97 @@
+#include "test_utils.h"
+
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <iostream>
+#include <sstream>
+#include <string>
+
+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;
+}

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)
 

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:

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