| { |
| "op": "com.microsoft.SparseAttention", |
| "fixtureArrays": { |
| "prompt_separate_two_layouts_input_blockColIndicesT": [0, 0, 1, 1, 2, 1, 2, 3, -1, 0, 0, 1, 0, 1, 2, 0, 2, 3] |
| }, |
| "cases": [ |
| { |
| "name": "prompt_separate_two_layouts", |
| "provenance": { |
| "notes": "The two-layout CSR example from the ONNX Runtime operator documentation, using separate Q/K/V tensors." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [2, 32, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [2, 32, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [2, 32, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [2, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [2, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 9], |
| "data": { |
| "kind": "values", |
| "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" } |
| } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [32, 32] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [2, 32, 32], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastKeyT": { "dtype": "float32", "shape": [2, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [2, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "prompt_packed_two_layouts", |
| "provenance": { "notes": "Packed [Q|K|V] form of the documented two-layout prompt case." }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 32, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.27, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 9], |
| "data": { |
| "kind": "values", |
| "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" } |
| } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 32, 32], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "prompt_separate_rotary_neox", |
| "provenance": { |
| "notes": "Full-width NeoX rotary with the minimum accepted head width; rotary_interleaved values other than 1 retain the NeoX layout." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16, "do_rotary": 1, "rotary_interleaved": 2 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 32, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 32, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 32, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 16], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 16], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 9], |
| "data": { |
| "kind": "values", |
| "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" } |
| } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }, |
| "cosCacheT": { "dtype": "float32", "shape": [64, 8], "data": { "kind": "constant", "value": 0.8 } }, |
| "sinCacheT": { "dtype": "float32", "shape": [64, 8], "data": { "kind": "constant", "value": 0.6 } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 32, 64], "tolerance": 0.00003, "relTolerance": 0.0003 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 16], "tolerance": 0.00003, "relTolerance": 0.0003 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 16], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "prompt_packed_partial_rotary_interleaved", |
| "provenance": { |
| "notes": "Packed QKV with interleaved rotary over the first 16 of 32 head channels; ONNX Runtime copies the remainder through." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16, "do_rotary": 1, "rotary_interleaved": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 32, 256], |
| "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 32], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 32], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 9], |
| "data": { |
| "kind": "values", |
| "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" } |
| } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }, |
| "cosCacheT": { "dtype": "float32", "shape": [64, 8], "data": { "kind": "constant", "value": 0.8 } }, |
| "sinCacheT": { "dtype": "float32", "shape": [64, 8], "data": { "kind": "constant", "value": 0.6 } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 32, 128], "tolerance": 0.00003, "relTolerance": 0.0003 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 32], "tolerance": 0.00003, "relTolerance": 0.0003 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 32], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "decode_separate_two_layouts", |
| "provenance": { |
| "notes": "Two decode rows with different history lengths, so the scalar maximum and per-row key lengths have distinct jobs." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [2, 1, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [2, 1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [2, 1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { |
| "dtype": "float32", |
| "shape": [2, 2, 64, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 } |
| }, |
| "pastValueT": { |
| "dtype": "float32", |
| "shape": [2, 2, 64, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 } |
| }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 9], |
| "data": { |
| "kind": "values", |
| "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" } |
| } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [41] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [41, 34] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [2, 1, 32], "tolerance": 0.00003, "relTolerance": 0.0003 }, |
| "pastKeyT": { "dtype": "float32", "shape": [2, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [2, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "decode_packed_rotary_neox", |
| "provenance": { "notes": "Packed decode with full-width NeoX rotary at each row's absolute position." }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16, "do_rotary": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [2, 1, 128], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "pastKeyT": { |
| "dtype": "float32", |
| "shape": [2, 2, 64, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 } |
| }, |
| "pastValueT": { |
| "dtype": "float32", |
| "shape": [2, 2, 64, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 } |
| }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 9], |
| "data": { |
| "kind": "values", |
| "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" } |
| } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [41] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [41, 34] } }, |
| "cosCacheT": { "dtype": "float32", "shape": [64, 8], "data": { "kind": "constant", "value": 0.8 } }, |
| "sinCacheT": { "dtype": "float32", "shape": [64, 8], "data": { "kind": "constant", "value": 0.6 } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [2, 1, 64], "tolerance": 0.00003, "relTolerance": 0.0003 }, |
| "pastKeyT": { "dtype": "float32", "shape": [2, 2, 64, 16], "tolerance": 0.00003, "relTolerance": 0.0003 }, |
| "pastValueT": { "dtype": "float32", "shape": [2, 2, 64, 16], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "chunked_separate_two_layouts", |
| "provenance": { |
| "notes": "Eight new tokens over 24 cached tokens, covering the non-prompt sequence_length > 1 path." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [2, 8, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [2, 8, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [2, 8, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { |
| "dtype": "float32", |
| "shape": [2, 2, 64, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 } |
| }, |
| "pastValueT": { |
| "dtype": "float32", |
| "shape": [2, 2, 64, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 } |
| }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 9], |
| "data": { |
| "kind": "values", |
| "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" } |
| } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [32, 32] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [2, 8, 32], "tolerance": 0.00003, "relTolerance": 0.0003 }, |
| "pastKeyT": { "dtype": "float32", "shape": [2, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [2, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "prompt_all_rows_gqa_four_to_one", |
| "provenance": { |
| "notes": "A 64-token prompt reaches every documented block row while four query heads share each KV head." |
| }, |
| "attrs": { "num_heads": 8, "kv_num_heads": 2, "sparse_block_size": 16 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 64, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 64, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 64, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 9], |
| "data": { |
| "kind": "values", |
| "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" } |
| } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 64, 64], "tolerance": 0.00003, "relTolerance": 0.0003 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "prompt_block32_head_dim16", |
| "provenance": { |
| "notes": "Block size 32 and head width 16 exercise independent key-tile and output-channel tails." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 32 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 64, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 64, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 64, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 128, 16], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 128, 16], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 9], |
| "data": { |
| "kind": "values", |
| "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" } |
| } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [64] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 64, 64], "tolerance": 0.00003, "relTolerance": 0.0003 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 128, 16], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 128, 16], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "prompt_explicit_scale", |
| "provenance": { "notes": "Explicit non-default score scale." }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16, "scale": 0.25 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 32, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 32, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 32, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 9], |
| "data": { |
| "kind": "values", |
| "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" } |
| } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 32, 32], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 64, 8], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "f16_prompt_separate", |
| "provenance": { "notes": "Meaningful float16 prompt; scores and online-softmax accumulators remain float32." }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [1, 32, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [1, 32, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [1, 32, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float16", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float16", "shape": [1, 2, 64, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 9], |
| "data": { |
| "kind": "values", |
| "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" } |
| } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float16", "shape": [1, 32, 32], "tolerance": 0.003, "relTolerance": 0.01 }, |
| "pastKeyT": { "dtype": "float16", "shape": [1, 2, 64, 8], "tolerance": 0.001, "relTolerance": 0.001 }, |
| "pastValueT": { "dtype": "float16", "shape": [1, 2, 64, 8], "tolerance": 0.001, "relTolerance": 0.001 } |
| } |
| }, |
| { |
| "name": "f16_decode_packed_partial_rotary", |
| "provenance": { "notes": "Float16 packed decode with partial NeoX rotary and unequal per-batch history lengths." }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 16, "do_rotary": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [2, 1, 256], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "pastKeyT": { |
| "dtype": "float16", |
| "shape": [2, 2, 64, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 } |
| }, |
| "pastValueT": { |
| "dtype": "float16", |
| "shape": [2, 2, 64, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 } |
| }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 8, 0, 1, 3, 6, 9] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 9], |
| "data": { |
| "kind": "values", |
| "values": { "$ref": "#/fixtureArrays/prompt_separate_two_layouts_input_blockColIndicesT" } |
| } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [41] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [41, 34] } }, |
| "cosCacheT": { "dtype": "float16", "shape": [64, 8], "data": { "kind": "constant", "value": 0.8 } }, |
| "sinCacheT": { "dtype": "float16", "shape": [64, 8], "data": { "kind": "constant", "value": 0.6 } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float16", "shape": [2, 1, 128], "tolerance": 0.003, "relTolerance": 0.01 }, |
| "pastKeyT": { "dtype": "float16", "shape": [2, 2, 64, 32], "tolerance": 0.002, "relTolerance": 0.005 }, |
| "pastValueT": { "dtype": "float16", "shape": [2, 2, 64, 32], "tolerance": 0.001, "relTolerance": 0.001 } |
| } |
| }, |
| { |
| "name": "dense_layout_ignores_malformed_columns", |
| "provenance": { |
| "notes": "ORT classifies a layout from its final CSR count. Once dense, repeated non-prefix column contents are ignored." |
| }, |
| "attrs": { "num_heads": 1, "kv_num_heads": 1, "sparse_block_size": 16 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 32, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 32, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 32, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 32, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 32, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 1, 3] } }, |
| "blockColIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [1, 1, 1] } }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [32] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 32, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 32, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 32, 8], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "sparse_duplicate_columns_are_idempotent", |
| "provenance": { |
| "notes": "A sparse row repeats block 0. CSR encodes a boolean mask, so the duplicate must not duplicate probability mass." |
| }, |
| "attrs": { "num_heads": 1, "kv_num_heads": 1, "sparse_block_size": 16 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 48, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 48, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 48, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 48, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 48, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { "dtype": "int32", "shape": [1, 4], "data": { "kind": "values", "values": [0, 1, 3, 5] } }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [1, 5], |
| "data": { "kind": "values", "values": [0, 0, 1, 0, 0] } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [48] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [48] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 48, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 48, 8], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 48, 8], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "empty_csr_row_uniform_scalar_total", |
| "provenance": { |
| "notes": "Pinned all-masked-row behavior from ORT's lowest-sentinel softmax: one causal key yields its value, not zero. Also covers scalar total_sequence_length and a zero-width column tensor." |
| }, |
| "attrs": { "num_heads": 1, "kv_num_heads": 1, "sparse_block_size": 16 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 1, 8], |
| "data": { "kind": "values", "values": [0.5, -0.5, 1.0, -1.0, 2.0, -2.0, 3.0, -3.0] } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 1, 8], |
| "data": { "kind": "values", "values": [8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0] } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 1, 8], |
| "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] } |
| }, |
| "pastKeyT": { |
| "dtype": "float32", |
| "shape": [1, 1, 16, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.25 } |
| }, |
| "pastValueT": { |
| "dtype": "float32", |
| "shape": [1, 1, 16, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.25 } |
| }, |
| "blockRowIndicesT": { "dtype": "int32", "shape": [1, 2], "data": { "kind": "values", "values": [0, 0] } }, |
| "blockColIndicesT": { "dtype": "int32", "shape": [1, 0], "data": { "kind": "values", "values": [] } }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [], "data": { "kind": "values", "values": [1] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1] } } |
| }, |
| "outputs": { |
| "outputT": { |
| "dtype": "float32", |
| "shape": [1, 1, 8], |
| "data": { "kind": "values", "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0] }, |
| "tolerance": 0 |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 16, 8], "tolerance": 0 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 16, 8], "tolerance": 0 } |
| } |
| }, |
| { |
| "name": "optional_rotary_input_ignored_when_attribute_is_not_one", |
| "provenance": { |
| "notes": "The CPU kernel enables do_rotary only for the exact integer value 1. Otherwise optional rotary inputs are ignored and need not be paired." |
| }, |
| "attrs": { "num_heads": 1, "kv_num_heads": 1, "sparse_block_size": 16, "do_rotary": 2, "rotary_interleaved": 2 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 16, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 16, 8], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { "dtype": "int32", "shape": [1, 2], "data": { "kind": "values", "values": [0, 1] } }, |
| "blockColIndicesT": { "dtype": "int32", "shape": [1, 1], "data": { "kind": "values", "values": [0] } }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1] } }, |
| "cosCacheT": { "dtype": "float32", "shape": [1, 1], "data": { "kind": "constant", "value": 0.0 } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 1, 8], "tolerance": 0.000001, "relTolerance": 0.000001 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 16, 8], "tolerance": 0.000001, "relTolerance": 0.000001 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 16, 8], "tolerance": 0.000001, "relTolerance": 0.000001 } |
| } |
| }, |
| { |
| "name": "zero_width_rotary_is_a_noop", |
| "provenance": { |
| "notes": "ONNX Runtime accepts zero-width paired caches with do_rotary=1; the resulting zero-dimensional rotary transform copies Q and K through." |
| }, |
| "attrs": { "num_heads": 1, "kv_num_heads": 1, "sparse_block_size": 16, "do_rotary": 1, "rotary_interleaved": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 16, 16], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 16, 16], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { "dtype": "int32", "shape": [1, 2], "data": { "kind": "values", "values": [0, 1] } }, |
| "blockColIndicesT": { "dtype": "int32", "shape": [1, 1], "data": { "kind": "values", "values": [0] } }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1] } }, |
| "cosCacheT": { "dtype": "float32", "shape": [1, 0], "data": { "kind": "values", "values": [] } }, |
| "sinCacheT": { "dtype": "float32", "shape": [1, 0], "data": { "kind": "values", "values": [] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 1, 16], "tolerance": 0.000001, "relTolerance": 0.000001 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 16, 16], "tolerance": 0.000001, "relTolerance": 0.000001 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 16, 16], "tolerance": 0.000001, "relTolerance": 0.000001 } |
| } |
| }, |
| { |
| "name": "sgmat_prompt_separate_blk64_duplicate_csr", |
| "provenance": { |
| "notes": "Prefill shape on the subgroup-matrix route: 64-token blocks, a duplicate CSR column whose repeated block must not repeat its mass, a dense-by-count row, and an off-diagonal-only row." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 64 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 192, 256], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 192, 128], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 192, 128], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 192, 64], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 192, 64], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 4], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 0, 1, 2, 3] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 0, 1, 1, 1, 0, 1, 2, 0, 0] } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [192] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [192] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 192, 256], "tolerance": 0.0001, "relTolerance": 0.0005 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 192, 64], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 192, 64], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "sgmat_prompt_packed_blk64", |
| "provenance": { |
| "notes": "Packed [Q|K|V] prefill on the subgroup-matrix route; layout 0 crosses the triangular dense-count threshold so its column values are ignored, layout 1 keeps a gapped sparse row." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "sparse_block_size": 64 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 128, 512], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 128, 64], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 128, 64], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 3], |
| "data": { "kind": "values", "values": [0, 1, 3, 0, 1, 2] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 3], |
| "data": { "kind": "values", "values": [0, 0, 1, 0, 1, 0] } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 128, 256], "tolerance": 0.0001, "relTolerance": 0.0005 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "sgmat_prompt_separate_rotary_neox_blk64", |
| "provenance": { |
| "notes": "Grouped-query prefill with partial NeoX rotary on the subgroup-matrix route; the attention pass reads the rotated-Q intermediate." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 1, "sparse_block_size": 64, "do_rotary": 1, "rotary_interleaved": 0 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 128, 128], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 128, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 128, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 64], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 64], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 3], |
| "data": { "kind": "values", "values": [0, 1, 3, 0, 1, 2] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 3], |
| "data": { "kind": "values", "values": [0, 0, 1, 0, 1, 0] } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }, |
| "cosCacheT": { "dtype": "float32", "shape": [128, 16], "data": { "kind": "constant", "value": 0.8 } }, |
| "sinCacheT": { "dtype": "float32", "shape": [128, 16], "data": { "kind": "constant", "value": 0.6 } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 128, 128], "tolerance": 0.0001, "relTolerance": 0.0005 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "sgmat_prompt_packed_rotary_interleaved_blk64", |
| "provenance": { |
| "notes": "Packed [Q|K|V] prefill with full interleaved rotary at head_size 32, the narrowest head the subgroup-matrix route admits (a single reduction chunk)." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64, "do_rotary": 1, "rotary_interleaved": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 128, 192], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 128, 32], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 128, 32], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 3], |
| "data": { "kind": "values", "values": [0, 1, 3, 0, 1, 2] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 3], |
| "data": { "kind": "values", "values": [0, 0, 1, 0, 1, 0] } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }, |
| "cosCacheT": { "dtype": "float32", "shape": [128, 16], "data": { "kind": "constant", "value": 0.8 } }, |
| "sinCacheT": { "dtype": "float32", "shape": [128, 16], "data": { "kind": "constant", "value": 0.6 } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 128, 64], "tolerance": 0.0001, "relTolerance": 0.0005 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 2, 128, 32], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 2, 128, 32], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "sgmat_prompt_tail_s100_explicit_scale", |
| "provenance": { |
| "notes": "A 100-token prefill leaves the second query tile 36 rows live, exercising the staged-Q zero padding and guarded stores; the explicit scale attribute overrides the head-size default." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 1, "sparse_block_size": 64, "scale": 0.25 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 100, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 100, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 100, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 32], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 32], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 1, 3] } }, |
| "blockColIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 0, 1] } }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [100] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [100] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 100, 64], "tolerance": 0.0001, "relTolerance": 0.0005 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 32], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 32], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "sgmat_prompt_blk128_subtiles", |
| "provenance": { |
| "notes": "128-token sparse blocks decompose into two 64-key tiles per block at head_size 128, the flagship geometry; the second mask row skips block 0 so its diagonal block carries the ragged causal edge." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 1, "sparse_block_size": 128 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 256, 256], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 256, 128], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 256, 128], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 256, 128], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 256, 128], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 1, 2] } }, |
| "blockColIndicesT": { "dtype": "int32", "shape": [1, 2], "data": { "kind": "values", "values": [0, 1] } }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 256, 256], "tolerance": 0.0001, "relTolerance": 0.0005 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 256, 128], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 256, 128], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "sgmat_chunked_two_mask_rows", |
| "provenance": { |
| "notes": "Chunked prefill with per-batch histories of 100 and 68 tokens: each 64-query tile straddles two mask rows at different split points, the misaligned path a runtime past length forces." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 2, "sparse_block_size": 64 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [2, 64, 128], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [2, 64, 128], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [2, 64, 128], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { |
| "dtype": "float32", |
| "shape": [2, 2, 256, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.011, "cosStep": 0.017, "scale": 0.25 } |
| }, |
| "pastValueT": { |
| "dtype": "float32", |
| "shape": [2, 2, 256, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.013, "scale": 0.25 } |
| }, |
| "blockRowIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 5], |
| "data": { "kind": "values", "values": [0, 1, 3, 5, 7, 0, 1, 2, 4, 5] } |
| }, |
| "blockColIndicesT": { |
| "dtype": "int32", |
| "shape": [2, 7], |
| "data": { "kind": "values", "values": [0, 0, 1, 0, 2, 1, 3, 0, 1, 0, 2, 3, 0, 0] } |
| }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }, |
| "keyTotalSequenceLengthsT": { |
| "dtype": "int32", |
| "shape": [2], |
| "data": { "kind": "values", "values": [164, 132] } |
| } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [2, 64, 128], "tolerance": 0.0001, "relTolerance": 0.0005 }, |
| "pastKeyT": { "dtype": "float32", "shape": [2, 2, 256, 64], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [2, 2, 256, 64], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| }, |
| { |
| "name": "sgmat_empty_row_uniform_fallback", |
| "provenance": { |
| "notes": "The second mask row selects no blocks, so all 64 of its queries carry zero probability mass and take the uniform average over causal history on the subgroup-matrix route." |
| }, |
| "attrs": { "num_heads": 1, "kv_num_heads": 1, "sparse_block_size": 64 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float32", |
| "shape": [1, 128, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.5 } |
| }, |
| "keyT": { |
| "dtype": "float32", |
| "shape": [1, 128, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.5 } |
| }, |
| "valueT": { |
| "dtype": "float32", |
| "shape": [1, 128, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 0.5 } |
| }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 64], "data": { "kind": "constant", "value": 0.0 } }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 64], "data": { "kind": "constant", "value": 0.0 } }, |
| "blockRowIndicesT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 1, 1] } }, |
| "blockColIndicesT": { "dtype": "int32", "shape": [1, 1], "data": { "kind": "values", "values": [0] } }, |
| "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } }, |
| "keyTotalSequenceLengthsT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [128] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float32", "shape": [1, 128, 64], "tolerance": 0.0001, "relTolerance": 0.0005 }, |
| "pastKeyT": { "dtype": "float32", "shape": [1, 1, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 }, |
| "pastValueT": { "dtype": "float32", "shape": [1, 1, 128, 64], "tolerance": 0.00002, "relTolerance": 0.0002 } |
| } |
| } |
| ] |
| } |
|
|
