From 46f6602fe880afdeb8fe48d32f3dbe01c70d8f94 Mon Sep 17 00:00:00 2001 From: Noah Cylich Date: Tue, 9 Jun 2026 22:15:33 -0700 Subject: [PATCH] transpiler: JAX attention causality + reduction precision, attention-output layout guard jax_semantic_rewrites.py: gate prefill is_causal to known causal-prefill components so a maskless self-attention isn't silently forced causal (full/bidirectional attention was being causal-masked); keep min/max reductions in native FP32 instead of legalizing through FP16, which overflows attention-mask sentinels and large logits and corrupts exact min/max selection (sum/mean stay FP16); strip the component-meta string. optimize_graph.py: don't fold an attention output permute when the attention output is itself a graph output, which left graph.outputs referencing a value no longer produced. Signed-off-by: Noah Cylich --- python/cactus/transpile/jax_semantic_rewrites.py | 8 +++++--- python/cactus/transpile/optimize_graph.py | 2 +- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/python/cactus/transpile/jax_semantic_rewrites.py b/python/cactus/transpile/jax_semantic_rewrites.py index 8ffe518f4..d0dd63f15 100644 --- a/python/cactus/transpile/jax_semantic_rewrites.py +++ b/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]) diff --git a/python/cactus/transpile/optimize_graph.py b/python/cactus/transpile/optimize_graph.py index c700c7128..245f25f62 100644 --- a/python/cactus/transpile/optimize_graph.py +++ b/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: