[Kernel][Helion][1/N] Add Helion kernel for fused_qk_norm_rope (#44010)
Signed-off-by: Sean Chen <seachen@redhat.com>
Xiaohong (Sean) Chen committed
Jun 29, 2026 at 10:54 UTC
49e28e8e91ad9ed102e88e50c190686408be5552
4 files changed
+5911
tests/kernels/helion/test_fused_qk_norm_rope.py
new
+261
@@ -0,0 +1,261 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the fused_qk_norm_rope helion kernel
4
+
5
+Run `pytest tests/kernels/helion/test_fused_qk_norm_rope.py`.
6
+"""
7
+
8
+from typing import Any
9
+
10
+import pytest
11
+import torch
12
+from torch._subclasses.fake_tensor import FakeTensorMode
13
+
14
+from tests.kernels.helion.utils import skip_if_platform_unsupported
15
+from vllm.benchmarks.lib.utils import default_vllm_config
16
+from vllm.kernels.helion.case_key import CaseKey
17
+from vllm.kernels.helion.config_manager import ConfigManager
18
+from vllm.kernels.helion.ops.fused_qk_norm_rope import (
19
+ _pick_cache,
20
+ baseline,
21
+ fused_qk_norm_rope,
22
+ pick_config,
23
+)
24
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
25
+from vllm.utils.import_utils import has_helion
26
+
27
+if not has_helion():
28
+ pytest.skip(
29
+ "Helion is not installed. Install with: pip install vllm[helion]",
30
+ allow_module_level=True,
31
+ )
32
+
33
+
34
+@default_vllm_config()
35
+def _generate_fake_input(
36
+ num_tokens: int, num_q_heads: int, num_kv_heads: int
37
+) -> tuple[Any, ...]:
38
+ with FakeTensorMode():
39
+ head_dim = 128
40
+ eps = 1e-6
41
+ is_neox = True
42
+ rotary_ratio = 1.0
43
+ device = "cuda"
44
+ dtype = torch.bfloat16
45
+ total_dim = (num_q_heads + 2 * num_kv_heads) * head_dim
46
+ qkv = torch.randn(num_tokens, total_dim, dtype=dtype, device=device)
47
+ positions = torch.arange(num_tokens, dtype=torch.long, device=device)
48
+ q_weight = torch.normal(
49
+ mean=1.0,
50
+ std=1.0,
51
+ size=(head_dim,),
52
+ dtype=qkv.dtype,
53
+ device=device,
54
+ )
55
+ k_weight = torch.normal(
56
+ mean=1.0,
57
+ std=1.0,
58
+ size=(head_dim,),
59
+ dtype=qkv.dtype,
60
+ device=device,
61
+ )
62
+ rotary_dim = int(head_dim * rotary_ratio)
63
+ rope = RotaryEmbedding(
64
+ head_size=head_dim,
65
+ rotary_dim=rotary_dim,
66
+ max_position_embeddings=4096,
67
+ base=10000.0,
68
+ is_neox_style=is_neox,
69
+ dtype=dtype,
70
+ ).to(device)
71
+ args = (
72
+ qkv,
73
+ num_q_heads,
74
+ num_kv_heads,
75
+ num_kv_heads,
76
+ head_dim,
77
+ eps,
78
+ q_weight,
79
+ k_weight,
80
+ rope.cos_sin_cache,
81
+ is_neox,
82
+ positions.view(-1),
83
+ )
84
+ return args
85
+
86
+
87
+@pytest.fixture(autouse=True)
88
+def reset_config_manager_singleton():
89
+ ConfigManager.reset_instance()
90
+ ConfigManager()
91
+ yield
92
+ ConfigManager.reset_instance()
93
+
94
+
95
+class TestFusedQkNormRopeConfigPicker:
96
+ def setup_method(self):
97
+ _pick_cache.clear()
98
+
99
+ def test_config_picker_exact_match(self):
100
+ config_keys = [
101
+ CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 16}),
102
+ CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}),
103
+ ]
104
+
105
+ args = _generate_fake_input(16, 4096, 128)
106
+ selected_key = pick_config(args, config_keys)
107
+ assert selected_key == CaseKey(
108
+ {"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}
109
+ )
110
+
111
+ def test_config_picker_closest_match(self):
112
+ config_keys = [
113
+ CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 16}),
114
+ CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 32}),
115
+ CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 16}),
116
+ CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 32}),
117
+ CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 16}),
118
+ CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 32}),
119
+ CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}),
120
+ CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 32}),
121
+ ]
122
+
123
+ args = _generate_fake_input(20, 3000, 70)
124
+ selected_key = pick_config(args, config_keys)
125
+ assert selected_key == CaseKey(
126
+ {"q_heads": 2048, "kv_heads": 64, "num_tokens": 32}
127
+ )
128
+
129
+ def test_config_picker_no_configs(self):
130
+ config_keys: list[dict] = []
131
+
132
+ args = _generate_fake_input(16, 4096, 128)
133
+ selected_key = pick_config(args, config_keys)
134
+ assert selected_key is None
135
+
136
+ def test_config_picker_fallback_to_largest(self):
137
+ config_keys = [
138
+ CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 16}),
139
+ CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 32}),
140
+ CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 16}),
141
+ CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 32}),
142
+ CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 16}),
143
+ CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 32}),
144
+ CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}),
145
+ CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 32}),
146
+ ]
147
+
148
+ args = _generate_fake_input(64, 8192, 256)
149
+ selected_key = pick_config(args, config_keys)
150
+ assert selected_key == CaseKey(
151
+ {"q_heads": 4096, "kv_heads": 128, "num_tokens": 32}
152
+ )
153
+
154
+
155
+class TestFusedQkNormRopeCorrectness:
156
+ @pytest.mark.parametrize(
157
+ "num_heads, num_kv_heads, head_dim", [(16, 4, 128), (64, 8, 128)]
158
+ )
159
+ @pytest.mark.parametrize("num_tokens", [1, 7, 1024, 1025])
160
+ @pytest.mark.parametrize("is_neox", [False, True])
161
+ @pytest.mark.parametrize("rotary_ratio", [1.0, 0.5, 0.25])
162
+ @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
163
+ @default_vllm_config()
164
+ def test_fused_qk_norm_rope(
165
+ self,
166
+ num_heads: int,
167
+ num_kv_heads: int,
168
+ head_dim: int,
169
+ num_tokens: int,
170
+ is_neox: bool,
171
+ rotary_ratio: float,
172
+ dtype: torch.dtype,
173
+ ):
174
+ skip_if_platform_unsupported("fused_qk_norm_rope")
175
+
176
+ torch.manual_seed(42)
177
+ eps = 1e-6
178
+ device = "cuda"
179
+ total_dim = (num_heads + 2 * num_kv_heads) * head_dim
180
+ ref_qkv = torch.empty(
181
+ num_tokens, total_dim, dtype=dtype, device=device
182
+ ).uniform_(-0.1, 0.1)
183
+ ops_qkv = ref_qkv.clone()
184
+ positions = torch.arange(num_tokens, dtype=torch.long, device=device)
185
+ q_weight = torch.empty(head_dim, dtype=dtype, device=device).uniform_(0.8, 1.2)
186
+ k_weight = torch.empty(head_dim, dtype=dtype, device=device).uniform_(0.8, 1.2)
187
+ rotary_dim = int(head_dim * rotary_ratio)
188
+ rope = RotaryEmbedding(
189
+ head_size=head_dim,
190
+ rotary_dim=rotary_dim,
191
+ max_position_embeddings=40960,
192
+ base=10000.0,
193
+ is_neox_style=is_neox,
194
+ dtype=dtype,
195
+ ).to(device)
196
+
197
+ baseline(
198
+ ref_qkv,
199
+ num_heads,
200
+ num_kv_heads,
201
+ num_kv_heads,
202
+ head_dim,
203
+ eps,
204
+ q_weight,
205
+ k_weight,
206
+ rope.cos_sin_cache,
207
+ is_neox,
208
+ positions.view(-1),
209
+ )
210
+
211
+ fused_qk_norm_rope(
212
+ ops_qkv,
213
+ num_heads,
214
+ num_kv_heads,
215
+ num_kv_heads,
216
+ head_dim,
217
+ eps,
218
+ q_weight,
219
+ k_weight,
220
+ rope.cos_sin_cache,
221
+ is_neox,
222
+ positions.view(-1),
223
+ )
224
+
225
+ if dtype == torch.bfloat16:
226
+ atol = 5e-2
227
+ rtol = 5e-2
228
+ else:
229
+ atol = 1e-2
230
+ rtol = 1e-2
231
+
232
+ torch.testing.assert_close(
233
+ ref_qkv,
234
+ ops_qkv,
235
+ atol=atol,
236
+ rtol=rtol,
237
+ )
238
+
239
+
240
+class TestFusedQkNormRopeIntegration:
241
+ def test_kernel_registration_integration(self):
242
+ from vllm.kernels.helion.register import get_registered_kernels
243
+
244
+ registered_kernels = get_registered_kernels()
245
+ assert "fused_qk_norm_rope" in registered_kernels
246
+
247
+ kernel_wrapper = registered_kernels["fused_qk_norm_rope"]
248
+ assert kernel_wrapper.op_name == "fused_qk_norm_rope"
249
+ assert kernel_wrapper._config_picker is not None
250
+ assert kernel_wrapper._mutates_args == ["qkv"]
251
+
252
+ def test_fake_impl_functionality(self):
253
+ skip_if_platform_unsupported("fused_qk_norm_rope")
254
+ from vllm.kernels.helion.register import get_registered_kernels
255
+
256
+ registered_kernels = get_registered_kernels()
257
+ kernel_wrapper = registered_kernels["fused_qk_norm_rope"]
258
+ fake_impl = kernel_wrapper._fake_impl
259
+
260
+ args = _generate_fake_input(16, 4096, 128)
261
+ assert fake_impl(*args) is None
vllm/kernels/helion/configs/fused_qk_norm_rope/nvidia_b200.json
new
+2612
@@ -0,0 +1,2612 @@
1
+[
2
+ {
3
+ "key": {
4
+ "q_heads": 16,
5
+ "kv_heads": 8,
6
+ "num_tokens": 1
7
+ },
8
+ "config": {
9
+ "block_sizes": [
10
+ 4
11
+ ],
12
+ "loop_orders": [
13
+ [
14
+ 2,
15
+ 1,
16
+ 0
17
+ ]
18
+ ],
19
+ "l2_groupings": [
20
+ 8
21
+ ],
22
+ "range_unroll_factors": [
23
+ 0
24
+ ],
25
+ "range_warp_specializes": [
26
+ null
27
+ ],
28
+ "range_num_stages": [],
29
+ "range_multi_buffers": [
30
+ null
31
+ ],
32
+ "range_flattens": [
33
+ null
34
+ ],
35
+ "load_eviction_policies": [
36
+ "",
37
+ "first",
38
+ "",
39
+ "last",
40
+ "first",
41
+ "",
42
+ "last",
43
+ "first"
44
+ ],
45
+ "num_warps": 1,
46
+ "num_stages": 5,
47
+ "indexing": [
48
+ "tensor_descriptor",
49
+ "pointer",
50
+ "pointer",
51
+ "tensor_descriptor",
52
+ "tensor_descriptor",
53
+ "pointer",
54
+ "tensor_descriptor",
55
+ "pointer",
56
+ "tensor_descriptor",
57
+ "tensor_descriptor",
58
+ "pointer"
59
+ ],
60
+ "atomic_indexing": [],
61
+ "pid_type": "flat"
62
+ }
63
+ },
64
+ {
65
+ "key": {
66
+ "q_heads": 32,
67
+ "kv_heads": 8,
68
+ "num_tokens": 1
69
+ },
70
+ "config": {
71
+ "block_sizes": [
72
+ 4
73
+ ],
74
+ "loop_orders": [
75
+ [
76
+ 0,
77
+ 1,
78
+ 2
79
+ ]
80
+ ],
81
+ "l2_groupings": [
82
+ 2
83
+ ],
84
+ "range_unroll_factors": [
85
+ 0
86
+ ],
87
+ "range_warp_specializes": [
88
+ null
89
+ ],
90
+ "range_num_stages": [],
91
+ "range_multi_buffers": [
92
+ null
93
+ ],
94
+ "range_flattens": [
95
+ null
96
+ ],
97
+ "load_eviction_policies": [
98
+ "",
99
+ "last",
100
+ "",
101
+ "last",
102
+ "last",
103
+ "first",
104
+ "first",
105
+ ""
106
+ ],
107
+ "num_warps": 4,
108
+ "num_stages": 3,
109
+ "indexing": [
110
+ "tensor_descriptor",
111
+ "tensor_descriptor",
112
+ "pointer",
113
+ "pointer",
114
+ "tensor_descriptor",
115
+ "tensor_descriptor",
116
+ "pointer",
117
+ "pointer",
118
+ "pointer",
119
+ "pointer",
120
+ "pointer"
121
+ ],
122
+ "atomic_indexing": [],
123
+ "pid_type": "flat"
124
+ }
125
+ },
126
+ {
127
+ "key": {
128
+ "q_heads": 64,
129
+ "kv_heads": 8,
130
+ "num_tokens": 1
131
+ },
132
+ "config": {
133
+ "block_sizes": [
134
+ 4
135
+ ],
136
+ "loop_orders": [
137
+ [
138
+ 2,
139
+ 1,
140
+ 0
141
+ ]
142
+ ],
143
+ "l2_groupings": [
144
+ 32
145
+ ],
146
+ "range_unroll_factors": [
147
+ 0
148
+ ],
149
+ "range_warp_specializes": [
150
+ null
151
+ ],
152
+ "range_num_stages": [],
153
+ "range_multi_buffers": [
154
+ null
155
+ ],
156
+ "range_flattens": [
157
+ null
158
+ ],
159
+ "load_eviction_policies": [
160
+ "last",
161
+ "first",
162
+ "",
163
+ "first",
164
+ "last",
165
+ "first",
166
+ "first",
167
+ "first"
168
+ ],
169
+ "num_warps": 1,
170
+ "num_stages": 5,
171
+ "indexing": [
172
+ "pointer",
173
+ "pointer",
174
+ "pointer",
175
+ "pointer",
176
+ "pointer",
177
+ "tensor_descriptor",
178
+ "pointer",
179
+ "pointer",
180
+ "tensor_descriptor",
181
+ "tensor_descriptor",
182
+ "tensor_descriptor"
183
+ ],
184
+ "atomic_indexing": [],
185
+ "pid_type": "flat"
186
+ }
187
+ },
188
+ {
189
+ "key": {
190
+ "q_heads": 16,
191
+ "kv_heads": 8,
192
+ "num_tokens": 2
193
+ },
194
+ "config": {
195
+ "block_sizes": [
196
+ 1
197
+ ],
198
+ "loop_orders": [
199
+ [
200
+ 1,
201
+ 2,
202
+ 0
203
+ ]
204
+ ],
205
+ "l2_groupings": [
206
+ 64
207
+ ],
208
+ "range_unroll_factors": [
209
+ 1
210
+ ],
211
+ "range_warp_specializes": [
212
+ false
213
+ ],
214
+ "range_multi_buffers": [
215
+ true
216
+ ],
217
+ "range_flattens": [
218
+ true
219
+ ],
220
+ "load_eviction_policies": [
221
+ "last",
222
+ "first",
223
+ "",
224
+ "last",
225
+ "first",
226
+ "last",
227
+ "",
228
+ "last"
229
+ ],
230
+ "num_warps": 1,
231
+ "num_stages": 1,
232
+ "indexing": [
233
+ "pointer",
234
+ "tensor_descriptor",
235
+ "tensor_descriptor",
236
+ "tensor_descriptor",
237
+ "tensor_descriptor",
238
+ "tensor_descriptor",
239
+ "pointer",
240
+ "tensor_descriptor",
241
+ "pointer",
242
+ "tensor_descriptor",
243
+ "tensor_descriptor"
244
+ ],
245
+ "atomic_indexing": [],
246
+ "pid_type": "persistent_interleaved",
247
+ "num_sm_multiplier": 4,
248
+ "maxnreg": 128
249
+ }
250
+ },
251
+ {
252
+ "key": {
253
+ "q_heads": 32,
254
+ "kv_heads": 8,
255
+ "num_tokens": 2
256
+ },
257
+ "config": {
258
+ "block_sizes": [
259
+ 4
260
+ ],
261
+ "loop_orders": [
262
+ [
263
+ 2,
264
+ 1,
265
+ 0
266
+ ]
267
+ ],
268
+ "l2_groupings": [
269
+ 32
270
+ ],
271
+ "range_unroll_factors": [
272
+ 0
273
+ ],
274
+ "range_warp_specializes": [
275
+ null
276
+ ],
277
+ "range_num_stages": [],
278
+ "range_multi_buffers": [
279
+ null
280
+ ],
281
+ "range_flattens": [
282
+ null
283
+ ],
284
+ "load_eviction_policies": [
285
+ "last",
286
+ "first",
287
+ "",
288
+ "first",
289
+ "last",
290
+ "first",
291
+ "first",
292
+ "first"
293
+ ],
294
+ "num_warps": 1,
295
+ "num_stages": 5,
296
+ "indexing": [
297
+ "pointer",
298
+ "pointer",
299
+ "pointer",
300
+ "pointer",
301
+ "pointer",
302
+ "tensor_descriptor",
303
+ "pointer",
304
+ "pointer",
305
+ "tensor_descriptor",
306
+ "tensor_descriptor",
307
+ "tensor_descriptor"
308
+ ],
309
+ "atomic_indexing": [],
310
+ "pid_type": "flat"
311
+ }
312
+ },
313
+ {
314
+ "key": {
315
+ "q_heads": 64,
316
+ "kv_heads": 8,
317
+ "num_tokens": 2
318
+ },
319
+ "config": {
320
+ "block_sizes": [
321
+ 4
322
+ ],
323
+ "loop_orders": [
324
+ [
325
+ 2,
326
+ 1,
327
+ 0
328
+ ]
329
+ ],
330
+ "l2_groupings": [
331
+ 32
332
+ ],
333
+ "range_unroll_factors": [
334
+ 0
335
+ ],
336
+ "range_warp_specializes": [
337
+ null
338
+ ],
339
+ "range_num_stages": [],
340
+ "range_multi_buffers": [
341
+ null
342
+ ],
343
+ "range_flattens": [
344
+ null
345
+ ],
346
+ "load_eviction_policies": [
347
+ "last",
348
+ "first",
349
+ "",
350
+ "first",
351
+ "last",
352
+ "first",
353
+ "first",
354
+ "first"
355
+ ],
356
+ "num_warps": 1,
357
+ "num_stages": 5,
358
+ "indexing": [
359
+ "pointer",
360
+ "pointer",
361
+ "pointer",
362
+ "pointer",
363
+ "pointer",
364
+ "tensor_descriptor",
365
+ "pointer",
366
+ "pointer",
367
+ "tensor_descriptor",
368
+ "tensor_descriptor",
369
+ "tensor_descriptor"
370
+ ],
371
+ "atomic_indexing": [],
372
+ "pid_type": "flat"
373
+ }
374
+ },
375
+ {
376
+ "key": {
377
+ "q_heads": 16,
378
+ "kv_heads": 8,
379
+ "num_tokens": 4
380
+ },
381
+ "config": {
382
+ "block_sizes": [
383
+ 8
384
+ ],
385
+ "loop_orders": [
386
+ [
387
+ 2,
388
+ 1,
389
+ 0
390
+ ]
391
+ ],
392
+ "l2_groupings": [
393
+ 64
394
+ ],
395
+ "range_unroll_factors": [
396
+ 0
397
+ ],
398
+ "range_warp_specializes": [
399
+ false
400
+ ],
401
+ "range_multi_buffers": [
402
+ true
403
+ ],
404
+ "range_flattens": [
405
+ false
406
+ ],
407
+ "load_eviction_policies": [
408
+ "last",
409
+ "first",
410
+ "last",
411
+ "last",
412
+ "first",
413
+ "last",
414
+ "",
415
+ "last"
416
+ ],
417
+ "num_warps": 8,
418
+ "num_stages": 2,
419
+ "indexing": [
420
+ "tensor_descriptor",
421
+ "pointer",
422
+ "tensor_descriptor",
423
+ "tensor_descriptor",
424
+ "tensor_descriptor",
425
+ "tensor_descriptor",
426
+ "pointer",
427
+ "tensor_descriptor",
428
+ "tensor_descriptor",
429
+ "pointer",
430
+ "tensor_descriptor"
431
+ ],
432
+ "atomic_indexing": [],
433
+ "pid_type": "persistent_blocked",
434
+ "num_sm_multiplier": 4,
435
+ "maxnreg": 128
436
+ }
437
+ },
438
+ {
439
+ "key": {
440
+ "q_heads": 32,
441
+ "kv_heads": 8,
442
+ "num_tokens": 4
443
+ },
444
+ "config": {
445
+ "block_sizes": [
446
+ 4
447
+ ],
448
+ "loop_orders": [
449
+ [
450
+ 0,
451
+ 1,
452
+ 2
453
+ ]
454
+ ],
455
+ "l2_groupings": [
456
+ 2
457
+ ],
458
+ "range_unroll_factors": [
459
+ 0
460
+ ],
461
+ "range_warp_specializes": [
462
+ null
463
+ ],
464
+ "range_num_stages": [],
465
+ "range_multi_buffers": [
466
+ null
467
+ ],
468
+ "range_flattens": [
469
+ null
470
+ ],
471
+ "load_eviction_policies": [
472
+ "",
473
+ "last",
474
+ "last",
475
+ "first",
476
+ "first",
477
+ "first",
478
+ "",
479
+ "first"
480
+ ],
481
+ "num_warps": 1,
482
+ "num_stages": 3,
483
+ "indexing": [
484
+ "pointer",
485
+ "tensor_descriptor",
486
+ "pointer",
487
+ "pointer",
488
+ "tensor_descriptor",
489
+ "pointer",
490
+ "tensor_descriptor",
491
+ "tensor_descriptor",
492
+ "tensor_descriptor",
493
+ "pointer",
494
+ "tensor_descriptor"
495
+ ],
496
+ "atomic_indexing": [],
497
+ "pid_type": "flat"
498
+ }
499
+ },
500
+ {
501
+ "key": {
502
+ "q_heads": 64,
503
+ "kv_heads": 8,
504
+ "num_tokens": 4
505
+ },
506
+ "config": {
507
+ "block_sizes": [
508
+ 4
509
+ ],
510
+ "loop_orders": [
511
+ [
512
+ 2,
513
+ 1,
514
+ 0
515
+ ]
516
+ ],
517
+ "l2_groupings": [
518
+ 32
519
+ ],
520
+ "range_unroll_factors": [
521
+ 0
522
+ ],
523
+ "range_warp_specializes": [
524
+ null
525
+ ],
526
+ "range_num_stages": [],
527
+ "range_multi_buffers": [
528
+ null
529
+ ],
530
+ "range_flattens": [
531
+ null
532
+ ],
533
+ "load_eviction_policies": [
534
+ "last",
535
+ "first",
536
+ "",
537
+ "first",
538
+ "last",
539
+ "first",
540
+ "first",
541
+ "first"
542
+ ],
543
+ "num_warps": 1,
544
+ "num_stages": 5,
545
+ "indexing": [
546
+ "pointer",
547
+ "pointer",
548
+ "pointer",
549
+ "pointer",
550
+ "pointer",
551
+ "tensor_descriptor",
552
+ "pointer",
553
+ "pointer",
554
+ "tensor_descriptor",
555
+ "tensor_descriptor",
556
+ "tensor_descriptor"
557
+ ],
558
+ "atomic_indexing": [],
559
+ "pid_type": "flat"
560
+ }
561
+ },
562
+ {
563
+ "key": {
564
+ "q_heads": 16,
565
+ "kv_heads": 8,
566
+ "num_tokens": 8
567
+ },
568
+ "config": {
569
+ "block_sizes": [
570
+ 4
571
+ ],
572
+ "loop_orders": [
573
+ [
574
+ 0,
575
+ 2,
576
+ 1
577
+ ]
578
+ ],
579
+ "l2_groupings": [
580
+ 1
581
+ ],
582
+ "range_unroll_factors": [
583
+ 0
584
+ ],
585
+ "range_warp_specializes": [
586
+ null
587
+ ],
588
+ "range_num_stages": [],
589
+ "range_multi_buffers": [
590
+ null
591
+ ],
592
+ "range_flattens": [
593
+ null
594
+ ],
595
+ "load_eviction_policies": [
596
+ "first",
597
+ "",
598
+ "",
599
+ "",
600
+ "last",
601
+ "last",
602
+ "first",
603
+ "first"
604
+ ],
605
+ "num_warps": 4,
606
+ "num_stages": 3,
607
+ "indexing": [
608
+ "pointer",
609
+ "pointer",
610
+ "pointer",
611
+ "tensor_descriptor",
612
+ "tensor_descriptor",
613
+ "tensor_descriptor",
614
+ "tensor_descriptor",
615
+ "tensor_descriptor",
616
+ "tensor_descriptor",
617
+ "tensor_descriptor",
618
+ "pointer"
619
+ ],
620
+ "atomic_indexing": [],
621
+ "pid_type": "flat"
622
+ }
623
+ },
624
+ {
625
+ "key": {
626
+ "q_heads": 32,
627
+ "kv_heads": 8,
628
+ "num_tokens": 8
629
+ },
630
+ "config": {
631
+ "block_sizes": [
632
+ 4
633
+ ],
634
+ "loop_orders": [
635
+ [
636
+ 2,
637
+ 1,
638
+ 0
639
+ ]
640
+ ],
641
+ "l2_groupings": [
642
+ 32
643
+ ],
644
+ "range_unroll_factors": [
645
+ 0
646
+ ],
647
+ "range_warp_specializes": [
648
+ null
649
+ ],
650
+ "range_num_stages": [],
651
+ "range_multi_buffers": [
652
+ null
653
+ ],
654
+ "range_flattens": [
655
+ null
656
+ ],
657
+ "load_eviction_policies": [
658
+ "last",
659
+ "first",
660
+ "",
661
+ "first",
662
+ "last",
663
+ "first",
664
+ "first",
665
+ "first"
666
+ ],
667
+ "num_warps": 1,
668
+ "num_stages": 5,
669
+ "indexing": [
670
+ "pointer",
671
+ "pointer",
672
+ "pointer",
673
+ "pointer",
674
+ "pointer",
675
+ "tensor_descriptor",
676
+ "pointer",
677
+ "pointer",
678
+ "tensor_descriptor",
679
+ "tensor_descriptor",
680
+ "tensor_descriptor"
681
+ ],
682
+ "atomic_indexing": [],
683
+ "pid_type": "flat"
684
+ }
685
+ },
686
+ {
687
+ "key": {
688
+ "q_heads": 64,
689
+ "kv_heads": 8,
690
+ "num_tokens": 8
691
+ },
692
+ "config": {
693
+ "block_sizes": [
694
+ 8
695
+ ],
696
+ "loop_orders": [
697
+ [
698
+ 2,
699
+ 1,
700
+ 0
701
+ ]
702
+ ],
703
+ "l2_groupings": [
704
+ 1
705
+ ],
706
+ "range_unroll_factors": [
707
+ 0
708
+ ],
709
+ "range_warp_specializes": [
710
+ null
711
+ ],
712
+ "range_num_stages": [],
713
+ "range_multi_buffers": [
714
+ null
715
+ ],
716
+ "range_flattens": [
717
+ null
718
+ ],
719
+ "load_eviction_policies": [
720
+ "last",
721
+ "first",
722
+ "first",
723
+ "first",
724
+ "last",
725
+ "last",
726
+ "last",
727
+ "last"
728
+ ],
729
+ "num_warps": 8,
730
+ "num_stages": 3,
731
+ "indexing": [
732
+ "pointer",
733
+ "pointer",
734
+ "pointer",
735
+ "pointer",
736
+ "tensor_descriptor",
737
+ "pointer",
738
+ "tensor_descriptor",
739
+ "tensor_descriptor",
740
+ "pointer",
741
+ "tensor_descriptor",
742
+ "pointer"
743
+ ],
744
+ "atomic_indexing": [],
745
+ "pid_type": "flat"
746
+ }
747
+ },
748
+ {
749
+ "key": {
750
+ "q_heads": 16,
751
+ "kv_heads": 8,
752
+ "num_tokens": 16
753
+ },
754
+ "config": {
755
+ "block_sizes": [
756
+ 4
757
+ ],
758
+ "loop_orders": [
759
+ [
760
+ 1,
761
+ 2,
762
+ 0
763
+ ]
764
+ ],
765
+ "l2_groupings": [
766
+ 4
767
+ ],
768
+ "range_unroll_factors": [
769
+ 0
770
+ ],
771
+ "range_warp_specializes": [
772
+ null
773
+ ],
774
+ "range_num_stages": [],
775
+ "range_multi_buffers": [
776
+ null
777
+ ],
778
+ "range_flattens": [
779
+ null
780
+ ],
781
+ "load_eviction_policies": [
782
+ "",
783
+ "last",
784
+ "first",
785
+ "",
786
+ "first",
787
+ "",
788
+ "",
789
+ "first"
790
+ ],
791
+ "num_warps": 2,
792
+ "num_stages": 8,
793
+ "indexing": [
794
+ "tensor_descriptor",
795
+ "pointer",
796
+ "pointer",
797
+ "pointer",
798
+ "pointer",
799
+ "pointer",
800
+ "tensor_descriptor",
801
+ "pointer",
802
+ "tensor_descriptor",
803
+ "tensor_descriptor",
804
+ "pointer"
805
+ ],
806
+ "atomic_indexing": [],
807
+ "pid_type": "flat"
808
+ }
809
+ },
810
+ {
811
+ "key": {
812
+ "q_heads": 32,
813
+ "kv_heads": 8,
814
+ "num_tokens": 16
815
+ },
816
+ "config": {
817
+ "block_sizes": [
818
+ 4
819
+ ],
820
+ "loop_orders": [
821
+ [
822
+ 2,
823
+ 1,
824
+ 0
825
+ ]
826
+ ],
827
+ "l2_groupings": [
828
+ 1
829
+ ],
830
+ "range_unroll_factors": [
831
+ 0
832
+ ],
833
+ "range_warp_specializes": [
834
+ null
835
+ ],
836
+ "range_num_stages": [],
837
+ "range_multi_buffers": [
838
+ null
839
+ ],
840
+ "range_flattens": [
841
+ null
842
+ ],
843
+ "load_eviction_policies": [
844
+ "first",
845
+ "last",
846
+ "first",
847
+ "",
848
+ "first",
849
+ "first",
850
+ "",
851
+ "first"
852
+ ],
853
+ "num_warps": 2,
854
+ "num_stages": 4,
855
+ "indexing": [
856
+ "pointer",
857
+ "tensor_descriptor",
858
+ "pointer",
859
+ "pointer",
860
+ "pointer",
861
+ "pointer",
862
+ "tensor_descriptor",
863
+ "tensor_descriptor",
864
+ "tensor_descriptor",
865
+ "pointer",
866
+ "tensor_descriptor"
867
+ ],
868
+ "atomic_indexing": [],
869
+ "pid_type": "flat"
870
+ }
871
+ },
872
+ {
873
+ "key": {
874
+ "q_heads": 64,
875
+ "kv_heads": 8,
876
+ "num_tokens": 16
877
+ },
878
+ "config": {
879
+ "block_sizes": [
880
+ 4
881
+ ],
882
+ "loop_orders": [
883
+ [
884
+ 0,
885
+ 2,
886
+ 1
887
+ ]
888
+ ],
889
+ "l2_groupings": [
890
+ 1
891
+ ],
892
+ "range_unroll_factors": [
893
+ 0
894
+ ],
895
+ "range_warp_specializes": [
896
+ null
897
+ ],
898
+ "range_num_stages": [],
899
+ "range_multi_buffers": [
900
+ null
901
+ ],
902
+ "range_flattens": [
903
+ null
904
+ ],
905
+ "load_eviction_policies": [
906
+ "first",
907
+ "",
908
+ "",
909
+ "",
910
+ "last",
911
+ "last",
912
+ "first",
913
+ "first"
914
+ ],
915
+ "num_warps": 4,
916
+ "num_stages": 3,
917
+ "indexing": [
918
+ "pointer",
919
+ "pointer",
920
+ "pointer",
921
+ "tensor_descriptor",
922
+ "tensor_descriptor",
923
+ "tensor_descriptor",
924
+ "tensor_descriptor",
925
+ "tensor_descriptor",
926
+ "tensor_descriptor",
927
+ "tensor_descriptor",
928
+ "pointer"
929
+ ],
930
+ "atomic_indexing": [],
931
+ "pid_type": "flat"
932
+ }
933
+ },
934
+ {
935
+ "key": {
936
+ "q_heads": 16,
937
+ "kv_heads": 8,
938
+ "num_tokens": 32
939
+ },
940
+ "config": {
941
+ "block_sizes": [
942
+ 4
943
+ ],
944
+ "loop_orders": [
945
+ [
946
+ 1,
947
+ 2,
948
+ 0
949
+ ]
950
+ ],
951
+ "l2_groupings": [
952
+ 4
953
+ ],
954
+ "range_unroll_factors": [
955
+ 0
956
+ ],
957
+ "range_warp_specializes": [
958
+ null
959
+ ],
960
+ "range_num_stages": [],
961
+ "range_multi_buffers": [
962
+ null
963
+ ],
964
+ "range_flattens": [
965
+ null
966
+ ],
967
+ "load_eviction_policies": [
968
+ "",
969
+ "last",
970
+ "first",
971
+ "",
972
+ "first",
973
+ "",
974
+ "",
975
+ "first"
976
+ ],
977
+ "num_warps": 2,
978
+ "num_stages": 8,
979
+ "indexing": [
980
+ "tensor_descriptor",
981
+ "pointer",
982
+ "pointer",
983
+ "pointer",
984
+ "pointer",
985
+ "pointer",
986
+ "tensor_descriptor",
987
+ "pointer",
988
+ "tensor_descriptor",
989
+ "tensor_descriptor",
990
+ "pointer"
991
+ ],
992
+ "atomic_indexing": [],
993
+ "pid_type": "flat"
994
+ }
995
+ },
996
+ {
997
+ "key": {
998
+ "q_heads": 32,
999
+ "kv_heads": 8,
1000
+ "num_tokens": 32
1001
+ },
1002
+ "config": {
1003
+ "block_sizes": [
1004
+ 4
1005
+ ],
1006
+ "loop_orders": [
1007
+ [
1008
+ 2,
1009
+ 1,
1010
+ 0
1011
+ ]
1012
+ ],
1013
+ "l2_groupings": [
1014
+ 32
1015
+ ],
1016
+ "range_unroll_factors": [
1017
+ 0
1018
+ ],
1019
+ "range_warp_specializes": [
1020
+ null
1021
+ ],
1022
+ "range_num_stages": [],
1023
+ "range_multi_buffers": [
1024
+ null
1025
+ ],
1026
+ "range_flattens": [
1027
+ null
1028
+ ],
1029
+ "load_eviction_policies": [
1030
+ "first",
1031
+ "last",
1032
+ "last",
1033
+ "",
1034
+ "last",
1035
+ "last",
1036
+ "",
1037
+ "first"
1038
+ ],
1039
+ "num_warps": 1,
1040
+ "num_stages": 2,
1041
+ "indexing": [
1042
+ "pointer",
1043
+ "tensor_descriptor",
1044
+ "pointer",
1045
+ "pointer",
1046
+ "pointer",
1047
+ "tensor_descriptor",
1048
+ "pointer",
1049
+ "pointer",
1050
+ "tensor_descriptor",
1051
+ "tensor_descriptor",
1052
+ "tensor_descriptor"
1053
+ ],
1054
+ "atomic_indexing": [],
1055
+ "pid_type": "flat"
1056
+ }
1057
+ },
1058
+ {
1059
+ "key": {
1060
+ "q_heads": 64,
1061
+ "kv_heads": 8,
1062
+ "num_tokens": 32
1063
+ },
1064
+ "config": {
1065
+ "block_sizes": [
1066
+ 4
1067
+ ],
1068
+ "loop_orders": [
1069
+ [
1070
+ 1,
1071
+ 2,
1072
+ 0
1073
+ ]
1074
+ ],
1075
+ "l2_groupings": [
1076
+ 4
1077
+ ],
1078
+ "range_unroll_factors": [
1079
+ 0
1080
+ ],
1081
+ "range_warp_specializes": [
1082
+ null
1083
+ ],
1084
+ "range_num_stages": [],
1085
+ "range_multi_buffers": [
1086
+ null
1087
+ ],
1088
+ "range_flattens": [
1089
+ null
1090
+ ],
1091
+ "load_eviction_policies": [
1092
+ "",
1093
+ "last",
1094
+ "first",
1095
+ "",
1096
+ "first",
1097
+ "",
1098
+ "",
1099
+ "first"
1100
+ ],
1101
+ "num_warps": 2,
1102
+ "num_stages": 8,
1103
+ "indexing": [
1104
+ "tensor_descriptor",
1105
+ "pointer",
1106
+ "pointer",
1107
+ "pointer",
1108
+ "pointer",
1109
+ "pointer",
1110
+ "tensor_descriptor",
1111
+ "pointer",
1112
+ "tensor_descriptor",
1113
+ "tensor_descriptor",
1114
+ "pointer"
1115
+ ],
1116
+ "atomic_indexing": [],
1117
+ "pid_type": "flat"
1118
+ }
1119
+ },
1120
+ {
1121
+ "key": {
1122
+ "q_heads": 16,
1123
+ "kv_heads": 8,
1124
+ "num_tokens": 64
1125
+ },
1126
+ "config": {
1127
+ "block_sizes": [
1128
+ 4
1129
+ ],
1130
+ "loop_orders": [
1131
+ [
1132
+ 1,
1133
+ 2,
1134
+ 0
1135
+ ]
1136
+ ],
1137
+ "l2_groupings": [
1138
+ 4
1139
+ ],
1140
+ "range_unroll_factors": [
1141
+ 0
1142
+ ],
1143
+ "range_warp_specializes": [
1144
+ null
1145
+ ],
1146
+ "range_num_stages": [],
1147
+ "range_multi_buffers": [
1148
+ null
1149
+ ],
1150
+ "range_flattens": [
1151
+ null
1152
+ ],
1153
+ "load_eviction_policies": [
1154
+ "",
1155
+ "last",
1156
+ "first",
1157
+ "",
1158
+ "first",
1159
+ "",
1160
+ "",
1161
+ "first"
1162
+ ],
1163
+ "num_warps": 2,
1164
+ "num_stages": 8,
1165
+ "indexing": [
1166
+ "tensor_descriptor",
1167
+ "pointer",
1168
+ "pointer",
1169
+ "pointer",
1170
+ "pointer",
1171
+ "pointer",
1172
+ "tensor_descriptor",
1173
+ "pointer",
1174
+ "tensor_descriptor",
1175
+ "tensor_descriptor",
1176
+ "pointer"
1177
+ ],
1178
+ "atomic_indexing": [],
1179
+ "pid_type": "flat"
1180
+ }
1181
+ },
1182
+ {
1183
+ "key": {
1184
+ "q_heads": 32,
1185
+ "kv_heads": 8,
1186
+ "num_tokens": 64
1187
+ },
1188
+ "config": {
1189
+ "block_sizes": [
1190
+ 4
1191
+ ],
1192
+ "loop_orders": [
1193
+ [
1194
+ 2,
1195
+ 1,
1196
+ 0
1197
+ ]
1198
+ ],
1199
+ "l2_groupings": [
1200
+ 32
1201
+ ],
1202
+ "range_unroll_factors": [
1203
+ 0
1204
+ ],
1205
+ "range_warp_specializes": [
1206
+ null
1207
+ ],
1208
+ "range_num_stages": [],
1209
+ "range_multi_buffers": [
1210
+ null
1211
+ ],
1212
+ "range_flattens": [
1213
+ null
1214
+ ],
1215
+ "load_eviction_policies": [
1216
+ "last",
1217
+ "first",
1218
+ "",
1219
+ "first",
1220
+ "last",
1221
+ "first",
1222
+ "first",
1223
+ "first"
1224
+ ],
1225
+ "num_warps": 1,
1226
+ "num_stages": 5,
1227
+ "indexing": [
1228
+ "pointer",
1229
+ "pointer",
1230
+ "pointer",
1231
+ "pointer",
1232
+ "pointer",
1233
+ "tensor_descriptor",
1234
+ "pointer",
1235
+ "pointer",
1236
+ "tensor_descriptor",
1237
+ "tensor_descriptor",
1238
+ "tensor_descriptor"
1239
+ ],
1240
+ "atomic_indexing": [],
1241
+ "pid_type": "flat"
1242
+ }
1243
+ },
1244
+ {
1245
+ "key": {
1246
+ "q_heads": 64,
1247
+ "kv_heads": 8,
1248
+ "num_tokens": 64
1249
+ },
1250
+ "config": {
1251
+ "block_sizes": [
1252
+ 4
1253
+ ],
1254
+ "loop_orders": [
1255
+ [
1256
+ 0,
1257
+ 2,
1258
+ 1
1259
+ ]
1260
+ ],
1261
+ "l2_groupings": [
1262
+ 1
1263
+ ],
1264
+ "range_unroll_factors": [
1265
+ 0
1266
+ ],
1267
+ "range_warp_specializes": [
1268
+ null
1269
+ ],
1270
+ "range_num_stages": [],
1271
+ "range_multi_buffers": [
1272
+ null
1273
+ ],
1274
+ "range_flattens": [
1275
+ null
1276
+ ],
1277
+ "load_eviction_policies": [
1278
+ "first",
1279
+ "last",
1280
+ "first",
1281
+ "",
1282
+ "last",
1283
+ "first",
1284
+ "first",
1285
+ "first"
1286
+ ],
1287
+ "num_warps": 2,
1288
+ "num_stages": 7,
1289
+ "indexing": [
1290
+ "tensor_descriptor",
1291
+ "pointer",
1292
+ "pointer",
1293
+ "tensor_descriptor",
1294
+ "tensor_descriptor",
1295
+ "pointer",
1296
+ "pointer",
1297
+ "tensor_descriptor",
1298
+ "tensor_descriptor",
1299
+ "tensor_descriptor",
1300
+ "tensor_descriptor"
1301
+ ],
1302
+ "atomic_indexing": [],
1303
+ "pid_type": "flat"
1304
+ }
1305
+ },
1306
+ {
1307
+ "key": {
1308
+ "q_heads": 16,
1309
+ "kv_heads": 8,
1310
+ "num_tokens": 128
1311
+ },
1312
+ "config": {
1313
+ "block_sizes": [
1314
+ 4
1315
+ ],
1316
+ "loop_orders": [
1317
+ [
1318
+ 1,
1319
+ 2,
1320
+ 0
1321
+ ]
1322
+ ],
1323
+ "l2_groupings": [
1324
+ 4
1325
+ ],
1326
+ "range_unroll_factors": [
1327
+ 0
1328
+ ],
1329
+ "range_warp_specializes": [
1330
+ null
1331
+ ],
1332
+ "range_num_stages": [],
1333
+ "range_multi_buffers": [
1334
+ null
1335
+ ],
1336
+ "range_flattens": [
1337
+ null
1338
+ ],
1339
+ "load_eviction_policies": [
1340
+ "",
1341
+ "last",
1342
+ "first",
1343
+ "",
1344
+ "first",
1345
+ "",
1346
+ "",
1347
+ "first"
1348
+ ],
1349
+ "num_warps": 2,
1350
+ "num_stages": 8,
1351
+ "indexing": [
1352
+ "tensor_descriptor",
1353
+ "pointer",
1354
+ "pointer",
1355
+ "pointer",
1356
+ "pointer",
1357
+ "pointer",
1358
+ "tensor_descriptor",
1359
+ "pointer",
1360
+ "tensor_descriptor",
1361
+ "tensor_descriptor",
1362
+ "pointer"
1363
+ ],
1364
+ "atomic_indexing": [],
1365
+ "pid_type": "flat"
1366
+ }
1367
+ },
1368
+ {
1369
+ "key": {
1370
+ "q_heads": 32,
1371
+ "kv_heads": 8,
1372
+ "num_tokens": 128
1373
+ },
1374
+ "config": {
1375
+ "block_sizes": [
1376
+ 4
1377
+ ],
1378
+ "loop_orders": [
1379
+ [
1380
+ 2,
1381
+ 1,
1382
+ 0
1383
+ ]
1384
+ ],
1385
+ "l2_groupings": [
1386
+ 8
1387
+ ],
1388
+ "range_unroll_factors": [
1389
+ 0
1390
+ ],
1391
+ "range_warp_specializes": [
1392
+ null
1393
+ ],
1394
+ "range_num_stages": [],
1395
+ "range_multi_buffers": [
1396
+ null
1397
+ ],
1398
+ "range_flattens": [
1399
+ null
1400
+ ],
1401
+ "load_eviction_policies": [
1402
+ "",
1403
+ "last",
1404
+ "",
1405
+ "",
1406
+ "last",
1407
+ "last",
1408
+ "first",
1409
+ ""
1410
+ ],
1411
+ "num_warps": 1,
1412
+ "num_stages": 7,
1413
+ "indexing": [
1414
+ "tensor_descriptor",
1415
+ "tensor_descriptor",
1416
+ "pointer",
1417
+ "pointer",
1418
+ "pointer",
1419
+ "tensor_descriptor",
1420
+ "pointer",
1421
+ "pointer",
1422
+ "tensor_descriptor",
1423
+ "pointer",
1424
+ "tensor_descriptor"
1425
+ ],
1426
+ "atomic_indexing": [],
1427
+ "pid_type": "flat"
1428
+ }
1429
+ },
1430
+ {
1431
+ "key": {
1432
+ "q_heads": 64,
1433
+ "kv_heads": 8,
1434
+ "num_tokens": 128
1435
+ },
1436
+ "config": {
1437
+ "block_sizes": [
1438
+ 4
1439
+ ],
1440
+ "loop_orders": [
1441
+ [
1442
+ 2,
1443
+ 1,
1444
+ 0
1445
+ ]
1446
+ ],
1447
+ "l2_groupings": [
1448
+ 8
1449
+ ],
1450
+ "range_unroll_factors": [
1451
+ 0
1452
+ ],
1453
+ "range_warp_specializes": [
1454
+ null
1455
+ ],
1456
+ "range_num_stages": [],
1457
+ "range_multi_buffers": [
1458
+ null
1459
+ ],
1460
+ "range_flattens": [
1461
+ null
1462
+ ],
1463
+ "load_eviction_policies": [
1464
+ "",
1465
+ "first",
1466
+ "",
1467
+ "last",
1468
+ "first",
1469
+ "",
1470
+ "last",
1471
+ "first"
1472
+ ],
1473
+ "num_warps": 1,
1474
+ "num_stages": 5,
1475
+ "indexing": [
1476
+ "tensor_descriptor",
1477
+ "pointer",
1478
+ "pointer",
1479
+ "tensor_descriptor",
1480
+ "tensor_descriptor",
1481
+ "pointer",
1482
+ "tensor_descriptor",
1483
+ "pointer",
1484
+ "tensor_descriptor",
1485
+ "tensor_descriptor",
1486
+ "pointer"
1487
+ ],
1488
+ "atomic_indexing": [],
1489
+ "pid_type": "flat"
1490
+ }
1491
+ },
1492
+ {
1493
+ "key": {
1494
+ "q_heads": 16,
1495
+ "kv_heads": 8,
1496
+ "num_tokens": 256
1497
+ },
1498
+ "config": {
1499
+ "block_sizes": [
1500
+ 4
1501
+ ],
1502
+ "loop_orders": [
1503
+ [
1504
+ 1,
1505
+ 2,
1506
+ 0
1507
+ ]
1508
+ ],
1509
+ "l2_groupings": [
1510
+ 4
1511
+ ],
1512
+ "range_unroll_factors": [
1513
+ 0
1514
+ ],
1515
+ "range_warp_specializes": [
1516
+ null
1517
+ ],
1518
+ "range_num_stages": [],
1519
+ "range_multi_buffers": [
1520
+ null
1521
+ ],
1522
+ "range_flattens": [
1523
+ null
1524
+ ],
1525
+ "load_eviction_policies": [
1526
+ "",
1527
+ "last",
1528
+ "first",
1529
+ "",
1530
+ "first",
1531
+ "",
1532
+ "",
1533
+ "first"
1534
+ ],
1535
+ "num_warps": 2,
1536
+ "num_stages": 8,
1537
+ "indexing": [
1538
+ "tensor_descriptor",
1539
+ "pointer",
1540
+ "pointer",
1541
+ "pointer",
1542
+ "pointer",
1543
+ "pointer",
1544
+ "tensor_descriptor",
1545
+ "pointer",
1546
+ "tensor_descriptor",
1547
+ "tensor_descriptor",
1548
+ "pointer"
1549
+ ],
1550
+ "atomic_indexing": [],
1551
+ "pid_type": "flat"
1552
+ }
1553
+ },
1554
+ {
1555
+ "key": {
1556
+ "q_heads": 32,
1557
+ "kv_heads": 8,
1558
+ "num_tokens": 256
1559
+ },
1560
+ "config": {
1561
+ "block_sizes": [
1562
+ 4
1563
+ ],
1564
+ "loop_orders": [
1565
+ [
1566
+ 2,
1567
+ 1,
1568
+ 0
1569
+ ]
1570
+ ],
1571
+ "l2_groupings": [
1572
+ 1
1573
+ ],
1574
+ "range_unroll_factors": [
1575
+ 0
1576
+ ],
1577
+ "range_warp_specializes": [
1578
+ null
1579
+ ],
1580
+ "range_num_stages": [],
1581
+ "range_multi_buffers": [
1582
+ null
1583
+ ],
1584
+ "range_flattens": [
1585
+ null
1586
+ ],
1587
+ "load_eviction_policies": [
1588
+ "first",
1589
+ "last",
1590
+ "first",
1591
+ "",
1592
+ "first",
1593
+ "first",
1594
+ "",
1595
+ "first"
1596
+ ],
1597
+ "num_warps": 2,
1598
+ "num_stages": 4,
1599
+ "indexing": [
1600
+ "pointer",
1601
+ "tensor_descriptor",
1602
+ "pointer",
1603
+ "pointer",
1604
+ "pointer",
1605
+ "pointer",
1606
+ "tensor_descriptor",
1607
+ "tensor_descriptor",
1608
+ "tensor_descriptor",
1609
+ "pointer",
1610
+ "tensor_descriptor"
1611
+ ],
1612
+ "atomic_indexing": [],
1613
+ "pid_type": "flat"
1614
+ }
1615
+ },
1616
+ {
1617
+ "key": {
1618
+ "q_heads": 64,
1619
+ "kv_heads": 8,
1620
+ "num_tokens": 256
1621
+ },
1622
+ "config": {
1623
+ "block_sizes": [
1624
+ 4
1625
+ ],
1626
+ "loop_orders": [
1627
+ [
1628
+ 2,
1629
+ 1,
1630
+ 0
1631
+ ]
1632
+ ],
1633
+ "l2_groupings": [
1634
+ 32
1635
+ ],
1636
+ "range_unroll_factors": [
1637
+ 0
1638
+ ],
1639
+ "range_warp_specializes": [
1640
+ null
1641
+ ],
1642
+ "range_num_stages": [],
1643
+ "range_multi_buffers": [
1644
+ null
1645
+ ],
1646
+ "range_flattens": [
1647
+ null
1648
+ ],
1649
+ "load_eviction_policies": [
1650
+ "first",
1651
+ "last",
1652
+ "last",
1653
+ "",
1654
+ "last",
1655
+ "last",
1656
+ "",
1657
+ "first"
1658
+ ],
1659
+ "num_warps": 1,
1660
+ "num_stages": 2,
1661
+ "indexing": [
1662
+ "pointer",
1663
+ "tensor_descriptor",
1664
+ "pointer",
1665
+ "pointer",
1666
+ "pointer",
1667
+ "tensor_descriptor",
1668
+ "pointer",
1669
+ "pointer",
1670
+ "tensor_descriptor",
1671
+ "tensor_descriptor",
1672
+ "tensor_descriptor"
1673
+ ],
1674
+ "atomic_indexing": [],
1675
+ "pid_type": "flat"
1676
+ }
1677
+ },
1678
+ {
1679
+ "key": {
1680
+ "q_heads": 16,
1681
+ "kv_heads": 8,
1682
+ "num_tokens": 512
1683
+ },
1684
+ "config": {
1685
+ "block_sizes": [
1686
+ 4
1687
+ ],
1688
+ "loop_orders": [
1689
+ [
1690
+ 1,
1691
+ 2,
1692
+ 0
1693
+ ]
1694
+ ],
1695
+ "l2_groupings": [
1696
+ 8
1697
+ ],
1698
+ "range_unroll_factors": [
1699
+ 0
1700
+ ],
1701
+ "range_warp_specializes": [],
1702
+ "range_num_stages": [],
1703
+ "range_multi_buffers": [
1704
+ null
1705
+ ],
1706
+ "range_flattens": [
1707
+ null
1708
+ ],
1709
+ "load_eviction_policies": [
1710
+ "",
1711
+ "first",
1712
+ "last",
1713
+ "",
1714
+ "last",
1715
+ "last",
1716
+ "first",
1717
+ "first"
1718
+ ],
1719
+ "num_warps": 1,
1720
+ "num_stages": 1,
1721
+ "indexing": [
1722
+ "pointer",
1723
+ "tensor_descriptor",
1724
+ "pointer",
1725
+ "tensor_descriptor",
1726
+ "pointer",
1727
+ "pointer",
1728
+ "pointer",
1729
+ "tensor_descriptor",
1730
+ "tensor_descriptor",
1731
+ "pointer",
1732
+ "pointer"
1733
+ ],
1734
+ "atomic_indexing": [],
1735
+ "pid_type": "flat"
1736
+ }
1737
+ },
1738
+ {
1739
+ "key": {
1740
+ "q_heads": 32,
1741
+ "kv_heads": 8,
1742
+ "num_tokens": 512
1743
+ },
1744
+ "config": {
1745
+ "block_sizes": [
1746
+ 2
1747
+ ],
1748
+ "loop_orders": [
1749
+ [
1750
+ 1,
1751
+ 2,
1752
+ 0
1753
+ ]
1754
+ ],
1755
+ "l2_groupings": [
1756
+ 64
1757
+ ],
1758
+ "range_unroll_factors": [
1759
+ 2
1760
+ ],
1761
+ "range_warp_specializes": [],
1762
+ "range_multi_buffers": [
1763
+ null
1764
+ ],
1765
+ "range_flattens": [
1766
+ null
1767
+ ],
1768
+ "load_eviction_policies": [
1769
+ "",
1770
+ "last",
1771
+ "",
1772
+ "last",
1773
+ "first",
1774
+ "first",
1775
+ "last",
1776
+ "first"
1777
+ ],
1778
+ "num_warps": 1,
1779
+ "num_stages": 5,
1780
+ "indexing": [
1781
+ "pointer",
1782
+ "tensor_descriptor",
1783
+ "pointer",
1784
+ "tensor_descriptor",
1785
+ "pointer",
1786
+ "tensor_descriptor",
1787
+ "tensor_descriptor",
1788
+ "tensor_descriptor",
1789
+ "tensor_descriptor",
1790
+ "tensor_descriptor",
1791
+ "pointer"
1792
+ ],
1793
+ "atomic_indexing": [],
1794
+ "pid_type": "persistent_blocked",
1795
+ "num_sm_multiplier": 32,
1796
+ "maxnreg": 64
1797
+ }
1798
+ },
1799
+ {
1800
+ "key": {
1801
+ "q_heads": 64,
1802
+ "kv_heads": 8,
1803
+ "num_tokens": 512
1804
+ },
1805
+ "config": {
1806
+ "block_sizes": [
1807
+ 2
1808
+ ],
1809
+ "loop_orders": [
1810
+ [
1811
+ 2,
1812
+ 1,
1813
+ 0
1814
+ ]
1815
+ ],
1816
+ "l2_groupings": [
1817
+ 64
1818
+ ],
1819
+ "range_unroll_factors": [
1820
+ 4
1821
+ ],
1822
+ "range_warp_specializes": [
1823
+ false
1824
+ ],
1825
+ "range_multi_buffers": [
1826
+ null
1827
+ ],
1828
+ "range_flattens": [
1829
+ false
1830
+ ],
1831
+ "load_eviction_policies": [
1832
+ "last",
1833
+ "",
1834
+ "first",
1835
+ "first",
1836
+ "",
1837
+ "first",
1838
+ "last",
1839
+ "last"
1840
+ ],
1841
+ "num_warps": 1,
1842
+ "num_stages": 1,
1843
+ "indexing": [
1844
+ "tensor_descriptor",
1845
+ "tensor_descriptor",
1846
+ "pointer",
1847
+ "tensor_descriptor",
1848
+ "tensor_descriptor",
1849
+ "pointer",
1850
+ "pointer",
1851
+ "tensor_descriptor",
1852
+ "tensor_descriptor",
1853
+ "tensor_descriptor",
1854
+ "tensor_descriptor"
1855
+ ],
1856
+ "atomic_indexing": [],
1857
+ "pid_type": "persistent_blocked",
1858
+ "num_sm_multiplier": 32
1859
+ }
1860
+ },
1861
+ {
1862
+ "key": {
1863
+ "q_heads": 16,
1864
+ "kv_heads": 8,
1865
+ "num_tokens": 1024
1866
+ },
1867
+ "config": {
1868
+ "block_sizes": [
1869
+ 2
1870
+ ],
1871
+ "loop_orders": [
1872
+ [
1873
+ 1,
1874
+ 2,
1875
+ 0
1876
+ ]
1877
+ ],
1878
+ "l2_groupings": [
1879
+ 64
1880
+ ],
1881
+ "range_unroll_factors": [
1882
+ 2
1883
+ ],
1884
+ "range_warp_specializes": [],
1885
+ "range_multi_buffers": [
1886
+ null
1887
+ ],
1888
+ "range_flattens": [
1889
+ null
1890
+ ],
1891
+ "load_eviction_policies": [
1892
+ "",
1893
+ "last",
1894
+ "",
1895
+ "last",
1896
+ "first",
1897
+ "first",
1898
+ "last",
1899
+ "first"
1900
+ ],
1901
+ "num_warps": 1,
1902
+ "num_stages": 5,
1903
+ "indexing": [
1904
+ "pointer",
1905
+ "tensor_descriptor",
1906
+ "pointer",
1907
+ "tensor_descriptor",
1908
+ "pointer",
1909
+ "tensor_descriptor",
1910
+ "tensor_descriptor",
1911
+ "tensor_descriptor",
1912
+ "tensor_descriptor",
1913
+ "tensor_descriptor",
1914
+ "pointer"
1915
+ ],
1916
+ "atomic_indexing": [],
1917
+ "pid_type": "persistent_blocked",
1918
+ "num_sm_multiplier": 32,
1919
+ "maxnreg": 64
1920
+ }
1921
+ },
1922
+ {
1923
+ "key": {
1924
+ "q_heads": 32,
1925
+ "kv_heads": 8,
1926
+ "num_tokens": 1024
1927
+ },
1928
+ "config": {
1929
+ "block_sizes": [
1930
+ 2
1931
+ ],
1932
+ "loop_orders": [
1933
+ [
1934
+ 1,
1935
+ 2,
1936
+ 0
1937
+ ]
1938
+ ],
1939
+ "l2_groupings": [
1940
+ 64
1941
+ ],
1942
+ "range_unroll_factors": [
1943
+ 2
1944
+ ],
1945
+ "range_warp_specializes": [
1946
+ false
1947
+ ],
1948
+ "range_multi_buffers": [
1949
+ true
1950
+ ],
1951
+ "range_flattens": [
1952
+ false
1953
+ ],
1954
+ "load_eviction_policies": [
1955
+ "",
1956
+ "last",
1957
+ "",
1958
+ "",
1959
+ "last",
1960
+ "last",
1961
+ "first",
1962
+ ""
1963
+ ],
1964
+ "num_warps": 1,
1965
+ "num_stages": 7,
1966
+ "indexing": [
1967
+ "tensor_descriptor",
1968
+ "tensor_descriptor",
1969
+ "pointer",
1970
+ "pointer",
1971
+ "tensor_descriptor",
1972
+ "tensor_descriptor",
1973
+ "pointer",
1974
+ "pointer",
1975
+ "pointer",
1976
+ "pointer",
1977
+ "tensor_descriptor"
1978
+ ],
1979
+ "atomic_indexing": [],
1980
+ "pid_type": "persistent_interleaved",
1981
+ "num_sm_multiplier": 32
1982
+ }
1983
+ },
1984
+ {
1985
+ "key": {
1986
+ "q_heads": 64,
1987
+ "kv_heads": 8,
1988
+ "num_tokens": 1024
1989
+ },
1990
+ "config": {
1991
+ "block_sizes": [
1992
+ 2
1993
+ ],
1994
+ "loop_orders": [
1995
+ [
1996
+ 2,
1997
+ 1,
1998
+ 0
1999
+ ]
2000
+ ],
2001
+ "l2_groupings": [
2002
+ 1
2003
+ ],
2004
+ "range_unroll_factors": [
2005
+ 0
2006
+ ],
2007
+ "range_warp_specializes": [
2008
+ false
2009
+ ],
2010
+ "range_multi_buffers": [
2011
+ null
2012
+ ],
2013
+ "range_flattens": [
2014
+ false
2015
+ ],
2016
+ "load_eviction_policies": [
2017
+ "",
2018
+ "",
2019
+ "",
2020
+ "",
2021
+ "",
2022
+ "",
2023
+ "last",
2024
+ ""
2025
+ ],
2026
+ "num_warps": 1,
2027
+ "num_stages": 6,
2028
+ "indexing": [
2029
+ "tensor_descriptor",
2030
+ "tensor_descriptor",
2031
+ "tensor_descriptor",
2032
+ "tensor_descriptor",
2033
+ "tensor_descriptor",
2034
+ "tensor_descriptor",
2035
+ "tensor_descriptor",
2036
+ "tensor_descriptor",
2037
+ "tensor_descriptor",
2038
+ "tensor_descriptor",
2039
+ "pointer"
2040
+ ],
2041
+ "atomic_indexing": [],
2042
+ "pid_type": "persistent_blocked",
2043
+ "num_sm_multiplier": 128,
2044
+ "maxnreg": 64
2045
+ }
2046
+ },
2047
+ {
2048
+ "key": {
2049
+ "q_heads": 16,
2050
+ "kv_heads": 8,
2051
+ "num_tokens": 2048
2052
+ },
2053
+ "config": {
2054
+ "block_sizes": [
2055
+ 2
2056
+ ],
2057
+ "loop_orders": [
2058
+ [
2059
+ 1,
2060
+ 2,
2061
+ 0
2062
+ ]
2063
+ ],
2064
+ "l2_groupings": [
2065
+ 64
2066
+ ],
2067
+ "range_unroll_factors": [
2068
+ 2
2069
+ ],
2070
+ "range_warp_specializes": [
2071
+ false
2072
+ ],
2073
+ "range_multi_buffers": [
2074
+ true
2075
+ ],
2076
+ "range_flattens": [
2077
+ false
2078
+ ],
2079
+ "load_eviction_policies": [
2080
+ "",
2081
+ "last",
2082
+ "",
2083
+ "",
2084
+ "last",
2085
+ "last",
2086
+ "first",
2087
+ ""
2088
+ ],
2089
+ "num_warps": 1,
2090
+ "num_stages": 7,
2091
+ "indexing": [
2092
+ "tensor_descriptor",
2093
+ "tensor_descriptor",
2094
+ "pointer",
2095
+ "pointer",
2096
+ "tensor_descriptor",
2097
+ "tensor_descriptor",
2098
+ "pointer",
2099
+ "pointer",
2100
+ "pointer",
2101
+ "pointer",
2102
+ "tensor_descriptor"
2103
+ ],
2104
+ "atomic_indexing": [],
2105
+ "pid_type": "persistent_interleaved",
2106
+ "num_sm_multiplier": 32
2107
+ }
2108
+ },
2109
+ {
2110
+ "key": {
2111
+ "q_heads": 32,
2112
+ "kv_heads": 8,
2113
+ "num_tokens": 2048
2114
+ },
2115
+ "config": {
2116
+ "block_sizes": [
2117
+ 2
2118
+ ],
2119
+ "loop_orders": [
2120
+ [
2121
+ 2,
2122
+ 1,
2123
+ 0
2124
+ ]
2125
+ ],
2126
+ "l2_groupings": [
2127
+ 64
2128
+ ],
2129
+ "range_unroll_factors": [
2130
+ 2
2131
+ ],
2132
+ "range_warp_specializes": [
2133
+ false
2134
+ ],
2135
+ "range_multi_buffers": [
2136
+ true
2137
+ ],
2138
+ "range_flattens": [
2139
+ false
2140
+ ],
2141
+ "load_eviction_policies": [
2142
+ "",
2143
+ "first",
2144
+ "last",
2145
+ "",
2146
+ "first",
2147
+ "",
2148
+ "first",
2149
+ ""
2150
+ ],
2151
+ "num_warps": 1,
2152
+ "num_stages": 8,
2153
+ "indexing": [
2154
+ "tensor_descriptor",
2155
+ "pointer",
2156
+ "pointer",
2157
+ "pointer",
2158
+ "pointer",
2159
+ "tensor_descriptor",
2160
+ "tensor_descriptor",
2161
+ "pointer",
2162
+ "pointer",
2163
+ "tensor_descriptor",
2164
+ "tensor_descriptor"
2165
+ ],
2166
+ "atomic_indexing": [],
2167
+ "pid_type": "persistent_interleaved",
2168
+ "num_sm_multiplier": 32
2169
+ }
2170
+ },
2171
+ {
2172
+ "key": {
2173
+ "q_heads": 64,
2174
+ "kv_heads": 8,
2175
+ "num_tokens": 2048
2176
+ },
2177
+ "config": {
2178
+ "block_sizes": [
2179
+ 2
2180
+ ],
2181
+ "loop_orders": [
2182
+ [
2183
+ 2,
2184
+ 1,
2185
+ 0
2186
+ ]
2187
+ ],
2188
+ "l2_groupings": [
2189
+ 1
2190
+ ],
2191
+ "range_unroll_factors": [
2192
+ 0
2193
+ ],
2194
+ "range_warp_specializes": [
2195
+ false
2196
+ ],
2197
+ "range_multi_buffers": [
2198
+ null
2199
+ ],
2200
+ "range_flattens": [
2201
+ false
2202
+ ],
2203
+ "load_eviction_policies": [
2204
+ "",
2205
+ "",
2206
+ "",
2207
+ "",
2208
+ "",
2209
+ "",
2210
+ "last",
2211
+ ""
2212
+ ],
2213
+ "num_warps": 1,
2214
+ "num_stages": 6,
2215
+ "indexing": [
2216
+ "tensor_descriptor",
2217
+ "tensor_descriptor",
2218
+ "tensor_descriptor",
2219
+ "tensor_descriptor",
2220
+ "tensor_descriptor",
2221
+ "tensor_descriptor",
2222
+ "tensor_descriptor",
2223
+ "tensor_descriptor",
2224
+ "tensor_descriptor",
2225
+ "tensor_descriptor",
2226
+ "pointer"
2227
+ ],
2228
+ "atomic_indexing": [],
2229
+ "pid_type": "persistent_blocked",
2230
+ "num_sm_multiplier": 128,
2231
+ "maxnreg": 64
2232
+ }
2233
+ },
2234
+ {
2235
+ "key": {
2236
+ "q_heads": 16,
2237
+ "kv_heads": 8,
2238
+ "num_tokens": 4096
2239
+ },
2240
+ "config": {
2241
+ "block_sizes": [
2242
+ 2
2243
+ ],
2244
+ "loop_orders": [
2245
+ [
2246
+ 2,
2247
+ 1,
2248
+ 0
2249
+ ]
2250
+ ],
2251
+ "l2_groupings": [
2252
+ 1
2253
+ ],
2254
+ "range_unroll_factors": [
2255
+ 0
2256
+ ],
2257
+ "range_warp_specializes": [
2258
+ false
2259
+ ],
2260
+ "range_multi_buffers": [
2261
+ null
2262
+ ],
2263
+ "range_flattens": [
2264
+ false
2265
+ ],
2266
+ "load_eviction_policies": [
2267
+ "",
2268
+ "",
2269
+ "",
2270
+ "",
2271
+ "",
2272
+ "",
2273
+ "last",
2274
+ ""
2275
+ ],
2276
+ "num_warps": 1,
2277
+ "num_stages": 6,
2278
+ "indexing": [
2279
+ "tensor_descriptor",
2280
+ "tensor_descriptor",
2281
+ "tensor_descriptor",
2282
+ "tensor_descriptor",
2283
+ "tensor_descriptor",
2284
+ "tensor_descriptor",
2285
+ "tensor_descriptor",
2286
+ "tensor_descriptor",
2287
+ "tensor_descriptor",
2288
+ "tensor_descriptor",
2289
+ "pointer"
2290
+ ],
2291
+ "atomic_indexing": [],
2292
+ "pid_type": "persistent_blocked",
2293
+ "num_sm_multiplier": 128,
2294
+ "maxnreg": 64
2295
+ }
2296
+ },
2297
+ {
2298
+ "key": {
2299
+ "q_heads": 32,
2300
+ "kv_heads": 8,
2301
+ "num_tokens": 4096
2302
+ },
2303
+ "config": {
2304
+ "block_sizes": [
2305
+ 2
2306
+ ],
2307
+ "loop_orders": [
2308
+ [
2309
+ 2,
2310
+ 1,
2311
+ 0
2312
+ ]
2313
+ ],
2314
+ "l2_groupings": [
2315
+ 1
2316
+ ],
2317
+ "range_unroll_factors": [
2318
+ 0
2319
+ ],
2320
+ "range_warp_specializes": [
2321
+ false
2322
+ ],
2323
+ "range_multi_buffers": [
2324
+ null
2325
+ ],
2326
+ "range_flattens": [
2327
+ false
2328
+ ],
2329
+ "load_eviction_policies": [
2330
+ "",
2331
+ "",
2332
+ "",
2333
+ "",
2334
+ "",
2335
+ "",
2336
+ "last",
2337
+ ""
2338
+ ],
2339
+ "num_warps": 1,
2340
+ "num_stages": 6,
2341
+ "indexing": [
2342
+ "tensor_descriptor",
2343
+ "tensor_descriptor",
2344
+ "tensor_descriptor",
2345
+ "tensor_descriptor",
2346
+ "tensor_descriptor",
2347
+ "tensor_descriptor",
2348
+ "tensor_descriptor",
2349
+ "tensor_descriptor",
2350
+ "tensor_descriptor",
2351
+ "tensor_descriptor",
2352
+ "pointer"
2353
+ ],
2354
+ "atomic_indexing": [],
2355
+ "pid_type": "persistent_blocked",
2356
+ "num_sm_multiplier": 128,
2357
+ "maxnreg": 64
2358
+ }
2359
+ },
2360
+ {
2361
+ "key": {
2362
+ "q_heads": 64,
2363
+ "kv_heads": 8,
2364
+ "num_tokens": 4096
2365
+ },
2366
+ "config": {
2367
+ "block_sizes": [
2368
+ 2
2369
+ ],
2370
+ "loop_orders": [
2371
+ [
2372
+ 2,
2373
+ 1,
2374
+ 0
2375
+ ]
2376
+ ],
2377
+ "l2_groupings": [
2378
+ 1
2379
+ ],
2380
+ "range_unroll_factors": [
2381
+ 0
2382
+ ],
2383
+ "range_warp_specializes": [
2384
+ false
2385
+ ],
2386
+ "range_multi_buffers": [
2387
+ null
2388
+ ],
2389
+ "range_flattens": [
2390
+ false
2391
+ ],
2392
+ "load_eviction_policies": [
2393
+ "",
2394
+ "",
2395
+ "",
2396
+ "",
2397
+ "",
2398
+ "",
2399
+ "last",
2400
+ ""
2401
+ ],
2402
+ "num_warps": 1,
2403
+ "num_stages": 6,
2404
+ "indexing": [
2405
+ "tensor_descriptor",
2406
+ "tensor_descriptor",
2407
+ "tensor_descriptor",
2408
+ "tensor_descriptor",
2409
+ "tensor_descriptor",
2410
+ "tensor_descriptor",
2411
+ "tensor_descriptor",
2412
+ "tensor_descriptor",
2413
+ "tensor_descriptor",
2414
+ "tensor_descriptor",
2415
+ "pointer"
2416
+ ],
2417
+ "atomic_indexing": [],
2418
+ "pid_type": "persistent_blocked",
2419
+ "num_sm_multiplier": 128,
2420
+ "maxnreg": 64
2421
+ }
2422
+ },
2423
+ {
2424
+ "key": {
2425
+ "q_heads": 16,
2426
+ "kv_heads": 8,
2427
+ "num_tokens": 8192
2428
+ },
2429
+ "config": {
2430
+ "block_sizes": [
2431
+ 2
2432
+ ],
2433
+ "loop_orders": [
2434
+ [
2435
+ 2,
2436
+ 1,
2437
+ 0
2438
+ ]
2439
+ ],
2440
+ "l2_groupings": [
2441
+ 1
2442
+ ],
2443
+ "range_unroll_factors": [
2444
+ 0
2445
+ ],
2446
+ "range_warp_specializes": [
2447
+ null
2448
+ ],
2449
+ "range_multi_buffers": [
2450
+ true
2451
+ ],
2452
+ "range_flattens": [
2453
+ null
2454
+ ],
2455
+ "load_eviction_policies": [
2456
+ "",
2457
+ "",
2458
+ "",
2459
+ "",
2460
+ "",
2461
+ "",
2462
+ "last",
2463
+ "last"
2464
+ ],
2465
+ "num_warps": 1,
2466
+ "num_stages": 6,
2467
+ "indexing": [
2468
+ "pointer",
2469
+ "tensor_descriptor",
2470
+ "tensor_descriptor",
2471
+ "tensor_descriptor",
2472
+ "tensor_descriptor",
2473
+ "tensor_descriptor",
2474
+ "tensor_descriptor",
2475
+ "tensor_descriptor",
2476
+ "tensor_descriptor",
2477
+ "tensor_descriptor",
2478
+ "pointer"
2479
+ ],
2480
+ "atomic_indexing": [],
2481
+ "pid_type": "persistent_blocked",
2482
+ "num_sm_multiplier": 128,
2483
+ "maxnreg": 64
2484
+ }
2485
+ },
2486
+ {
2487
+ "key": {
2488
+ "q_heads": 32,
2489
+ "kv_heads": 8,
2490
+ "num_tokens": 8192
2491
+ },
2492
+ "config": {
2493
+ "block_sizes": [
2494
+ 2
2495
+ ],
2496
+ "loop_orders": [
2497
+ [
2498
+ 2,
2499
+ 1,
2500
+ 0
2501
+ ]
2502
+ ],
2503
+ "l2_groupings": [
2504
+ 1
2505
+ ],
2506
+ "range_unroll_factors": [
2507
+ 0
2508
+ ],
2509
+ "range_warp_specializes": [
2510
+ false
2511
+ ],
2512
+ "range_multi_buffers": [
2513
+ null
2514
+ ],
2515
+ "range_flattens": [
2516
+ false
2517
+ ],
2518
+ "load_eviction_policies": [
2519
+ "",
2520
+ "",
2521
+ "",
2522
+ "",
2523
+ "",
2524
+ "",
2525
+ "last",
2526
+ ""
2527
+ ],
2528
+ "num_warps": 1,
2529
+ "num_stages": 6,
2530
+ "indexing": [
2531
+ "tensor_descriptor",
2532
+ "tensor_descriptor",
2533
+ "tensor_descriptor",
2534
+ "tensor_descriptor",
2535
+ "tensor_descriptor",
2536
+ "tensor_descriptor",
2537
+ "tensor_descriptor",
2538
+ "tensor_descriptor",
2539
+ "tensor_descriptor",
2540
+ "tensor_descriptor",
2541
+ "pointer"
2542
+ ],
2543
+ "atomic_indexing": [],
2544
+ "pid_type": "persistent_blocked",
2545
+ "num_sm_multiplier": 128,
2546
+ "maxnreg": 64
2547
+ }
2548
+ },
2549
+ {
2550
+ "key": {
2551
+ "q_heads": 64,
2552
+ "kv_heads": 8,
2553
+ "num_tokens": 8192
2554
+ },
2555
+ "config": {
2556
+ "block_sizes": [
2557
+ 2
2558
+ ],
2559
+ "loop_orders": [
2560
+ [
2561
+ 2,
2562
+ 1,
2563
+ 0
2564
+ ]
2565
+ ],
2566
+ "l2_groupings": [
2567
+ 1
2568
+ ],
2569
+ "range_unroll_factors": [
2570
+ 0
2571
+ ],
2572
+ "range_warp_specializes": [
2573
+ false
2574
+ ],
2575
+ "range_multi_buffers": [
2576
+ null
2577
+ ],
2578
+ "range_flattens": [
2579
+ false
2580
+ ],
2581
+ "load_eviction_policies": [
2582
+ "",
2583
+ "",
2584
+ "",
2585
+ "",
2586
+ "",
2587
+ "",
2588
+ "last",
2589
+ ""
2590
+ ],
2591
+ "num_warps": 1,
2592
+ "num_stages": 6,
2593
+ "indexing": [
2594
+ "tensor_descriptor",
2595
+ "tensor_descriptor",
2596
+ "tensor_descriptor",
2597
+ "tensor_descriptor",
2598
+ "tensor_descriptor",
2599
+ "tensor_descriptor",
2600
+ "tensor_descriptor",
2601
+ "tensor_descriptor",
2602
+ "tensor_descriptor",
2603
+ "tensor_descriptor",
2604
+ "pointer"
2605
+ ],
2606
+ "atomic_indexing": [],
2607
+ "pid_type": "persistent_blocked",
2608
+ "num_sm_multiplier": 128,
2609
+ "maxnreg": 64
2610
+ }
2611
+ }
2612
+]
\ No newline at end of file
vllm/kernels/helion/configs/fused_qk_norm_rope/nvidia_h100.json
new
+2722
@@ -0,0 +1,2722 @@
1
+[
2
+ {
3
+ "key": {
4
+ "q_heads": 16,
5
+ "kv_heads": 8,
6
+ "num_tokens": 1
7
+ },
8
+ "config": {
9
+ "block_sizes": [
10
+ 1
11
+ ],
12
+ "loop_orders": [
13
+ [
14
+ 1,
15
+ 0,
16
+ 2
17
+ ]
18
+ ],
19
+ "l2_groupings": [
20
+ 64
21
+ ],
22
+ "range_unroll_factors": [
23
+ 4
24
+ ],
25
+ "range_warp_specializes": [],
26
+ "range_multi_buffers": [
27
+ false
28
+ ],
29
+ "range_flattens": [
30
+ true
31
+ ],
32
+ "load_eviction_policies": [
33
+ "",
34
+ "",
35
+ "first",
36
+ "",
37
+ "",
38
+ "last",
39
+ "",
40
+ "first"
41
+ ],
42
+ "num_warps": 1,
43
+ "num_stages": 3,
44
+ "indexing": [
45
+ "tensor_descriptor",
46
+ "pointer",
47
+ "tensor_descriptor",
48
+ "pointer",
49
+ "pointer",
50
+ "pointer",
51
+ "pointer",
52
+ "tensor_descriptor",
53
+ "pointer",
54
+ "tensor_descriptor",
55
+ "pointer"
56
+ ],
57
+ "atomic_indexing": [],
58
+ "pid_type": "persistent_blocked",
59
+ "num_sm_multiplier": 2
60
+ }
61
+ },
62
+ {
63
+ "key": {
64
+ "q_heads": 32,
65
+ "kv_heads": 8,
66
+ "num_tokens": 1
67
+ },
68
+ "config": {
69
+ "block_sizes": [
70
+ 1
71
+ ],
72
+ "loop_orders": [
73
+ [
74
+ 0,
75
+ 2,
76
+ 1
77
+ ]
78
+ ],
79
+ "l2_groupings": [
80
+ 64
81
+ ],
82
+ "range_unroll_factors": [
83
+ 0
84
+ ],
85
+ "range_warp_specializes": [],
86
+ "range_num_stages": [],
87
+ "range_multi_buffers": [
88
+ null
89
+ ],
90
+ "range_flattens": [
91
+ null
92
+ ],
93
+ "load_eviction_policies": [
94
+ "",
95
+ "",
96
+ "first",
97
+ "first",
98
+ "last",
99
+ "",
100
+ "last",
101
+ ""
102
+ ],
103
+ "num_warps": 1,
104
+ "num_stages": 1,
105
+ "indexing": [
106
+ "tensor_descriptor",
107
+ "tensor_descriptor",
108
+ "tensor_descriptor",
109
+ "pointer",
110
+ "tensor_descriptor",
111
+ "pointer",
112
+ "pointer",
113
+ "tensor_descriptor",
114
+ "pointer",
115
+ "tensor_descriptor",
116
+ "tensor_descriptor"
117
+ ],
118
+ "atomic_indexing": [],
119
+ "pid_type": "flat"
120
+ }
121
+ },
122
+ {
123
+ "key": {
124
+ "q_heads": 64,
125
+ "kv_heads": 8,
126
+ "num_tokens": 1
127
+ },
128
+ "config": {
129
+ "block_sizes": [
130
+ 2
131
+ ],
132
+ "loop_orders": [
133
+ [
134
+ 2,
135
+ 0,
136
+ 1
137
+ ]
138
+ ],
139
+ "l2_groupings": [
140
+ 8
141
+ ],
142
+ "range_unroll_factors": [
143
+ 0
144
+ ],
145
+ "range_warp_specializes": [],
146
+ "range_num_stages": [],
147
+ "range_multi_buffers": [
148
+ null
149
+ ],
150
+ "range_flattens": [
151
+ null
152
+ ],
153
+ "load_eviction_policies": [
154
+ "last",
155
+ "last",
156
+ "last",
157
+ "first",
158
+ "first",
159
+ "last",
160
+ "last",
161
+ "first"
162
+ ],
163
+ "num_warps": 2,
164
+ "num_stages": 1,
165
+ "indexing": [
166
+ "tensor_descriptor",
167
+ "tensor_descriptor",
168
+ "pointer",
169
+ "tensor_descriptor",
170
+ "pointer",
171
+ "pointer",
172
+ "pointer",
173
+ "tensor_descriptor",
174
+ "pointer",
175
+ "tensor_descriptor",
176
+ "pointer"
177
+ ],
178
+ "atomic_indexing": [],
179
+ "pid_type": "flat"
180
+ }
181
+ },
182
+ {
183
+ "key": {
184
+ "q_heads": 16,
185
+ "kv_heads": 8,
186
+ "num_tokens": 2
187
+ },
188
+ "config": {
189
+ "block_sizes": [
190
+ 1
191
+ ],
192
+ "loop_orders": [
193
+ [
194
+ 1,
195
+ 2,
196
+ 0
197
+ ]
198
+ ],
199
+ "l2_groupings": [
200
+ 16
201
+ ],
202
+ "range_unroll_factors": [
203
+ 0
204
+ ],
205
+ "range_warp_specializes": [],
206
+ "range_num_stages": [],
207
+ "range_multi_buffers": [
208
+ null
209
+ ],
210
+ "range_flattens": [
211
+ null
212
+ ],
213
+ "load_eviction_policies": [
214
+ "",
215
+ "",
216
+ "",
217
+ "first",
218
+ "first",
219
+ "last",
220
+ "first",
221
+ "last"
222
+ ],
223
+ "num_warps": 1,
224
+ "num_stages": 5,
225
+ "indexing": [
226
+ "tensor_descriptor",
227
+ "pointer",
228
+ "tensor_descriptor",
229
+ "tensor_descriptor",
230
+ "pointer",
231
+ "pointer",
232
+ "tensor_descriptor",
233
+ "tensor_descriptor",
234
+ "tensor_descriptor",
235
+ "tensor_descriptor",
236
+ "pointer"
237
+ ],
238
+ "atomic_indexing": [],
239
+ "pid_type": "flat"
240
+ }
241
+ },
242
+ {
243
+ "key": {
244
+ "q_heads": 32,
245
+ "kv_heads": 8,
246
+ "num_tokens": 2
247
+ },
248
+ "config": {
249
+ "block_sizes": [
250
+ 1
251
+ ],
252
+ "loop_orders": [
253
+ [
254
+ 2,
255
+ 1,
256
+ 0
257
+ ]
258
+ ],
259
+ "l2_groupings": [
260
+ 32
261
+ ],
262
+ "range_unroll_factors": [
263
+ 0
264
+ ],
265
+ "range_warp_specializes": [],
266
+ "range_num_stages": [],
267
+ "range_multi_buffers": [
268
+ null
269
+ ],
270
+ "range_flattens": [
271
+ null
272
+ ],
273
+ "load_eviction_policies": [
274
+ "",
275
+ "",
276
+ "last",
277
+ "first",
278
+ "first",
279
+ "first",
280
+ "last",
281
+ ""
282
+ ],
283
+ "num_warps": 1,
284
+ "num_stages": 8,
285
+ "indexing": [
286
+ "pointer",
287
+ "pointer",
288
+ "pointer",
289
+ "pointer",
290
+ "tensor_descriptor",
291
+ "tensor_descriptor",
292
+ "pointer",
293
+ "pointer",
294
+ "pointer",
295
+ "tensor_descriptor",
296
+ "pointer"
297
+ ],
298
+ "atomic_indexing": [],
299
+ "pid_type": "flat"
300
+ }
301
+ },
302
+ {
303
+ "key": {
304
+ "q_heads": 64,
305
+ "kv_heads": 8,
306
+ "num_tokens": 2
307
+ },
308
+ "config": {
309
+ "block_sizes": [
310
+ 2
311
+ ],
312
+ "loop_orders": [
313
+ [
314
+ 2,
315
+ 1,
316
+ 0
317
+ ]
318
+ ],
319
+ "l2_groupings": [
320
+ 2
321
+ ],
322
+ "range_unroll_factors": [
323
+ 0
324
+ ],
325
+ "range_warp_specializes": [],
326
+ "range_num_stages": [],
327
+ "range_multi_buffers": [
328
+ null
329
+ ],
330
+ "range_flattens": [
331
+ null
332
+ ],
333
+ "load_eviction_policies": [
334
+ "last",
335
+ "first",
336
+ "",
337
+ "",
338
+ "last",
339
+ "",
340
+ "",
341
+ "first"
342
+ ],
343
+ "num_warps": 1,
344
+ "num_stages": 8,
345
+ "indexing": [
346
+ "tensor_descriptor",
347
+ "tensor_descriptor",
348
+ "pointer",
349
+ "pointer",
350
+ "tensor_descriptor",
351
+ "pointer",
352
+ "tensor_descriptor",
353
+ "pointer",
354
+ "pointer",
355
+ "pointer",
356
+ "tensor_descriptor"
357
+ ],
358
+ "atomic_indexing": [],
359
+ "pid_type": "flat"
360
+ }
361
+ },
362
+ {
363
+ "key": {
364
+ "q_heads": 16,
365
+ "kv_heads": 8,
366
+ "num_tokens": 4
367
+ },
368
+ "config": {
369
+ "block_sizes": [
370
+ 1
371
+ ],
372
+ "loop_orders": [
373
+ [
374
+ 1,
375
+ 2,
376
+ 0
377
+ ]
378
+ ],
379
+ "l2_groupings": [
380
+ 16
381
+ ],
382
+ "range_unroll_factors": [
383
+ 0
384
+ ],
385
+ "range_warp_specializes": [],
386
+ "range_num_stages": [],
387
+ "range_multi_buffers": [
388
+ null
389
+ ],
390
+ "range_flattens": [
391
+ null
392
+ ],
393
+ "load_eviction_policies": [
394
+ "last",
395
+ "last",
396
+ "",
397
+ "",
398
+ "last",
399
+ "last",
400
+ "",
401
+ "last"
402
+ ],
403
+ "num_warps": 1,
404
+ "num_stages": 4,
405
+ "indexing": [
406
+ "tensor_descriptor",
407
+ "tensor_descriptor",
408
+ "pointer",
409
+ "tensor_descriptor",
410
+ "pointer",
411
+ "pointer",
412
+ "pointer",
413
+ "tensor_descriptor",
414
+ "tensor_descriptor",
415
+ "pointer",
416
+ "pointer"
417
+ ],
418
+ "atomic_indexing": [],
419
+ "pid_type": "flat"
420
+ }
421
+ },
422
+ {
423
+ "key": {
424
+ "q_heads": 32,
425
+ "kv_heads": 8,
426
+ "num_tokens": 4
427
+ },
428
+ "config": {
429
+ "block_sizes": [
430
+ 2
431
+ ],
432
+ "loop_orders": [
433
+ [
434
+ 2,
435
+ 1,
436
+ 0
437
+ ]
438
+ ],
439
+ "l2_groupings": [
440
+ 32
441
+ ],
442
+ "range_unroll_factors": [
443
+ 0
444
+ ],
445
+ "range_warp_specializes": [],
446
+ "range_num_stages": [],
447
+ "range_multi_buffers": [
448
+ null
449
+ ],
450
+ "range_flattens": [
451
+ null
452
+ ],
453
+ "load_eviction_policies": [
454
+ "first",
455
+ "",
456
+ "last",
457
+ "",
458
+ "first",
459
+ "",
460
+ "first",
461
+ "last"
462
+ ],
463
+ "num_warps": 1,
464
+ "num_stages": 5,
465
+ "indexing": [
466
+ "pointer",
467
+ "pointer",
468
+ "tensor_descriptor",
469
+ "tensor_descriptor",
470
+ "tensor_descriptor",
471
+ "tensor_descriptor",
472
+ "tensor_descriptor",
473
+ "pointer",
474
+ "pointer",
475
+ "pointer",
476
+ "tensor_descriptor"
477
+ ],
478
+ "atomic_indexing": [],
479
+ "pid_type": "flat"
480
+ }
481
+ },
482
+ {
483
+ "key": {
484
+ "q_heads": 64,
485
+ "kv_heads": 8,
486
+ "num_tokens": 4
487
+ },
488
+ "config": {
489
+ "block_sizes": [
490
+ 2
491
+ ],
492
+ "loop_orders": [
493
+ [
494
+ 2,
495
+ 1,
496
+ 0
497
+ ]
498
+ ],
499
+ "l2_groupings": [
500
+ 64
501
+ ],
502
+ "range_unroll_factors": [
503
+ 0
504
+ ],
505
+ "range_warp_specializes": [],
506
+ "range_num_stages": [],
507
+ "range_multi_buffers": [
508
+ null
509
+ ],
510
+ "range_flattens": [
511
+ null
512
+ ],
513
+ "load_eviction_policies": [
514
+ "last",
515
+ "last",
516
+ "last",
517
+ "first",
518
+ "last",
519
+ "last",
520
+ "last",
521
+ ""
522
+ ],
523
+ "num_warps": 1,
524
+ "num_stages": 2,
525
+ "indexing": [
526
+ "tensor_descriptor",
527
+ "tensor_descriptor",
528
+ "tensor_descriptor",
529
+ "tensor_descriptor",
530
+ "tensor_descriptor",
531
+ "pointer",
532
+ "pointer",
533
+ "pointer",
534
+ "tensor_descriptor",
535
+ "pointer",
536
+ "pointer"
537
+ ],
538
+ "atomic_indexing": [],
539
+ "pid_type": "flat"
540
+ }
541
+ },
542
+ {
543
+ "key": {
544
+ "q_heads": 16,
545
+ "kv_heads": 8,
546
+ "num_tokens": 8
547
+ },
548
+ "config": {
549
+ "block_sizes": [
550
+ 1
551
+ ],
552
+ "loop_orders": [
553
+ [
554
+ 1,
555
+ 2,
556
+ 0
557
+ ]
558
+ ],
559
+ "l2_groupings": [
560
+ 16
561
+ ],
562
+ "range_unroll_factors": [
563
+ 0
564
+ ],
565
+ "range_multi_buffers": [
566
+ null
567
+ ],
568
+ "range_flattens": [
569
+ null
570
+ ],
571
+ "load_eviction_policies": [
572
+ "last",
573
+ "first",
574
+ "first",
575
+ "last",
576
+ "first",
577
+ "first",
578
+ "last",
579
+ ""
580
+ ],
581
+ "num_warps": 1,
582
+ "num_stages": 2,
583
+ "indexing": [
584
+ "pointer",
585
+ "pointer",
586
+ "pointer",
587
+ "pointer",
588
+ "pointer",
589
+ "pointer",
590
+ "pointer",
591
+ "pointer",
592
+ "pointer",
593
+ "pointer",
594
+ "pointer"
595
+ ],
596
+ "pid_type": "flat",
597
+ "atomic_indexing": [],
598
+ "range_warp_specializes": [],
599
+ "range_num_stages": []
600
+ }
601
+ },
602
+ {
603
+ "key": {
604
+ "q_heads": 32,
605
+ "kv_heads": 8,
606
+ "num_tokens": 8
607
+ },
608
+ "config": {
609
+ "block_sizes": [
610
+ 2
611
+ ],
612
+ "loop_orders": [
613
+ [
614
+ 2,
615
+ 1,
616
+ 0
617
+ ]
618
+ ],
619
+ "l2_groupings": [
620
+ 64
621
+ ],
622
+ "range_unroll_factors": [
623
+ 0
624
+ ],
625
+ "range_warp_specializes": [],
626
+ "range_num_stages": [],
627
+ "range_multi_buffers": [
628
+ null
629
+ ],
630
+ "range_flattens": [
631
+ null
632
+ ],
633
+ "load_eviction_policies": [
634
+ "last",
635
+ "last",
636
+ "last",
637
+ "first",
638
+ "last",
639
+ "last",
640
+ "last",
641
+ ""
642
+ ],
643
+ "num_warps": 1,
644
+ "num_stages": 2,
645
+ "indexing": [
646
+ "tensor_descriptor",
647
+ "tensor_descriptor",
648
+ "tensor_descriptor",
649
+ "tensor_descriptor",
650
+ "tensor_descriptor",
651
+ "pointer",
652
+ "pointer",
653
+ "pointer",
654
+ "tensor_descriptor",
655
+ "pointer",
656
+ "pointer"
657
+ ],
658
+ "atomic_indexing": [],
659
+ "pid_type": "flat"
660
+ }
661
+ },
662
+ {
663
+ "key": {
664
+ "q_heads": 64,
665
+ "kv_heads": 8,
666
+ "num_tokens": 8
667
+ },
668
+ "config": {
669
+ "block_sizes": [
670
+ 2
671
+ ],
672
+ "loop_orders": [
673
+ [
674
+ 2,
675
+ 1,
676
+ 0
677
+ ]
678
+ ],
679
+ "l2_groupings": [
680
+ 64
681
+ ],
682
+ "range_unroll_factors": [
683
+ 0
684
+ ],
685
+ "range_warp_specializes": [],
686
+ "range_num_stages": [],
687
+ "range_multi_buffers": [
688
+ null
689
+ ],
690
+ "range_flattens": [
691
+ null
692
+ ],
693
+ "load_eviction_policies": [
694
+ "last",
695
+ "last",
696
+ "last",
697
+ "first",
698
+ "last",
699
+ "last",
700
+ "last",
701
+ ""
702
+ ],
703
+ "num_warps": 1,
704
+ "num_stages": 2,
705
+ "indexing": [
706
+ "tensor_descriptor",
707
+ "tensor_descriptor",
708
+ "tensor_descriptor",
709
+ "tensor_descriptor",
710
+ "tensor_descriptor",
711
+ "pointer",
712
+ "pointer",
713
+ "pointer",
714
+ "tensor_descriptor",
715
+ "pointer",
716
+ "pointer"
717
+ ],
718
+ "atomic_indexing": [],
719
+ "pid_type": "flat"
720
+ }
721
+ },
722
+ {
723
+ "key": {
724
+ "q_heads": 16,
725
+ "kv_heads": 8,
726
+ "num_tokens": 16
727
+ },
728
+ "config": {
729
+ "block_sizes": [
730
+ 1
731
+ ],
732
+ "loop_orders": [
733
+ [
734
+ 1,
735
+ 2,
736
+ 0
737
+ ]
738
+ ],
739
+ "l2_groupings": [
740
+ 16
741
+ ],
742
+ "range_unroll_factors": [
743
+ 0
744
+ ],
745
+ "range_warp_specializes": [],
746
+ "range_num_stages": [],
747
+ "range_multi_buffers": [
748
+ null
749
+ ],
750
+ "range_flattens": [
751
+ null
752
+ ],
753
+ "load_eviction_policies": [
754
+ "last",
755
+ "last",
756
+ "",
757
+ "",
758
+ "last",
759
+ "last",
760
+ "",
761
+ "last"
762
+ ],
763
+ "num_warps": 1,
764
+ "num_stages": 4,
765
+ "indexing": [
766
+ "tensor_descriptor",
767
+ "tensor_descriptor",
768
+ "pointer",
769
+ "tensor_descriptor",
770
+ "pointer",
771
+ "pointer",
772
+ "pointer",
773
+ "tensor_descriptor",
774
+ "tensor_descriptor",
775
+ "pointer",
776
+ "pointer"
777
+ ],
778
+ "atomic_indexing": [],
779
+ "pid_type": "flat"
780
+ }
781
+ },
782
+ {
783
+ "key": {
784
+ "q_heads": 32,
785
+ "kv_heads": 8,
786
+ "num_tokens": 16
787
+ },
788
+ "config": {
789
+ "block_sizes": [
790
+ 2
791
+ ],
792
+ "loop_orders": [
793
+ [
794
+ 1,
795
+ 2,
796
+ 0
797
+ ]
798
+ ],
799
+ "l2_groupings": [
800
+ 64
801
+ ],
802
+ "range_unroll_factors": [
803
+ 0
804
+ ],
805
+ "range_warp_specializes": [],
806
+ "range_num_stages": [],
807
+ "range_multi_buffers": [
808
+ null
809
+ ],
810
+ "range_flattens": [
811
+ null
812
+ ],
813
+ "load_eviction_policies": [
814
+ "last",
815
+ "",
816
+ "last",
817
+ "first",
818
+ "last",
819
+ "last",
820
+ "last",
821
+ ""
822
+ ],
823
+ "num_warps": 1,
824
+ "num_stages": 8,
825
+ "indexing": [
826
+ "pointer",
827
+ "tensor_descriptor",
828
+ "pointer",
829
+ "pointer",
830
+ "pointer",
831
+ "pointer",
832
+ "pointer",
833
+ "tensor_descriptor",
834
+ "pointer",
835
+ "tensor_descriptor",
836
+ "pointer"
837
+ ],
838
+ "atomic_indexing": [],
839
+ "pid_type": "flat"
840
+ }
841
+ },
842
+ {
843
+ "key": {
844
+ "q_heads": 64,
845
+ "kv_heads": 8,
846
+ "num_tokens": 16
847
+ },
848
+ "config": {
849
+ "block_sizes": [
850
+ 4
851
+ ],
852
+ "loop_orders": [
853
+ [
854
+ 2,
855
+ 1,
856
+ 0
857
+ ]
858
+ ],
859
+ "l2_groupings": [
860
+ 32
861
+ ],
862
+ "range_unroll_factors": [
863
+ 0
864
+ ],
865
+ "range_warp_specializes": [],
866
+ "range_num_stages": [],
867
+ "range_multi_buffers": [
868
+ null
869
+ ],
870
+ "range_flattens": [
871
+ null
872
+ ],
873
+ "load_eviction_policies": [
874
+ "first",
875
+ "last",
876
+ "",
877
+ "first",
878
+ "last",
879
+ "last",
880
+ "first",
881
+ "first"
882
+ ],
883
+ "num_warps": 1,
884
+ "num_stages": 4,
885
+ "indexing": [
886
+ "tensor_descriptor",
887
+ "tensor_descriptor",
888
+ "pointer",
889
+ "pointer",
890
+ "pointer",
891
+ "tensor_descriptor",
892
+ "tensor_descriptor",
893
+ "tensor_descriptor",
894
+ "pointer",
895
+ "tensor_descriptor",
896
+ "tensor_descriptor"
897
+ ],
898
+ "atomic_indexing": [],
899
+ "pid_type": "flat"
900
+ }
901
+ },
902
+ {
903
+ "key": {
904
+ "q_heads": 16,
905
+ "kv_heads": 8,
906
+ "num_tokens": 32
907
+ },
908
+ "config": {
909
+ "block_sizes": [
910
+ 2
911
+ ],
912
+ "loop_orders": [
913
+ [
914
+ 1,
915
+ 2,
916
+ 0
917
+ ]
918
+ ],
919
+ "l2_groupings": [
920
+ 32
921
+ ],
922
+ "range_unroll_factors": [
923
+ 0
924
+ ],
925
+ "range_warp_specializes": [],
926
+ "range_num_stages": [],
927
+ "range_multi_buffers": [
928
+ null
929
+ ],
930
+ "range_flattens": [
931
+ null
932
+ ],
933
+ "load_eviction_policies": [
934
+ "",
935
+ "",
936
+ "",
937
+ "first",
938
+ "last",
939
+ "",
940
+ "first",
941
+ "first"
942
+ ],
943
+ "num_warps": 1,
944
+ "num_stages": 3,
945
+ "indexing": [
946
+ "pointer",
947
+ "pointer",
948
+ "tensor_descriptor",
949
+ "tensor_descriptor",
950
+ "pointer",
951
+ "pointer",
952
+ "pointer",
953
+ "tensor_descriptor",
954
+ "pointer",
955
+ "pointer",
956
+ "tensor_descriptor"
957
+ ],
958
+ "atomic_indexing": [],
959
+ "pid_type": "flat"
960
+ }
961
+ },
962
+ {
963
+ "key": {
964
+ "q_heads": 32,
965
+ "kv_heads": 8,
966
+ "num_tokens": 32
967
+ },
968
+ "config": {
969
+ "block_sizes": [
970
+ 4
971
+ ],
972
+ "loop_orders": [
973
+ [
974
+ 1,
975
+ 2,
976
+ 0
977
+ ]
978
+ ],
979
+ "l2_groupings": [
980
+ 32
981
+ ],
982
+ "range_unroll_factors": [
983
+ 0
984
+ ],
985
+ "range_multi_buffers": [
986
+ null
987
+ ],
988
+ "range_flattens": [
989
+ null
990
+ ],
991
+ "load_eviction_policies": [
992
+ "",
993
+ "first",
994
+ "",
995
+ "first",
996
+ "last",
997
+ "",
998
+ "last",
999
+ "last"
1000
+ ],
1001
+ "num_warps": 1,
1002
+ "num_stages": 2,
1003
+ "indexing": [
1004
+ "pointer",
1005
+ "pointer",
1006
+ "pointer",
1007
+ "pointer",
1008
+ "pointer",
1009
+ "pointer",
1010
+ "pointer",
1011
+ "pointer",
1012
+ "pointer",
1013
+ "pointer",
1014
+ "pointer"
1015
+ ],
1016
+ "pid_type": "flat",
1017
+ "atomic_indexing": [],
1018
+ "range_warp_specializes": [],
1019
+ "range_num_stages": []
1020
+ }
1021
+ },
1022
+ {
1023
+ "key": {
1024
+ "q_heads": 64,
1025
+ "kv_heads": 8,
1026
+ "num_tokens": 32
1027
+ },
1028
+ "config": {
1029
+ "block_sizes": [
1030
+ 4
1031
+ ],
1032
+ "loop_orders": [
1033
+ [
1034
+ 2,
1035
+ 1,
1036
+ 0
1037
+ ]
1038
+ ],
1039
+ "l2_groupings": [
1040
+ 8
1041
+ ],
1042
+ "range_unroll_factors": [
1043
+ 0
1044
+ ],
1045
+ "range_warp_specializes": [],
1046
+ "range_num_stages": [],
1047
+ "range_multi_buffers": [
1048
+ null
1049
+ ],
1050
+ "range_flattens": [
1051
+ null
1052
+ ],
1053
+ "load_eviction_policies": [
1054
+ "last",
1055
+ "",
1056
+ "last",
1057
+ "last",
1058
+ "first",
1059
+ "last",
1060
+ "first",
1061
+ ""
1062
+ ],
1063
+ "num_warps": 1,
1064
+ "num_stages": 2,
1065
+ "indexing": [
1066
+ "tensor_descriptor",
1067
+ "pointer",
1068
+ "tensor_descriptor",
1069
+ "tensor_descriptor",
1070
+ "tensor_descriptor",
1071
+ "pointer",
1072
+ "pointer",
1073
+ "pointer",
1074
+ "pointer",
1075
+ "tensor_descriptor",
1076
+ "tensor_descriptor"
1077
+ ],
1078
+ "atomic_indexing": [],
1079
+ "pid_type": "flat"
1080
+ }
1081
+ },
1082
+ {
1083
+ "key": {
1084
+ "q_heads": 16,
1085
+ "kv_heads": 8,
1086
+ "num_tokens": 64
1087
+ },
1088
+ "config": {
1089
+ "block_sizes": [
1090
+ 4
1091
+ ],
1092
+ "loop_orders": [
1093
+ [
1094
+ 1,
1095
+ 2,
1096
+ 0
1097
+ ]
1098
+ ],
1099
+ "l2_groupings": [
1100
+ 8
1101
+ ],
1102
+ "range_unroll_factors": [
1103
+ 0
1104
+ ],
1105
+ "range_warp_specializes": [],
1106
+ "range_num_stages": [],
1107
+ "range_multi_buffers": [
1108
+ null
1109
+ ],
1110
+ "range_flattens": [
1111
+ null
1112
+ ],
1113
+ "load_eviction_policies": [
1114
+ "",
1115
+ "first",
1116
+ "last",
1117
+ "",
1118
+ "last",
1119
+ "last",
1120
+ "first",
1121
+ "first"
1122
+ ],
1123
+ "num_warps": 1,
1124
+ "num_stages": 1,
1125
+ "indexing": [
1126
+ "pointer",
1127
+ "tensor_descriptor",
1128
+ "pointer",
1129
+ "tensor_descriptor",
1130
+ "pointer",
1131
+ "pointer",
1132
+ "pointer",
1133
+ "tensor_descriptor",
1134
+ "tensor_descriptor",
1135
+ "pointer",
1136
+ "pointer"
1137
+ ],
1138
+ "atomic_indexing": [],
1139
+ "pid_type": "flat"
1140
+ }
1141
+ },
1142
+ {
1143
+ "key": {
1144
+ "q_heads": 32,
1145
+ "kv_heads": 8,
1146
+ "num_tokens": 64
1147
+ },
1148
+ "config": {
1149
+ "block_sizes": [
1150
+ 4
1151
+ ],
1152
+ "loop_orders": [
1153
+ [
1154
+ 2,
1155
+ 1,
1156
+ 0
1157
+ ]
1158
+ ],
1159
+ "l2_groupings": [
1160
+ 4
1161
+ ],
1162
+ "range_unroll_factors": [
1163
+ 0
1164
+ ],
1165
+ "range_warp_specializes": [],
1166
+ "range_num_stages": [],
1167
+ "range_multi_buffers": [
1168
+ null
1169
+ ],
1170
+ "range_flattens": [
1171
+ null
1172
+ ],
1173
+ "load_eviction_policies": [
1174
+ "first",
1175
+ "first",
1176
+ "last",
1177
+ "",
1178
+ "last",
1179
+ "last",
1180
+ "last",
1181
+ "first"
1182
+ ],
1183
+ "num_warps": 2,
1184
+ "num_stages": 2,
1185
+ "indexing": [
1186
+ "pointer",
1187
+ "pointer",
1188
+ "pointer",
1189
+ "tensor_descriptor",
1190
+ "tensor_descriptor",
1191
+ "tensor_descriptor",
1192
+ "tensor_descriptor",
1193
+ "pointer",
1194
+ "pointer",
1195
+ "pointer",
1196
+ "pointer"
1197
+ ],
1198
+ "atomic_indexing": [],
1199
+ "pid_type": "flat"
1200
+ }
1201
+ },
1202
+ {
1203
+ "key": {
1204
+ "q_heads": 64,
1205
+ "kv_heads": 8,
1206
+ "num_tokens": 64
1207
+ },
1208
+ "config": {
1209
+ "block_sizes": [
1210
+ 4
1211
+ ],
1212
+ "loop_orders": [
1213
+ [
1214
+ 2,
1215
+ 1,
1216
+ 0
1217
+ ]
1218
+ ],
1219
+ "l2_groupings": [
1220
+ 2
1221
+ ],
1222
+ "range_unroll_factors": [
1223
+ 0
1224
+ ],
1225
+ "range_warp_specializes": [],
1226
+ "range_num_stages": [],
1227
+ "range_multi_buffers": [
1228
+ null
1229
+ ],
1230
+ "range_flattens": [
1231
+ null
1232
+ ],
1233
+ "load_eviction_policies": [
1234
+ "first",
1235
+ "last",
1236
+ "last",
1237
+ "last",
1238
+ "last",
1239
+ "first",
1240
+ "first",
1241
+ ""
1242
+ ],
1243
+ "num_warps": 2,
1244
+ "num_stages": 6,
1245
+ "indexing": [
1246
+ "tensor_descriptor",
1247
+ "tensor_descriptor",
1248
+ "pointer",
1249
+ "tensor_descriptor",
1250
+ "tensor_descriptor",
1251
+ "pointer",
1252
+ "tensor_descriptor",
1253
+ "tensor_descriptor",
1254
+ "pointer",
1255
+ "pointer",
1256
+ "tensor_descriptor"
1257
+ ],
1258
+ "atomic_indexing": [],
1259
+ "pid_type": "flat"
1260
+ }
1261
+ },
1262
+ {
1263
+ "key": {
1264
+ "q_heads": 16,
1265
+ "kv_heads": 8,
1266
+ "num_tokens": 128
1267
+ },
1268
+ "config": {
1269
+ "block_sizes": [
1270
+ 4
1271
+ ],
1272
+ "loop_orders": [
1273
+ [
1274
+ 1,
1275
+ 2,
1276
+ 0
1277
+ ]
1278
+ ],
1279
+ "l2_groupings": [
1280
+ 8
1281
+ ],
1282
+ "range_unroll_factors": [
1283
+ 0
1284
+ ],
1285
+ "range_warp_specializes": [],
1286
+ "range_num_stages": [],
1287
+ "range_multi_buffers": [
1288
+ null
1289
+ ],
1290
+ "range_flattens": [
1291
+ null
1292
+ ],
1293
+ "load_eviction_policies": [
1294
+ "",
1295
+ "first",
1296
+ "last",
1297
+ "",
1298
+ "last",
1299
+ "last",
1300
+ "first",
1301
+ "first"
1302
+ ],
1303
+ "num_warps": 1,
1304
+ "num_stages": 1,
1305
+ "indexing": [
1306
+ "pointer",
1307
+ "tensor_descriptor",
1308
+ "pointer",
1309
+ "tensor_descriptor",
1310
+ "pointer",
1311
+ "pointer",
1312
+ "pointer",
1313
+ "tensor_descriptor",
1314
+ "tensor_descriptor",
1315
+ "pointer",
1316
+ "pointer"
1317
+ ],
1318
+ "atomic_indexing": [],
1319
+ "pid_type": "flat"
1320
+ }
1321
+ },
1322
+ {
1323
+ "key": {
1324
+ "q_heads": 32,
1325
+ "kv_heads": 8,
1326
+ "num_tokens": 128
1327
+ },
1328
+ "config": {
1329
+ "block_sizes": [
1330
+ 4
1331
+ ],
1332
+ "loop_orders": [
1333
+ [
1334
+ 2,
1335
+ 1,
1336
+ 0
1337
+ ]
1338
+ ],
1339
+ "l2_groupings": [
1340
+ 64
1341
+ ],
1342
+ "range_unroll_factors": [
1343
+ 0
1344
+ ],
1345
+ "range_warp_specializes": [],
1346
+ "range_num_stages": [],
1347
+ "range_multi_buffers": [
1348
+ null
1349
+ ],
1350
+ "range_flattens": [
1351
+ null
1352
+ ],
1353
+ "load_eviction_policies": [
1354
+ "",
1355
+ "",
1356
+ "",
1357
+ "",
1358
+ "last",
1359
+ "last",
1360
+ "last",
1361
+ ""
1362
+ ],
1363
+ "num_warps": 2,
1364
+ "num_stages": 5,
1365
+ "indexing": [
1366
+ "pointer",
1367
+ "pointer",
1368
+ "pointer",
1369
+ "pointer",
1370
+ "pointer",
1371
+ "tensor_descriptor",
1372
+ "pointer",
1373
+ "tensor_descriptor",
1374
+ "tensor_descriptor",
1375
+ "tensor_descriptor",
1376
+ "pointer"
1377
+ ],
1378
+ "atomic_indexing": [],
1379
+ "pid_type": "flat"
1380
+ }
1381
+ },
1382
+ {
1383
+ "key": {
1384
+ "q_heads": 64,
1385
+ "kv_heads": 8,
1386
+ "num_tokens": 128
1387
+ },
1388
+ "config": {
1389
+ "block_sizes": [
1390
+ 4
1391
+ ],
1392
+ "loop_orders": [
1393
+ [
1394
+ 2,
1395
+ 1,
1396
+ 0
1397
+ ]
1398
+ ],
1399
+ "l2_groupings": [
1400
+ 8
1401
+ ],
1402
+ "range_unroll_factors": [
1403
+ 0
1404
+ ],
1405
+ "range_warp_specializes": [],
1406
+ "range_num_stages": [],
1407
+ "range_multi_buffers": [
1408
+ null
1409
+ ],
1410
+ "range_flattens": [
1411
+ null
1412
+ ],
1413
+ "load_eviction_policies": [
1414
+ "",
1415
+ "first",
1416
+ "",
1417
+ "",
1418
+ "first",
1419
+ "",
1420
+ "first",
1421
+ "first"
1422
+ ],
1423
+ "num_warps": 2,
1424
+ "num_stages": 2,
1425
+ "indexing": [
1426
+ "tensor_descriptor",
1427
+ "tensor_descriptor",
1428
+ "tensor_descriptor",
1429
+ "tensor_descriptor",
1430
+ "tensor_descriptor",
1431
+ "pointer",
1432
+ "pointer",
1433
+ "tensor_descriptor",
1434
+ "tensor_descriptor",
1435
+ "pointer",
1436
+ "pointer"
1437
+ ],
1438
+ "atomic_indexing": [],
1439
+ "pid_type": "flat"
1440
+ }
1441
+ },
1442
+ {
1443
+ "key": {
1444
+ "q_heads": 16,
1445
+ "kv_heads": 8,
1446
+ "num_tokens": 256
1447
+ },
1448
+ "config": {
1449
+ "block_sizes": [
1450
+ 4
1451
+ ],
1452
+ "loop_orders": [
1453
+ [
1454
+ 2,
1455
+ 1,
1456
+ 0
1457
+ ]
1458
+ ],
1459
+ "l2_groupings": [
1460
+ 64
1461
+ ],
1462
+ "range_unroll_factors": [
1463
+ 0
1464
+ ],
1465
+ "range_warp_specializes": [],
1466
+ "range_num_stages": [],
1467
+ "range_multi_buffers": [
1468
+ null
1469
+ ],
1470
+ "range_flattens": [
1471
+ null
1472
+ ],
1473
+ "load_eviction_policies": [
1474
+ "",
1475
+ "",
1476
+ "",
1477
+ "",
1478
+ "last",
1479
+ "last",
1480
+ "last",
1481
+ ""
1482
+ ],
1483
+ "num_warps": 2,
1484
+ "num_stages": 5,
1485
+ "indexing": [
1486
+ "pointer",
1487
+ "pointer",
1488
+ "pointer",
1489
+ "pointer",
1490
+ "pointer",
1491
+ "tensor_descriptor",
1492
+ "pointer",
1493
+ "tensor_descriptor",
1494
+ "tensor_descriptor",
1495
+ "tensor_descriptor",
1496
+ "pointer"
1497
+ ],
1498
+ "atomic_indexing": [],
1499
+ "pid_type": "flat"
1500
+ }
1501
+ },
1502
+ {
1503
+ "key": {
1504
+ "q_heads": 32,
1505
+ "kv_heads": 8,
1506
+ "num_tokens": 256
1507
+ },
1508
+ "config": {
1509
+ "block_sizes": [
1510
+ 2
1511
+ ],
1512
+ "loop_orders": [
1513
+ [
1514
+ 1,
1515
+ 2,
1516
+ 0
1517
+ ]
1518
+ ],
1519
+ "l2_groupings": [
1520
+ 64
1521
+ ],
1522
+ "range_unroll_factors": [
1523
+ 2
1524
+ ],
1525
+ "range_warp_specializes": [],
1526
+ "range_multi_buffers": [
1527
+ null
1528
+ ],
1529
+ "range_flattens": [
1530
+ null
1531
+ ],
1532
+ "load_eviction_policies": [
1533
+ "",
1534
+ "last",
1535
+ "",
1536
+ "last",
1537
+ "first",
1538
+ "first",
1539
+ "last",
1540
+ "first"
1541
+ ],
1542
+ "num_warps": 1,
1543
+ "num_stages": 5,
1544
+ "indexing": [
1545
+ "pointer",
1546
+ "tensor_descriptor",
1547
+ "pointer",
1548
+ "tensor_descriptor",
1549
+ "pointer",
1550
+ "tensor_descriptor",
1551
+ "tensor_descriptor",
1552
+ "tensor_descriptor",
1553
+ "tensor_descriptor",
1554
+ "tensor_descriptor",
1555
+ "pointer"
1556
+ ],
1557
+ "atomic_indexing": [],
1558
+ "pid_type": "persistent_blocked",
1559
+ "num_sm_multiplier": 32,
1560
+ "maxnreg": 64
1561
+ }
1562
+ },
1563
+ {
1564
+ "key": {
1565
+ "q_heads": 64,
1566
+ "kv_heads": 8,
1567
+ "num_tokens": 256
1568
+ },
1569
+ "config": {
1570
+ "block_sizes": [
1571
+ 2
1572
+ ],
1573
+ "loop_orders": [
1574
+ [
1575
+ 1,
1576
+ 2,
1577
+ 0
1578
+ ]
1579
+ ],
1580
+ "l2_groupings": [
1581
+ 64
1582
+ ],
1583
+ "range_unroll_factors": [
1584
+ 2
1585
+ ],
1586
+ "range_warp_specializes": [],
1587
+ "range_multi_buffers": [
1588
+ null
1589
+ ],
1590
+ "range_flattens": [
1591
+ null
1592
+ ],
1593
+ "load_eviction_policies": [
1594
+ "",
1595
+ "last",
1596
+ "",
1597
+ "last",
1598
+ "first",
1599
+ "first",
1600
+ "last",
1601
+ "first"
1602
+ ],
1603
+ "num_warps": 1,
1604
+ "num_stages": 5,
1605
+ "indexing": [
1606
+ "pointer",
1607
+ "tensor_descriptor",
1608
+ "pointer",
1609
+ "tensor_descriptor",
1610
+ "pointer",
1611
+ "tensor_descriptor",
1612
+ "tensor_descriptor",
1613
+ "tensor_descriptor",
1614
+ "tensor_descriptor",
1615
+ "tensor_descriptor",
1616
+ "pointer"
1617
+ ],
1618
+ "atomic_indexing": [],
1619
+ "pid_type": "persistent_blocked",
1620
+ "num_sm_multiplier": 32,
1621
+ "maxnreg": 64
1622
+ }
1623
+ },
1624
+ {
1625
+ "key": {
1626
+ "q_heads": 16,
1627
+ "kv_heads": 8,
1628
+ "num_tokens": 512
1629
+ },
1630
+ "config": {
1631
+ "block_sizes": [
1632
+ 2
1633
+ ],
1634
+ "loop_orders": [
1635
+ [
1636
+ 1,
1637
+ 2,
1638
+ 0
1639
+ ]
1640
+ ],
1641
+ "l2_groupings": [
1642
+ 64
1643
+ ],
1644
+ "range_unroll_factors": [
1645
+ 2
1646
+ ],
1647
+ "range_warp_specializes": [],
1648
+ "range_multi_buffers": [
1649
+ null
1650
+ ],
1651
+ "range_flattens": [
1652
+ null
1653
+ ],
1654
+ "load_eviction_policies": [
1655
+ "",
1656
+ "last",
1657
+ "",
1658
+ "last",
1659
+ "first",
1660
+ "first",
1661
+ "last",
1662
+ "first"
1663
+ ],
1664
+ "num_warps": 1,
1665
+ "num_stages": 5,
1666
+ "indexing": [
1667
+ "pointer",
1668
+ "tensor_descriptor",
1669
+ "pointer",
1670
+ "tensor_descriptor",
1671
+ "pointer",
1672
+ "tensor_descriptor",
1673
+ "tensor_descriptor",
1674
+ "tensor_descriptor",
1675
+ "tensor_descriptor",
1676
+ "tensor_descriptor",
1677
+ "pointer"
1678
+ ],
1679
+ "atomic_indexing": [],
1680
+ "pid_type": "persistent_blocked",
1681
+ "num_sm_multiplier": 32,
1682
+ "maxnreg": 64
1683
+ }
1684
+ },
1685
+ {
1686
+ "key": {
1687
+ "q_heads": 32,
1688
+ "kv_heads": 8,
1689
+ "num_tokens": 512
1690
+ },
1691
+ "config": {
1692
+ "block_sizes": [
1693
+ 2
1694
+ ],
1695
+ "loop_orders": [
1696
+ [
1697
+ 1,
1698
+ 2,
1699
+ 0
1700
+ ]
1701
+ ],
1702
+ "l2_groupings": [
1703
+ 64
1704
+ ],
1705
+ "range_unroll_factors": [
1706
+ 2
1707
+ ],
1708
+ "range_warp_specializes": [],
1709
+ "range_multi_buffers": [
1710
+ null
1711
+ ],
1712
+ "range_flattens": [
1713
+ null
1714
+ ],
1715
+ "load_eviction_policies": [
1716
+ "",
1717
+ "last",
1718
+ "",
1719
+ "last",
1720
+ "first",
1721
+ "first",
1722
+ "last",
1723
+ "first"
1724
+ ],
1725
+ "num_warps": 1,
1726
+ "num_stages": 5,
1727
+ "indexing": [
1728
+ "pointer",
1729
+ "tensor_descriptor",
1730
+ "pointer",
1731
+ "tensor_descriptor",
1732
+ "pointer",
1733
+ "tensor_descriptor",
1734
+ "tensor_descriptor",
1735
+ "tensor_descriptor",
1736
+ "tensor_descriptor",
1737
+ "tensor_descriptor",
1738
+ "pointer"
1739
+ ],
1740
+ "atomic_indexing": [],
1741
+ "pid_type": "persistent_blocked",
1742
+ "num_sm_multiplier": 32,
1743
+ "maxnreg": 64
1744
+ }
1745
+ },
1746
+ {
1747
+ "key": {
1748
+ "q_heads": 64,
1749
+ "kv_heads": 8,
1750
+ "num_tokens": 512
1751
+ },
1752
+ "config": {
1753
+ "block_sizes": [
1754
+ 2
1755
+ ],
1756
+ "loop_orders": [
1757
+ [
1758
+ 2,
1759
+ 1,
1760
+ 0
1761
+ ]
1762
+ ],
1763
+ "l2_groupings": [
1764
+ 2
1765
+ ],
1766
+ "range_unroll_factors": [
1767
+ 4
1768
+ ],
1769
+ "range_warp_specializes": [],
1770
+ "range_multi_buffers": [
1771
+ null
1772
+ ],
1773
+ "range_flattens": [
1774
+ true
1775
+ ],
1776
+ "load_eviction_policies": [
1777
+ "first",
1778
+ "",
1779
+ "first",
1780
+ "first",
1781
+ "last",
1782
+ "last",
1783
+ "",
1784
+ ""
1785
+ ],
1786
+ "num_warps": 1,
1787
+ "num_stages": 5,
1788
+ "indexing": [
1789
+ "pointer",
1790
+ "tensor_descriptor",
1791
+ "tensor_descriptor",
1792
+ "pointer",
1793
+ "tensor_descriptor",
1794
+ "tensor_descriptor",
1795
+ "pointer",
1796
+ "pointer",
1797
+ "tensor_descriptor",
1798
+ "tensor_descriptor",
1799
+ "tensor_descriptor"
1800
+ ],
1801
+ "atomic_indexing": [],
1802
+ "pid_type": "persistent_interleaved",
1803
+ "num_sm_multiplier": 64,
1804
+ "maxnreg": 64
1805
+ }
1806
+ },
1807
+ {
1808
+ "key": {
1809
+ "q_heads": 16,
1810
+ "kv_heads": 8,
1811
+ "num_tokens": 1024
1812
+ },
1813
+ "config": {
1814
+ "block_sizes": [
1815
+ 2
1816
+ ],
1817
+ "loop_orders": [
1818
+ [
1819
+ 1,
1820
+ 2,
1821
+ 0
1822
+ ]
1823
+ ],
1824
+ "l2_groupings": [
1825
+ 64
1826
+ ],
1827
+ "range_unroll_factors": [
1828
+ 2
1829
+ ],
1830
+ "range_warp_specializes": [],
1831
+ "range_multi_buffers": [
1832
+ null
1833
+ ],
1834
+ "range_flattens": [
1835
+ null
1836
+ ],
1837
+ "load_eviction_policies": [
1838
+ "",
1839
+ "last",
1840
+ "",
1841
+ "last",
1842
+ "first",
1843
+ "first",
1844
+ "last",
1845
+ "first"
1846
+ ],
1847
+ "num_warps": 1,
1848
+ "num_stages": 5,
1849
+ "indexing": [
1850
+ "pointer",
1851
+ "tensor_descriptor",
1852
+ "pointer",
1853
+ "tensor_descriptor",
1854
+ "pointer",
1855
+ "tensor_descriptor",
1856
+ "tensor_descriptor",
1857
+ "tensor_descriptor",
1858
+ "tensor_descriptor",
1859
+ "tensor_descriptor",
1860
+ "pointer"
1861
+ ],
1862
+ "atomic_indexing": [],
1863
+ "pid_type": "persistent_blocked",
1864
+ "num_sm_multiplier": 32,
1865
+ "maxnreg": 64
1866
+ }
1867
+ },
1868
+ {
1869
+ "key": {
1870
+ "q_heads": 32,
1871
+ "kv_heads": 8,
1872
+ "num_tokens": 1024
1873
+ },
1874
+ "config": {
1875
+ "block_sizes": [
1876
+ 2
1877
+ ],
1878
+ "loop_orders": [
1879
+ [
1880
+ 1,
1881
+ 2,
1882
+ 0
1883
+ ]
1884
+ ],
1885
+ "l2_groupings": [
1886
+ 64
1887
+ ],
1888
+ "range_unroll_factors": [
1889
+ 2
1890
+ ],
1891
+ "range_warp_specializes": [],
1892
+ "range_multi_buffers": [
1893
+ null
1894
+ ],
1895
+ "range_flattens": [
1896
+ null
1897
+ ],
1898
+ "load_eviction_policies": [
1899
+ "",
1900
+ "last",
1901
+ "",
1902
+ "last",
1903
+ "first",
1904
+ "first",
1905
+ "last",
1906
+ "first"
1907
+ ],
1908
+ "num_warps": 1,
1909
+ "num_stages": 5,
1910
+ "indexing": [
1911
+ "pointer",
1912
+ "tensor_descriptor",
1913
+ "pointer",
1914
+ "tensor_descriptor",
1915
+ "pointer",
1916
+ "tensor_descriptor",
1917
+ "tensor_descriptor",
1918
+ "tensor_descriptor",
1919
+ "tensor_descriptor",
1920
+ "tensor_descriptor",
1921
+ "pointer"
1922
+ ],
1923
+ "atomic_indexing": [],
1924
+ "pid_type": "persistent_blocked",
1925
+ "num_sm_multiplier": 32,
1926
+ "maxnreg": 64
1927
+ }
1928
+ },
1929
+ {
1930
+ "key": {
1931
+ "q_heads": 64,
1932
+ "kv_heads": 8,
1933
+ "num_tokens": 1024
1934
+ },
1935
+ "config": {
1936
+ "block_sizes": [
1937
+ 2
1938
+ ],
1939
+ "loop_orders": [
1940
+ [
1941
+ 2,
1942
+ 1,
1943
+ 0
1944
+ ]
1945
+ ],
1946
+ "l2_groupings": [
1947
+ 2
1948
+ ],
1949
+ "range_unroll_factors": [
1950
+ 0
1951
+ ],
1952
+ "range_multi_buffers": [
1953
+ null
1954
+ ],
1955
+ "range_flattens": [
1956
+ null
1957
+ ],
1958
+ "load_eviction_policies": [
1959
+ "",
1960
+ "last",
1961
+ "first",
1962
+ "last",
1963
+ "",
1964
+ "last",
1965
+ "",
1966
+ "first"
1967
+ ],
1968
+ "num_warps": 1,
1969
+ "num_stages": 6,
1970
+ "indexing": [
1971
+ "pointer",
1972
+ "pointer",
1973
+ "pointer",
1974
+ "pointer",
1975
+ "pointer",
1976
+ "pointer",
1977
+ "pointer",
1978
+ "pointer",
1979
+ "pointer",
1980
+ "pointer",
1981
+ "pointer"
1982
+ ],
1983
+ "pid_type": "persistent_interleaved",
1984
+ "num_sm_multiplier": 128,
1985
+ "maxnreg": 64,
1986
+ "atomic_indexing": [],
1987
+ "range_warp_specializes": []
1988
+ }
1989
+ },
1990
+ {
1991
+ "key": {
1992
+ "q_heads": 16,
1993
+ "kv_heads": 8,
1994
+ "num_tokens": 2048
1995
+ },
1996
+ "config": {
1997
+ "block_sizes": [
1998
+ 2
1999
+ ],
2000
+ "loop_orders": [
2001
+ [
2002
+ 1,
2003
+ 2,
2004
+ 0
2005
+ ]
2006
+ ],
2007
+ "l2_groupings": [
2008
+ 16
2009
+ ],
2010
+ "range_unroll_factors": [
2011
+ 1
2012
+ ],
2013
+ "range_warp_specializes": [],
2014
+ "range_multi_buffers": [
2015
+ false
2016
+ ],
2017
+ "range_flattens": [
2018
+ false
2019
+ ],
2020
+ "load_eviction_policies": [
2021
+ "",
2022
+ "",
2023
+ "",
2024
+ "last",
2025
+ "",
2026
+ "last",
2027
+ "first",
2028
+ ""
2029
+ ],
2030
+ "num_warps": 1,
2031
+ "num_stages": 4,
2032
+ "indexing": [
2033
+ "pointer",
2034
+ "tensor_descriptor",
2035
+ "pointer",
2036
+ "pointer",
2037
+ "pointer",
2038
+ "tensor_descriptor",
2039
+ "pointer",
2040
+ "pointer",
2041
+ "pointer",
2042
+ "pointer",
2043
+ "pointer"
2044
+ ],
2045
+ "atomic_indexing": [],
2046
+ "pid_type": "persistent_blocked",
2047
+ "num_sm_multiplier": 64,
2048
+ "maxnreg": 256
2049
+ }
2050
+ },
2051
+ {
2052
+ "key": {
2053
+ "q_heads": 32,
2054
+ "kv_heads": 8,
2055
+ "num_tokens": 2048
2056
+ },
2057
+ "config": {
2058
+ "block_sizes": [
2059
+ 2
2060
+ ],
2061
+ "loop_orders": [
2062
+ [
2063
+ 2,
2064
+ 1,
2065
+ 0
2066
+ ]
2067
+ ],
2068
+ "l2_groupings": [
2069
+ 16
2070
+ ],
2071
+ "range_unroll_factors": [
2072
+ 1
2073
+ ],
2074
+ "range_warp_specializes": [],
2075
+ "range_multi_buffers": [
2076
+ null
2077
+ ],
2078
+ "range_flattens": [
2079
+ true
2080
+ ],
2081
+ "load_eviction_policies": [
2082
+ "",
2083
+ "first",
2084
+ "first",
2085
+ "last",
2086
+ "",
2087
+ "",
2088
+ "last",
2089
+ "first"
2090
+ ],
2091
+ "num_warps": 1,
2092
+ "num_stages": 1,
2093
+ "indexing": [
2094
+ "pointer",
2095
+ "tensor_descriptor",
2096
+ "tensor_descriptor",
2097
+ "tensor_descriptor",
2098
+ "pointer",
2099
+ "tensor_descriptor",
2100
+ "pointer",
2101
+ "pointer",
2102
+ "pointer",
2103
+ "pointer",
2104
+ "tensor_descriptor"
2105
+ ],
2106
+ "atomic_indexing": [],
2107
+ "pid_type": "persistent_blocked",
2108
+ "num_sm_multiplier": 64,
2109
+ "maxnreg": 128
2110
+ }
2111
+ },
2112
+ {
2113
+ "key": {
2114
+ "q_heads": 64,
2115
+ "kv_heads": 8,
2116
+ "num_tokens": 2048
2117
+ },
2118
+ "config": {
2119
+ "block_sizes": [
2120
+ 2
2121
+ ],
2122
+ "loop_orders": [
2123
+ [
2124
+ 2,
2125
+ 1,
2126
+ 0
2127
+ ]
2128
+ ],
2129
+ "l2_groupings": [
2130
+ 16
2131
+ ],
2132
+ "range_unroll_factors": [
2133
+ 1
2134
+ ],
2135
+ "range_warp_specializes": [],
2136
+ "range_multi_buffers": [
2137
+ null
2138
+ ],
2139
+ "range_flattens": [
2140
+ true
2141
+ ],
2142
+ "load_eviction_policies": [
2143
+ "",
2144
+ "first",
2145
+ "first",
2146
+ "last",
2147
+ "",
2148
+ "",
2149
+ "last",
2150
+ "first"
2151
+ ],
2152
+ "num_warps": 1,
2153
+ "num_stages": 1,
2154
+ "indexing": [
2155
+ "pointer",
2156
+ "tensor_descriptor",
2157
+ "tensor_descriptor",
2158
+ "tensor_descriptor",
2159
+ "pointer",
2160
+ "tensor_descriptor",
2161
+ "pointer",
2162
+ "pointer",
2163
+ "pointer",
2164
+ "pointer",
2165
+ "tensor_descriptor"
2166
+ ],
2167
+ "atomic_indexing": [],
2168
+ "pid_type": "persistent_blocked",
2169
+ "num_sm_multiplier": 64,
2170
+ "maxnreg": 128
2171
+ }
2172
+ },
2173
+ {
2174
+ "key": {
2175
+ "q_heads": 16,
2176
+ "kv_heads": 8,
2177
+ "num_tokens": 4096
2178
+ },
2179
+ "config": {
2180
+ "block_sizes": [
2181
+ 2
2182
+ ],
2183
+ "loop_orders": [
2184
+ [
2185
+ 1,
2186
+ 2,
2187
+ 0
2188
+ ]
2189
+ ],
2190
+ "l2_groupings": [
2191
+ 16
2192
+ ],
2193
+ "range_unroll_factors": [
2194
+ 1
2195
+ ],
2196
+ "range_warp_specializes": [],
2197
+ "range_multi_buffers": [
2198
+ false
2199
+ ],
2200
+ "range_flattens": [
2201
+ false
2202
+ ],
2203
+ "load_eviction_policies": [
2204
+ "",
2205
+ "",
2206
+ "",
2207
+ "last",
2208
+ "",
2209
+ "last",
2210
+ "first",
2211
+ ""
2212
+ ],
2213
+ "num_warps": 1,
2214
+ "num_stages": 4,
2215
+ "indexing": [
2216
+ "pointer",
2217
+ "tensor_descriptor",
2218
+ "pointer",
2219
+ "pointer",
2220
+ "pointer",
2221
+ "tensor_descriptor",
2222
+ "pointer",
2223
+ "pointer",
2224
+ "pointer",
2225
+ "pointer",
2226
+ "pointer"
2227
+ ],
2228
+ "atomic_indexing": [],
2229
+ "pid_type": "persistent_blocked",
2230
+ "num_sm_multiplier": 64,
2231
+ "maxnreg": 256
2232
+ }
2233
+ },
2234
+ {
2235
+ "key": {
2236
+ "q_heads": 32,
2237
+ "kv_heads": 8,
2238
+ "num_tokens": 4096
2239
+ },
2240
+ "config": {
2241
+ "block_sizes": [
2242
+ 2
2243
+ ],
2244
+ "loop_orders": [
2245
+ [
2246
+ 2,
2247
+ 1,
2248
+ 0
2249
+ ]
2250
+ ],
2251
+ "l2_groupings": [
2252
+ 2
2253
+ ],
2254
+ "range_unroll_factors": [
2255
+ 1
2256
+ ],
2257
+ "range_warp_specializes": [],
2258
+ "range_multi_buffers": [
2259
+ false
2260
+ ],
2261
+ "range_flattens": [
2262
+ null
2263
+ ],
2264
+ "load_eviction_policies": [
2265
+ "",
2266
+ "first",
2267
+ "last",
2268
+ "",
2269
+ "",
2270
+ "",
2271
+ "",
2272
+ "first"
2273
+ ],
2274
+ "num_warps": 1,
2275
+ "num_stages": 2,
2276
+ "indexing": [
2277
+ "pointer",
2278
+ "pointer",
2279
+ "pointer",
2280
+ "tensor_descriptor",
2281
+ "tensor_descriptor",
2282
+ "pointer",
2283
+ "pointer",
2284
+ "tensor_descriptor",
2285
+ "tensor_descriptor",
2286
+ "tensor_descriptor",
2287
+ "tensor_descriptor"
2288
+ ],
2289
+ "atomic_indexing": [],
2290
+ "pid_type": "persistent_blocked",
2291
+ "num_sm_multiplier": 128,
2292
+ "maxnreg": 64
2293
+ }
2294
+ },
2295
+ {
2296
+ "key": {
2297
+ "q_heads": 64,
2298
+ "kv_heads": 8,
2299
+ "num_tokens": 4096
2300
+ },
2301
+ "config": {
2302
+ "block_sizes": [
2303
+ 2
2304
+ ],
2305
+ "loop_orders": [
2306
+ [
2307
+ 2,
2308
+ 1,
2309
+ 0
2310
+ ]
2311
+ ],
2312
+ "l2_groupings": [
2313
+ 2
2314
+ ],
2315
+ "range_unroll_factors": [
2316
+ 1
2317
+ ],
2318
+ "range_warp_specializes": [],
2319
+ "range_multi_buffers": [
2320
+ false
2321
+ ],
2322
+ "range_flattens": [
2323
+ null
2324
+ ],
2325
+ "load_eviction_policies": [
2326
+ "",
2327
+ "first",
2328
+ "last",
2329
+ "",
2330
+ "",
2331
+ "",
2332
+ "",
2333
+ "first"
2334
+ ],
2335
+ "num_warps": 1,
2336
+ "num_stages": 2,
2337
+ "indexing": [
2338
+ "pointer",
2339
+ "pointer",
2340
+ "pointer",
2341
+ "tensor_descriptor",
2342
+ "tensor_descriptor",
2343
+ "pointer",
2344
+ "pointer",
2345
+ "tensor_descriptor",
2346
+ "tensor_descriptor",
2347
+ "tensor_descriptor",
2348
+ "tensor_descriptor"
2349
+ ],
2350
+ "atomic_indexing": [],
2351
+ "pid_type": "persistent_blocked",
2352
+ "num_sm_multiplier": 128,
2353
+ "maxnreg": 64
2354
+ }
2355
+ },
2356
+ {
2357
+ "key": {
2358
+ "q_heads": 16,
2359
+ "kv_heads": 8,
2360
+ "num_tokens": 8192
2361
+ },
2362
+ "config": {
2363
+ "block_sizes": [
2364
+ 2
2365
+ ],
2366
+ "loop_orders": [
2367
+ [
2368
+ 2,
2369
+ 1,
2370
+ 0
2371
+ ]
2372
+ ],
2373
+ "l2_groupings": [
2374
+ 2
2375
+ ],
2376
+ "range_unroll_factors": [
2377
+ 1
2378
+ ],
2379
+ "range_warp_specializes": [],
2380
+ "range_multi_buffers": [
2381
+ false
2382
+ ],
2383
+ "range_flattens": [
2384
+ null
2385
+ ],
2386
+ "load_eviction_policies": [
2387
+ "",
2388
+ "first",
2389
+ "last",
2390
+ "",
2391
+ "",
2392
+ "",
2393
+ "",
2394
+ "first"
2395
+ ],
2396
+ "num_warps": 1,
2397
+ "num_stages": 2,
2398
+ "indexing": [
2399
+ "pointer",
2400
+ "pointer",
2401
+ "pointer",
2402
+ "tensor_descriptor",
2403
+ "tensor_descriptor",
2404
+ "pointer",
2405
+ "pointer",
2406
+ "tensor_descriptor",
2407
+ "tensor_descriptor",
2408
+ "tensor_descriptor",
2409
+ "tensor_descriptor"
2410
+ ],
2411
+ "atomic_indexing": [],
2412
+ "pid_type": "persistent_blocked",
2413
+ "num_sm_multiplier": 128,
2414
+ "maxnreg": 64
2415
+ }
2416
+ },
2417
+ {
2418
+ "key": {
2419
+ "q_heads": 32,
2420
+ "kv_heads": 8,
2421
+ "num_tokens": 8192
2422
+ },
2423
+ "config": {
2424
+ "block_sizes": [
2425
+ 2
2426
+ ],
2427
+ "loop_orders": [
2428
+ [
2429
+ 2,
2430
+ 1,
2431
+ 0
2432
+ ]
2433
+ ],
2434
+ "l2_groupings": [
2435
+ 2
2436
+ ],
2437
+ "range_unroll_factors": [
2438
+ 1
2439
+ ],
2440
+ "range_warp_specializes": [],
2441
+ "range_multi_buffers": [
2442
+ false
2443
+ ],
2444
+ "range_flattens": [
2445
+ null
2446
+ ],
2447
+ "load_eviction_policies": [
2448
+ "",
2449
+ "first",
2450
+ "last",
2451
+ "",
2452
+ "",
2453
+ "",
2454
+ "",
2455
+ "first"
2456
+ ],
2457
+ "num_warps": 1,
2458
+ "num_stages": 2,
2459
+ "indexing": [
2460
+ "pointer",
2461
+ "pointer",
2462
+ "pointer",
2463
+ "tensor_descriptor",
2464
+ "tensor_descriptor",
2465
+ "pointer",
2466
+ "pointer",
2467
+ "tensor_descriptor",
2468
+ "tensor_descriptor",
2469
+ "tensor_descriptor",
2470
+ "tensor_descriptor"
2471
+ ],
2472
+ "atomic_indexing": [],
2473
+ "pid_type": "persistent_blocked",
2474
+ "num_sm_multiplier": 128,
2475
+ "maxnreg": 64
2476
+ }
2477
+ },
2478
+ {
2479
+ "key": {
2480
+ "q_heads": 64,
2481
+ "kv_heads": 8,
2482
+ "num_tokens": 8192
2483
+ },
2484
+ "config": {
2485
+ "block_sizes": [
2486
+ 2
2487
+ ],
2488
+ "loop_orders": [
2489
+ [
2490
+ 2,
2491
+ 1,
2492
+ 0
2493
+ ]
2494
+ ],
2495
+ "l2_groupings": [
2496
+ 2
2497
+ ],
2498
+ "range_unroll_factors": [
2499
+ 1
2500
+ ],
2501
+ "range_warp_specializes": [],
2502
+ "range_multi_buffers": [
2503
+ false
2504
+ ],
2505
+ "range_flattens": [
2506
+ null
2507
+ ],
2508
+ "load_eviction_policies": [
2509
+ "",
2510
+ "first",
2511
+ "last",
2512
+ "",
2513
+ "",
2514
+ "",
2515
+ "",
2516
+ "first"
2517
+ ],
2518
+ "num_warps": 1,
2519
+ "num_stages": 2,
2520
+ "indexing": [
2521
+ "pointer",
2522
+ "pointer",
2523
+ "pointer",
2524
+ "tensor_descriptor",
2525
+ "tensor_descriptor",
2526
+ "pointer",
2527
+ "pointer",
2528
+ "tensor_descriptor",
2529
+ "tensor_descriptor",
2530
+ "tensor_descriptor",
2531
+ "tensor_descriptor"
2532
+ ],
2533
+ "atomic_indexing": [],
2534
+ "pid_type": "persistent_blocked",
2535
+ "num_sm_multiplier": 128,
2536
+ "maxnreg": 64
2537
+ }
2538
+ },
2539
+ {
2540
+ "key": {
2541
+ "q_heads": 16,
2542
+ "kv_heads": 8,
2543
+ "num_tokens": 16384
2544
+ },
2545
+ "config": {
2546
+ "block_sizes": [
2547
+ 2
2548
+ ],
2549
+ "loop_orders": [
2550
+ [
2551
+ 1,
2552
+ 2,
2553
+ 0
2554
+ ]
2555
+ ],
2556
+ "l2_groupings": [
2557
+ 1
2558
+ ],
2559
+ "range_unroll_factors": [
2560
+ 0
2561
+ ],
2562
+ "range_multi_buffers": [
2563
+ null
2564
+ ],
2565
+ "range_flattens": [
2566
+ true
2567
+ ],
2568
+ "load_eviction_policies": [
2569
+ "",
2570
+ "last",
2571
+ "last",
2572
+ "",
2573
+ "last",
2574
+ "last",
2575
+ "",
2576
+ ""
2577
+ ],
2578
+ "num_warps": 1,
2579
+ "num_stages": 7,
2580
+ "indexing": [
2581
+ "pointer",
2582
+ "pointer",
2583
+ "pointer",
2584
+ "pointer",
2585
+ "pointer",
2586
+ "pointer",
2587
+ "pointer",
2588
+ "pointer",
2589
+ "pointer",
2590
+ "pointer",
2591
+ "pointer"
2592
+ ],
2593
+ "pid_type": "persistent_blocked",
2594
+ "num_sm_multiplier": 128,
2595
+ "maxnreg": 64,
2596
+ "atomic_indexing": [],
2597
+ "range_warp_specializes": []
2598
+ }
2599
+ },
2600
+ {
2601
+ "key": {
2602
+ "q_heads": 32,
2603
+ "kv_heads": 8,
2604
+ "num_tokens": 16384
2605
+ },
2606
+ "config": {
2607
+ "block_sizes": [
2608
+ 2
2609
+ ],
2610
+ "loop_orders": [
2611
+ [
2612
+ 2,
2613
+ 1,
2614
+ 0
2615
+ ]
2616
+ ],
2617
+ "l2_groupings": [
2618
+ 2
2619
+ ],
2620
+ "range_unroll_factors": [
2621
+ 1
2622
+ ],
2623
+ "range_warp_specializes": [],
2624
+ "range_multi_buffers": [
2625
+ false
2626
+ ],
2627
+ "range_flattens": [
2628
+ null
2629
+ ],
2630
+ "load_eviction_policies": [
2631
+ "",
2632
+ "first",
2633
+ "last",
2634
+ "",
2635
+ "",
2636
+ "",
2637
+ "",
2638
+ "first"
2639
+ ],
2640
+ "num_warps": 1,
2641
+ "num_stages": 2,
2642
+ "indexing": [
2643
+ "pointer",
2644
+ "pointer",
2645
+ "pointer",
2646
+ "tensor_descriptor",
2647
+ "tensor_descriptor",
2648
+ "pointer",
2649
+ "pointer",
2650
+ "tensor_descriptor",
2651
+ "tensor_descriptor",
2652
+ "tensor_descriptor",
2653
+ "tensor_descriptor"
2654
+ ],
2655
+ "atomic_indexing": [],
2656
+ "pid_type": "persistent_blocked",
2657
+ "num_sm_multiplier": 128,
2658
+ "maxnreg": 64
2659
+ }
2660
+ },
2661
+ {
2662
+ "key": {
2663
+ "q_heads": 64,
2664
+ "kv_heads": 8,
2665
+ "num_tokens": 16384
2666
+ },
2667
+ "config": {
2668
+ "block_sizes": [
2669
+ 2
2670
+ ],
2671
+ "loop_orders": [
2672
+ [
2673
+ 2,
2674
+ 1,
2675
+ 0
2676
+ ]
2677
+ ],
2678
+ "l2_groupings": [
2679
+ 2
2680
+ ],
2681
+ "range_unroll_factors": [
2682
+ 1
2683
+ ],
2684
+ "range_warp_specializes": [],
2685
+ "range_multi_buffers": [
2686
+ false
2687
+ ],
2688
+ "range_flattens": [
2689
+ null
2690
+ ],
2691
+ "load_eviction_policies": [
2692
+ "",
2693
+ "first",
2694
+ "last",
2695
+ "",
2696
+ "",
2697
+ "",
2698
+ "",
2699
+ "first"
2700
+ ],
2701
+ "num_warps": 1,
2702
+ "num_stages": 2,
2703
+ "indexing": [
2704
+ "pointer",
2705
+ "pointer",
2706
+ "pointer",
2707
+ "tensor_descriptor",
2708
+ "tensor_descriptor",
2709
+ "pointer",
2710
+ "pointer",
2711
+ "tensor_descriptor",
2712
+ "tensor_descriptor",
2713
+ "tensor_descriptor",
2714
+ "tensor_descriptor"
2715
+ ],
2716
+ "atomic_indexing": [],
2717
+ "pid_type": "persistent_blocked",
2718
+ "num_sm_multiplier": 128,
2719
+ "maxnreg": 64
2720
+ }
2721
+ }
2722
+]
\ No newline at end of file
vllm/kernels/helion/ops/fused_qk_norm_rope.py
new
+316
@@ -0,0 +1,316 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from itertools import product
5
+from typing import Any
6
+
7
+import torch
8
+
9
+from vllm.kernels.helion.case_key import CaseKey
10
+from vllm.logger import init_logger
11
+from vllm.utils.import_utils import has_helion
12
+
13
+if not has_helion():
14
+ raise ImportError(
15
+ "Helion kernel requires helion to be installed. "
16
+ "Install it with: pip install helion"
17
+ )
18
+
19
+import helion
20
+import helion.language as hl
21
+
22
+from vllm import ir
23
+from vllm.kernels.helion.register import register_kernel
24
+from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
25
+
26
+logger = init_logger(__name__)
27
+
28
+
29
+def _compute_cos_sin_cache(
30
+ max_position_embeddings, rotary_dim, device="cuda", dtype=torch.float
31
+):
32
+ inv_freq = 1.0 / (
33
+ 10000
34
+ ** (torch.arange(0, rotary_dim, 2, device=device, dtype=dtype) / rotary_dim)
35
+ )
36
+
37
+ t = torch.arange(max_position_embeddings, device=device, dtype=dtype)
38
+
39
+ freqs = torch.einsum("i,j -> ij", t, inv_freq)
40
+ cos = freqs.cos()
41
+ sin = freqs.sin()
42
+ cache = torch.cat((cos, sin), dim=-1)
43
+ return cache
44
+
45
+
46
+def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]:
47
+ # TODO(xiaohongchen1991): it is difficult for kernel author to cover
48
+ # all input property combination. Currently, dtypes are fixed. We need
49
+ # optimization to bucket/skip some combinations
50
+ num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
51
+ num_heads_pair = [
52
+ (16, 8),
53
+ (32, 8),
54
+ (64, 8),
55
+ ]
56
+ head_dim = 128
57
+ in_dtype: torch.dtype = torch.bfloat16
58
+ rotary_ratio = 1.0
59
+ is_neox = True
60
+ eps = 1e-6
61
+ device = "cuda"
62
+ inputs = {}
63
+
64
+ for num_tokens, (num_q_heads, num_kv_heads) in product(
65
+ num_tokens_list, num_heads_pair
66
+ ):
67
+ total_dim = (num_q_heads + 2 * num_kv_heads) * head_dim
68
+ qkv = torch.empty(
69
+ num_tokens, total_dim, dtype=in_dtype, device=device
70
+ ).uniform_(-0.1, 0.1)
71
+ positions = torch.arange(num_tokens, dtype=torch.long, device=device)
72
+ q_weight = torch.empty(head_dim, dtype=in_dtype, device=device).uniform_(
73
+ 0.8, 1.2
74
+ )
75
+ k_weight = torch.empty(head_dim, dtype=in_dtype, device=device).uniform_(
76
+ 0.8, 1.2
77
+ )
78
+ rotary_dim = int(head_dim * rotary_ratio)
79
+ cos_sin_cache = _compute_cos_sin_cache(40960, rotary_dim)
80
+ cos_sin_cache = cos_sin_cache.to(in_dtype)
81
+
82
+ config_key = CaseKey(
83
+ {
84
+ "q_heads": num_q_heads,
85
+ "kv_heads": num_kv_heads,
86
+ "num_tokens": num_tokens,
87
+ }
88
+ )
89
+ inputs[config_key] = (
90
+ qkv,
91
+ num_q_heads,
92
+ num_kv_heads,
93
+ num_kv_heads,
94
+ head_dim,
95
+ eps,
96
+ q_weight,
97
+ k_weight,
98
+ cos_sin_cache,
99
+ is_neox,
100
+ positions.view(-1),
101
+ )
102
+
103
+ return inputs
104
+
105
+
106
+_pick_cache: dict[tuple[int, int, int], CaseKey | None] = {}
107
+
108
+
109
+def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None:
110
+ """Pick the best pre-tuned config for the given input shape.
111
+
112
+ Selection strategy:
113
+ 1. Find the closest q_heads among available configs
114
+ (exact match preferred).
115
+ 2. Find the closest kv_heads among available configs
116
+ (exact match preferred).
117
+ 3. Among the num_tokens values tuned for that q_heads and q_heads, pick
118
+ the smallest num_tokens >= the input's num_tokens. If the input is
119
+ larger than all available num_tokens, fall back to the largest.
120
+ """
121
+
122
+ if not config_keys:
123
+ return None
124
+
125
+ qkv, q_heads, kv_heads, *_ = args
126
+ num_tokens = qkv.shape[0]
127
+
128
+ cache_key = (num_tokens, q_heads, kv_heads)
129
+ cached = _pick_cache.get(cache_key)
130
+ if cached is not None:
131
+ return cached
132
+
133
+ configs: dict[int, dict[int, list[int]]] = {}
134
+ for key in config_keys:
135
+ if key.is_default():
136
+ continue
137
+ configs.setdefault(key["q_heads"], {}).setdefault(key["kv_heads"], []).append(
138
+ key["num_tokens"]
139
+ )
140
+
141
+ if not configs:
142
+ return None
143
+
144
+ best_q_heads = min(configs, key=lambda s: abs(s - q_heads))
145
+ best_kv_heads = min(configs[best_q_heads], key=lambda s: abs(s - kv_heads))
146
+ available_num_tokens = sorted(configs[best_q_heads][best_kv_heads])
147
+ best_num_tokens = next(
148
+ (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1]
149
+ )
150
+
151
+ result = CaseKey(
152
+ {
153
+ "q_heads": best_q_heads,
154
+ "kv_heads": best_kv_heads,
155
+ "num_tokens": best_num_tokens,
156
+ }
157
+ )
158
+ _pick_cache[cache_key] = result
159
+ return result
160
+
161
+
162
+def fake_impl(
163
+ qkv: torch.Tensor, # [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]
164
+ num_heads_q: int,
165
+ num_heads_k: int,
166
+ num_heads_v: int,
167
+ head_dim: int,
168
+ eps: float,
169
+ q_weight: torch.Tensor,
170
+ k_weight: torch.Tensor,
171
+ cos_sin_cache: torch.Tensor, # [max_position, rotary_dim]
172
+ is_neox: bool,
173
+ position_ids: torch.Tensor, # [num_tokens],
174
+ forced_token_heads_per_warp: int = -1, # dummy
175
+) -> None:
176
+ return
177
+
178
+
179
+def baseline(
180
+ qkv: torch.Tensor, # [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]
181
+ num_heads_q: int,
182
+ num_heads_k: int,
183
+ num_heads_v: int,
184
+ head_dim: int,
185
+ eps: float,
186
+ q_weight: torch.Tensor,
187
+ k_weight: torch.Tensor,
188
+ cos_sin_cache: torch.Tensor, # [max_position, rotary_dim]
189
+ is_neox: bool,
190
+ position_ids: torch.Tensor, # [num_tokens],
191
+ forced_token_heads_per_warp: int = -1, # dummy
192
+) -> None:
193
+ q_size = num_heads_q * head_dim
194
+ kv_size = num_heads_k * head_dim
195
+
196
+ q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
197
+
198
+ q_by_head = q.view(*q.shape[:-1], q.shape[-1] // head_dim, head_dim)
199
+ q_by_head = ir.ops.rms_norm(q_by_head, q_weight, eps)
200
+ q = q_by_head.view(q.shape)
201
+
202
+ k_by_head = k.view(*k.shape[:-1], k.shape[-1] // head_dim, head_dim)
203
+ k_by_head = ir.ops.rms_norm(k_by_head, k_weight, eps)
204
+ k = k_by_head.view(k.shape)
205
+
206
+ q, k = RotaryEmbedding.forward_static(
207
+ position_ids, q, k, head_dim, cos_sin_cache.shape[1], cos_sin_cache, is_neox
208
+ )
209
+ qkv[:, :q_size].copy_(q)
210
+ qkv[:, q_size : q_size + kv_size].copy_(k)
211
+
212
+
213
+# Overwrite autotune_baseline_atol and autotune_baseline_rtol
214
+# if too many configs failed due to baseline check during autotuning
215
+@register_kernel(
216
+ mutates_args=["qkv"],
217
+ config_picker=pick_config,
218
+ input_generator=generate_inputs,
219
+ fake_impl=fake_impl,
220
+ helion_settings=helion.Settings(
221
+ autotune_baseline_fn=baseline,
222
+ autotune_baseline_atol=5e-2,
223
+ autotune_baseline_rtol=5e-2,
224
+ ignore_warnings=[helion.exc.TensorOperationInWrapper],
225
+ ),
226
+) # type: ignore[misc]
227
+def fused_qk_norm_rope(
228
+ qkv: torch.Tensor, # [num_tokens, (num_heads_q+num_heads_k+num_heads_v)*head_dim]
229
+ num_heads_q: int,
230
+ num_heads_k: int,
231
+ num_heads_v: int,
232
+ head_dim: int,
233
+ eps: float,
234
+ q_weight: torch.Tensor,
235
+ k_weight: torch.Tensor,
236
+ cos_sin_cache: torch.Tensor, # [max_position, rotary_dim]
237
+ is_neox: bool,
238
+ position_ids: torch.Tensor, # [num_tokens],
239
+ forced_token_heads_per_warp: int = -1, # dummy
240
+) -> None:
241
+ assert qkv.ndim == 2
242
+ num_tokens = qkv.shape[0]
243
+ total_heads = num_heads_q + num_heads_k + num_heads_v
244
+ assert qkv.shape[1] == total_heads * head_dim
245
+ hl.specialize(qkv.shape[1])
246
+
247
+ assert cos_sin_cache.ndim == 2
248
+ max_position, rotary_dim = cos_sin_cache.shape
249
+ hl.specialize(max_position)
250
+ hl.specialize(rotary_dim)
251
+ assert rotary_dim % 2 == 0
252
+ assert rotary_dim <= head_dim
253
+ embed_dim = rotary_dim // 2
254
+
255
+ hl.specialize(num_heads_q)
256
+ hl.specialize(num_heads_k)
257
+ hl.specialize(num_heads_v)
258
+ hl.specialize(head_dim)
259
+
260
+ assert position_ids.ndim == 1 and position_ids.shape[0] == num_tokens
261
+ hl.specialize(position_ids.shape[0])
262
+
263
+ assert q_weight.ndim == 1 and q_weight.shape[0] == head_dim
264
+ hl.specialize(q_weight.shape[0])
265
+ assert k_weight.ndim == 1 and k_weight.shape[0] == head_dim
266
+ hl.specialize(k_weight.shape[0])
267
+
268
+ assert qkv.dtype == q_weight.dtype and q_weight.dtype == k_weight.dtype
269
+ assert position_ids.dtype == torch.int64
270
+
271
+ assert qkv.is_contiguous()
272
+ assert position_ids.is_contiguous()
273
+ assert q_weight.is_contiguous()
274
+ assert k_weight.is_contiguous()
275
+ assert cos_sin_cache.is_contiguous()
276
+
277
+ qk_heads = num_heads_q + num_heads_k
278
+
279
+ qkv = qkv.view(num_tokens, -1, head_dim)
280
+
281
+ for tile_m, tile_gn, tile_n in hl.tile(
282
+ [num_tokens, qk_heads, head_dim], block_size=[1, None, head_dim]
283
+ ):
284
+ x_blk = qkv[tile_m, tile_gn, tile_n].to(dtype=torch.float32)
285
+
286
+ rms = x_blk.pow(2).sum(dim=-1)
287
+ rms = torch.rsqrt(rms * (1.0 / head_dim) + eps)
288
+
289
+ use_q_weight = (tile_gn.index < num_heads_q)[None, :, None]
290
+ w_blk = torch.where(
291
+ use_q_weight, q_weight[None, None, tile_n], k_weight[None, None, tile_n]
292
+ )
293
+
294
+ x_blk = (x_blk * rms[:, :, None]).to(qkv.dtype) * w_blk
295
+
296
+ qkv[tile_m, tile_gn, tile_n] = x_blk
297
+
298
+ pos_id = position_ids[tile_m]
299
+ cos_blk = cos_sin_cache[pos_id, hl.arange(embed_dim)]
300
+ sin_blk = cos_sin_cache[pos_id, hl.arange(embed_dim) + embed_dim]
301
+
302
+ if is_neox:
303
+ x1_offset = hl.arange(embed_dim)
304
+ x2_offset = x1_offset + embed_dim
305
+ else:
306
+ x1_offset = hl.arange(embed_dim) * 2
307
+ x2_offset = x1_offset + 1
308
+
309
+ x1_blk = qkv[tile_m, tile_gn, x1_offset]
310
+ x2_blk = qkv[tile_m, tile_gn, x2_offset]
311
+
312
+ o1_blk = x1_blk * cos_blk[:, None, :] - x2_blk * sin_blk[:, None, :]
313
+ o2_blk = x2_blk * cos_blk[:, None, :] + x1_blk * sin_blk[:, None, :]
314
+
315
+ qkv[tile_m, tile_gn, x1_offset] = o1_blk
316
+ qkv[tile_m, tile_gn, x2_offset] = o2_blk