A high-throughput and memory-efficient inference and serving engine for LLMs
[Kernel][Helion][1/N] Add Helion kernel for dynamic_per_token_scaled_fp8_quant (#33790)
Signed-off-by: Sean Chen <seachen@redhat.com> Co-authored-by: Yanan Cao <gmagogsfm@gmail.com>
Xiaohong (Sean) Chen committed
Jun 12, 2026 at 12:50 UTC
272c16953eac7c46db7719d284d8a0ff19e63446
4 files changed
+7540
tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.py
new
+165
@@ -0,0 +1,165 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the dynamic_per_token_scaled_fp8_quant helion kernel
4
+
5
+Run `pytest tests/kernels/helion/test_dynamic_per_token_scaled_fp8_quant.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 tests.kernels.quant_utils import FP8_DTYPE
16
+from vllm.kernels.helion.case_key import CaseKey
17
+from vllm.kernels.helion.config_manager import ConfigManager
18
+from vllm.kernels.helion.ops.dynamic_per_token_scaled_fp8_quant import (
19
+ _pick_cache,
20
+ baseline,
21
+ dynamic_per_token_scaled_fp8_quant,
22
+ pick_config,
23
+)
24
+from vllm.platforms import current_platform
25
+from vllm.utils.import_utils import has_helion
26
+from vllm.utils.torch_utils import set_random_seed
27
+
28
+if not has_helion():
29
+ pytest.skip(
30
+ "Helion is not installed. Install with: pip install vllm[helion]",
31
+ allow_module_level=True,
32
+ )
33
+
34
+
35
+def _generate_fake_input(num_tokens: int, hidden_size: int) -> tuple[Any, ...]:
36
+ with FakeTensorMode():
37
+ input = torch.randn(
38
+ num_tokens, hidden_size, device="cuda", dtype=torch.bfloat16
39
+ )
40
+ result = torch.empty(
41
+ input.shape, device=input.device, dtype=current_platform.fp8_dtype()
42
+ )
43
+ scale = torch.empty((num_tokens, 1), device=input.device, dtype=torch.float32)
44
+ scale_ub = torch.mean(input).to(torch.float32)
45
+ args = (result, input, scale, scale_ub)
46
+ return args
47
+
48
+
49
+@pytest.fixture(autouse=True)
50
+def reset_config_manager_singleton():
51
+ ConfigManager.reset_instance()
52
+ ConfigManager()
53
+ yield
54
+ ConfigManager.reset_instance()
55
+
56
+
57
+class TestDynamicPerTokenScaledFp8QuantConfigPicker:
58
+ def setup_method(self):
59
+ _pick_cache.clear()
60
+
61
+ def test_config_picker_exact_match(self):
62
+ config_keys = [
63
+ CaseKey({"hidden_size": 2048, "num_tokens": 16}),
64
+ CaseKey({"hidden_size": 4096, "num_tokens": 16}),
65
+ ]
66
+
67
+ args = _generate_fake_input(16, 4096)
68
+ selected_key = pick_config(args, config_keys)
69
+ assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16})
70
+
71
+ def test_config_picker_closest_match(self):
72
+ config_keys = [
73
+ CaseKey({"hidden_size": 2048, "num_tokens": 16}),
74
+ CaseKey({"hidden_size": 2048, "num_tokens": 32}),
75
+ CaseKey({"hidden_size": 4096, "num_tokens": 16}),
76
+ CaseKey({"hidden_size": 4096, "num_tokens": 32}),
77
+ ]
78
+
79
+ args = _generate_fake_input(20, 3000)
80
+ selected_key = pick_config(args, config_keys)
81
+ assert selected_key == CaseKey({"hidden_size": 2048, "num_tokens": 32})
82
+
83
+ def test_config_picker_no_configs(self):
84
+ config_keys: list[dict] = []
85
+
86
+ args = _generate_fake_input(16, 4096)
87
+ selected_key = pick_config(args, config_keys)
88
+ assert selected_key is None
89
+
90
+ def test_config_picker_fallback_to_largest(self):
91
+ config_keys = [
92
+ CaseKey({"hidden_size": 2048, "num_tokens": 16}),
93
+ CaseKey({"hidden_size": 4096, "num_tokens": 16}),
94
+ ]
95
+
96
+ args = _generate_fake_input(32, 8192)
97
+ selected_key = pick_config(args, config_keys)
98
+ assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16})
99
+
100
+
101
+class TestDynamicPerTokenScaledFp8QuantCorrectness:
102
+ @pytest.mark.parametrize("num_tokens", [1, 7, 4096])
103
+ @pytest.mark.parametrize("hidden_size", [17, 1024, 1025, 1026, 5137, 8193])
104
+ @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float])
105
+ @pytest.mark.parametrize("has_scale_ub", [True, False])
106
+ @pytest.mark.parametrize("seed", [0])
107
+ def test_dynamic_per_token_fp8_quant(
108
+ self,
109
+ num_tokens: int,
110
+ hidden_size: int,
111
+ dtype: torch.dtype,
112
+ has_scale_ub: bool,
113
+ seed: int,
114
+ ) -> None:
115
+ skip_if_platform_unsupported("dynamic_per_token_scaled_fp8_quant")
116
+ set_random_seed(seed)
117
+
118
+ x = (
119
+ torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") + 1e-6
120
+ ) # avoid nans
121
+
122
+ scale_ub = (
123
+ torch.mean(x).to(dtype=torch.float32, device="cuda")
124
+ if has_scale_ub
125
+ else None
126
+ )
127
+
128
+ ref_out = torch.empty(x.shape, device="cuda", dtype=FP8_DTYPE)
129
+ ref_scales = torch.empty((x.shape[0], 1), device="cuda", dtype=torch.float32)
130
+ baseline(ref_out, x, ref_scales, scale_ub)
131
+
132
+ ops_out = torch.empty(x.shape, device="cuda", dtype=FP8_DTYPE)
133
+ ops_scales = torch.empty((x.shape[0], 1), device="cuda", dtype=torch.float32)
134
+ dynamic_per_token_scaled_fp8_quant(ops_out, x, ops_scales, scale_ub)
135
+
136
+ torch.testing.assert_close(ref_scales, ops_scales)
137
+ # allow 1 ULP difference
138
+ assert (
139
+ ref_out.view(torch.uint8).to(torch.int16)
140
+ - ops_out.view(torch.uint8).to(torch.int16)
141
+ ).abs().max() <= 1
142
+
143
+
144
+class TestDynamicPerTokenScaledFp8QuantIntegration:
145
+ def test_kernel_registration_integration(self):
146
+ from vllm.kernels.helion.register import get_registered_kernels
147
+
148
+ registered_kernels = get_registered_kernels()
149
+ assert "dynamic_per_token_scaled_fp8_quant" in registered_kernels
150
+
151
+ kernel_wrapper = registered_kernels["dynamic_per_token_scaled_fp8_quant"]
152
+ assert kernel_wrapper.op_name == "dynamic_per_token_scaled_fp8_quant"
153
+ assert kernel_wrapper._config_picker is not None
154
+ assert kernel_wrapper._mutates_args == ["result", "scale"]
155
+
156
+ def test_fake_impl_functionality(self):
157
+ skip_if_platform_unsupported("dynamic_per_token_scaled_fp8_quant")
158
+ from vllm.kernels.helion.register import get_registered_kernels
159
+
160
+ registered_kernels = get_registered_kernels()
161
+ kernel_wrapper = registered_kernels["dynamic_per_token_scaled_fp8_quant"]
162
+ fake_impl = kernel_wrapper._fake_impl
163
+
164
+ args = _generate_fake_input(16, 4096)
165
+ assert fake_impl(*args) is None
vllm/kernels/helion/configs/dynamic_per_token_scaled_fp8_quant/nvidia_b200.json
new
+2025
@@ -0,0 +1,2025 @@
1
+[
2
+ {
3
+ "key": {
4
+ "hidden_size": 2048,
5
+ "num_tokens": 1
6
+ },
7
+ "config": {
8
+ "block_sizes": [
9
+ 16384,
10
+ 4096
11
+ ],
12
+ "range_unroll_factors": [
13
+ 0,
14
+ 1,
15
+ 2
16
+ ],
17
+ "range_warp_specializes": [],
18
+ "range_num_stages": [],
19
+ "range_multi_buffers": [
20
+ null,
21
+ true,
22
+ false
23
+ ],
24
+ "range_flattens": [
25
+ null,
26
+ true,
27
+ null
28
+ ],
29
+ "load_eviction_policies": [
30
+ "",
31
+ "last",
32
+ "last"
33
+ ],
34
+ "num_warps": 16,
35
+ "num_stages": 6,
36
+ "indexing": [
37
+ "tensor_descriptor",
38
+ "tensor_descriptor",
39
+ "pointer",
40
+ "tensor_descriptor",
41
+ "tensor_descriptor"
42
+ ],
43
+ "atomic_indexing": [],
44
+ "pid_type": "flat"
45
+ }
46
+ },
47
+ {
48
+ "key": {
49
+ "hidden_size": 4096,
50
+ "num_tokens": 1
51
+ },
52
+ "config": {
53
+ "block_sizes": [
54
+ 16384,
55
+ 4096
56
+ ],
57
+ "range_unroll_factors": [
58
+ 0,
59
+ 1,
60
+ 2
61
+ ],
62
+ "range_warp_specializes": [],
63
+ "range_num_stages": [],
64
+ "range_multi_buffers": [
65
+ null,
66
+ true,
67
+ false
68
+ ],
69
+ "range_flattens": [
70
+ null,
71
+ true,
72
+ null
73
+ ],
74
+ "load_eviction_policies": [
75
+ "",
76
+ "last",
77
+ "last"
78
+ ],
79
+ "num_warps": 16,
80
+ "num_stages": 6,
81
+ "indexing": [
82
+ "tensor_descriptor",
83
+ "tensor_descriptor",
84
+ "pointer",
85
+ "tensor_descriptor",
86
+ "tensor_descriptor"
87
+ ],
88
+ "atomic_indexing": [],
89
+ "pid_type": "flat"
90
+ }
91
+ },
92
+ {
93
+ "key": {
94
+ "hidden_size": 5120,
95
+ "num_tokens": 1
96
+ },
97
+ "config": {
98
+ "block_sizes": [
99
+ 16384,
100
+ 4096
101
+ ],
102
+ "range_unroll_factors": [
103
+ 0,
104
+ 1,
105
+ 2
106
+ ],
107
+ "range_warp_specializes": [],
108
+ "range_num_stages": [],
109
+ "range_multi_buffers": [
110
+ null,
111
+ true,
112
+ false
113
+ ],
114
+ "range_flattens": [
115
+ null,
116
+ true,
117
+ null
118
+ ],
119
+ "load_eviction_policies": [
120
+ "",
121
+ "last",
122
+ "last"
123
+ ],
124
+ "num_warps": 16,
125
+ "num_stages": 6,
126
+ "indexing": [
127
+ "tensor_descriptor",
128
+ "tensor_descriptor",
129
+ "pointer",
130
+ "tensor_descriptor",
131
+ "tensor_descriptor"
132
+ ],
133
+ "atomic_indexing": [],
134
+ "pid_type": "flat"
135
+ }
136
+ },
137
+ {
138
+ "key": {
139
+ "hidden_size": 2048,
140
+ "num_tokens": 2
141
+ },
142
+ "config": {
143
+ "block_sizes": [
144
+ 16384,
145
+ 4096
146
+ ],
147
+ "range_unroll_factors": [
148
+ 0,
149
+ 1,
150
+ 2
151
+ ],
152
+ "range_warp_specializes": [],
153
+ "range_num_stages": [],
154
+ "range_multi_buffers": [
155
+ null,
156
+ true,
157
+ false
158
+ ],
159
+ "range_flattens": [
160
+ null,
161
+ true,
162
+ null
163
+ ],
164
+ "load_eviction_policies": [
165
+ "",
166
+ "last",
167
+ "last"
168
+ ],
169
+ "num_warps": 16,
170
+ "num_stages": 6,
171
+ "indexing": [
172
+ "tensor_descriptor",
173
+ "tensor_descriptor",
174
+ "pointer",
175
+ "tensor_descriptor",
176
+ "tensor_descriptor"
177
+ ],
178
+ "atomic_indexing": [],
179
+ "pid_type": "flat"
180
+ }
181
+ },
182
+ {
183
+ "key": {
184
+ "hidden_size": 4096,
185
+ "num_tokens": 2
186
+ },
187
+ "config": {
188
+ "block_sizes": [
189
+ 4096,
190
+ 4096
191
+ ],
192
+ "range_unroll_factors": [
193
+ 0,
194
+ 4,
195
+ 3
196
+ ],
197
+ "range_warp_specializes": [
198
+ null,
199
+ false,
200
+ null
201
+ ],
202
+ "range_num_stages": [],
203
+ "range_multi_buffers": [
204
+ null,
205
+ false,
206
+ true
207
+ ],
208
+ "range_flattens": [
209
+ null,
210
+ true,
211
+ false
212
+ ],
213
+ "load_eviction_policies": [
214
+ "first",
215
+ "first",
216
+ "last"
217
+ ],
218
+ "num_warps": 16,
219
+ "num_stages": 7,
220
+ "indexing": [
221
+ "tensor_descriptor",
222
+ "tensor_descriptor",
223
+ "tensor_descriptor",
224
+ "tensor_descriptor",
225
+ "pointer"
226
+ ],
227
+ "atomic_indexing": [],
228
+ "pid_type": "flat"
229
+ }
230
+ },
231
+ {
232
+ "key": {
233
+ "hidden_size": 5120,
234
+ "num_tokens": 2
235
+ },
236
+ "config": {
237
+ "block_sizes": [
238
+ 8192,
239
+ 1024
240
+ ],
241
+ "range_unroll_factors": [
242
+ 0,
243
+ 3,
244
+ 2
245
+ ],
246
+ "range_warp_specializes": [
247
+ null,
248
+ false,
249
+ false
250
+ ],
251
+ "range_num_stages": [],
252
+ "range_multi_buffers": [
253
+ null,
254
+ true,
255
+ null
256
+ ],
257
+ "range_flattens": [
258
+ null,
259
+ null,
260
+ null
261
+ ],
262
+ "load_eviction_policies": [
263
+ "last",
264
+ "first",
265
+ "last"
266
+ ],
267
+ "num_warps": 16,
268
+ "num_stages": 6,
269
+ "indexing": [
270
+ "pointer",
271
+ "pointer",
272
+ "tensor_descriptor",
273
+ "pointer",
274
+ "tensor_descriptor"
275
+ ],
276
+ "atomic_indexing": [],
277
+ "pid_type": "flat"
278
+ }
279
+ },
280
+ {
281
+ "key": {
282
+ "hidden_size": 2048,
283
+ "num_tokens": 4
284
+ },
285
+ "config": {
286
+ "block_sizes": [
287
+ 2048,
288
+ 1024
289
+ ],
290
+ "range_unroll_factors": [
291
+ 0,
292
+ 4,
293
+ 1
294
+ ],
295
+ "range_warp_specializes": [
296
+ null,
297
+ null,
298
+ true
299
+ ],
300
+ "range_num_stages": [],
301
+ "range_multi_buffers": [
302
+ null,
303
+ false,
304
+ null
305
+ ],
306
+ "range_flattens": [
307
+ null,
308
+ false,
309
+ true
310
+ ],
311
+ "load_eviction_policies": [
312
+ "",
313
+ "last",
314
+ "first"
315
+ ],
316
+ "num_warps": 2,
317
+ "num_stages": 1,
318
+ "indexing": [
319
+ "pointer",
320
+ "tensor_descriptor",
321
+ "tensor_descriptor",
322
+ "tensor_descriptor",
323
+ "pointer"
324
+ ],
325
+ "atomic_indexing": [],
326
+ "pid_type": "flat"
327
+ }
328
+ },
329
+ {
330
+ "key": {
331
+ "hidden_size": 4096,
332
+ "num_tokens": 4
333
+ },
334
+ "config": {
335
+ "block_sizes": [
336
+ 4096,
337
+ 2048
338
+ ],
339
+ "range_unroll_factors": [
340
+ 0,
341
+ 2,
342
+ 4
343
+ ],
344
+ "range_warp_specializes": [
345
+ null,
346
+ null,
347
+ false
348
+ ],
349
+ "range_num_stages": [],
350
+ "range_multi_buffers": [
351
+ null,
352
+ true,
353
+ null
354
+ ],
355
+ "range_flattens": [
356
+ null,
357
+ true,
358
+ null
359
+ ],
360
+ "load_eviction_policies": [
361
+ "first",
362
+ "",
363
+ "first"
364
+ ],
365
+ "num_warps": 8,
366
+ "num_stages": 6,
367
+ "indexing": [
368
+ "pointer",
369
+ "pointer",
370
+ "tensor_descriptor",
371
+ "pointer",
372
+ "pointer"
373
+ ],
374
+ "atomic_indexing": [],
375
+ "pid_type": "flat"
376
+ }
377
+ },
378
+ {
379
+ "key": {
380
+ "hidden_size": 5120,
381
+ "num_tokens": 4
382
+ },
383
+ "config": {
384
+ "block_sizes": [
385
+ 8192,
386
+ 512
387
+ ],
388
+ "range_unroll_factors": [
389
+ 0,
390
+ 2,
391
+ 0
392
+ ],
393
+ "range_warp_specializes": [
394
+ null,
395
+ false,
396
+ true
397
+ ],
398
+ "range_num_stages": [],
399
+ "range_multi_buffers": [
400
+ null,
401
+ null,
402
+ true
403
+ ],
404
+ "range_flattens": [
405
+ null,
406
+ null,
407
+ false
408
+ ],
409
+ "load_eviction_policies": [
410
+ "first",
411
+ "first",
412
+ ""
413
+ ],
414
+ "num_warps": 8,
415
+ "num_stages": 8,
416
+ "indexing": [
417
+ "pointer",
418
+ "pointer",
419
+ "pointer",
420
+ "pointer",
421
+ "tensor_descriptor"
422
+ ],
423
+ "atomic_indexing": [],
424
+ "pid_type": "flat"
425
+ }
426
+ },
427
+ {
428
+ "key": {
429
+ "hidden_size": 2048,
430
+ "num_tokens": 8
431
+ },
432
+ "config": {
433
+ "block_sizes": [
434
+ 2048,
435
+ 1024
436
+ ],
437
+ "range_unroll_factors": [
438
+ 0,
439
+ 4,
440
+ 1
441
+ ],
442
+ "range_warp_specializes": [
443
+ null,
444
+ null,
445
+ true
446
+ ],
447
+ "range_num_stages": [],
448
+ "range_multi_buffers": [
449
+ null,
450
+ false,
451
+ null
452
+ ],
453
+ "range_flattens": [
454
+ null,
455
+ false,
456
+ true
457
+ ],
458
+ "load_eviction_policies": [
459
+ "",
460
+ "last",
461
+ "first"
462
+ ],
463
+ "num_warps": 2,
464
+ "num_stages": 1,
465
+ "indexing": [
466
+ "pointer",
467
+ "tensor_descriptor",
468
+ "tensor_descriptor",
469
+ "tensor_descriptor",
470
+ "pointer"
471
+ ],
472
+ "atomic_indexing": [],
473
+ "pid_type": "flat"
474
+ }
475
+ },
476
+ {
477
+ "key": {
478
+ "hidden_size": 4096,
479
+ "num_tokens": 8
480
+ },
481
+ "config": {
482
+ "block_sizes": [
483
+ 4096,
484
+ 2048
485
+ ],
486
+ "range_unroll_factors": [
487
+ 0,
488
+ 4,
489
+ 1
490
+ ],
491
+ "range_warp_specializes": [
492
+ null,
493
+ null,
494
+ null
495
+ ],
496
+ "range_num_stages": [],
497
+ "range_multi_buffers": [
498
+ null,
499
+ true,
500
+ null
501
+ ],
502
+ "range_flattens": [
503
+ null,
504
+ null,
505
+ true
506
+ ],
507
+ "load_eviction_policies": [
508
+ "first",
509
+ "last",
510
+ "first"
511
+ ],
512
+ "num_warps": 4,
513
+ "num_stages": 3,
514
+ "indexing": [
515
+ "tensor_descriptor",
516
+ "pointer",
517
+ "tensor_descriptor",
518
+ "tensor_descriptor",
519
+ "pointer"
520
+ ],
521
+ "atomic_indexing": [],
522
+ "pid_type": "flat"
523
+ }
524
+ },
525
+ {
526
+ "key": {
527
+ "hidden_size": 5120,
528
+ "num_tokens": 8
529
+ },
530
+ "config": {
531
+ "block_sizes": [
532
+ 8192,
533
+ 1024
534
+ ],
535
+ "range_unroll_factors": [
536
+ 0,
537
+ 3,
538
+ 2
539
+ ],
540
+ "range_warp_specializes": [
541
+ null,
542
+ false,
543
+ false
544
+ ],
545
+ "range_num_stages": [],
546
+ "range_multi_buffers": [
547
+ null,
548
+ true,
549
+ null
550
+ ],
551
+ "range_flattens": [
552
+ null,
553
+ null,
554
+ null
555
+ ],
556
+ "load_eviction_policies": [
557
+ "last",
558
+ "first",
559
+ "last"
560
+ ],
561
+ "num_warps": 16,
562
+ "num_stages": 6,
563
+ "indexing": [
564
+ "pointer",
565
+ "pointer",
566
+ "tensor_descriptor",
567
+ "pointer",
568
+ "tensor_descriptor"
569
+ ],
570
+ "atomic_indexing": [],
571
+ "pid_type": "flat"
572
+ }
573
+ },
574
+ {
575
+ "key": {
576
+ "hidden_size": 2048,
577
+ "num_tokens": 16
578
+ },
579
+ "config": {
580
+ "block_sizes": [
581
+ 2048,
582
+ 1024
583
+ ],
584
+ "range_unroll_factors": [
585
+ 0,
586
+ 4,
587
+ 1
588
+ ],
589
+ "range_warp_specializes": [
590
+ null,
591
+ null,
592
+ true
593
+ ],
594
+ "range_num_stages": [],
595
+ "range_multi_buffers": [
596
+ null,
597
+ false,
598
+ null
599
+ ],
600
+ "range_flattens": [
601
+ null,
602
+ false,
603
+ true
604
+ ],
605
+ "load_eviction_policies": [
606
+ "",
607
+ "last",
608
+ "first"
609
+ ],
610
+ "num_warps": 2,
611
+ "num_stages": 1,
612
+ "indexing": [
613
+ "pointer",
614
+ "tensor_descriptor",
615
+ "tensor_descriptor",
616
+ "tensor_descriptor",
617
+ "pointer"
618
+ ],
619
+ "atomic_indexing": [],
620
+ "pid_type": "flat"
621
+ }
622
+ },
623
+ {
624
+ "key": {
625
+ "hidden_size": 4096,
626
+ "num_tokens": 16
627
+ },
628
+ "config": {
629
+ "block_sizes": [
630
+ 4096,
631
+ 4096
632
+ ],
633
+ "range_unroll_factors": [
634
+ 0,
635
+ 0,
636
+ 4
637
+ ],
638
+ "range_warp_specializes": [
639
+ false,
640
+ null,
641
+ false
642
+ ],
643
+ "range_multi_buffers": [
644
+ true,
645
+ true,
646
+ true
647
+ ],
648
+ "range_flattens": [
649
+ false,
650
+ true,
651
+ true
652
+ ],
653
+ "load_eviction_policies": [
654
+ "last",
655
+ "first",
656
+ "last"
657
+ ],
658
+ "num_warps": 16,
659
+ "num_stages": 2,
660
+ "indexing": [
661
+ "tensor_descriptor",
662
+ "tensor_descriptor",
663
+ "tensor_descriptor",
664
+ "tensor_descriptor",
665
+ "pointer"
666
+ ],
667
+ "atomic_indexing": [],
668
+ "pid_type": "persistent_interleaved",
669
+ "num_sm_multiplier": 1,
670
+ "maxnreg": 128
671
+ }
672
+ },
673
+ {
674
+ "key": {
675
+ "hidden_size": 5120,
676
+ "num_tokens": 16
677
+ },
678
+ "config": {
679
+ "block_sizes": [
680
+ 8192,
681
+ 1024
682
+ ],
683
+ "range_unroll_factors": [
684
+ 3,
685
+ 4,
686
+ 1
687
+ ],
688
+ "range_warp_specializes": [
689
+ false,
690
+ false,
691
+ true
692
+ ],
693
+ "range_multi_buffers": [
694
+ false,
695
+ false,
696
+ null
697
+ ],
698
+ "range_flattens": [
699
+ null,
700
+ null,
701
+ null
702
+ ],
703
+ "load_eviction_policies": [
704
+ "last",
705
+ "",
706
+ "first"
707
+ ],
708
+ "num_warps": 8,
709
+ "num_stages": 8,
710
+ "indexing": [
711
+ "pointer",
712
+ "pointer",
713
+ "tensor_descriptor",
714
+ "tensor_descriptor",
715
+ "pointer"
716
+ ],
717
+ "atomic_indexing": [],
718
+ "pid_type": "persistent_interleaved",
719
+ "num_sm_multiplier": 8,
720
+ "maxnreg": 64
721
+ }
722
+ },
723
+ {
724
+ "key": {
725
+ "hidden_size": 2048,
726
+ "num_tokens": 32
727
+ },
728
+ "config": {
729
+ "block_sizes": [
730
+ 2048,
731
+ 1024
732
+ ],
733
+ "range_unroll_factors": [
734
+ 0,
735
+ 4,
736
+ 1
737
+ ],
738
+ "range_warp_specializes": [
739
+ null,
740
+ null,
741
+ true
742
+ ],
743
+ "range_num_stages": [],
744
+ "range_multi_buffers": [
745
+ null,
746
+ false,
747
+ null
748
+ ],
749
+ "range_flattens": [
750
+ null,
751
+ false,
752
+ true
753
+ ],
754
+ "load_eviction_policies": [
755
+ "",
756
+ "last",
757
+ "first"
758
+ ],
759
+ "num_warps": 2,
760
+ "num_stages": 1,
761
+ "indexing": [
762
+ "pointer",
763
+ "tensor_descriptor",
764
+ "tensor_descriptor",
765
+ "tensor_descriptor",
766
+ "pointer"
767
+ ],
768
+ "atomic_indexing": [],
769
+ "pid_type": "flat"
770
+ }
771
+ },
772
+ {
773
+ "key": {
774
+ "hidden_size": 4096,
775
+ "num_tokens": 32
776
+ },
777
+ "config": {
778
+ "block_sizes": [
779
+ 4096,
780
+ 4096
781
+ ],
782
+ "range_unroll_factors": [
783
+ 0,
784
+ 0,
785
+ 4
786
+ ],
787
+ "range_warp_specializes": [
788
+ false,
789
+ null,
790
+ false
791
+ ],
792
+ "range_multi_buffers": [
793
+ true,
794
+ true,
795
+ true
796
+ ],
797
+ "range_flattens": [
798
+ false,
799
+ true,
800
+ true
801
+ ],
802
+ "load_eviction_policies": [
803
+ "last",
804
+ "first",
805
+ "last"
806
+ ],
807
+ "num_warps": 16,
808
+ "num_stages": 2,
809
+ "indexing": [
810
+ "tensor_descriptor",
811
+ "tensor_descriptor",
812
+ "tensor_descriptor",
813
+ "tensor_descriptor",
814
+ "pointer"
815
+ ],
816
+ "atomic_indexing": [],
817
+ "pid_type": "persistent_interleaved",
818
+ "num_sm_multiplier": 1,
819
+ "maxnreg": 128
820
+ }
821
+ },
822
+ {
823
+ "key": {
824
+ "hidden_size": 5120,
825
+ "num_tokens": 32
826
+ },
827
+ "config": {
828
+ "block_sizes": [
829
+ 8192,
830
+ 1024
831
+ ],
832
+ "range_unroll_factors": [
833
+ 3,
834
+ 4,
835
+ 1
836
+ ],
837
+ "range_warp_specializes": [
838
+ false,
839
+ false,
840
+ true
841
+ ],
842
+ "range_multi_buffers": [
843
+ false,
844
+ false,
845
+ null
846
+ ],
847
+ "range_flattens": [
848
+ null,
849
+ null,
850
+ null
851
+ ],
852
+ "load_eviction_policies": [
853
+ "last",
854
+ "",
855
+ "first"
856
+ ],
857
+ "num_warps": 8,
858
+ "num_stages": 8,
859
+ "indexing": [
860
+ "pointer",
861
+ "pointer",
862
+ "tensor_descriptor",
863
+ "tensor_descriptor",
864
+ "pointer"
865
+ ],
866
+ "atomic_indexing": [],
867
+ "pid_type": "persistent_interleaved",
868
+ "num_sm_multiplier": 8,
869
+ "maxnreg": 64
870
+ }
871
+ },
872
+ {
873
+ "key": {
874
+ "hidden_size": 2048,
875
+ "num_tokens": 64
876
+ },
877
+ "config": {
878
+ "block_sizes": [
879
+ 2048,
880
+ 1024
881
+ ],
882
+ "range_unroll_factors": [
883
+ 0,
884
+ 4,
885
+ 1
886
+ ],
887
+ "range_warp_specializes": [
888
+ null,
889
+ null,
890
+ true
891
+ ],
892
+ "range_num_stages": [],
893
+ "range_multi_buffers": [
894
+ null,
895
+ false,
896
+ null
897
+ ],
898
+ "range_flattens": [
899
+ null,
900
+ false,
901
+ true
902
+ ],
903
+ "load_eviction_policies": [
904
+ "",
905
+ "last",
906
+ "first"
907
+ ],
908
+ "num_warps": 2,
909
+ "num_stages": 1,
910
+ "indexing": [
911
+ "pointer",
912
+ "tensor_descriptor",
913
+ "tensor_descriptor",
914
+ "tensor_descriptor",
915
+ "pointer"
916
+ ],
917
+ "atomic_indexing": [],
918
+ "pid_type": "flat"
919
+ }
920
+ },
921
+ {
922
+ "key": {
923
+ "hidden_size": 4096,
924
+ "num_tokens": 64
925
+ },
926
+ "config": {
927
+ "block_sizes": [
928
+ 16384,
929
+ 4096
930
+ ],
931
+ "range_unroll_factors": [
932
+ 0,
933
+ 1,
934
+ 2
935
+ ],
936
+ "range_warp_specializes": [],
937
+ "range_num_stages": [],
938
+ "range_multi_buffers": [
939
+ null,
940
+ null,
941
+ null
942
+ ],
943
+ "range_flattens": [
944
+ null,
945
+ false,
946
+ null
947
+ ],
948
+ "load_eviction_policies": [
949
+ "last",
950
+ "last",
951
+ "first"
952
+ ],
953
+ "num_warps": 16,
954
+ "num_stages": 8,
955
+ "indexing": [
956
+ "tensor_descriptor",
957
+ "tensor_descriptor",
958
+ "tensor_descriptor",
959
+ "pointer",
960
+ "tensor_descriptor"
961
+ ],
962
+ "atomic_indexing": [],
963
+ "pid_type": "flat"
964
+ }
965
+ },
966
+ {
967
+ "key": {
968
+ "hidden_size": 5120,
969
+ "num_tokens": 64
970
+ },
971
+ "config": {
972
+ "block_sizes": [
973
+ 8192,
974
+ 1024
975
+ ],
976
+ "range_unroll_factors": [
977
+ 0,
978
+ 3,
979
+ 2
980
+ ],
981
+ "range_warp_specializes": [
982
+ null,
983
+ false,
984
+ false
985
+ ],
986
+ "range_num_stages": [],
987
+ "range_multi_buffers": [
988
+ null,
989
+ true,
990
+ null
991
+ ],
992
+ "range_flattens": [
993
+ null,
994
+ null,
995
+ null
996
+ ],
997
+ "load_eviction_policies": [
998
+ "last",
999
+ "first",
1000
+ "last"
1001
+ ],
1002
+ "num_warps": 16,
1003
+ "num_stages": 6,
1004
+ "indexing": [
1005
+ "pointer",
1006
+ "pointer",
1007
+ "tensor_descriptor",
1008
+ "pointer",
1009
+ "tensor_descriptor"
1010
+ ],
1011
+ "atomic_indexing": [],
1012
+ "pid_type": "flat"
1013
+ }
1014
+ },
1015
+ {
1016
+ "key": {
1017
+ "hidden_size": 2048,
1018
+ "num_tokens": 128
1019
+ },
1020
+ "config": {
1021
+ "block_sizes": [
1022
+ 2048,
1023
+ 1024
1024
+ ],
1025
+ "range_unroll_factors": [
1026
+ 0,
1027
+ 4,
1028
+ 1
1029
+ ],
1030
+ "range_warp_specializes": [
1031
+ null,
1032
+ null,
1033
+ true
1034
+ ],
1035
+ "range_num_stages": [],
1036
+ "range_multi_buffers": [
1037
+ null,
1038
+ false,
1039
+ null
1040
+ ],
1041
+ "range_flattens": [
1042
+ null,
1043
+ false,
1044
+ true
1045
+ ],
1046
+ "load_eviction_policies": [
1047
+ "",
1048
+ "last",
1049
+ "first"
1050
+ ],
1051
+ "num_warps": 2,
1052
+ "num_stages": 1,
1053
+ "indexing": [
1054
+ "pointer",
1055
+ "tensor_descriptor",
1056
+ "tensor_descriptor",
1057
+ "tensor_descriptor",
1058
+ "pointer"
1059
+ ],
1060
+ "atomic_indexing": [],
1061
+ "pid_type": "flat"
1062
+ }
1063
+ },
1064
+ {
1065
+ "key": {
1066
+ "hidden_size": 4096,
1067
+ "num_tokens": 128
1068
+ },
1069
+ "config": {
1070
+ "block_sizes": [
1071
+ 16384,
1072
+ 4096
1073
+ ],
1074
+ "range_unroll_factors": [
1075
+ 0,
1076
+ 1,
1077
+ 2
1078
+ ],
1079
+ "range_warp_specializes": [],
1080
+ "range_num_stages": [],
1081
+ "range_multi_buffers": [
1082
+ null,
1083
+ null,
1084
+ null
1085
+ ],
1086
+ "range_flattens": [
1087
+ null,
1088
+ false,
1089
+ null
1090
+ ],
1091
+ "load_eviction_policies": [
1092
+ "last",
1093
+ "last",
1094
+ "first"
1095
+ ],
1096
+ "num_warps": 16,
1097
+ "num_stages": 8,
1098
+ "indexing": [
1099
+ "tensor_descriptor",
1100
+ "tensor_descriptor",
1101
+ "tensor_descriptor",
1102
+ "pointer",
1103
+ "tensor_descriptor"
1104
+ ],
1105
+ "atomic_indexing": [],
1106
+ "pid_type": "flat"
1107
+ }
1108
+ },
1109
+ {
1110
+ "key": {
1111
+ "hidden_size": 5120,
1112
+ "num_tokens": 128
1113
+ },
1114
+ "config": {
1115
+ "block_sizes": [
1116
+ 8192,
1117
+ 1024
1118
+ ],
1119
+ "range_unroll_factors": [
1120
+ 0,
1121
+ 3,
1122
+ 2
1123
+ ],
1124
+ "range_warp_specializes": [
1125
+ null,
1126
+ false,
1127
+ false
1128
+ ],
1129
+ "range_num_stages": [],
1130
+ "range_multi_buffers": [
1131
+ null,
1132
+ true,
1133
+ null
1134
+ ],
1135
+ "range_flattens": [
1136
+ null,
1137
+ null,
1138
+ null
1139
+ ],
1140
+ "load_eviction_policies": [
1141
+ "last",
1142
+ "first",
1143
+ "last"
1144
+ ],
1145
+ "num_warps": 16,
1146
+ "num_stages": 6,
1147
+ "indexing": [
1148
+ "pointer",
1149
+ "pointer",
1150
+ "tensor_descriptor",
1151
+ "pointer",
1152
+ "tensor_descriptor"
1153
+ ],
1154
+ "atomic_indexing": [],
1155
+ "pid_type": "flat"
1156
+ }
1157
+ },
1158
+ {
1159
+ "key": {
1160
+ "hidden_size": 2048,
1161
+ "num_tokens": 256
1162
+ },
1163
+ "config": {
1164
+ "block_sizes": [
1165
+ 2048,
1166
+ 1024
1167
+ ],
1168
+ "range_unroll_factors": [
1169
+ 0,
1170
+ 4,
1171
+ 1
1172
+ ],
1173
+ "range_warp_specializes": [
1174
+ null,
1175
+ null,
1176
+ true
1177
+ ],
1178
+ "range_num_stages": [],
1179
+ "range_multi_buffers": [
1180
+ null,
1181
+ false,
1182
+ null
1183
+ ],
1184
+ "range_flattens": [
1185
+ null,
1186
+ false,
1187
+ true
1188
+ ],
1189
+ "load_eviction_policies": [
1190
+ "",
1191
+ "last",
1192
+ "first"
1193
+ ],
1194
+ "num_warps": 2,
1195
+ "num_stages": 1,
1196
+ "indexing": [
1197
+ "pointer",
1198
+ "tensor_descriptor",
1199
+ "tensor_descriptor",
1200
+ "tensor_descriptor",
1201
+ "pointer"
1202
+ ],
1203
+ "atomic_indexing": [],
1204
+ "pid_type": "flat"
1205
+ }
1206
+ },
1207
+ {
1208
+ "key": {
1209
+ "hidden_size": 4096,
1210
+ "num_tokens": 256
1211
+ },
1212
+ "config": {
1213
+ "block_sizes": [
1214
+ 16384,
1215
+ 4096
1216
+ ],
1217
+ "range_unroll_factors": [
1218
+ 0,
1219
+ 1,
1220
+ 2
1221
+ ],
1222
+ "range_warp_specializes": [],
1223
+ "range_num_stages": [],
1224
+ "range_multi_buffers": [
1225
+ null,
1226
+ null,
1227
+ null
1228
+ ],
1229
+ "range_flattens": [
1230
+ null,
1231
+ false,
1232
+ null
1233
+ ],
1234
+ "load_eviction_policies": [
1235
+ "last",
1236
+ "last",
1237
+ "first"
1238
+ ],
1239
+ "num_warps": 16,
1240
+ "num_stages": 8,
1241
+ "indexing": [
1242
+ "tensor_descriptor",
1243
+ "tensor_descriptor",
1244
+ "tensor_descriptor",
1245
+ "pointer",
1246
+ "tensor_descriptor"
1247
+ ],
1248
+ "atomic_indexing": [],
1249
+ "pid_type": "flat"
1250
+ }
1251
+ },
1252
+ {
1253
+ "key": {
1254
+ "hidden_size": 5120,
1255
+ "num_tokens": 256
1256
+ },
1257
+ "config": {
1258
+ "block_sizes": [
1259
+ 8192,
1260
+ 1024
1261
+ ],
1262
+ "range_unroll_factors": [
1263
+ 0,
1264
+ 0,
1265
+ 4
1266
+ ],
1267
+ "range_warp_specializes": [
1268
+ null,
1269
+ true,
1270
+ null
1271
+ ],
1272
+ "range_num_stages": [],
1273
+ "range_multi_buffers": [
1274
+ null,
1275
+ true,
1276
+ false
1277
+ ],
1278
+ "range_flattens": [
1279
+ null,
1280
+ null,
1281
+ false
1282
+ ],
1283
+ "load_eviction_policies": [
1284
+ "",
1285
+ "",
1286
+ ""
1287
+ ],
1288
+ "num_warps": 4,
1289
+ "num_stages": 2,
1290
+ "indexing": [
1291
+ "tensor_descriptor",
1292
+ "pointer",
1293
+ "pointer",
1294
+ "tensor_descriptor",
1295
+ "tensor_descriptor"
1296
+ ],
1297
+ "atomic_indexing": [],
1298
+ "pid_type": "flat"
1299
+ }
1300
+ },
1301
+ {
1302
+ "key": {
1303
+ "hidden_size": 2048,
1304
+ "num_tokens": 512
1305
+ },
1306
+ "config": {
1307
+ "block_sizes": [
1308
+ 2048,
1309
+ 1024
1310
+ ],
1311
+ "range_unroll_factors": [
1312
+ 0,
1313
+ 4,
1314
+ 1
1315
+ ],
1316
+ "range_warp_specializes": [
1317
+ null,
1318
+ null,
1319
+ true
1320
+ ],
1321
+ "range_num_stages": [],
1322
+ "range_multi_buffers": [
1323
+ null,
1324
+ false,
1325
+ null
1326
+ ],
1327
+ "range_flattens": [
1328
+ null,
1329
+ false,
1330
+ true
1331
+ ],
1332
+ "load_eviction_policies": [
1333
+ "",
1334
+ "last",
1335
+ "first"
1336
+ ],
1337
+ "num_warps": 2,
1338
+ "num_stages": 1,
1339
+ "indexing": [
1340
+ "pointer",
1341
+ "tensor_descriptor",
1342
+ "tensor_descriptor",
1343
+ "tensor_descriptor",
1344
+ "pointer"
1345
+ ],
1346
+ "atomic_indexing": [],
1347
+ "pid_type": "flat"
1348
+ }
1349
+ },
1350
+ {
1351
+ "key": {
1352
+ "hidden_size": 4096,
1353
+ "num_tokens": 512
1354
+ },
1355
+ "config": {
1356
+ "block_sizes": [
1357
+ 16384,
1358
+ 4096
1359
+ ],
1360
+ "range_unroll_factors": [
1361
+ 0,
1362
+ 1,
1363
+ 3
1364
+ ],
1365
+ "range_warp_specializes": [],
1366
+ "range_num_stages": [],
1367
+ "range_multi_buffers": [
1368
+ null,
1369
+ null,
1370
+ false
1371
+ ],
1372
+ "range_flattens": [
1373
+ null,
1374
+ true,
1375
+ true
1376
+ ],
1377
+ "load_eviction_policies": [
1378
+ "first",
1379
+ "first",
1380
+ "first"
1381
+ ],
1382
+ "num_warps": 16,
1383
+ "num_stages": 3,
1384
+ "indexing": [
1385
+ "pointer",
1386
+ "pointer",
1387
+ "tensor_descriptor",
1388
+ "tensor_descriptor",
1389
+ "tensor_descriptor"
1390
+ ],
1391
+ "atomic_indexing": [],
1392
+ "pid_type": "flat"
1393
+ }
1394
+ },
1395
+ {
1396
+ "key": {
1397
+ "hidden_size": 5120,
1398
+ "num_tokens": 512
1399
+ },
1400
+ "config": {
1401
+ "block_sizes": [
1402
+ 8192,
1403
+ 1024
1404
+ ],
1405
+ "range_unroll_factors": [
1406
+ 0,
1407
+ 1,
1408
+ 0
1409
+ ],
1410
+ "range_warp_specializes": [
1411
+ null,
1412
+ false,
1413
+ true
1414
+ ],
1415
+ "range_num_stages": [],
1416
+ "range_multi_buffers": [
1417
+ null,
1418
+ true,
1419
+ null
1420
+ ],
1421
+ "range_flattens": [
1422
+ null,
1423
+ false,
1424
+ null
1425
+ ],
1426
+ "load_eviction_policies": [
1427
+ "first",
1428
+ "last",
1429
+ "first"
1430
+ ],
1431
+ "num_warps": 1,
1432
+ "num_stages": 3,
1433
+ "indexing": [
1434
+ "pointer",
1435
+ "tensor_descriptor",
1436
+ "tensor_descriptor",
1437
+ "tensor_descriptor",
1438
+ "pointer"
1439
+ ],
1440
+ "atomic_indexing": [],
1441
+ "pid_type": "flat"
1442
+ }
1443
+ },
1444
+ {
1445
+ "key": {
1446
+ "hidden_size": 2048,
1447
+ "num_tokens": 1024
1448
+ },
1449
+ "config": {
1450
+ "block_sizes": [
1451
+ 2048,
1452
+ 1024
1453
+ ],
1454
+ "range_unroll_factors": [
1455
+ 0,
1456
+ 4,
1457
+ 1
1458
+ ],
1459
+ "range_warp_specializes": [
1460
+ null,
1461
+ null,
1462
+ true
1463
+ ],
1464
+ "range_num_stages": [],
1465
+ "range_multi_buffers": [
1466
+ null,
1467
+ false,
1468
+ null
1469
+ ],
1470
+ "range_flattens": [
1471
+ null,
1472
+ false,
1473
+ true
1474
+ ],
1475
+ "load_eviction_policies": [
1476
+ "",
1477
+ "last",
1478
+ "first"
1479
+ ],
1480
+ "num_warps": 2,
1481
+ "num_stages": 1,
1482
+ "indexing": [
1483
+ "pointer",
1484
+ "tensor_descriptor",
1485
+ "tensor_descriptor",
1486
+ "tensor_descriptor",
1487
+ "pointer"
1488
+ ],
1489
+ "atomic_indexing": [],
1490
+ "pid_type": "flat"
1491
+ }
1492
+ },
1493
+ {
1494
+ "key": {
1495
+ "hidden_size": 4096,
1496
+ "num_tokens": 1024
1497
+ },
1498
+ "config": {
1499
+ "block_sizes": [
1500
+ 4096,
1501
+ 2048
1502
+ ],
1503
+ "range_unroll_factors": [
1504
+ 0,
1505
+ 4,
1506
+ 1
1507
+ ],
1508
+ "range_warp_specializes": [
1509
+ null,
1510
+ null,
1511
+ null
1512
+ ],
1513
+ "range_num_stages": [],
1514
+ "range_multi_buffers": [
1515
+ null,
1516
+ true,
1517
+ null
1518
+ ],
1519
+ "range_flattens": [
1520
+ null,
1521
+ null,
1522
+ true
1523
+ ],
1524
+ "load_eviction_policies": [
1525
+ "first",
1526
+ "last",
1527
+ "first"
1528
+ ],
1529
+ "num_warps": 4,
1530
+ "num_stages": 3,
1531
+ "indexing": [
1532
+ "tensor_descriptor",
1533
+ "pointer",
1534
+ "tensor_descriptor",
1535
+ "tensor_descriptor",
1536
+ "pointer"
1537
+ ],
1538
+ "atomic_indexing": [],
1539
+ "pid_type": "flat"
1540
+ }
1541
+ },
1542
+ {
1543
+ "key": {
1544
+ "hidden_size": 5120,
1545
+ "num_tokens": 1024
1546
+ },
1547
+ "config": {
1548
+ "block_sizes": [
1549
+ 8192,
1550
+ 8192
1551
+ ],
1552
+ "range_unroll_factors": [
1553
+ 0,
1554
+ 3,
1555
+ 0
1556
+ ],
1557
+ "range_warp_specializes": [],
1558
+ "range_num_stages": [],
1559
+ "range_multi_buffers": [
1560
+ null,
1561
+ true,
1562
+ true
1563
+ ],
1564
+ "range_flattens": [
1565
+ null,
1566
+ true,
1567
+ true
1568
+ ],
1569
+ "load_eviction_policies": [
1570
+ "",
1571
+ "last",
1572
+ "first"
1573
+ ],
1574
+ "num_warps": 2,
1575
+ "num_stages": 8,
1576
+ "indexing": [
1577
+ "tensor_descriptor",
1578
+ "pointer",
1579
+ "pointer",
1580
+ "tensor_descriptor",
1581
+ "tensor_descriptor"
1582
+ ],
1583
+ "atomic_indexing": [],
1584
+ "pid_type": "flat"
1585
+ }
1586
+ },
1587
+ {
1588
+ "key": {
1589
+ "hidden_size": 2048,
1590
+ "num_tokens": 2048
1591
+ },
1592
+ "config": {
1593
+ "block_sizes": [
1594
+ 2048,
1595
+ 1024
1596
+ ],
1597
+ "range_unroll_factors": [
1598
+ 0,
1599
+ 0,
1600
+ 3
1601
+ ],
1602
+ "range_warp_specializes": [
1603
+ null,
1604
+ false,
1605
+ null
1606
+ ],
1607
+ "range_num_stages": [],
1608
+ "range_multi_buffers": [
1609
+ null,
1610
+ null,
1611
+ true
1612
+ ],
1613
+ "range_flattens": [
1614
+ null,
1615
+ false,
1616
+ true
1617
+ ],
1618
+ "load_eviction_policies": [
1619
+ "",
1620
+ "last",
1621
+ "first"
1622
+ ],
1623
+ "num_warps": 4,
1624
+ "num_stages": 6,
1625
+ "indexing": [
1626
+ "pointer",
1627
+ "tensor_descriptor",
1628
+ "pointer",
1629
+ "tensor_descriptor",
1630
+ "tensor_descriptor"
1631
+ ],
1632
+ "atomic_indexing": [],
1633
+ "pid_type": "flat"
1634
+ }
1635
+ },
1636
+ {
1637
+ "key": {
1638
+ "hidden_size": 4096,
1639
+ "num_tokens": 2048
1640
+ },
1641
+ "config": {
1642
+ "block_sizes": [
1643
+ 8192,
1644
+ 8192
1645
+ ],
1646
+ "range_unroll_factors": [
1647
+ 0,
1648
+ 1,
1649
+ 3
1650
+ ],
1651
+ "range_warp_specializes": [
1652
+ null,
1653
+ null,
1654
+ null
1655
+ ],
1656
+ "range_num_stages": [],
1657
+ "range_multi_buffers": [
1658
+ null,
1659
+ false,
1660
+ false
1661
+ ],
1662
+ "range_flattens": [
1663
+ null,
1664
+ true,
1665
+ null
1666
+ ],
1667
+ "load_eviction_policies": [
1668
+ "last",
1669
+ "first",
1670
+ "first"
1671
+ ],
1672
+ "num_warps": 2,
1673
+ "num_stages": 8,
1674
+ "indexing": [
1675
+ "pointer",
1676
+ "tensor_descriptor",
1677
+ "pointer",
1678
+ "tensor_descriptor",
1679
+ "tensor_descriptor"
1680
+ ],
1681
+ "atomic_indexing": [],
1682
+ "pid_type": "flat"
1683
+ }
1684
+ },
1685
+ {
1686
+ "key": {
1687
+ "hidden_size": 5120,
1688
+ "num_tokens": 2048
1689
+ },
1690
+ "config": {
1691
+ "block_sizes": [
1692
+ 8192,
1693
+ 8192
1694
+ ],
1695
+ "range_unroll_factors": [
1696
+ 0,
1697
+ 1,
1698
+ 3
1699
+ ],
1700
+ "range_warp_specializes": [
1701
+ null,
1702
+ null,
1703
+ null
1704
+ ],
1705
+ "range_num_stages": [],
1706
+ "range_multi_buffers": [
1707
+ null,
1708
+ false,
1709
+ false
1710
+ ],
1711
+ "range_flattens": [
1712
+ null,
1713
+ true,
1714
+ null
1715
+ ],
1716
+ "load_eviction_policies": [
1717
+ "last",
1718
+ "first",
1719
+ "first"
1720
+ ],
1721
+ "num_warps": 2,
1722
+ "num_stages": 8,
1723
+ "indexing": [
1724
+ "pointer",
1725
+ "tensor_descriptor",
1726
+ "pointer",
1727
+ "tensor_descriptor",
1728
+ "tensor_descriptor"
1729
+ ],
1730
+ "atomic_indexing": [],
1731
+ "pid_type": "flat"
1732
+ }
1733
+ },
1734
+ {
1735
+ "key": {
1736
+ "hidden_size": 2048,
1737
+ "num_tokens": 4096
1738
+ },
1739
+ "config": {
1740
+ "block_sizes": [
1741
+ 512,
1742
+ 4096
1743
+ ],
1744
+ "range_unroll_factors": [
1745
+ 0,
1746
+ 0,
1747
+ 4
1748
+ ],
1749
+ "range_warp_specializes": [
1750
+ null,
1751
+ null,
1752
+ null
1753
+ ],
1754
+ "range_num_stages": [],
1755
+ "range_multi_buffers": [
1756
+ null,
1757
+ false,
1758
+ true
1759
+ ],
1760
+ "range_flattens": [
1761
+ null,
1762
+ null,
1763
+ false
1764
+ ],
1765
+ "load_eviction_policies": [
1766
+ "",
1767
+ "",
1768
+ "first"
1769
+ ],
1770
+ "num_warps": 1,
1771
+ "num_stages": 4,
1772
+ "indexing": [
1773
+ "pointer",
1774
+ "pointer",
1775
+ "tensor_descriptor",
1776
+ "pointer",
1777
+ "tensor_descriptor"
1778
+ ],
1779
+ "atomic_indexing": [],
1780
+ "pid_type": "flat"
1781
+ }
1782
+ },
1783
+ {
1784
+ "key": {
1785
+ "hidden_size": 4096,
1786
+ "num_tokens": 4096
1787
+ },
1788
+ "config": {
1789
+ "block_sizes": [
1790
+ 4096,
1791
+ 2048
1792
+ ],
1793
+ "range_unroll_factors": [
1794
+ 0,
1795
+ 2,
1796
+ 2
1797
+ ],
1798
+ "range_warp_specializes": [
1799
+ null,
1800
+ false,
1801
+ false
1802
+ ],
1803
+ "range_num_stages": [],
1804
+ "range_multi_buffers": [
1805
+ null,
1806
+ true,
1807
+ true
1808
+ ],
1809
+ "range_flattens": [
1810
+ null,
1811
+ null,
1812
+ true
1813
+ ],
1814
+ "load_eviction_policies": [
1815
+ "",
1816
+ "",
1817
+ ""
1818
+ ],
1819
+ "num_warps": 4,
1820
+ "num_stages": 1,
1821
+ "indexing": [
1822
+ "pointer",
1823
+ "tensor_descriptor",
1824
+ "pointer",
1825
+ "pointer",
1826
+ "tensor_descriptor"
1827
+ ],
1828
+ "atomic_indexing": [],
1829
+ "pid_type": "flat"
1830
+ }
1831
+ },
1832
+ {
1833
+ "key": {
1834
+ "hidden_size": 5120,
1835
+ "num_tokens": 4096
1836
+ },
1837
+ "config": {
1838
+ "block_sizes": [
1839
+ 8192,
1840
+ 8192
1841
+ ],
1842
+ "range_unroll_factors": [
1843
+ 0,
1844
+ 3,
1845
+ 0
1846
+ ],
1847
+ "range_warp_specializes": [],
1848
+ "range_num_stages": [],
1849
+ "range_multi_buffers": [
1850
+ null,
1851
+ true,
1852
+ true
1853
+ ],
1854
+ "range_flattens": [
1855
+ null,
1856
+ true,
1857
+ true
1858
+ ],
1859
+ "load_eviction_policies": [
1860
+ "",
1861
+ "last",
1862
+ "first"
1863
+ ],
1864
+ "num_warps": 2,
1865
+ "num_stages": 8,
1866
+ "indexing": [
1867
+ "tensor_descriptor",
1868
+ "pointer",
1869
+ "pointer",
1870
+ "tensor_descriptor",
1871
+ "tensor_descriptor"
1872
+ ],
1873
+ "atomic_indexing": [],
1874
+ "pid_type": "flat"
1875
+ }
1876
+ },
1877
+ {
1878
+ "key": {
1879
+ "hidden_size": 2048,
1880
+ "num_tokens": 8192
1881
+ },
1882
+ "config": {
1883
+ "block_sizes": [
1884
+ 512,
1885
+ 512
1886
+ ],
1887
+ "range_unroll_factors": [
1888
+ 0,
1889
+ 0,
1890
+ 4
1891
+ ],
1892
+ "range_warp_specializes": [
1893
+ null,
1894
+ false,
1895
+ null
1896
+ ],
1897
+ "range_num_stages": [],
1898
+ "range_multi_buffers": [
1899
+ null,
1900
+ null,
1901
+ false
1902
+ ],
1903
+ "range_flattens": [
1904
+ null,
1905
+ false,
1906
+ false
1907
+ ],
1908
+ "load_eviction_policies": [
1909
+ "",
1910
+ "last",
1911
+ ""
1912
+ ],
1913
+ "num_warps": 1,
1914
+ "num_stages": 4,
1915
+ "indexing": [
1916
+ "pointer",
1917
+ "tensor_descriptor",
1918
+ "pointer",
1919
+ "tensor_descriptor",
1920
+ "tensor_descriptor"
1921
+ ],
1922
+ "atomic_indexing": [],
1923
+ "pid_type": "flat"
1924
+ }
1925
+ },
1926
+ {
1927
+ "key": {
1928
+ "hidden_size": 4096,
1929
+ "num_tokens": 8192
1930
+ },
1931
+ "config": {
1932
+ "block_sizes": [
1933
+ 512,
1934
+ 4096
1935
+ ],
1936
+ "range_unroll_factors": [
1937
+ 0,
1938
+ 0,
1939
+ 4
1940
+ ],
1941
+ "range_warp_specializes": [
1942
+ null,
1943
+ null,
1944
+ null
1945
+ ],
1946
+ "range_num_stages": [],
1947
+ "range_multi_buffers": [
1948
+ null,
1949
+ false,
1950
+ true
1951
+ ],
1952
+ "range_flattens": [
1953
+ null,
1954
+ null,
1955
+ false
1956
+ ],
1957
+ "load_eviction_policies": [
1958
+ "",
1959
+ "",
1960
+ "first"
1961
+ ],
1962
+ "num_warps": 1,
1963
+ "num_stages": 4,
1964
+ "indexing": [
1965
+ "pointer",
1966
+ "pointer",
1967
+ "tensor_descriptor",
1968
+ "pointer",
1969
+ "tensor_descriptor"
1970
+ ],
1971
+ "atomic_indexing": [],
1972
+ "pid_type": "flat"
1973
+ }
1974
+ },
1975
+ {
1976
+ "key": {
1977
+ "hidden_size": 5120,
1978
+ "num_tokens": 8192
1979
+ },
1980
+ "config": {
1981
+ "block_sizes": [
1982
+ 1024,
1983
+ 512
1984
+ ],
1985
+ "range_unroll_factors": [
1986
+ 2,
1987
+ 3,
1988
+ 2
1989
+ ],
1990
+ "range_warp_specializes": [
1991
+ false,
1992
+ false,
1993
+ null
1994
+ ],
1995
+ "range_multi_buffers": [
1996
+ false,
1997
+ null,
1998
+ null
1999
+ ],
2000
+ "range_flattens": [
2001
+ null,
2002
+ false,
2003
+ null
2004
+ ],
2005
+ "load_eviction_policies": [
2006
+ "last",
2007
+ "first",
2008
+ "first"
2009
+ ],
2010
+ "num_warps": 1,
2011
+ "num_stages": 1,
2012
+ "indexing": [
2013
+ "pointer",
2014
+ "pointer",
2015
+ "pointer",
2016
+ "tensor_descriptor",
2017
+ "pointer"
2018
+ ],
2019
+ "atomic_indexing": [],
2020
+ "pid_type": "persistent_interleaved",
2021
+ "num_sm_multiplier": 64,
2022
+ "maxnreg": 128
2023
+ }
2024
+ }
2025
+]
\ No newline at end of file
vllm/kernels/helion/configs/dynamic_per_token_scaled_fp8_quant/nvidia_h100.json
new
+5185
@@ -0,0 +1,5185 @@
1
+[
2
+ {
3
+ "key": {
4
+ "hidden_size": 512,
5
+ "num_tokens": 1
6
+ },
7
+ "config": {
8
+ "block_sizes": [
9
+ 256,
10
+ 512
11
+ ],
12
+ "range_unroll_factors": [
13
+ 0,
14
+ 1,
15
+ 3
16
+ ],
17
+ "range_warp_specializes": [],
18
+ "range_num_stages": [],
19
+ "range_multi_buffers": [
20
+ null,
21
+ null,
22
+ true
23
+ ],
24
+ "range_flattens": [
25
+ null,
26
+ false,
27
+ true
28
+ ],
29
+ "load_eviction_policies": [
30
+ "last",
31
+ "",
32
+ ""
33
+ ],
34
+ "num_warps": 1,
35
+ "num_stages": 7,
36
+ "indexing": [
37
+ "pointer",
38
+ "pointer",
39
+ "tensor_descriptor",
40
+ "tensor_descriptor",
41
+ "pointer"
42
+ ],
43
+ "atomic_indexing": [],
44
+ "pid_type": "flat"
45
+ }
46
+ },
47
+ {
48
+ "key": {
49
+ "hidden_size": 2048,
50
+ "num_tokens": 1
51
+ },
52
+ "config": {
53
+ "block_sizes": [
54
+ 2048,
55
+ 2048
56
+ ],
57
+ "range_unroll_factors": [
58
+ 0,
59
+ 4,
60
+ 1
61
+ ],
62
+ "range_warp_specializes": [],
63
+ "range_num_stages": [],
64
+ "range_multi_buffers": [
65
+ null,
66
+ true,
67
+ null
68
+ ],
69
+ "range_flattens": [
70
+ null,
71
+ false,
72
+ null
73
+ ],
74
+ "load_eviction_policies": [
75
+ "",
76
+ "first",
77
+ ""
78
+ ],
79
+ "num_warps": 8,
80
+ "num_stages": 8,
81
+ "indexing": [
82
+ "tensor_descriptor",
83
+ "pointer",
84
+ "pointer",
85
+ "tensor_descriptor",
86
+ "pointer"
87
+ ],
88
+ "atomic_indexing": [],
89
+ "pid_type": "flat"
90
+ }
91
+ },
92
+ {
93
+ "key": {
94
+ "hidden_size": 4096,
95
+ "num_tokens": 1
96
+ },
97
+ "config": {
98
+ "block_sizes": [
99
+ 4096,
100
+ 4096
101
+ ],
102
+ "range_unroll_factors": [
103
+ 0,
104
+ 3,
105
+ 4
106
+ ],
107
+ "range_warp_specializes": [],
108
+ "range_num_stages": [],
109
+ "range_multi_buffers": [
110
+ null,
111
+ true,
112
+ false
113
+ ],
114
+ "range_flattens": [
115
+ null,
116
+ false,
117
+ true
118
+ ],
119
+ "load_eviction_policies": [
120
+ "first",
121
+ "first",
122
+ "last"
123
+ ],
124
+ "num_warps": 8,
125
+ "num_stages": 1,
126
+ "indexing": [
127
+ "pointer",
128
+ "pointer",
129
+ "pointer",
130
+ "pointer",
131
+ "tensor_descriptor"
132
+ ],
133
+ "atomic_indexing": [],
134
+ "pid_type": "flat"
135
+ }
136
+ },
137
+ {
138
+ "key": {
139
+ "hidden_size": 6144,
140
+ "num_tokens": 1
141
+ },
142
+ "config": {
143
+ "block_sizes": [
144
+ 8192,
145
+ 2048
146
+ ],
147
+ "range_unroll_factors": [
148
+ 0,
149
+ 1,
150
+ 0
151
+ ],
152
+ "range_warp_specializes": [],
153
+ "range_num_stages": [],
154
+ "range_multi_buffers": [
155
+ null,
156
+ true,
157
+ false
158
+ ],
159
+ "range_flattens": [
160
+ null,
161
+ true,
162
+ true
163
+ ],
164
+ "load_eviction_policies": [
165
+ "last",
166
+ "first",
167
+ "first"
168
+ ],
169
+ "num_warps": 8,
170
+ "num_stages": 2,
171
+ "indexing": [
172
+ "pointer",
173
+ "tensor_descriptor",
174
+ "pointer",
175
+ "pointer",
176
+ "pointer"
177
+ ],
178
+ "atomic_indexing": [],
179
+ "pid_type": "flat"
180
+ }
181
+ },
182
+ {
183
+ "key": {
184
+ "hidden_size": 8192,
185
+ "num_tokens": 1
186
+ },
187
+ "config": {
188
+ "block_sizes": [
189
+ 8192,
190
+ 4096
191
+ ],
192
+ "range_unroll_factors": [
193
+ 4,
194
+ 1,
195
+ 0
196
+ ],
197
+ "range_warp_specializes": [],
198
+ "range_multi_buffers": [
199
+ true,
200
+ null,
201
+ true
202
+ ],
203
+ "range_flattens": [
204
+ false,
205
+ false,
206
+ true
207
+ ],
208
+ "load_eviction_policies": [
209
+ "last",
210
+ "",
211
+ "last"
212
+ ],
213
+ "num_warps": 16,
214
+ "num_stages": 3,
215
+ "indexing": [
216
+ "pointer",
217
+ "pointer",
218
+ "tensor_descriptor",
219
+ "tensor_descriptor",
220
+ "pointer"
221
+ ],
222
+ "atomic_indexing": [],
223
+ "pid_type": "persistent_interleaved",
224
+ "num_sm_multiplier": 1
225
+ }
226
+ },
227
+ {
228
+ "key": {
229
+ "hidden_size": 12288,
230
+ "num_tokens": 1
231
+ },
232
+ "config": {
233
+ "block_sizes": [
234
+ 16384,
235
+ 4096
236
+ ],
237
+ "range_unroll_factors": [
238
+ 0,
239
+ 1,
240
+ 2
241
+ ],
242
+ "range_warp_specializes": [],
243
+ "range_num_stages": [],
244
+ "range_multi_buffers": [
245
+ null,
246
+ true,
247
+ false
248
+ ],
249
+ "range_flattens": [
250
+ null,
251
+ true,
252
+ null
253
+ ],
254
+ "load_eviction_policies": [
255
+ "",
256
+ "last",
257
+ "last"
258
+ ],
259
+ "num_warps": 16,
260
+ "num_stages": 6,
261
+ "indexing": [
262
+ "tensor_descriptor",
263
+ "tensor_descriptor",
264
+ "pointer",
265
+ "tensor_descriptor",
266
+ "tensor_descriptor"
267
+ ],
268
+ "atomic_indexing": [],
269
+ "pid_type": "flat"
270
+ }
271
+ },
272
+ {
273
+ "key": {
274
+ "hidden_size": 28672,
275
+ "num_tokens": 1
276
+ },
277
+ "config": {
278
+ "block_sizes": [
279
+ 32768,
280
+ 4096
281
+ ],
282
+ "range_unroll_factors": [
283
+ 0,
284
+ 4,
285
+ 1
286
+ ],
287
+ "range_warp_specializes": [],
288
+ "range_num_stages": [],
289
+ "range_multi_buffers": [
290
+ null,
291
+ false,
292
+ false
293
+ ],
294
+ "range_flattens": [
295
+ null,
296
+ true,
297
+ true
298
+ ],
299
+ "load_eviction_policies": [
300
+ "",
301
+ "first",
302
+ "last"
303
+ ],
304
+ "num_warps": 32,
305
+ "num_stages": 5,
306
+ "indexing": [
307
+ "tensor_descriptor",
308
+ "tensor_descriptor",
309
+ "pointer",
310
+ "tensor_descriptor",
311
+ "pointer"
312
+ ],
313
+ "atomic_indexing": [],
314
+ "pid_type": "flat"
315
+ }
316
+ },
317
+ {
318
+ "key": {
319
+ "hidden_size": 512,
320
+ "num_tokens": 2
321
+ },
322
+ "config": {
323
+ "block_sizes": [
324
+ 512,
325
+ 512
326
+ ],
327
+ "range_unroll_factors": [
328
+ 0,
329
+ 1,
330
+ 4
331
+ ],
332
+ "range_warp_specializes": [],
333
+ "range_num_stages": [],
334
+ "range_multi_buffers": [
335
+ null,
336
+ false,
337
+ false
338
+ ],
339
+ "range_flattens": [
340
+ null,
341
+ null,
342
+ false
343
+ ],
344
+ "load_eviction_policies": [
345
+ "",
346
+ "first",
347
+ "last"
348
+ ],
349
+ "num_warps": 1,
350
+ "num_stages": 5,
351
+ "indexing": [
352
+ "pointer",
353
+ "pointer",
354
+ "pointer",
355
+ "pointer",
356
+ "tensor_descriptor"
357
+ ],
358
+ "atomic_indexing": [],
359
+ "pid_type": "flat"
360
+ }
361
+ },
362
+ {
363
+ "key": {
364
+ "hidden_size": 2048,
365
+ "num_tokens": 2
366
+ },
367
+ "config": {
368
+ "block_sizes": [
369
+ 2048,
370
+ 2048
371
+ ],
372
+ "range_unroll_factors": [
373
+ 0,
374
+ 2,
375
+ 0
376
+ ],
377
+ "range_warp_specializes": [],
378
+ "range_num_stages": [],
379
+ "range_multi_buffers": [
380
+ null,
381
+ null,
382
+ true
383
+ ],
384
+ "range_flattens": [
385
+ null,
386
+ true,
387
+ true
388
+ ],
389
+ "load_eviction_policies": [
390
+ "last",
391
+ "",
392
+ ""
393
+ ],
394
+ "num_warps": 8,
395
+ "num_stages": 2,
396
+ "indexing": [
397
+ "pointer",
398
+ "pointer",
399
+ "tensor_descriptor",
400
+ "tensor_descriptor",
401
+ "tensor_descriptor"
402
+ ],
403
+ "atomic_indexing": [],
404
+ "pid_type": "flat"
405
+ }
406
+ },
407
+ {
408
+ "key": {
409
+ "hidden_size": 4096,
410
+ "num_tokens": 2
411
+ },
412
+ "config": {
413
+ "block_sizes": [
414
+ 4096,
415
+ 2048
416
+ ],
417
+ "range_unroll_factors": [
418
+ 0,
419
+ 4,
420
+ 4
421
+ ],
422
+ "range_warp_specializes": [],
423
+ "range_num_stages": [],
424
+ "range_multi_buffers": [
425
+ null,
426
+ null,
427
+ null
428
+ ],
429
+ "range_flattens": [
430
+ null,
431
+ false,
432
+ true
433
+ ],
434
+ "load_eviction_policies": [
435
+ "last",
436
+ "",
437
+ "last"
438
+ ],
439
+ "num_warps": 8,
440
+ "num_stages": 3,
441
+ "indexing": [
442
+ "pointer",
443
+ "pointer",
444
+ "pointer",
445
+ "tensor_descriptor",
446
+ "pointer"
447
+ ],
448
+ "atomic_indexing": [],
449
+ "pid_type": "flat"
450
+ }
451
+ },
452
+ {
453
+ "key": {
454
+ "hidden_size": 6144,
455
+ "num_tokens": 2
456
+ },
457
+ "config": {
458
+ "block_sizes": [
459
+ 8192,
460
+ 2048
461
+ ],
462
+ "range_unroll_factors": [
463
+ 0,
464
+ 2,
465
+ 2
466
+ ],
467
+ "range_warp_specializes": [],
468
+ "range_num_stages": [],
469
+ "range_multi_buffers": [
470
+ null,
471
+ false,
472
+ false
473
+ ],
474
+ "range_flattens": [
475
+ null,
476
+ false,
477
+ null
478
+ ],
479
+ "load_eviction_policies": [
480
+ "",
481
+ "last",
482
+ ""
483
+ ],
484
+ "num_warps": 8,
485
+ "num_stages": 7,
486
+ "indexing": [
487
+ "tensor_descriptor",
488
+ "tensor_descriptor",
489
+ "tensor_descriptor",
490
+ "tensor_descriptor",
491
+ "tensor_descriptor"
492
+ ],
493
+ "atomic_indexing": [],
494
+ "pid_type": "flat"
495
+ }
496
+ },
497
+ {
498
+ "key": {
499
+ "hidden_size": 8192,
500
+ "num_tokens": 2
501
+ },
502
+ "config": {
503
+ "block_sizes": [
504
+ 8192,
505
+ 4096
506
+ ],
507
+ "range_unroll_factors": [
508
+ 0,
509
+ 1,
510
+ 1
511
+ ],
512
+ "range_warp_specializes": [],
513
+ "range_num_stages": [],
514
+ "range_multi_buffers": [
515
+ null,
516
+ true,
517
+ false
518
+ ],
519
+ "range_flattens": [
520
+ null,
521
+ false,
522
+ true
523
+ ],
524
+ "load_eviction_policies": [
525
+ "",
526
+ "",
527
+ ""
528
+ ],
529
+ "num_warps": 16,
530
+ "num_stages": 1,
531
+ "indexing": [
532
+ "pointer",
533
+ "pointer",
534
+ "tensor_descriptor",
535
+ "tensor_descriptor",
536
+ "tensor_descriptor"
537
+ ],
538
+ "atomic_indexing": [],
539
+ "pid_type": "flat"
540
+ }
541
+ },
542
+ {
543
+ "key": {
544
+ "hidden_size": 12288,
545
+ "num_tokens": 2
546
+ },
547
+ "config": {
548
+ "block_sizes": [
549
+ 16384,
550
+ 4096
551
+ ],
552
+ "range_unroll_factors": [
553
+ 0,
554
+ 1,
555
+ 2
556
+ ],
557
+ "range_warp_specializes": [],
558
+ "range_num_stages": [],
559
+ "range_multi_buffers": [
560
+ null,
561
+ null,
562
+ null
563
+ ],
564
+ "range_flattens": [
565
+ null,
566
+ false,
567
+ null
568
+ ],
569
+ "load_eviction_policies": [
570
+ "last",
571
+ "last",
572
+ "first"
573
+ ],
574
+ "num_warps": 16,
575
+ "num_stages": 8,
576
+ "indexing": [
577
+ "tensor_descriptor",
578
+ "tensor_descriptor",
579
+ "tensor_descriptor",
580
+ "pointer",
581
+ "tensor_descriptor"
582
+ ],
583
+ "atomic_indexing": [],
584
+ "pid_type": "flat"
585
+ }
586
+ },
587
+ {
588
+ "key": {
589
+ "hidden_size": 28672,
590
+ "num_tokens": 2
591
+ },
592
+ "config": {
593
+ "block_sizes": [
594
+ 32768,
595
+ 4096
596
+ ],
597
+ "range_unroll_factors": [
598
+ 0,
599
+ 1,
600
+ 4
601
+ ],
602
+ "range_warp_specializes": [],
603
+ "range_num_stages": [],
604
+ "range_multi_buffers": [
605
+ null,
606
+ true,
607
+ true
608
+ ],
609
+ "range_flattens": [
610
+ null,
611
+ true,
612
+ true
613
+ ],
614
+ "load_eviction_policies": [
615
+ "last",
616
+ "last",
617
+ ""
618
+ ],
619
+ "num_warps": 16,
620
+ "num_stages": 4,
621
+ "indexing": [
622
+ "pointer",
623
+ "pointer",
624
+ "tensor_descriptor",
625
+ "tensor_descriptor",
626
+ "pointer"
627
+ ],
628
+ "atomic_indexing": [],
629
+ "pid_type": "flat"
630
+ }
631
+ },
632
+ {
633
+ "key": {
634
+ "hidden_size": 512,
635
+ "num_tokens": 4
636
+ },
637
+ "config": {
638
+ "block_sizes": [
639
+ 512,
640
+ 512
641
+ ],
642
+ "range_unroll_factors": [
643
+ 0,
644
+ 4,
645
+ 2
646
+ ],
647
+ "range_warp_specializes": [],
648
+ "range_num_stages": [],
649
+ "range_multi_buffers": [
650
+ null,
651
+ true,
652
+ true
653
+ ],
654
+ "range_flattens": [
655
+ null,
656
+ true,
657
+ true
658
+ ],
659
+ "load_eviction_policies": [
660
+ "last",
661
+ "first",
662
+ "last"
663
+ ],
664
+ "num_warps": 1,
665
+ "num_stages": 5,
666
+ "indexing": [
667
+ "tensor_descriptor",
668
+ "tensor_descriptor",
669
+ "tensor_descriptor",
670
+ "pointer",
671
+ "tensor_descriptor"
672
+ ],
673
+ "atomic_indexing": [],
674
+ "pid_type": "flat"
675
+ }
676
+ },
677
+ {
678
+ "key": {
679
+ "hidden_size": 2048,
680
+ "num_tokens": 4
681
+ },
682
+ "config": {
683
+ "block_sizes": [
684
+ 2048,
685
+ 2048
686
+ ],
687
+ "range_unroll_factors": [
688
+ 0,
689
+ 0,
690
+ 3
691
+ ],
692
+ "range_warp_specializes": [],
693
+ "range_num_stages": [],
694
+ "range_multi_buffers": [
695
+ null,
696
+ false,
697
+ false
698
+ ],
699
+ "range_flattens": [
700
+ null,
701
+ false,
702
+ null
703
+ ],
704
+ "load_eviction_policies": [
705
+ "last",
706
+ "last",
707
+ "first"
708
+ ],
709
+ "num_warps": 8,
710
+ "num_stages": 3,
711
+ "indexing": [
712
+ "tensor_descriptor",
713
+ "pointer",
714
+ "pointer",
715
+ "tensor_descriptor",
716
+ "tensor_descriptor"
717
+ ],
718
+ "atomic_indexing": [],
719
+ "pid_type": "flat"
720
+ }
721
+ },
722
+ {
723
+ "key": {
724
+ "hidden_size": 4096,
725
+ "num_tokens": 4
726
+ },
727
+ "config": {
728
+ "block_sizes": [
729
+ 4096,
730
+ 2048
731
+ ],
732
+ "range_unroll_factors": [
733
+ 0,
734
+ 1,
735
+ 0
736
+ ],
737
+ "range_warp_specializes": [],
738
+ "range_num_stages": [],
739
+ "range_multi_buffers": [
740
+ null,
741
+ true,
742
+ false
743
+ ],
744
+ "range_flattens": [
745
+ null,
746
+ true,
747
+ null
748
+ ],
749
+ "load_eviction_policies": [
750
+ "last",
751
+ "",
752
+ ""
753
+ ],
754
+ "num_warps": 8,
755
+ "num_stages": 5,
756
+ "indexing": [
757
+ "tensor_descriptor",
758
+ "tensor_descriptor",
759
+ "pointer",
760
+ "pointer",
761
+ "tensor_descriptor"
762
+ ],
763
+ "atomic_indexing": [],
764
+ "pid_type": "flat"
765
+ }
766
+ },
767
+ {
768
+ "key": {
769
+ "hidden_size": 6144,
770
+ "num_tokens": 4
771
+ },
772
+ "config": {
773
+ "block_sizes": [
774
+ 8192,
775
+ 2048
776
+ ],
777
+ "range_unroll_factors": [
778
+ 0,
779
+ 1,
780
+ 0
781
+ ],
782
+ "range_warp_specializes": [],
783
+ "range_num_stages": [],
784
+ "range_multi_buffers": [
785
+ null,
786
+ false,
787
+ null
788
+ ],
789
+ "range_flattens": [
790
+ null,
791
+ false,
792
+ false
793
+ ],
794
+ "load_eviction_policies": [
795
+ "first",
796
+ "",
797
+ ""
798
+ ],
799
+ "num_warps": 8,
800
+ "num_stages": 4,
801
+ "indexing": [
802
+ "tensor_descriptor",
803
+ "tensor_descriptor",
804
+ "tensor_descriptor",
805
+ "tensor_descriptor",
806
+ "tensor_descriptor"
807
+ ],
808
+ "atomic_indexing": [],
809
+ "pid_type": "flat"
810
+ }
811
+ },
812
+ {
813
+ "key": {
814
+ "hidden_size": 8192,
815
+ "num_tokens": 4
816
+ },
817
+ "config": {
818
+ "block_sizes": [
819
+ 8192,
820
+ 4096
821
+ ],
822
+ "range_unroll_factors": [
823
+ 0,
824
+ 0,
825
+ 4
826
+ ],
827
+ "range_warp_specializes": [],
828
+ "range_num_stages": [],
829
+ "range_multi_buffers": [
830
+ null,
831
+ true,
832
+ false
833
+ ],
834
+ "range_flattens": [
835
+ null,
836
+ true,
837
+ true
838
+ ],
839
+ "load_eviction_policies": [
840
+ "first",
841
+ "last",
842
+ "last"
843
+ ],
844
+ "num_warps": 16,
845
+ "num_stages": 1,
846
+ "indexing": [
847
+ "pointer",
848
+ "pointer",
849
+ "tensor_descriptor",
850
+ "pointer",
851
+ "tensor_descriptor"
852
+ ],
853
+ "atomic_indexing": [],
854
+ "pid_type": "flat"
855
+ }
856
+ },
857
+ {
858
+ "key": {
859
+ "hidden_size": 12288,
860
+ "num_tokens": 4
861
+ },
862
+ "config": {
863
+ "block_sizes": [
864
+ 16384,
865
+ 4096
866
+ ],
867
+ "range_unroll_factors": [
868
+ 0,
869
+ 3,
870
+ 3
871
+ ],
872
+ "range_warp_specializes": [],
873
+ "range_num_stages": [],
874
+ "range_multi_buffers": [
875
+ null,
876
+ true,
877
+ false
878
+ ],
879
+ "range_flattens": [
880
+ null,
881
+ false,
882
+ null
883
+ ],
884
+ "load_eviction_policies": [
885
+ "",
886
+ "",
887
+ ""
888
+ ],
889
+ "num_warps": 16,
890
+ "num_stages": 6,
891
+ "indexing": [
892
+ "tensor_descriptor",
893
+ "pointer",
894
+ "tensor_descriptor",
895
+ "tensor_descriptor",
896
+ "tensor_descriptor"
897
+ ],
898
+ "atomic_indexing": [],
899
+ "pid_type": "flat"
900
+ }
901
+ },
902
+ {
903
+ "key": {
904
+ "hidden_size": 28672,
905
+ "num_tokens": 4
906
+ },
907
+ "config": {
908
+ "block_sizes": [
909
+ 32768,
910
+ 4096
911
+ ],
912
+ "range_unroll_factors": [
913
+ 0,
914
+ 2,
915
+ 2
916
+ ],
917
+ "range_warp_specializes": [],
918
+ "range_num_stages": [],
919
+ "range_multi_buffers": [
920
+ null,
921
+ false,
922
+ false
923
+ ],
924
+ "range_flattens": [
925
+ null,
926
+ null,
927
+ null
928
+ ],
929
+ "load_eviction_policies": [
930
+ "last",
931
+ "first",
932
+ ""
933
+ ],
934
+ "num_warps": 16,
935
+ "num_stages": 4,
936
+ "indexing": [
937
+ "pointer",
938
+ "pointer",
939
+ "pointer",
940
+ "pointer",
941
+ "pointer"
942
+ ],
943
+ "atomic_indexing": [],
944
+ "pid_type": "flat"
945
+ }
946
+ },
947
+ {
948
+ "key": {
949
+ "hidden_size": 512,
950
+ "num_tokens": 8
951
+ },
952
+ "config": {
953
+ "block_sizes": [
954
+ 256,
955
+ 512
956
+ ],
957
+ "range_unroll_factors": [
958
+ 0,
959
+ 4,
960
+ 0
961
+ ],
962
+ "range_warp_specializes": [],
963
+ "range_num_stages": [],
964
+ "range_multi_buffers": [
965
+ null,
966
+ null,
967
+ null
968
+ ],
969
+ "range_flattens": [
970
+ null,
971
+ true,
972
+ false
973
+ ],
974
+ "load_eviction_policies": [
975
+ "last",
976
+ "first",
977
+ ""
978
+ ],
979
+ "num_warps": 1,
980
+ "num_stages": 5,
981
+ "indexing": [
982
+ "tensor_descriptor",
983
+ "tensor_descriptor",
984
+ "tensor_descriptor",
985
+ "pointer",
986
+ "tensor_descriptor"
987
+ ],
988
+ "atomic_indexing": [],
989
+ "pid_type": "flat"
990
+ }
991
+ },
992
+ {
993
+ "key": {
994
+ "hidden_size": 2048,
995
+ "num_tokens": 8
996
+ },
997
+ "config": {
998
+ "block_sizes": [
999
+ 2048,
1000
+ 2048
1001
+ ],
1002
+ "range_unroll_factors": [
1003
+ 0,
1004
+ 0,
1005
+ 2
1006
+ ],
1007
+ "range_warp_specializes": [],
1008
+ "range_num_stages": [],
1009
+ "range_multi_buffers": [
1010
+ null,
1011
+ true,
1012
+ false
1013
+ ],
1014
+ "range_flattens": [
1015
+ null,
1016
+ false,
1017
+ false
1018
+ ],
1019
+ "load_eviction_policies": [
1020
+ "",
1021
+ "",
1022
+ ""
1023
+ ],
1024
+ "num_warps": 8,
1025
+ "num_stages": 6,
1026
+ "indexing": [
1027
+ "tensor_descriptor",
1028
+ "pointer",
1029
+ "tensor_descriptor",
1030
+ "pointer",
1031
+ "pointer"
1032
+ ],
1033
+ "atomic_indexing": [],
1034
+ "pid_type": "flat"
1035
+ }
1036
+ },
1037
+ {
1038
+ "key": {
1039
+ "hidden_size": 4096,
1040
+ "num_tokens": 8
1041
+ },
1042
+ "config": {
1043
+ "block_sizes": [
1044
+ 4096,
1045
+ 2048
1046
+ ],
1047
+ "range_unroll_factors": [
1048
+ 0,
1049
+ 4,
1050
+ 0
1051
+ ],
1052
+ "range_warp_specializes": [],
1053
+ "range_num_stages": [],
1054
+ "range_multi_buffers": [
1055
+ null,
1056
+ null,
1057
+ null
1058
+ ],
1059
+ "range_flattens": [
1060
+ null,
1061
+ true,
1062
+ null
1063
+ ],
1064
+ "load_eviction_policies": [
1065
+ "last",
1066
+ "first",
1067
+ ""
1068
+ ],
1069
+ "num_warps": 8,
1070
+ "num_stages": 4,
1071
+ "indexing": [
1072
+ "pointer",
1073
+ "pointer",
1074
+ "tensor_descriptor",
1075
+ "pointer",
1076
+ "pointer"
1077
+ ],
1078
+ "atomic_indexing": [],
1079
+ "pid_type": "flat"
1080
+ }
1081
+ },
1082
+ {
1083
+ "key": {
1084
+ "hidden_size": 6144,
1085
+ "num_tokens": 8
1086
+ },
1087
+ "config": {
1088
+ "block_sizes": [
1089
+ 8192,
1090
+ 2048
1091
+ ],
1092
+ "range_unroll_factors": [
1093
+ 0,
1094
+ 4,
1095
+ 2
1096
+ ],
1097
+ "range_warp_specializes": [],
1098
+ "range_num_stages": [],
1099
+ "range_multi_buffers": [
1100
+ null,
1101
+ false,
1102
+ null
1103
+ ],
1104
+ "range_flattens": [
1105
+ null,
1106
+ false,
1107
+ null
1108
+ ],
1109
+ "load_eviction_policies": [
1110
+ "last",
1111
+ "last",
1112
+ "first"
1113
+ ],
1114
+ "num_warps": 8,
1115
+ "num_stages": 2,
1116
+ "indexing": [
1117
+ "tensor_descriptor",
1118
+ "tensor_descriptor",
1119
+ "tensor_descriptor",
1120
+ "tensor_descriptor",
1121
+ "pointer"
1122
+ ],
1123
+ "atomic_indexing": [],
1124
+ "pid_type": "flat"
1125
+ }
1126
+ },
1127
+ {
1128
+ "key": {
1129
+ "hidden_size": 8192,
1130
+ "num_tokens": 8
1131
+ },
1132
+ "config": {
1133
+ "block_sizes": [
1134
+ 8192,
1135
+ 4096
1136
+ ],
1137
+ "range_unroll_factors": [
1138
+ 0,
1139
+ 2,
1140
+ 2
1141
+ ],
1142
+ "range_warp_specializes": [],
1143
+ "range_num_stages": [],
1144
+ "range_multi_buffers": [
1145
+ null,
1146
+ null,
1147
+ true
1148
+ ],
1149
+ "range_flattens": [
1150
+ null,
1151
+ false,
1152
+ null
1153
+ ],
1154
+ "load_eviction_policies": [
1155
+ "last",
1156
+ "last",
1157
+ "first"
1158
+ ],
1159
+ "num_warps": 16,
1160
+ "num_stages": 4,
1161
+ "indexing": [
1162
+ "tensor_descriptor",
1163
+ "pointer",
1164
+ "pointer",
1165
+ "tensor_descriptor",
1166
+ "pointer"
1167
+ ],
1168
+ "atomic_indexing": [],
1169
+ "pid_type": "flat"
1170
+ }
1171
+ },
1172
+ {
1173
+ "key": {
1174
+ "hidden_size": 12288,
1175
+ "num_tokens": 8
1176
+ },
1177
+ "config": {
1178
+ "block_sizes": [
1179
+ 16384,
1180
+ 4096
1181
+ ],
1182
+ "range_unroll_factors": [
1183
+ 0,
1184
+ 0,
1185
+ 1
1186
+ ],
1187
+ "range_warp_specializes": [],
1188
+ "range_num_stages": [],
1189
+ "range_multi_buffers": [
1190
+ null,
1191
+ false,
1192
+ false
1193
+ ],
1194
+ "range_flattens": [
1195
+ null,
1196
+ false,
1197
+ null
1198
+ ],
1199
+ "load_eviction_policies": [
1200
+ "",
1201
+ "last",
1202
+ "last"
1203
+ ],
1204
+ "num_warps": 16,
1205
+ "num_stages": 5,
1206
+ "indexing": [
1207
+ "pointer",
1208
+ "tensor_descriptor",
1209
+ "tensor_descriptor",
1210
+ "pointer",
1211
+ "tensor_descriptor"
1212
+ ],
1213
+ "atomic_indexing": [],
1214
+ "pid_type": "flat"
1215
+ }
1216
+ },
1217
+ {
1218
+ "key": {
1219
+ "hidden_size": 28672,
1220
+ "num_tokens": 8
1221
+ },
1222
+ "config": {
1223
+ "block_sizes": [
1224
+ 32768,
1225
+ 4096
1226
+ ],
1227
+ "range_unroll_factors": [
1228
+ 0,
1229
+ 0,
1230
+ 2
1231
+ ],
1232
+ "range_warp_specializes": [],
1233
+ "range_num_stages": [],
1234
+ "range_multi_buffers": [
1235
+ null,
1236
+ false,
1237
+ true
1238
+ ],
1239
+ "range_flattens": [
1240
+ null,
1241
+ false,
1242
+ false
1243
+ ],
1244
+ "load_eviction_policies": [
1245
+ "",
1246
+ "last",
1247
+ "first"
1248
+ ],
1249
+ "num_warps": 16,
1250
+ "num_stages": 1,
1251
+ "indexing": [
1252
+ "pointer",
1253
+ "pointer",
1254
+ "pointer",
1255
+ "tensor_descriptor",
1256
+ "tensor_descriptor"
1257
+ ],
1258
+ "atomic_indexing": [],
1259
+ "pid_type": "flat"
1260
+ }
1261
+ },
1262
+ {
1263
+ "key": {
1264
+ "hidden_size": 512,
1265
+ "num_tokens": 16
1266
+ },
1267
+ "config": {
1268
+ "block_sizes": [
1269
+ 512,
1270
+ 128
1271
+ ],
1272
+ "range_unroll_factors": [
1273
+ 0,
1274
+ 1,
1275
+ 2
1276
+ ],
1277
+ "range_warp_specializes": [],
1278
+ "range_num_stages": [],
1279
+ "range_multi_buffers": [
1280
+ null,
1281
+ null,
1282
+ true
1283
+ ],
1284
+ "range_flattens": [
1285
+ null,
1286
+ true,
1287
+ true
1288
+ ],
1289
+ "load_eviction_policies": [
1290
+ "",
1291
+ "last",
1292
+ "last"
1293
+ ],
1294
+ "num_warps": 1,
1295
+ "num_stages": 7,
1296
+ "indexing": [
1297
+ "pointer",
1298
+ "pointer",
1299
+ "pointer",
1300
+ "pointer",
1301
+ "pointer"
1302
+ ],
1303
+ "atomic_indexing": [],
1304
+ "pid_type": "flat"
1305
+ }
1306
+ },
1307
+ {
1308
+ "key": {
1309
+ "hidden_size": 2048,
1310
+ "num_tokens": 16
1311
+ },
1312
+ "config": {
1313
+ "block_sizes": [
1314
+ 2048,
1315
+ 2048
1316
+ ],
1317
+ "range_unroll_factors": [
1318
+ 0,
1319
+ 2,
1320
+ 0
1321
+ ],
1322
+ "range_warp_specializes": [],
1323
+ "range_num_stages": [],
1324
+ "range_multi_buffers": [
1325
+ null,
1326
+ true,
1327
+ true
1328
+ ],
1329
+ "range_flattens": [
1330
+ null,
1331
+ true,
1332
+ false
1333
+ ],
1334
+ "load_eviction_policies": [
1335
+ "last",
1336
+ "",
1337
+ ""
1338
+ ],
1339
+ "num_warps": 16,
1340
+ "num_stages": 7,
1341
+ "indexing": [
1342
+ "tensor_descriptor",
1343
+ "pointer",
1344
+ "tensor_descriptor",
1345
+ "tensor_descriptor",
1346
+ "tensor_descriptor"
1347
+ ],
1348
+ "atomic_indexing": [],
1349
+ "pid_type": "flat"
1350
+ }
1351
+ },
1352
+ {
1353
+ "key": {
1354
+ "hidden_size": 4096,
1355
+ "num_tokens": 16
1356
+ },
1357
+ "config": {
1358
+ "block_sizes": [
1359
+ 4096,
1360
+ 2048
1361
+ ],
1362
+ "range_unroll_factors": [
1363
+ 0,
1364
+ 2,
1365
+ 2
1366
+ ],
1367
+ "range_warp_specializes": [],
1368
+ "range_num_stages": [],
1369
+ "range_multi_buffers": [
1370
+ null,
1371
+ null,
1372
+ null
1373
+ ],
1374
+ "range_flattens": [
1375
+ null,
1376
+ false,
1377
+ true
1378
+ ],
1379
+ "load_eviction_policies": [
1380
+ "",
1381
+ "",
1382
+ "last"
1383
+ ],
1384
+ "num_warps": 8,
1385
+ "num_stages": 3,
1386
+ "indexing": [
1387
+ "tensor_descriptor",
1388
+ "tensor_descriptor",
1389
+ "pointer",
1390
+ "tensor_descriptor",
1391
+ "tensor_descriptor"
1392
+ ],
1393
+ "atomic_indexing": [],
1394
+ "pid_type": "flat"
1395
+ }
1396
+ },
1397
+ {
1398
+ "key": {
1399
+ "hidden_size": 6144,
1400
+ "num_tokens": 16
1401
+ },
1402
+ "config": {
1403
+ "block_sizes": [
1404
+ 8192,
1405
+ 2048
1406
+ ],
1407
+ "range_unroll_factors": [
1408
+ 0,
1409
+ 4,
1410
+ 1
1411
+ ],
1412
+ "range_warp_specializes": [],
1413
+ "range_num_stages": [],
1414
+ "range_multi_buffers": [
1415
+ null,
1416
+ null,
1417
+ true
1418
+ ],
1419
+ "range_flattens": [
1420
+ null,
1421
+ true,
1422
+ null
1423
+ ],
1424
+ "load_eviction_policies": [
1425
+ "",
1426
+ "first",
1427
+ "last"
1428
+ ],
1429
+ "num_warps": 8,
1430
+ "num_stages": 5,
1431
+ "indexing": [
1432
+ "pointer",
1433
+ "pointer",
1434
+ "pointer",
1435
+ "pointer",
1436
+ "tensor_descriptor"
1437
+ ],
1438
+ "atomic_indexing": [],
1439
+ "pid_type": "flat"
1440
+ }
1441
+ },
1442
+ {
1443
+ "key": {
1444
+ "hidden_size": 8192,
1445
+ "num_tokens": 16
1446
+ },
1447
+ "config": {
1448
+ "block_sizes": [
1449
+ 8192,
1450
+ 4096
1451
+ ],
1452
+ "range_unroll_factors": [
1453
+ 0,
1454
+ 1,
1455
+ 0
1456
+ ],
1457
+ "range_warp_specializes": [],
1458
+ "range_num_stages": [],
1459
+ "range_multi_buffers": [
1460
+ null,
1461
+ true,
1462
+ true
1463
+ ],
1464
+ "range_flattens": [
1465
+ null,
1466
+ false,
1467
+ null
1468
+ ],
1469
+ "load_eviction_policies": [
1470
+ "",
1471
+ "last",
1472
+ "last"
1473
+ ],
1474
+ "num_warps": 16,
1475
+ "num_stages": 8,
1476
+ "indexing": [
1477
+ "tensor_descriptor",
1478
+ "pointer",
1479
+ "pointer",
1480
+ "pointer",
1481
+ "tensor_descriptor"
1482
+ ],
1483
+ "atomic_indexing": [],
1484
+ "pid_type": "flat"
1485
+ }
1486
+ },
1487
+ {
1488
+ "key": {
1489
+ "hidden_size": 12288,
1490
+ "num_tokens": 16
1491
+ },
1492
+ "config": {
1493
+ "block_sizes": [
1494
+ 16384,
1495
+ 4096
1496
+ ],
1497
+ "range_unroll_factors": [
1498
+ 0,
1499
+ 0,
1500
+ 1
1501
+ ],
1502
+ "range_warp_specializes": [],
1503
+ "range_num_stages": [],
1504
+ "range_multi_buffers": [
1505
+ null,
1506
+ false,
1507
+ true
1508
+ ],
1509
+ "range_flattens": [
1510
+ null,
1511
+ true,
1512
+ false
1513
+ ],
1514
+ "load_eviction_policies": [
1515
+ "",
1516
+ "last",
1517
+ "last"
1518
+ ],
1519
+ "num_warps": 16,
1520
+ "num_stages": 4,
1521
+ "indexing": [
1522
+ "tensor_descriptor",
1523
+ "tensor_descriptor",
1524
+ "pointer",
1525
+ "pointer",
1526
+ "pointer"
1527
+ ],
1528
+ "atomic_indexing": [],
1529
+ "pid_type": "flat"
1530
+ }
1531
+ },
1532
+ {
1533
+ "key": {
1534
+ "hidden_size": 28672,
1535
+ "num_tokens": 16
1536
+ },
1537
+ "config": {
1538
+ "block_sizes": [
1539
+ 32768,
1540
+ 4096
1541
+ ],
1542
+ "range_unroll_factors": [
1543
+ 0,
1544
+ 3,
1545
+ 3
1546
+ ],
1547
+ "range_warp_specializes": [],
1548
+ "range_num_stages": [],
1549
+ "range_multi_buffers": [
1550
+ null,
1551
+ false,
1552
+ false
1553
+ ],
1554
+ "range_flattens": [
1555
+ null,
1556
+ false,
1557
+ null
1558
+ ],
1559
+ "load_eviction_policies": [
1560
+ "",
1561
+ "",
1562
+ ""
1563
+ ],
1564
+ "num_warps": 16,
1565
+ "num_stages": 1,
1566
+ "indexing": [
1567
+ "pointer",
1568
+ "tensor_descriptor",
1569
+ "pointer",
1570
+ "tensor_descriptor",
1571
+ "tensor_descriptor"
1572
+ ],
1573
+ "atomic_indexing": [],
1574
+ "pid_type": "flat"
1575
+ }
1576
+ },
1577
+ {
1578
+ "key": {
1579
+ "hidden_size": 512,
1580
+ "num_tokens": 32
1581
+ },
1582
+ "config": {
1583
+ "block_sizes": [
1584
+ 512,
1585
+ 256
1586
+ ],
1587
+ "range_unroll_factors": [
1588
+ 0,
1589
+ 3,
1590
+ 4
1591
+ ],
1592
+ "range_warp_specializes": [],
1593
+ "range_num_stages": [],
1594
+ "range_multi_buffers": [
1595
+ null,
1596
+ true,
1597
+ false
1598
+ ],
1599
+ "range_flattens": [
1600
+ null,
1601
+ true,
1602
+ true
1603
+ ],
1604
+ "load_eviction_policies": [
1605
+ "first",
1606
+ "last",
1607
+ "last"
1608
+ ],
1609
+ "num_warps": 1,
1610
+ "num_stages": 6,
1611
+ "indexing": [
1612
+ "tensor_descriptor",
1613
+ "pointer",
1614
+ "pointer",
1615
+ "tensor_descriptor",
1616
+ "tensor_descriptor"
1617
+ ],
1618
+ "atomic_indexing": [],
1619
+ "pid_type": "flat"
1620
+ }
1621
+ },
1622
+ {
1623
+ "key": {
1624
+ "hidden_size": 2048,
1625
+ "num_tokens": 32
1626
+ },
1627
+ "config": {
1628
+ "block_sizes": [
1629
+ 2048,
1630
+ 2048
1631
+ ],
1632
+ "range_unroll_factors": [
1633
+ 0,
1634
+ 0,
1635
+ 1
1636
+ ],
1637
+ "range_warp_specializes": [],
1638
+ "range_num_stages": [],
1639
+ "range_multi_buffers": [
1640
+ null,
1641
+ true,
1642
+ null
1643
+ ],
1644
+ "range_flattens": [
1645
+ null,
1646
+ null,
1647
+ true
1648
+ ],
1649
+ "load_eviction_policies": [
1650
+ "",
1651
+ "first",
1652
+ "first"
1653
+ ],
1654
+ "num_warps": 8,
1655
+ "num_stages": 5,
1656
+ "indexing": [
1657
+ "pointer",
1658
+ "pointer",
1659
+ "tensor_descriptor",
1660
+ "pointer",
1661
+ "pointer"
1662
+ ],
1663
+ "atomic_indexing": [],
1664
+ "pid_type": "flat"
1665
+ }
1666
+ },
1667
+ {
1668
+ "key": {
1669
+ "hidden_size": 4096,
1670
+ "num_tokens": 32
1671
+ },
1672
+ "config": {
1673
+ "block_sizes": [
1674
+ 4096,
1675
+ 2048
1676
+ ],
1677
+ "range_unroll_factors": [
1678
+ 0,
1679
+ 1,
1680
+ 2
1681
+ ],
1682
+ "range_warp_specializes": [],
1683
+ "range_num_stages": [],
1684
+ "range_multi_buffers": [
1685
+ null,
1686
+ true,
1687
+ null
1688
+ ],
1689
+ "range_flattens": [
1690
+ null,
1691
+ null,
1692
+ true
1693
+ ],
1694
+ "load_eviction_policies": [
1695
+ "last",
1696
+ "",
1697
+ "last"
1698
+ ],
1699
+ "num_warps": 8,
1700
+ "num_stages": 8,
1701
+ "indexing": [
1702
+ "tensor_descriptor",
1703
+ "tensor_descriptor",
1704
+ "pointer",
1705
+ "tensor_descriptor",
1706
+ "pointer"
1707
+ ],
1708
+ "atomic_indexing": [],
1709
+ "pid_type": "flat"
1710
+ }
1711
+ },
1712
+ {
1713
+ "key": {
1714
+ "hidden_size": 6144,
1715
+ "num_tokens": 32
1716
+ },
1717
+ "config": {
1718
+ "block_sizes": [
1719
+ 8192,
1720
+ 2048
1721
+ ],
1722
+ "range_unroll_factors": [
1723
+ 0,
1724
+ 1,
1725
+ 1
1726
+ ],
1727
+ "range_warp_specializes": [],
1728
+ "range_num_stages": [],
1729
+ "range_multi_buffers": [
1730
+ null,
1731
+ true,
1732
+ false
1733
+ ],
1734
+ "range_flattens": [
1735
+ null,
1736
+ true,
1737
+ null
1738
+ ],
1739
+ "load_eviction_policies": [
1740
+ "last",
1741
+ "",
1742
+ ""
1743
+ ],
1744
+ "num_warps": 8,
1745
+ "num_stages": 2,
1746
+ "indexing": [
1747
+ "pointer",
1748
+ "tensor_descriptor",
1749
+ "tensor_descriptor",
1750
+ "tensor_descriptor",
1751
+ "pointer"
1752
+ ],
1753
+ "atomic_indexing": [],
1754
+ "pid_type": "flat"
1755
+ }
1756
+ },
1757
+ {
1758
+ "key": {
1759
+ "hidden_size": 8192,
1760
+ "num_tokens": 32
1761
+ },
1762
+ "config": {
1763
+ "block_sizes": [
1764
+ 8192,
1765
+ 4096
1766
+ ],
1767
+ "range_unroll_factors": [
1768
+ 0,
1769
+ 3,
1770
+ 4
1771
+ ],
1772
+ "range_warp_specializes": [],
1773
+ "range_num_stages": [],
1774
+ "range_multi_buffers": [
1775
+ null,
1776
+ null,
1777
+ null
1778
+ ],
1779
+ "range_flattens": [
1780
+ null,
1781
+ null,
1782
+ null
1783
+ ],
1784
+ "load_eviction_policies": [
1785
+ "last",
1786
+ "first",
1787
+ "last"
1788
+ ],
1789
+ "num_warps": 16,
1790
+ "num_stages": 2,
1791
+ "indexing": [
1792
+ "pointer",
1793
+ "pointer",
1794
+ "pointer",
1795
+ "tensor_descriptor",
1796
+ "tensor_descriptor"
1797
+ ],
1798
+ "atomic_indexing": [],
1799
+ "pid_type": "flat"
1800
+ }
1801
+ },
1802
+ {
1803
+ "key": {
1804
+ "hidden_size": 12288,
1805
+ "num_tokens": 32
1806
+ },
1807
+ "config": {
1808
+ "block_sizes": [
1809
+ 16384,
1810
+ 4096
1811
+ ],
1812
+ "range_unroll_factors": [
1813
+ 0,
1814
+ 1,
1815
+ 0
1816
+ ],
1817
+ "range_warp_specializes": [],
1818
+ "range_num_stages": [],
1819
+ "range_multi_buffers": [
1820
+ null,
1821
+ null,
1822
+ true
1823
+ ],
1824
+ "range_flattens": [
1825
+ null,
1826
+ null,
1827
+ true
1828
+ ],
1829
+ "load_eviction_policies": [
1830
+ "",
1831
+ "last",
1832
+ "last"
1833
+ ],
1834
+ "num_warps": 16,
1835
+ "num_stages": 5,
1836
+ "indexing": [
1837
+ "tensor_descriptor",
1838
+ "pointer",
1839
+ "tensor_descriptor",
1840
+ "pointer",
1841
+ "pointer"
1842
+ ],
1843
+ "atomic_indexing": [],
1844
+ "pid_type": "flat"
1845
+ }
1846
+ },
1847
+ {
1848
+ "key": {
1849
+ "hidden_size": 28672,
1850
+ "num_tokens": 32
1851
+ },
1852
+ "config": {
1853
+ "block_sizes": [
1854
+ 32768,
1855
+ 4096
1856
+ ],
1857
+ "range_unroll_factors": [
1858
+ 0,
1859
+ 3,
1860
+ 3
1861
+ ],
1862
+ "range_warp_specializes": [],
1863
+ "range_num_stages": [],
1864
+ "range_multi_buffers": [
1865
+ null,
1866
+ false,
1867
+ false
1868
+ ],
1869
+ "range_flattens": [
1870
+ null,
1871
+ false,
1872
+ null
1873
+ ],
1874
+ "load_eviction_policies": [
1875
+ "",
1876
+ "",
1877
+ ""
1878
+ ],
1879
+ "num_warps": 16,
1880
+ "num_stages": 1,
1881
+ "indexing": [
1882
+ "pointer",
1883
+ "tensor_descriptor",
1884
+ "pointer",
1885
+ "tensor_descriptor",
1886
+ "tensor_descriptor"
1887
+ ],
1888
+ "atomic_indexing": [],
1889
+ "pid_type": "flat"
1890
+ }
1891
+ },
1892
+ {
1893
+ "key": {
1894
+ "hidden_size": 512,
1895
+ "num_tokens": 64
1896
+ },
1897
+ "config": {
1898
+ "block_sizes": [
1899
+ 512,
1900
+ 256
1901
+ ],
1902
+ "range_unroll_factors": [
1903
+ 0,
1904
+ 2,
1905
+ 2
1906
+ ],
1907
+ "range_warp_specializes": [],
1908
+ "range_num_stages": [],
1909
+ "range_multi_buffers": [
1910
+ null,
1911
+ false,
1912
+ true
1913
+ ],
1914
+ "range_flattens": [
1915
+ null,
1916
+ false,
1917
+ false
1918
+ ],
1919
+ "load_eviction_policies": [
1920
+ "first",
1921
+ "first",
1922
+ "last"
1923
+ ],
1924
+ "num_warps": 1,
1925
+ "num_stages": 2,
1926
+ "indexing": [
1927
+ "tensor_descriptor",
1928
+ "pointer",
1929
+ "pointer",
1930
+ "pointer",
1931
+ "tensor_descriptor"
1932
+ ],
1933
+ "atomic_indexing": [],
1934
+ "pid_type": "flat"
1935
+ }
1936
+ },
1937
+ {
1938
+ "key": {
1939
+ "hidden_size": 2048,
1940
+ "num_tokens": 64
1941
+ },
1942
+ "config": {
1943
+ "block_sizes": [
1944
+ 2048,
1945
+ 2048
1946
+ ],
1947
+ "range_unroll_factors": [
1948
+ 0,
1949
+ 4,
1950
+ 2
1951
+ ],
1952
+ "range_warp_specializes": [],
1953
+ "range_num_stages": [],
1954
+ "range_multi_buffers": [
1955
+ null,
1956
+ false,
1957
+ false
1958
+ ],
1959
+ "range_flattens": [
1960
+ null,
1961
+ false,
1962
+ null
1963
+ ],
1964
+ "load_eviction_policies": [
1965
+ "",
1966
+ "first",
1967
+ "last"
1968
+ ],
1969
+ "num_warps": 8,
1970
+ "num_stages": 3,
1971
+ "indexing": [
1972
+ "pointer",
1973
+ "tensor_descriptor",
1974
+ "tensor_descriptor",
1975
+ "pointer",
1976
+ "tensor_descriptor"
1977
+ ],
1978
+ "atomic_indexing": [],
1979
+ "pid_type": "flat"
1980
+ }
1981
+ },
1982
+ {
1983
+ "key": {
1984
+ "hidden_size": 4096,
1985
+ "num_tokens": 64
1986
+ },
1987
+ "config": {
1988
+ "block_sizes": [
1989
+ 4096,
1990
+ 4096
1991
+ ],
1992
+ "range_unroll_factors": [
1993
+ 0,
1994
+ 2,
1995
+ 3
1996
+ ],
1997
+ "range_warp_specializes": [],
1998
+ "range_num_stages": [],
1999
+ "range_multi_buffers": [
2000
+ null,
2001
+ true,
2002
+ true
2003
+ ],
2004
+ "range_flattens": [
2005
+ null,
2006
+ false,
2007
+ false
2008
+ ],
2009
+ "load_eviction_policies": [
2010
+ "",
2011
+ "first",
2012
+ "first"
2013
+ ],
2014
+ "num_warps": 16,
2015
+ "num_stages": 2,
2016
+ "indexing": [
2017
+ "tensor_descriptor",
2018
+ "pointer",
2019
+ "tensor_descriptor",
2020
+ "tensor_descriptor",
2021
+ "pointer"
2022
+ ],
2023
+ "atomic_indexing": [],
2024
+ "pid_type": "flat"
2025
+ }
2026
+ },
2027
+ {
2028
+ "key": {
2029
+ "hidden_size": 6144,
2030
+ "num_tokens": 64
2031
+ },
2032
+ "config": {
2033
+ "block_sizes": [
2034
+ 8192,
2035
+ 2048
2036
+ ],
2037
+ "range_unroll_factors": [
2038
+ 0,
2039
+ 2,
2040
+ 4
2041
+ ],
2042
+ "range_warp_specializes": [],
2043
+ "range_num_stages": [],
2044
+ "range_multi_buffers": [
2045
+ null,
2046
+ true,
2047
+ true
2048
+ ],
2049
+ "range_flattens": [
2050
+ null,
2051
+ true,
2052
+ null
2053
+ ],
2054
+ "load_eviction_policies": [
2055
+ "",
2056
+ "",
2057
+ ""
2058
+ ],
2059
+ "num_warps": 8,
2060
+ "num_stages": 7,
2061
+ "indexing": [
2062
+ "tensor_descriptor",
2063
+ "pointer",
2064
+ "pointer",
2065
+ "pointer",
2066
+ "tensor_descriptor"
2067
+ ],
2068
+ "atomic_indexing": [],
2069
+ "pid_type": "flat"
2070
+ }
2071
+ },
2072
+ {
2073
+ "key": {
2074
+ "hidden_size": 8192,
2075
+ "num_tokens": 64
2076
+ },
2077
+ "config": {
2078
+ "block_sizes": [
2079
+ 8192,
2080
+ 4096
2081
+ ],
2082
+ "range_unroll_factors": [
2083
+ 0,
2084
+ 1,
2085
+ 2
2086
+ ],
2087
+ "range_warp_specializes": [],
2088
+ "range_num_stages": [],
2089
+ "range_multi_buffers": [
2090
+ null,
2091
+ false,
2092
+ null
2093
+ ],
2094
+ "range_flattens": [
2095
+ null,
2096
+ null,
2097
+ false
2098
+ ],
2099
+ "load_eviction_policies": [
2100
+ "first",
2101
+ "last",
2102
+ "last"
2103
+ ],
2104
+ "num_warps": 16,
2105
+ "num_stages": 3,
2106
+ "indexing": [
2107
+ "pointer",
2108
+ "tensor_descriptor",
2109
+ "pointer",
2110
+ "pointer",
2111
+ "tensor_descriptor"
2112
+ ],
2113
+ "atomic_indexing": [],
2114
+ "pid_type": "flat"
2115
+ }
2116
+ },
2117
+ {
2118
+ "key": {
2119
+ "hidden_size": 12288,
2120
+ "num_tokens": 64
2121
+ },
2122
+ "config": {
2123
+ "block_sizes": [
2124
+ 16384,
2125
+ 4096
2126
+ ],
2127
+ "range_unroll_factors": [
2128
+ 0,
2129
+ 4,
2130
+ 3
2131
+ ],
2132
+ "range_warp_specializes": [],
2133
+ "range_num_stages": [],
2134
+ "range_multi_buffers": [
2135
+ null,
2136
+ true,
2137
+ true
2138
+ ],
2139
+ "range_flattens": [
2140
+ null,
2141
+ null,
2142
+ true
2143
+ ],
2144
+ "load_eviction_policies": [
2145
+ "",
2146
+ "",
2147
+ "first"
2148
+ ],
2149
+ "num_warps": 16,
2150
+ "num_stages": 2,
2151
+ "indexing": [
2152
+ "tensor_descriptor",
2153
+ "tensor_descriptor",
2154
+ "pointer",
2155
+ "tensor_descriptor",
2156
+ "pointer"
2157
+ ],
2158
+ "atomic_indexing": [],
2159
+ "pid_type": "flat"
2160
+ }
2161
+ },
2162
+ {
2163
+ "key": {
2164
+ "hidden_size": 28672,
2165
+ "num_tokens": 64
2166
+ },
2167
+ "config": {
2168
+ "block_sizes": [
2169
+ 32768,
2170
+ 4096
2171
+ ],
2172
+ "range_unroll_factors": [
2173
+ 0,
2174
+ 3,
2175
+ 3
2176
+ ],
2177
+ "range_warp_specializes": [],
2178
+ "range_num_stages": [],
2179
+ "range_multi_buffers": [
2180
+ null,
2181
+ false,
2182
+ false
2183
+ ],
2184
+ "range_flattens": [
2185
+ null,
2186
+ false,
2187
+ null
2188
+ ],
2189
+ "load_eviction_policies": [
2190
+ "",
2191
+ "",
2192
+ ""
2193
+ ],
2194
+ "num_warps": 16,
2195
+ "num_stages": 1,
2196
+ "indexing": [
2197
+ "pointer",
2198
+ "tensor_descriptor",
2199
+ "pointer",
2200
+ "tensor_descriptor",
2201
+ "tensor_descriptor"
2202
+ ],
2203
+ "atomic_indexing": [],
2204
+ "pid_type": "flat"
2205
+ }
2206
+ },
2207
+ {
2208
+ "key": {
2209
+ "hidden_size": 512,
2210
+ "num_tokens": 128
2211
+ },
2212
+ "config": {
2213
+ "block_sizes": [
2214
+ 256,
2215
+ 128
2216
+ ],
2217
+ "range_unroll_factors": [
2218
+ 0,
2219
+ 0,
2220
+ 3
2221
+ ],
2222
+ "range_warp_specializes": [],
2223
+ "range_num_stages": [],
2224
+ "range_multi_buffers": [
2225
+ null,
2226
+ true,
2227
+ null
2228
+ ],
2229
+ "range_flattens": [
2230
+ null,
2231
+ true,
2232
+ false
2233
+ ],
2234
+ "load_eviction_policies": [
2235
+ "last",
2236
+ "first",
2237
+ "first"
2238
+ ],
2239
+ "num_warps": 1,
2240
+ "num_stages": 7,
2241
+ "indexing": [
2242
+ "pointer",
2243
+ "pointer",
2244
+ "tensor_descriptor",
2245
+ "tensor_descriptor",
2246
+ "pointer"
2247
+ ],
2248
+ "atomic_indexing": [],
2249
+ "pid_type": "flat"
2250
+ }
2251
+ },
2252
+ {
2253
+ "key": {
2254
+ "hidden_size": 2048,
2255
+ "num_tokens": 128
2256
+ },
2257
+ "config": {
2258
+ "block_sizes": [
2259
+ 2048,
2260
+ 1024
2261
+ ],
2262
+ "range_unroll_factors": [
2263
+ 0,
2264
+ 3,
2265
+ 3
2266
+ ],
2267
+ "range_warp_specializes": [],
2268
+ "range_num_stages": [],
2269
+ "range_multi_buffers": [
2270
+ null,
2271
+ true,
2272
+ false
2273
+ ],
2274
+ "range_flattens": [
2275
+ null,
2276
+ false,
2277
+ true
2278
+ ],
2279
+ "load_eviction_policies": [
2280
+ "",
2281
+ "",
2282
+ "last"
2283
+ ],
2284
+ "num_warps": 8,
2285
+ "num_stages": 6,
2286
+ "indexing": [
2287
+ "tensor_descriptor",
2288
+ "pointer",
2289
+ "tensor_descriptor",
2290
+ "pointer",
2291
+ "pointer"
2292
+ ],
2293
+ "atomic_indexing": [],
2294
+ "pid_type": "flat"
2295
+ }
2296
+ },
2297
+ {
2298
+ "key": {
2299
+ "hidden_size": 4096,
2300
+ "num_tokens": 128
2301
+ },
2302
+ "config": {
2303
+ "block_sizes": [
2304
+ 4096,
2305
+ 2048
2306
+ ],
2307
+ "range_unroll_factors": [
2308
+ 0,
2309
+ 1,
2310
+ 4
2311
+ ],
2312
+ "range_warp_specializes": [],
2313
+ "range_num_stages": [],
2314
+ "range_multi_buffers": [
2315
+ null,
2316
+ true,
2317
+ null
2318
+ ],
2319
+ "range_flattens": [
2320
+ null,
2321
+ false,
2322
+ false
2323
+ ],
2324
+ "load_eviction_policies": [
2325
+ "last",
2326
+ "first",
2327
+ "last"
2328
+ ],
2329
+ "num_warps": 16,
2330
+ "num_stages": 8,
2331
+ "indexing": [
2332
+ "tensor_descriptor",
2333
+ "tensor_descriptor",
2334
+ "tensor_descriptor",
2335
+ "tensor_descriptor",
2336
+ "tensor_descriptor"
2337
+ ],
2338
+ "atomic_indexing": [],
2339
+ "pid_type": "flat"
2340
+ }
2341
+ },
2342
+ {
2343
+ "key": {
2344
+ "hidden_size": 6144,
2345
+ "num_tokens": 128
2346
+ },
2347
+ "config": {
2348
+ "block_sizes": [
2349
+ 8192,
2350
+ 2048
2351
+ ],
2352
+ "range_unroll_factors": [
2353
+ 0,
2354
+ 3,
2355
+ 3
2356
+ ],
2357
+ "range_warp_specializes": [],
2358
+ "range_num_stages": [],
2359
+ "range_multi_buffers": [
2360
+ null,
2361
+ false,
2362
+ true
2363
+ ],
2364
+ "range_flattens": [
2365
+ null,
2366
+ null,
2367
+ true
2368
+ ],
2369
+ "load_eviction_policies": [
2370
+ "last",
2371
+ "first",
2372
+ "last"
2373
+ ],
2374
+ "num_warps": 16,
2375
+ "num_stages": 4,
2376
+ "indexing": [
2377
+ "pointer",
2378
+ "tensor_descriptor",
2379
+ "tensor_descriptor",
2380
+ "tensor_descriptor",
2381
+ "pointer"
2382
+ ],
2383
+ "atomic_indexing": [],
2384
+ "pid_type": "flat"
2385
+ }
2386
+ },
2387
+ {
2388
+ "key": {
2389
+ "hidden_size": 8192,
2390
+ "num_tokens": 128
2391
+ },
2392
+ "config": {
2393
+ "block_sizes": [
2394
+ 16384,
2395
+ 4096
2396
+ ],
2397
+ "range_unroll_factors": [
2398
+ 0,
2399
+ 0,
2400
+ 2
2401
+ ],
2402
+ "range_warp_specializes": [],
2403
+ "range_num_stages": [],
2404
+ "range_multi_buffers": [
2405
+ null,
2406
+ false,
2407
+ false
2408
+ ],
2409
+ "range_flattens": [
2410
+ null,
2411
+ null,
2412
+ true
2413
+ ],
2414
+ "load_eviction_policies": [
2415
+ "",
2416
+ "",
2417
+ "first"
2418
+ ],
2419
+ "num_warps": 16,
2420
+ "num_stages": 1,
2421
+ "indexing": [
2422
+ "tensor_descriptor",
2423
+ "pointer",
2424
+ "pointer",
2425
+ "tensor_descriptor",
2426
+ "pointer"
2427
+ ],
2428
+ "atomic_indexing": [],
2429
+ "pid_type": "flat"
2430
+ }
2431
+ },
2432
+ {
2433
+ "key": {
2434
+ "hidden_size": 12288,
2435
+ "num_tokens": 128
2436
+ },
2437
+ "config": {
2438
+ "block_sizes": [
2439
+ 16384,
2440
+ 4096
2441
+ ],
2442
+ "range_unroll_factors": [
2443
+ 0,
2444
+ 1,
2445
+ 3
2446
+ ],
2447
+ "range_warp_specializes": [],
2448
+ "range_num_stages": [],
2449
+ "range_multi_buffers": [
2450
+ null,
2451
+ null,
2452
+ false
2453
+ ],
2454
+ "range_flattens": [
2455
+ null,
2456
+ true,
2457
+ true
2458
+ ],
2459
+ "load_eviction_policies": [
2460
+ "first",
2461
+ "first",
2462
+ "first"
2463
+ ],
2464
+ "num_warps": 16,
2465
+ "num_stages": 3,
2466
+ "indexing": [
2467
+ "pointer",
2468
+ "pointer",
2469
+ "tensor_descriptor",
2470
+ "tensor_descriptor",
2471
+ "tensor_descriptor"
2472
+ ],
2473
+ "atomic_indexing": [],
2474
+ "pid_type": "flat"
2475
+ }
2476
+ },
2477
+ {
2478
+ "key": {
2479
+ "hidden_size": 28672,
2480
+ "num_tokens": 128
2481
+ },
2482
+ "config": {
2483
+ "block_sizes": [
2484
+ 32768,
2485
+ 4096
2486
+ ],
2487
+ "range_unroll_factors": [
2488
+ 0,
2489
+ 3,
2490
+ 0
2491
+ ],
2492
+ "range_warp_specializes": [],
2493
+ "range_num_stages": [],
2494
+ "range_multi_buffers": [
2495
+ null,
2496
+ true,
2497
+ false
2498
+ ],
2499
+ "range_flattens": [
2500
+ null,
2501
+ null,
2502
+ true
2503
+ ],
2504
+ "load_eviction_policies": [
2505
+ "",
2506
+ "first",
2507
+ ""
2508
+ ],
2509
+ "num_warps": 32,
2510
+ "num_stages": 7,
2511
+ "indexing": [
2512
+ "pointer",
2513
+ "pointer",
2514
+ "tensor_descriptor",
2515
+ "tensor_descriptor",
2516
+ "pointer"
2517
+ ],
2518
+ "atomic_indexing": [],
2519
+ "pid_type": "flat"
2520
+ }
2521
+ },
2522
+ {
2523
+ "key": {
2524
+ "hidden_size": 512,
2525
+ "num_tokens": 256
2526
+ },
2527
+ "config": {
2528
+ "block_sizes": [
2529
+ 512,
2530
+ 512
2531
+ ],
2532
+ "range_unroll_factors": [
2533
+ 0,
2534
+ 4,
2535
+ 0
2536
+ ],
2537
+ "range_warp_specializes": [],
2538
+ "range_num_stages": [],
2539
+ "range_multi_buffers": [
2540
+ null,
2541
+ true,
2542
+ false
2543
+ ],
2544
+ "range_flattens": [
2545
+ null,
2546
+ null,
2547
+ null
2548
+ ],
2549
+ "load_eviction_policies": [
2550
+ "",
2551
+ "",
2552
+ ""
2553
+ ],
2554
+ "num_warps": 1,
2555
+ "num_stages": 6,
2556
+ "indexing": [
2557
+ "pointer",
2558
+ "pointer",
2559
+ "tensor_descriptor",
2560
+ "pointer",
2561
+ "pointer"
2562
+ ],
2563
+ "atomic_indexing": [],
2564
+ "pid_type": "flat"
2565
+ }
2566
+ },
2567
+ {
2568
+ "key": {
2569
+ "hidden_size": 2048,
2570
+ "num_tokens": 256
2571
+ },
2572
+ "config": {
2573
+ "block_sizes": [
2574
+ 8192,
2575
+ 2048
2576
+ ],
2577
+ "range_unroll_factors": [
2578
+ 0,
2579
+ 4,
2580
+ 1
2581
+ ],
2582
+ "range_warp_specializes": [],
2583
+ "range_num_stages": [],
2584
+ "range_multi_buffers": [
2585
+ null,
2586
+ null,
2587
+ true
2588
+ ],
2589
+ "range_flattens": [
2590
+ null,
2591
+ true,
2592
+ null
2593
+ ],
2594
+ "load_eviction_policies": [
2595
+ "",
2596
+ "first",
2597
+ "last"
2598
+ ],
2599
+ "num_warps": 8,
2600
+ "num_stages": 5,
2601
+ "indexing": [
2602
+ "pointer",
2603
+ "pointer",
2604
+ "pointer",
2605
+ "pointer",
2606
+ "tensor_descriptor"
2607
+ ],
2608
+ "atomic_indexing": [],
2609
+ "pid_type": "flat"
2610
+ }
2611
+ },
2612
+ {
2613
+ "key": {
2614
+ "hidden_size": 4096,
2615
+ "num_tokens": 256
2616
+ },
2617
+ "config": {
2618
+ "block_sizes": [
2619
+ 4096,
2620
+ 2048
2621
+ ],
2622
+ "range_unroll_factors": [
2623
+ 0,
2624
+ 3,
2625
+ 2
2626
+ ],
2627
+ "range_warp_specializes": [],
2628
+ "range_num_stages": [],
2629
+ "range_multi_buffers": [
2630
+ null,
2631
+ false,
2632
+ true
2633
+ ],
2634
+ "range_flattens": [
2635
+ null,
2636
+ false,
2637
+ null
2638
+ ],
2639
+ "load_eviction_policies": [
2640
+ "first",
2641
+ "last",
2642
+ ""
2643
+ ],
2644
+ "num_warps": 8,
2645
+ "num_stages": 2,
2646
+ "indexing": [
2647
+ "pointer",
2648
+ "tensor_descriptor",
2649
+ "tensor_descriptor",
2650
+ "tensor_descriptor",
2651
+ "pointer"
2652
+ ],
2653
+ "atomic_indexing": [],
2654
+ "pid_type": "flat"
2655
+ }
2656
+ },
2657
+ {
2658
+ "key": {
2659
+ "hidden_size": 6144,
2660
+ "num_tokens": 256
2661
+ },
2662
+ "config": {
2663
+ "block_sizes": [
2664
+ 8192,
2665
+ 2048
2666
+ ],
2667
+ "range_unroll_factors": [
2668
+ 0,
2669
+ 3,
2670
+ 1
2671
+ ],
2672
+ "range_warp_specializes": [],
2673
+ "range_num_stages": [],
2674
+ "range_multi_buffers": [
2675
+ null,
2676
+ null,
2677
+ false
2678
+ ],
2679
+ "range_flattens": [
2680
+ null,
2681
+ false,
2682
+ true
2683
+ ],
2684
+ "load_eviction_policies": [
2685
+ "last",
2686
+ "last",
2687
+ ""
2688
+ ],
2689
+ "num_warps": 8,
2690
+ "num_stages": 1,
2691
+ "indexing": [
2692
+ "pointer",
2693
+ "tensor_descriptor",
2694
+ "tensor_descriptor",
2695
+ "tensor_descriptor",
2696
+ "tensor_descriptor"
2697
+ ],
2698
+ "atomic_indexing": [],
2699
+ "pid_type": "flat"
2700
+ }
2701
+ },
2702
+ {
2703
+ "key": {
2704
+ "hidden_size": 8192,
2705
+ "num_tokens": 256
2706
+ },
2707
+ "config": {
2708
+ "block_sizes": [
2709
+ 8192,
2710
+ 2048
2711
+ ],
2712
+ "range_unroll_factors": [
2713
+ 0,
2714
+ 2,
2715
+ 3
2716
+ ],
2717
+ "range_warp_specializes": [],
2718
+ "range_num_stages": [],
2719
+ "range_multi_buffers": [
2720
+ null,
2721
+ false,
2722
+ false
2723
+ ],
2724
+ "range_flattens": [
2725
+ null,
2726
+ null,
2727
+ true
2728
+ ],
2729
+ "load_eviction_policies": [
2730
+ "",
2731
+ "",
2732
+ ""
2733
+ ],
2734
+ "num_warps": 4,
2735
+ "num_stages": 7,
2736
+ "indexing": [
2737
+ "pointer",
2738
+ "pointer",
2739
+ "pointer",
2740
+ "pointer",
2741
+ "pointer"
2742
+ ],
2743
+ "atomic_indexing": [],
2744
+ "pid_type": "flat"
2745
+ }
2746
+ },
2747
+ {
2748
+ "key": {
2749
+ "hidden_size": 12288,
2750
+ "num_tokens": 256
2751
+ },
2752
+ "config": {
2753
+ "block_sizes": [
2754
+ 16384,
2755
+ 4096
2756
+ ],
2757
+ "range_unroll_factors": [
2758
+ 0,
2759
+ 3,
2760
+ 2
2761
+ ],
2762
+ "range_warp_specializes": [],
2763
+ "range_num_stages": [],
2764
+ "range_multi_buffers": [
2765
+ null,
2766
+ null,
2767
+ true
2768
+ ],
2769
+ "range_flattens": [
2770
+ null,
2771
+ null,
2772
+ null
2773
+ ],
2774
+ "load_eviction_policies": [
2775
+ "last",
2776
+ "first",
2777
+ "first"
2778
+ ],
2779
+ "num_warps": 16,
2780
+ "num_stages": 8,
2781
+ "indexing": [
2782
+ "tensor_descriptor",
2783
+ "tensor_descriptor",
2784
+ "pointer",
2785
+ "pointer",
2786
+ "tensor_descriptor"
2787
+ ],
2788
+ "atomic_indexing": [],
2789
+ "pid_type": "flat"
2790
+ }
2791
+ },
2792
+ {
2793
+ "key": {
2794
+ "hidden_size": 28672,
2795
+ "num_tokens": 256
2796
+ },
2797
+ "config": {
2798
+ "block_sizes": [
2799
+ 32768,
2800
+ 16384
2801
+ ],
2802
+ "range_unroll_factors": [
2803
+ 0,
2804
+ 4,
2805
+ 2
2806
+ ],
2807
+ "range_warp_specializes": [],
2808
+ "range_num_stages": [],
2809
+ "range_multi_buffers": [
2810
+ null,
2811
+ false,
2812
+ true
2813
+ ],
2814
+ "range_flattens": [
2815
+ null,
2816
+ null,
2817
+ null
2818
+ ],
2819
+ "load_eviction_policies": [
2820
+ "",
2821
+ "",
2822
+ "first"
2823
+ ],
2824
+ "num_warps": 8,
2825
+ "num_stages": 2,
2826
+ "indexing": [
2827
+ "pointer",
2828
+ "pointer",
2829
+ "pointer",
2830
+ "pointer",
2831
+ "pointer"
2832
+ ],
2833
+ "atomic_indexing": [],
2834
+ "pid_type": "flat"
2835
+ }
2836
+ },
2837
+ {
2838
+ "key": {
2839
+ "hidden_size": 512,
2840
+ "num_tokens": 512
2841
+ },
2842
+ "config": {
2843
+ "block_sizes": [
2844
+ 512,
2845
+ 128
2846
+ ],
2847
+ "range_unroll_factors": [
2848
+ 0,
2849
+ 3,
2850
+ 3
2851
+ ],
2852
+ "range_warp_specializes": [],
2853
+ "range_num_stages": [],
2854
+ "range_multi_buffers": [
2855
+ null,
2856
+ true,
2857
+ null
2858
+ ],
2859
+ "range_flattens": [
2860
+ null,
2861
+ false,
2862
+ true
2863
+ ],
2864
+ "load_eviction_policies": [
2865
+ "first",
2866
+ "",
2867
+ ""
2868
+ ],
2869
+ "num_warps": 1,
2870
+ "num_stages": 7,
2871
+ "indexing": [
2872
+ "tensor_descriptor",
2873
+ "tensor_descriptor",
2874
+ "tensor_descriptor",
2875
+ "pointer",
2876
+ "pointer"
2877
+ ],
2878
+ "atomic_indexing": [],
2879
+ "pid_type": "flat"
2880
+ }
2881
+ },
2882
+ {
2883
+ "key": {
2884
+ "hidden_size": 2048,
2885
+ "num_tokens": 512
2886
+ },
2887
+ "config": {
2888
+ "block_sizes": [
2889
+ 512,
2890
+ 1024
2891
+ ],
2892
+ "range_unroll_factors": [
2893
+ 0,
2894
+ 2,
2895
+ 1
2896
+ ],
2897
+ "range_warp_specializes": [],
2898
+ "range_num_stages": [],
2899
+ "range_multi_buffers": [
2900
+ null,
2901
+ null,
2902
+ false
2903
+ ],
2904
+ "range_flattens": [
2905
+ null,
2906
+ true,
2907
+ null
2908
+ ],
2909
+ "load_eviction_policies": [
2910
+ "",
2911
+ "first",
2912
+ "first"
2913
+ ],
2914
+ "num_warps": 1,
2915
+ "num_stages": 7,
2916
+ "indexing": [
2917
+ "tensor_descriptor",
2918
+ "pointer",
2919
+ "tensor_descriptor",
2920
+ "pointer",
2921
+ "tensor_descriptor"
2922
+ ],
2923
+ "atomic_indexing": [],
2924
+ "pid_type": "flat"
2925
+ }
2926
+ },
2927
+ {
2928
+ "key": {
2929
+ "hidden_size": 4096,
2930
+ "num_tokens": 512
2931
+ },
2932
+ "config": {
2933
+ "block_sizes": [
2934
+ 4096,
2935
+ 2048
2936
+ ],
2937
+ "range_unroll_factors": [
2938
+ 0,
2939
+ 2,
2940
+ 2
2941
+ ],
2942
+ "range_warp_specializes": [],
2943
+ "range_num_stages": [],
2944
+ "range_multi_buffers": [
2945
+ null,
2946
+ null,
2947
+ null
2948
+ ],
2949
+ "range_flattens": [
2950
+ null,
2951
+ false,
2952
+ null
2953
+ ],
2954
+ "load_eviction_policies": [
2955
+ "first",
2956
+ "first",
2957
+ "last"
2958
+ ],
2959
+ "num_warps": 4,
2960
+ "num_stages": 1,
2961
+ "indexing": [
2962
+ "tensor_descriptor",
2963
+ "tensor_descriptor",
2964
+ "pointer",
2965
+ "tensor_descriptor",
2966
+ "pointer"
2967
+ ],
2968
+ "atomic_indexing": [],
2969
+ "pid_type": "flat"
2970
+ }
2971
+ },
2972
+ {
2973
+ "key": {
2974
+ "hidden_size": 6144,
2975
+ "num_tokens": 512
2976
+ },
2977
+ "config": {
2978
+ "block_sizes": [
2979
+ 8192,
2980
+ 4096
2981
+ ],
2982
+ "range_unroll_factors": [
2983
+ 0,
2984
+ 4,
2985
+ 2
2986
+ ],
2987
+ "range_warp_specializes": [],
2988
+ "range_num_stages": [],
2989
+ "range_multi_buffers": [
2990
+ null,
2991
+ false,
2992
+ false
2993
+ ],
2994
+ "range_flattens": [
2995
+ null,
2996
+ null,
2997
+ null
2998
+ ],
2999
+ "load_eviction_policies": [
3000
+ "",
3001
+ "",
3002
+ ""
3003
+ ],
3004
+ "num_warps": 2,
3005
+ "num_stages": 1,
3006
+ "indexing": [
3007
+ "pointer",
3008
+ "tensor_descriptor",
3009
+ "tensor_descriptor",
3010
+ "pointer",
3011
+ "tensor_descriptor"
3012
+ ],
3013
+ "atomic_indexing": [],
3014
+ "pid_type": "flat"
3015
+ }
3016
+ },
3017
+ {
3018
+ "key": {
3019
+ "hidden_size": 8192,
3020
+ "num_tokens": 512
3021
+ },
3022
+ "config": {
3023
+ "block_sizes": [
3024
+ 8192,
3025
+ 8192
3026
+ ],
3027
+ "range_unroll_factors": [
3028
+ 0,
3029
+ 3,
3030
+ 1
3031
+ ],
3032
+ "range_warp_specializes": [],
3033
+ "range_num_stages": [],
3034
+ "range_multi_buffers": [
3035
+ null,
3036
+ true,
3037
+ null
3038
+ ],
3039
+ "range_flattens": [
3040
+ null,
3041
+ null,
3042
+ true
3043
+ ],
3044
+ "load_eviction_policies": [
3045
+ "last",
3046
+ "last",
3047
+ ""
3048
+ ],
3049
+ "num_warps": 2,
3050
+ "num_stages": 1,
3051
+ "indexing": [
3052
+ "tensor_descriptor",
3053
+ "pointer",
3054
+ "pointer",
3055
+ "tensor_descriptor",
3056
+ "tensor_descriptor"
3057
+ ],
3058
+ "atomic_indexing": [],
3059
+ "pid_type": "flat"
3060
+ }
3061
+ },
3062
+ {
3063
+ "key": {
3064
+ "hidden_size": 12288,
3065
+ "num_tokens": 512
3066
+ },
3067
+ "config": {
3068
+ "block_sizes": [
3069
+ 16384,
3070
+ 4096
3071
+ ],
3072
+ "range_unroll_factors": [
3073
+ 0,
3074
+ 1,
3075
+ 4
3076
+ ],
3077
+ "range_warp_specializes": [],
3078
+ "range_num_stages": [],
3079
+ "range_multi_buffers": [
3080
+ null,
3081
+ null,
3082
+ true
3083
+ ],
3084
+ "range_flattens": [
3085
+ null,
3086
+ true,
3087
+ null
3088
+ ],
3089
+ "load_eviction_policies": [
3090
+ "",
3091
+ "",
3092
+ "first"
3093
+ ],
3094
+ "num_warps": 16,
3095
+ "num_stages": 6,
3096
+ "indexing": [
3097
+ "pointer",
3098
+ "tensor_descriptor",
3099
+ "pointer",
3100
+ "pointer",
3101
+ "pointer"
3102
+ ],
3103
+ "atomic_indexing": [],
3104
+ "pid_type": "flat"
3105
+ }
3106
+ },
3107
+ {
3108
+ "key": {
3109
+ "hidden_size": 28672,
3110
+ "num_tokens": 512
3111
+ },
3112
+ "config": {
3113
+ "block_sizes": [
3114
+ 32768,
3115
+ 4096
3116
+ ],
3117
+ "range_unroll_factors": [
3118
+ 0,
3119
+ 3,
3120
+ 1
3121
+ ],
3122
+ "range_warp_specializes": [],
3123
+ "range_num_stages": [],
3124
+ "range_multi_buffers": [
3125
+ null,
3126
+ null,
3127
+ true
3128
+ ],
3129
+ "range_flattens": [
3130
+ null,
3131
+ true,
3132
+ false
3133
+ ],
3134
+ "load_eviction_policies": [
3135
+ "",
3136
+ "last",
3137
+ "first"
3138
+ ],
3139
+ "num_warps": 32,
3140
+ "num_stages": 6,
3141
+ "indexing": [
3142
+ "pointer",
3143
+ "tensor_descriptor",
3144
+ "tensor_descriptor",
3145
+ "pointer",
3146
+ "tensor_descriptor"
3147
+ ],
3148
+ "atomic_indexing": [],
3149
+ "pid_type": "flat"
3150
+ }
3151
+ },
3152
+ {
3153
+ "key": {
3154
+ "hidden_size": 512,
3155
+ "num_tokens": 1024
3156
+ },
3157
+ "config": {
3158
+ "block_sizes": [
3159
+ 512,
3160
+ 512
3161
+ ],
3162
+ "range_unroll_factors": [
3163
+ 0,
3164
+ 2,
3165
+ 2
3166
+ ],
3167
+ "range_warp_specializes": [],
3168
+ "range_num_stages": [],
3169
+ "range_multi_buffers": [
3170
+ null,
3171
+ null,
3172
+ true
3173
+ ],
3174
+ "range_flattens": [
3175
+ null,
3176
+ null,
3177
+ false
3178
+ ],
3179
+ "load_eviction_policies": [
3180
+ "last",
3181
+ "first",
3182
+ "first"
3183
+ ],
3184
+ "num_warps": 1,
3185
+ "num_stages": 4,
3186
+ "indexing": [
3187
+ "pointer",
3188
+ "tensor_descriptor",
3189
+ "pointer",
3190
+ "pointer",
3191
+ "tensor_descriptor"
3192
+ ],
3193
+ "atomic_indexing": [],
3194
+ "pid_type": "flat"
3195
+ }
3196
+ },
3197
+ {
3198
+ "key": {
3199
+ "hidden_size": 2048,
3200
+ "num_tokens": 1024
3201
+ },
3202
+ "config": {
3203
+ "block_sizes": [
3204
+ 8192,
3205
+ 8192
3206
+ ],
3207
+ "range_unroll_factors": [
3208
+ 0,
3209
+ 3,
3210
+ 0
3211
+ ],
3212
+ "range_warp_specializes": [],
3213
+ "range_num_stages": [],
3214
+ "range_multi_buffers": [
3215
+ null,
3216
+ true,
3217
+ true
3218
+ ],
3219
+ "range_flattens": [
3220
+ null,
3221
+ true,
3222
+ true
3223
+ ],
3224
+ "load_eviction_policies": [
3225
+ "",
3226
+ "last",
3227
+ "first"
3228
+ ],
3229
+ "num_warps": 2,
3230
+ "num_stages": 8,
3231
+ "indexing": [
3232
+ "tensor_descriptor",
3233
+ "pointer",
3234
+ "pointer",
3235
+ "tensor_descriptor",
3236
+ "tensor_descriptor"
3237
+ ],
3238
+ "atomic_indexing": [],
3239
+ "pid_type": "flat"
3240
+ }
3241
+ },
3242
+ {
3243
+ "key": {
3244
+ "hidden_size": 4096,
3245
+ "num_tokens": 1024
3246
+ },
3247
+ "config": {
3248
+ "block_sizes": [
3249
+ 4096,
3250
+ 2048
3251
+ ],
3252
+ "range_unroll_factors": [
3253
+ 2,
3254
+ 2,
3255
+ 4
3256
+ ],
3257
+ "range_warp_specializes": [],
3258
+ "range_multi_buffers": [
3259
+ false,
3260
+ null,
3261
+ true
3262
+ ],
3263
+ "range_flattens": [
3264
+ false,
3265
+ null,
3266
+ true
3267
+ ],
3268
+ "load_eviction_policies": [
3269
+ "",
3270
+ "first",
3271
+ "first"
3272
+ ],
3273
+ "num_warps": 8,
3274
+ "num_stages": 6,
3275
+ "indexing": [
3276
+ "pointer",
3277
+ "pointer",
3278
+ "tensor_descriptor",
3279
+ "tensor_descriptor",
3280
+ "pointer"
3281
+ ],
3282
+ "atomic_indexing": [],
3283
+ "pid_type": "persistent_interleaved",
3284
+ "num_sm_multiplier": 4,
3285
+ "maxnreg": 32
3286
+ }
3287
+ },
3288
+ {
3289
+ "key": {
3290
+ "hidden_size": 6144,
3291
+ "num_tokens": 1024
3292
+ },
3293
+ "config": {
3294
+ "block_sizes": [
3295
+ 8192,
3296
+ 8192
3297
+ ],
3298
+ "range_unroll_factors": [
3299
+ 0,
3300
+ 3,
3301
+ 0
3302
+ ],
3303
+ "range_warp_specializes": [],
3304
+ "range_num_stages": [],
3305
+ "range_multi_buffers": [
3306
+ null,
3307
+ true,
3308
+ true
3309
+ ],
3310
+ "range_flattens": [
3311
+ null,
3312
+ true,
3313
+ true
3314
+ ],
3315
+ "load_eviction_policies": [
3316
+ "",
3317
+ "last",
3318
+ "first"
3319
+ ],
3320
+ "num_warps": 2,
3321
+ "num_stages": 8,
3322
+ "indexing": [
3323
+ "tensor_descriptor",
3324
+ "pointer",
3325
+ "pointer",
3326
+ "tensor_descriptor",
3327
+ "tensor_descriptor"
3328
+ ],
3329
+ "atomic_indexing": [],
3330
+ "pid_type": "flat"
3331
+ }
3332
+ },
3333
+ {
3334
+ "key": {
3335
+ "hidden_size": 8192,
3336
+ "num_tokens": 1024
3337
+ },
3338
+ "config": {
3339
+ "block_sizes": [
3340
+ 8192,
3341
+ 8192
3342
+ ],
3343
+ "range_unroll_factors": [
3344
+ 1,
3345
+ 2,
3346
+ 2
3347
+ ],
3348
+ "range_warp_specializes": [],
3349
+ "range_multi_buffers": [
3350
+ true,
3351
+ false,
3352
+ null
3353
+ ],
3354
+ "range_flattens": [
3355
+ true,
3356
+ false,
3357
+ false
3358
+ ],
3359
+ "load_eviction_policies": [
3360
+ "",
3361
+ "",
3362
+ "first"
3363
+ ],
3364
+ "num_warps": 16,
3365
+ "num_stages": 5,
3366
+ "indexing": [
3367
+ "pointer",
3368
+ "tensor_descriptor",
3369
+ "tensor_descriptor",
3370
+ "pointer",
3371
+ "tensor_descriptor"
3372
+ ],
3373
+ "atomic_indexing": [],
3374
+ "pid_type": "persistent_blocked",
3375
+ "num_sm_multiplier": 4,
3376
+ "maxnreg": 128
3377
+ }
3378
+ },
3379
+ {
3380
+ "key": {
3381
+ "hidden_size": 12288,
3382
+ "num_tokens": 1024
3383
+ },
3384
+ "config": {
3385
+ "block_sizes": [
3386
+ 16384,
3387
+ 4096
3388
+ ],
3389
+ "range_unroll_factors": [
3390
+ 0,
3391
+ 3,
3392
+ 2
3393
+ ],
3394
+ "range_warp_specializes": [],
3395
+ "range_num_stages": [],
3396
+ "range_multi_buffers": [
3397
+ null,
3398
+ false,
3399
+ false
3400
+ ],
3401
+ "range_flattens": [
3402
+ null,
3403
+ true,
3404
+ true
3405
+ ],
3406
+ "load_eviction_policies": [
3407
+ "",
3408
+ "last",
3409
+ "first"
3410
+ ],
3411
+ "num_warps": 16,
3412
+ "num_stages": 6,
3413
+ "indexing": [
3414
+ "pointer",
3415
+ "tensor_descriptor",
3416
+ "tensor_descriptor",
3417
+ "tensor_descriptor",
3418
+ "pointer"
3419
+ ],
3420
+ "atomic_indexing": [],
3421
+ "pid_type": "flat"
3422
+ }
3423
+ },
3424
+ {
3425
+ "key": {
3426
+ "hidden_size": 28672,
3427
+ "num_tokens": 1024
3428
+ },
3429
+ "config": {
3430
+ "block_sizes": [
3431
+ 32768,
3432
+ 4096
3433
+ ],
3434
+ "range_unroll_factors": [
3435
+ 0,
3436
+ 3,
3437
+ 2
3438
+ ],
3439
+ "range_warp_specializes": [],
3440
+ "range_num_stages": [],
3441
+ "range_multi_buffers": [
3442
+ null,
3443
+ false,
3444
+ null
3445
+ ],
3446
+ "range_flattens": [
3447
+ null,
3448
+ null,
3449
+ false
3450
+ ],
3451
+ "load_eviction_policies": [
3452
+ "",
3453
+ "last",
3454
+ "first"
3455
+ ],
3456
+ "num_warps": 32,
3457
+ "num_stages": 8,
3458
+ "indexing": [
3459
+ "pointer",
3460
+ "tensor_descriptor",
3461
+ "pointer",
3462
+ "pointer",
3463
+ "tensor_descriptor"
3464
+ ],
3465
+ "atomic_indexing": [],
3466
+ "pid_type": "flat"
3467
+ }
3468
+ },
3469
+ {
3470
+ "key": {
3471
+ "hidden_size": 512,
3472
+ "num_tokens": 2048
3473
+ },
3474
+ "config": {
3475
+ "block_sizes": [
3476
+ 512,
3477
+ 512
3478
+ ],
3479
+ "range_unroll_factors": [
3480
+ 0,
3481
+ 1,
3482
+ 2
3483
+ ],
3484
+ "range_warp_specializes": [],
3485
+ "range_num_stages": [],
3486
+ "range_multi_buffers": [
3487
+ null,
3488
+ true,
3489
+ false
3490
+ ],
3491
+ "range_flattens": [
3492
+ null,
3493
+ true,
3494
+ null
3495
+ ],
3496
+ "load_eviction_policies": [
3497
+ "last",
3498
+ "first",
3499
+ "last"
3500
+ ],
3501
+ "num_warps": 1,
3502
+ "num_stages": 7,
3503
+ "indexing": [
3504
+ "tensor_descriptor",
3505
+ "pointer",
3506
+ "tensor_descriptor",
3507
+ "pointer",
3508
+ "pointer"
3509
+ ],
3510
+ "atomic_indexing": [],
3511
+ "pid_type": "flat"
3512
+ }
3513
+ },
3514
+ {
3515
+ "key": {
3516
+ "hidden_size": 2048,
3517
+ "num_tokens": 2048
3518
+ },
3519
+ "config": {
3520
+ "block_sizes": [
3521
+ 1024,
3522
+ 1024
3523
+ ],
3524
+ "range_unroll_factors": [
3525
+ 0,
3526
+ 2,
3527
+ 2
3528
+ ],
3529
+ "range_warp_specializes": [],
3530
+ "range_num_stages": [],
3531
+ "range_multi_buffers": [
3532
+ null,
3533
+ false,
3534
+ false
3535
+ ],
3536
+ "range_flattens": [
3537
+ null,
3538
+ false,
3539
+ false
3540
+ ],
3541
+ "load_eviction_policies": [
3542
+ "",
3543
+ "",
3544
+ "last"
3545
+ ],
3546
+ "num_warps": 2,
3547
+ "num_stages": 1,
3548
+ "indexing": [
3549
+ "tensor_descriptor",
3550
+ "pointer",
3551
+ "pointer",
3552
+ "pointer",
3553
+ "tensor_descriptor"
3554
+ ],
3555
+ "atomic_indexing": [],
3556
+ "pid_type": "flat"
3557
+ }
3558
+ },
3559
+ {
3560
+ "key": {
3561
+ "hidden_size": 4096,
3562
+ "num_tokens": 2048
3563
+ },
3564
+ "config": {
3565
+ "block_sizes": [
3566
+ 4096,
3567
+ 4096
3568
+ ],
3569
+ "range_unroll_factors": [
3570
+ 0,
3571
+ 0,
3572
+ 2
3573
+ ],
3574
+ "range_warp_specializes": [],
3575
+ "range_num_stages": [],
3576
+ "range_multi_buffers": [
3577
+ null,
3578
+ false,
3579
+ null
3580
+ ],
3581
+ "range_flattens": [
3582
+ null,
3583
+ true,
3584
+ null
3585
+ ],
3586
+ "load_eviction_policies": [
3587
+ "",
3588
+ "last",
3589
+ ""
3590
+ ],
3591
+ "num_warps": 8,
3592
+ "num_stages": 7,
3593
+ "indexing": [
3594
+ "tensor_descriptor",
3595
+ "pointer",
3596
+ "tensor_descriptor",
3597
+ "pointer",
3598
+ "pointer"
3599
+ ],
3600
+ "atomic_indexing": [],
3601
+ "pid_type": "flat"
3602
+ }
3603
+ },
3604
+ {
3605
+ "key": {
3606
+ "hidden_size": 6144,
3607
+ "num_tokens": 2048
3608
+ },
3609
+ "config": {
3610
+ "block_sizes": [
3611
+ 8192,
3612
+ 2048
3613
+ ],
3614
+ "range_unroll_factors": [
3615
+ 2,
3616
+ 3,
3617
+ 0
3618
+ ],
3619
+ "range_warp_specializes": [],
3620
+ "range_multi_buffers": [
3621
+ false,
3622
+ true,
3623
+ true
3624
+ ],
3625
+ "range_flattens": [
3626
+ null,
3627
+ false,
3628
+ false
3629
+ ],
3630
+ "load_eviction_policies": [
3631
+ "",
3632
+ "last",
3633
+ "first"
3634
+ ],
3635
+ "num_warps": 8,
3636
+ "num_stages": 8,
3637
+ "indexing": [
3638
+ "pointer",
3639
+ "pointer",
3640
+ "pointer",
3641
+ "tensor_descriptor",
3642
+ "tensor_descriptor"
3643
+ ],
3644
+ "atomic_indexing": [],
3645
+ "pid_type": "persistent_interleaved",
3646
+ "num_sm_multiplier": 8,
3647
+ "maxnreg": 32
3648
+ }
3649
+ },
3650
+ {
3651
+ "key": {
3652
+ "hidden_size": 8192,
3653
+ "num_tokens": 2048
3654
+ },
3655
+ "config": {
3656
+ "block_sizes": [
3657
+ 16384,
3658
+ 4096
3659
+ ],
3660
+ "range_unroll_factors": [
3661
+ 0,
3662
+ 0,
3663
+ 1
3664
+ ],
3665
+ "range_warp_specializes": [],
3666
+ "range_num_stages": [],
3667
+ "range_multi_buffers": [
3668
+ null,
3669
+ false,
3670
+ true
3671
+ ],
3672
+ "range_flattens": [
3673
+ null,
3674
+ true,
3675
+ false
3676
+ ],
3677
+ "load_eviction_policies": [
3678
+ "",
3679
+ "last",
3680
+ "last"
3681
+ ],
3682
+ "num_warps": 16,
3683
+ "num_stages": 4,
3684
+ "indexing": [
3685
+ "tensor_descriptor",
3686
+ "tensor_descriptor",
3687
+ "pointer",
3688
+ "pointer",
3689
+ "pointer"
3690
+ ],
3691
+ "atomic_indexing": [],
3692
+ "pid_type": "flat"
3693
+ }
3694
+ },
3695
+ {
3696
+ "key": {
3697
+ "hidden_size": 12288,
3698
+ "num_tokens": 2048
3699
+ },
3700
+ "config": {
3701
+ "block_sizes": [
3702
+ 16384,
3703
+ 4096
3704
+ ],
3705
+ "range_unroll_factors": [
3706
+ 0,
3707
+ 0,
3708
+ 2
3709
+ ],
3710
+ "range_warp_specializes": [],
3711
+ "range_num_stages": [],
3712
+ "range_multi_buffers": [
3713
+ null,
3714
+ false,
3715
+ false
3716
+ ],
3717
+ "range_flattens": [
3718
+ null,
3719
+ null,
3720
+ true
3721
+ ],
3722
+ "load_eviction_policies": [
3723
+ "",
3724
+ "",
3725
+ "first"
3726
+ ],
3727
+ "num_warps": 16,
3728
+ "num_stages": 1,
3729
+ "indexing": [
3730
+ "tensor_descriptor",
3731
+ "pointer",
3732
+ "pointer",
3733
+ "tensor_descriptor",
3734
+ "pointer"
3735
+ ],
3736
+ "atomic_indexing": [],
3737
+ "pid_type": "flat"
3738
+ }
3739
+ },
3740
+ {
3741
+ "key": {
3742
+ "hidden_size": 28672,
3743
+ "num_tokens": 2048
3744
+ },
3745
+ "config": {
3746
+ "block_sizes": [
3747
+ 32768,
3748
+ 4096
3749
+ ],
3750
+ "range_unroll_factors": [
3751
+ 0,
3752
+ 3,
3753
+ 2
3754
+ ],
3755
+ "range_warp_specializes": [],
3756
+ "range_num_stages": [],
3757
+ "range_multi_buffers": [
3758
+ null,
3759
+ false,
3760
+ null
3761
+ ],
3762
+ "range_flattens": [
3763
+ null,
3764
+ false,
3765
+ null
3766
+ ],
3767
+ "load_eviction_policies": [
3768
+ "",
3769
+ "last",
3770
+ "first"
3771
+ ],
3772
+ "num_warps": 32,
3773
+ "num_stages": 2,
3774
+ "indexing": [
3775
+ "pointer",
3776
+ "pointer",
3777
+ "tensor_descriptor",
3778
+ "tensor_descriptor",
3779
+ "tensor_descriptor"
3780
+ ],
3781
+ "atomic_indexing": [],
3782
+ "pid_type": "flat"
3783
+ }
3784
+ },
3785
+ {
3786
+ "key": {
3787
+ "hidden_size": 512,
3788
+ "num_tokens": 4096
3789
+ },
3790
+ "config": {
3791
+ "block_sizes": [
3792
+ 512,
3793
+ 256
3794
+ ],
3795
+ "range_unroll_factors": [
3796
+ 1,
3797
+ 1,
3798
+ 4
3799
+ ],
3800
+ "range_warp_specializes": [],
3801
+ "range_multi_buffers": [
3802
+ false,
3803
+ false,
3804
+ true
3805
+ ],
3806
+ "range_flattens": [
3807
+ false,
3808
+ null,
3809
+ false
3810
+ ],
3811
+ "load_eviction_policies": [
3812
+ "",
3813
+ "first",
3814
+ "first"
3815
+ ],
3816
+ "num_warps": 1,
3817
+ "num_stages": 5,
3818
+ "indexing": [
3819
+ "pointer",
3820
+ "pointer",
3821
+ "tensor_descriptor",
3822
+ "pointer",
3823
+ "pointer"
3824
+ ],
3825
+ "atomic_indexing": [],
3826
+ "pid_type": "persistent_interleaved",
3827
+ "num_sm_multiplier": 16,
3828
+ "maxnreg": 128
3829
+ }
3830
+ },
3831
+ {
3832
+ "key": {
3833
+ "hidden_size": 2048,
3834
+ "num_tokens": 4096
3835
+ },
3836
+ "config": {
3837
+ "block_sizes": [
3838
+ 2048,
3839
+ 2048
3840
+ ],
3841
+ "range_unroll_factors": [
3842
+ 3,
3843
+ 2,
3844
+ 2
3845
+ ],
3846
+ "range_warp_specializes": [],
3847
+ "range_multi_buffers": [
3848
+ true,
3849
+ false,
3850
+ false
3851
+ ],
3852
+ "range_flattens": [
3853
+ false,
3854
+ false,
3855
+ true
3856
+ ],
3857
+ "load_eviction_policies": [
3858
+ "",
3859
+ "last",
3860
+ "first"
3861
+ ],
3862
+ "num_warps": 4,
3863
+ "num_stages": 2,
3864
+ "indexing": [
3865
+ "pointer",
3866
+ "pointer",
3867
+ "tensor_descriptor",
3868
+ "pointer",
3869
+ "tensor_descriptor"
3870
+ ],
3871
+ "atomic_indexing": [],
3872
+ "pid_type": "persistent_blocked",
3873
+ "num_sm_multiplier": 16
3874
+ }
3875
+ },
3876
+ {
3877
+ "key": {
3878
+ "hidden_size": 4096,
3879
+ "num_tokens": 4096
3880
+ },
3881
+ "config": {
3882
+ "block_sizes": [
3883
+ 4096,
3884
+ 4096
3885
+ ],
3886
+ "range_unroll_factors": [
3887
+ 0,
3888
+ 3,
3889
+ 1
3890
+ ],
3891
+ "range_warp_specializes": [],
3892
+ "range_num_stages": [],
3893
+ "range_multi_buffers": [
3894
+ null,
3895
+ true,
3896
+ null
3897
+ ],
3898
+ "range_flattens": [
3899
+ null,
3900
+ null,
3901
+ true
3902
+ ],
3903
+ "load_eviction_policies": [
3904
+ "",
3905
+ "last",
3906
+ "first"
3907
+ ],
3908
+ "num_warps": 8,
3909
+ "num_stages": 5,
3910
+ "indexing": [
3911
+ "tensor_descriptor",
3912
+ "pointer",
3913
+ "tensor_descriptor",
3914
+ "tensor_descriptor",
3915
+ "tensor_descriptor"
3916
+ ],
3917
+ "atomic_indexing": [],
3918
+ "pid_type": "flat"
3919
+ }
3920
+ },
3921
+ {
3922
+ "key": {
3923
+ "hidden_size": 6144,
3924
+ "num_tokens": 4096
3925
+ },
3926
+ "config": {
3927
+ "block_sizes": [
3928
+ 8192,
3929
+ 2048
3930
+ ],
3931
+ "range_unroll_factors": [
3932
+ 0,
3933
+ 0,
3934
+ 1
3935
+ ],
3936
+ "range_warp_specializes": [],
3937
+ "range_num_stages": [],
3938
+ "range_multi_buffers": [
3939
+ null,
3940
+ null,
3941
+ true
3942
+ ],
3943
+ "range_flattens": [
3944
+ null,
3945
+ null,
3946
+ false
3947
+ ],
3948
+ "load_eviction_policies": [
3949
+ "",
3950
+ "",
3951
+ "first"
3952
+ ],
3953
+ "num_warps": 8,
3954
+ "num_stages": 6,
3955
+ "indexing": [
3956
+ "pointer",
3957
+ "pointer",
3958
+ "tensor_descriptor",
3959
+ "tensor_descriptor",
3960
+ "tensor_descriptor"
3961
+ ],
3962
+ "atomic_indexing": [],
3963
+ "pid_type": "flat"
3964
+ }
3965
+ },
3966
+ {
3967
+ "key": {
3968
+ "hidden_size": 8192,
3969
+ "num_tokens": 4096
3970
+ },
3971
+ "config": {
3972
+ "block_sizes": [
3973
+ 8192,
3974
+ 8192
3975
+ ],
3976
+ "range_unroll_factors": [
3977
+ 0,
3978
+ 1,
3979
+ 0
3980
+ ],
3981
+ "range_warp_specializes": [],
3982
+ "range_num_stages": [],
3983
+ "range_multi_buffers": [
3984
+ null,
3985
+ true,
3986
+ false
3987
+ ],
3988
+ "range_flattens": [
3989
+ null,
3990
+ null,
3991
+ false
3992
+ ],
3993
+ "load_eviction_policies": [
3994
+ "",
3995
+ "",
3996
+ "first"
3997
+ ],
3998
+ "num_warps": 16,
3999
+ "num_stages": 4,
4000
+ "indexing": [
4001
+ "tensor_descriptor",
4002
+ "tensor_descriptor",
4003
+ "tensor_descriptor",
4004
+ "tensor_descriptor",
4005
+ "tensor_descriptor"
4006
+ ],
4007
+ "atomic_indexing": [],
4008
+ "pid_type": "flat"
4009
+ }
4010
+ },
4011
+ {
4012
+ "key": {
4013
+ "hidden_size": 12288,
4014
+ "num_tokens": 4096
4015
+ },
4016
+ "config": {
4017
+ "block_sizes": [
4018
+ 16384,
4019
+ 4096
4020
+ ],
4021
+ "range_unroll_factors": [
4022
+ 0,
4023
+ 3,
4024
+ 0
4025
+ ],
4026
+ "range_warp_specializes": [],
4027
+ "range_num_stages": [],
4028
+ "range_multi_buffers": [
4029
+ null,
4030
+ false,
4031
+ false
4032
+ ],
4033
+ "range_flattens": [
4034
+ null,
4035
+ true,
4036
+ true
4037
+ ],
4038
+ "load_eviction_policies": [
4039
+ "",
4040
+ "last",
4041
+ "first"
4042
+ ],
4043
+ "num_warps": 16,
4044
+ "num_stages": 6,
4045
+ "indexing": [
4046
+ "tensor_descriptor",
4047
+ "tensor_descriptor",
4048
+ "tensor_descriptor",
4049
+ "tensor_descriptor",
4050
+ "tensor_descriptor"
4051
+ ],
4052
+ "atomic_indexing": [],
4053
+ "pid_type": "flat"
4054
+ }
4055
+ },
4056
+ {
4057
+ "key": {
4058
+ "hidden_size": 28672,
4059
+ "num_tokens": 4096
4060
+ },
4061
+ "config": {
4062
+ "block_sizes": [
4063
+ 32768,
4064
+ 4096
4065
+ ],
4066
+ "range_unroll_factors": [
4067
+ 2,
4068
+ 2,
4069
+ 3
4070
+ ],
4071
+ "range_warp_specializes": [],
4072
+ "range_multi_buffers": [
4073
+ false,
4074
+ null,
4075
+ true
4076
+ ],
4077
+ "range_flattens": [
4078
+ null,
4079
+ true,
4080
+ true
4081
+ ],
4082
+ "load_eviction_policies": [
4083
+ "last",
4084
+ "",
4085
+ "last"
4086
+ ],
4087
+ "num_warps": 16,
4088
+ "num_stages": 2,
4089
+ "indexing": [
4090
+ "pointer",
4091
+ "tensor_descriptor",
4092
+ "tensor_descriptor",
4093
+ "tensor_descriptor",
4094
+ "tensor_descriptor"
4095
+ ],
4096
+ "atomic_indexing": [],
4097
+ "pid_type": "persistent_blocked",
4098
+ "num_sm_multiplier": 32,
4099
+ "maxnreg": 64
4100
+ }
4101
+ },
4102
+ {
4103
+ "key": {
4104
+ "hidden_size": 512,
4105
+ "num_tokens": 8192
4106
+ },
4107
+ "config": {
4108
+ "block_sizes": [
4109
+ 512,
4110
+ 512
4111
+ ],
4112
+ "range_unroll_factors": [
4113
+ 3,
4114
+ 4,
4115
+ 0
4116
+ ],
4117
+ "range_warp_specializes": [],
4118
+ "range_multi_buffers": [
4119
+ false,
4120
+ true,
4121
+ null
4122
+ ],
4123
+ "range_flattens": [
4124
+ false,
4125
+ false,
4126
+ false
4127
+ ],
4128
+ "load_eviction_policies": [
4129
+ "",
4130
+ "",
4131
+ "first"
4132
+ ],
4133
+ "num_warps": 1,
4134
+ "num_stages": 6,
4135
+ "indexing": [
4136
+ "pointer",
4137
+ "pointer",
4138
+ "tensor_descriptor",
4139
+ "tensor_descriptor",
4140
+ "pointer"
4141
+ ],
4142
+ "atomic_indexing": [],
4143
+ "pid_type": "persistent_blocked",
4144
+ "num_sm_multiplier": 32,
4145
+ "maxnreg": 128
4146
+ }
4147
+ },
4148
+ {
4149
+ "key": {
4150
+ "hidden_size": 2048,
4151
+ "num_tokens": 8192
4152
+ },
4153
+ "config": {
4154
+ "block_sizes": [
4155
+ 512,
4156
+ 512
4157
+ ],
4158
+ "range_unroll_factors": [
4159
+ 0,
4160
+ 1,
4161
+ 4
4162
+ ],
4163
+ "range_warp_specializes": [],
4164
+ "range_num_stages": [],
4165
+ "range_multi_buffers": [
4166
+ null,
4167
+ false,
4168
+ false
4169
+ ],
4170
+ "range_flattens": [
4171
+ null,
4172
+ null,
4173
+ false
4174
+ ],
4175
+ "load_eviction_policies": [
4176
+ "",
4177
+ "first",
4178
+ "last"
4179
+ ],
4180
+ "num_warps": 1,
4181
+ "num_stages": 5,
4182
+ "indexing": [
4183
+ "pointer",
4184
+ "pointer",
4185
+ "pointer",
4186
+ "pointer",
4187
+ "tensor_descriptor"
4188
+ ],
4189
+ "atomic_indexing": [],
4190
+ "pid_type": "flat"
4191
+ }
4192
+ },
4193
+ {
4194
+ "key": {
4195
+ "hidden_size": 4096,
4196
+ "num_tokens": 8192
4197
+ },
4198
+ "config": {
4199
+ "block_sizes": [
4200
+ 4096,
4201
+ 4096
4202
+ ],
4203
+ "range_unroll_factors": [
4204
+ 0,
4205
+ 3,
4206
+ 2
4207
+ ],
4208
+ "range_warp_specializes": [],
4209
+ "range_num_stages": [],
4210
+ "range_multi_buffers": [
4211
+ null,
4212
+ true,
4213
+ false
4214
+ ],
4215
+ "range_flattens": [
4216
+ null,
4217
+ null,
4218
+ true
4219
+ ],
4220
+ "load_eviction_policies": [
4221
+ "",
4222
+ "last",
4223
+ "first"
4224
+ ],
4225
+ "num_warps": 8,
4226
+ "num_stages": 3,
4227
+ "indexing": [
4228
+ "tensor_descriptor",
4229
+ "tensor_descriptor",
4230
+ "tensor_descriptor",
4231
+ "pointer",
4232
+ "tensor_descriptor"
4233
+ ],
4234
+ "atomic_indexing": [],
4235
+ "pid_type": "flat"
4236
+ }
4237
+ },
4238
+ {
4239
+ "key": {
4240
+ "hidden_size": 6144,
4241
+ "num_tokens": 8192
4242
+ },
4243
+ "config": {
4244
+ "block_sizes": [
4245
+ 8192,
4246
+ 2048
4247
+ ],
4248
+ "range_unroll_factors": [
4249
+ 0,
4250
+ 4,
4251
+ 0
4252
+ ],
4253
+ "range_warp_specializes": [],
4254
+ "range_num_stages": [],
4255
+ "range_multi_buffers": [
4256
+ null,
4257
+ null,
4258
+ false
4259
+ ],
4260
+ "range_flattens": [
4261
+ null,
4262
+ false,
4263
+ null
4264
+ ],
4265
+ "load_eviction_policies": [
4266
+ "",
4267
+ "last",
4268
+ "first"
4269
+ ],
4270
+ "num_warps": 8,
4271
+ "num_stages": 2,
4272
+ "indexing": [
4273
+ "tensor_descriptor",
4274
+ "tensor_descriptor",
4275
+ "pointer",
4276
+ "tensor_descriptor",
4277
+ "pointer"
4278
+ ],
4279
+ "atomic_indexing": [],
4280
+ "pid_type": "flat"
4281
+ }
4282
+ },
4283
+ {
4284
+ "key": {
4285
+ "hidden_size": 8192,
4286
+ "num_tokens": 8192
4287
+ },
4288
+ "config": {
4289
+ "block_sizes": [
4290
+ 8192,
4291
+ 8192
4292
+ ],
4293
+ "range_unroll_factors": [
4294
+ 0,
4295
+ 4,
4296
+ 1
4297
+ ],
4298
+ "range_warp_specializes": [],
4299
+ "range_num_stages": [],
4300
+ "range_multi_buffers": [
4301
+ null,
4302
+ true,
4303
+ null
4304
+ ],
4305
+ "range_flattens": [
4306
+ null,
4307
+ false,
4308
+ null
4309
+ ],
4310
+ "load_eviction_policies": [
4311
+ "",
4312
+ "",
4313
+ "first"
4314
+ ],
4315
+ "num_warps": 16,
4316
+ "num_stages": 1,
4317
+ "indexing": [
4318
+ "pointer",
4319
+ "tensor_descriptor",
4320
+ "tensor_descriptor",
4321
+ "tensor_descriptor",
4322
+ "pointer"
4323
+ ],
4324
+ "atomic_indexing": [],
4325
+ "pid_type": "flat"
4326
+ }
4327
+ },
4328
+ {
4329
+ "key": {
4330
+ "hidden_size": 12288,
4331
+ "num_tokens": 8192
4332
+ },
4333
+ "config": {
4334
+ "block_sizes": [
4335
+ 16384,
4336
+ 4096
4337
+ ],
4338
+ "range_unroll_factors": [
4339
+ 0,
4340
+ 1,
4341
+ 3
4342
+ ],
4343
+ "range_warp_specializes": [],
4344
+ "range_num_stages": [],
4345
+ "range_multi_buffers": [
4346
+ null,
4347
+ null,
4348
+ true
4349
+ ],
4350
+ "range_flattens": [
4351
+ null,
4352
+ null,
4353
+ false
4354
+ ],
4355
+ "load_eviction_policies": [
4356
+ "last",
4357
+ "last",
4358
+ "first"
4359
+ ],
4360
+ "num_warps": 16,
4361
+ "num_stages": 4,
4362
+ "indexing": [
4363
+ "pointer",
4364
+ "tensor_descriptor",
4365
+ "tensor_descriptor",
4366
+ "tensor_descriptor",
4367
+ "tensor_descriptor"
4368
+ ],
4369
+ "atomic_indexing": [],
4370
+ "pid_type": "flat"
4371
+ }
4372
+ },
4373
+ {
4374
+ "key": {
4375
+ "hidden_size": 28672,
4376
+ "num_tokens": 8192
4377
+ },
4378
+ "config": {
4379
+ "block_sizes": [
4380
+ 32768,
4381
+ 4096
4382
+ ],
4383
+ "range_unroll_factors": [
4384
+ 1,
4385
+ 0,
4386
+ 1
4387
+ ],
4388
+ "range_warp_specializes": [],
4389
+ "range_multi_buffers": [
4390
+ true,
4391
+ null,
4392
+ false
4393
+ ],
4394
+ "range_flattens": [
4395
+ false,
4396
+ false,
4397
+ null
4398
+ ],
4399
+ "load_eviction_policies": [
4400
+ "last",
4401
+ "",
4402
+ "first"
4403
+ ],
4404
+ "num_warps": 16,
4405
+ "num_stages": 5,
4406
+ "indexing": [
4407
+ "pointer",
4408
+ "pointer",
4409
+ "tensor_descriptor",
4410
+ "pointer",
4411
+ "pointer"
4412
+ ],
4413
+ "atomic_indexing": [],
4414
+ "pid_type": "persistent_blocked",
4415
+ "num_sm_multiplier": 64,
4416
+ "maxnreg": 128
4417
+ }
4418
+ },
4419
+ {
4420
+ "key": {
4421
+ "hidden_size": 5120,
4422
+ "num_tokens": 1
4423
+ },
4424
+ "config": {
4425
+ "block_sizes": [
4426
+ 8192,
4427
+ 1024
4428
+ ],
4429
+ "range_unroll_factors": [
4430
+ 0,
4431
+ 0,
4432
+ 3
4433
+ ],
4434
+ "range_warp_specializes": [],
4435
+ "range_num_stages": [],
4436
+ "range_multi_buffers": [
4437
+ null,
4438
+ null,
4439
+ false
4440
+ ],
4441
+ "range_flattens": [
4442
+ null,
4443
+ null,
4444
+ false
4445
+ ],
4446
+ "load_eviction_policies": [
4447
+ "",
4448
+ "last",
4449
+ ""
4450
+ ],
4451
+ "num_warps": 16,
4452
+ "num_stages": 2,
4453
+ "indexing": [
4454
+ "tensor_descriptor",
4455
+ "tensor_descriptor",
4456
+ "tensor_descriptor",
4457
+ "tensor_descriptor",
4458
+ "tensor_descriptor"
4459
+ ],
4460
+ "atomic_indexing": [],
4461
+ "pid_type": "flat"
4462
+ }
4463
+ },
4464
+ {
4465
+ "key": {
4466
+ "hidden_size": 5120,
4467
+ "num_tokens": 2
4468
+ },
4469
+ "config": {
4470
+ "block_sizes": [
4471
+ 8192,
4472
+ 2048
4473
+ ],
4474
+ "range_unroll_factors": [
4475
+ 0,
4476
+ 4,
4477
+ 1
4478
+ ],
4479
+ "range_warp_specializes": [],
4480
+ "range_num_stages": [],
4481
+ "range_multi_buffers": [
4482
+ null,
4483
+ null,
4484
+ true
4485
+ ],
4486
+ "range_flattens": [
4487
+ null,
4488
+ true,
4489
+ null
4490
+ ],
4491
+ "load_eviction_policies": [
4492
+ "",
4493
+ "first",
4494
+ "last"
4495
+ ],
4496
+ "num_warps": 8,
4497
+ "num_stages": 5,
4498
+ "indexing": [
4499
+ "pointer",
4500
+ "pointer",
4501
+ "pointer",
4502
+ "pointer",
4503
+ "tensor_descriptor"
4504
+ ],
4505
+ "atomic_indexing": [],
4506
+ "pid_type": "flat"
4507
+ }
4508
+ },
4509
+ {
4510
+ "key": {
4511
+ "hidden_size": 5120,
4512
+ "num_tokens": 4
4513
+ },
4514
+ "config": {
4515
+ "block_sizes": [
4516
+ 8192,
4517
+ 1024
4518
+ ],
4519
+ "range_unroll_factors": [
4520
+ 0,
4521
+ 1,
4522
+ 0
4523
+ ],
4524
+ "range_warp_specializes": [],
4525
+ "range_num_stages": [],
4526
+ "range_multi_buffers": [
4527
+ null,
4528
+ false,
4529
+ true
4530
+ ],
4531
+ "range_flattens": [
4532
+ null,
4533
+ false,
4534
+ false
4535
+ ],
4536
+ "load_eviction_policies": [
4537
+ "",
4538
+ "",
4539
+ "first"
4540
+ ],
4541
+ "num_warps": 8,
4542
+ "num_stages": 6,
4543
+ "indexing": [
4544
+ "tensor_descriptor",
4545
+ "pointer",
4546
+ "pointer",
4547
+ "tensor_descriptor",
4548
+ "tensor_descriptor"
4549
+ ],
4550
+ "atomic_indexing": [],
4551
+ "pid_type": "flat"
4552
+ }
4553
+ },
4554
+ {
4555
+ "key": {
4556
+ "hidden_size": 5120,
4557
+ "num_tokens": 8
4558
+ },
4559
+ "config": {
4560
+ "block_sizes": [
4561
+ 8192,
4562
+ 2048
4563
+ ],
4564
+ "range_unroll_factors": [
4565
+ 0,
4566
+ 0,
4567
+ 3
4568
+ ],
4569
+ "range_warp_specializes": [],
4570
+ "range_num_stages": [],
4571
+ "range_multi_buffers": [
4572
+ null,
4573
+ false,
4574
+ true
4575
+ ],
4576
+ "range_flattens": [
4577
+ null,
4578
+ false,
4579
+ false
4580
+ ],
4581
+ "load_eviction_policies": [
4582
+ "",
4583
+ "last",
4584
+ "last"
4585
+ ],
4586
+ "num_warps": 16,
4587
+ "num_stages": 4,
4588
+ "indexing": [
4589
+ "pointer",
4590
+ "tensor_descriptor",
4591
+ "tensor_descriptor",
4592
+ "pointer",
4593
+ "tensor_descriptor"
4594
+ ],
4595
+ "atomic_indexing": [],
4596
+ "pid_type": "flat"
4597
+ }
4598
+ },
4599
+ {
4600
+ "key": {
4601
+ "hidden_size": 5120,
4602
+ "num_tokens": 16
4603
+ },
4604
+ "config": {
4605
+ "block_sizes": [
4606
+ 8192,
4607
+ 1024
4608
+ ],
4609
+ "range_unroll_factors": [
4610
+ 0,
4611
+ 0,
4612
+ 3
4613
+ ],
4614
+ "range_warp_specializes": [],
4615
+ "range_num_stages": [],
4616
+ "range_multi_buffers": [
4617
+ null,
4618
+ null,
4619
+ false
4620
+ ],
4621
+ "range_flattens": [
4622
+ null,
4623
+ null,
4624
+ false
4625
+ ],
4626
+ "load_eviction_policies": [
4627
+ "",
4628
+ "last",
4629
+ ""
4630
+ ],
4631
+ "num_warps": 16,
4632
+ "num_stages": 2,
4633
+ "indexing": [
4634
+ "tensor_descriptor",
4635
+ "tensor_descriptor",
4636
+ "tensor_descriptor",
4637
+ "tensor_descriptor",
4638
+ "tensor_descriptor"
4639
+ ],
4640
+ "atomic_indexing": [],
4641
+ "pid_type": "flat"
4642
+ }
4643
+ },
4644
+ {
4645
+ "key": {
4646
+ "hidden_size": 5120,
4647
+ "num_tokens": 32
4648
+ },
4649
+ "config": {
4650
+ "block_sizes": [
4651
+ 8192,
4652
+ 1024
4653
+ ],
4654
+ "range_unroll_factors": [
4655
+ 0,
4656
+ 1,
4657
+ 1
4658
+ ],
4659
+ "range_warp_specializes": [],
4660
+ "range_num_stages": [],
4661
+ "range_multi_buffers": [
4662
+ null,
4663
+ true,
4664
+ null
4665
+ ],
4666
+ "range_flattens": [
4667
+ null,
4668
+ null,
4669
+ true
4670
+ ],
4671
+ "load_eviction_policies": [
4672
+ "",
4673
+ "first",
4674
+ "last"
4675
+ ],
4676
+ "num_warps": 16,
4677
+ "num_stages": 5,
4678
+ "indexing": [
4679
+ "tensor_descriptor",
4680
+ "pointer",
4681
+ "tensor_descriptor",
4682
+ "tensor_descriptor",
4683
+ "pointer"
4684
+ ],
4685
+ "atomic_indexing": [],
4686
+ "pid_type": "flat"
4687
+ }
4688
+ },
4689
+ {
4690
+ "key": {
4691
+ "hidden_size": 5120,
4692
+ "num_tokens": 64
4693
+ },
4694
+ "config": {
4695
+ "block_sizes": [
4696
+ 8192,
4697
+ 2048
4698
+ ],
4699
+ "range_unroll_factors": [
4700
+ 0,
4701
+ 2,
4702
+ 1
4703
+ ],
4704
+ "range_warp_specializes": [],
4705
+ "range_num_stages": [],
4706
+ "range_multi_buffers": [
4707
+ null,
4708
+ true,
4709
+ false
4710
+ ],
4711
+ "range_flattens": [
4712
+ null,
4713
+ false,
4714
+ false
4715
+ ],
4716
+ "load_eviction_policies": [
4717
+ "",
4718
+ "last",
4719
+ "last"
4720
+ ],
4721
+ "num_warps": 16,
4722
+ "num_stages": 8,
4723
+ "indexing": [
4724
+ "pointer",
4725
+ "tensor_descriptor",
4726
+ "pointer",
4727
+ "tensor_descriptor",
4728
+ "pointer"
4729
+ ],
4730
+ "atomic_indexing": [],
4731
+ "pid_type": "flat"
4732
+ }
4733
+ },
4734
+ {
4735
+ "key": {
4736
+ "hidden_size": 5120,
4737
+ "num_tokens": 128
4738
+ },
4739
+ "config": {
4740
+ "block_sizes": [
4741
+ 8192,
4742
+ 1024
4743
+ ],
4744
+ "range_unroll_factors": [
4745
+ 0,
4746
+ 1,
4747
+ 1
4748
+ ],
4749
+ "range_warp_specializes": [],
4750
+ "range_num_stages": [],
4751
+ "range_multi_buffers": [
4752
+ null,
4753
+ true,
4754
+ null
4755
+ ],
4756
+ "range_flattens": [
4757
+ null,
4758
+ null,
4759
+ true
4760
+ ],
4761
+ "load_eviction_policies": [
4762
+ "",
4763
+ "first",
4764
+ "last"
4765
+ ],
4766
+ "num_warps": 16,
4767
+ "num_stages": 5,
4768
+ "indexing": [
4769
+ "tensor_descriptor",
4770
+ "pointer",
4771
+ "tensor_descriptor",
4772
+ "tensor_descriptor",
4773
+ "pointer"
4774
+ ],
4775
+ "atomic_indexing": [],
4776
+ "pid_type": "flat"
4777
+ }
4778
+ },
4779
+ {
4780
+ "key": {
4781
+ "hidden_size": 5120,
4782
+ "num_tokens": 256
4783
+ },
4784
+ "config": {
4785
+ "block_sizes": [
4786
+ 8192,
4787
+ 1024
4788
+ ],
4789
+ "range_unroll_factors": [
4790
+ 0,
4791
+ 1,
4792
+ 3
4793
+ ],
4794
+ "range_warp_specializes": [],
4795
+ "range_num_stages": [],
4796
+ "range_multi_buffers": [
4797
+ null,
4798
+ false,
4799
+ false
4800
+ ],
4801
+ "range_flattens": [
4802
+ null,
4803
+ true,
4804
+ true
4805
+ ],
4806
+ "load_eviction_policies": [
4807
+ "first",
4808
+ "first",
4809
+ "last"
4810
+ ],
4811
+ "num_warps": 8,
4812
+ "num_stages": 2,
4813
+ "indexing": [
4814
+ "pointer",
4815
+ "tensor_descriptor",
4816
+ "tensor_descriptor",
4817
+ "pointer",
4818
+ "pointer"
4819
+ ],
4820
+ "atomic_indexing": [],
4821
+ "pid_type": "flat"
4822
+ }
4823
+ },
4824
+ {
4825
+ "key": {
4826
+ "hidden_size": 5120,
4827
+ "num_tokens": 512
4828
+ },
4829
+ "config": {
4830
+ "block_sizes": [
4831
+ 8192,
4832
+ 4096
4833
+ ],
4834
+ "range_unroll_factors": [
4835
+ 0,
4836
+ 0,
4837
+ 2
4838
+ ],
4839
+ "range_warp_specializes": [],
4840
+ "range_num_stages": [],
4841
+ "range_multi_buffers": [
4842
+ null,
4843
+ false,
4844
+ false
4845
+ ],
4846
+ "range_flattens": [
4847
+ null,
4848
+ true,
4849
+ null
4850
+ ],
4851
+ "load_eviction_policies": [
4852
+ "",
4853
+ "last",
4854
+ ""
4855
+ ],
4856
+ "num_warps": 2,
4857
+ "num_stages": 1,
4858
+ "indexing": [
4859
+ "pointer",
4860
+ "pointer",
4861
+ "pointer",
4862
+ "tensor_descriptor",
4863
+ "pointer"
4864
+ ],
4865
+ "atomic_indexing": [],
4866
+ "pid_type": "flat"
4867
+ }
4868
+ },
4869
+ {
4870
+ "key": {
4871
+ "hidden_size": 5120,
4872
+ "num_tokens": 1024
4873
+ },
4874
+ "config": {
4875
+ "block_sizes": [
4876
+ 8192,
4877
+ 8192
4878
+ ],
4879
+ "range_unroll_factors": [
4880
+ 0,
4881
+ 3,
4882
+ 0
4883
+ ],
4884
+ "range_warp_specializes": [],
4885
+ "range_num_stages": [],
4886
+ "range_multi_buffers": [
4887
+ null,
4888
+ true,
4889
+ true
4890
+ ],
4891
+ "range_flattens": [
4892
+ null,
4893
+ true,
4894
+ true
4895
+ ],
4896
+ "load_eviction_policies": [
4897
+ "",
4898
+ "last",
4899
+ "first"
4900
+ ],
4901
+ "num_warps": 2,
4902
+ "num_stages": 8,
4903
+ "indexing": [
4904
+ "tensor_descriptor",
4905
+ "pointer",
4906
+ "pointer",
4907
+ "tensor_descriptor",
4908
+ "tensor_descriptor"
4909
+ ],
4910
+ "atomic_indexing": [],
4911
+ "pid_type": "flat"
4912
+ }
4913
+ },
4914
+ {
4915
+ "key": {
4916
+ "hidden_size": 5120,
4917
+ "num_tokens": 2048
4918
+ },
4919
+ "config": {
4920
+ "block_sizes": [
4921
+ 8192,
4922
+ 8192
4923
+ ],
4924
+ "range_unroll_factors": [
4925
+ 0,
4926
+ 3,
4927
+ 0
4928
+ ],
4929
+ "range_warp_specializes": [],
4930
+ "range_num_stages": [],
4931
+ "range_multi_buffers": [
4932
+ null,
4933
+ true,
4934
+ true
4935
+ ],
4936
+ "range_flattens": [
4937
+ null,
4938
+ true,
4939
+ true
4940
+ ],
4941
+ "load_eviction_policies": [
4942
+ "",
4943
+ "last",
4944
+ "first"
4945
+ ],
4946
+ "num_warps": 2,
4947
+ "num_stages": 8,
4948
+ "indexing": [
4949
+ "tensor_descriptor",
4950
+ "pointer",
4951
+ "pointer",
4952
+ "tensor_descriptor",
4953
+ "tensor_descriptor"
4954
+ ],
4955
+ "atomic_indexing": [],
4956
+ "pid_type": "flat"
4957
+ }
4958
+ },
4959
+ {
4960
+ "key": {
4961
+ "hidden_size": 5120,
4962
+ "num_tokens": 4096
4963
+ },
4964
+ "config": {
4965
+ "block_sizes": [
4966
+ 8192,
4967
+ 1024
4968
+ ],
4969
+ "range_unroll_factors": [
4970
+ 3,
4971
+ 4,
4972
+ 2
4973
+ ],
4974
+ "range_warp_specializes": [],
4975
+ "range_multi_buffers": [
4976
+ true,
4977
+ false,
4978
+ true
4979
+ ],
4980
+ "range_flattens": [
4981
+ false,
4982
+ true,
4983
+ true
4984
+ ],
4985
+ "load_eviction_policies": [
4986
+ "",
4987
+ "",
4988
+ "first"
4989
+ ],
4990
+ "num_warps": 4,
4991
+ "num_stages": 6,
4992
+ "indexing": [
4993
+ "pointer",
4994
+ "pointer",
4995
+ "pointer",
4996
+ "pointer",
4997
+ "pointer"
4998
+ ],
4999
+ "atomic_indexing": [],
This file is too large to show in full.
vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py
new
+165
@@ -0,0 +1,165 @@
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.kernels.helion.register import register_kernel
11
+from vllm.logger import init_logger
12
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ get_fp8_min_max,
14
+)
15
+from vllm.platforms import current_platform
16
+from vllm.utils.import_utils import has_helion
17
+
18
+if not has_helion():
19
+ raise ImportError(
20
+ "Helion kernel requires helion to be installed. "
21
+ "Install it with: pip install helion"
22
+ )
23
+
24
+import helion
25
+import helion.language as hl
26
+
27
+logger = init_logger(__name__)
28
+
29
+
30
+def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]:
31
+ # TODO(xiaohongchen1991): it is difficult for kernel author to cover all input
32
+ # property combination. Currently, dtypes are fixed. We need optimization to
33
+ # bucket/skip some combinations
34
+ num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
35
+ hidden_size_list = [2048, 4096, 5120]
36
+ in_dtype: torch.dtype = torch.bfloat16
37
+ out_dtype: torch.dtype = current_platform.fp8_dtype()
38
+ scale_dtype: torch.dtype = torch.float32
39
+ inputs = {}
40
+ for num_tokens, hidden_size in product(num_tokens_list, hidden_size_list):
41
+ input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype)
42
+ result = torch.empty(input.shape, device=input.device, dtype=out_dtype)
43
+ scale = torch.empty((num_tokens, 1), device=input.device, dtype=scale_dtype)
44
+ scale_ub = torch.mean(input).to(scale_dtype)
45
+
46
+ config_key = CaseKey({"hidden_size": hidden_size, "num_tokens": num_tokens})
47
+ inputs[config_key] = (result, input, scale, scale_ub)
48
+
49
+ return inputs
50
+
51
+
52
+_pick_cache: dict[tuple[int, int], CaseKey | None] = {}
53
+
54
+
55
+def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None:
56
+ """Pick the best pre-tuned config for the given input shape.
57
+
58
+ Selection strategy:
59
+ 1. Find the closest hidden_size among available configs
60
+ (exact match preferred).
61
+ 2. Among the num_tokens values tuned for that hidden_size, pick
62
+ the smallest num_tokens >= the input's num_tokens. If the input is
63
+ larger than all available num_tokens, fall back to the largest.
64
+ """
65
+
66
+ if not config_keys:
67
+ return None
68
+
69
+ _, input, *_ = args
70
+ num_tokens, hidden_size = input.shape
71
+
72
+ cache_key = (num_tokens, hidden_size)
73
+ cached = _pick_cache.get(cache_key)
74
+ if cached is not None:
75
+ return cached
76
+
77
+ configs: dict[int, list[int]] = {}
78
+ for key in config_keys:
79
+ if key.is_default():
80
+ continue
81
+ configs.setdefault(key["hidden_size"], []).append(key["num_tokens"])
82
+
83
+ if not configs:
84
+ return None
85
+
86
+ best_hidden_size = min(configs, key=lambda s: abs(s - hidden_size))
87
+ available_num_tokens = sorted(configs[best_hidden_size])
88
+ best_num_tokens = next(
89
+ (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1]
90
+ )
91
+
92
+ result = CaseKey({"hidden_size": best_hidden_size, "num_tokens": best_num_tokens})
93
+ _pick_cache[cache_key] = result
94
+ return result
95
+
96
+
97
+def fake_impl(
98
+ result: torch.Tensor, # [num_tokens, hidden_size]
99
+ input: torch.Tensor, # [num_tokens, hidden_size]
100
+ scale: torch.Tensor, # [num_tokens, 1]
101
+ scale_ub: torch.Tensor | None = None, # scalar tensor
102
+) -> None:
103
+ return
104
+
105
+
106
+def baseline(
107
+ result: torch.Tensor, # [num_tokens, hidden_size]
108
+ input: torch.Tensor, # [num_tokens, hidden_size]
109
+ scale: torch.Tensor, # [num_tokens, 1]
110
+ scale_ub: torch.Tensor | None = None, # scalar tensor
111
+) -> None:
112
+ torch.ops._C.dynamic_per_token_scaled_fp8_quant(result, input, scale, scale_ub)
113
+
114
+
115
+# Overwrite autotune_baseline_atol and autotune_baseline_rtol
116
+# if too many configs failed due to baseline check during autotuning
117
+@register_kernel(
118
+ mutates_args=["result", "scale"],
119
+ config_picker=pick_config,
120
+ input_generator=generate_inputs,
121
+ fake_impl=fake_impl,
122
+ helion_settings=helion.Settings(
123
+ autotune_baseline_fn=baseline,
124
+ ignore_warnings=[helion.exc.TensorOperationInWrapper],
125
+ ),
126
+)
127
+def dynamic_per_token_scaled_fp8_quant(
128
+ result: torch.Tensor, # [num_tokens, hidden_size]
129
+ input: torch.Tensor, # [num_tokens, hidden_size]
130
+ scale: torch.Tensor, # [num_tokens, 1]
131
+ scale_ub: torch.Tensor | None = None, # scalar tensor
132
+) -> None:
133
+ # This code assumes batch_dim and num_tokens are flattened
134
+ assert input.ndim == 2
135
+ num_tokens, hidden_size = input.shape
136
+ hl.specialize(hidden_size)
137
+
138
+ assert result.shape == input.shape
139
+ assert scale.shape[0] == num_tokens
140
+ assert scale.dtype == torch.float32
141
+ assert input.stride()[-1] == 1
142
+ assert result.stride()[-1] == 1
143
+
144
+ fp8_min, fp8_max = get_fp8_min_max()
145
+ min_scaling_factor = 1.0 / (fp8_max * 512.0)
146
+
147
+ for tile_m in hl.tile(num_tokens, block_size=1):
148
+ s_blk = hl.zeros([tile_m], dtype=torch.float32)
149
+ for tile_n in hl.tile(hidden_size):
150
+ x_blk = input[tile_m, tile_n].to(dtype=torch.float32)
151
+ tmp_blk = torch.amax(torch.abs(x_blk), dim=-1)
152
+ s_blk = torch.maximum(s_blk, tmp_blk)
153
+
154
+ if scale_ub is not None:
155
+ scale_ub_s = hl.load(scale_ub, [])
156
+ s_blk = s_blk.clamp(max=scale_ub_s)
157
+ s_blk = s_blk * (1.0 / fp8_max)
158
+ s_blk = s_blk.clamp(min=min_scaling_factor)
159
+ scale[tile_m, 0] = s_blk
160
+
161
+ for tile_n in hl.tile(hidden_size):
162
+ x_blk = input[tile_m, tile_n].to(torch.float32)
163
+ y_blk = x_blk * (1.0 / s_blk[:, None])
164
+
165
+ result[tile_m, tile_n] = y_blk.clamp(fp8_min, fp8_max).to(result.dtype)