cactus-compute 2.2.0__tar.gz → 2.2.2__tar.gz

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cactus-compute
-Version: 2.2.0
+Version: 2.2.2
 Summary: On-device AI inference — LLM, vision, speech, embeddings, and RAG
 Author: Cactus Compute
 License: Cactus Compute License

cactus_compute-2.2.2/cactus/_version.py

@@ -0,0 +1 @@
+__version__ = "2.2.2"

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus/bindings/cactus.py

@@ -568,6 +568,11 @@ _lib.cactus_graph_rel_pos_bias.argtypes = [
     cactus_graph_t, cactus_node_t, cactus_node_t, ctypes.c_float, ctypes.POINTER(cactus_node_t)
 ]
 _lib.cactus_graph_rel_pos_bias.restype = ctypes.c_int
+_lib.cactus_graph_rel_pos_attention.argtypes = [
+    cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t, cactus_node_t,
+    ctypes.c_bool, cactus_node_t, ctypes.c_float, ctypes.c_size_t, ctypes.POINTER(cactus_node_t)
+]
+_lib.cactus_graph_rel_pos_attention.restype = ctypes.c_int
 _lib.cactus_graph_attention_int8_hybrid.argtypes = [
     cactus_graph_t, cactus_node_t, cactus_node_t, cactus_node_t, ctypes.c_float, ctypes.c_size_t,
     ctypes.POINTER(ctypes.c_int8), ctypes.POINTER(ctypes.c_int8),
@@ -2570,6 +2575,31 @@ class Graph:
             raise RuntimeError(_err("graph_rel_pos_bias failed"))
         return self._apply_backend(self._tensor_from_node(out.value), backend)
 
+    def rel_pos_attention(self, query, key, value, rel_query, relative_key, key_mask, scale, window_size=0, backend=None):
+        query = self._ensure_tensor(query)
+        key = self._ensure_tensor(key)
+        value = self._ensure_tensor(value)
+        rel_query = self._ensure_tensor(rel_query)
+        relative_key = self._ensure_tensor(relative_key)
+        key_mask = None if key_mask is None else self._ensure_tensor(key_mask)
+        out = cactus_node_t()
+        rc = _lib.cactus_graph_rel_pos_attention(
+            self.h,
+            cactus_node_t(query.id),
+            cactus_node_t(key.id),
+            cactus_node_t(value.id),
+            cactus_node_t(rel_query.id),
+            cactus_node_t(relative_key.id),
+            ctypes.c_bool(key_mask is not None),
+            cactus_node_t(0 if key_mask is None else key_mask.id),
+            ctypes.c_float(float(scale)),
+            ctypes.c_size_t(int(window_size)),
+            ctypes.byref(out),
+        )
+        if rc != 0:
+            raise RuntimeError(_err("graph_rel_pos_attention failed"))
+        return self._apply_backend(self._tensor_from_node(out.value), backend)
+
     def attention_int8_hybrid(self, query, key_new, value_new, scale, position_offset,
                               cached_keys, cached_values, k_scales, v_scales,
                               cache_len, num_kv_heads, head_dim, window_size=0, backend=None):

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus/convert/tests/test_audio_buckets.py

@@ -1,5 +1,6 @@
 from cactus.transpile.audio_preprocess import audio_bucket_frames
 from cactus.transpile.audio_preprocess import audio_bucket_seconds
+from cactus.transpile.audio_preprocess import audio_capture_frames
 
 
 def test_fibonacci_ladder_capped_at_model_window():
@@ -26,8 +27,18 @@ def test_buckets_are_sorted_and_unique():
     assert len(buckets) == len(set(buckets))
 
 
-def test_default_ladder_is_fibonacci_seconds():
-    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21)
+def test_default_ladder_seconds():
+    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 300, 377)
+
+
+def test_capture_window_is_420_seconds():
+    assert audio_capture_frames() == 42000
+
+
+def test_env_override_keeps_capture_window(monkeypatch):
+    monkeypatch.setenv("CACTUS_TRANSPILER_AUDIO_BUCKETS", "30,60")
+    assert audio_capture_frames() == 42000
+    assert audio_bucket_frames(audio_capture_frames()) == [3000, 6000, 42000]
 
 
 def test_env_override_replaces_ladder(monkeypatch):
