| { |
| "op": "com.microsoft.PagedAttention", |
| "cases": [ |
| { |
| "name": "two_sequences_gqa", |
| "provenance": { |
| "notes": "Two packed sequences with different histories. The first sequence's two new tokens land at absolute slots 3 and 4, so the scatter crosses a block boundary and the attention reads two physical blocks." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [5, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [5, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [5, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { |
| "dtype": "int32", |
| "shape": [3], |
| "data": { "kind": "values", "values": [0, 2, 5] } |
| }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } }, |
| "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 } } |
| }, |
| { |
| "name": "two_sequences_packed_qkv", |
| "provenance": { "notes": "Same schedule with Q, K and V packed on one query row." }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [5, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { |
| "dtype": "int32", |
| "shape": [3], |
| "data": { "kind": "values", "values": [0, 2, 5] } |
| }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } }, |
| "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 } } |
| }, |
| { |
| "name": "two_sequences_key_cache_output", |
| "provenance": { |
| "notes": "Requests only the updated key-cache alias; the paired schema value-cache output remains internal." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [5, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [5, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [5, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { |
| "dtype": "int32", |
| "shape": [3], |
| "data": { "kind": "values", "values": [0, 2, 5] } |
| }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } }, |
| "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 }, |
| "keyCacheT": { "dtype": "float16", "shape": [6, 4, 2, 8], "tolerance": 0.00001, "relTolerance": 0.0001 } |
| } |
| }, |
| { |
| "name": "explicit_scale", |
| "attrs": { "num_heads": 4, "kv_num_heads": 2, "scale": 0.125, "is_causal": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [5, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [5, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [5, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { |
| "dtype": "int32", |
| "shape": [3], |
| "data": { "kind": "values", "values": [0, 2, 5] } |
| }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } }, |
| "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 } } |
| }, |
| { |
| "name": "slot_mapping_explicit", |
| "provenance": { |
| "notes": "slot_mapping names each token's flat cache slot directly; these values reproduce the derived mapping, and this case requests only the updated value-cache alias." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [5, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [5, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [5, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { |
| "dtype": "int32", |
| "shape": [3], |
| "data": { "kind": "values", "values": [0, 2, 5] } |
| }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } }, |
| "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } }, |
| "slotMappingT": { "dtype": "int32", "shape": [5], "data": { "kind": "values", "values": [3, 4, 12, 13, 14] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 }, |
| "valueCacheT": { "dtype": "float16", "shape": [6, 4, 2, 8], "tolerance": 0.00001, "relTolerance": 0.0001 } |
| } |
| }, |
| { |
| "name": "slot_mapping_skips_token", |
| "provenance": { |
| "notes": "A slot of -1 suppresses that token's cache write, which a scheduler uses for a prefix-cache hit or a rejected speculative token. The cache slot keeps its prior contents and attention still reads it." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [5, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [5, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [5, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { |
| "dtype": "int32", |
| "shape": [3], |
| "data": { "kind": "values", "values": [0, 2, 5] } |
| }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } }, |
| "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } }, |
| "slotMappingT": { "dtype": "int32", "shape": [5], "data": { "kind": "values", "values": [3, -1, 12, 13, 14] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 }, |
| "keyCacheT": { "dtype": "float16", "shape": [6, 4, 2, 8], "tolerance": 0.00001, "relTolerance": 0.0001 }, |
| "valueCacheT": { "dtype": "float16", "shape": [6, 4, 2, 8], "tolerance": 0.00001, "relTolerance": 0.0001 } |
| } |
| }, |
| { |
| "name": "packed_qkv_slot_mapping", |
| "provenance": { |
| "notes": "Packed [Q|K|V] with an explicit slot mapping, so the scatter reads K and V out of the query row and binds no schedule tensors at all." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [5, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { |
| "dtype": "int32", |
| "shape": [3], |
| "data": { "kind": "values", "values": [0, 2, 5] } |
| }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } }, |
| "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } }, |
| "slotMappingT": { "dtype": "int32", "shape": [5], "data": { "kind": "values", "values": [3, 4, 12, 13, 14] } } |
| }, |
| "outputs": { |
| "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 }, |
| "keyCacheT": { "dtype": "float16", "shape": [6, 4, 2, 8], "tolerance": 0.00001, "relTolerance": 0.0001 }, |
| "valueCacheT": { "dtype": "float16", "shape": [6, 4, 2, 8], "tolerance": 0.00001, "relTolerance": 0.0001 } |
| } |
| }, |
| { |
| "name": "single_sequence_decode", |
| "provenance": { |
| "notes": "One decode token whose history spans two full blocks and lands at the start of the third." |
| }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 2 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [1, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [3, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [3, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 1] } }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [7] } }, |
| "blockTableT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [0, 1, 2] } } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [1, 32], "tolerance": 0.00001, "relTolerance": 0.0001 } } |
| }, |
| { |
| "name": "equal_head_counts", |
| "provenance": { "notes": "num_heads equals kv_num_heads, so the grouped-query mapping is the identity." }, |
| "attrs": { "num_heads": 4, "kv_num_heads": 4 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [5, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [5, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [5, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 4, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 4, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { |
| "dtype": "int32", |
| "shape": [3], |
| "data": { "kind": "values", "values": [0, 2, 5] } |
| }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } }, |
| "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.00001, "relTolerance": 0.0001 } } |
| }, |
| { |
| "name": "long_history_multi_tile", |
| "provenance": { |
| "notes": "A 71-key history exceeds one workgroup tile, and the block table is deliberately not the identity, so the online-softmax merge runs over multiple tiles rather than a single one." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [9, 8, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [9, 8, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 1] } }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [70] } }, |
| "blockTableT": { |
| "dtype": "int32", |
| "shape": [1, 9], |
| "data": { "kind": "values", "values": [5, 2, 8, 0, 7, 1, 6, 3, 4] } |
| } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } } |
| }, |
| { |
| "name": "head_size_16", |
| "provenance": { |
| "notes": "A head dimension wider than the tile's per-thread share, so the accumulator loop strides." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [3, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [3, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [3, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [4, 4, 1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [4, 4, 1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { |
| "dtype": "int32", |
| "shape": [3], |
| "data": { "kind": "values", "values": [0, 2, 3] } |
| }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [1, 5] } }, |
| "blockTableT": { "dtype": "int32", "shape": [2, 2], "data": { "kind": "values", "values": [0, 1, 2, 3] } } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [3, 32], "tolerance": 0.00001, "relTolerance": 0.0001 } } |
| }, |
| { |
| "name": "f16_two_sequences", |
| "attrs": { "num_heads": 4, "kv_num_heads": 2 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [5, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [5, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [5, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { |
| "dtype": "int32", |
| "shape": [3], |
| "data": { "kind": "values", "values": [0, 2, 5] } |
| }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } }, |
| "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.004, "relTolerance": 0.01 } } |
| }, |
| { |
| "name": "f16_packed_qkv", |
| "attrs": { "num_heads": 4, "kv_num_heads": 2 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [5, 64], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [6, 4, 2, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { |
| "dtype": "int32", |
| "shape": [3], |
| "data": { "kind": "values", "values": [0, 2, 5] } |
| }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [3, 0] } }, |
| "blockTableT": { "dtype": "int32", "shape": [2, 3], "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5] } } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [5, 32], "tolerance": 0.004, "relTolerance": 0.01 } } |
| }, |
| { |
| "name": "pinned_paged_history", |
| "provenance": { |
| "notes": "Expected values computed by an independent implementation written from the ONNX Runtime schema text alone, with every tensor value rounded to its declared float16 dtype before evaluation, so this case checks the trusted reference as well as the kernels. The block table [2, 0, 1] is deliberately not the identity, so a kernel that ignored it would read the wrong blocks; the two new tokens land in blocks 0 and 1 respectively." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [2, 4], |
| "data": { "kind": "values", "values": [0.7854, 1.0234, 1.039, 0.8123, 0.3859, -0.1448, -0.6562, -1.028] } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [2, 2], |
| "data": { "kind": "values", "values": [-0.3046, -0.0781, 0.1597, 0.3877] } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [2, 2], |
| "data": { "kind": "values", "values": [0.5647, 0.756, 0.8391, 0.7958] } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [3, 2, 1, 2], |
| "data": { |
| "kind": "values", |
| "values": [0.2751, 0.3295, 0.3637, 0.3794, 0.3798, 0.369, 0.3514, 0.3311, 0.3118, 0.296, 0.2849, 0.2785] |
| } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [3, 2, 1, 2], |
| "data": { |
| "kind": "values", |
| "values": [-0.0755, 0.1377, 0.2257, 0.1768, 0.0339, -0.1232, -0.2093, -0.1675, 0.0079, 0.2677, 0.5238, 0.6804] |
| } |
| }, |
| "cumulativeSequenceLengthT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 2] } }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2] } }, |
| "blockTableT": { "dtype": "int32", "shape": [1, 3], "data": { "kind": "values", "values": [2, 0, 1] } } |
| }, |
| "outputs": { |
| "outputT": { |
| "dtype": "float16", |
| "shape": [2, 4], |
| "data": { |
| "kind": "values", |
| "values": [0.32763671875, 0.5322265625, 0.32470703125, 0.52978515625, 0.4765625, 0.6181640625, 0.4970703125, 0.64599609375] |
| }, |
| "tolerance": 0.000001, |
| "relTolerance": 0.00001 |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [3, 2, 1, 2], |
| "data": { |
| "kind": "values", |
| "values": [-0.3046, -0.0781, 0.1597, 0.3877, 0.3798, 0.369, 0.3514, 0.3311, 0.3118, 0.296, 0.2849, 0.2785] |
| }, |
| "tolerance": 0.000001, |
| "relTolerance": 0.000001 |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [3, 2, 1, 2], |
| "data": { |
| "kind": "values", |
| "values": [0.5647, 0.756, 0.8391, 0.7958, 0.0339, -0.1232, -0.2093, -0.1675, 0.0079, 0.2677, 0.5238, 0.6804] |
| }, |
| "tolerance": 0.000001, |
| "relTolerance": 0.000001 |
| } |
| } |
| }, |
| { |
| "name": "splitk_multi_split_history", |
| "provenance": { |
| "notes": "A 301-key history over a 512-key block table, which is what puts the split-K schedule above its two-split floor: four contiguous key ranges of 76 keys each cross a workgroup tile boundary, so every range runs the tile loop more than once and the merge folds four non-empty online states. Grouped-query (2 query heads per KV head) so the split epilogue writes both streams." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [64, 8, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [64, 8, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 1] } }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [300] } }, |
| "blockTableT": { |
| "dtype": "int32", |
| "shape": [1, 64], |
| "data": { |
| "kind": "values", |
| "values": [3, 10, 17, 24, 31, 38, 45, 52, 59, 2, 9, 16, 23, 30, 37, 44, 51, 58, 1, 8, 15, 22, 29, 36, 43, 50, 57, 0, 7, 14, 21, 28, 35, 42, 49, 56, 63, 6, 13, 20, 27, 34, 41, 48, 55, 62, 5, 12, 19, 26, 33, 40, 47, 54, 61, 4, 11, 18, 25, 32, 39, 46, 53, 60] |
| } |
| } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } } |
| }, |
| { |
| "name": "splitk_empty_trailing_split", |
| "provenance": { |
| "notes": "One key of history over a 256-key block table: the split schedule is two ranges of one key each, so the second range covers nothing. It is the split that must write the (-FLT_MAX, 0) identity rather than a normalized row, and the merge must fold that identity away." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [32, 8, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [32, 8, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 1] } }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [0] } }, |
| "blockTableT": { |
| "dtype": "int32", |
| "shape": [1, 32], |
| "data": { |
| "kind": "values", |
| "values": [5, 12, 19, 26, 1, 8, 15, 22, 29, 4, 11, 18, 25, 0, 7, 14, 21, 28, 3, 10, 17, 24, 31, 6, 13, 20, 27, 2, 9, 16, 23, 30] |
| } |
| } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } } |
| }, |
| { |
| "name": "splitk_slot_mapping", |
| "provenance": { |
| "notes": "The split-K schedule with an explicit slot mapping instead of a derived write slot: same 201-key history over a 256-key block table, so the two split ranges both run and the scatter pass takes the slot-mapping arm." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [1, 16], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyT": { |
| "dtype": "float16", |
| "shape": [1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.13, "scale": 0.8 } |
| }, |
| "valueT": { |
| "dtype": "float16", |
| "shape": [1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.07, "scale": 0.7 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [32, 8, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [32, 8, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 1] } }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [200] } }, |
| "blockTableT": { |
| "dtype": "int32", |
| "shape": [1, 32], |
| "data": { |
| "kind": "values", |
| "values": [7, 14, 21, 28, 3, 10, 17, 24, 31, 6, 13, 20, 27, 2, 9, 16, 23, 30, 5, 12, 19, 26, 1, 8, 15, 22, 29, 4, 11, 18, 25, 0] |
| } |
| }, |
| "slotMappingT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [176] } } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } } |
| }, |
| { |
| "name": "splitk_packed_qkv", |
| "provenance": { |
| "notes": "The split-K schedule reading a packed [Q; K; V] query row, so the split pass strides the packed query layout while the merge writes the unpacked output." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [1, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [32, 8, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [32, 8, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 1] } }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [200] } }, |
| "blockTableT": { |
| "dtype": "int32", |
| "shape": [1, 32], |
| "data": { |
| "kind": "values", |
| "values": [9, 16, 23, 30, 5, 12, 19, 26, 1, 8, 15, 22, 29, 4, 11, 18, 25, 0, 7, 14, 21, 28, 3, 10, 17, 24, 31, 6, 13, 20, 27, 2] |
| } |
| } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } } |
| }, |
| { |
| "name": "splitk_packed_qkv_slot_mapping", |
| "provenance": { |
| "notes": "The split-K schedule with both a packed query row and an explicit slot mapping, the last of the four scatter/query combinations the family covers." |
| }, |
| "attrs": { "num_heads": 2, "kv_num_heads": 1 }, |
| "inputs": { |
| "queryT": { |
| "dtype": "float16", |
| "shape": [1, 32], |
| "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31, "scale": 0.9 } |
| }, |
| "keyCacheT": { |
| "dtype": "float16", |
| "shape": [32, 8, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.37, "scale": 0.6 } |
| }, |
| "valueCacheT": { |
| "dtype": "float16", |
| "shape": [32, 8, 1, 8], |
| "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 0.5 } |
| }, |
| "cumulativeSequenceLengthT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 1] } }, |
| "pastSeqlensT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [200] } }, |
| "blockTableT": { |
| "dtype": "int32", |
| "shape": [1, 32], |
| "data": { |
| "kind": "values", |
| "values": [11, 18, 25, 0, 7, 14, 21, 28, 3, 10, 17, 24, 31, 6, 13, 20, 27, 2, 9, 16, 23, 30, 5, 12, 19, 26, 1, 8, 15, 22, 29, 4] |
| } |
| }, |
| "slotMappingT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [208] } } |
| }, |
| "outputs": { "outputT": { "dtype": "float16", "shape": [1, 16], "tolerance": 0.00001, "relTolerance": 0.0001 } } |
| } |
| ] |
| } |
|
|
