File size: 12,268 Bytes
b8b1d7c 28b65f6 b8b1d7c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 | {
"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 }
}
}
]
}
|