@@ -43,9 +54,9 @@ def test_env_override_is_sorted_and_deduplicated(monkeypatch):
 
 def test_malformed_env_override_falls_back_to_default(monkeypatch):
     monkeypatch.setenv("CACTUS_TRANSPILER_AUDIO_BUCKETS", "not-a-number")
-    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21)
+    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 300, 377)
 
 
 def test_empty_env_override_falls_back_to_default(monkeypatch):
     monkeypatch.setenv("CACTUS_TRANSPILER_AUDIO_BUCKETS", "0,-3")
-    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21)
+    assert audio_bucket_seconds() == (1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 300, 377)

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus/convert/tests/test_parakeet_capture.py

@@ -19,9 +19,11 @@ def test_capture_window_is_independent_of_validation_audio():
     with patch.dict("os.environ", {"CACTUS_TRANSPILER_AUDIO_BUCKETS": ""}):
         specs = build_parakeet_tdt_component_specs(model, named_tensors={"input_features": features})
     encoders = specs[:-1]
-    assert [int(spec.metadata["audio_frames"]) for spec in encoders] == [100, 200, 300, 500, 800, 1300, 2100, 3000]
+    assert [int(spec.metadata["audio_frames"]) for spec in encoders] == [
+        100, 200, 300, 500, 800, 1300, 2100, 3400, 5500, 8900, 14400, 23300, 30000, 37700, 42000,
+    ]
     assert encoders[-1].component == "audio_encoder"
-    assert tuple(encoders[-1].example_inputs[0].shape) == (1, 3000, 128)
+    assert tuple(encoders[-1].example_inputs[0].shape) == (1, 42000, 128)
     assert torch.all(features == 2.0)
     assert not torch.all(encoders[-1].example_inputs[0] == 2.0)
 

cactus_compute-2.2.2/cactus/convert/tests/test_parakeet_local_attention.py

