diff --git a/cactus-engine/src/model.cpp b/cactus-engine/src/model.cpp index 393dfb32d..443c7f1e9 100644 --- a/cactus-engine/src/model.cpp +++ b/cactus-engine/src/model.cpp @@ -1072,26 +1072,31 @@ Model::ChunkedPrefillResult Model::run_chunked_prefill(const std::vector 0 ? std::min(chunk_size, component_tokens) : component_tokens; if (effective_chunk != component_tokens) effective_chunk = component_tokens; size_t whole_chunks_end = (tokens.size() / effective_chunk) * effective_chunk; - const bool has_recurrent_state = [&]() { + auto any_cache_node = [&](auto predicate) { if (!decoder_prefill_->graph) return false; 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; - if (decoder_prefill_->graph->get_node_op_type(static_cast(node_id)) - == OpType::RECURRENT_CACHE_STATE) { - return true; - } + if (predicate(static_cast(node_id))) return true; } } return false; - }(); + }; + const bool has_recurrent_state = any_cache_node([&](size_t id) { + return decoder_prefill_->graph->get_node_op_type(id) == OpType::RECURRENT_CACHE_STATE; + }); if (has_recurrent_state && whole_chunks_end > effective_chunk) { whole_chunks_end = effective_chunk; } + const bool has_sliding_window_cache = any_cache_node([&](size_t id) { + return decoder_prefill_->graph->get_node_op_type(id) == OpType::KV_CACHE_STATE + && decoder_prefill_->graph->get_node_window_size(id) > 0; + }); const size_t tail_tokens = tokens.size() - whole_chunks_end; const size_t padding_cutoff = std::max(1, effective_chunk / 16); const bool pad_tail = family_ != "lfm2_vl" && !has_recurrent_state + && !has_sliding_window_cache && tail_tokens >= padding_cutoff; const size_t executable_tokens = whole_chunks_end + (pad_tail ? effective_chunk : 0); if (executable_tokens == 0) { diff --git a/cactus-graph/cactus_graph.h b/cactus-graph/cactus_graph.h index 7c77adf30..acb88baba 100644 --- a/cactus-graph/cactus_graph.h +++ b/cactus-graph/cactus_graph.h @@ -707,6 +707,7 @@ class CactusGraph { const std::vector& output_shape, const OpParams& params = {}); 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; void allocate_buffers(); size_t get_node_count() const; diff --git a/cactus-graph/src/builder.cpp b/cactus-graph/src/builder.cpp index 3f513d0d5..b54b0a92a 100644 --- a/cactus-graph/src/builder.cpp +++ b/cactus-graph/src/builder.cpp @@ -1257,6 +1257,10 @@ OpType CactusGraph::get_node_op_type(size_t node_id) const { return nodes_[node_index_map_.at(node_id)]->op_type; } +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::persistent(size_t source_node) { const auto& source_buffer = get_output_buffer(source_node); OpParams params; diff --git a/python/cactus/transpile/optimize_graph.py b/python/cactus/transpile/optimize_graph.py index 430fd0e7d..f26cbecb9 100644 --- a/python/cactus/transpile/optimize_graph.py +++ b/python/cactus/transpile/optimize_graph.py @@ -20,8 +20,10 @@ from cactus.transpile.fusion import match_rope from cactus.transpile.fusion import match_self_attention_block from cactus.transpile.fusion.common import producer +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 _find_constant_ancestor from cactus.transpile.graph_ir import IRGraph from cactus.transpile.graph_ir import IRNode from cactus.transpile.graph_ir import IRValue @@ -106,6 +108,9 @@ def optimize_graph(graph: IRGraph, *, max_passes: int = 8, config: FusionConfig break canonicalize_exported_graph(graph) + if precompute_rope_tables(graph): + canonicalize_exported_graph(graph) + annotate_gold_patterns(graph) _prune_unused_inputs(graph) verify_ir(graph) @@ -361,6 +366,122 @@ def fuse_rope(graph: IRGraph) -> bool: return changed +def precompute_rope_tables(graph: IRGraph) -> bool: + """Replace the runtime ``cos/sin(position * inv_freq)`` rope angle with an + fp64-precomputed fp16 table gathered by position id (the fp16 angle matmul + cannot represent positions > 2048, randomising cos/sin past long context).""" + + changed = False + for node_id in list(graph.order): + node = graph.nodes.get(node_id) + if node is None or node.op not in {"scalar_cos", "scalar_sin"} or len(node.inputs) != 1: + continue + # Only the text decoder's rope. Vision/audio encoders carry cos/sin with + # the same shape but an axial/2D position this scalar table cannot model. + if node.meta.get("component") != "decoder": + continue + + output_id = node.outputs[0] + output_value = graph.values.get(output_id) + if output_value is None or output_value.shape is None: + continue + index_shape = tuple(int(dim) for dim in output_value.shape[:-1]) + + match = _match_rope_angle_table_source(graph, node.inputs[0], index_shape) + if match is None: + continue + inv_freq_value_id, position_value_id = match + + inv_freq = graph.constants[inv_freq_value_id].detach().cpu().to(torch.float64).reshape(-1) + head_dim = int(inv_freq.numel()) * 2 + max_seq = _rope_table_max_seq(graph) + positions = torch.arange(max_seq, dtype=torch.float64).reshape(max_seq, 1) + freqs = positions * inv_freq.reshape(1, -1) + emb = torch.cat((freqs, freqs), dim=-1) + table = (torch.cos(emb) if node.op == "scalar_cos" else torch.sin(emb)).to(torch.float16) + assert table.shape[-1] == head_dim + + table_value_id = _materialize_rope_table_constant(graph, node.id, table) + output_value.dtype = "fp16" + node.op = "embedding" + node.inputs = [table_value_id, position_value_id] + node.attrs = {} + node.kind = "generic" + changed = True + + if changed: + rebuild_graph(graph) + return changed + + +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: + return None + if strip_layout_passthrough(graph, cat_node.inputs[0]) != strip_layout_passthrough(graph, cat_node.inputs[1]): + return None + angle_node = producer(graph, strip_layout_passthrough(graph, cat_node.inputs[0])) + if angle_node is not None and angle_node.op == "permute": + angle_node = producer(graph, strip_layout_passthrough(graph, angle_node.inputs[0])) + if angle_node is None or angle_node.op != "matmul" or len(angle_node.inputs) != 2: + return None + + 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) + if const_id is not None: + inv_freq_value_id = const_id + continue + position_value_id = _trace_integer_position_source(graph, input_id, index_shape) + + if inv_freq_value_id is None or position_value_id is None: + return None + return inv_freq_value_id, position_value_id + + +def _trace_integer_position_source(graph: IRGraph, value_id: str, index_shape: tuple[int, ...]) -> str | None: + current = value_id + visited: set[str] = set() + while current not in visited: + visited.add(current) + value = graph.values.get(current) + if ( + value is not None + and str(value.dtype) in {"int32", "int64", "i32", "i64"} + and value.shape is not None + and tuple(int(dim) for dim in value.shape) == index_shape + ): + return current + node = producer(graph, current) + if node is None or node.op not in {"precision_cast", "view", "reshape", "unsqueeze", "squeeze", "contiguous"} or len(node.inputs) != 1: + return None + current = node.inputs[0] + return None + + +def _rope_table_max_seq(graph: IRGraph) -> int: + max_seq = _coerce_optional_int(graph.meta.get("max_cache_seq_len")) + if max_seq is None or max_seq <= 0: + max_seq = _coerce_optional_int(graph.meta.get("max_position_embeddings")) + if max_seq is None or max_seq <= 0: + raise ValueError("rope table precompute requires max_cache_seq_len or max_position_embeddings in graph.meta") + return int(max_seq) + + +def _materialize_rope_table_constant(graph: IRGraph, node_id: str, table: torch.Tensor) -> str: + value_id = f"c_rope_table_{node_id}" + graph.constants[value_id] = table + graph.values[value_id] = IRValue( + id=value_id, + shape=tuple(int(dim) for dim in table.shape), + dtype=dtype_to_ir(table.dtype), + producer=None, + users=[], + ) + return value_id + + def fuse_attention(graph: IRGraph) -> bool: changed = False for node_id in list(graph.order): diff --git a/python/tests/test_precompute_rope_tables.py b/python/tests/test_precompute_rope_tables.py new file mode 100644 index 000000000..ae8af1153 --- /dev/null +++ b/python/tests/test_precompute_rope_tables.py @@ -0,0 +1,390 @@ +import json +from pathlib import Path + +import pytest +import torch + +from cactus.transpile.canonicalize.utils import rebuild_graph +from cactus.transpile.graph_ir import IRGraph +from cactus.transpile.graph_ir import IRNode +from cactus.transpile.graph_ir import IRValue +from cactus.transpile.optimize_graph import precompute_rope_tables + + +WEIGHTS = Path(__file__).resolve().parents[2] / "weights" +GEMMA_BUNDLE = WEIGHTS / "gemma-4-e2b-it" +GEMMA_16K_BUNDLE = WEIGHTS / "gemma-4-e2b-16k" +QWEN_BUNDLE = WEIGHTS / "qwen3-0.6b" + +TEST_MAX_SEQ = 16384 + +# inv_freq = 1 / theta ** (arange(0, dim, 2) / dim), exactly the constant the +# real graph carries. The serialized optimized_ir drops multi-element tensor +# payloads, so we re-materialize the genuine per-theta inv_freq. The (theta, +# head_dim) spec is read from the bundle's config.json at test time (not +# hardcoded) so a model shipped with the wrong rope_theta / head_dim fails here +# instead of being matched against a self-referential reconstruction. +def _gemma_inv_freq_spec(bundle: Path) -> dict[str, tuple[float, int]]: + text = json.loads((bundle / "config.json").read_text())["text_config"] + rope = text["rope_parameters"] + return { + "v_module_backbone_rotary_emb_sliding_attention_inv_freq": ( + float(rope["sliding_attention"]["rope_theta"]), + int(text["head_dim"]), + ), + "v_module_backbone_rotary_emb_full_attention_inv_freq": ( + float(rope["full_attention"]["rope_theta"]), + int(text["global_head_dim"]), + ), + } + + +def _qwen_inv_freq_spec(bundle: Path) -> dict[str, tuple[float, int]]: + cfg = json.loads((bundle / "config.json").read_text()) + return { + "v_module_backbone_rotary_emb_inv_freq": ( + float(cfg["rope_theta"]), + int(cfg["head_dim"]), + ), + } + + +# The IR these tests load lives in transpiled weight bundles under weights/ +# (gitignored, ~9MB), so skip when they are absent (e.g. CI) rather than failing +# collection; run locally after transpiling the bundles. +_BUNDLES_PRESENT = (GEMMA_BUNDLE / "config.json").exists() and (QWEN_BUNDLE / "config.json").exists() +pytestmark = pytest.mark.skipif( + not _BUNDLES_PRESENT, reason="requires transpiled weight bundles under weights/ (gitignored)" +) + +GEMMA_INV_FREQ = _gemma_inv_freq_spec(GEMMA_BUNDLE) if _BUNDLES_PRESENT else {} +GEMMA_16K_INV_FREQ = _gemma_inv_freq_spec(GEMMA_16K_BUNDLE) if (GEMMA_16K_BUNDLE / "config.json").exists() else {} +QWEN_INV_FREQ = _qwen_inv_freq_spec(QWEN_BUNDLE) if _BUNDLES_PRESENT else {} + + +def _inv_freq(theta: float, head_dim: int) -> torch.Tensor: + return 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + + +def test_gemma_inv_freq_spec_matches_config() -> None: + """Guard the reconstruction source: the (theta, head_dim) the parity tests + feed the pass must equal the bundle's real config.json rope_parameters, so a + wrong-theta or wrong-head_dim bundle fails instead of self-referentially + passing.""" + text = json.loads((GEMMA_BUNDLE / "config.json").read_text())["text_config"] + rope = text["rope_parameters"] + assert GEMMA_INV_FREQ == { + "v_module_backbone_rotary_emb_sliding_attention_inv_freq": ( + float(rope["sliding_attention"]["rope_theta"]), + int(text["head_dim"]), + ), + "v_module_backbone_rotary_emb_full_attention_inv_freq": ( + float(rope["full_attention"]["rope_theta"]), + int(text["global_head_dim"]), + ), + } + sliding = GEMMA_INV_FREQ["v_module_backbone_rotary_emb_sliding_attention_inv_freq"] + full = GEMMA_INV_FREQ["v_module_backbone_rotary_emb_full_attention_inv_freq"] + assert sliding[0] != full[0], "sliding and full thetas must be distinct (dual-theta)" + assert sliding[1] != full[1], "sliding and full head_dims must be distinct" + + +def _load_ir(path: Path, inv_freq_spec: dict[str, tuple[float, int]], *, max_seq: int = TEST_MAX_SEQ) -> IRGraph: + payload = json.loads(path.read_text()) + graph = payload["graph"] + values = { + vid: IRValue( + id=vid, + shape=tuple(v["shape"]) if v.get("shape") is not None else None, + dtype=v.get("dtype"), + ) + for vid, v in graph["values"].items() + } + nodes: dict[str, IRNode] = {} + order: list[str] = [] + for n in graph["nodes"]: + nodes[n["id"]] = IRNode( + id=n["id"], + op=n["op"], + inputs=list(n["inputs"]), + outputs=list(n["outputs"]), + attrs=dict(n.get("attrs", {})), + meta=dict(n.get("meta", {})), + kind=n.get("kind", "generic"), + ) + order.append(n["id"]) + constants = { + cid: _inv_freq(*inv_freq_spec[cid]) + for cid in graph["constants"] + if cid in inv_freq_spec + } + meta = dict(graph["meta"]) + meta["max_cache_seq_len"] = max_seq + ir = IRGraph( + values=values, + nodes=nodes, + order=order, + inputs=list(graph["inputs"]), + outputs=list(graph["outputs"]), + constants=constants, + meta=meta, + ) + rebuild_graph(ir) + return ir + + +def _cos_sin_inputs(payload_path: Path) -> dict[str, str]: + """Map each scalar_cos/scalar_sin output id -> the int position graph input the + original angle matmul consumed, mirroring the pass's index_shape constraint + (the int value whose shape == cos/sin output shape minus its last dim).""" + graph = json.loads(payload_path.read_text())["graph"] + producer = {o: n for n in graph["nodes"] for o in n["outputs"]} + values = graph["values"] + + def find_position(value_id: str, index_shape: tuple[int, ...]) -> str | None: + seen: set[str] = set() + stack = [value_id] + while stack: + cur = stack.pop() + if cur in seen: + continue + seen.add(cur) + val = values.get(cur, {}) + dtype = str(val.get("dtype")) + shape = val.get("shape") + if ( + dtype in {"int32", "int64", "i32", "i64"} + and shape is not None + and tuple(int(d) for d in shape) == index_shape + ): + return cur + node = producer.get(cur) + if node is None: + continue + stack.extend(node["inputs"]) + return None + + result: dict[str, str] = {} + for n in graph["nodes"]: + if n["op"] in {"scalar_cos", "scalar_sin"}: + out_shape = values[n["outputs"][0]]["shape"] + index_shape = tuple(int(d) for d in out_shape[:-1]) + result[n["outputs"][0]] = find_position(n["inputs"][0], index_shape) + return result + + +def _rope_nodes(ir: IRGraph) -> tuple[list[IRNode], list[str]]: + embeddings = [ + n + for n in ir.nodes.values() + if n.op == "embedding" and n.inputs and n.inputs[0].startswith("c_rope_table_") + ] + tables = [c for c in ir.constants if c.startswith("c_rope_table_")] + return embeddings, tables + + +# --------------------------------------------------------------------------- # +# fires +# --------------------------------------------------------------------------- # + + +def test_gemma_decoder_step_rewrites_four_cos_sin() -> None: + ir = _load_ir(GEMMA_BUNDLE / "optimized_ir_decoder_step.json", GEMMA_INV_FREQ) + assert precompute_rope_tables(ir) is True + embeddings, tables = _rope_nodes(ir) + assert len(embeddings) == 4 + assert len(tables) == 4 + assert all(n.op != "scalar_cos" and n.op != "scalar_sin" for n in ir.nodes.values()) + + +def test_gemma_prefill_chunk_rewrites_eight_cos_sin() -> None: + ir = _load_ir(GEMMA_BUNDLE / "optimized_ir_decoder_prefill_chunk.json", GEMMA_INV_FREQ) + assert precompute_rope_tables(ir) is True + embeddings, tables = _rope_nodes(ir) + assert len(embeddings) == 8 + assert len(tables) == 8 + + +def test_gemma_media_step_rewrites_its_rope_tables() -> None: + path = GEMMA_16K_BUNDLE / "components" / "decoder_media_step" / "optimized_ir.json" + if not path.exists(): + path = GEMMA_16K_BUNDLE / "optimized_ir_decoder_media_step.json" + if not path.exists(): + pytest.skip( + "gemma4 component pipeline emits no decoder_media_step component " + "(only lm_encoder_media_step, which carries no rope angle path); " + "this guards future models that do emit a decoder media step" + ) + ir = _load_ir(path, GEMMA_16K_INV_FREQ) + assert precompute_rope_tables(ir) is True + embeddings, tables = _rope_nodes(ir) + assert len(embeddings) >= 4 + assert len(tables) >= 4 + widths = {int(ir.constants[t].shape[-1]) for t in tables} + assert 256 in widths and 512 in widths + + +# --------------------------------------------------------------------------- # +# two thetas: both widths exist and each table is the correct fp64 cos/sin +# --------------------------------------------------------------------------- # + + +def _table_by_width(ir: IRGraph) -> dict[int, dict[str, torch.Tensor]]: + out: dict[int, dict[str, torch.Tensor]] = {} + for node in ir.nodes.values(): + if node.op != "embedding" or not node.inputs[0].startswith("c_rope_table_"): + continue + table = ir.constants[node.inputs[0]] + width = int(table.shape[-1]) + kind = "cos" if "cos" in node.inputs[0] else "sin" + out.setdefault(width, {})[kind] = table + return out + + +def test_gemma_two_thetas_tables_match_fresh_fp64_cos_sin() -> None: + ir = _load_ir(GEMMA_BUNDLE / "optimized_ir_decoder_step.json", GEMMA_INV_FREQ) + precompute_rope_tables(ir) + by_width = _table_by_width(ir) + assert set(by_width.keys()) == {256, 512} + + expected = { + head_dim: _inv_freq(theta, head_dim).to(torch.float64) + for theta, head_dim in GEMMA_INV_FREQ.values() + } + assert set(expected.keys()) == {256, 512} + probe_positions = [0, 1, 2000, 2049, TEST_MAX_SEQ - 1] + for width, inv in expected.items(): + positions = torch.tensor(probe_positions, dtype=torch.float64).reshape(-1, 1) + freqs = positions * inv.reshape(1, -1) + emb = torch.cat((freqs, freqs), dim=-1) + exp_cos = torch.cos(emb).to(torch.float16) + exp_sin = torch.sin(emb).to(torch.float16) + got_cos = by_width[width]["cos"][probe_positions] + got_sin = by_width[width]["sin"][probe_positions] + assert torch.equal(got_cos, exp_cos), f"cos width {width}" + assert torch.equal(got_sin, exp_sin), f"sin width {width}" + + +def test_gemma_width256_uses_sliding_not_global_inv_freq() -> None: + """Cross-check the per-theta binding: the width-256 table must NOT match the + global theta and vice versa.""" + ir = _load_ir(GEMMA_BUNDLE / "optimized_ir_decoder_step.json", GEMMA_INV_FREQ) + precompute_rope_tables(ir) + by_width = _table_by_width(ir) + + wrong = _inv_freq(1000000.0, 256).to(torch.float64) # global theta on sliding width + positions = torch.tensor([2049.0], dtype=torch.float64).reshape(-1, 1) + freqs = positions * wrong.reshape(1, -1) + emb = torch.cat((freqs, freqs), dim=-1) + wrong_cos = torch.cos(emb).to(torch.float16) + assert not torch.equal(by_width[256]["cos"][[2049]], wrong_cos) + + +# --------------------------------------------------------------------------- # +# layout +# --------------------------------------------------------------------------- # + + +def test_gemma_table_width_matches_output_and_halves_are_byte_equal() -> None: + payload = json.loads((GEMMA_BUNDLE / "optimized_ir_decoder_step.json").read_text()) + out_width = { + n["outputs"][0]: int(payload["graph"]["values"][n["outputs"][0]]["shape"][-1]) + for n in payload["graph"]["nodes"] + if n["op"] in {"scalar_cos", "scalar_sin"} + } + ir = _load_ir(GEMMA_BUNDLE / "optimized_ir_decoder_step.json", GEMMA_INV_FREQ) + precompute_rope_tables(ir) + for node in ir.nodes.values(): + if node.op != "embedding" or not node.inputs[0].startswith("c_rope_table_"): + continue + out_id = node.outputs[0] + table = ir.constants[node.inputs[0]] + width = int(table.shape[-1]) + assert width == out_width[out_id] + half = width // 2 + assert torch.equal(table[:, :half], table[:, half:]) + + +# --------------------------------------------------------------------------- # +# position binding +# --------------------------------------------------------------------------- # + + +def test_gemma_embedding_index_is_original_position_input() -> None: + path = GEMMA_BUNDLE / "optimized_ir_decoder_step.json" + pos_for_output = _cos_sin_inputs(path) + ir = _load_ir(path, GEMMA_INV_FREQ) + precompute_rope_tables(ir) + for node in ir.nodes.values(): + if node.op != "embedding" or not node.inputs[0].startswith("c_rope_table_"): + continue + out_id = node.outputs[0] + assert node.inputs[1] == pos_for_output[out_id] + assert set(pos_for_output.values()) == {"v_args_2"} + + +# --------------------------------------------------------------------------- # +# Qwen regression: single width table +# --------------------------------------------------------------------------- # + + +def test_qwen_single_width_table_still_correct() -> None: + path = QWEN_BUNDLE / "optimized_ir_decoder_step.json" + pos_for_output = _cos_sin_inputs(path) + ir = _load_ir(path, QWEN_INV_FREQ) + assert precompute_rope_tables(ir) is True + embeddings, tables = _rope_nodes(ir) + assert len(embeddings) == 2 + assert len(tables) == 2 + + widths = {int(ir.constants[t].shape[-1]) for t in tables} + qwen_theta, qwen_head_dim = QWEN_INV_FREQ["v_module_backbone_rotary_emb_inv_freq"] + assert widths == {qwen_head_dim} + + inv = _inv_freq(qwen_theta, qwen_head_dim).to(torch.float64) + positions = torch.tensor([0.0, 1.0, 2049.0, TEST_MAX_SEQ - 1], dtype=torch.float64).reshape(-1, 1) + freqs = positions * inv.reshape(1, -1) + emb = torch.cat((freqs, freqs), dim=-1) + for node in ir.nodes.values(): + if node.op != "embedding" or not node.inputs[0].startswith("c_rope_table_"): + continue + table = ir.constants[node.inputs[0]] + ref = (torch.cos(emb) if "cos" in node.inputs[0] else torch.sin(emb)).to(torch.float16) + idx = [0, 1, 2049, TEST_MAX_SEQ - 1] + assert torch.equal(table[idx], ref) + assert node.inputs[1] == pos_for_output[node.outputs[0]] + + +# --------------------------------------------------------------------------- # +# inert: lm_encoder has no rope angle path +# --------------------------------------------------------------------------- # + + +def test_lm_encoder_is_inert() -> None: + ir = _load_ir(GEMMA_BUNDLE / "optimized_ir_lm_encoder_step.json", GEMMA_INV_FREQ) + assert precompute_rope_tables(ir) is False + embeddings, tables = _rope_nodes(ir) + assert tables == [] + + +def test_vision_encoder_is_inert() -> None: + """The vision tower's axial rope has the same cat(halves) <- matmul(inv_freq, + position) shape as decoder rope, but its position is a mid-graph ``select`` + (not a graph input). The pass must NOT fire, or it corrupts vision RoPE. The + inv_freq constant is materialized so the pass is rejected by the position + graph-input check, not by a missing constant.""" + path = GEMMA_BUNDLE / "optimized_ir_vision_encoder.json" + if not path.exists(): + pytest.skip("no vision encoder component for this bundle") + graph = json.loads(path.read_text())["graph"] + spec = { + cid: (10000.0, int(graph["values"][cid]["shape"][0]) * 2) + for cid in graph["constants"] + if "inv_freq" in cid + } + assert spec, "expected a vision rope inv_freq constant" + assert any(n["op"] in {"scalar_cos", "scalar_sin"} for n in graph["nodes"]), "vision encoder should have rope trig" + ir = _load_ir(path, spec) + assert precompute_rope_tables(ir) is False + _, tables = _rope_nodes(ir) + assert tables == []