{ "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 } } } ] }