@@ -0,0 +1,125 @@
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+
+from cactus.transpile.tdt_runtime import ParakeetTDTSelfAttention
+from cactus.transpile.tdt_runtime import rel_pos_attention
+
+
+def _random_inputs(seq_len, rel_len, *, heads=2, dim=16, padded=3, seed=0):
+    generator = torch.Generator().manual_seed(seed)
+    shape = (1, seq_len, heads, dim)
+    tensors = [torch.randn(shape, generator=generator) for _ in range(4)]
+    relative_key = torch.randn((1, rel_len, heads, dim), generator=generator)
+    key_mask = torch.zeros((1, seq_len))
+    key_mask[:, seq_len - padded:] = -10000.0
+    return (*tensors, relative_key, key_mask)
+
+
+def _dense_reference(query, key, value, rel_query, relative_key, key_mask, scale, window):
+    seq_len = query.shape[1]
+    center = (relative_key.shape[1] - 1) // 2
+    out = torch.empty_like(query)
+    for t in range(seq_len):
+        for h in range(query.shape[2]):
+            scores = []
+            for j in range(seq_len):
+                if window > 0 and abs(t - j) > window:
+                    scores.append(float("-inf"))
+                    continue
+                content = torch.dot(query[0, t, h], key[0, j, h])
+                position = torch.dot(rel_query[0, t, h], relative_key[0, center - (t - j), h])
+                scores.append(float((content + position) * scale + key_mask[0, j]))
+            weights = torch.softmax(torch.tensor(scores), dim=0)
+            out[0, t, h] = weights @ value[0, :, h]
+    return out
+
+
+def test_rel_pos_attention_matches_dense_reference():
+    for seq_len, window, rel_len in ((40, 6, 13), (40, 0, 79), (300, 20, 41)):
+        inputs = _random_inputs(seq_len, rel_len)
+        expected = _dense_reference(*inputs, 0.25, window)
+        actual = rel_pos_attention(*inputs, 0.25, window)
+        torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
+
+
+def _attention_layer(hidden=32, heads=4):
+    torch.manual_seed(0)
+    state_dict = {f"layer.{name}.weight": torch.randn(hidden, hidden) * 0.2
+                  for name in ("linear_q", "linear_k", "linear_v", "linear_out", "linear_pos")}
+    state_dict["layer.pos_bias_u"] = torch.randn(heads, hidden // heads)
+    state_dict["layer.pos_bias_v"] = torch.randn(heads, hidden // heads)
+    config = SimpleNamespace(
+        hidden_dim=hidden,
+        attention_heads=heads,
+        attention_head_dim=hidden // heads,
+        attention_scale=(hidden // heads) ** -0.5,
+    )
+    return ParakeetTDTSelfAttention(config, "layer", state_dict).eval()
+
+
+def test_local_attention_exports_single_semantic_op():
+    from cactus.transpile.capture_pytorch import capture_model
+
+    class _Windowed(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.attention = _attention_layer()
+
+        def forward(self, x, mask):
+            return self.attention(x, mask, 4)
+
+    captured = capture_model(_Windowed().eval(), (torch.randn(1, 12, 32), torch.zeros(1, 1, 1, 12)), strict=False)
+    nodes = [
+        captured.ir_graph.nodes[node_id]
+        for node_id in captured.ir_graph.order
+        if captured.ir_graph.nodes[node_id].op == "rel_pos_attention"
+    ]
+    assert len(nodes) == 1
+    assert nodes[0].attrs == {"scale": 8 ** -0.5, "window_size": 4}
+    assert len(nodes[0].inputs) == 6
+
+
+def test_engine_op_matches_torch_reference():
+    from cactus.bindings.cactus import Graph
+
+    for seq_len, window, rel_len in ((150, 20, 41), (70, 0, 139)):
+        inputs = _random_inputs(seq_len, rel_len, dim=32, padded=5, seed=1)
+        inputs = tuple(tensor.half().float() for tensor in inputs)
+        expected = rel_pos_attention(*inputs, 32 ** -0.5, window)
+
+        graph = Graph()
+        nodes = [graph.input(tuple(tensor.shape)) for tensor in inputs]
+        out = graph.rel_pos_attention(*nodes, 32 ** -0.5, window)
+        for node, tensor in zip(nodes, inputs):
+            graph.set_input(node, tensor.numpy().astype(np.float16))
+        graph.execute()
+        np.testing.assert_allclose(out.numpy().astype(np.float32), expected.numpy(), atol=3e-3)
+
+
+def test_chunked_subsampling_matches_single_pass():
+    from cactus.transpile.tdt_runtime import ParakeetTDTPreEncode
+
+    torch.manual_seed(0)
+    channels, mels, hidden = 4, 16, 8
+    state_dict = {
+        "encoder.pre_encode.conv.0.weight": torch.randn(channels, 1, 3, 3),
+        "encoder.pre_encode.conv.0.bias": torch.randn(channels),
+        "encoder.pre_encode.out.weight": torch.randn(hidden, channels * 2),
+        "encoder.pre_encode.out.bias": torch.randn(hidden),
+    }
+    for index in (2, 5):
+        state_dict[f"encoder.pre_encode.conv.{index}.weight"] = torch.randn(channels, 1, 3, 3)
+        state_dict[f"encoder.pre_encode.conv.{index}.bias"] = torch.randn(channels)
+        state_dict[f"encoder.pre_encode.conv.{index + 1}.weight"] = torch.randn(channels, channels, 1, 1)
+        state_dict[f"encoder.pre_encode.conv.{index + 1}.bias"] = torch.randn(channels)
+    state_dict = {name: tensor * 0.3 for name, tensor in state_dict.items()}
+    config = SimpleNamespace(subsampling_conv_channels=channels, num_mel_bins=mels, hidden_dim=hidden)
+    pre_encode = ParakeetTDTPreEncode(config, state_dict).eval()
+    features = torch.randn(1, 7003, mels)
+    with torch.no_grad():
+        chunked = pre_encode(features)
+        single = pre_encode._subsample(features)
+    assert chunked.shape == single.shape
+    torch.testing.assert_close(chunked, single, atol=1e-5, rtol=1e-5)

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus/transpile/aten_ops.py

@@ -54,6 +54,7 @@ _EXACT_RULES: dict[str, str] = {
     "qwen2_moe_layer_gated.default": "qwen2_moe_layer_gated",
     "cactus_transpile.gemma4_moe_layer_gated.default": "gemma4_moe_layer_gated",
     "gemma4_moe_layer_gated.default": "gemma4_moe_layer_gated",
+    "rel_pos_attention.default": "rel_pos_attention",
     # Keep the legacy importer key until the importer itself is renamed.
     "aten.diff.default": "aten.diff.default",
 }

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus/transpile/audio_preprocess.py

@@ -21,7 +21,8 @@ _PARAKEET_HOP_LENGTH = 160
 _PARAKEET_PREEMPHASIS = 0.97
 _PARAKEET_LOG_FLOOR = np.float32(2**-24)
 _DEFAULT_MAX_AUDIO_SECONDS = 30.0
-_DEFAULT_AUDIO_BUCKET_SECONDS = (1, 2, 3, 5, 8, 13, 21)
+_DEFAULT_AUDIO_BUCKET_SECONDS = (1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 300, 377)
+_AUDIO_CAPTURE_SECONDS = 420
 
 
 def audio_duration_limit_seconds() -> float:
@@ -47,6 +48,10 @@ def audio_bucket_seconds() -> tuple[int, ...]:
     return positive
 
 
+def audio_capture_frames() -> int:
+    return _AUDIO_CAPTURE_SECONDS * (_PARAKEET_SAMPLE_RATE // _PARAKEET_HOP_LENGTH)
+
+
 def audio_bucket_frames(max_frames: int) -> list[int]:
     if max_frames <= 0:
         return []

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus/transpile/canonicalize/cleanup.py

@@ -132,6 +132,7 @@ FP16_ONLY_OUTPUT_OPS = {
     "lfm2_moe_layer_gated",
     "qwen2_moe_layer_gated",
     "gemma4_moe_layer_gated",
+    "rel_pos_attention",
     "rope",
 }
 

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus/transpile/importers.py

@@ -1860,6 +1860,23 @@ def import_gemma4_moe_layer_gated(ir, node, ctx, *, shape, dtype, torch_op):
                             op="gemma4_moe_layer_gated", has_expert_bias=False, activation="gelu")
 
 
+def import_rel_pos_attention(ir, node, ctx, *, shape, dtype, torch_op):
+    if len(node.args) != 8:
+        ctx.fail(f"rel_pos_attention expected 8 args, got {len(node.args)}")
+    ir_node = IRNode(
+        id=node_id(node),
+        op="rel_pos_attention",
+        inputs=[value_id(node.args[i], ctx) for i in range(6)],
+        outputs=[value_id(node, ctx)],
+        attrs={
+            "scale": float(extract_literals(node.args[6])),
+            "window_size": int(extract_literals(node.args[7])),
+        },
+        meta=_base_meta(shape, dtype, torch_op, node),
+    )
+    register_node(ir, ir_node, shape=shape, dtype=dtype)
+
+
 OP_IMPORTERS = {
     "arange": import_arange,
     "identity": import_identity,
@@ -1957,4 +1974,5 @@ OP_IMPORTERS = {
     "lfm2_moe_layer_gated": import_lfm2_moe_layer_gated,
     "qwen2_moe_layer_gated": import_qwen2_moe_layer_gated,
     "gemma4_moe_layer_gated": import_gemma4_moe_layer_gated,
+    "rel_pos_attention": import_rel_pos_attention,
 }

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus/transpile/lower.py

@@ -2218,6 +2218,26 @@ def _lower_ir_node(g: Graph, node: IRNode, env: dict[str, Any], ir: IRGraph) ->
             )
         return [g.rel_pos_bias(query, relative_key, float(node.attrs.get("scale", 1.0)))]
 
+    if op == "rel_pos_attention":
+        query, key, value, rel_query, relative_key, key_mask = (
+            _ensure_fp16_tensor(g, _tensor(env, input_id)) for input_id in node.inputs
+        )
+        mask_shape = (int(query.shape[0]), int(query.shape[1]))
+        if tuple(int(dim) for dim in key_mask.shape) != mask_shape:
+            key_mask = g.reshape(key_mask, mask_shape)
+        return [
+            g.rel_pos_attention(
+                query,
+                key,
+                value,
+                rel_query,
+                relative_key,
+                key_mask,
+                float(node.attrs["scale"]),
+                int(node.attrs.get("window_size", 0)),
+            )
+        ]
+
     if op in {"gated_deltanet_prefill", "gated_deltanet_decode"}:
         x = _tensor(env, node.inputs[0])
         qkv_weight = _tensor(env, node.inputs[1])

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus/transpile/ops.py

@@ -147,6 +147,7 @@ OPS: tuple[OpSchema, ...] = (
     OpSchema("rms_norm", num_inputs=2, attrs=("eps",), backend_op="rms_norm"),
     OpSchema("rope", num_inputs=1, attrs=("theta", "position_offset"), backend_op="rope"),
     OpSchema("rel_pos_bias", num_inputs=2, attrs=("scale",), backend_op="rel_pos_bias"),
+    OpSchema("rel_pos_attention", num_inputs=6, attrs=("scale", "window_size"), backend_op="rel_pos_attention"),
     OpSchema(
         "attention",
         num_inputs=3,

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus/transpile/tdt_runtime.py

@@ -14,11 +14,63 @@ import torch.nn.functional as F
 
 from cactus.transpile.component_pipeline import ComponentModuleSpec
 from cactus.transpile.audio_preprocess import audio_bucket_frames
+from cactus.transpile.audio_preprocess import audio_capture_frames
 from cactus.transpile.audio_preprocess import prepare_native_parakeet_audio_features
 from cactus.transpile.model_profiles import add_tensor_aliases
 from cactus.transpile.model_profiles import PARAKEET_TDT_PROFILE
 
 
+_FULL_ATTENTION_MAX_ENCODER_FRAMES = 688
+_LOCAL_ATTENTION_WINDOW = 256
+_REL_POS_ATTENTION_QUERY_BLOCK = 256
+_PRE_ENCODE_CHUNK_FRAMES = 3000
+_PRE_ENCODE_HALO_FRAMES = 16
+
+
+@torch.library.custom_op("cactus_transpile::rel_pos_attention", mutates_args=())
+def rel_pos_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    rel_query: torch.Tensor,
+    relative_key: torch.Tensor,
+    key_mask: torch.Tensor,
+    scale: float,
+    window: int,
+) -> torch.Tensor:
+    seq_len = int(query.shape[1])
+    rel_len = int(relative_key.shape[1])
+    center = (rel_len - 1) // 2
+    reach = seq_len - 1 if window <= 0 or window >= seq_len else window
+    q = query.permute(0, 2, 1, 3).float()
+    k = key.permute(0, 2, 1, 3).float()
+    v = value.permute(0, 2, 1, 3).float()
+    q_rel = rel_query.permute(0, 2, 1, 3).float()
+    r = relative_key.permute(0, 2, 1, 3).float()
+    mask = key_mask.float()
+    out = torch.empty_like(q)
+    for t0 in range(0, seq_len, _REL_POS_ATTENTION_QUERY_BLOCK):
+        t1 = min(seq_len, t0 + _REL_POS_ATTENTION_QUERY_BLOCK)
+        k0 = max(0, t0 - reach)
+        k1 = min(seq_len, t1 + reach)
+        offset = (
+            torch.arange(k0, k1, device=query.device).view(1, -1)
+            - torch.arange(t0, t1, device=query.device).view(-1, 1)
+        )
+        rel_index = (center + offset).clamp(0, rel_len - 1).expand(q.shape[0], q.shape[1], -1, -1)
+        content = q[:, :, t0:t1] @ k[:, :, k0:k1].transpose(-1, -2)
+        position = (q_rel[:, :, t0:t1] @ r.transpose(-1, -2)).gather(-1, rel_index)
+        scores = (content + position) * scale + mask[:, None, None, k0:k1]
+        scores = scores.masked_fill(offset.abs() > reach, float("-inf"))
+        out[:, :, t0:t1] = scores.softmax(dim=-1) @ v[:, :, k0:k1]
+    return out.permute(0, 2, 1, 3).to(query.dtype).contiguous()
+
+
+@rel_pos_attention.register_fake
+def _rel_pos_attention_fake(query, key, value, rel_query, relative_key, key_mask, scale, window):
+    return torch.empty_like(query)
+
+
 def _cfg_get(config: dict[str, Any], key: str, default: Any = None) -> Any:
     value = config.get(key, default)
     return default if value is None else value
@@ -403,19 +455,6 @@ def _relative_position_embeddings(*, seq_len: int, hidden_dim: int, device: torc
     return embeddings.to(dtype=dtype)
 
 
-def _relative_position_bias(query: torch.Tensor, relative_key: torch.Tensor, *, scale: float) -> torch.Tensor:
-    batch, heads, seq_len, _ = query.shape
-    scores = torch.matmul(query, relative_key.transpose(-1, -2))
-    rel_index = (
-        torch.arange(seq_len, device=query.device).view(seq_len, 1)
-        - torch.arange(seq_len, device=query.device).view(1, seq_len)
-        + (seq_len - 1)
-    )
-    rel_index = rel_index.view(1, 1, seq_len, seq_len).expand(batch, heads, seq_len, seq_len)
-    gathered = scores.gather(-1, rel_index)
-    return gathered * float(scale)
-
-
 class ParakeetTDTFeedForward(nn.Module):
     def __init__(self, config: ParakeetTDTConfig, prefix: str, state_dict: dict[str, torch.Tensor]):
         super().__init__()
@@ -457,7 +496,7 @@ class ParakeetTDTSelfAttention(nn.Module):
         self.pos_bias_v = nn.Parameter(state_dict[f"{prefix}.pos_bias_v"].clone())
         self.config = config
 
-    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, attn_mask: torch.Tensor, window: int = 0) -> torch.Tensor:
         batch, seq_len, _ = x.shape
         num_heads = self.config.attention_heads
         head_dim = self.config.attention_head_dim
@@ -466,39 +505,29 @@ class ParakeetTDTSelfAttention(nn.Module):
         k = self.linear_k(x).view(batch, seq_len, num_heads, head_dim)
         v = self.linear_v(x).view(batch, seq_len, num_heads, head_dim)
 
+        reach = window if window > 0 else seq_len - 1
         rel_pos = _relative_position_embeddings(
-            seq_len=seq_len,
+            seq_len=reach + 1,
             hidden_dim=self.config.hidden_dim,
             device=x.device,
             dtype=x.dtype,
         )
-        rel_k = self.linear_pos(rel_pos).view(1, 2 * seq_len - 1, num_heads, head_dim)
+        rel_k = self.linear_pos(rel_pos).view(1, 2 * reach + 1, num_heads, head_dim)
 
         q_u = q + self.pos_bias_u.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device)
         q_v = q + self.pos_bias_v.view(1, 1, num_heads, head_dim).to(dtype=x.dtype, device=x.device)
 
-        q_u_heads = q_u.permute(0, 2, 1, 3)
-        q_v_heads = q_v.permute(0, 2, 1, 3)
-        k_heads = k.permute(0, 2, 1, 3)
-        v_heads = v.permute(0, 2, 1, 3)
-        rel_k_heads = rel_k.permute(0, 2, 1, 3)
-
-        rel_bias = _relative_position_bias(
-            q_v_heads,
-            rel_k_heads,
-            scale=self.config.attention_scale,
-        )
-        rel_bias = rel_bias + attn_mask
-        attn = F.scaled_dot_product_attention(
-            q_u_heads,
-            k_heads,
-            v_heads,
-            attn_mask=rel_bias,
-            dropout_p=0.0,
-            is_causal=False,
+        attn = rel_pos_attention(
+            q_u,
+            k,
+            v,
+            q_v,
+            rel_k,
+            attn_mask.reshape(batch, seq_len),
+            self.config.attention_scale,
+            window,
         )
-        attn = attn.permute(0, 2, 1, 3).reshape(batch, seq_len, self.config.hidden_dim)
-        return self.linear_out(attn)
+        return self.linear_out(attn.reshape(batch, seq_len, self.config.hidden_dim))
 
 
 class ParakeetTDTConformerConv(nn.Module):
@@ -583,9 +612,9 @@ class ParakeetTDTEncoderLayer(nn.Module):
             module.bias.data.copy_(state_dict[f"{prefix}.{name}.bias"].to(dtype=module.bias.dtype))
         self.config = config
 
-    def forward(self, x: torch.Tensor, conv_mask: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, conv_mask: torch.Tensor, attn_mask: torch.Tensor, window: int = 0) -> torch.Tensor:
         x = x + 0.5 * self.feed_forward1(self.norm_feed_forward1(x), activation=self.config.encoder_hidden_act)
-        x = x + self.self_attn(self.norm_self_att(x), attn_mask)
+        x = x + self.self_attn(self.norm_self_att(x), attn_mask, window)
         x = x + self.conv(self.norm_conv(x), conv_mask)
         x = x + 0.5 * self.feed_forward2(self.norm_feed_forward2(x), activation=self.config.encoder_hidden_act)
         return self.norm_out(x)
@@ -624,6 +653,19 @@ class ParakeetTDTPreEncode(nn.Module):
     def forward(self, input_features: torch.Tensor) -> torch.Tensor:
         if input_features.ndim != 3:
             raise ValueError(f"expected input_features [batch, frames, mels], got {tuple(input_features.shape)}")
+        frames = int(input_features.shape[1])
+        if frames <= _PRE_ENCODE_CHUNK_FRAMES:
+            return self._subsample(input_features)
+        chunks = []
+        for start in range(0, frames, _PRE_ENCODE_CHUNK_FRAMES):
+            end = min(frames, start + _PRE_ENCODE_CHUNK_FRAMES)
+            low = max(0, start - _PRE_ENCODE_HALO_FRAMES)
+            out = self._subsample(input_features[:, low : min(frames, end + _PRE_ENCODE_HALO_FRAMES)])
+            first = (start - low) // 8
+            chunks.append(out[:, first : first + (end - start + 7) // 8])
+        return torch.cat(chunks, dim=1)
+
+    def _subsample(self, input_features: torch.Tensor) -> torch.Tensor:
         x = input_features.unsqueeze(1)
         x = F.relu(self.conv[0](x))
         x = F.relu(self.conv[3](self.conv[2](x)))
@@ -646,8 +688,9 @@ class ParakeetTDTEncoder(nn.Module):
         pad_mask = input_features_mask[:, :: self.config.subsampling_factor]
         conv_mask = pad_mask.unsqueeze(1)
         attn_mask = ((pad_mask - 1.0) * 10000.0).unsqueeze(1).unsqueeze(1)
+        window = 0 if int(x.shape[1]) <= _FULL_ATTENTION_MAX_ENCODER_FRAMES else _LOCAL_ATTENTION_WINDOW
         for layer in self.layers:
-            x = layer(x, conv_mask, attn_mask)
+            x = layer(x, conv_mask, attn_mask, window)
         return x
 
 
@@ -850,11 +893,12 @@ def build_parakeet_tdt_component_specs(
 ) -> list[ComponentModuleSpec]:
     features = named_tensors["input_features"]
     input_features = torch.randn(
-        (int(features.shape[0]), 3000, int(features.shape[2])),
+        (int(features.shape[0]), audio_capture_frames(), int(features.shape[2])),
         device=features.device,
         dtype=features.dtype,
     )
-    example_hidden = model.encoder(input_features, torch.ones_like(input_features[:, :, 0]))
+    example_features = input_features[:, : audio_bucket_frames(int(input_features.shape[1]))[0], :]
+    example_hidden = model.encoder(example_features, torch.ones_like(example_features[:, :, 0]))
     batch_size = int(example_hidden.shape[0])
     initial_states = model.initial_decoder_state(
         batch_size=batch_size,

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus/transpiler/Converter/input_processor.py

@@ -95,6 +95,7 @@ PROCESSOR_MAP = {
     "google/gemma-4-E2B": gemma4_processor,
     "openai/whisper-tiny": default_processor,
     "nvidia/parakeet-tdt-0.6b-v3": parakeet_processor,
+    "nvidia/parakeet-tdt-0.6b-v2": parakeet_processor,
     "LiquidAI/LFM2-VL-3B": default_processor,
     "Qwen/Qwen2.5-0.5B": default_processor,
     "LiquidAI/LFM2.5-8B-A1B": text_tokenizer_processor,

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus/transpiler/ModelProfiles/profiles.py

@@ -586,6 +586,7 @@ MODEL_ID_MAP = {
     "openai/whisper-tiny": WHISPER_PROFILE,
     "openai/whisper-small": WHISPER_PROFILE,
     "nvidia/parakeet-tdt-0.6b-v3": PARAKEET_PROFILE,
+    "nvidia/parakeet-tdt-0.6b-v2": PARAKEET_PROFILE,
     "LiquidAI/LFM2-VL-450M": LFM_VLM_PROFILE,
     "LiquidAI/LFM2-VL-3B": LFM_VLM_PROFILE,
     "Qwen/Qwen2.5-0.5B": QWEN2_5_0_5B_PROFILE,

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus_compute.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cactus-compute
-Version: 2.2.0
+Version: 2.2.2
 Summary: On-device AI inference — LLM, vision, speech, embeddings, and RAG
 Author: Cactus Compute
 License: Cactus Compute License

{cactus_compute-2.2.0 → cactus_compute-2.2.2}/cactus_compute.egg-info/SOURCES.txt

@@ -63,6 +63,7 @@ cactus/convert/tests/test_hessian.py
 cactus/convert/tests/test_lfm2_adapter.py
 cactus/convert/tests/test_naming_qdq.py
 cactus/convert/tests/test_parakeet_capture.py
+cactus/convert/tests/test_parakeet_local_attention.py
 cactus/convert/tests/test_policy.py
 cactus/convert/tests/test_qwen_adapter.py
 cactus/convert/tests/test_reports.py

cactus_compute-2.2.0/cactus/_version.py

@@ -1 +0,0 @@
-__version__ = "2.2.0"