| { |
| "cases": [ |
| { |
| "name": "linear_gqa_no_state_nondefault_chunk", |
| "provenance": { |
| "source": "onnx/reference/ops/op_linear_attention.py", |
| "notes": "Standard grouped-query attention with odd key width and a non-default correctness-neutral chunk_size checks the recurrence without past state." |
| }, |
| "attrs": { "q_num_heads": 4, "kv_num_heads": 2, "update_rule": "linear", "chunk_size": 17 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 3, 12], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 3, 6], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 3, 4], |
| "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 } |
| } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 3, 8], "tolerance": 0.0001 }, |
| "presentStateT": { "dtype": "float32", "shape": [1, 2, 3, 2], "tolerance": 0.0001 } |
| } |
| }, |
| { |
| "name": "linear_mqa_mixed_f16_state", |
| "attrs": { "q_num_heads": 4, "kv_num_heads": 1, "update_rule": "linear", "scale": 0.5 }, |
| "requires": { "features": ["shader-f16"] }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 2, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 2, 4], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.37 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 2, 3], |
| "data": { "kind": "fillFloat32", "sinStep": 0.09, "cosStep": 0.43 } |
| }, |
| "pastStateT": { |
| "dtype": "float16", |
| "shape": [1, 1, 4, 3], |
| "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.21, "scale": 0.2 } |
| } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 2, 12], "tolerance": 0.0002 }, |
| "presentStateT": { "dtype": "float16", "shape": [1, 1, 4, 3], "tolerance": 0.003 } |
| } |
| }, |
| { |
| "name": "gated_per_head_decay_no_state", |
| "attrs": { "q_num_heads": 2, "kv_num_heads": 2, "update_rule": "gated", "scale": 0.4 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.27 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.12, "cosStep": 0.35 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 2, 4], |
| "data": { "kind": "fillFloat32", "sinStep": 0.08, "cosStep": 0.39 } |
| }, |
| "decayT": { |
| "dtype": "float32", |
| "shape": [1, 2, 2], |
| "data": { "kind": "values", "values": [-0.1, -0.2, -0.05, -0.15] } |
| } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 2, 4], "tolerance": 0.0001 }, |
| "presentStateT": { "dtype": "float32", "shape": [1, 2, 4, 2], "tolerance": 0.0001 } |
| } |
| }, |
| { |
| "name": "gated_per_key_decay_f16_with_f32_state", |
| "attrs": { "q_num_heads": 4, "kv_num_heads": 2, "update_rule": "gated" }, |
| "requires": { "features": ["shader-f16"] }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [1, 2, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.15, "cosStep": 0.33, "scale": 0.3 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [1, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.1, "cosStep": 0.25, "scale": 0.3 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [1, 2, 4], |
| "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41, "scale": 0.3 } |
| }, |
| "pastStateT": { "dtype": "float32", "shape": [1, 2, 4, 2], "data": { "kind": "constant", "value": 0.1 } }, |
| "decayT": { |
| "dtype": "float16", |
| "shape": [1, 2, 8], |
| "data": { "kind": "linspace", "start": -0.25, "end": -0.02 } |
| } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float16", "shape": [1, 2, 8], "tolerance": 0.006 }, |
| "presentStateT": { "dtype": "float32", "shape": [1, 2, 4, 2], "tolerance": 0.001 } |
| } |
| }, |
| { |
| "name": "delta_broadcast_beta_no_state", |
| "attrs": { "q_num_heads": 2, "kv_num_heads": 1, "update_rule": "delta", "scale": 0.6 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [2, 2, 6], |
| "data": { "kind": "fillFloat32", "sinStep": 0.16, "cosStep": 0.22 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [2, 2, 3], |
| "data": { "kind": "fillFloat32", "sinStep": 0.09, "cosStep": 0.34 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [2, 2, 2], |
| "data": { "kind": "fillFloat32", "sinStep": 0.06, "cosStep": 0.38 } |
| }, |
| "betaT": { |
| "dtype": "float32", |
| "shape": [2, 2, 1], |
| "data": { "kind": "values", "values": [0.8, 0.6, 0.7, 0.9] } |
| } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [2, 2, 4], "tolerance": 0.0001 }, |
| "presentStateT": { "dtype": "float32", "shape": [2, 1, 3, 2], "tolerance": 0.0001 } |
| } |
| }, |
| { |
| "name": "delta_vec4_no_state", |
| "attrs": { "q_num_heads": 2, "kv_num_heads": 1, "update_rule": "delta", "scale": 0.6 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.16, "cosStep": 0.22 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 2, 4], |
| "data": { "kind": "fillFloat32", "sinStep": 0.09, "cosStep": 0.34 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 2, 4], |
| "data": { "kind": "fillFloat32", "sinStep": 0.06, "cosStep": 0.38 } |
| }, |
| "betaT": { "dtype": "float32", "shape": [1, 2, 1], "data": { "kind": "values", "values": [0.8, 0.6] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 2, 8], "tolerance": 0.0001 }, |
| "presentStateT": { "dtype": "float32", "shape": [1, 1, 4, 4], "tolerance": 0.0001 } |
| } |
| }, |
| { |
| "name": "delta_per_head_beta_f16_state", |
| "attrs": { "q_num_heads": 4, "kv_num_heads": 2, "update_rule": "delta", "scale": 0.5 }, |
| "requires": { "features": ["shader-f16"] }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [1, 2, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.18, "cosStep": 0.28, "scale": 0.2 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [1, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.14, "cosStep": 0.32, "scale": 0.2 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [1, 2, 4], |
| "data": { "kind": "fillFloat32", "sinStep": 0.1, "cosStep": 0.36, "scale": 0.2 } |
| }, |
| "pastStateT": { "dtype": "float16", "shape": [1, 2, 4, 2], "data": { "kind": "constant", "value": 0.125 } }, |
| "betaT": { |
| "dtype": "float16", |
| "shape": [1, 2, 2], |
| "data": { "kind": "values", "values": [0.75, 0.5, 0.625, 0.875] } |
| } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float16", "shape": [1, 2, 8], "tolerance": 0.003 }, |
| "presentStateT": { "dtype": "float16", "shape": [1, 2, 4, 2], "tolerance": 0.003 } |
| } |
| }, |
| { |
| "name": "default_gated_delta_no_state", |
| "attrs": { "q_num_heads": 2, "kv_num_heads": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 2, 4], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 2, 3], |
| "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 } |
| }, |
| "decayT": { |
| "dtype": "float32", |
| "shape": [1, 2, 4], |
| "data": { "kind": "linspace", "start": -0.2, "end": -0.01 } |
| }, |
| "betaT": { "dtype": "float32", "shape": [1, 2, 1], "data": { "kind": "values", "values": [0.8, 0.6] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 2, 6], "tolerance": 0.0001 }, |
| "presentStateT": { "dtype": "float32", "shape": [1, 1, 4, 3], "tolerance": 0.0001 } |
| } |
| }, |
| { |
| "name": "gated_delta_gqa_with_state", |
| "attrs": { "q_num_heads": 4, "kv_num_heads": 2, "update_rule": "gated_delta", "scale": 0.7, "chunk_size": -7 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 3, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 3, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.37 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 3, 6], |
| "data": { "kind": "fillFloat32", "sinStep": 0.09, "cosStep": 0.43 } |
| }, |
| "pastStateT": { |
| "dtype": "float32", |
| "shape": [1, 2, 4, 3], |
| "data": { "kind": "fillFloat32", "sinStep": 0.05, "cosStep": 0.21, "scale": 0.2 } |
| }, |
| "decayT": { |
| "dtype": "float32", |
| "shape": [1, 3, 2], |
| "data": { "kind": "linspace", "start": -0.2, "end": -0.03 } |
| }, |
| "betaT": { "dtype": "float32", "shape": [1, 3, 2], "data": { "kind": "linspace", "start": 0.45, "end": 0.9 } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 3, 12], "tolerance": 0.0001 }, |
| "presentStateT": { "dtype": "float32", "shape": [1, 2, 4, 3], "tolerance": 0.0001 } |
| } |
| }, |
| { |
| "name": "linear_zero_sequence_carries_state", |
| "attrs": { "q_num_heads": 2, "kv_num_heads": 1, "update_rule": "linear" }, |
| "inputs": { |
| "queryT": { "dtype": "float32", "shape": [1, 0, 6], "data": { "kind": "values", "values": [] } }, |
| "keyT": { "dtype": "float32", "shape": [1, 0, 3], "data": { "kind": "values", "values": [] } }, |
| "valueT": { "dtype": "float32", "shape": [1, 0, 2], "data": { "kind": "values", "values": [] } }, |
| "pastStateT": { |
| "dtype": "float32", |
| "shape": [1, 1, 3, 2], |
| "data": { "kind": "values", "values": [0.1, -0.2, 0.3, -0.4, 0.5, -0.6] } |
| } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 0, 4], "data": { "kind": "values", "values": [] } }, |
| "presentStateT": { |
| "dtype": "float32", |
| "shape": [1, 1, 3, 2], |
| "data": { "kind": "values", "values": [0.1, -0.2, 0.3, -0.4, 0.5, -0.6] } |
| } |
| } |
| }, |
| { |
| "name": "linear_wide_vec4_head", |
| "attrs": { "q_num_heads": 2, "kv_num_heads": 1, "update_rule": "linear", "scale": 0.04 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 1, 1024], |
| "data": { "kind": "fillFloat32", "sinStep": 0.013, "cosStep": 0.021, "scale": 0.05 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 1, 512], |
| "data": { "kind": "fillFloat32", "sinStep": 0.017, "cosStep": 0.027, "scale": 0.05 } |
| }, |
| "valueT": { "dtype": "float32", "shape": [1, 1, 1], "data": { "kind": "values", "values": [0.25] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 1, 2], "tolerance": 0.0002 }, |
| "presentStateT": { "dtype": "float32", "shape": [1, 1, 512, 1], "tolerance": 0.0001 } |
| } |
| } |
| ] |
| } |
|
|
