transpiler: JAX attention causality + reduction precision, attention-output layout guard

python/cactus/transpile/jax_semantic_rewrites.py

@@ -685,7 +685,9 @@ def _rewrite_jax_prefill_attentions(graph: IRGraph) -> None:
         node = graph.nodes[node_id]
         if node.op not in {"permute", "view", "reshape"}:
             continue
-        is_decoder_step = str(graph.meta.get("component", "") or "").lower() == "decoder_step"
+        component = str(graph.meta.get("component", "") or "").strip().lower()
+        is_decoder_step = component == "decoder_step"
+        is_causal_prefill = component in {"decoder_prefill_chunk", "decoder_media_step"}
         if node.op == "permute" and tuple(int(dim) for dim in node.attrs.get("permutation", ())) != (0, 2, 1, 3):
             continue
         if node.op in {"view", "reshape"} and not is_decoder_step:
@@ -733,7 +735,7 @@ def _rewrite_jax_prefill_attentions(graph: IRGraph) -> None:
             additive_mask=additive_mask,
             scale=scale,
             output_layout="bthd",
-            is_causal=mask_id is None and is_self_attention,
+            is_causal=is_causal_prefill and mask_id is None and is_self_attention,
             rewritten_from="jax_prefill_attention",
         )
 
@@ -787,7 +789,7 @@ def _legalize_remaining_jax_reductions(graph: IRGraph) -> None:
     new_order: list[str] = []
     for node_id in list(graph.order):
         node = graph.nodes[node_id]
-        if node.op not in {"sum", "mean", "min", "max"} or len(node.inputs) != 1 or len(node.outputs) != 1:
+        if node.op not in {"sum", "mean"} or len(node.inputs) != 1 or len(node.outputs) != 1:
             new_order.append(node_id)
             continue
         input_value = graph.values.get(node.inputs[0])

python/cactus/transpile/optimize_graph.py

@@ -669,7 +669,7 @@ def normalize_attention_layouts(graph: IRGraph) -> bool:
         output_id = node.outputs[0]
         output_value = graph.values.get(output_id)
         output_users = list(output_value.users) if output_value is not None else []
-        if len(output_users) != 1:
+        if len(output_users) != 1 or output_id in graph.outputs:
             continue
         output_user = graph.nodes.get(output_users[0])
         if output_user is None or output_user.op != "permute" or len(output_user.inputs) != 1 or len(output_user.outputs) != 1: