From 708a0204e8a3b347f5c06ceee1abd22dc77614b1 Mon Sep 17 00:00:00 2001 From: G Ramalingam Date: Wed, 22 Apr 2026 16:48:25 +0000 Subject: [PATCH 1/4] Refactor ir.Function builders to use build_function from onnxscript 0.7.0 Replace manual ir.Graph + GraphBuilder + ir.Function construction boilerplate with the new builder.build_function() utility. This eliminates ~25 lines of setup per function and automatically handles initializer lifting to Constant nodes. Updated files: - causal_conv.py: body extracted to trace function - linear_attention.py: body extracted, conditional inputs preserved - packed_multi_head_attention.py: body extracted, raw ir.Node via op._builder._graph.append() for ref_attr_name forwarding - skip_layer_normalization.py: both functions refactored similarly - pyproject.toml: bumped onnxscript requirement to >=0.7.0 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam --- pyproject.toml | 2 +- src/mobius/functions/causal_conv.py | 124 ++++----- src/mobius/functions/linear_attention.py | 263 +++++++++--------- .../functions/packed_multi_head_attention.py | 228 ++++++++------- .../functions/skip_layer_normalization.py | 160 +++++------ 5 files changed, 376 insertions(+), 401 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 458ec521..c90b906f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,7 @@ dependencies = [ "numpy>=1.24.0", "onnx_ir>=0.1.0", "onnx-shape-inference==0.1.7", - "onnxscript>=0.6.0.dev", + "onnxscript>=0.7.0", "torch>=2.1.0", "tqdm", ] diff --git a/src/mobius/functions/causal_conv.py b/src/mobius/functions/causal_conv.py index e9699161..f759fd7b 100644 --- a/src/mobius/functions/causal_conv.py +++ b/src/mobius/functions/causal_conv.py @@ -20,9 +20,6 @@ DOMAIN = "com.microsoft" -# TODO(justinchuby): Simplify function creation boilerplate - - def causal_conv_nd_with_state( *, kernel_size: int, @@ -94,80 +91,71 @@ def causal_conv_nd_with_state( # The temporal axis = last spatial dim = index (ndim + 1) in (B, D, *spatial). temporal_axis = ndim + 1 - input_val = ir.Value(name="input") - weight_val = ir.Value(name="weight") - bias_val = ir.Value(name="bias") - conv_state_val = ir.Value(name="conv_state") - - graph = ir.Graph( - inputs=[input_val, weight_val, bias_val, conv_state_val], - outputs=[], - nodes=[], - name="CausalConvWithState_body", - opset_imports={"": OPSET_VERSION}, - ) - gb = builder.GraphBuilder(graph) - op = gb.op - - # Step 1: Prepend conv_state along the temporal axis. - conv_input = op.Concat(conv_state_val, input_val, axis=temporal_axis) - - # Step 2: Extract new carry state — last K-1 positions of conv_input. - total_len = op.Gather(op.Shape(conv_input), op.Constant(value_int=temporal_axis), axis=0) - state_start = op.Sub(total_len, op.Constant(value_int=state_width)) - present_state = op.Slice( - conv_input, - op.Reshape(state_start, op.Constant(value_ints=[1])), - op.Reshape(total_len, op.Constant(value_ints=[1])), - op.Constant(value_ints=[temporal_axis]), - ) - present_state.name = "present_state" - - # Step 3: Depthwise N-d Conv (group = channels, no padding — already prepended). - kernel_shape = [kernel_size] * ndim - dilations = [1] * ndim - pads = [0] * (2 * ndim) - conv_out = op.Conv( - conv_input, - weight_val, - kernel_shape=kernel_shape, - dilations=dilations, - pads=pads, - group=channels, - ) + def body(op, input_val, weight_val, bias_val, conv_state_val): + # Step 1: Prepend conv_state along the temporal axis. + conv_input = op.Concat(conv_state_val, input_val, axis=temporal_axis) - # Step 4: Add bias — reshape to (1, D, *[1]*ndim) for broadcasting. - bias_shape = [1, -1] + [1] * ndim - bias_reshaped = op.Reshape(bias_val, op.Constant(value_ints=bias_shape)) - conv_out = op.Add(conv_out, bias_reshaped) - - # Step 5: Apply activation. - if activation in ("silu", "swish"): - output = op.Mul(conv_out, op.Sigmoid(conv_out)) - elif activation == "none": - output = conv_out - else: - raise ValueError( - f"Unsupported activation: {activation!r}. Expected 'silu', 'swish', or 'none'." + # Step 2: Extract new carry state — last K-1 positions of conv_input. + total_len = op.Gather( + op.Shape(conv_input), op.Constant(value_int=temporal_axis), axis=0 + ) + state_start = op.Sub(total_len, op.Constant(value_int=state_width)) + present_state = op.Slice( + conv_input, + op.Reshape(state_start, op.Constant(value_ints=[1])), + op.Reshape(total_len, op.Constant(value_ints=[1])), + op.Constant(value_ints=[temporal_axis]), + ) + present_state.name = "present_state" + + # Step 3: Depthwise N-d Conv (group = channels, no padding — already prepended). + kernel_shape = [kernel_size] * ndim + dilations = [1] * ndim + pads = [0] * (2 * ndim) + conv_out = op.Conv( + conv_input, + weight_val, + kernel_shape=kernel_shape, + dilations=dilations, + pads=pads, + group=channels, ) - output.name = "output" - graph.outputs.extend([output, present_state]) + # Step 4: Add bias — reshape to (1, D, *[1]*ndim) for broadcasting. + bias_shape = [1, -1] + [1] * ndim + bias_reshaped = op.Reshape(bias_val, op.Constant(value_ints=bias_shape)) + conv_out = op.Add(conv_out, bias_reshaped) + + # Step 5: Apply activation. + if activation in ("silu", "swish"): + output = op.Mul(conv_out, op.Sigmoid(conv_out)) + elif activation == "none": + output = conv_out + else: + raise ValueError( + f"Unsupported activation: {activation!r}. Expected 'silu', 'swish', or 'none'." + ) + output.name = "output" + + return output, present_state # NOTE: Do not set ``overload`` here — call sites (op.CausalConvWithState) # do not set an overload, so setting one on the function would prevent the # serializer from matching nodes to this function definition. - return ir.Function( + return builder.build_function( + body, + [ + ir.Value(name="input"), + ir.Value(name="weight"), + ir.Value(name="bias"), + ir.Value(name="conv_state"), + ], domain=DOMAIN, name="CausalConvWithState", - graph=graph, - attributes={ - "activation": ir.Attr( - "activation", - ir.AttributeType.STRING, - activation, - ), - }, + attributes=[ + ir.Attr("activation", ir.AttributeType.STRING, activation), + ], + opset_imports={"": OPSET_VERSION}, ) diff --git a/src/mobius/functions/linear_attention.py b/src/mobius/functions/linear_attention.py index 7daa5e39..a5007fbc 100644 --- a/src/mobius/functions/linear_attention.py +++ b/src/mobius/functions/linear_attention.py @@ -114,7 +114,7 @@ def linear_attention( key = ir.Value(name="key") # (B, T, q_num_heads * d_k) value = ir.Value(name="value") # (B, T, kv_num_heads * d_v) past_state = ir.Value(name="past_state") - inputs = [query, key, value, past_state] + inputs: list[ir.Value] = [query, key, value, past_state] decay: ir.Value | None = None beta: ir.Value | None = None if uses_decay: @@ -124,145 +124,146 @@ def linear_attention( beta = ir.Value(name="beta") inputs.append(beta) - # --- Build function body graph --- - graph = ir.Graph( - inputs=inputs, - outputs=[], - nodes=[], - name=f"LinearAttention_{update_rule}_body", - opset_imports={"": OPSET_VERSION}, - ) - gb = builder.GraphBuilder(graph) - op = gb.op - - # --- Reshape 3D → 4D using head counts --- - b_dim = op.Shape(query, start=0, end=1) - t_dim = op.Shape(query, start=1, end=2) - - # Q/K: [B, T, q_num_heads*d_k] → [B, T, q_num_heads, d_k] - # → transpose to [B, q_num_heads, T, d_k] - qk_4d_shape = op.Concat( - b_dim, - t_dim, - op.Constant(value_ints=[q_num_heads, -1]), - axis=0, - ) - query_4d = op.Transpose( - op.Reshape(query, qk_4d_shape), perm=[0, 2, 1, 3] - ) # [B, q_num_heads, T, d_k] - key_4d = op.Transpose( - op.Reshape(key, qk_4d_shape), perm=[0, 2, 1, 3] - ) # [B, q_num_heads, T, d_k] - - # V: [B, T, kv_num_heads*d_v] → [B, kv_num_heads, T, d_v] - # Reuse kv_4d_shape for both V and decay (same [B, T, kv_num_heads, -1]). - kv_4d_shape = op.Concat( - b_dim, - t_dim, - op.Constant(value_ints=[kv_num_heads, -1]), - axis=0, - ) - value_4d = op.Transpose( - op.Reshape(value, kv_4d_shape), perm=[0, 2, 1, 3] - ) # [B, kv_num_heads, T, d_v] - - # --- GQA: expand Q/K heads to match V head count --- - if kv_num_heads % q_num_heads != 0: - raise ValueError( - f"kv_num_heads ({kv_num_heads}) must be divisible by q_num_heads ({q_num_heads})" + def body(op, *args): + # Unpack positional args matching the conditional input list. + query_v, key_v, value_v, past_state_v = args[:4] + idx = 4 + decay_v: ir.Value | None = None + beta_v: ir.Value | None = None + if uses_decay: + decay_v = args[idx] + idx += 1 + if uses_beta: + beta_v = args[idx] + + # --- Reshape 3D → 4D using head counts --- + b_dim = op.Shape(query_v, start=0, end=1) + t_dim = op.Shape(query_v, start=1, end=2) + + # Q/K: [B, T, q_num_heads*d_k] → [B, T, q_num_heads, d_k] + # → transpose to [B, q_num_heads, T, d_k] + qk_4d_shape = op.Concat( + b_dim, + t_dim, + op.Constant(value_ints=[q_num_heads, -1]), + axis=0, + ) + query_4d = op.Transpose( + op.Reshape(query_v, qk_4d_shape), perm=[0, 2, 1, 3] + ) # [B, q_num_heads, T, d_k] + key_4d = op.Transpose( + op.Reshape(key_v, qk_4d_shape), perm=[0, 2, 1, 3] + ) # [B, q_num_heads, T, d_k] + + # V: [B, T, kv_num_heads*d_v] → [B, kv_num_heads, T, d_v] + # Reuse kv_4d_shape for both V and decay (same [B, T, kv_num_heads, -1]). + kv_4d_shape = op.Concat( + b_dim, + t_dim, + op.Constant(value_ints=[kv_num_heads, -1]), + axis=0, + ) + value_4d = op.Transpose( + op.Reshape(value_v, kv_4d_shape), perm=[0, 2, 1, 3] + ) # [B, kv_num_heads, T, d_v] + + # --- GQA: expand Q/K heads to match V head count --- + if kv_num_heads % q_num_heads != 0: + raise ValueError( + f"kv_num_heads ({kv_num_heads}) must be divisible by q_num_heads ({q_num_heads})" + ) + gqa_ratio = kv_num_heads // q_num_heads + query_expanded, key_expanded = _expand_kv_heads( + op, query_4d, key_4d, gqa_ratio=gqa_ratio ) - gqa_ratio = kv_num_heads // q_num_heads - query_expanded, key_expanded = _expand_kv_heads(op, query_4d, key_4d, gqa_ratio=gqa_ratio) - - # --- Reshape decay/beta 3D → 4D (only when used) --- - # decay: (B, T, kv_num_heads * d_k) → (B, T, kv_num_heads, d_k) - # → transpose to (B, kv_num_heads, T, d_k) - if uses_decay: - decay_4d = op.Transpose( - op.Reshape(decay, kv_4d_shape), perm=[0, 2, 1, 3] - ) # [B, kv_num_heads, T, d_k] - if uses_beta: - # beta: (B, T, kv_num_heads) → transpose to (B, kv_num_heads, T) - beta_3d = op.Transpose(beta, perm=[0, 2, 1]) # [B, kv_num_heads, T] - - # --- Apply query scale (matches op spec default of 1/sqrt(d_k)) --- - # CastLike ensures scale constant matches the input dtype. - scaled_query = op.Mul( - query_expanded, - op.CastLike(op.Constant(value_float=scale), query_expanded), - ) - - # --- Build Scan for sequential recurrence --- - scan_body = _build_recurrence_body(uses_decay, uses_beta, stash_type=stash_type) - # Transpose to T-first for Scan: (B, H, T, D) -> (T, B, H, D) - q_t = op.Transpose(scaled_query, perm=[2, 0, 1, 3]) - k_t = op.Transpose(key_expanded, perm=[2, 0, 1, 3]) - v_t = op.Transpose(value_4d, perm=[2, 0, 1, 3]) + # --- Reshape decay/beta 3D → 4D (only when used) --- + # decay: (B, T, kv_num_heads * d_k) → (B, T, kv_num_heads, d_k) + # → transpose to (B, kv_num_heads, T, d_k) + if uses_decay: + decay_4d = op.Transpose( + op.Reshape(decay_v, kv_4d_shape), perm=[0, 2, 1, 3] + ) # [B, kv_num_heads, T, d_k] + if uses_beta: + # beta: (B, T, kv_num_heads) → transpose to (B, kv_num_heads, T) + beta_3d = op.Transpose(beta_v, perm=[0, 2, 1]) # [B, kv_num_heads, T] + + # --- Apply query scale (matches op spec default of 1/sqrt(d_k)) --- + # CastLike ensures scale constant matches the input dtype. + scaled_query = op.Mul( + query_expanded, + op.CastLike(op.Constant(value_float=scale), query_expanded), + ) - scan_inputs = [q_t, k_t, v_t] - if uses_decay: - decay_t = op.Transpose(decay_4d, perm=[2, 0, 1, 3]) # (T, B, H, d_k) - scan_inputs.append(decay_t) - if uses_beta: - beta_t = op.Transpose(beta_3d, perm=[2, 0, 1]) # (T, B, H) - scan_inputs.append(beta_t) - - present_state, output_t = op.Scan( - past_state, # carry: (B, H, d_k, d_v) - *scan_inputs, - body=scan_body, - num_scan_inputs=len(scan_inputs), - _outputs=2, - ) - # present_state: (B, H, d_k, d_v) - # output_t: (T, B, H, d_v) + # --- Build Scan for sequential recurrence --- + scan_body = _build_recurrence_body(uses_decay, uses_beta, stash_type=stash_type) + + # Transpose to T-first for Scan: (B, H, T, D) -> (T, B, H, D) + q_t = op.Transpose(scaled_query, perm=[2, 0, 1, 3]) + k_t = op.Transpose(key_expanded, perm=[2, 0, 1, 3]) + v_t = op.Transpose(value_4d, perm=[2, 0, 1, 3]) + + scan_inputs = [q_t, k_t, v_t] + if uses_decay: + decay_t = op.Transpose(decay_4d, perm=[2, 0, 1, 3]) # (T, B, H, d_k) + scan_inputs.append(decay_t) + if uses_beta: + beta_t = op.Transpose(beta_3d, perm=[2, 0, 1]) # (T, B, H) + scan_inputs.append(beta_t) + + present_state_v, output_t = op.Scan( + past_state_v, # carry: (B, H, d_k, d_v) + *scan_inputs, + body=scan_body, + num_scan_inputs=len(scan_inputs), + _outputs=2, + ) + # present_state_v: (B, H, d_k, d_v) + # output_t: (T, B, H, d_v) - # --- Reshape output 4D → 3D --- - # (T, B, H, d_v) → (B, T, H, d_v) → (B, T, H*d_v) - output_bthd = op.Transpose(output_t, perm=[1, 0, 2, 3]) - out_3d_shape = op.Concat(b_dim, t_dim, op.Constant(value_ints=[-1]), axis=0) - output = op.Reshape(output_bthd, out_3d_shape) # [B, T, H*d_v] + # --- Reshape output 4D → 3D --- + # (T, B, H, d_v) → (B, T, H, d_v) → (B, T, H*d_v) + output_bthd = op.Transpose(output_t, perm=[1, 0, 2, 3]) + out_3d_shape = op.Concat(b_dim, t_dim, op.Constant(value_ints=[-1]), axis=0) + output = op.Reshape(output_bthd, out_3d_shape) # [B, T, H*d_v] - output.name = "output" - present_state.name = "present_state" - graph.outputs.extend([output, present_state]) + output.name = "output" + present_state_v.name = "present_state" + return output, present_state_v # --- Build the ir.Function --- - update_rule_attr = ir.Attr( - "update_rule", - ir.AttributeType.STRING, - update_rule, - ref_attr_name="update_rule", - ) - scale_attr = ir.Attr( - "scale", - ir.AttributeType.FLOAT, - scale, - ref_attr_name="scale", - ) - q_heads_attr = ir.Attr( - "q_num_heads", - ir.AttributeType.INT, - q_num_heads, - ref_attr_name="q_num_heads", - ) - kv_heads_attr = ir.Attr( - "kv_num_heads", - ir.AttributeType.INT, - kv_num_heads, - ref_attr_name="kv_num_heads", - ) - return ir.Function( + return builder.build_function( + body, + inputs, domain=DOMAIN, name="LinearAttention", - graph=graph, - attributes={ - "update_rule": update_rule_attr, - "scale": scale_attr, - "q_num_heads": q_heads_attr, - "kv_num_heads": kv_heads_attr, - }, + attributes=[ + ir.Attr( + "update_rule", + ir.AttributeType.STRING, + update_rule, + ref_attr_name="update_rule", + ), + ir.Attr( + "scale", + ir.AttributeType.FLOAT, + scale, + ref_attr_name="scale", + ), + ir.Attr( + "q_num_heads", + ir.AttributeType.INT, + q_num_heads, + ref_attr_name="q_num_heads", + ), + ir.Attr( + "kv_num_heads", + ir.AttributeType.INT, + kv_num_heads, + ref_attr_name="kv_num_heads", + ), + ], + opset_imports={"": OPSET_VERSION}, ) diff --git a/src/mobius/functions/packed_multi_head_attention.py b/src/mobius/functions/packed_multi_head_attention.py index 7c106380..12eef05b 100644 --- a/src/mobius/functions/packed_multi_head_attention.py +++ b/src/mobius/functions/packed_multi_head_attention.py @@ -63,126 +63,118 @@ def packed_multi_head_attention() -> ir.Function: num_heads (int): Number of attention heads. scale (float): Attention scale (default 1.0). """ - # --- Graph inputs --- - query_input = ir.Value(name="query") - key_input = ir.Value(name="key") - value_input = ir.Value(name="value") - token_offset_input = ir.Value(name="token_offset") - cumulative_sequence_length_input = ir.Value(name="cumulative_sequence_length") - - graph = ir.Graph( - inputs=[ - query_input, - key_input, - value_input, - token_offset_input, - cumulative_sequence_length_input, - ], - outputs=[], - nodes=[], - name="PackedMultiHeadAttention_body", - opset_imports={"": OPSET_VERSION}, - ) - graph_builder = builder.GraphBuilder(graph) - op = graph_builder.op - - # --- Compute sequence length from query shape --- - # query: (token_count, hidden_size) - token_count = op.Shape(query_input, start=0, end=1) - token_count_scalar = op.Squeeze(token_count) - - # --- Build block-diagonal attention bias from cu_seqlens --- - # Create range [0, 1, ..., token_count - 1] - positions = op.Range( - op.Constant(value_int=0), - token_count_scalar, - op.Constant(value_int=1), - ) - # Compute segment IDs: for each position i, count how many - # cu_seqlens boundaries it has passed. - # segment_ids[i] = sum(i >= cu_seqlens[j] for all j) - 1 - positions_column = op.Unsqueeze(positions, [1]) # (N, 1) - cu_seqlens_int64 = op.Cast(cumulative_sequence_length_input, to=7) # INT64 - cu_seqlens_row = op.Unsqueeze(cu_seqlens_int64, [0]) # (1, S+1) - - # ge_mask[i, j] = (position_i >= cu_seqlens_j) - greater_or_equal_mask = op.GreaterOrEqual(positions_column, cu_seqlens_row) - greater_or_equal_int = op.Cast(greater_or_equal_mask, to=7) - - segment_ids = op.Sub( - op.ReduceSum(greater_or_equal_int, [1], keepdims=False), - op.Constant(value_int=1), - ) # (token_count,) - - # Build same-segment mask: same_segment[i, j] = (seg[i] == seg[j]) - segment_ids_row = op.Unsqueeze(segment_ids, [1]) # (N, 1) - segment_ids_column = op.Unsqueeze(segment_ids, [0]) # (1, N) - same_segment = op.Equal(segment_ids_row, segment_ids_column) # (N, N) - - # Convert to attention bias: 0 for same segment, -10000 for different - attention_bias = op.Where( - same_segment, - op.Constant(value_float=0.0), - op.Constant(value_float=-10000.0), - ) - # Reshape for Attention: (1, 1, N, N) - attention_bias = op.Unsqueeze(attention_bias, [0, 1]) - - # --- Add batch dimension for Attention op --- - # (token_count, hidden) → (1, token_count, hidden) - query_batched = op.Unsqueeze(query_input, [0]) - key_batched = op.Unsqueeze(key_input, [0]) - value_batched = op.Unsqueeze(value_input, [0]) - - # --- Call standard Attention op --- - # Use ir.Node for ref_attr_name forwarding of num_heads and scale. - attention_node = ir.Node( - "", - "Attention", - inputs=[ - query_batched, - key_batched, - value_batched, - attention_bias, - ], - attributes=[ - ir.Attr( - "q_num_heads", - ir.AttributeType.INT, - 1, - ref_attr_name="num_heads", - ), - ir.Attr( - "kv_num_heads", - ir.AttributeType.INT, - 1, - ref_attr_name="num_heads", - ), - ir.Attr( - "scale", - ir.AttributeType.FLOAT, - 1.0, - ref_attr_name="scale", - ), + def body( + op, + query_input, + key_input, + value_input, + token_offset_input, + cumulative_sequence_length_input, + ): + # --- Compute sequence length from query shape --- + # query: (token_count, hidden_size) + token_count = op.Shape(query_input, start=0, end=1) + token_count_scalar = op.Squeeze(token_count) + + # --- Build block-diagonal attention bias from cu_seqlens --- + # Create range [0, 1, ..., token_count - 1] + positions = op.Range( + op.Constant(value_int=0), + token_count_scalar, + op.Constant(value_int=1), + ) + + # Compute segment IDs: for each position i, count how many + # cu_seqlens boundaries it has passed. + # segment_ids[i] = sum(i >= cu_seqlens[j] for all j) - 1 + positions_column = op.Unsqueeze(positions, [1]) # (N, 1) + cu_seqlens_int64 = op.Cast(cumulative_sequence_length_input, to=7) # INT64 + cu_seqlens_row = op.Unsqueeze(cu_seqlens_int64, [0]) # (1, S+1) + + # ge_mask[i, j] = (position_i >= cu_seqlens_j) + greater_or_equal_mask = op.GreaterOrEqual(positions_column, cu_seqlens_row) + greater_or_equal_int = op.Cast(greater_or_equal_mask, to=7) + + segment_ids = op.Sub( + op.ReduceSum(greater_or_equal_int, [1], keepdims=False), + op.Constant(value_int=1), + ) # (token_count,) + + # Build same-segment mask: same_segment[i, j] = (seg[i] == seg[j]) + segment_ids_row = op.Unsqueeze(segment_ids, [1]) # (N, 1) + segment_ids_column = op.Unsqueeze(segment_ids, [0]) # (1, N) + same_segment = op.Equal(segment_ids_row, segment_ids_column) # (N, N) + + # Convert to attention bias: 0 for same segment, -10000 for different + attention_bias = op.Where( + same_segment, + op.Constant(value_float=0.0), + op.Constant(value_float=-10000.0), + ) + # Reshape for Attention: (1, 1, N, N) + attention_bias = op.Unsqueeze(attention_bias, [0, 1]) + + # --- Add batch dimension for Attention op --- + # (token_count, hidden) → (1, token_count, hidden) + query_batched = op.Unsqueeze(query_input, [0]) + key_batched = op.Unsqueeze(key_input, [0]) + value_batched = op.Unsqueeze(value_input, [0]) + + # --- Call standard Attention op --- + # Use ir.Node for ref_attr_name forwarding of num_heads and scale. + attention_node = ir.Node( + "", + "Attention", + inputs=[ + query_batched, + key_batched, + value_batched, + attention_bias, + ], + attributes=[ + ir.Attr( + "q_num_heads", + ir.AttributeType.INT, + 1, + ref_attr_name="num_heads", + ), + ir.Attr( + "kv_num_heads", + ir.AttributeType.INT, + 1, + ref_attr_name="num_heads", + ), + ir.Attr( + "scale", + ir.AttributeType.FLOAT, + 1.0, + ref_attr_name="scale", + ), + ], + num_outputs=1, + ) + op._builder._graph.append(attention_node) + attention_output = attention_node.outputs[0] + + # --- Remove batch dimension --- + # (1, token_count, v_hidden) → (token_count, v_hidden) + return op.Squeeze(attention_output, [0]) + + return builder.build_function( + body, + [ + ir.Value(name="query"), + ir.Value(name="key"), + ir.Value(name="value"), + ir.Value(name="token_offset"), + ir.Value(name="cumulative_sequence_length"), ], - num_outputs=1, - ) - graph.append(attention_node) - attention_output = attention_node.outputs[0] - - # --- Remove batch dimension --- - # (1, token_count, v_hidden) → (token_count, v_hidden) - output = op.Squeeze(attention_output, [0]) - - graph.outputs.append(output) - - return ir.Function( domain=DOMAIN, name="PackedMultiHeadAttention", - graph=graph, - attributes={ - "num_heads": ir.Attr("num_heads", ir.AttributeType.INT, 1), - "scale": ir.Attr("scale", ir.AttributeType.FLOAT, 1.0), - }, + attributes=[ + ir.Attr("num_heads", ir.AttributeType.INT, 1), + ir.Attr("scale", ir.AttributeType.FLOAT, 1.0), + ], + opset_imports={"": OPSET_VERSION}, ) diff --git a/src/mobius/functions/skip_layer_normalization.py b/src/mobius/functions/skip_layer_normalization.py index c5e2d8c4..0271fd8c 100644 --- a/src/mobius/functions/skip_layer_normalization.py +++ b/src/mobius/functions/skip_layer_normalization.py @@ -50,53 +50,51 @@ def skip_layer_normalization() -> ir.Function: Outputs: ``[norm_out, mean_out, inv_std_out, add_out]`` Attr: ``epsilon`` (float) """ - v_input = ir.Value(name="input") - v_skip = ir.Value(name="skip") - v_weight = ir.Value(name="weight") - v_bias = ir.Value(name="bias") - - graph = ir.Graph( - inputs=[v_input, v_skip, v_weight, v_bias], - outputs=[], - nodes=[], - name="SkipLayerNormalization_body", - opset_imports={"": OPSET_VERSION}, - ) - gb = builder.GraphBuilder(graph) - op = gb.op - - add_out = op.Add(v_input, v_skip) - - # ir.Node is required here: ref_attr_name forwards the caller's - # epsilon value at InlinePass expand time (OpBuilder doesn't - # support ref_attr_name). - ln_node = ir.Node( - "", - "LayerNormalization", - inputs=[add_out, v_weight, v_bias], - attributes=[ - ir.Attr("axis", ir.AttributeType.INT, -1), - ir.Attr( - "epsilon", - ir.AttributeType.FLOAT, - 1e-5, - ref_attr_name="epsilon", - ), - ], - num_outputs=3, - ) - graph.append(ln_node) - norm_out, mean_out, inv_std_out = ln_node.outputs - - graph.outputs.extend([norm_out, mean_out, inv_std_out, add_out]) - return ir.Function( + def body(op, v_input, v_skip, v_weight, v_bias): + add_out = op.Add(v_input, v_skip) + + # ir.Node is required here: ref_attr_name forwards the caller's + # epsilon value at InlinePass expand time (OpBuilder doesn't + # support ref_attr_name). + ln_node = ir.Node( + "", + "LayerNormalization", + inputs=[add_out, v_weight, v_bias], + attributes=[ + ir.Attr("axis", ir.AttributeType.INT, -1), + ir.Attr( + "epsilon", + ir.AttributeType.FLOAT, + 1e-5, + ref_attr_name="epsilon", + ), + ], + num_outputs=3, + ) + op._builder._graph.append(ln_node) + norm_out, mean_out, inv_std_out = ln_node.outputs + + norm_out.name = "norm_out" + mean_out.name = "mean_out" + inv_std_out.name = "inv_std_out" + add_out.name = "add_out" + return norm_out, mean_out, inv_std_out, add_out + + return builder.build_function( + body, + [ + ir.Value(name="input"), + ir.Value(name="skip"), + ir.Value(name="weight"), + ir.Value(name="bias"), + ], domain=DOMAIN, name="SkipLayerNormalization", - graph=graph, - attributes={ - "epsilon": ir.Attr("epsilon", ir.AttributeType.FLOAT, 1e-5), - }, + attributes=[ + ir.Attr("epsilon", ir.AttributeType.FLOAT, 1e-5), + ], + opset_imports={"": OPSET_VERSION}, ) @@ -112,47 +110,43 @@ def skip_simplified_layer_normalization() -> ir.Function: Outputs: ``[norm_out, add_out]`` Attr: ``epsilon`` (float) """ - v_input = ir.Value(name="input") - v_skip = ir.Value(name="skip") - v_weight = ir.Value(name="weight") - - graph = ir.Graph( - inputs=[v_input, v_skip, v_weight], - outputs=[], - nodes=[], - name="SkipSimplifiedLayerNormalization_body", - opset_imports={"": OPSET_VERSION}, - ) - gb = builder.GraphBuilder(graph) - op = gb.op - - add_out = op.Add(v_input, v_skip) - # ir.Node required: ref_attr_name forwards caller's epsilon. - rms_node = ir.Node( - "", - "RMSNormalization", - inputs=[add_out, v_weight], - attributes=[ - ir.Attr( - "epsilon", - ir.AttributeType.FLOAT, - 1e-5, - ref_attr_name="epsilon", - ), + def body(op, v_input, v_skip, v_weight): + add_out = op.Add(v_input, v_skip) + + # ir.Node required: ref_attr_name forwards caller's epsilon. + rms_node = ir.Node( + "", + "RMSNormalization", + inputs=[add_out, v_weight], + attributes=[ + ir.Attr( + "epsilon", + ir.AttributeType.FLOAT, + 1e-5, + ref_attr_name="epsilon", + ), + ], + num_outputs=1, + ) + op._builder._graph.append(rms_node) + norm_out = rms_node.outputs[0] + + norm_out.name = "norm_out" + add_out.name = "add_out" + return norm_out, add_out + + return builder.build_function( + body, + [ + ir.Value(name="input"), + ir.Value(name="skip"), + ir.Value(name="weight"), ], - num_outputs=1, - ) - graph.append(rms_node) - norm_out = rms_node.outputs[0] - - graph.outputs.extend([norm_out, add_out]) - - return ir.Function( domain=DOMAIN, name="SkipSimplifiedLayerNormalization", - graph=graph, - attributes={ - "epsilon": ir.Attr("epsilon", ir.AttributeType.FLOAT, 1e-5), - }, + attributes=[ + ir.Attr("epsilon", ir.AttributeType.FLOAT, 1e-5), + ], + opset_imports={"": OPSET_VERSION}, ) From deb829d48ce2cea750d9d4432ff872685542d07d Mon Sep 17 00:00:00 2001 From: G Ramalingam Date: Wed, 22 Apr 2026 17:22:00 +0000 Subject: [PATCH 2/4] Use None pattern for optional LinearAttention inputs Replace *args unpacking with named parameters and None entries. The function signature is now stable (6 formals) across all update_rule variants. Absent optional inputs (decay, beta) get placeholder formals via build_function; the trace function receives None and branches naturally with 'if decay_v is not None'. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam --- src/mobius/functions/linear_attention.py | 42 +++++++++--------------- 1 file changed, 15 insertions(+), 27 deletions(-) diff --git a/src/mobius/functions/linear_attention.py b/src/mobius/functions/linear_attention.py index a5007fbc..a89b6681 100644 --- a/src/mobius/functions/linear_attention.py +++ b/src/mobius/functions/linear_attention.py @@ -109,33 +109,21 @@ def linear_attention( uses_decay = update_rule in ("gated", "gated_delta") uses_beta = update_rule in ("delta", "gated_delta") - # --- Define function inputs (conditional on update_rule) --- - query = ir.Value(name="query") # (B, T, q_num_heads * d_k) - key = ir.Value(name="key") # (B, T, q_num_heads * d_k) - value = ir.Value(name="value") # (B, T, kv_num_heads * d_v) - past_state = ir.Value(name="past_state") - inputs: list[ir.Value] = [query, key, value, past_state] - decay: ir.Value | None = None - beta: ir.Value | None = None - if uses_decay: - decay = ir.Value(name="decay") - inputs.append(decay) - if uses_beta: - beta = ir.Value(name="beta") - inputs.append(beta) - - def body(op, *args): - # Unpack positional args matching the conditional input list. - query_v, key_v, value_v, past_state_v = args[:4] - idx = 4 - decay_v: ir.Value | None = None - beta_v: ir.Value | None = None - if uses_decay: - decay_v = args[idx] - idx += 1 - if uses_beta: - beta_v = args[idx] - + # --- Define function inputs --- + # All 6 formal parameters are always declared so the function + # signature is stable across update_rule variants. Absent optional + # inputs are ``None``; build_function creates placeholder formals + # for them and passes ``None`` to the trace function. + inputs: list[ir.Value | None] = [ + ir.Value(name="query"), # (B, T, q_num_heads * d_k) + ir.Value(name="key"), # (B, T, q_num_heads * d_k) + ir.Value(name="value"), # (B, T, kv_num_heads * d_v) + ir.Value(name="past_state"), + ir.Value(name="decay") if uses_decay else None, + ir.Value(name="beta") if uses_beta else None, + ] + + def body(op, query_v, key_v, value_v, past_state_v, decay_v, beta_v): # --- Reshape 3D → 4D using head counts --- b_dim = op.Shape(query_v, start=0, end=1) t_dim = op.Shape(query_v, start=1, end=2) From 15b0d866b409e4d01e5113bde667a1eebc16962e Mon Sep 17 00:00:00 2001 From: G Ramalingam Date: Wed, 22 Apr 2026 17:41:01 +0000 Subject: [PATCH 3/4] Use OpBuilder for ref_attr_name; fix LinearAttention arity - skip_layer_normalization: Replace raw ir.Node with op.LayerNormalization and op.RMSNormalization, passing ir.Attr objects with ref_attr_name directly as kwargs. OpBuilder passes them through to ir.node() correctly. Remove outdated docstring note about OpBuilder limitation. - packed_multi_head_attention: Replace raw ir.Node with op.Attention, passing ir.Attr objects with ref_attr_name for num_heads and scale. - linear_attention: Revert to per-variant input arity. The function is specialized per-model, so only inputs used by the update_rule are declared. This matches call sites (Mamba2/Bamba pass 5 args for 'gated', GatedDeltaNet passes 6 for 'gated_delta'). Shape inference does strict arity checking. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam --- src/mobius/functions/linear_attention.py | 26 +++++--- .../functions/packed_multi_head_attention.py | 58 ++++++++--------- .../functions/skip_layer_normalization.py | 64 ++++++------------- 3 files changed, 63 insertions(+), 85 deletions(-) diff --git a/src/mobius/functions/linear_attention.py b/src/mobius/functions/linear_attention.py index a89b6681..8129c744 100644 --- a/src/mobius/functions/linear_attention.py +++ b/src/mobius/functions/linear_attention.py @@ -109,21 +109,29 @@ def linear_attention( uses_decay = update_rule in ("gated", "gated_delta") uses_beta = update_rule in ("delta", "gated_delta") - # --- Define function inputs --- - # All 6 formal parameters are always declared so the function - # signature is stable across update_rule variants. Absent optional - # inputs are ``None``; build_function creates placeholder formals - # for them and passes ``None`` to the trace function. - inputs: list[ir.Value | None] = [ + # --- Define function inputs (conditional on update_rule) --- + # The function is specialized per-model: only inputs actually used + # by this update_rule variant are declared. Call sites (e.g. + # Mamba2Block with "gated") pass exactly the declared number of args. + inputs: list[ir.Value] = [ ir.Value(name="query"), # (B, T, q_num_heads * d_k) ir.Value(name="key"), # (B, T, q_num_heads * d_k) ir.Value(name="value"), # (B, T, kv_num_heads * d_v) ir.Value(name="past_state"), - ir.Value(name="decay") if uses_decay else None, - ir.Value(name="beta") if uses_beta else None, ] + if uses_decay: + inputs.append(ir.Value(name="decay")) + if uses_beta: + inputs.append(ir.Value(name="beta")) + + def body(op, *args): + query_v, key_v, value_v, past_state_v = args[:4] + idx = 4 + decay_v = args[idx] if uses_decay else None + if uses_decay: + idx += 1 + beta_v = args[idx] if uses_beta else None - def body(op, query_v, key_v, value_v, past_state_v, decay_v, beta_v): # --- Reshape 3D → 4D using head counts --- b_dim = op.Shape(query_v, start=0, end=1) t_dim = op.Shape(query_v, start=1, end=2) diff --git a/src/mobius/functions/packed_multi_head_attention.py b/src/mobius/functions/packed_multi_head_attention.py index 12eef05b..7ac88096 100644 --- a/src/mobius/functions/packed_multi_head_attention.py +++ b/src/mobius/functions/packed_multi_head_attention.py @@ -122,40 +122,32 @@ def body( value_batched = op.Unsqueeze(value_input, [0]) # --- Call standard Attention op --- - # Use ir.Node for ref_attr_name forwarding of num_heads and scale. - attention_node = ir.Node( - "", - "Attention", - inputs=[ - query_batched, - key_batched, - value_batched, - attention_bias, - ], - attributes=[ - ir.Attr( - "q_num_heads", - ir.AttributeType.INT, - 1, - ref_attr_name="num_heads", - ), - ir.Attr( - "kv_num_heads", - ir.AttributeType.INT, - 1, - ref_attr_name="num_heads", - ), - ir.Attr( - "scale", - ir.AttributeType.FLOAT, - 1.0, - ref_attr_name="scale", - ), - ], - num_outputs=1, + # ir.Attr with ref_attr_name forwards num_heads and scale from + # the function's formal attributes to the inner Attention node. + attention_output = op.Attention( + query_batched, + key_batched, + value_batched, + attention_bias, + q_num_heads=ir.Attr( + "q_num_heads", + ir.AttributeType.INT, + 1, + ref_attr_name="num_heads", + ), + kv_num_heads=ir.Attr( + "kv_num_heads", + ir.AttributeType.INT, + 1, + ref_attr_name="num_heads", + ), + scale=ir.Attr( + "scale", + ir.AttributeType.FLOAT, + 1.0, + ref_attr_name="scale", + ), ) - op._builder._graph.append(attention_node) - attention_output = attention_node.outputs[0] # --- Remove batch dimension --- # (1, token_count, v_hidden) → (token_count, v_hidden) diff --git a/src/mobius/functions/skip_layer_normalization.py b/src/mobius/functions/skip_layer_normalization.py index 0271fd8c..e1ac8096 100644 --- a/src/mobius/functions/skip_layer_normalization.py +++ b/src/mobius/functions/skip_layer_normalization.py @@ -15,14 +15,6 @@ discoverability, matching the convention in :mod:`~mobius.functions.causal_conv` and :mod:`~mobius.functions.linear_attention`. - -.. note:: - - These function bodies use raw ``ir.Node`` construction for - ``LayerNormalization`` and ``RMSNormalization`` nodes because - ``ref_attr_name`` (which tells InlinePass to forward the caller's - ``epsilon`` value) is not supported by the ``OpBuilder`` API. - All other ops use the standard ``OpBuilder`` (``op.Add``, etc.). """ from __future__ import annotations @@ -54,26 +46,21 @@ def skip_layer_normalization() -> ir.Function: def body(op, v_input, v_skip, v_weight, v_bias): add_out = op.Add(v_input, v_skip) - # ir.Node is required here: ref_attr_name forwards the caller's - # epsilon value at InlinePass expand time (OpBuilder doesn't - # support ref_attr_name). - ln_node = ir.Node( - "", - "LayerNormalization", - inputs=[add_out, v_weight, v_bias], - attributes=[ - ir.Attr("axis", ir.AttributeType.INT, -1), - ir.Attr( - "epsilon", - ir.AttributeType.FLOAT, - 1e-5, - ref_attr_name="epsilon", - ), - ], - num_outputs=3, + # ref_attr_name forwards the caller's epsilon at InlinePass expand time. + epsilon_attr = ir.Attr( + "epsilon", + ir.AttributeType.FLOAT, + 1e-5, + ref_attr_name="epsilon", + ) + norm_out, mean_out, inv_std_out = op.LayerNormalization( + add_out, + v_weight, + v_bias, + axis=-1, + epsilon=epsilon_attr, + _outputs=3, ) - op._builder._graph.append(ln_node) - norm_out, mean_out, inv_std_out = ln_node.outputs norm_out.name = "norm_out" mean_out.name = "mean_out" @@ -114,23 +101,14 @@ def skip_simplified_layer_normalization() -> ir.Function: def body(op, v_input, v_skip, v_weight): add_out = op.Add(v_input, v_skip) - # ir.Node required: ref_attr_name forwards caller's epsilon. - rms_node = ir.Node( - "", - "RMSNormalization", - inputs=[add_out, v_weight], - attributes=[ - ir.Attr( - "epsilon", - ir.AttributeType.FLOAT, - 1e-5, - ref_attr_name="epsilon", - ), - ], - num_outputs=1, + # ref_attr_name forwards caller's epsilon at InlinePass expand time. + epsilon_attr = ir.Attr( + "epsilon", + ir.AttributeType.FLOAT, + 1e-5, + ref_attr_name="epsilon", ) - op._builder._graph.append(rms_node) - norm_out = rms_node.outputs[0] + norm_out = op.RMSNormalization(add_out, v_weight, epsilon=epsilon_attr) norm_out.name = "norm_out" add_out.name = "add_out" From 40601c150509411e5b5a3dee3cc6c6b75fbc4a72 Mon Sep 17 00:00:00 2001 From: G Ramalingam Date: Wed, 22 Apr 2026 17:47:56 +0000 Subject: [PATCH 4/4] Add TODO for variadic function arity in LinearAttention The strict arity check is in the Python onnx-shape-inference package (onnx_shape_inference._functions.infer_function_call_output_shapes), not in ONNX's C++ shape inference. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: G Ramalingam --- src/mobius/functions/linear_attention.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/mobius/functions/linear_attention.py b/src/mobius/functions/linear_attention.py index 8129c744..73525796 100644 --- a/src/mobius/functions/linear_attention.py +++ b/src/mobius/functions/linear_attention.py @@ -110,6 +110,8 @@ def linear_attention( uses_beta = update_rule in ("delta", "gated_delta") # --- Define function inputs (conditional on update_rule) --- + # TODO: Investigate relaxing onnx-shape-inference's strict arity check + # to allow trailing optional inputs, enabling a single 6-input signature. # The function is specialized per-model: only inputs actually used # by this update_rule variant are declared. Call sites (e.g. # Mamba2Block with "gated") pass exactly the declared number of args.