File size: 7,184 Bytes
2da2397
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fd6745d
2da2397
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
{
  "tunableSpace": {
    "WORKGROUP_SIZE": [32, 64, 128],
    "SCATTER_WORKGROUP_SIZE": [64, 128, 256],
    "SPLIT_TARGET_WORKGROUPS": [512, 1024, 2048, 4096],
    "SPLIT_MIN_KEYS": [64, 128, 256],
    "MAX_SPLITS": [8, 16, 32]
  },
  "cases": [
    {
      "name": "paged-decode-b8-past1024-h32kv8-d128",
      "preset": "smoke",
      "vars": { "dtype": "float16", "batch": 8, "past": 1024, "heads": 32, "kvHeads": 8, "headDim": 128 },
      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
      "inputs": {
        "queryT": { "shape": [8, 4096], "dtype": "float16", "dist": "normal", "seed": 9201, "scale": 1 },
        "keyT": { "shape": [8, 1024], "dtype": "float16", "dist": "normal", "seed": 9202, "scale": 1 },
        "valueT": { "shape": [8, 1024], "dtype": "float16", "dist": "normal", "seed": 9203, "scale": 1 },
        "keyCacheT": { "shape": [512, 16, 8, 128], "dtype": "float16", "dist": "normal", "seed": 9204, "scale": 1 },
        "valueCacheT": { "shape": [512, 16, 8, 128], "dtype": "float16", "dist": "normal", "seed": 9205, "scale": 1 },
        "cumulativeSequenceLengthT": {
          "shape": [9],
          "dtype": "int32",
          "data": { "kind": "values", "values": [0, 1, 2, 3, 4, 5, 6, 7, 8] }
        },
        "pastSeqlensT": { "shape": [8], "dtype": "int32", "dist": "constant", "value": 1023 },
        "blockTableT": { "shape": [8, 64], "dtype": "int32", "dist": "linearMod", "seed": 9208, "step": 1, "mod": 512 }
      },
      "outputs": { "outputT": { "shape": [8, 4096], "dtype": "float16" } },
      "bench": {
        "metrics": [{ "type": "bandwidth", "value": "2 * args.batch * args.kvHeads * args.past * args.headDim * 2" }]
      },
      "provenance": {
        "notes": "The bandwidth metric derives KV traffic from the key/value head count and the two-byte float16 cache element size."
      }
    },
    {
      "name": "paged-decode-llama-b16-past2048-h32kv32-d128",
      "preset": "model",
      "provenance": {
        "notes": "Llama class defaults (32 query heads, no GQA, head_dim 128) decoding 16 concurrent sequences with 2048 tokens of paged KV each."
      },
      "vars": { "dtype": "float16", "batch": 16, "past": 2048, "heads": 32, "kvHeads": 32, "headDim": 128 },
      "attrs": { "num_heads": 32, "kv_num_heads": 32 },
      "inputs": {
        "queryT": { "shape": [16, 4096], "dtype": "float16", "dist": "normal", "seed": 7600, "scale": 0.3 },
        "keyT": { "shape": [16, 4096], "dtype": "float16", "dist": "normal", "seed": 7601, "scale": 0.3 },
        "valueT": { "shape": [16, 4096], "dtype": "float16", "dist": "normal", "seed": 7602, "scale": 0.3 },
        "keyCacheT": { "shape": [2048, 16, 32, 128], "dtype": "float16", "dist": "normal", "seed": 7603, "scale": 0.3 },
        "valueCacheT": {
          "shape": [2048, 16, 32, 128],
          "dtype": "float16",
          "dist": "normal",
          "seed": 7604,
          "scale": 0.3
        },
        "cumulativeSequenceLengthT": {
          "shape": [17],
          "dtype": "int32",
          "dist": "linearMod",
          "seed": 7605,
          "step": 1,
          "mod": 17
        },
        "pastSeqlensT": { "shape": [16], "dtype": "int32", "dist": "constant", "value": 2048 },
        "blockTableT": {
          "shape": [16, 128],
          "dtype": "int32",
          "dist": "linearMod",
          "seed": 7606,
          "step": 1,
          "mod": 2048
        }
      },
      "outputs": { "outputT": { "shape": [16, 4096], "dtype": "float16" } },
      "bench": {
        "metrics": [{ "type": "bandwidth", "value": "2 * args.batch * args.kvHeads * args.past * args.headDim * 2" }]
      }
    },
    {
      "name": "paged-decode-mistral-b16-past4096-h32kv8-d128",
      "preset": "model",
      "provenance": {
        "notes": "Mistral class defaults (32 query heads, 8 key/value heads, head_dim 128) at 4096 tokens of history -- the GQA case, where the cache is a quarter the size for the same context."
      },
      "vars": { "dtype": "float16", "batch": 16, "past": 4096, "heads": 32, "kvHeads": 8, "headDim": 128 },
      "attrs": { "num_heads": 32, "kv_num_heads": 8 },
      "inputs": {
        "queryT": { "shape": [16, 4096], "dtype": "float16", "dist": "normal", "seed": 7700, "scale": 0.3 },
        "keyT": { "shape": [16, 1024], "dtype": "float16", "dist": "normal", "seed": 7701, "scale": 0.3 },
        "valueT": { "shape": [16, 1024], "dtype": "float16", "dist": "normal", "seed": 7702, "scale": 0.3 },
        "keyCacheT": { "shape": [4096, 16, 8, 128], "dtype": "float16", "dist": "normal", "seed": 7703, "scale": 0.3 },
        "valueCacheT": { "shape": [4096, 16, 8, 128], "dtype": "float16", "dist": "normal", "seed": 7704, "scale": 0.3 },
        "cumulativeSequenceLengthT": {
          "shape": [17],
          "dtype": "int32",
          "dist": "linearMod",
          "seed": 7705,
          "step": 1,
          "mod": 17
        },
        "pastSeqlensT": { "shape": [16], "dtype": "int32", "dist": "constant", "value": 4096 },
        "blockTableT": {
          "shape": [16, 256],
          "dtype": "int32",
          "dist": "linearMod",
          "seed": 7706,
          "step": 1,
          "mod": 4096
        }
      },
      "outputs": { "outputT": { "shape": [16, 4096], "dtype": "float16" } },
      "bench": {
        "metrics": [{ "type": "bandwidth", "value": "2 * args.batch * args.kvHeads * args.past * args.headDim * 2" }]
      }
    },
    {
      "name": "paged-decode-gemma3-b8-past2048-h8kv4-d256",
      "preset": "model",
      "provenance": {
        "notes": "Gemma3 class defaults (8 query heads, 4 key/value heads, head_dim 256) -- few heads but the widest head dimension in this set."
      },
      "vars": { "dtype": "float16", "batch": 8, "past": 2048, "heads": 8, "kvHeads": 4, "headDim": 256 },
      "attrs": { "num_heads": 8, "kv_num_heads": 4 },
      "inputs": {
        "queryT": { "shape": [8, 2048], "dtype": "float16", "dist": "normal", "seed": 7800, "scale": 0.3 },
        "keyT": { "shape": [8, 1024], "dtype": "float16", "dist": "normal", "seed": 7801, "scale": 0.3 },
        "valueT": { "shape": [8, 1024], "dtype": "float16", "dist": "normal", "seed": 7802, "scale": 0.3 },
        "keyCacheT": { "shape": [1024, 16, 4, 256], "dtype": "float16", "dist": "normal", "seed": 7803, "scale": 0.3 },
        "valueCacheT": { "shape": [1024, 16, 4, 256], "dtype": "float16", "dist": "normal", "seed": 7804, "scale": 0.3 },
        "cumulativeSequenceLengthT": {
          "shape": [9],
          "dtype": "int32",
          "dist": "linearMod",
          "seed": 7805,
          "step": 1,
          "mod": 9
        },
        "pastSeqlensT": { "shape": [8], "dtype": "int32", "dist": "constant", "value": 2048 },
        "blockTableT": {
          "shape": [8, 128],
          "dtype": "int32",
          "dist": "linearMod",
          "seed": 7806,
          "step": 1,
          "mod": 1024
        }
      },
      "outputs": { "outputT": { "shape": [8, 2048], "dtype": "float16" } },
      "bench": {
        "metrics": [{ "type": "bandwidth", "value": "2 * args.batch * args.kvHeads * args.past * args.headDim * 2" }]
      }
    }
  ]
}