linear attention signature

[guschmue]

Proposal for CausalConvWithState and LinearAttention onnxruntime custom operator.
This follows the proposal in onnx/onnx#7767.

[guschmue]

a working end to end implementation for webgpu with this signature can be found here:
https://github.com/microsoft/onnxruntime/tree/gs/wgpu-lattn

Possible changes at this point:

1. maybe inputs should be transposed, else the model will have transpose operators in front of LinearAttention
2. maybe CausalConv1DWithState should be CausalConvWithState

[github-actions]

You can commit the suggested changes from lintrunner.

[justinchuby]

Cross-referencing with the ONNX proposal (onnx/onnx#7767) — review by AI agent team

I compared the LinearAttention and CausalConvWithState schemas between this ORT contrib PR and the ONNX proposal in onnx/onnx#7767. Summary below.

---

LinearAttention — comparison

✅ Matches

Item	Both PRs
Input names (0–5)	query, key, value, past_state, decay, beta — identical order and names
Output names (0–1)	output, present_state — identical
Input shapes	All match: 3D packed [B,T,H*D] for Q/K/V/decay/beta; 4D [B,H_kv,d_k,d_v] for state
Attribute names	q_num_heads, kv_num_heads, update_rule, scale, chunk_size — all five present
Attribute types & defaults	update_rule="gated_delta", scale=0.0, chunk_size=64, both heads required
TypeConstraint T	{tensor(float), tensor(float16), tensor(bfloat16)} — identical
Mathematical semantics	linear, gated, delta, gated_delta update rules described identically
Optional input semantics	past_state defaults to zeros; decay/beta required by respective modes
Namespace	ORT uses com.microsoft, ONNX proposal uses ai.onnx — expected and correct per the 3-phase adoption path

❌ Key mismatch — state type precision (S vs T)

The ONNX proposal introduces a second type parameter S for the recurrent state tensors, following the stash_type convention used by LayerNormalization and GroupNormalization:

• past_state → type S (not T)
• present_state → type S (not T)
• S is constrained to float32, or the same as T

This ORT PR uses a single type T for all inputs/outputs including state. This means:

• In this ORT schema, running with T=float16 forces the recurrent state to also be float16
• The ONNX proposal explicitly supports T=float16, S=float32 — float16 activations with float32 state accumulation — which is important for numerical stability during long-sequence decoding (the state matrix is accumulated over hundreds or thousands of tokens)
• The ONNX proposal notes: "Using S = float32 with T = float16/bfloat16 is the recommended configuration for long sequences; runtimes handle any necessary casting internally"

Recommendation: Consider adding a stash_type attribute (integer, default 1 = float32) analogous to LayerNormalization, so callers can opt into float32 state accumulation independently of the activation dtype. This would align with the ONNX proposal and avoid a breaking schema change later.

⚠️ Minor gap — input validation rules not in doc string

The ONNX proposal has an explicit table of required/forbidden optional inputs per update_rule:

update_rule	decay	beta
"linear"	must be omitted	must be omitted
"gated"	required	must be omitted
"delta"	must be omitted	required
"gated_delta"	required	required

The ORT doc string describes the semantics but doesn't explicitly state that providing a forbidden input is a model validation error. Worth adding to the doc for clarity — helps implementors know they should validate at model-load time, not silently ignore extra inputs.

---

CausalConvWithState — comparison

✅ Full match

All inputs (input, weight, bias, past_state), outputs (output, present_state), attributes (ndim, activation), type constraints, and shapes are identical between the two PRs. No differences found.

---

Summary

Op	Status
LinearAttention	Mostly aligned — 1 structural gap (state type S vs T), 1 minor doc gap
CausalConvWithState	Fully aligned ✅

The state precision gap is the only item that would cause a future breaking schema change if not addressed here. Everything else looks well-aligned with the ONNX proposal.

[justinchuby]

Cross-referencing with the ONNX proposal (onnx/onnx#7767) — review by AI agent team

I systematically compared the LinearAttention and CausalConvWithState op schemas in this PR against the formal schema defined in onnx/onnx#7767. The namespace difference (com.microsoft here vs ai.onnx in the proposal) is expected per the contrib-first adoption path and is not flagged below.

---

LinearAttention

✅ Matches

Element	ORT	ONNX proposal
Input names	query, key, value, past_state, decay, beta	same
Input order	0..5	same
Input optionality	past_state, decay, beta optional	same
Input shapes (3D packed)	(B, T, H*D) for Q/K/V; (B, H_kv, d_k, d_v) for state	same
Decay shape variants	(B, T, H_kv*d_k) or (B, T, H_kv)	same
Output names	output, present_state	same
Output shapes	(B, T, H_q*d_v) and (B, H_kv, d_k, d_v)	same
Attribute names	update_rule, scale, q_num_heads, kv_num_heads, chunk_size	same
Attribute types	all match	same
Attribute defaults	update_rule="gated_delta", scale=0.0, chunk_size=64	same
q_num_heads/kv_num_heads required	no default value	same
Activation dtype constraint	float, float16, bfloat16	same
Zero-initialized state semantics	implied (optional input)	explicitly stated

❌ Mismatch: State type parameter (S vs single T)

This is the most significant schema difference.

The ONNX proposal defines two separate type parameters:

• T — activation dtype for query, key, value, decay, beta, output
• S — state dtype for past_state and present_state; must be float32 or same as T

This allows the recommended configuration of fp32 state + fp16/bf16 activations — important for numerical accuracy in long sequences where state accumulation in fp16 diverges. The ONNX proposal explicitly calls this out: "Using S = float32 with T = float16/bfloat16 is the recommended configuration for long sequences."

This PR uses a single T type for all inputs including state, which prevents expressing this mixed-precision configuration in the op signature.

Suggestion: Consider adding a second type constraint S for past_state/present_state, or at minimum add an attribute state_dtype that allows the runtime to accumulate state in float32 even when activations are fp16.

⚠️ Minor: Input combination validation not documented in schema

The ONNX proposal formally specifies that the presence/absence of decay and beta depends on update_rule and must be validated at model-load time (e.g., "linear" requires neither, "gated_delta" requires both; providing a forbidden input is a schema error).

This PR marks both as Optional in the schema without documenting this constraint. The kernel may still validate at runtime, but adding a note to the docstring would make the contract explicit to model builders.

---

CausalConvWithState

✅ Matches (essentially identical)

Element	ORT	ONNX proposal
Input names	input, weight, bias, past_state	same
Input shapes	(B, C, ...) input, (C, 1, k, ...) weight, (C,) bias, (B, C, k-1) state	same
Input optionality	bias and past_state optional	same
Output names	output, present_state	same
Output shapes	same as input; same as past_state	same
Attributes	activation (default "none"), ndim (default 1)	same
Activation values	"silu", "swish", "none" ("silu" and "swish" are aliases)	same
Type constraint	single T: float, float16, bfloat16	same

No divergence found for CausalConvWithState. The single-type constraint is appropriate here since the conv state dtype naturally matches the input dtype.

---

Summary

The two schemas are structurally aligned. The one actionable difference is the missing S state-dtype type parameter in LinearAttention. Everything else matches: all 5 attributes, all 6 inputs, both outputs, all defaults, and the full CausalConvWithState schema.

The ONNX proposal is at onnx/onnx#7767 if you want to review the reference-level pseudocode and the formal input-combination validation table.

[justinchuby]

LGTM w/ agreement w/ the AI comments